Add array_slice and array_unshift function expr (#7950)

* add array_slice and array_unshift function expr

* feedback address

core/src/main/java/org/apache/druid/math/expr/Function.java

@@ -258,7 +258,7 @@ interface Function
     public void validateArguments(List<Expr> 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<Expr> args)
+    {
+      if (args.size() != 2 && args.size() != 3) {
+        throw new IAE("Function[%s] needs 2 or 3 arguments", name());
+      }
+    }
+
+    @Override
+    public ExprEval apply(List<Expr> 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<Expr> getScalarInputs(List<Expr> args)
+    {
+      if (args.size() == 3) {
+        return ImmutableSet.of(args.get(1), args.get(2));
+      } else {
+        return ImmutableSet.of(args.get(1));
+      }
+    }
+
+    @Override
+    public Set<Expr> getArrayInputs(List<Expr> args)
+    {
+      return ImmutableSet.of(args.get(0));
+    }
+  }
+
+  class ArrayPrependFunction implements Function
+  {
+    @Override
+    public String name()
+    {
+      return "array_prepend";
+    }
+
+    @Override
+    public void validateArguments(List<Expr> args)
+    {
+      if (args.size() != 2) {
+        throw new IAE("Function[%s] needs 2 arguments", name());
+      }
+    }
+
+    @Override
+    public ExprEval apply(List<Expr> 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 <T> Stream<T> prepend(T val, T[] array)
+    {
+      List<T> l = new ArrayList<>(Arrays.asList(array));
+      l.add(0, val);
+      return l.stream();
+    }
+    @Override
+    public Set<Expr> getScalarInputs(List<Expr> args)
+    {
+      return ImmutableSet.of(args.get(0));
+    }
+
+    @Override
+    public Set<Expr> getArrayInputs(List<Expr> args)
+    {
+      return ImmutableSet.of(args.get(1));
+    }
+  }
 }

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());

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