61yao commented on code in PR #10594:
URL: https://github.com/apache/pinot/pull/10594#discussion_r1178576003
##########
pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/CoalesceTransformFunctionTest.java:
##########
Review Comment:
good point. and it actually catches a bug..
##########
pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/CoalesceTransformFunctionTest.java:
##########
@@ -18,456 +18,78 @@
*/
package org.apache.pinot.core.operator.transform.function;
-import java.io.File;
-import java.math.BigDecimal;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Random;
-import java.util.Set;
-import org.apache.commons.io.FileUtils;
-import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.RequestContextUtils;
-import org.apache.pinot.core.operator.DocIdSetOperator;
-import org.apache.pinot.core.operator.ProjectionOperator;
-import org.apache.pinot.core.operator.blocks.ProjectionBlock;
-import org.apache.pinot.core.operator.filter.MatchAllFilterOperator;
-import org.apache.pinot.core.plan.DocIdSetPlanNode;
-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.IndexSegment;
-import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
-import org.apache.pinot.segment.spi.datasource.DataSource;
-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.roaringbitmap.RoaringBitmap;
import org.testng.Assert;
-import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
public class CoalesceTransformFunctionTest extends BaseTransformFunctionTest {
- private static final String ENABLE_NULL_SEGMENT_NAME = "testSegment1";
- private static final String DISABLE_NULL_SEGMENT_NAME = "testSegment2";
- private static final Random RANDOM = new Random();
-
- private static final int NUM_ROWS = 1000;
- private static final String INT_SV_COLUMN1 = "intSV1";
- private static final String INT_SV_COLUMN2 = "intSV2";
- private static final String STRING_SV_COLUMN1 = "StringSV1";
- private static final String STRING_SV_COLUMN2 = "StringSV2";
- private static final String BIG_DECIMAL_SV_COLUMN1 = "BigDecimalSV1";
- private static final String BIG_DECIMAL_SV_COLUMN2 = "BigDecimalSV2";
- private static final String LONG_SV_COLUMN1 = "LongSV1";
- private static final String LONG_SV_COLUMN2 = "LongSV2";
- private static final String DOUBLE_SV_COLUMN1 = "DoubleSV1";
- private static final String DOUBLE_SV_COLUMN2 = "DoubleSV2";
-
- private static final String FLOAT_SV_COLUMN1 = "FloatSV1";
- private static final String FLOAT_SV_COLUMN2 = "FLoatSV2";
- private final int[] _intSVValues = new int[NUM_ROWS];
- private final double[] _doubleValues = new double[NUM_ROWS];
- private final float[] _floatValues = new float[NUM_ROWS];
- private final String[] _stringSVValues = new String[NUM_ROWS];
- private Map<String, DataSource> _enableNullDataSourceMap;
- private Map<String, DataSource> _disableNullDataSourceMap;
- private ProjectionBlock _enableNullProjectionBlock;
- private ProjectionBlock _disableNullProjectionBlock;
- // Mod decides whether the first column of the same type should be null.
- private static final int NULL_MOD1 = 3;
- // Mod decides whether the second column of the same type should be null.
- private static final int NULL_MOD2 = 5;
- // Difference between two same type numeric columns.
- private static final int INT_VALUE_SHIFT = 2;
- private static final double DOUBLE_VALUE_SHIFT = 0.1;
- private static final float FLOAT_VALUE_SHIFT = 0.1f;
-
- // Suffix for second string column.
- private static final String SUFFIX = "column2";
-
- private static String getIndexDirPath(String segmentName) {
- return FileUtils.getTempDirectoryPath() + File.separator + segmentName;
- }
-
- private static Map<String, DataSource> getDataSourceMap(Schema schema,
List<GenericRow> rows, String segmentName)
- throws Exception {
- TableConfig tableConfig =
- new
TableConfigBuilder(TableType.OFFLINE).setTableName(segmentName).setNullHandlingEnabled(true).build();
- SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig,
schema);
- config.setOutDir(getIndexDirPath(segmentName));
- config.setSegmentName(segmentName);
- SegmentIndexCreationDriverImpl driver = new
SegmentIndexCreationDriverImpl();
- driver.init(config, new GenericRowRecordReader(rows));
- driver.build();
- IndexSegment indexSegment =
- ImmutableSegmentLoader.load(new File(getIndexDirPath(segmentName),
segmentName), ReadMode.heap);
- Set<String> columnNames = indexSegment.getPhysicalColumnNames();
- Map<String, DataSource> enableNullDataSourceMap = new
HashMap<>(columnNames.size());
- for (String columnName : columnNames) {
- enableNullDataSourceMap.put(columnName,
indexSegment.getDataSource(columnName));
- }
- return enableNullDataSourceMap;
- }
-
- private static ProjectionBlock getProjectionBlock(Map<String, DataSource>
dataSourceMap) {
- return new ProjectionOperator(dataSourceMap,
- new DocIdSetOperator(new MatchAllFilterOperator(NUM_ROWS),
DocIdSetPlanNode.MAX_DOC_PER_CALL)).nextBlock();
- }
-
- private static boolean isColumn1Null(int i) {
- return i % NULL_MOD1 == 0;
- }
-
- private static boolean isColumn2Null(int i) {
- return i % NULL_MOD2 == 0;
- }
-
- @BeforeClass
- public void setup()
- throws Exception {
- // Set up two tables: one with null option enable, the other with null
option disable.
- // Each table one string column, and one int column with some rows set to
null.
- FileUtils.deleteQuietly(new
File(getIndexDirPath(DISABLE_NULL_SEGMENT_NAME)));
- FileUtils.deleteQuietly(new
File(getIndexDirPath(ENABLE_NULL_SEGMENT_NAME)));
- for (int i = 0; i < NUM_ROWS; i++) {
- _intSVValues[i] = RANDOM.nextInt();
- _doubleValues[i] = RANDOM.nextDouble();
- _floatValues[i] = RANDOM.nextFloat();
- _stringSVValues[i] = "a" + RANDOM.nextInt();
- }
- List<GenericRow> rows = new ArrayList<>(NUM_ROWS);
- for (int i = 0; i < NUM_ROWS; i++) {
- Map<String, Object> map = new HashMap<>();
- map.put(INT_SV_COLUMN1, _intSVValues[i]);
- map.put(INT_SV_COLUMN2, _intSVValues[i] + INT_VALUE_SHIFT);
- map.put(DOUBLE_SV_COLUMN1, _doubleValues[i]);
- map.put(DOUBLE_SV_COLUMN2, _doubleValues[i] + DOUBLE_VALUE_SHIFT);
- map.put(FLOAT_SV_COLUMN1, _floatValues[i]);
- map.put(FLOAT_SV_COLUMN2, _floatValues[i] + FLOAT_VALUE_SHIFT);
- map.put(STRING_SV_COLUMN1, _stringSVValues[i]);
- map.put(STRING_SV_COLUMN2, _stringSVValues[i] + SUFFIX);
- map.put(BIG_DECIMAL_SV_COLUMN1, BigDecimal.valueOf(_intSVValues[i]));
- map.put(BIG_DECIMAL_SV_COLUMN2, BigDecimal.valueOf(_intSVValues[i] +
INT_VALUE_SHIFT));
- map.put(LONG_SV_COLUMN1, _intSVValues[i]);
- map.put(LONG_SV_COLUMN2, _intSVValues[i] + INT_VALUE_SHIFT);
-
- if (isColumn1Null(i)) {
- map.put(INT_SV_COLUMN1, null);
- map.put(STRING_SV_COLUMN1, null);
- map.put(BIG_DECIMAL_SV_COLUMN1, null);
- map.put(LONG_SV_COLUMN1, null);
- map.put(DOUBLE_SV_COLUMN1, null);
- map.put(FLOAT_SV_COLUMN1, null);
- }
- if (isColumn2Null(i)) {
- map.put(INT_SV_COLUMN2, null);
- map.put(STRING_SV_COLUMN2, null);
- map.put(LONG_SV_COLUMN2, null);
- map.put(BIG_DECIMAL_SV_COLUMN2, null);
- map.put(DOUBLE_SV_COLUMN2, null);
- map.put(FLOAT_SV_COLUMN2, null);
- }
- GenericRow row = new GenericRow();
- row.init(map);
- rows.add(row);
- }
- Schema schema = new
Schema.SchemaBuilder().addSingleValueDimension(INT_SV_COLUMN1,
FieldSpec.DataType.INT)
- .addSingleValueDimension(INT_SV_COLUMN2, FieldSpec.DataType.INT)
- .addSingleValueDimension(STRING_SV_COLUMN1, FieldSpec.DataType.STRING)
- .addSingleValueDimension(STRING_SV_COLUMN2, FieldSpec.DataType.STRING)
- .addSingleValueDimension(LONG_SV_COLUMN1, FieldSpec.DataType.LONG)
- .addSingleValueDimension(LONG_SV_COLUMN2, FieldSpec.DataType.LONG)
- .addSingleValueDimension(DOUBLE_SV_COLUMN1, FieldSpec.DataType.DOUBLE)
- .addSingleValueDimension(DOUBLE_SV_COLUMN2, FieldSpec.DataType.DOUBLE)
- .addSingleValueDimension(FLOAT_SV_COLUMN1, FieldSpec.DataType.FLOAT)
- .addSingleValueDimension(FLOAT_SV_COLUMN2, FieldSpec.DataType.FLOAT)
- .addMetric(BIG_DECIMAL_SV_COLUMN1, FieldSpec.DataType.BIG_DECIMAL)
- .addMetric(BIG_DECIMAL_SV_COLUMN2,
FieldSpec.DataType.BIG_DECIMAL).build();
- _enableNullDataSourceMap = getDataSourceMap(schema, rows,
ENABLE_NULL_SEGMENT_NAME);
- _enableNullProjectionBlock = getProjectionBlock(_enableNullDataSourceMap);
- _disableNullDataSourceMap = getDataSourceMap(schema, rows,
DISABLE_NULL_SEGMENT_NAME);
- _disableNullProjectionBlock =
getProjectionBlock(_disableNullDataSourceMap);
- }
-
- private static void testIntTransformFunction(ExpressionContext expression,
int[] expectedValues,
- ProjectionBlock projectionBlock, Map<String, DataSource> dataSourceMap)
- throws Exception {
- int[] actualValues =
- TransformFunctionFactory.get(expression,
dataSourceMap).transformToIntValuesSV(projectionBlock);
- for (int i = 0; i < NUM_ROWS; i++) {
- Assert.assertEquals(actualValues[i], expectedValues[i]);
- }
- }
-
- private static void testStringTransformFunction(ExpressionContext
expression, String[] expectedValues,
- ProjectionBlock projectionBlock, Map<String, DataSource> dataSourceMap)
- throws Exception {
- String[] actualValues =
- TransformFunctionFactory.get(expression,
dataSourceMap).transformToStringValuesSV(projectionBlock);
- for (int i = 0; i < NUM_ROWS; i++) {
- Assert.assertEquals(actualValues[i], expectedValues[i]);
- }
- }
-
- private static void testLongTransformFunction(ExpressionContext expression,
long[] expectedValues,
- ProjectionBlock projectionBlock, Map<String, DataSource> dataSourceMap)
- throws Exception {
- long[] actualValues =
- TransformFunctionFactory.get(expression,
dataSourceMap).transformToLongValuesSV(projectionBlock);
- for (int i = 0; i < NUM_ROWS; i++) {
- Assert.assertEquals(actualValues[i], expectedValues[i]);
- }
- }
-
- private static void testDoubleTransformFunction(ExpressionContext
expression, double[] expectedValues,
- ProjectionBlock projectionBlock, Map<String, DataSource> dataSourceMap)
- throws Exception {
- double[] actualValues =
- TransformFunctionFactory.get(expression,
dataSourceMap).transformToDoubleValuesSV(projectionBlock);
- for (int i = 0; i < NUM_ROWS; i++) {
- Assert.assertEquals(actualValues[i], expectedValues[i]);
- }
- }
-
- private static void testFloatTransformFunction(ExpressionContext expression,
float[] expectedValues,
- ProjectionBlock projectionBlock, Map<String, DataSource> dataSourceMap)
- throws Exception {
- float[] actualValues =
- TransformFunctionFactory.get(expression,
dataSourceMap).transformToFloatValuesSV(projectionBlock);
- for (int i = 0; i < NUM_ROWS; i++) {
- Assert.assertEquals(actualValues[i], expectedValues[i]);
- }
- }
-
- private static void testBigDecimalTransformFunction(ExpressionContext
expression, BigDecimal[] expectedValues,
- ProjectionBlock projectionBlock, Map<String, DataSource> dataSourceMap)
- throws Exception {
- BigDecimal[] actualValues =
- TransformFunctionFactory.get(expression,
dataSourceMap).transformToBigDecimalValuesSV(projectionBlock);
- for (int i = 0; i < NUM_ROWS; i++) {
- Assert.assertEquals(actualValues[i], expectedValues[i]);
- }
- }
-
- // Test the Coalesce on two Int columns where one or the other or both can
be null.
@Test
- public void testCoalesceIntColumns()
- throws Exception {
- ExpressionContext coalesceExpr =
- RequestContextUtils.getExpression(String.format("COALESCE(%s,%s)",
INT_SV_COLUMN1, INT_SV_COLUMN2));
- TransformFunction coalesceTransformFunction =
TransformFunctionFactory.get(coalesceExpr, _enableNullDataSourceMap);
- Assert.assertEquals(coalesceTransformFunction.getName(), "coalesce");
- int[] expectedResults = new int[NUM_ROWS];
- for (int i = 0; i < NUM_ROWS; i++) {
- if (isColumn1Null(i) && isColumn2Null(i)) {
- expectedResults[i] = CoalesceTransformFunction.NULL_INT;
- } else if (isColumn1Null(i)) {
- expectedResults[i] = _intSVValues[i] + INT_VALUE_SHIFT;
- } else if (isColumn2Null(i)) {
- expectedResults[i] = _intSVValues[i];
- } else {
- expectedResults[i] = _intSVValues[i];
- }
- }
- testIntTransformFunction(coalesceExpr, expectedResults,
_enableNullProjectionBlock, _enableNullDataSourceMap);
- testIntTransformFunction(coalesceExpr, expectedResults,
_disableNullProjectionBlock, _disableNullDataSourceMap);
- }
+ public void testCoalesceIntColumns() {
+ TransformFunction coalesceFunc = TransformFunctionFactory.get(
+ RequestContextUtils.getExpression(String.format("COALESCE(%s,%s)",
INT_SV_NULL_COLUMN, LONG_SV_COLUMN)),
+ _dataSourceMap);
- // Test the Coalesce on two long columns where one or the other or both can
be null.
- @Test
- public void testCoalesceLongColumns()
- throws Exception {
- ExpressionContext coalesceExpr =
- RequestContextUtils.getExpression(String.format("COALESCE(%s,%s)",
LONG_SV_COLUMN1, LONG_SV_COLUMN2));
- TransformFunction coalesceTransformFunction =
TransformFunctionFactory.get(coalesceExpr, _enableNullDataSourceMap);
- Assert.assertEquals(coalesceTransformFunction.getName(), "coalesce");
long[] expectedResults = new long[NUM_ROWS];
for (int i = 0; i < NUM_ROWS; i++) {
- if (isColumn1Null(i) && isColumn2Null(i)) {
- expectedResults[i] = CoalesceTransformFunction.NULL_LONG;
- } else if (isColumn1Null(i)) {
- expectedResults[i] = _intSVValues[i] + INT_VALUE_SHIFT;
- } else if (isColumn2Null(i)) {
- expectedResults[i] = _intSVValues[i];
- } else {
+ if (i % 2 == 0) {
expectedResults[i] = _intSVValues[i];
- }
- }
- testLongTransformFunction(coalesceExpr, expectedResults,
_enableNullProjectionBlock, _enableNullDataSourceMap);
- testLongTransformFunction(coalesceExpr, expectedResults,
_disableNullProjectionBlock, _disableNullDataSourceMap);
- }
-
- // Test the Coalesce on two float columns where one or the other or both can
be null.
- @Test
- public void testCoalesceFloatColumns()
- throws Exception {
- ExpressionContext coalesceExpr =
- RequestContextUtils.getExpression(String.format("COALESCE(%s,%s)",
FLOAT_SV_COLUMN1, FLOAT_SV_COLUMN2));
- TransformFunction coalesceTransformFunction =
TransformFunctionFactory.get(coalesceExpr, _enableNullDataSourceMap);
- Assert.assertEquals(coalesceTransformFunction.getName(), "coalesce");
- float[] expectedResults = new float[NUM_ROWS];
- for (int i = 0; i < NUM_ROWS; i++) {
- if (isColumn1Null(i) && isColumn2Null(i)) {
- expectedResults[i] = CoalesceTransformFunction.NULL_FLOAT;
- } else if (isColumn1Null(i)) {
- expectedResults[i] = _floatValues[i] + FLOAT_VALUE_SHIFT;
- } else if (isColumn2Null(i)) {
- expectedResults[i] = _floatValues[i];
- } else {
- expectedResults[i] = _floatValues[i];
- }
- }
- testFloatTransformFunction(coalesceExpr, expectedResults,
_enableNullProjectionBlock, _enableNullDataSourceMap);
- testFloatTransformFunction(coalesceExpr, expectedResults,
_disableNullProjectionBlock, _disableNullDataSourceMap);
- }
-
- // Test the Coalesce on two double columns where one or the other or both
can be null.
- @Test
- public void testCoalesceDoubleColumns()
- throws Exception {
- ExpressionContext coalesceExpr =
- RequestContextUtils.getExpression(String.format("COALESCE(%s,%s)",
DOUBLE_SV_COLUMN1, DOUBLE_SV_COLUMN2));
- TransformFunction coalesceTransformFunction =
TransformFunctionFactory.get(coalesceExpr, _enableNullDataSourceMap);
- Assert.assertEquals(coalesceTransformFunction.getName(), "coalesce");
- double[] expectedResults = new double[NUM_ROWS];
- for (int i = 0; i < NUM_ROWS; i++) {
- if (isColumn1Null(i) && isColumn2Null(i)) {
- expectedResults[i] = CoalesceTransformFunction.NULL_DOUBLE;
- } else if (isColumn1Null(i)) {
- expectedResults[i] = _doubleValues[i] + DOUBLE_VALUE_SHIFT;
- } else if (isColumn2Null(i)) {
- expectedResults[i] = _doubleValues[i];
} else {
- expectedResults[i] = _doubleValues[i];
+ expectedResults[i] = _longSVValues[i];
}
}
- testDoubleTransformFunction(coalesceExpr, expectedResults,
_enableNullProjectionBlock, _enableNullDataSourceMap);
- testDoubleTransformFunction(coalesceExpr, expectedResults,
_disableNullProjectionBlock, _disableNullDataSourceMap);
+ testTransformFunction(coalesceFunc, expectedResults);
}
- // Test the Coalesce on two big decimal columns where one or the other or
both can be null.
@Test
- public void testCoalesceBigDecimalColumns()
- throws Exception {
- ExpressionContext coalesceExpr = RequestContextUtils.getExpression(
- String.format("COALESCE(%s,%s)", BIG_DECIMAL_SV_COLUMN1,
BIG_DECIMAL_SV_COLUMN2));
- TransformFunction coalesceTransformFunction =
TransformFunctionFactory.get(coalesceExpr, _enableNullDataSourceMap);
- Assert.assertEquals(coalesceTransformFunction.getName(), "coalesce");
- BigDecimal[] expectedResults = new BigDecimal[NUM_ROWS];
+ public void testCoalesceIntColumnsAndLiterals() {
+ TransformFunction coalesceFunc = TransformFunctionFactory.get(
+ RequestContextUtils.getExpression(String.format("COALESCE(%s,%s)",
INT_SV_NULL_COLUMN, 314)),
Review Comment:
done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org
