Jackie-Jiang commented on a change in pull request #8426: URL: https://github.com/apache/pinot/pull/8426#discussion_r836885719
########## File path: pinot-core/src/test/java/org/apache/pinot/queries/JsonIngestionFromAvroQueriesTest.java ########## @@ -0,0 +1,305 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.queries; + +import com.google.common.collect.Lists; +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.IntermediateResultsBlock; +import org.apache.pinot.plugin.inputformat.avro.AvroRecordReader; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.utils.Pair; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.apache.avro.Schema.*; + + +/** + * Test if ComplexType (RECORD, ARRAY, MAP, and UNION) field from an AVRO file can be ingested into a JSON column in + * a Pinot segment. This class tests. Ingestion from ENUM (symbol) and FIXED (binary) is not supported. Review comment: This class tests ... ########## File path: pinot-common/src/main/java/org/apache/pinot/common/utils/PinotDataType.java ########## @@ -829,18 +829,23 @@ public String toString(Object value) { public String toJson(Object value) { if (value instanceof String) { try { + // Try to parse the string as JSON first return JsonUtils.stringToJsonNode((String) value).toString(); + } catch (com.fasterxml.jackson.core.JsonParseException jpe) { Review comment: Import this class ########## File path: pinot-core/src/test/java/org/apache/pinot/queries/JsonIngestionFromAvroQueriesTest.java ########## @@ -0,0 +1,305 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.queries; + +import com.google.common.collect.Lists; +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.IntermediateResultsBlock; +import org.apache.pinot.plugin.inputformat.avro.AvroRecordReader; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.utils.Pair; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.apache.avro.Schema.*; + + +/** + * Test if ComplexType (RECORD, ARRAY, MAP, and UNION) field from an AVRO file can be ingested into a JSON column in + * a Pinot segment. This class tests. Ingestion from ENUM (symbol) and FIXED (binary) is not supported. Review comment: IIRC, `ENUM` is already extracted as `STRING` data type; `FIXED` can be extracted as `BYTES` data type. Since we are working on this, I'd suggest support all the avro data types. ########## File path: pinot-core/src/test/java/org/apache/pinot/queries/JsonIngestionFromAvroQueriesTest.java ########## @@ -0,0 +1,305 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.queries; + +import com.google.common.collect.Lists; +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.IntermediateResultsBlock; +import org.apache.pinot.plugin.inputformat.avro.AvroRecordReader; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.utils.Pair; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.apache.avro.Schema.*; Review comment: Seems there are more usage of avro Schema than pinot Schema. Consider import avro Schema and use full path for pinot Schema ########## File path: pinot-core/src/test/java/org/apache/pinot/queries/JsonIngestionFromAvroQueriesTest.java ########## @@ -0,0 +1,305 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.queries; + +import com.google.common.collect.Lists; +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.IntermediateResultsBlock; +import org.apache.pinot.plugin.inputformat.avro.AvroRecordReader; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.utils.Pair; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.apache.avro.Schema.*; + + +/** + * Test if ComplexType (RECORD, ARRAY, MAP, and UNION) field from an AVRO file can be ingested into a JSON column in + * a Pinot segment. This class tests. Ingestion from ENUM (symbol) and FIXED (binary) is not supported. + */ +public class JsonIngestionFromAvroQueriesTest extends BaseQueriesTest { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "JsonIngestionFromAvroTest"); + private static final File AVRO_DATA_FILE = new File(INDEX_DIR, "JsonIngestionFromAvroTest.avro"); + private static final String RAW_TABLE_NAME = "testTable"; + private static final String SEGMENT_NAME = "testSegment"; + + private static final String INT_COLUMN = "intColumn"; + private static final String JSON_COLUMN = "jsonColumn"; + private static final String STRING_COLUMN = "stringColumn"; + private static final Schema SCHEMA = + new Schema.SchemaBuilder().addSingleValueDimension(INT_COLUMN, FieldSpec.DataType.INT) + .addSingleValueDimension(JSON_COLUMN, FieldSpec.DataType.JSON) + .addSingleValueDimension(STRING_COLUMN, FieldSpec.DataType.STRING).build(); + private static final TableConfig TABLE_CONFIG = + new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build(); + + private IndexSegment _indexSegment; + private List<IndexSegment> _indexSegments; + + @Override + protected String getFilter() { + return ""; + } + + @Override + protected IndexSegment getIndexSegment() { + return _indexSegment; + } + + @Override + protected List<IndexSegment> getIndexSegments() { + return _indexSegments; + } + + /** @return {@link GenericRow} representing a row in Pinot table. */ + private static GenericRow createTableRecord(int intValue, String stringValue, Object jsonValue) { + GenericRow record = new GenericRow(); + record.putValue(INT_COLUMN, intValue); + record.putValue(STRING_COLUMN, stringValue); + record.putValue(JSON_COLUMN, jsonValue); + + return record; + } + + private static Map<String, String> createMapField(Pair<String, String>[] pairs) { + Map<String, String> map = new HashMap<>(); + for (Pair<String, String> pair : pairs) { + map.put(pair.getFirst(), pair.getSecond()); + } + return map; + } + + private static org.apache.avro.Schema createRecordSchema() { + List<org.apache.avro.Schema.Field> fields = new ArrayList<>(); + fields.add(new org.apache.avro.Schema.Field("id", create(Type.INT))); + fields.add(new org.apache.avro.Schema.Field("name", create(Type.STRING))); + return createRecord("record", "doc", JsonIngestionFromAvroQueriesTest.class.getCanonicalName(), false, fields); + } + + private static GenericData.Record createRecordField(String k1, int v1, String k2, String v2) { + GenericData.Record record = new GenericData.Record(createRecordSchema()); + record.put(k1, v1); + record.put(k2, v2); + return record; + } + + private static void createInputFile() + throws IOException { + INDEX_DIR.mkdir(); + org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("eventsRecord", null, null, false); + List<Field> fields = Arrays + .asList(new Field(INT_COLUMN, createUnion(Lists.newArrayList(create(Type.INT), create(Type.NULL))), null, null), + new Field(STRING_COLUMN, createUnion(Lists.newArrayList(create(Type.STRING), create(Type.NULL))), null, + null), new Field(JSON_COLUMN, + createUnion(createArray(create(Type.STRING)), createMap(create(Type.STRING)), createRecordSchema(), + create(Type.STRING), create(Type.NULL)))); + avroSchema.setFields(fields); + List<GenericRow> inputRecords = new ArrayList<>(); + // Insert ARRAY + inputRecords.add(createTableRecord(1, "daffy duck", Arrays.asList("this", "is", "a", "test"))); + + // Insert MAP + inputRecords + .add(createTableRecord(2, "mickey mouse", createMapField(new Pair[]{new Pair("a", "1"), new Pair("b", "2")}))); + inputRecords + .add(createTableRecord(3, "donald duck", createMapField(new Pair[]{new Pair("a", "1"), new Pair("b", "2")}))); + inputRecords.add( + createTableRecord(4, "scrooge mcduck", createMapField(new Pair[]{new Pair("a", "1"), new Pair("b", "2")}))); + + // insert RECORD + inputRecords.add(createTableRecord(5, "minney mouse", createRecordField("id", 1, "name", "minney"))); + + // Insert simple Java String (gets converted into JSON value) + inputRecords.add(createTableRecord(6, "pluto", "test")); + + // Insert JSON string (gets converted into JSON document) + inputRecords.add(createTableRecord(7, "scooby doo", "{\"name\":\"scooby\",\"id\":7}")); + + try (DataFileWriter<GenericData.Record> fileWriter = new DataFileWriter<>(new GenericDatumWriter<>(avroSchema))) { + fileWriter.create(avroSchema, AVRO_DATA_FILE); + for (GenericRow inputRecord : inputRecords) { + GenericData.Record record = new GenericData.Record(avroSchema); + record.put(INT_COLUMN, inputRecord.getValue(INT_COLUMN)); + record.put(STRING_COLUMN, inputRecord.getValue(STRING_COLUMN)); + record.put(JSON_COLUMN, inputRecord.getValue(JSON_COLUMN)); + fileWriter.append(record); + } + } + } + + private static RecordReader createRecordReader() + throws IOException { + Set<String> set = new HashSet<>(); + set.add(INT_COLUMN); + set.add(STRING_COLUMN); + set.add(JSON_COLUMN); + AvroRecordReader avroRecordReader = new AvroRecordReader(); + avroRecordReader.init(AVRO_DATA_FILE, set, null); + return avroRecordReader; + } + + /** Create an AVRO file and then ingest it into Pinot while creating a JsonIndex. */ + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteDirectory(INDEX_DIR); + createInputFile(); + + List<String> jsonIndexColumns = new ArrayList<>(); + jsonIndexColumns.add("jsonColumn"); + TABLE_CONFIG.getIndexingConfig().setJsonIndexColumns(jsonIndexColumns); + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + segmentGeneratorConfig.setTableName(RAW_TABLE_NAME); + segmentGeneratorConfig.setSegmentName(SEGMENT_NAME); + segmentGeneratorConfig.setOutDir(INDEX_DIR.getPath()); + segmentGeneratorConfig.setInputFilePath(AVRO_DATA_FILE.getPath()); + + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, createRecordReader()); + driver.build(); + + IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(); + indexLoadingConfig.setTableConfig(TABLE_CONFIG); + indexLoadingConfig.setJsonIndexColumns(new HashSet<String>(jsonIndexColumns)); + indexLoadingConfig.setReadMode(ReadMode.mmap); + + ImmutableSegment immutableSegment = + ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), indexLoadingConfig); + _indexSegment = immutableSegment; + _indexSegments = Arrays.asList(immutableSegment, immutableSegment); + } + + /** Verify that we can query the JSON column that ingested ComplexType data from an AVRO file (see setUp). */ + @Test + public void testSimpleSelectOnJsonColumn() { + try { + Operator operator = getOperatorForSqlQuery("select intColumn, stringColumn, jsonColumn FROM testTable limit 100"); + IntermediateResultsBlock block = (IntermediateResultsBlock) operator.nextBlock(); + Collection<Object[]> rows = block.getSelectionResult(); + Assert.assertEquals(block.getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.INT); + Assert.assertEquals(block.getDataSchema().getColumnDataType(1), DataSchema.ColumnDataType.STRING); + Assert.assertEquals(block.getDataSchema().getColumnDataType(2), DataSchema.ColumnDataType.JSON); + + List<String> expecteds = Arrays + .asList("[1, daffy duck, [\"this\",\"is\",\"a\",\"test\"]]", "[2, mickey mouse, {\"a\":\"1\",\"b\":\"2\"}]", + "[3, donald duck, {\"a\":\"1\",\"b\":\"2\"}]", "[4, scrooge mcduck, {\"a\":\"1\",\"b\":\"2\"}]", + "[5, minney mouse, {\"name\":\"minney\",\"id\":1}]", "[6, pluto, \"test\"]", + "[7, scooby doo, {\"name\":\"scooby\",\"id\":7}]"); + int index = 0; + + Iterator<Object[]> iterator = rows.iterator(); + while (iterator.hasNext()) { + Object[] row = iterator.next(); + System.out.println(Arrays.toString(row)); + Assert.assertEquals(Arrays.toString(row), expecteds.get(index++)); + } + } catch (IllegalStateException ise) { + Assert.assertTrue(true); + } + } + + /** Verify simple path expression query on ingested Avro file. */ + @Test + public void testJsonPathSelectOnJsonColumn() { + try { + Operator operator = getOperatorForSqlQuery("select intColumn, json_extract_scalar(jsonColumn, '$.name', " + + "'STRING', 'null') FROM testTable"); + IntermediateResultsBlock block = (IntermediateResultsBlock) operator.nextBlock(); + Collection<Object[]> rows = block.getSelectionResult(); + Assert.assertEquals(block.getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.INT); + Assert.assertEquals(block.getDataSchema().getColumnDataType(1), DataSchema.ColumnDataType.STRING); + + List<String> expecteds = + Arrays.asList("[1, null]", "[2, null]", "[3, null]", "[4, null]", "[5, minney]", "[6, null]", "[7, scooby]"); + int index = 0; + + Iterator<Object[]> iterator = rows.iterator(); + while (iterator.hasNext()) { + Object[] row = iterator.next(); + System.out.println(Arrays.toString(row)); + Assert.assertEquals(Arrays.toString(row), expecteds.get(index++)); + } + } catch (IllegalStateException ise) { + Assert.assertTrue(true); Review comment: This assert is an no-op. Do we expect exception here? ########## File path: pinot-core/src/test/java/org/apache/pinot/queries/JsonIngestionFromAvroQueriesTest.java ########## @@ -0,0 +1,305 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.queries; + +import com.google.common.collect.Lists; +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.IntermediateResultsBlock; +import org.apache.pinot.plugin.inputformat.avro.AvroRecordReader; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.utils.Pair; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.apache.avro.Schema.*; + + +/** + * Test if ComplexType (RECORD, ARRAY, MAP, and UNION) field from an AVRO file can be ingested into a JSON column in + * a Pinot segment. This class tests. Ingestion from ENUM (symbol) and FIXED (binary) is not supported. + */ +public class JsonIngestionFromAvroQueriesTest extends BaseQueriesTest { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "JsonIngestionFromAvroTest"); + private static final File AVRO_DATA_FILE = new File(INDEX_DIR, "JsonIngestionFromAvroTest.avro"); + private static final String RAW_TABLE_NAME = "testTable"; + private static final String SEGMENT_NAME = "testSegment"; + + private static final String INT_COLUMN = "intColumn"; + private static final String JSON_COLUMN = "jsonColumn"; + private static final String STRING_COLUMN = "stringColumn"; + private static final Schema SCHEMA = + new Schema.SchemaBuilder().addSingleValueDimension(INT_COLUMN, FieldSpec.DataType.INT) + .addSingleValueDimension(JSON_COLUMN, FieldSpec.DataType.JSON) + .addSingleValueDimension(STRING_COLUMN, FieldSpec.DataType.STRING).build(); + private static final TableConfig TABLE_CONFIG = + new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build(); + + private IndexSegment _indexSegment; + private List<IndexSegment> _indexSegments; + + @Override + protected String getFilter() { + return ""; + } + + @Override + protected IndexSegment getIndexSegment() { + return _indexSegment; + } + + @Override + protected List<IndexSegment> getIndexSegments() { + return _indexSegments; + } + + /** @return {@link GenericRow} representing a row in Pinot table. */ + private static GenericRow createTableRecord(int intValue, String stringValue, Object jsonValue) { + GenericRow record = new GenericRow(); + record.putValue(INT_COLUMN, intValue); + record.putValue(STRING_COLUMN, stringValue); + record.putValue(JSON_COLUMN, jsonValue); + + return record; + } + + private static Map<String, String> createMapField(Pair<String, String>[] pairs) { + Map<String, String> map = new HashMap<>(); + for (Pair<String, String> pair : pairs) { + map.put(pair.getFirst(), pair.getSecond()); + } + return map; + } + + private static org.apache.avro.Schema createRecordSchema() { + List<org.apache.avro.Schema.Field> fields = new ArrayList<>(); + fields.add(new org.apache.avro.Schema.Field("id", create(Type.INT))); + fields.add(new org.apache.avro.Schema.Field("name", create(Type.STRING))); + return createRecord("record", "doc", JsonIngestionFromAvroQueriesTest.class.getCanonicalName(), false, fields); + } + + private static GenericData.Record createRecordField(String k1, int v1, String k2, String v2) { + GenericData.Record record = new GenericData.Record(createRecordSchema()); + record.put(k1, v1); + record.put(k2, v2); + return record; + } + + private static void createInputFile() + throws IOException { + INDEX_DIR.mkdir(); + org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("eventsRecord", null, null, false); + List<Field> fields = Arrays + .asList(new Field(INT_COLUMN, createUnion(Lists.newArrayList(create(Type.INT), create(Type.NULL))), null, null), + new Field(STRING_COLUMN, createUnion(Lists.newArrayList(create(Type.STRING), create(Type.NULL))), null, + null), new Field(JSON_COLUMN, + createUnion(createArray(create(Type.STRING)), createMap(create(Type.STRING)), createRecordSchema(), + create(Type.STRING), create(Type.NULL)))); + avroSchema.setFields(fields); + List<GenericRow> inputRecords = new ArrayList<>(); + // Insert ARRAY + inputRecords.add(createTableRecord(1, "daffy duck", Arrays.asList("this", "is", "a", "test"))); + + // Insert MAP + inputRecords + .add(createTableRecord(2, "mickey mouse", createMapField(new Pair[]{new Pair("a", "1"), new Pair("b", "2")}))); + inputRecords + .add(createTableRecord(3, "donald duck", createMapField(new Pair[]{new Pair("a", "1"), new Pair("b", "2")}))); + inputRecords.add( + createTableRecord(4, "scrooge mcduck", createMapField(new Pair[]{new Pair("a", "1"), new Pair("b", "2")}))); + + // insert RECORD + inputRecords.add(createTableRecord(5, "minney mouse", createRecordField("id", 1, "name", "minney"))); + + // Insert simple Java String (gets converted into JSON value) + inputRecords.add(createTableRecord(6, "pluto", "test")); + + // Insert JSON string (gets converted into JSON document) + inputRecords.add(createTableRecord(7, "scooby doo", "{\"name\":\"scooby\",\"id\":7}")); + + try (DataFileWriter<GenericData.Record> fileWriter = new DataFileWriter<>(new GenericDatumWriter<>(avroSchema))) { + fileWriter.create(avroSchema, AVRO_DATA_FILE); + for (GenericRow inputRecord : inputRecords) { + GenericData.Record record = new GenericData.Record(avroSchema); + record.put(INT_COLUMN, inputRecord.getValue(INT_COLUMN)); + record.put(STRING_COLUMN, inputRecord.getValue(STRING_COLUMN)); + record.put(JSON_COLUMN, inputRecord.getValue(JSON_COLUMN)); + fileWriter.append(record); + } + } + } + + private static RecordReader createRecordReader() + throws IOException { + Set<String> set = new HashSet<>(); + set.add(INT_COLUMN); + set.add(STRING_COLUMN); + set.add(JSON_COLUMN); + AvroRecordReader avroRecordReader = new AvroRecordReader(); + avroRecordReader.init(AVRO_DATA_FILE, set, null); + return avroRecordReader; + } + + /** Create an AVRO file and then ingest it into Pinot while creating a JsonIndex. */ + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteDirectory(INDEX_DIR); + createInputFile(); + + List<String> jsonIndexColumns = new ArrayList<>(); + jsonIndexColumns.add("jsonColumn"); + TABLE_CONFIG.getIndexingConfig().setJsonIndexColumns(jsonIndexColumns); + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + segmentGeneratorConfig.setTableName(RAW_TABLE_NAME); + segmentGeneratorConfig.setSegmentName(SEGMENT_NAME); + segmentGeneratorConfig.setOutDir(INDEX_DIR.getPath()); + segmentGeneratorConfig.setInputFilePath(AVRO_DATA_FILE.getPath()); + + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, createRecordReader()); + driver.build(); + + IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(); + indexLoadingConfig.setTableConfig(TABLE_CONFIG); + indexLoadingConfig.setJsonIndexColumns(new HashSet<String>(jsonIndexColumns)); + indexLoadingConfig.setReadMode(ReadMode.mmap); + + ImmutableSegment immutableSegment = + ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), indexLoadingConfig); + _indexSegment = immutableSegment; + _indexSegments = Arrays.asList(immutableSegment, immutableSegment); + } + + /** Verify that we can query the JSON column that ingested ComplexType data from an AVRO file (see setUp). */ + @Test + public void testSimpleSelectOnJsonColumn() { + try { + Operator operator = getOperatorForSqlQuery("select intColumn, stringColumn, jsonColumn FROM testTable limit 100"); + IntermediateResultsBlock block = (IntermediateResultsBlock) operator.nextBlock(); + Collection<Object[]> rows = block.getSelectionResult(); + Assert.assertEquals(block.getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.INT); + Assert.assertEquals(block.getDataSchema().getColumnDataType(1), DataSchema.ColumnDataType.STRING); + Assert.assertEquals(block.getDataSchema().getColumnDataType(2), DataSchema.ColumnDataType.JSON); + + List<String> expecteds = Arrays + .asList("[1, daffy duck, [\"this\",\"is\",\"a\",\"test\"]]", "[2, mickey mouse, {\"a\":\"1\",\"b\":\"2\"}]", + "[3, donald duck, {\"a\":\"1\",\"b\":\"2\"}]", "[4, scrooge mcduck, {\"a\":\"1\",\"b\":\"2\"}]", + "[5, minney mouse, {\"name\":\"minney\",\"id\":1}]", "[6, pluto, \"test\"]", + "[7, scooby doo, {\"name\":\"scooby\",\"id\":7}]"); + int index = 0; + + Iterator<Object[]> iterator = rows.iterator(); + while (iterator.hasNext()) { + Object[] row = iterator.next(); + System.out.println(Arrays.toString(row)); + Assert.assertEquals(Arrays.toString(row), expecteds.get(index++)); + } + } catch (IllegalStateException ise) { + Assert.assertTrue(true); + } + } + + /** Verify simple path expression query on ingested Avro file. */ + @Test + public void testJsonPathSelectOnJsonColumn() { + try { + Operator operator = getOperatorForSqlQuery("select intColumn, json_extract_scalar(jsonColumn, '$.name', " + + "'STRING', 'null') FROM testTable"); + IntermediateResultsBlock block = (IntermediateResultsBlock) operator.nextBlock(); + Collection<Object[]> rows = block.getSelectionResult(); + Assert.assertEquals(block.getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.INT); + Assert.assertEquals(block.getDataSchema().getColumnDataType(1), DataSchema.ColumnDataType.STRING); + + List<String> expecteds = + Arrays.asList("[1, null]", "[2, null]", "[3, null]", "[4, null]", "[5, minney]", "[6, null]", "[7, scooby]"); + int index = 0; + + Iterator<Object[]> iterator = rows.iterator(); + while (iterator.hasNext()) { + Object[] row = iterator.next(); + System.out.println(Arrays.toString(row)); Review comment: Avoid console output ########## File path: pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroIngestionSchemaValidator.java ########## @@ -124,13 +124,16 @@ private void validateSchemas() { "The Pinot column: %s is 'single-value' column but the column: %s from input %s is 'multi-value' column.", columnName, avroColumnName, getInputSchemaType())); } - FieldSpec.DataType dataTypeForSVColumn = AvroUtils.extractFieldDataType(avroColumnField); - // check data type mismatch - if (fieldSpec.getDataType() != dataTypeForSVColumn) { - _dataTypeMismatch.addMismatchReason(String - .format("The Pinot column: (%s: %s) doesn't match with the column (%s: %s) in input %s schema.", - columnName, fieldSpec.getDataType().name(), avroColumnName, avroColumnType.name(), - getInputSchemaType())); + + if (fieldSpec.getDataType() != FieldSpec.DataType.JSON) { Review comment: We should skip the validation completely for JSON pinot type after validating the avro column field existence ########## File path: pinot-core/src/test/java/org/apache/pinot/queries/JsonIngestionFromAvroQueriesTest.java ########## @@ -0,0 +1,305 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.queries; + +import com.google.common.collect.Lists; +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.IntermediateResultsBlock; +import org.apache.pinot.plugin.inputformat.avro.AvroRecordReader; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.utils.Pair; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.apache.avro.Schema.*; + + +/** + * Test if ComplexType (RECORD, ARRAY, MAP, and UNION) field from an AVRO file can be ingested into a JSON column in + * a Pinot segment. This class tests. Ingestion from ENUM (symbol) and FIXED (binary) is not supported. + */ +public class JsonIngestionFromAvroQueriesTest extends BaseQueriesTest { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "JsonIngestionFromAvroTest"); + private static final File AVRO_DATA_FILE = new File(INDEX_DIR, "JsonIngestionFromAvroTest.avro"); + private static final String RAW_TABLE_NAME = "testTable"; + private static final String SEGMENT_NAME = "testSegment"; + + private static final String INT_COLUMN = "intColumn"; + private static final String JSON_COLUMN = "jsonColumn"; + private static final String STRING_COLUMN = "stringColumn"; + private static final Schema SCHEMA = + new Schema.SchemaBuilder().addSingleValueDimension(INT_COLUMN, FieldSpec.DataType.INT) + .addSingleValueDimension(JSON_COLUMN, FieldSpec.DataType.JSON) + .addSingleValueDimension(STRING_COLUMN, FieldSpec.DataType.STRING).build(); + private static final TableConfig TABLE_CONFIG = + new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build(); + + private IndexSegment _indexSegment; + private List<IndexSegment> _indexSegments; + + @Override + protected String getFilter() { + return ""; + } + + @Override + protected IndexSegment getIndexSegment() { + return _indexSegment; + } + + @Override + protected List<IndexSegment> getIndexSegments() { + return _indexSegments; + } + + /** @return {@link GenericRow} representing a row in Pinot table. */ + private static GenericRow createTableRecord(int intValue, String stringValue, Object jsonValue) { + GenericRow record = new GenericRow(); + record.putValue(INT_COLUMN, intValue); + record.putValue(STRING_COLUMN, stringValue); + record.putValue(JSON_COLUMN, jsonValue); + + return record; + } + + private static Map<String, String> createMapField(Pair<String, String>[] pairs) { + Map<String, String> map = new HashMap<>(); + for (Pair<String, String> pair : pairs) { + map.put(pair.getFirst(), pair.getSecond()); + } + return map; + } + + private static org.apache.avro.Schema createRecordSchema() { + List<org.apache.avro.Schema.Field> fields = new ArrayList<>(); Review comment: This can be simplified ```suggestion List<Field> fields = new ArrayList<>(); ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org
