You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by ro...@apache.org on 2022/10/09 17:47:46 UTC
[pinot] branch master updated: add array element at function (#9554)
This is an automated email from the ASF dual-hosted git repository.
rongr pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new fd95f75d3e add array element at function (#9554)
fd95f75d3e is described below
commit fd95f75d3ea2f85b00c174e1f60a50649293ebd2
Author: Rong Rong <ro...@apache.org>
AuthorDate: Sun Oct 9 10:47:36 2022 -0700
add array element at function (#9554)
* add array element at function
* fix test
Co-authored-by: Rong Rong <ro...@startree.ai>
---
.../common/function/scalar/ArrayFunctions.java | 26 ++++++
.../ScalarTransformFunctionWrapperTest.java | 92 ++++++++++++++++++++++
2 files changed, 118 insertions(+)
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java
index 85033860ac..454fbe5912 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java
@@ -25,6 +25,7 @@ import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenHashSet;
import it.unimi.dsi.fastutil.objects.ObjectSet;
import java.util.Arrays;
import org.apache.commons.lang3.ArrayUtils;
+import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.spi.annotations.ScalarFunction;
@@ -151,4 +152,29 @@ public class ArrayFunctions {
public static String[] arrayConcatString(String[] values1, String[] values2) {
return ArrayUtils.addAll(values1, values2);
}
+
+ @ScalarFunction
+ public static int arrayElementAtInt(int[] arr, int idx) {
+ return idx > 0 && idx <= arr.length ? arr[idx - 1] : (int) DataSchema.ColumnDataType.INT.getNullPlaceholder();
+ }
+
+ @ScalarFunction
+ public static long arrayElementAtLong(long[] arr, int idx) {
+ return idx > 0 && idx <= arr.length ? arr[idx - 1] : (long) DataSchema.ColumnDataType.LONG.getNullPlaceholder();
+ }
+
+ @ScalarFunction
+ public static float arrayElementAtFloat(float[] arr, int idx) {
+ return idx > 0 && idx <= arr.length ? arr[idx - 1] : (float) DataSchema.ColumnDataType.FLOAT.getNullPlaceholder();
+ }
+
+ @ScalarFunction
+ public static double arrayElementAtDouble(double[] arr, int idx) {
+ return idx > 0 && idx <= arr.length ? arr[idx - 1] : (double) DataSchema.ColumnDataType.DOUBLE.getNullPlaceholder();
+ }
+
+ @ScalarFunction
+ public static String arrayElementAtString(String[] arr, int idx) {
+ return idx > 0 && idx <= arr.length ? arr[idx - 1] : (String) DataSchema.ColumnDataType.STRING.getNullPlaceholder();
+ }
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
index 844fc208de..7053d9a901 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
@@ -24,11 +24,13 @@ import java.nio.charset.StandardCharsets;
import java.text.Normalizer;
import java.util.Arrays;
import java.util.Base64;
+import java.util.Random;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.RequestContextUtils;
+import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.spi.utils.ArrayCopyUtils;
import org.apache.pinot.spi.utils.BigDecimalUtils;
@@ -865,6 +867,96 @@ public class ScalarTransformFunctionWrapperTest extends BaseTransformFunctionTes
testTransformFunctionMV(transformFunction, expectedValues);
}
+ @Test
+ public void testArrayElementAtInt() {
+ Random rand = new Random();
+ int index = rand.nextInt(MAX_NUM_MULTI_VALUES);
+ ExpressionContext expression = RequestContextUtils.getExpression(
+ String.format("array_element_at_int(%s, %d)", INT_MV_COLUMN, index + 1));
+ TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
+ assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
+ assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.INT);
+ assertTrue(transformFunction.getResultMetadata().isSingleValue());
+ int[] expectedValues = new int[NUM_ROWS];
+ for (int i = 0; i < NUM_ROWS; i++) {
+ expectedValues[i] = _intMVValues[i].length > index ? _intMVValues[i][index]
+ : (Integer) DataSchema.ColumnDataType.INT.getNullPlaceholder();
+ }
+ testTransformFunction(transformFunction, expectedValues);
+ }
+
+ @Test
+ public void testArrayElementAtLong() {
+ Random rand = new Random();
+ int index = rand.nextInt(MAX_NUM_MULTI_VALUES);
+ ExpressionContext expression = RequestContextUtils.getExpression(
+ String.format("array_element_at_long(%s, %d)", LONG_MV_COLUMN, index + 1));
+ TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
+ assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
+ assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.LONG);
+ assertTrue(transformFunction.getResultMetadata().isSingleValue());
+ long[] expectedValues = new long[NUM_ROWS];
+ for (int i = 0; i < NUM_ROWS; i++) {
+ expectedValues[i] = _longMVValues[i].length > index ? _longMVValues[i][index]
+ : (Long) DataSchema.ColumnDataType.LONG.getNullPlaceholder();
+ }
+ testTransformFunction(transformFunction, expectedValues);
+ }
+
+ @Test
+ public void testArrayElementAtFloat() {
+ Random rand = new Random();
+ int index = rand.nextInt(MAX_NUM_MULTI_VALUES);
+ ExpressionContext expression = RequestContextUtils.getExpression(
+ String.format("array_element_at_float(%s, %d)", FLOAT_MV_COLUMN, index + 1));
+ TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
+ assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
+ assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.FLOAT);
+ assertTrue(transformFunction.getResultMetadata().isSingleValue());
+ float[] expectedValues = new float[NUM_ROWS];
+ for (int i = 0; i < NUM_ROWS; i++) {
+ expectedValues[i] = _floatMVValues[i].length > index ? _floatMVValues[i][index]
+ : (Float) DataSchema.ColumnDataType.FLOAT.getNullPlaceholder();
+ }
+ testTransformFunction(transformFunction, expectedValues);
+ }
+
+ @Test
+ public void testArrayElementAtDouble() {
+ Random rand = new Random();
+ int index = rand.nextInt(MAX_NUM_MULTI_VALUES);
+ ExpressionContext expression = RequestContextUtils.getExpression(
+ String.format("array_element_at_double(%s, %d)", DOUBLE_MV_COLUMN, index + 1));
+ TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
+ assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
+ assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.DOUBLE);
+ assertTrue(transformFunction.getResultMetadata().isSingleValue());
+ double[] expectedValues = new double[NUM_ROWS];
+ for (int i = 0; i < NUM_ROWS; i++) {
+ expectedValues[i] = _doubleMVValues[i].length > index ? _doubleMVValues[i][index]
+ : (Double) DataSchema.ColumnDataType.DOUBLE.getNullPlaceholder();
+ }
+ testTransformFunction(transformFunction, expectedValues);
+ }
+
+ @Test
+ public void testArrayElementAtString() {
+ Random rand = new Random();
+ int index = rand.nextInt(MAX_NUM_MULTI_VALUES);
+ ExpressionContext expression = RequestContextUtils.getExpression(
+ String.format("array_element_at_string(%s, %d)", STRING_MV_COLUMN, index + 1));
+ TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
+ assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
+ assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.STRING);
+ assertTrue(transformFunction.getResultMetadata().isSingleValue());
+ String[] expectedValues = new String[NUM_ROWS];
+ for (int i = 0; i < NUM_ROWS; i++) {
+ expectedValues[i] = _stringMVValues[i].length > index ? _stringMVValues[i][index]
+ : (String) DataSchema.ColumnDataType.STRING.getNullPlaceholder();
+ }
+ testTransformFunction(transformFunction, expectedValues);
+ }
+
@Test
public void testBase64TransformFunction() {
ExpressionContext expression = RequestContextUtils.getExpression(String.format("toBase64(%s)", BYTES_SV_COLUMN));
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org