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-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org
