diff --git a/core/src/main/java/org/apache/druid/math/expr/Function.java b/core/src/main/java/org/apache/druid/math/expr/Function.java index 65643e226d7..c81790b3bfc 100644 --- a/core/src/main/java/org/apache/druid/math/expr/Function.java +++ b/core/src/main/java/org/apache/druid/math/expr/Function.java @@ -258,7 +258,7 @@ interface Function public void validateArguments(List args) { if (args.size() != 2) { - throw new IAE("Function[%s] needs 2 argument", name()); + throw new IAE("Function[%s] needs 2 arguments", name()); } } @@ -2062,4 +2062,139 @@ interface Function return ExprEval.bestEffortOf(any); } } + + class ArraySliceFunction implements Function + { + @Override + public String name() + { + return "array_slice"; + } + + @Override + public void validateArguments(List args) + { + if (args.size() != 2 && args.size() != 3) { + throw new IAE("Function[%s] needs 2 or 3 arguments", name()); + } + } + + @Override + public ExprEval apply(List args, Expr.ObjectBinding bindings) + { + final ExprEval expr = args.get(0).eval(bindings); + final Object[] array = expr.asArray(); + if (array == null) { + return ExprEval.of(null); + } + + final int start = args.get(1).eval(bindings).asInt(); + int end = array.length; + if (args.size() == 3) { + end = args.get(2).eval(bindings).asInt(); + } + + if (start < 0 || start > array.length || start > end) { + // Arrays.copyOfRange will throw exception in these cases + return ExprEval.of(null); + } + + switch (expr.type()) { + case STRING: + case STRING_ARRAY: + return ExprEval.ofStringArray(Arrays.copyOfRange(expr.asStringArray(), start, end)); + case LONG: + case LONG_ARRAY: + return ExprEval.ofLongArray(Arrays.copyOfRange(expr.asLongArray(), start, end)); + case DOUBLE: + case DOUBLE_ARRAY: + return ExprEval.ofDoubleArray(Arrays.copyOfRange(expr.asDoubleArray(), start, end)); + } + throw new RE("Unable to slice to unknown type %s", expr.type()); + } + + @Override + public Set getScalarInputs(List args) + { + if (args.size() == 3) { + return ImmutableSet.of(args.get(1), args.get(2)); + } else { + return ImmutableSet.of(args.get(1)); + } + } + + @Override + public Set getArrayInputs(List args) + { + return ImmutableSet.of(args.get(0)); + } + } + + class ArrayPrependFunction implements Function + { + @Override + public String name() + { + return "array_prepend"; + } + + @Override + public void validateArguments(List args) + { + if (args.size() != 2) { + throw new IAE("Function[%s] needs 2 arguments", name()); + } + } + + @Override + public ExprEval apply(List args, Expr.ObjectBinding bindings) + { + final ExprEval scalarExpr = args.get(0).eval(bindings); + final ExprEval arrayExpr = args.get(1).eval(bindings); + if (arrayExpr.asArray() == null) { + return ExprEval.of(null); + } + switch (arrayExpr.type()) { + case STRING: + case STRING_ARRAY: + return ExprEval.ofStringArray(this.prepend(scalarExpr.asString(), arrayExpr.asStringArray()).toArray(String[]::new)); + case LONG: + case LONG_ARRAY: + return ExprEval.ofLongArray( + this.prepend( + scalarExpr.isNumericNull() ? null : scalarExpr.asLong(), + arrayExpr.asLongArray()).toArray(Long[]::new + ) + ); + case DOUBLE: + case DOUBLE_ARRAY: + return ExprEval.ofDoubleArray( + this.prepend( + scalarExpr.isNumericNull() ? null : scalarExpr.asDouble(), + arrayExpr.asDoubleArray()).toArray(Double[]::new + ) + ); + } + + throw new RE("Unable to prepend to unknown type %s", arrayExpr.type()); + } + + private Stream prepend(T val, T[] array) + { + List l = new ArrayList<>(Arrays.asList(array)); + l.add(0, val); + return l.stream(); + } + @Override + public Set getScalarInputs(List args) + { + return ImmutableSet.of(args.get(0)); + } + + @Override + public Set getArrayInputs(List args) + { + return ImmutableSet.of(args.get(1)); + } + } } diff --git a/core/src/test/java/org/apache/druid/math/expr/FunctionTest.java b/core/src/test/java/org/apache/druid/math/expr/FunctionTest.java index ec0884cefcf..51df20c6c9b 100644 --- a/core/src/test/java/org/apache/druid/math/expr/FunctionTest.java +++ b/core/src/test/java/org/apache/druid/math/expr/FunctionTest.java @@ -260,6 +260,26 @@ public class FunctionTest assertExpr("cast(['1.0', '2.0', '3.0'], 'LONG_ARRAY')", new Long[]{1L, 2L, 3L}); } + @Test + public void testArraySlice() + { + assertExpr("array_slice([1, 2, 3, 4], 1, 3)", new Long[] {2L, 3L}); + assertExpr("array_slice([1.0, 2.1, 3.2, 4.3], 2)", new Double[] {3.2, 4.3}); + assertExpr("array_slice(['a', 'b', 'c', 'd'], 4, 6)", new String[] {null, null}); + assertExpr("array_slice([1, 2, 3, 4], 2, 2)", new Long[] {}); + assertExpr("array_slice([1, 2, 3, 4], 5, 7)", null); + assertExpr("array_slice([1, 2, 3, 4], 2, 1)", null); + } + + @Test + public void testArrayPrepend() + { + assertExpr("array_prepend(4, [1, 2, 3])", new Long[]{4L, 1L, 2L, 3L}); + assertExpr("array_prepend('bar', [1, 2, 3])", new Long[]{null, 1L, 2L, 3L}); + assertExpr("array_prepend(1, [])", new String[]{"1"}); + assertExpr("array_prepend(1, cast([], 'LONG_ARRAY'))", new Long[]{1L}); + } + private void assertExpr(final String expression, final Object expectedResult) { final Expr expr = Parser.parse(expression, ExprMacroTable.nil()); diff --git a/docs/content/misc/math-expr.md b/docs/content/misc/math-expr.md index 57427a988c3..9a686a1dc42 100644 --- a/docs/content/misc/math-expr.md +++ b/docs/content/misc/math-expr.md @@ -179,6 +179,8 @@ See javadoc of java.lang.Math for detailed explanation for each function. | `array_concat(arr1,arr2)` | concatenates 2 arrays, the resulting array type determined by the type of the first array | | `array_to_string(arr,str)` | joins all elements of arr by the delimiter specified by str | | `string_to_array(str1,str2)` | splits str1 into an array on the delimiter specified by str2 | +| `array_slice(arr,start,end)` | return the subarray of arr from the 0 based index start(inclusive) to end(exclusive), or `null`, if start is less than 0, greater than length of arr or less than end| +| `array_prepend(expr,arr)` | adds expr to arr at the beginning, the resulting array type determined by the type of the array | ## Apply Functions