pinot-core/src/test/java/org/apache/pinot/queries/FirstWithTimeQueriesTest.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 com.google.common.collect.Lists;
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
 import org.apache.commons.io.FileUtils;
 import org.apache.pinot.common.response.broker.BrokerResponseNative;
 import org.apache.pinot.common.response.broker.ResultTable;
 import org.apache.pinot.common.utils.DataSchema;
 import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
 import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock;
 import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
 import org.apache.pinot.core.operator.query.AggregationOperator;
 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.customobject.DoubleLongPair;
 import org.apache.pinot.segment.local.customobject.FloatLongPair;
 import org.apache.pinot.segment.local.customobject.IntLongPair;
 import org.apache.pinot.segment.local.customobject.LongLongPair;
 import org.apache.pinot.segment.local.customobject.StringLongPair;
 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.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.DataType;
 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.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;

 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.assertTrue;


 /**
  * Queries test for FIRSTWITHTIME queries.
  */
 public class FirstWithTimeQueriesTest extends BaseQueriesTest {
   private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "FirstQueriesTest");
   private static final String RAW_TABLE_NAME = "testTable";
   private static final String SEGMENT_NAME = "testSegment";
   private static final Random RANDOM = new Random();

   private static final int NUM_RECORDS = 2000;
   private static final int MAX_VALUE = 1000;

   private static final String BOOL_COLUMN = "boolColumn";
   private static final String BOOL_NO_DICT_COLUMN = "boolNoDictColumn";
   private static final String INT_COLUMN = "intColumn";
   private static final String INT_MV_COLUMN = "intMVColumn";
   private static final String INT_NO_DICT_COLUMN = "intNoDictColumn";
   private static final String LONG_COLUMN = "longColumn";
   private static final String LONG_NO_DICT_COLUMN = "longNoDictColumn";
   private static final String FLOAT_COLUMN = "floatColumn";
   private static final String FLOAT_NO_DICT_COLUMN = "floatNoDictColumn";
   private static final String DOUBLE_COLUMN = "doubleColumn";
   private static final String DOUBLE_NO_DICT_COLUMN = "doubleNoDictColumn";
   private static final String STRING_COLUMN = "stringColumn";
   private static final String STRING_NO_DICT_COLUMN = "stringNoDictColumn";
   private static final String TIME_COLUMN = "timestampColumn";
   private static final Schema SCHEMA = new Schema.SchemaBuilder()
       .addSingleValueDimension(BOOL_COLUMN, DataType.BOOLEAN)
       .addSingleValueDimension(BOOL_NO_DICT_COLUMN, DataType.BOOLEAN)
       .addSingleValueDimension(INT_COLUMN, DataType.INT)
       .addMultiValueDimension(INT_MV_COLUMN, DataType.INT)
       .addSingleValueDimension(INT_NO_DICT_COLUMN, DataType.INT)
       .addSingleValueDimension(LONG_COLUMN, DataType.LONG)
       .addSingleValueDimension(LONG_NO_DICT_COLUMN, DataType.LONG)
       .addSingleValueDimension(FLOAT_COLUMN, DataType.FLOAT)
       .addSingleValueDimension(FLOAT_NO_DICT_COLUMN, DataType.FLOAT)
       .addSingleValueDimension(DOUBLE_COLUMN, DataType.DOUBLE)
       .addSingleValueDimension(DOUBLE_NO_DICT_COLUMN, DataType.DOUBLE)
       .addSingleValueDimension(STRING_COLUMN, DataType.STRING)
       .addSingleValueDimension(STRING_NO_DICT_COLUMN, DataType.STRING)
       .addSingleValueDimension(TIME_COLUMN, DataType.LONG).build();
   private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME)
       .setNoDictionaryColumns(
           Lists.newArrayList(INT_NO_DICT_COLUMN, LONG_NO_DICT_COLUMN, FLOAT_NO_DICT_COLUMN, DOUBLE_NO_DICT_COLUMN))
       .build();

   private Boolean _expectedResultFirstBoolean;
   private Integer _expectedResultFirstInt;
   private Long _expectedResultFirstLong;
   private Float _expectedResultFirstFloat;
   private Double _expectedResultFirstDouble;
   private String _expectedResultFirstString;
   private Map<Integer, Boolean> _boolGroupValues;
   private Map<Integer, Integer> _intGroupValues;
   private Map<Integer, Long> _longGroupValues;
   private Map<Integer, Float> _floatGroupValues;
   private Map<Integer, Double> _doubleGroupValues;
   private Map<Integer, String> _stringGroupValues;
   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;
   }

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

     List<GenericRow> records = new ArrayList<>(NUM_RECORDS);
     _boolGroupValues = new HashMap<>();
     _intGroupValues = new HashMap<>();
     _longGroupValues = new HashMap<>();
     _floatGroupValues = new HashMap<>();
     _doubleGroupValues = new HashMap<>();
     _stringGroupValues = new HashMap<>();
     for (int i = 0; i < NUM_RECORDS; i++) {
       boolean boolValue = RANDOM.nextBoolean();
       int intValue = RANDOM.nextInt(MAX_VALUE);
       long longValue = RANDOM.nextLong();
       float floatValue = RANDOM.nextFloat();
       double doubleValue = RANDOM.nextDouble();
       String strValue = String.valueOf(RANDOM.nextDouble());
       GenericRow record = new GenericRow();
       record.putValue(BOOL_COLUMN, boolValue);
       record.putValue(BOOL_NO_DICT_COLUMN, boolValue);
       record.putValue(INT_COLUMN, intValue);
       record.putValue(INT_MV_COLUMN, new Integer[]{intValue, intValue});
       record.putValue(INT_NO_DICT_COLUMN, intValue);
       record.putValue(LONG_COLUMN, longValue);
       record.putValue(LONG_NO_DICT_COLUMN, longValue);
       record.putValue(FLOAT_COLUMN, floatValue);
       record.putValue(FLOAT_NO_DICT_COLUMN, floatValue);
       record.putValue(DOUBLE_COLUMN, doubleValue);
       record.putValue(DOUBLE_NO_DICT_COLUMN, doubleValue);
       record.putValue(STRING_COLUMN, strValue);
       record.putValue(STRING_NO_DICT_COLUMN, strValue);
       record.putValue(TIME_COLUMN, (long) i);
       if (i == 0) {
         _expectedResultFirstBoolean = boolValue;
         _expectedResultFirstInt = intValue;
         _expectedResultFirstLong = longValue;
         _expectedResultFirstFloat = floatValue;
         _expectedResultFirstDouble = doubleValue;
         _expectedResultFirstString = strValue;
       }
       _boolGroupValues.putIfAbsent(intValue, boolValue);
       _intGroupValues.putIfAbsent(intValue, intValue);
       _longGroupValues.putIfAbsent(intValue, longValue);
       _floatGroupValues.putIfAbsent(intValue, floatValue);
       _doubleGroupValues.putIfAbsent(intValue, doubleValue);
       _stringGroupValues.putIfAbsent(intValue, strValue);
       records.add(record);
     }

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

     ImmutableSegment immutableSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), ReadMode.mmap);
     _indexSegment = immutableSegment;
     _indexSegments = Arrays.asList(immutableSegment, immutableSegment);
   }

   @Test
   public void testAggregationOnly() {
     String query = "SELECT "
         + "FIRSTWITHTIME(boolColumn, timestampColumn, 'BOOLEAN'), "
         + "FIRSTWITHTIME(intColumn, timestampColumn, 'INT'), "
         + "FIRSTWITHTIME(longColumn, timestampColumn, 'LONG'), "
         + "FIRSTWITHTIME(floatColumn, timestampColumn, 'FLOAT'), "
         + "FIRSTWITHTIME(doubleColumn, timestampColumn, 'DOUBLE'), "
         + "FIRSTWITHTIME(stringColumn, timestampColumn, 'STRING') "
         + "FROM testTable";

     // Inner segment
     AggregationOperator aggregationOperator = getOperator(query);
     AggregationResultsBlock resultsBlock = aggregationOperator.nextBlock();
     QueriesTestUtils.testInnerSegmentExecutionStatistics(aggregationOperator.getExecutionStatistics(), NUM_RECORDS, 0,
         7 * NUM_RECORDS, NUM_RECORDS);
     List<Object> aggregationResult = resultsBlock.getResults();
     assertNotNull(aggregationResult);
     assertEquals(((IntLongPair) aggregationResult.get(0)).getValue() != 0, _expectedResultFirstBoolean.booleanValue());
     assertEquals(((IntLongPair) aggregationResult.get(1)).getValue(), _expectedResultFirstInt);
     assertEquals(((LongLongPair) aggregationResult.get(2)).getValue(), _expectedResultFirstLong);
     assertEquals(((FloatLongPair) aggregationResult.get(3)).getValue(), _expectedResultFirstFloat);
     assertEquals(((DoubleLongPair) aggregationResult.get(4)).getValue(), _expectedResultFirstDouble);
     assertEquals(((StringLongPair) aggregationResult.get(5)).getValue(), _expectedResultFirstString);

     // Inter segments (expect 4 * inner segment result)
     BrokerResponseNative brokerResponse = getBrokerResponse(query);
     DataSchema expectedDataSchema = new DataSchema(new String[]{
         "firstwithtime(boolColumn,timestampColumn,'BOOLEAN')",
         "firstwithtime(intColumn,timestampColumn,'INT')",
         "firstwithtime(longColumn,timestampColumn,'LONG')",
         "firstwithtime(floatColumn,timestampColumn,'FLOAT')",
         "firstwithtime(doubleColumn,timestampColumn,'DOUBLE')",
         "firstwithtime(stringColumn,timestampColumn,'STRING')"
     }, new ColumnDataType[]{
         ColumnDataType.BOOLEAN,
         ColumnDataType.INT,
         ColumnDataType.LONG,
         ColumnDataType.FLOAT,
         ColumnDataType.DOUBLE,
         ColumnDataType.STRING
     });
     Object[] expectedResults = new Object[]{
         _expectedResultFirstBoolean,
         _expectedResultFirstInt,
         _expectedResultFirstLong,
         _expectedResultFirstFloat,
         _expectedResultFirstDouble,
         _expectedResultFirstString
     };
     QueriesTestUtils.testInterSegmentsResult(brokerResponse, 4 * NUM_RECORDS, 0L, 4 * 7 * NUM_RECORDS,
         4 * NUM_RECORDS, new ResultTable(expectedDataSchema, Collections.singletonList(expectedResults)));
   }

   @Test
   public void testAggregationOnlyNoDictionary() {
     String query = "SELECT "
         + "FIRSTWITHTIME(boolNoDictColumn,timestampColumn,'boolean'), "
         + "FIRSTWITHTIME(intNoDictColumn,timestampColumn,'int'), "
         + "FIRSTWITHTIME(longNoDictColumn,timestampColumn,'long'), "
         + "FIRSTWITHTIME(floatNoDictColumn,timestampColumn,'float'), "
         + "FIRSTWITHTIME(doubleNoDictColumn,timestampColumn,'double'), "
         + "FIRSTWITHTIME(stringNoDictColumn,timestampColumn,'string') "
         + "FROM testTable";

     // Inner segment
     AggregationOperator aggregationOperator = getOperator(query);
     AggregationResultsBlock resultsBlock = aggregationOperator.nextBlock();
     QueriesTestUtils.testInnerSegmentExecutionStatistics(aggregationOperator.getExecutionStatistics(), NUM_RECORDS, 0,
         7 * NUM_RECORDS, NUM_RECORDS);
     List<Object> aggregationResult = resultsBlock.getResults();
     assertNotNull(aggregationResult);
     assertEquals(((IntLongPair) aggregationResult.get(0)).getValue() != 0, _expectedResultFirstBoolean.booleanValue());
     assertEquals(((IntLongPair) aggregationResult.get(1)).getValue(), _expectedResultFirstInt);
     assertEquals(((LongLongPair) aggregationResult.get(2)).getValue(), _expectedResultFirstLong);
     assertEquals(((FloatLongPair) aggregationResult.get(3)).getValue(), _expectedResultFirstFloat);
     assertEquals(((DoubleLongPair) aggregationResult.get(4)).getValue(), _expectedResultFirstDouble);
     assertEquals(((StringLongPair) aggregationResult.get(5)).getValue(), _expectedResultFirstString);

     // Inter segments (expect 4 * inner segment result)
     BrokerResponseNative brokerResponse = getBrokerResponse(query);
     DataSchema expectedDataSchema = new DataSchema(new String[]{
         "firstwithtime(boolNoDictColumn,timestampColumn,'BOOLEAN')",
         "firstwithtime(intNoDictColumn,timestampColumn,'INT')",
         "firstwithtime(longNoDictColumn,timestampColumn,'LONG')",
         "firstwithtime(floatNoDictColumn,timestampColumn,'FLOAT')",
         "firstwithtime(doubleNoDictColumn,timestampColumn,'DOUBLE')",
         "firstwithtime(stringNoDictColumn,timestampColumn,'STRING')"
     }, new ColumnDataType[]{
         ColumnDataType.BOOLEAN,
         ColumnDataType.INT,
         ColumnDataType.LONG,
         ColumnDataType.FLOAT,
         ColumnDataType.DOUBLE,
         ColumnDataType.STRING
     });
     Object[] expectedResults = new Object[]{
         _expectedResultFirstBoolean,
         _expectedResultFirstInt,
         _expectedResultFirstLong,
         _expectedResultFirstFloat,
         _expectedResultFirstDouble,
         _expectedResultFirstString
     };
     QueriesTestUtils.testInterSegmentsResult(brokerResponse, 4 * NUM_RECORDS, 0L, 4 * 7 * NUM_RECORDS,
         4 * NUM_RECORDS, new ResultTable(expectedDataSchema, Collections.singletonList(expectedResults)));
   }

   @Test
   public void testAggregationGroupBySV() {
     String query = "SELECT intColumn AS key, "
         + "FIRSTWITHTIME(boolColumn,timestampColumn,'boolean') AS v1, "
         + "FIRSTWITHTIME(intColumn,timestampColumn,'int') AS v2, "
         + "FIRSTWITHTIME(longColumn,timestampColumn,'long') AS v3, "
         + "FIRSTWITHTIME(floatColumn,timestampColumn,'float') AS v4, "
         + "FIRSTWITHTIME(doubleColumn,timestampColumn,'double') AS v5, "
         + "FIRSTWITHTIME(stringColumn,timestampColumn,'string') AS v6 "
         + "FROM testTable GROUP BY key";
     verifyAggregationGroupBy(query, 7);
   }

   @Test
   public void testAggregationGroupBySVNoDictionary() {
     String query = "SELECT intNoDictColumn AS key, "
         + "FIRSTWITHTIME(boolNoDictColumn,timestampColumn,'boolean') AS v1, "
         + "FIRSTWITHTIME(intNoDictColumn,timestampColumn,'int') AS v2, "
         + "FIRSTWITHTIME(longNoDictColumn,timestampColumn,'long') AS v3, "
         + "FIRSTWITHTIME(floatNoDictColumn,timestampColumn,'float') AS v4, "
         + "FIRSTWITHTIME(doubleNoDictColumn,timestampColumn,'double') AS v5, "
         + "FIRSTWITHTIME(stringNoDictColumn,timestampColumn,'string') AS v6 "
         + "FROM testTable GROUP BY key";
     verifyAggregationGroupBy(query, 7);
   }

   @Test
   public void testAggregationGroupByMV() {
     String query = "SELECT intMVColumn AS key, "
         + "FIRSTWITHTIME(boolColumn,timestampColumn,'boolean') AS v1, "
         + "FIRSTWITHTIME(intColumn,timestampColumn,'int') AS v2, "
         + "FIRSTWITHTIME(longColumn,timestampColumn,'long') AS v3, "
         + "FIRSTWITHTIME(floatColumn,timestampColumn,'float') AS v4, "
         + "FIRSTWITHTIME(doubleColumn,timestampColumn,'double') AS v5, "
         + "FIRSTWITHTIME(stringColumn,timestampColumn,'string') AS v6 "
         + "FROM testTable GROUP BY key";
     verifyAggregationGroupBy(query, 8);
   }

   @Test
   public void testAggregationGroupByMVNoDictionary() {
     String query = "SELECT intMVColumn AS key, "
         + "FIRSTWITHTIME(boolNoDictColumn,timestampColumn,'boolean') AS v1, "
         + "FIRSTWITHTIME(intNoDictColumn,timestampColumn,'int') AS v2, "
         + "FIRSTWITHTIME(longNoDictColumn,timestampColumn,'long') AS v3, "
         + "FIRSTWITHTIME(floatNoDictColumn,timestampColumn,'float') AS v4, "
         + "FIRSTWITHTIME(doubleNoDictColumn,timestampColumn,'double') AS v5, "
         + "FIRSTWITHTIME(stringNoDictColumn,timestampColumn,'string') AS v6 "
         + "FROM testTable GROUP BY key";
     verifyAggregationGroupBy(query, 8);
   }

   private void verifyAggregationGroupBy(String query, int numProjectedColumns) {
     // Inner segment
     GroupByOperator groupByOperator = getOperator(query);
     GroupByResultsBlock resultsBlock = groupByOperator.nextBlock();
     QueriesTestUtils.testInnerSegmentExecutionStatistics(groupByOperator.getExecutionStatistics(), NUM_RECORDS, 0,
         numProjectedColumns * (long) NUM_RECORDS, NUM_RECORDS);
     AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
     assertNotNull(aggregationGroupByResult);
     int numGroups = 0;
     Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
     while (groupKeyIterator.hasNext()) {
       numGroups++;
       GroupKeyGenerator.GroupKey groupKey = groupKeyIterator.next();
       Integer key = (Integer) groupKey._keys[0];
       assertTrue(_intGroupValues.containsKey(key));
       assertEquals(((IntLongPair) aggregationGroupByResult.getResultForGroupId(0, groupKey._groupId)).getValue() != 0,
           (boolean) _boolGroupValues.get(key));
       assertEquals(((IntLongPair) aggregationGroupByResult.getResultForGroupId(1, groupKey._groupId)).getValue(),
           _intGroupValues.get(key));
       assertEquals(((LongLongPair) aggregationGroupByResult.getResultForGroupId(2, groupKey._groupId)).getValue(),
           _longGroupValues.get(key));
       assertEquals(((FloatLongPair) aggregationGroupByResult.getResultForGroupId(3, groupKey._groupId)).getValue(),
           _floatGroupValues.get(key));
       assertEquals(((DoubleLongPair) aggregationGroupByResult.getResultForGroupId(4, groupKey._groupId)).getValue(),
           _doubleGroupValues.get(key));
       assertEquals(((StringLongPair) aggregationGroupByResult.getResultForGroupId(5, groupKey._groupId)).getValue(),
           _stringGroupValues.get(key));
     }
     assertEquals(numGroups, _intGroupValues.size());

     // Inter segments (expect 4 * inner segment result)
     BrokerResponseNative brokerResponse = getBrokerResponse(query);
     assertEquals(brokerResponse.getNumDocsScanned(), 4 * NUM_RECORDS);
     assertEquals(brokerResponse.getNumEntriesScannedInFilter(), 0);
     assertEquals(brokerResponse.getNumEntriesScannedPostFilter(), 4 * numProjectedColumns * (long) NUM_RECORDS);
     assertEquals(brokerResponse.getTotalDocs(), 4 * NUM_RECORDS);

     ResultTable resultTable = brokerResponse.getResultTable();
     DataSchema expectedDataSchema =
         new DataSchema(new String[]{"key", "v1", "v2", "v3", "v4", "v5", "v6"}, new ColumnDataType[]{
             ColumnDataType.INT, ColumnDataType.BOOLEAN, ColumnDataType.INT, ColumnDataType.LONG,
             ColumnDataType.FLOAT, ColumnDataType.DOUBLE, ColumnDataType.STRING
         });
     assertEquals(resultTable.getDataSchema(), expectedDataSchema);
     List<Object[]> rows = resultTable.getRows();
     assertEquals(rows.size(), 10);
     for (Object[] row : rows) {
       assertEquals(row.length, 7);
       int key = (Integer) row[0];
       assertEquals(row[1], _boolGroupValues.get(key));
       assertEquals(row[2], _intGroupValues.get(key));
       assertEquals(row[3], _longGroupValues.get(key));
       assertEquals(row[4], _floatGroupValues.get(key));
       assertEquals(row[5], _doubleGroupValues.get(key));
       assertEquals(row[6], _stringGroupValues.get(key));
     }
   }

   @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 com.google.common.collect.Lists;
	import java.io.File;
	import java.io.IOException;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.Collections;
	import java.util.HashMap;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Map;
	import java.util.Random;
	import org.apache.commons.io.FileUtils;
	import org.apache.pinot.common.response.broker.BrokerResponseNative;
	import org.apache.pinot.common.response.broker.ResultTable;
	import org.apache.pinot.common.utils.DataSchema;
	import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
	import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock;
	import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
	import org.apache.pinot.core.operator.query.AggregationOperator;
	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.customobject.DoubleLongPair;
	import org.apache.pinot.segment.local.customobject.FloatLongPair;
	import org.apache.pinot.segment.local.customobject.IntLongPair;
	import org.apache.pinot.segment.local.customobject.LongLongPair;
	import org.apache.pinot.segment.local.customobject.StringLongPair;
	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.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.DataType;
	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.annotations.AfterClass;
	import org.testng.annotations.BeforeClass;
	import org.testng.annotations.Test;

	import static org.testng.Assert.assertEquals;
	import static org.testng.Assert.assertNotNull;
	import static org.testng.Assert.assertTrue;


	/**
	* Queries test for FIRSTWITHTIME queries.
	*/
	public class FirstWithTimeQueriesTest extends BaseQueriesTest {
	private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "FirstQueriesTest");
	private static final String RAW_TABLE_NAME = "testTable";
	private static final String SEGMENT_NAME = "testSegment";
	private static final Random RANDOM = new Random();

	private static final int NUM_RECORDS = 2000;
	private static final int MAX_VALUE = 1000;

	private static final String BOOL_COLUMN = "boolColumn";
	private static final String BOOL_NO_DICT_COLUMN = "boolNoDictColumn";
	private static final String INT_COLUMN = "intColumn";
	private static final String INT_MV_COLUMN = "intMVColumn";
	private static final String INT_NO_DICT_COLUMN = "intNoDictColumn";
	private static final String LONG_COLUMN = "longColumn";
	private static final String LONG_NO_DICT_COLUMN = "longNoDictColumn";
	private static final String FLOAT_COLUMN = "floatColumn";
	private static final String FLOAT_NO_DICT_COLUMN = "floatNoDictColumn";
	private static final String DOUBLE_COLUMN = "doubleColumn";
	private static final String DOUBLE_NO_DICT_COLUMN = "doubleNoDictColumn";
	private static final String STRING_COLUMN = "stringColumn";
	private static final String STRING_NO_DICT_COLUMN = "stringNoDictColumn";
	private static final String TIME_COLUMN = "timestampColumn";
	private static final Schema SCHEMA = new Schema.SchemaBuilder()
	.addSingleValueDimension(BOOL_COLUMN, DataType.BOOLEAN)
	.addSingleValueDimension(BOOL_NO_DICT_COLUMN, DataType.BOOLEAN)
	.addSingleValueDimension(INT_COLUMN, DataType.INT)
	.addMultiValueDimension(INT_MV_COLUMN, DataType.INT)
	.addSingleValueDimension(INT_NO_DICT_COLUMN, DataType.INT)
	.addSingleValueDimension(LONG_COLUMN, DataType.LONG)
	.addSingleValueDimension(LONG_NO_DICT_COLUMN, DataType.LONG)
	.addSingleValueDimension(FLOAT_COLUMN, DataType.FLOAT)
	.addSingleValueDimension(FLOAT_NO_DICT_COLUMN, DataType.FLOAT)
	.addSingleValueDimension(DOUBLE_COLUMN, DataType.DOUBLE)
	.addSingleValueDimension(DOUBLE_NO_DICT_COLUMN, DataType.DOUBLE)
	.addSingleValueDimension(STRING_COLUMN, DataType.STRING)
	.addSingleValueDimension(STRING_NO_DICT_COLUMN, DataType.STRING)
	.addSingleValueDimension(TIME_COLUMN, DataType.LONG).build();
	private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME)
	.setNoDictionaryColumns(
	Lists.newArrayList(INT_NO_DICT_COLUMN, LONG_NO_DICT_COLUMN, FLOAT_NO_DICT_COLUMN, DOUBLE_NO_DICT_COLUMN))
	.build();

	private Boolean _expectedResultFirstBoolean;
	private Integer _expectedResultFirstInt;
	private Long _expectedResultFirstLong;
	private Float _expectedResultFirstFloat;
	private Double _expectedResultFirstDouble;
	private String _expectedResultFirstString;
	private Map<Integer, Boolean> _boolGroupValues;
	private Map<Integer, Integer> _intGroupValues;
	private Map<Integer, Long> _longGroupValues;
	private Map<Integer, Float> _floatGroupValues;
	private Map<Integer, Double> _doubleGroupValues;
	private Map<Integer, String> _stringGroupValues;
	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;
	}

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

	List<GenericRow> records = new ArrayList<>(NUM_RECORDS);
	_boolGroupValues = new HashMap<>();
	_intGroupValues = new HashMap<>();
	_longGroupValues = new HashMap<>();
	_floatGroupValues = new HashMap<>();
	_doubleGroupValues = new HashMap<>();
	_stringGroupValues = new HashMap<>();
	for (int i = 0; i < NUM_RECORDS; i++) {
	boolean boolValue = RANDOM.nextBoolean();
	int intValue = RANDOM.nextInt(MAX_VALUE);
	long longValue = RANDOM.nextLong();
	float floatValue = RANDOM.nextFloat();
	double doubleValue = RANDOM.nextDouble();
	String strValue = String.valueOf(RANDOM.nextDouble());
	GenericRow record = new GenericRow();
	record.putValue(BOOL_COLUMN, boolValue);
	record.putValue(BOOL_NO_DICT_COLUMN, boolValue);
	record.putValue(INT_COLUMN, intValue);
	record.putValue(INT_MV_COLUMN, new Integer[]{intValue, intValue});
	record.putValue(INT_NO_DICT_COLUMN, intValue);
	record.putValue(LONG_COLUMN, longValue);
	record.putValue(LONG_NO_DICT_COLUMN, longValue);
	record.putValue(FLOAT_COLUMN, floatValue);
	record.putValue(FLOAT_NO_DICT_COLUMN, floatValue);
	record.putValue(DOUBLE_COLUMN, doubleValue);
	record.putValue(DOUBLE_NO_DICT_COLUMN, doubleValue);
	record.putValue(STRING_COLUMN, strValue);
	record.putValue(STRING_NO_DICT_COLUMN, strValue);
	record.putValue(TIME_COLUMN, (long) i);
	if (i == 0) {
	_expectedResultFirstBoolean = boolValue;
	_expectedResultFirstInt = intValue;
	_expectedResultFirstLong = longValue;
	_expectedResultFirstFloat = floatValue;
	_expectedResultFirstDouble = doubleValue;
	_expectedResultFirstString = strValue;
	}
	_boolGroupValues.putIfAbsent(intValue, boolValue);
	_intGroupValues.putIfAbsent(intValue, intValue);
	_longGroupValues.putIfAbsent(intValue, longValue);
	_floatGroupValues.putIfAbsent(intValue, floatValue);
	_doubleGroupValues.putIfAbsent(intValue, doubleValue);
	_stringGroupValues.putIfAbsent(intValue, strValue);
	records.add(record);
	}

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

	ImmutableSegment immutableSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), ReadMode.mmap);
	_indexSegment = immutableSegment;
	_indexSegments = Arrays.asList(immutableSegment, immutableSegment);
	}

	@Test
	public void testAggregationOnly() {
	String query = "SELECT "
	+ "FIRSTWITHTIME(boolColumn, timestampColumn, 'BOOLEAN'), "
	+ "FIRSTWITHTIME(intColumn, timestampColumn, 'INT'), "
	+ "FIRSTWITHTIME(longColumn, timestampColumn, 'LONG'), "
	+ "FIRSTWITHTIME(floatColumn, timestampColumn, 'FLOAT'), "
	+ "FIRSTWITHTIME(doubleColumn, timestampColumn, 'DOUBLE'), "
	+ "FIRSTWITHTIME(stringColumn, timestampColumn, 'STRING') "
	+ "FROM testTable";

	// Inner segment
	AggregationOperator aggregationOperator = getOperator(query);
	AggregationResultsBlock resultsBlock = aggregationOperator.nextBlock();
	QueriesTestUtils.testInnerSegmentExecutionStatistics(aggregationOperator.getExecutionStatistics(), NUM_RECORDS, 0,
	7 * NUM_RECORDS, NUM_RECORDS);
	List<Object> aggregationResult = resultsBlock.getResults();
	assertNotNull(aggregationResult);
	assertEquals(((IntLongPair) aggregationResult.get(0)).getValue() != 0, _expectedResultFirstBoolean.booleanValue());
	assertEquals(((IntLongPair) aggregationResult.get(1)).getValue(), _expectedResultFirstInt);
	assertEquals(((LongLongPair) aggregationResult.get(2)).getValue(), _expectedResultFirstLong);
	assertEquals(((FloatLongPair) aggregationResult.get(3)).getValue(), _expectedResultFirstFloat);
	assertEquals(((DoubleLongPair) aggregationResult.get(4)).getValue(), _expectedResultFirstDouble);
	assertEquals(((StringLongPair) aggregationResult.get(5)).getValue(), _expectedResultFirstString);

	// Inter segments (expect 4 * inner segment result)
	BrokerResponseNative brokerResponse = getBrokerResponse(query);
	DataSchema expectedDataSchema = new DataSchema(new String[]{
	"firstwithtime(boolColumn,timestampColumn,'BOOLEAN')",
	"firstwithtime(intColumn,timestampColumn,'INT')",
	"firstwithtime(longColumn,timestampColumn,'LONG')",
	"firstwithtime(floatColumn,timestampColumn,'FLOAT')",
	"firstwithtime(doubleColumn,timestampColumn,'DOUBLE')",
	"firstwithtime(stringColumn,timestampColumn,'STRING')"
	}, new ColumnDataType[]{
	ColumnDataType.BOOLEAN,
	ColumnDataType.INT,
	ColumnDataType.LONG,
	ColumnDataType.FLOAT,
	ColumnDataType.DOUBLE,
	ColumnDataType.STRING
	});
	Object[] expectedResults = new Object[]{
	_expectedResultFirstBoolean,
	_expectedResultFirstInt,
	_expectedResultFirstLong,
	_expectedResultFirstFloat,
	_expectedResultFirstDouble,
	_expectedResultFirstString
	};
	QueriesTestUtils.testInterSegmentsResult(brokerResponse, 4 * NUM_RECORDS, 0L, 4 * 7 * NUM_RECORDS,
	4 * NUM_RECORDS, new ResultTable(expectedDataSchema, Collections.singletonList(expectedResults)));
	}

	@Test
	public void testAggregationOnlyNoDictionary() {
	String query = "SELECT "
	+ "FIRSTWITHTIME(boolNoDictColumn,timestampColumn,'boolean'), "
	+ "FIRSTWITHTIME(intNoDictColumn,timestampColumn,'int'), "
	+ "FIRSTWITHTIME(longNoDictColumn,timestampColumn,'long'), "
	+ "FIRSTWITHTIME(floatNoDictColumn,timestampColumn,'float'), "
	+ "FIRSTWITHTIME(doubleNoDictColumn,timestampColumn,'double'), "
	+ "FIRSTWITHTIME(stringNoDictColumn,timestampColumn,'string') "
	+ "FROM testTable";

	// Inner segment
	AggregationOperator aggregationOperator = getOperator(query);
	AggregationResultsBlock resultsBlock = aggregationOperator.nextBlock();
	QueriesTestUtils.testInnerSegmentExecutionStatistics(aggregationOperator.getExecutionStatistics(), NUM_RECORDS, 0,
	7 * NUM_RECORDS, NUM_RECORDS);
	List<Object> aggregationResult = resultsBlock.getResults();
	assertNotNull(aggregationResult);
	assertEquals(((IntLongPair) aggregationResult.get(0)).getValue() != 0, _expectedResultFirstBoolean.booleanValue());
	assertEquals(((IntLongPair) aggregationResult.get(1)).getValue(), _expectedResultFirstInt);
	assertEquals(((LongLongPair) aggregationResult.get(2)).getValue(), _expectedResultFirstLong);
	assertEquals(((FloatLongPair) aggregationResult.get(3)).getValue(), _expectedResultFirstFloat);
	assertEquals(((DoubleLongPair) aggregationResult.get(4)).getValue(), _expectedResultFirstDouble);
	assertEquals(((StringLongPair) aggregationResult.get(5)).getValue(), _expectedResultFirstString);

	// Inter segments (expect 4 * inner segment result)
	BrokerResponseNative brokerResponse = getBrokerResponse(query);
	DataSchema expectedDataSchema = new DataSchema(new String[]{
	"firstwithtime(boolNoDictColumn,timestampColumn,'BOOLEAN')",
	"firstwithtime(intNoDictColumn,timestampColumn,'INT')",
	"firstwithtime(longNoDictColumn,timestampColumn,'LONG')",
	"firstwithtime(floatNoDictColumn,timestampColumn,'FLOAT')",
	"firstwithtime(doubleNoDictColumn,timestampColumn,'DOUBLE')",
	"firstwithtime(stringNoDictColumn,timestampColumn,'STRING')"
	}, new ColumnDataType[]{
	ColumnDataType.BOOLEAN,
	ColumnDataType.INT,
	ColumnDataType.LONG,
	ColumnDataType.FLOAT,
	ColumnDataType.DOUBLE,
	ColumnDataType.STRING
	});
	Object[] expectedResults = new Object[]{
	_expectedResultFirstBoolean,
	_expectedResultFirstInt,
	_expectedResultFirstLong,
	_expectedResultFirstFloat,
	_expectedResultFirstDouble,
	_expectedResultFirstString
	};
	QueriesTestUtils.testInterSegmentsResult(brokerResponse, 4 * NUM_RECORDS, 0L, 4 * 7 * NUM_RECORDS,
	4 * NUM_RECORDS, new ResultTable(expectedDataSchema, Collections.singletonList(expectedResults)));
	}

	@Test
	public void testAggregationGroupBySV() {
	String query = "SELECT intColumn AS key, "
	+ "FIRSTWITHTIME(boolColumn,timestampColumn,'boolean') AS v1, "
	+ "FIRSTWITHTIME(intColumn,timestampColumn,'int') AS v2, "
	+ "FIRSTWITHTIME(longColumn,timestampColumn,'long') AS v3, "
	+ "FIRSTWITHTIME(floatColumn,timestampColumn,'float') AS v4, "
	+ "FIRSTWITHTIME(doubleColumn,timestampColumn,'double') AS v5, "
	+ "FIRSTWITHTIME(stringColumn,timestampColumn,'string') AS v6 "
	+ "FROM testTable GROUP BY key";
	verifyAggregationGroupBy(query, 7);
	}

	@Test
	public void testAggregationGroupBySVNoDictionary() {
	String query = "SELECT intNoDictColumn AS key, "
	+ "FIRSTWITHTIME(boolNoDictColumn,timestampColumn,'boolean') AS v1, "
	+ "FIRSTWITHTIME(intNoDictColumn,timestampColumn,'int') AS v2, "
	+ "FIRSTWITHTIME(longNoDictColumn,timestampColumn,'long') AS v3, "
	+ "FIRSTWITHTIME(floatNoDictColumn,timestampColumn,'float') AS v4, "
	+ "FIRSTWITHTIME(doubleNoDictColumn,timestampColumn,'double') AS v5, "
	+ "FIRSTWITHTIME(stringNoDictColumn,timestampColumn,'string') AS v6 "
	+ "FROM testTable GROUP BY key";
	verifyAggregationGroupBy(query, 7);
	}

	@Test
	public void testAggregationGroupByMV() {
	String query = "SELECT intMVColumn AS key, "
	+ "FIRSTWITHTIME(boolColumn,timestampColumn,'boolean') AS v1, "
	+ "FIRSTWITHTIME(intColumn,timestampColumn,'int') AS v2, "
	+ "FIRSTWITHTIME(longColumn,timestampColumn,'long') AS v3, "
	+ "FIRSTWITHTIME(floatColumn,timestampColumn,'float') AS v4, "
	+ "FIRSTWITHTIME(doubleColumn,timestampColumn,'double') AS v5, "
	+ "FIRSTWITHTIME(stringColumn,timestampColumn,'string') AS v6 "
	+ "FROM testTable GROUP BY key";
	verifyAggregationGroupBy(query, 8);
	}

	@Test
	public void testAggregationGroupByMVNoDictionary() {
	String query = "SELECT intMVColumn AS key, "
	+ "FIRSTWITHTIME(boolNoDictColumn,timestampColumn,'boolean') AS v1, "
	+ "FIRSTWITHTIME(intNoDictColumn,timestampColumn,'int') AS v2, "
	+ "FIRSTWITHTIME(longNoDictColumn,timestampColumn,'long') AS v3, "
	+ "FIRSTWITHTIME(floatNoDictColumn,timestampColumn,'float') AS v4, "
	+ "FIRSTWITHTIME(doubleNoDictColumn,timestampColumn,'double') AS v5, "
	+ "FIRSTWITHTIME(stringNoDictColumn,timestampColumn,'string') AS v6 "
	+ "FROM testTable GROUP BY key";
	verifyAggregationGroupBy(query, 8);
	}

	private void verifyAggregationGroupBy(String query, int numProjectedColumns) {
	// Inner segment
	GroupByOperator groupByOperator = getOperator(query);
	GroupByResultsBlock resultsBlock = groupByOperator.nextBlock();
	QueriesTestUtils.testInnerSegmentExecutionStatistics(groupByOperator.getExecutionStatistics(), NUM_RECORDS, 0,
	numProjectedColumns * (long) NUM_RECORDS, NUM_RECORDS);
	AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
	assertNotNull(aggregationGroupByResult);
	int numGroups = 0;
	Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
	while (groupKeyIterator.hasNext()) {
	numGroups++;
	GroupKeyGenerator.GroupKey groupKey = groupKeyIterator.next();
	Integer key = (Integer) groupKey._keys[0];
	assertTrue(_intGroupValues.containsKey(key));
	assertEquals(((IntLongPair) aggregationGroupByResult.getResultForGroupId(0, groupKey._groupId)).getValue() != 0,
	(boolean) _boolGroupValues.get(key));
	assertEquals(((IntLongPair) aggregationGroupByResult.getResultForGroupId(1, groupKey._groupId)).getValue(),
	_intGroupValues.get(key));
	assertEquals(((LongLongPair) aggregationGroupByResult.getResultForGroupId(2, groupKey._groupId)).getValue(),
	_longGroupValues.get(key));
	assertEquals(((FloatLongPair) aggregationGroupByResult.getResultForGroupId(3, groupKey._groupId)).getValue(),
	_floatGroupValues.get(key));
	assertEquals(((DoubleLongPair) aggregationGroupByResult.getResultForGroupId(4, groupKey._groupId)).getValue(),
	_doubleGroupValues.get(key));
	assertEquals(((StringLongPair) aggregationGroupByResult.getResultForGroupId(5, groupKey._groupId)).getValue(),
	_stringGroupValues.get(key));
	}
	assertEquals(numGroups, _intGroupValues.size());

	// Inter segments (expect 4 * inner segment result)
	BrokerResponseNative brokerResponse = getBrokerResponse(query);
	assertEquals(brokerResponse.getNumDocsScanned(), 4 * NUM_RECORDS);
	assertEquals(brokerResponse.getNumEntriesScannedInFilter(), 0);
	assertEquals(brokerResponse.getNumEntriesScannedPostFilter(), 4 * numProjectedColumns * (long) NUM_RECORDS);
	assertEquals(brokerResponse.getTotalDocs(), 4 * NUM_RECORDS);

	ResultTable resultTable = brokerResponse.getResultTable();
	DataSchema expectedDataSchema =
	new DataSchema(new String[]{"key", "v1", "v2", "v3", "v4", "v5", "v6"}, new ColumnDataType[]{
	ColumnDataType.INT, ColumnDataType.BOOLEAN, ColumnDataType.INT, ColumnDataType.LONG,
	ColumnDataType.FLOAT, ColumnDataType.DOUBLE, ColumnDataType.STRING
	});
	assertEquals(resultTable.getDataSchema(), expectedDataSchema);
	List<Object[]> rows = resultTable.getRows();
	assertEquals(rows.size(), 10);
	for (Object[] row : rows) {
	assertEquals(row.length, 7);
	int key = (Integer) row[0];
	assertEquals(row[1], _boolGroupValues.get(key));
	assertEquals(row[2], _intGroupValues.get(key));
	assertEquals(row[3], _longGroupValues.get(key));
	assertEquals(row[4], _floatGroupValues.get(key));
	assertEquals(row[5], _doubleGroupValues.get(key));
	assertEquals(row[6], _stringGroupValues.get(key));
	}
	}

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