pinot-core/src/test/java/org/apache/pinot/queries/CastQueriesTest.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.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Random;
 import org.apache.commons.io.FileUtils;
 import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock;
 import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
 import org.apache.pinot.core.operator.query.AggregationOperator;
 import org.apache.pinot.core.operator.query.GroupByOperator;
 import org.apache.pinot.core.operator.query.SelectionOnlyOperator;
 import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
 import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
 import org.apache.pinot.core.query.utils.idset.Roaring64NavigableMapIdSet;
 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;
 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.BeforeClass;
 import org.testng.annotations.Test;

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


 public class CastQueriesTest extends BaseQueriesTest {

   private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "CastQueriesTest");
   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 = 1000;
   private static final int BUCKET_SIZE = 8;
   private static final int NUM_ENTRY_MV = 2;
   private static final String CLASSIFICATION_COLUMN = "class";
   private static final String X_COL = "x";
   private static final String Y_COL = "y";
   private static final String STRING_MV_COL = "stringMvCol";

   private long[][] _expectedLongValues = new long[NUM_RECORDS][];

   private static final Schema SCHEMA = new Schema.SchemaBuilder()
       .addSingleValueDimension(X_COL, FieldSpec.DataType.DOUBLE)
       .addSingleValueDimension(Y_COL, FieldSpec.DataType.DOUBLE)
       .addSingleValueDimension(CLASSIFICATION_COLUMN, FieldSpec.DataType.STRING)
       .addMultiValueDimension(STRING_MV_COL, 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;
   }

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

     List<GenericRow> records = new ArrayList<>(NUM_RECORDS);
     for (int i = 0; i < NUM_RECORDS; i++) {
       GenericRow record = new GenericRow();
       record.putValue(X_COL, 0.5);
       record.putValue(Y_COL, 0.25);
       record.putValue(CLASSIFICATION_COLUMN, "" + (i % BUCKET_SIZE));
       String[] mvCol = new String[NUM_ENTRY_MV];
       _expectedLongValues[i] = new long[NUM_ENTRY_MV];
       for (int j = 0; j < mvCol.length; j++) {
         long longVal = RANDOM.nextLong();
         _expectedLongValues[i][j] = longVal;
         mvCol[j] = String.valueOf(longVal);
       }
       record.putValue(STRING_MV_COL, mvCol);
       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 testCastSum() {
     String query = "select cast(sum(" + X_COL + ") as int), "
         + "cast(sum(" + Y_COL + ") as int) "
         + "from " + RAW_TABLE_NAME;
     AggregationOperator aggregationOperator = getOperator(query);
     List<Object> aggregationResult = aggregationOperator.nextBlock().getResults();
     assertNotNull(aggregationResult);
     assertEquals(aggregationResult.size(), 2);
     assertEquals(((Number) aggregationResult.get(0)).intValue(), NUM_RECORDS / 2);
     assertEquals(((Number) aggregationResult.get(1)).intValue(), NUM_RECORDS / 4);
   }

   @Test
   public void testCastMV() {
     String query = String.format("select cast(%s as LONG) as col1 from %s limit 100", STRING_MV_COL, RAW_TABLE_NAME);
     SelectionResultsBlock block = ((SelectionOnlyOperator) getOperator(query)).nextBlock();
     List<Object[]> results = (List<Object[]>) block.getRows();
     for (int i = 0; i < 100; i++) {
       Object[] row = results.get(i);
       long[] cast = (long[]) row[0];
       assertEquals(cast, _expectedLongValues[i]);
     }

     query = String.format("select id_set(cast(%s as LONG)) as col1 from %s limit 100", STRING_MV_COL, RAW_TABLE_NAME);
     AggregationResultsBlock aggregationResultsBlock = ((AggregationOperator) getOperator(query)).nextBlock();
     List<Object> aggregationResultsBlockResults = aggregationResultsBlock.getResults();
     assertTrue(aggregationResultsBlockResults.get(0) instanceof Roaring64NavigableMapIdSet);
     Roaring64NavigableMapIdSet idSet = (Roaring64NavigableMapIdSet) aggregationResultsBlockResults.get(0);
     for (int i = 0; i < NUM_RECORDS; i++) {
       for (int j = 0; j < NUM_ENTRY_MV; j++) {
         assertTrue(idSet.contains(_expectedLongValues[i][j]));
       }
     }
   }

   @Test
   public void testCastSumGroupBy() {
     String query = "select cast(sum(" + X_COL + ") as int), "
         + "cast(sum(" + Y_COL + ") as int) "
         + "from " + RAW_TABLE_NAME + " "
         + "group by " + CLASSIFICATION_COLUMN;
     GroupByOperator groupByOperator = getOperator(query);
     AggregationGroupByResult result = groupByOperator.nextBlock().getAggregationGroupByResult();
     assertNotNull(result);
     Iterator<GroupKeyGenerator.GroupKey> it = result.getGroupKeyIterator();
     while (it.hasNext()) {
       GroupKeyGenerator.GroupKey groupKey = it.next();
       Object aggregate = result.getResultForGroupId(0, groupKey._groupId);
       assertEquals(((Number) aggregate).intValue(), NUM_RECORDS / (2 * BUCKET_SIZE));
       aggregate = result.getResultForGroupId(1, groupKey._groupId);
       assertEquals(((Number) aggregate).intValue(), NUM_RECORDS / (4 * BUCKET_SIZE));
     }
   }

   @Test
   public void testCastFilterAndProject() {
     String query = "select cast(" + CLASSIFICATION_COLUMN + " as int)"
         + " from " + RAW_TABLE_NAME
         + " where " + CLASSIFICATION_COLUMN + " = cast(0 as string) limit " + NUM_RECORDS;
     SelectionOnlyOperator selectionOperator = getOperator(query);
     Collection<Object[]> result = selectionOperator.nextBlock().getRows();
     assertNotNull(result);
     assertEquals(result.size(), NUM_RECORDS / BUCKET_SIZE);
     for (Object[] row : result) {
       assertEquals(row.length, 1);
       assertEquals(row[0], 0);
     }
   }
 }
	/**
	* 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.util.ArrayList;
	import java.util.Arrays;
	import java.util.Collection;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Random;
	import org.apache.commons.io.FileUtils;
	import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock;
	import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
	import org.apache.pinot.core.operator.query.AggregationOperator;
	import org.apache.pinot.core.operator.query.GroupByOperator;
	import org.apache.pinot.core.operator.query.SelectionOnlyOperator;
	import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
	import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
	import org.apache.pinot.core.query.utils.idset.Roaring64NavigableMapIdSet;
	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;
	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.BeforeClass;
	import org.testng.annotations.Test;

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


	public class CastQueriesTest extends BaseQueriesTest {

	private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "CastQueriesTest");
	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 = 1000;
	private static final int BUCKET_SIZE = 8;
	private static final int NUM_ENTRY_MV = 2;
	private static final String CLASSIFICATION_COLUMN = "class";
	private static final String X_COL = "x";
	private static final String Y_COL = "y";
	private static final String STRING_MV_COL = "stringMvCol";

	private long[][] _expectedLongValues = new long[NUM_RECORDS][];

	private static final Schema SCHEMA = new Schema.SchemaBuilder()
	.addSingleValueDimension(X_COL, FieldSpec.DataType.DOUBLE)
	.addSingleValueDimension(Y_COL, FieldSpec.DataType.DOUBLE)
	.addSingleValueDimension(CLASSIFICATION_COLUMN, FieldSpec.DataType.STRING)
	.addMultiValueDimension(STRING_MV_COL, 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;
	}

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

	List<GenericRow> records = new ArrayList<>(NUM_RECORDS);
	for (int i = 0; i < NUM_RECORDS; i++) {
	GenericRow record = new GenericRow();
	record.putValue(X_COL, 0.5);
	record.putValue(Y_COL, 0.25);
	record.putValue(CLASSIFICATION_COLUMN, "" + (i % BUCKET_SIZE));
	String[] mvCol = new String[NUM_ENTRY_MV];
	_expectedLongValues[i] = new long[NUM_ENTRY_MV];
	for (int j = 0; j < mvCol.length; j++) {
	long longVal = RANDOM.nextLong();
	_expectedLongValues[i][j] = longVal;
	mvCol[j] = String.valueOf(longVal);
	}
	record.putValue(STRING_MV_COL, mvCol);
	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 testCastSum() {
	String query = "select cast(sum(" + X_COL + ") as int), "
	+ "cast(sum(" + Y_COL + ") as int) "
	+ "from " + RAW_TABLE_NAME;
	AggregationOperator aggregationOperator = getOperator(query);
	List<Object> aggregationResult = aggregationOperator.nextBlock().getResults();
	assertNotNull(aggregationResult);
	assertEquals(aggregationResult.size(), 2);
	assertEquals(((Number) aggregationResult.get(0)).intValue(), NUM_RECORDS / 2);
	assertEquals(((Number) aggregationResult.get(1)).intValue(), NUM_RECORDS / 4);
	}

	@Test
	public void testCastMV() {
	String query = String.format("select cast(%s as LONG) as col1 from %s limit 100", STRING_MV_COL, RAW_TABLE_NAME);
	SelectionResultsBlock block = ((SelectionOnlyOperator) getOperator(query)).nextBlock();
	List<Object[]> results = (List<Object[]>) block.getRows();
	for (int i = 0; i < 100; i++) {
	Object[] row = results.get(i);
	long[] cast = (long[]) row[0];
	assertEquals(cast, _expectedLongValues[i]);
	}

	query = String.format("select id_set(cast(%s as LONG)) as col1 from %s limit 100", STRING_MV_COL, RAW_TABLE_NAME);
	AggregationResultsBlock aggregationResultsBlock = ((AggregationOperator) getOperator(query)).nextBlock();
	List<Object> aggregationResultsBlockResults = aggregationResultsBlock.getResults();
	assertTrue(aggregationResultsBlockResults.get(0) instanceof Roaring64NavigableMapIdSet);
	Roaring64NavigableMapIdSet idSet = (Roaring64NavigableMapIdSet) aggregationResultsBlockResults.get(0);
	for (int i = 0; i < NUM_RECORDS; i++) {
	for (int j = 0; j < NUM_ENTRY_MV; j++) {
	assertTrue(idSet.contains(_expectedLongValues[i][j]));
	}
	}
	}

	@Test
	public void testCastSumGroupBy() {
	String query = "select cast(sum(" + X_COL + ") as int), "
	+ "cast(sum(" + Y_COL + ") as int) "
	+ "from " + RAW_TABLE_NAME + " "
	+ "group by " + CLASSIFICATION_COLUMN;
	GroupByOperator groupByOperator = getOperator(query);
	AggregationGroupByResult result = groupByOperator.nextBlock().getAggregationGroupByResult();
	assertNotNull(result);
	Iterator<GroupKeyGenerator.GroupKey> it = result.getGroupKeyIterator();
	while (it.hasNext()) {
	GroupKeyGenerator.GroupKey groupKey = it.next();
	Object aggregate = result.getResultForGroupId(0, groupKey._groupId);
	assertEquals(((Number) aggregate).intValue(), NUM_RECORDS / (2 * BUCKET_SIZE));
	aggregate = result.getResultForGroupId(1, groupKey._groupId);
	assertEquals(((Number) aggregate).intValue(), NUM_RECORDS / (4 * BUCKET_SIZE));
	}
	}

	@Test
	public void testCastFilterAndProject() {
	String query = "select cast(" + CLASSIFICATION_COLUMN + " as int)"
	+ " from " + RAW_TABLE_NAME
	+ " where " + CLASSIFICATION_COLUMN + " = cast(0 as string) limit " + NUM_RECORDS;
	SelectionOnlyOperator selectionOperator = getOperator(query);
	Collection<Object[]> result = selectionOperator.nextBlock().getRows();
	assertNotNull(result);
	assertEquals(result.size(), NUM_RECORDS / BUCKET_SIZE);
	for (Object[] row : result) {
	assertEquals(row.length, 1);
	assertEquals(row[0], 0);
	}
	}
	}