Add type coercion and null check to left, right, repeat exprs. (apach…

…e#16480) * Add type coercion and null check to left, right, repeat exprs. These exprs shouldn't validate types; they should coerce types. Coercion is typical behavior for functions because it enables schema evolution. The functions are also modified to check isNumericNull on the right-hand argument. This was missing previously, which would erroneously cause nulls to be treated as zeroes. * Fix tests.
hevansDev · Aug 29, 2024 · e8750bc · e8750bc
1 parent f8a45a9
commit e8750bc
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 87 deletions.
diff --git a/processing/src/main/java/org/apache/druid/math/expr/Function.java b/processing/src/main/java/org/apache/druid/math/expr/Function.java
@@ -368,20 +368,24 @@ public boolean canVectorize(Expr.InputBindingInspector inspector, List<Expr> arg
 
   /**
    * Base class for a 2 variable input {@link Function} whose first argument is a {@link ExprType#STRING} and second
-   * argument is {@link ExprType#LONG}
+   * argument is {@link ExprType#LONG}. These functions return null if either argument is null.
    */
   abstract class StringLongFunction extends BivariateFunction
   {
     @Override
     protected final ExprEval eval(ExprEval x, ExprEval y)
     {
-      if (!x.type().is(ExprType.STRING) || !y.type().is(ExprType.LONG)) {
-        throw validationFailed("needs a STRING as first argument and a LONG as second argument");
+      final String xString = x.asString();
+      if (xString == null) {
+        return ExprEval.of(null);
+      }
+      if (y.isNumericNull()) {
+        return ExprEval.of(null);
       }
-      return eval(x.asString(), y.asInt());
+      return eval(xString, y.asLong());
     }
 
-    protected abstract ExprEval eval(@Nullable String x, int y);
+    protected abstract ExprEval eval(String x, long y);
   }
 
   /**
@@ -2824,16 +2828,14 @@ public ExpressionType getOutputType(Expr.InputBindingInspector inspector, List<E
     }
 
     @Override
-    protected ExprEval eval(@Nullable String x, int y)
+    protected ExprEval eval(String x, long y)
     {
-      if (y < 0) {
+      int yInt = (int) y;
+      if (y < 0 || yInt != y) {
         throw validationFailed("needs a positive integer as the second argument");
       }
-      if (x == null) {
-        return ExprEval.of(null);
-      }
       int len = x.length();
-      return ExprEval.of(y < len ? x.substring(len - y) : x);
+      return ExprEval.of(y < len ? x.substring(len - yInt) : x);
     }
   }
 
@@ -2853,15 +2855,13 @@ public ExpressionType getOutputType(Expr.InputBindingInspector inspector, List<E
     }
 
     @Override
-    protected ExprEval eval(@Nullable String x, int y)
+    protected ExprEval eval(String x, long y)
     {
-      if (y < 0) {
-        throw validationFailed("needs a postive integer as second argument");
-      }
-      if (x == null) {
-        return ExprEval.of(null);
+      int yInt = (int) y;
+      if (yInt < 0 || yInt != y) {
+        throw validationFailed("needs a positive integer as the second argument");
       }
-      return ExprEval.of(y < x.length() ? x.substring(0, y) : x);
+      return ExprEval.of(y < x.length() ? x.substring(0, yInt) : x);
     }
   }
 
@@ -3005,12 +3005,13 @@ public ExpressionType getOutputType(Expr.InputBindingInspector inspector, List<E
     }
 
     @Override
-    protected ExprEval eval(String x, int y)
+    protected ExprEval eval(String x, long y)
     {
-      if (x == null) {
-        return ExprEval.of(null);
+      int yInt = (int) y;
+      if (yInt != y) {
+        throw validationFailed("needs an integer as the second argument");
       }
-      return ExprEval.of(y < 1 ? NullHandling.defaultStringValue() : StringUtils.repeat(x, y));
+      return ExprEval.of(y < 1 ? NullHandling.defaultStringValue() : StringUtils.repeat(x, yInt));
     }
   }
 

diff --git a/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java b/processing/src/test/java/org/apache/druid/math/expr/FunctionTest.java
@@ -40,7 +40,6 @@
 import org.junit.Test;
 
 import javax.annotation.Nullable;
-
 import java.math.BigDecimal;
 import java.math.RoundingMode;
 import java.nio.ByteBuffer;
@@ -962,13 +961,65 @@ public void testBitwise()
     assertExpr("bitwiseConvertDoubleToLongBits(null)", null);
   }
 
+  @Test
+  public void testLeft()
+  {
+    assertExpr("left('hello', 0)", NullHandling.sqlCompatible() ? "" : null);
+    assertExpr("left('hello', 2)", "he");
+    assertExpr("left('hello', '2')", "he");
+    assertExpr("left('hello', 'hello')", null);
+    assertExpr("left('hello', 10)", "hello");
+    assertExpr("left('hello', null)", null);
+    assertExpr("left(31337, 2)", "31");
+    assertExpr("left(null, 10)", null);
+    assertExpr("left(nonexistent, 10)", null);
+
+    Throwable t1 = Assert.assertThrows(
+        DruidException.class,
+        () -> assertExpr("left('foo', -2)", null)
+    );
+    Assert.assertEquals("Function[left] needs a positive integer as the second argument", t1.getMessage());
+  }
+
+  @Test
+  public void testRight()
+  {
+    assertExpr("right('hello', 0)", NullHandling.sqlCompatible() ? "" : null);
+    assertExpr("right('hello', 2)", "lo");
+    assertExpr("right('hello', '2')", "lo");
+    assertExpr("right('hello', 'hello')", null);
+    assertExpr("right('hello', 10)", "hello");
+    assertExpr("right('hello', null)", null);
+    assertExpr("right(31337, 2)", "37");
+    assertExpr("right(null, 10)", null);
+    assertExpr("right(nonexistent, 10)", null);
+
+    Throwable t1 = Assert.assertThrows(
+        DruidException.class,
+        () -> assertExpr("right('foo', -2)", null)
+    );
+    Assert.assertEquals("Function[right] needs a positive integer as the second argument", t1.getMessage());
+  }
+
   @Test
   public void testRepeat()
   {
+    assertExpr("repeat('hello', 0)", null);
     assertExpr("repeat('hello', 2)", "hellohello");
+    assertExpr("repeat('hello', '2')", "hellohello");
+    assertExpr("repeat('hello', 'hello')", null);
     assertExpr("repeat('hello', -1)", null);
+    assertExpr("repeat('hello', null)", null);
     assertExpr("repeat(null, 10)", null);
     assertExpr("repeat(nonexistent, 10)", null);
+    assertExpr("repeat(4, 3)", "444");
+    assertExpr("repeat(4.1, 3)", "4.14.14.1");
+
+    Throwable t = Assert.assertThrows(
+        DruidException.class,
+        () -> assertExpr("repeat('foo', 9999999999)", null)
+    );
+    Assert.assertEquals("Function[repeat] needs an integer as the second argument", t.getMessage());
   }
 
   @Test

diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionsTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionsTest.java
@@ -2371,27 +2371,6 @@ public void testAbnormalRightWithNegativeNumber()
     Assert.assertEquals("Function[right] needs a positive integer as the second argument", t.getMessage());
   }
 
-  @Test
-  public void testAbnormalRightWithWrongType()
-  {
-    Throwable t = Assert.assertThrows(
-        DruidException.class,
-        () -> testHelper.testExpressionString(
-            new RightOperatorConversion().calciteOperator(),
-            ImmutableList.of(
-                testHelper.makeInputRef("s"),
-                testHelper.makeInputRef("s")
-            ),
-            makeExpression("right(\"s\",\"s\")"),
-            null
-        )
-    );
-    Assert.assertEquals(
-        "Function[right] needs a STRING as first argument and a LONG as second argument",
-        t.getMessage()
-    );
-  }
-
   @Test
   public void testLeft()
   {
@@ -2461,28 +2440,7 @@ public void testAbnormalLeftWithNegativeNumber()
             null
         )
     );
-    Assert.assertEquals("Function[left] needs a postive integer as second argument", t.getMessage());
-  }
-
-  @Test
-  public void testAbnormalLeftWithWrongType()
-  {
-    Throwable t = Assert.assertThrows(
-        DruidException.class,
-        () -> testHelper.testExpressionString(
-            new LeftOperatorConversion().calciteOperator(),
-            ImmutableList.of(
-                testHelper.makeInputRef("s"),
-                testHelper.makeInputRef("s")
-            ),
-            makeExpression("left(\"s\",\"s\")"),
-            null
-        )
-    );
-    Assert.assertEquals(
-        "Function[left] needs a STRING as first argument and a LONG as second argument",
-        t.getMessage()
-    );
+    Assert.assertEquals("Function[left] needs a positive integer as the second argument", t.getMessage());
   }
 
   @Test
@@ -2519,27 +2477,6 @@ public void testRepeat()
     );
   }
 
-  @Test
-  public void testAbnormalRepeatWithWrongType()
-  {
-    Throwable t = Assert.assertThrows(
-        DruidException.class,
-        () -> testHelper.testExpressionString(
-            new RepeatOperatorConversion().calciteOperator(),
-            ImmutableList.of(
-                testHelper.makeInputRef("s"),
-                testHelper.makeInputRef("s")
-            ),
-            makeExpression("repeat(\"s\",\"s\")"),
-            null
-        )
-    );
-    Assert.assertEquals(
-        "Function[repeat] needs a STRING as first argument and a LONG as second argument",
-        t.getMessage()
-    );
-  }
-
   @Test
   public void testOperatorConversionsDruidUnaryLongFn()
   {