pinot-core/src/test/java/org/apache/pinot/queries/JsonDataTypeQueriesTest.java - pinot - Git at Google

 /**
  * 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 java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 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.results.SelectionResultsBlock;
 import org.apache.pinot.core.operator.query.GroupByOperator;
 import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
 import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
 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.local.segment.readers.GenericRowRecordReader;
 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.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;


 public class JsonDataTypeQueriesTest extends BaseQueriesTest {
   private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "JsonDataTypeQueriesTest");
   private static final String RAW_TABLE_NAME = "testTable";
   private static final String SEGMENT_NAME = "testSegment";
   private static final int NUM_RECORDS = 10;

   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;
   }

   GenericRow createRecord(int intValue, String stringValue, String jsonValue) {
     GenericRow record = new GenericRow();
     record.putValue(INT_COLUMN, intValue);
     record.putValue(STRING_COLUMN, stringValue);
     record.putValue(JSON_COLUMN, jsonValue);

     return record;
   }

   @Override
   protected List<IndexSegment> getIndexSegments() {
     return _indexSegments;
   }

   @BeforeClass
   public void setUp()
       throws Exception {
     FileUtils.deleteDirectory(INDEX_DIR);

     List<GenericRow> records = new ArrayList<>(NUM_RECORDS);
     records.add(createRecord(1, "daffy duck",
         "{\"name\": {\"first\": \"daffy\", \"last\": \"duck\"}, \"id\": 101, \"data\": [\"a\", \"b\", \"c\", \"d\"]}"));
     records.add(createRecord(2, "mickey mouse",
         "{\"name\": {\"first\": \"mickey\", \"last\": \"mouse\"}, \"id\": 111, \"data\": [\"e\", \"b\", \"c\", "
             + "\"d\"]}"));
     records.add(createRecord(3, "donald duck",
         "{\"name\": {\"first\": \"donald\", \"last\": \"duck\"}, \"id\": 121, \"data\": [\"f\", \"b\", \"c\", "
             + "\"d\"]}"));
     records.add(createRecord(4, "scrooge mcduck",
         "{\"name\": {\"first\": \"scrooge\", \"last\": \"mcduck\"}, \"id\": 131, \"data\": [\"g\", \"b\", \"c\", "
             + "\"d\"]}"));
     records.add(createRecord(5, "minnie mouse",
         "{\"name\": {\"first\": \"minnie\", \"last\": \"mouse\"}, \"id\": 141, \"data\": [\"h\", \"b\", \"c\", "
             + "\"d\"]}"));
     records.add(createRecord(6, "daisy duck",
         "{\"name\": {\"first\": \"daisy\", \"last\": \"duck\"}, \"id\": 161.5, \"data\": [\"i\", \"b\", \"c\", "
             + "\"d\"]}"));
     records.add(createRecord(7, "pluto dog",
         "{\"name\": {\"first\": \"pluto\", \"last\": \"dog\"}, \"id\": 161, \"data\": [\"j\", \"b\", \"c\", \"d\"]}"));
     records.add(createRecord(8, "goofy dwag",
         "{\"name\": {\"first\": \"goofy\", \"last\": \"dwag\"}, \"id\": 171, \"data\": [\"k\", \"b\", \"c\", \"d\"]}"));
     records.add(createRecord(9, "ludwik von drake",
         "{\"name\": {\"first\": \"ludwik\", \"last\": \"von drake\"}, \"id\": 181, \"data\": [\"l\", \"b\", \"c\", "
             + "\"d\"]}"));
     records.add(createRecord(10, "nested array",
         "{\"name\":{\"first\":\"nested\",\"last\":\"array\"},\"id\":111,\"data\":[{\"e\":[{\"x\":[{\"i1\":1,"
             + "\"i2\":2}]},{\"y\":[{\"i1\":1,\"i2\":2}]},{\"z\":[{\"i1\":1,\"i2\":2}]}]},{\"b\":[{\"x\":[{\"i1\":1,"
             + "\"i2\":2}]},{\"y\":[{\"i1\":1,\"i2\":2}]},{\"z\":[{\"i1\":10,\"i2\":20}]}]}]}"));
     records.add(createRecord(11, "multi-dimensional-1 array",
         "{\"name\": {\"first\": \"multi-dimensional-1\",\"last\": \"array\"},\"id\": 111,\"data\": [[[1,2],[3,4]],"
             + "[[\"a\",\"b\"],[\"c\",\"d\"]]]}"));
     records.add(createRecord(12, "multi-dimensional-2 array",
         "{\"name\": {\"first\": \"multi-dimensional-2\",\"last\": \"array\"},\"id\": 111,\"data\": [[[1,2],[3,4]],"
             + "[[\"a\",\"b\"],[\"c\",\"d\"]]]}"));
     records.add(createRecord(13, "multi-dimensional-1 array",
         "{\"name\": {\"first\": \"multi-dimensional-1\",\"last\": \"array\"},\"id\": 111,\"data\": [[[1,2],[3,4]],"
             + "[[\"a\",\"b\"],[\"c\",\"d\"]]]}"));

     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());

     SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
     driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records));
     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 result column type of a simple select query against JSON column */
   @Test
   public void testSimpleSelectOnJsonColumn() {
     try {
       Operator<SelectionResultsBlock> operator = getOperator("select jsonColumn FROM testTable limit 100");
       SelectionResultsBlock block = operator.nextBlock();
       Collection<Object[]> rows = block.getRows();
       Assert.assertEquals(rows.size(), 13);
       Assert.assertEquals(block.getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.JSON);
     } catch (IllegalStateException ise) {
       Assert.assertTrue(true);
     }
   }

   /** Test filtering on string value associated with  JSON key*/
   @Test
   public void testExtractScalarWithStringFilter() {
     Operator<SelectionResultsBlock> operator = getOperator(
         "select intColumn, json_extract_scalar(jsonColumn, '$.name.last', 'STRING') FROM testTable WHERE "
             + "json_extract_scalar(jsonColumn, '$.name.first', 'STRING') = 'daffy'");
     SelectionResultsBlock block = operator.nextBlock();
     Collection<Object[]> rows = block.getRows();
     Assert.assertEquals(rows.size(), 1);

     Iterator<Object[]> iterator = rows.iterator();
     Assert.assertTrue(iterator.hasNext());
     Object[] row = iterator.next();
     Assert.assertEquals(row[0], 1);
     Assert.assertEquals(row[1], "duck");
   }

   /** Test filtering on number value associated with  JSON key*/
   @Test
   public void testExtractScalarWithNumericIntFilter() {
     Operator<SelectionResultsBlock> operator = getOperator(
         "select json_extract_scalar(jsonColumn, '$.name.last', 'STRING') FROM testTable WHERE json_extract_scalar"
             + "(jsonColumn, '$.id', 'INT') = 171");
     SelectionResultsBlock block = operator.nextBlock();
     Collection<Object[]> rows = block.getRows();
     Assert.assertEquals(rows.size(), 1);

     Iterator<Object[]> iterator = rows.iterator();
     Assert.assertTrue(iterator.hasNext());
     Assert.assertEquals(iterator.next()[0], "dwag");
   }

   /** Test filtering on float value associated with  JSON key*/
   @Test
   public void testExtractScalarWithNumericFloatFilter() {

     // query to retrieve result as INT
     Operator<SelectionResultsBlock> operator1 = getOperator(
         "select json_extract_scalar(jsonColumn, '$.id', 'INT') FROM testTable WHERE json_extract_scalar(jsonColumn, "
             + "'$.id', 'FLOAT') = 161.5");
     SelectionResultsBlock block1 = operator1.nextBlock();
     Collection<Object[]> rows1 = block1.getRows();
     Assert.assertEquals(rows1.size(), 1);

     Iterator<Object[]> iterator1 = rows1.iterator();
     Assert.assertTrue(iterator1.hasNext());
     Assert.assertEquals(iterator1.next()[0], 161);

     // query to retrieve result as DOUBLE
     Operator<SelectionResultsBlock> operator2 = getOperator(
         "select json_extract_scalar(jsonColumn, '$.id', 'DOUBLE') FROM testTable WHERE json_extract_scalar"
             + "(jsonColumn, '$.id', 'FLOAT') = 161.5");
     SelectionResultsBlock block2 = operator2.nextBlock();
     Collection<Object[]> rows2 = block2.getRows();
     Assert.assertEquals(rows2.size(), 1);

     Iterator<Object[]> iterator2 = rows2.iterator();
     Assert.assertTrue(iterator2.hasNext());
     Assert.assertEquals(iterator2.next()[0], 161.5d);
   }

   /** Retrieve JSON array after filtering on string value associated with  JSON key*/
   @Test
   public void testExtractScalarArrayWithStringFilter() {
     Operator<SelectionResultsBlock> operator = getOperator(
         "select json_extract_scalar(jsonColumn, '$.data', 'STRING') FROM testTable WHERE json_extract_scalar"
             + "(jsonColumn, '$.name.first', 'STRING') = 'daffy'");
     SelectionResultsBlock block = operator.nextBlock();
     Collection<Object[]> rows = block.getRows();
     Assert.assertEquals(rows.size(), 1);

     Iterator<Object[]> iterator = rows.iterator();
     Assert.assertTrue(iterator.hasNext());
     Assert.assertEquals(iterator.next()[0], "[\"a\",\"b\",\"c\",\"d\"]");
   }

   /** Test filtering on string value within a JSON array*/
   @Test
   public void testExtractScalarWithArrayFilter() {
     Operator<SelectionResultsBlock> operator = getOperator(
         "select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_extract_scalar"
             + "(jsonColumn, '$.data[0]', 'STRING') IN ('i', 'k')");
     SelectionResultsBlock block = operator.nextBlock();
     Collection<Object[]> rows = block.getRows();
     Assert.assertEquals(rows.size(), 2);

     Iterator<Object[]> iterator = rows.iterator();
     Assert.assertTrue(iterator.hasNext());
     Assert.assertEquals(iterator.next()[0], "daisy");
     Assert.assertEquals(iterator.next()[0], "goofy");
   }

   @Test
   public void testJsonMatchWithoutIndex() {
     try {
       getOperator("select json_extract_scalar(stringColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match"
           + "(stringColumn, '\"$.id\"=101')");
       Assert.fail();
     } catch (IllegalStateException ise) {
       Assert.assertTrue(true);
     }
   }

   @Test
   public void testJsonMatchAtLevel1WithIndex() {
     Operator<SelectionResultsBlock> operator = getOperator(
         "select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
             + " '\"$.id\"=101')");

     SelectionResultsBlock block = operator.nextBlock();
     Collection<Object[]> rows = block.getRows();
     Assert.assertEquals(rows.size(), 1);

     Iterator<Object[]> iterator = rows.iterator();
     Assert.assertTrue(iterator.hasNext());
     Assert.assertEquals(iterator.next()[0], "daffy");
   }

   @Test
   public void testJsonMatchAtLevel2WithIndex() {
     Operator<SelectionResultsBlock> operator = getOperator(
         "select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
             + " '\"$.name.first\" = ''daffy''')");

     SelectionResultsBlock block = operator.nextBlock();
     Collection<Object[]> rows = block.getRows();
     Assert.assertEquals(rows.size(), 1);

     Iterator<Object[]> iterator = rows.iterator();
     Assert.assertTrue(iterator.hasNext());
     Assert.assertEquals(iterator.next()[0], "daffy");
   }

   @Test
   public void testJsonMatchArrayWithIndex() {
     Operator<SelectionResultsBlock> operator = getOperator(
         "select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
             + " '\"$.data[0]\" IN (''k'', ''j'')')");

     SelectionResultsBlock block = operator.nextBlock();
     Collection<Object[]> rows = block.getRows();
     Assert.assertEquals(rows.size(), 2);

     Iterator<Object[]> iterator = rows.iterator();
     Assert.assertTrue(iterator.hasNext());
     Assert.assertEquals(iterator.next()[0], "pluto");
     Assert.assertTrue(iterator.hasNext());
     Assert.assertEquals(iterator.next()[0], "goofy");
   }

   @Test
   public void testJsonMatchNestedArrayWithIndex() {
     Operator<SelectionResultsBlock> operator = getOperator(
         "select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
             + " '\"$.data[0].e[1].y[0].i1\" = 1')");

     SelectionResultsBlock block = operator.nextBlock();
     Collection<Object[]> rows = block.getRows();
     Assert.assertEquals(rows.size(), 1);

     Iterator<Object[]> iterator = rows.iterator();
     Assert.assertTrue(iterator.hasNext());
     Assert.assertEquals(iterator.next()[0], "nested");
   }

   @Test
   public void testJsonMatchMultidimensionalArrayWithIndex() {
     Operator<SelectionResultsBlock> operator = getOperator(
         "select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
             + " '\"$.data[0][1][0]\" = ''3''')");

     SelectionResultsBlock block = operator.nextBlock();
     Collection<Object[]> rows = block.getRows();
     Assert.assertEquals(rows.size(), 3);

     Iterator<Object[]> iterator = rows.iterator();
     Assert.assertTrue(iterator.hasNext());
     Assert.assertEquals(iterator.next()[0], "multi-dimensional-1");
   }

   @Test
   public void testJsonMatchMultidimensionalArrayGroupByWithIndex() {
     GroupByOperator groupByOperator = getOperator(
         "select json_extract_scalar(jsonColumn, '$.name.first', 'STRING'), count(*) FROM testTable WHERE json_match"
             + "(jsonColumn, '\"$.data[0][1][0]\" = ''3''') GROUP BY json_extract_scalar(jsonColumn, '$.name.first', "
             + "'STRING')");
     AggregationGroupByResult groupByResult = groupByOperator.nextBlock().getAggregationGroupByResult();
     Assert.assertNotNull(groupByResult);
     Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator = groupByResult.getGroupKeyIterator();

     Assert.assertTrue(groupKeyIterator.hasNext());
     GroupKeyGenerator.GroupKey groupKey1 = groupKeyIterator.next();
     Assert.assertEquals(groupKey1._keys, new Object[]{"multi-dimensional-2"});
     Assert.assertEquals(groupByResult.getResultForGroupId(0, groupKey1._groupId), 1L);

     Assert.assertTrue(groupKeyIterator.hasNext());
     GroupKeyGenerator.GroupKey groupKey2 = groupKeyIterator.next();
     Assert.assertEquals(groupKey2._keys, new Object[]{"multi-dimensional-1"});
     Assert.assertEquals(groupByResult.getResultForGroupId(0, groupKey1._groupId), 2L);

     Assert.assertFalse(groupKeyIterator.hasNext());
   }

   @AfterClass
   public void tearDown()
       throws IOException {
     _indexSegment.destroy();
     FileUtils.deleteDirectory(INDEX_DIR);
   }
 }
	/**
	* 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 java.io.File;
	import java.io.IOException;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.Collection;
	import java.util.HashSet;
	import java.util.Iterator;
	import java.util.List;
	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.results.SelectionResultsBlock;
	import org.apache.pinot.core.operator.query.GroupByOperator;
	import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
	import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
	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.local.segment.readers.GenericRowRecordReader;
	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.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;


	public class JsonDataTypeQueriesTest extends BaseQueriesTest {
	private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "JsonDataTypeQueriesTest");
	private static final String RAW_TABLE_NAME = "testTable";
	private static final String SEGMENT_NAME = "testSegment";
	private static final int NUM_RECORDS = 10;

	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;
	}

	GenericRow createRecord(int intValue, String stringValue, String jsonValue) {
	GenericRow record = new GenericRow();
	record.putValue(INT_COLUMN, intValue);
	record.putValue(STRING_COLUMN, stringValue);
	record.putValue(JSON_COLUMN, jsonValue);

	return record;
	}

	@Override
	protected List<IndexSegment> getIndexSegments() {
	return _indexSegments;
	}

	@BeforeClass
	public void setUp()
	throws Exception {
	FileUtils.deleteDirectory(INDEX_DIR);

	List<GenericRow> records = new ArrayList<>(NUM_RECORDS);
	records.add(createRecord(1, "daffy duck",
	"{\"name\": {\"first\": \"daffy\", \"last\": \"duck\"}, \"id\": 101, \"data\": [\"a\", \"b\", \"c\", \"d\"]}"));
	records.add(createRecord(2, "mickey mouse",
	"{\"name\": {\"first\": \"mickey\", \"last\": \"mouse\"}, \"id\": 111, \"data\": [\"e\", \"b\", \"c\", "
	+ "\"d\"]}"));
	records.add(createRecord(3, "donald duck",
	"{\"name\": {\"first\": \"donald\", \"last\": \"duck\"}, \"id\": 121, \"data\": [\"f\", \"b\", \"c\", "
	+ "\"d\"]}"));
	records.add(createRecord(4, "scrooge mcduck",
	"{\"name\": {\"first\": \"scrooge\", \"last\": \"mcduck\"}, \"id\": 131, \"data\": [\"g\", \"b\", \"c\", "
	+ "\"d\"]}"));
	records.add(createRecord(5, "minnie mouse",
	"{\"name\": {\"first\": \"minnie\", \"last\": \"mouse\"}, \"id\": 141, \"data\": [\"h\", \"b\", \"c\", "
	+ "\"d\"]}"));
	records.add(createRecord(6, "daisy duck",
	"{\"name\": {\"first\": \"daisy\", \"last\": \"duck\"}, \"id\": 161.5, \"data\": [\"i\", \"b\", \"c\", "
	+ "\"d\"]}"));
	records.add(createRecord(7, "pluto dog",
	"{\"name\": {\"first\": \"pluto\", \"last\": \"dog\"}, \"id\": 161, \"data\": [\"j\", \"b\", \"c\", \"d\"]}"));
	records.add(createRecord(8, "goofy dwag",
	"{\"name\": {\"first\": \"goofy\", \"last\": \"dwag\"}, \"id\": 171, \"data\": [\"k\", \"b\", \"c\", \"d\"]}"));
	records.add(createRecord(9, "ludwik von drake",
	"{\"name\": {\"first\": \"ludwik\", \"last\": \"von drake\"}, \"id\": 181, \"data\": [\"l\", \"b\", \"c\", "
	+ "\"d\"]}"));
	records.add(createRecord(10, "nested array",
	"{\"name\":{\"first\":\"nested\",\"last\":\"array\"},\"id\":111,\"data\":[{\"e\":[{\"x\":[{\"i1\":1,"
	+ "\"i2\":2}]},{\"y\":[{\"i1\":1,\"i2\":2}]},{\"z\":[{\"i1\":1,\"i2\":2}]}]},{\"b\":[{\"x\":[{\"i1\":1,"
	+ "\"i2\":2}]},{\"y\":[{\"i1\":1,\"i2\":2}]},{\"z\":[{\"i1\":10,\"i2\":20}]}]}]}"));
	records.add(createRecord(11, "multi-dimensional-1 array",
	"{\"name\": {\"first\": \"multi-dimensional-1\",\"last\": \"array\"},\"id\": 111,\"data\": [[[1,2],[3,4]],"
	+ "[[\"a\",\"b\"],[\"c\",\"d\"]]]}"));
	records.add(createRecord(12, "multi-dimensional-2 array",
	"{\"name\": {\"first\": \"multi-dimensional-2\",\"last\": \"array\"},\"id\": 111,\"data\": [[[1,2],[3,4]],"
	+ "[[\"a\",\"b\"],[\"c\",\"d\"]]]}"));
	records.add(createRecord(13, "multi-dimensional-1 array",
	"{\"name\": {\"first\": \"multi-dimensional-1\",\"last\": \"array\"},\"id\": 111,\"data\": [[[1,2],[3,4]],"
	+ "[[\"a\",\"b\"],[\"c\",\"d\"]]]}"));

	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());

	SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
	driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records));
	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 result column type of a simple select query against JSON column */
	@Test
	public void testSimpleSelectOnJsonColumn() {
	try {
	Operator<SelectionResultsBlock> operator = getOperator("select jsonColumn FROM testTable limit 100");
	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 13);
	Assert.assertEquals(block.getDataSchema().getColumnDataType(0), DataSchema.ColumnDataType.JSON);
	} catch (IllegalStateException ise) {
	Assert.assertTrue(true);
	}
	}

	/** Test filtering on string value associated with JSON key*/
	@Test
	public void testExtractScalarWithStringFilter() {
	Operator<SelectionResultsBlock> operator = getOperator(
	"select intColumn, json_extract_scalar(jsonColumn, '$.name.last', 'STRING') FROM testTable WHERE "
	+ "json_extract_scalar(jsonColumn, '$.name.first', 'STRING') = 'daffy'");
	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 1);

	Iterator<Object[]> iterator = rows.iterator();
	Assert.assertTrue(iterator.hasNext());
	Object[] row = iterator.next();
	Assert.assertEquals(row[0], 1);
	Assert.assertEquals(row[1], "duck");
	}

	/** Test filtering on number value associated with JSON key*/
	@Test
	public void testExtractScalarWithNumericIntFilter() {
	Operator<SelectionResultsBlock> operator = getOperator(
	"select json_extract_scalar(jsonColumn, '$.name.last', 'STRING') FROM testTable WHERE json_extract_scalar"
	+ "(jsonColumn, '$.id', 'INT') = 171");
	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 1);

	Iterator<Object[]> iterator = rows.iterator();
	Assert.assertTrue(iterator.hasNext());
	Assert.assertEquals(iterator.next()[0], "dwag");
	}

	/** Test filtering on float value associated with JSON key*/
	@Test
	public void testExtractScalarWithNumericFloatFilter() {

	// query to retrieve result as INT
	Operator<SelectionResultsBlock> operator1 = getOperator(
	"select json_extract_scalar(jsonColumn, '$.id', 'INT') FROM testTable WHERE json_extract_scalar(jsonColumn, "
	+ "'$.id', 'FLOAT') = 161.5");
	SelectionResultsBlock block1 = operator1.nextBlock();
	Collection<Object[]> rows1 = block1.getRows();
	Assert.assertEquals(rows1.size(), 1);

	Iterator<Object[]> iterator1 = rows1.iterator();
	Assert.assertTrue(iterator1.hasNext());
	Assert.assertEquals(iterator1.next()[0], 161);

	// query to retrieve result as DOUBLE
	Operator<SelectionResultsBlock> operator2 = getOperator(
	"select json_extract_scalar(jsonColumn, '$.id', 'DOUBLE') FROM testTable WHERE json_extract_scalar"
	+ "(jsonColumn, '$.id', 'FLOAT') = 161.5");
	SelectionResultsBlock block2 = operator2.nextBlock();
	Collection<Object[]> rows2 = block2.getRows();
	Assert.assertEquals(rows2.size(), 1);

	Iterator<Object[]> iterator2 = rows2.iterator();
	Assert.assertTrue(iterator2.hasNext());
	Assert.assertEquals(iterator2.next()[0], 161.5d);
	}

	/** Retrieve JSON array after filtering on string value associated with JSON key*/
	@Test
	public void testExtractScalarArrayWithStringFilter() {
	Operator<SelectionResultsBlock> operator = getOperator(
	"select json_extract_scalar(jsonColumn, '$.data', 'STRING') FROM testTable WHERE json_extract_scalar"
	+ "(jsonColumn, '$.name.first', 'STRING') = 'daffy'");
	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 1);

	Iterator<Object[]> iterator = rows.iterator();
	Assert.assertTrue(iterator.hasNext());
	Assert.assertEquals(iterator.next()[0], "[\"a\",\"b\",\"c\",\"d\"]");
	}

	/** Test filtering on string value within a JSON array*/
	@Test
	public void testExtractScalarWithArrayFilter() {
	Operator<SelectionResultsBlock> operator = getOperator(
	"select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_extract_scalar"
	+ "(jsonColumn, '$.data[0]', 'STRING') IN ('i', 'k')");
	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 2);

	Iterator<Object[]> iterator = rows.iterator();
	Assert.assertTrue(iterator.hasNext());
	Assert.assertEquals(iterator.next()[0], "daisy");
	Assert.assertEquals(iterator.next()[0], "goofy");
	}

	@Test
	public void testJsonMatchWithoutIndex() {
	try {
	getOperator("select json_extract_scalar(stringColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match"
	+ "(stringColumn, '\"$.id\"=101')");
	Assert.fail();
	} catch (IllegalStateException ise) {
	Assert.assertTrue(true);
	}
	}

	@Test
	public void testJsonMatchAtLevel1WithIndex() {
	Operator<SelectionResultsBlock> operator = getOperator(
	"select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
	+ " '\"$.id\"=101')");

	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 1);

	Iterator<Object[]> iterator = rows.iterator();
	Assert.assertTrue(iterator.hasNext());
	Assert.assertEquals(iterator.next()[0], "daffy");
	}

	@Test
	public void testJsonMatchAtLevel2WithIndex() {
	Operator<SelectionResultsBlock> operator = getOperator(
	"select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
	+ " '\"$.name.first\" = ''daffy''')");

	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 1);

	Iterator<Object[]> iterator = rows.iterator();
	Assert.assertTrue(iterator.hasNext());
	Assert.assertEquals(iterator.next()[0], "daffy");
	}

	@Test
	public void testJsonMatchArrayWithIndex() {
	Operator<SelectionResultsBlock> operator = getOperator(
	"select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
	+ " '\"$.data[0]\" IN (''k'', ''j'')')");

	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 2);

	Iterator<Object[]> iterator = rows.iterator();
	Assert.assertTrue(iterator.hasNext());
	Assert.assertEquals(iterator.next()[0], "pluto");
	Assert.assertTrue(iterator.hasNext());
	Assert.assertEquals(iterator.next()[0], "goofy");
	}

	@Test
	public void testJsonMatchNestedArrayWithIndex() {
	Operator<SelectionResultsBlock> operator = getOperator(
	"select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
	+ " '\"$.data[0].e[1].y[0].i1\" = 1')");

	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 1);

	Iterator<Object[]> iterator = rows.iterator();
	Assert.assertTrue(iterator.hasNext());
	Assert.assertEquals(iterator.next()[0], "nested");
	}

	@Test
	public void testJsonMatchMultidimensionalArrayWithIndex() {
	Operator<SelectionResultsBlock> operator = getOperator(
	"select json_extract_scalar(jsonColumn, '$.name.first', 'STRING') FROM testTable WHERE json_match(jsonColumn,"
	+ " '\"$.data[0][1][0]\" = ''3''')");

	SelectionResultsBlock block = operator.nextBlock();
	Collection<Object[]> rows = block.getRows();
	Assert.assertEquals(rows.size(), 3);

	Iterator<Object[]> iterator = rows.iterator();
	Assert.assertTrue(iterator.hasNext());
	Assert.assertEquals(iterator.next()[0], "multi-dimensional-1");
	}

	@Test
	public void testJsonMatchMultidimensionalArrayGroupByWithIndex() {
	GroupByOperator groupByOperator = getOperator(
	"select json_extract_scalar(jsonColumn, '$.name.first', 'STRING'), count(*) FROM testTable WHERE json_match"
	+ "(jsonColumn, '\"$.data[0][1][0]\" = ''3''') GROUP BY json_extract_scalar(jsonColumn, '$.name.first', "
	+ "'STRING')");
	AggregationGroupByResult groupByResult = groupByOperator.nextBlock().getAggregationGroupByResult();
	Assert.assertNotNull(groupByResult);
	Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator = groupByResult.getGroupKeyIterator();

	Assert.assertTrue(groupKeyIterator.hasNext());
	GroupKeyGenerator.GroupKey groupKey1 = groupKeyIterator.next();
	Assert.assertEquals(groupKey1._keys, new Object[]{"multi-dimensional-2"});
	Assert.assertEquals(groupByResult.getResultForGroupId(0, groupKey1._groupId), 1L);

	Assert.assertTrue(groupKeyIterator.hasNext());
	GroupKeyGenerator.GroupKey groupKey2 = groupKeyIterator.next();
	Assert.assertEquals(groupKey2._keys, new Object[]{"multi-dimensional-1"});
	Assert.assertEquals(groupByResult.getResultForGroupId(0, groupKey1._groupId), 2L);

	Assert.assertFalse(groupKeyIterator.hasNext());
	}

	@AfterClass
	public void tearDown()
	throws IOException {
	_indexSegment.destroy();
	FileUtils.deleteDirectory(INDEX_DIR);
	}
	}