Addressing comments

apache · Oct 1, 2023 · bd0cac9 · bd0cac9
1 parent 757d8ac
commit bd0cac9
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 15 deletions.
diff --git a/docs/querying/sql-scalar.md b/docs/querying/sql-scalar.md
@@ -98,7 +98,7 @@ String functions accept strings, and return a type appropriate to the function.
 |`CHAR_LENGTH(expr)`|Alias for `LENGTH`.|
 |`CHARACTER_LENGTH(expr)`|Alias for `LENGTH`.|
 |`STRLEN(expr)`|Alias for `LENGTH`.|
-|`LOOKUP(expr, lookupName, [replaceMissingValueWith])`|Look up `expr` in a registered [query-time lookup table](lookups.md). Note that lookups can also be queried directly using the [`lookup` schema](sql.md#from).|
+|`LOOKUP(expr, lookupName, [replaceMissingValueWith])`|Look up `expr` in a registered [query-time lookup table](lookups.md). Note that lookups can also be queried directly using the [`lookup` schema](sql.md#from). Optional constant replaceMissingValueWith can be passed as 3rd argument to be returned when value is missing from lookup.|
 |`LOWER(expr)`|Returns `expr` in all lowercase.|
 |`UPPER(expr)`|Returns `expr` in all uppercase.|
 |`PARSE_LONG(string, [radix])`|Parses a string into a long (BIGINT) with the given radix, or 10 (decimal) if a radix is not provided.|

diff --git a/processing/src/main/java/org/apache/druid/query/expression/LookupExprMacro.java b/processing/src/main/java/org/apache/druid/query/expression/LookupExprMacro.java
@@ -58,7 +58,7 @@ public Expr apply(final List<Expr> args)
 
     final Expr arg = args.get(0);
     final Expr lookupExpr = args.get(1);
-    final String replaceMissingValueWith = getReplaceMissingValueWith(args);
+    final Expr replaceMissingValueWith = getReplaceMissingValueWith(args);
 
     validationHelperCheckArgIsLiteral(lookupExpr, "second argument");
     if (lookupExpr.getLiteralValue() == null) {
@@ -70,7 +70,9 @@ public Expr apply(final List<Expr> args)
         lookupExtractorFactoryContainerProvider,
         lookupName,
         false,
-        replaceMissingValueWith,
+        replaceMissingValueWith != null && replaceMissingValueWith.isLiteral()
+        ? (String) replaceMissingValueWith.getLiteralValue()
+        : null,
         false,
         null
     );
@@ -107,11 +109,11 @@ public String stringify()
       {
         if (replaceMissingValueWith != null) {
           return StringUtils.format(
-              "%s(%s, %s, '%s')",
+              "%s(%s, %s, %s)",
               FN_NAME,
               arg.stringify(),
               lookupExpr.stringify(),
-              replaceMissingValueWith
+              replaceMissingValueWith.stringify()
           );
         }
         return StringUtils.format("%s(%s, %s)", FN_NAME, arg.stringify(), lookupExpr.stringify());
@@ -127,12 +129,12 @@ public void decorateCacheKeyBuilder(CacheKeyBuilder builder)
     return new LookupExpr(arg);
   }
 
-  private String getReplaceMissingValueWith(final List<Expr> args)
+  private Expr getReplaceMissingValueWith(final List<Expr> args)
   {
     if (args.size() > 2) {
       final Expr missingValExpr = args.get(2);
       validationHelperCheckArgIsLiteral(missingValExpr, "third argument");
-      return String.valueOf(missingValExpr.getLiteralValue());
+      return missingValExpr;
     }
     return null;
   }

diff --git a/server/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java b/server/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java
@@ -51,7 +51,7 @@ public void testLookup()
   public void testLookupMissingValue()
   {
     assertExpr("lookup(y, 'lookyloo', 'N/A')", "N/A");
-    assertExpr("lookup(y, 'lookyloo', null)", "null");
+    assertExpr("lookup(y, 'lookyloo', null)", null);
   }
   @Test
   public void testLookupNotFound()

diff --git a/...n/java/org/apache/druid/sql/calcite/expression/builtin/QueryLookupOperatorConversion.java b/...n/java/org/apache/druid/sql/calcite/expression/builtin/QueryLookupOperatorConversion.java
@@ -50,6 +50,7 @@ public class QueryLookupOperatorConversion implements SqlOperatorConversion
                                      SqlTypeFamily.CHARACTER
                                  )
                                  .requiredOperandCount(2)
+                                 .literalOperands(2)
                                  .build())
       .returnTypeNullable(SqlTypeName.VARCHAR)
       .functionCategory(SqlFunctionCategory.STRING)
@@ -112,7 +113,7 @@ private String getReplaceMissingValueWith(
     if (inputExpressions.size() > 2) {
       final Expr missingValExpr = plannerContext.parseExpression(inputExpressions.get(2).getExpression());
       if (missingValExpr.isLiteral()) {
-        return missingValExpr.getLiteralValue().toString();
+        return (String) missingValExpr.getLiteralValue();
       }
     }
     return null;

diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
@@ -8973,17 +8973,24 @@ public void testLookupReplaceMissingValueWith()
     // Cannot vectorize due to extraction dimension specs.
     cannotVectorize();
 
-    final RegisteredLookupExtractionFn extractionFn = new RegisteredLookupExtractionFn(
+    final RegisteredLookupExtractionFn extractionFn1 = new RegisteredLookupExtractionFn(
         null,
         "lookyloo",
         false,
         "Missing_Value",
         null,
         true
     );
-
+    final RegisteredLookupExtractionFn extractionFnRMVNull = new RegisteredLookupExtractionFn(
+        null,
+        "lookyloo",
+        false,
+        null,
+        null,
+        true
+    );
     testQuery(
-        "SELECT LOOKUP(dim1, 'lookyloo', 'Missing_Value'), COUNT(*) FROM foo group by 1",
+        "SELECT LOOKUP(dim1, 'lookyloo', 'Missing_Value'), LOOKUP(dim1, 'lookyloo', null) as rmvNull, COUNT(*) FROM foo group by 1,2",
         ImmutableList.of(
             GroupByQuery.builder()
                         .setDataSource(CalciteTests.DATASOURCE1)
@@ -8995,7 +9002,13 @@ public void testLookupReplaceMissingValueWith()
                                     "dim1",
                                     "d0",
                                     ColumnType.STRING,
-                                    extractionFn
+                                    extractionFn1
+                                ),
+                                new ExtractionDimensionSpec(
+                                    "dim1",
+                                    "d1",
+                                    ColumnType.STRING,
+                                    extractionFnRMVNull
                                 )
                             )
                         )
@@ -9008,8 +9021,8 @@ public void testLookupReplaceMissingValueWith()
                         .build()
         ),
         ImmutableList.of(
-            new Object[]{"Missing_Value", 5L},
-            new Object[]{"xabc", 1L}
+            new Object[]{"Missing_Value", NullHandling.sqlCompatible() ? null : "", 5L},
+            new Object[]{"xabc", "xabc", 1L}
         )
     );
   }