Exposing optional replaceMissingValueWith in lookup function and macros

docs/querying/lookups.md

@@ -62,6 +62,12 @@ SELECT
 FROM sales
 GROUP BY 1
 ```
+The lookup function also accepts the 3rd argument called `replaceMissingValueWith` as a constant string. If your value is missing a lookup for the queried key, the lookup function returns the result value from `replaceMissingValueWith`
+For example:
+```
+LOOKUP(store, 'store_to_country', 'NA')
+```
+If value is missing from `store_to_country` lookup for given key 'store' then it will return `NA`.
 
 They can also be queried using the [JOIN operator](datasource.md#join):
 

processing/src/main/java/org/apache/druid/query/expression/LookupExprMacro.java

@@ -54,10 +54,11 @@ public String name()
   @Override
   public Expr apply(final List<Expr> args)
   {
-    validationHelperCheckArgumentCount(args, 2);
+    validationHelperCheckMinArgumentCount(args, 2);
 
     final Expr arg = args.get(0);
     final Expr lookupExpr = args.get(1);
+    final String replaceMissingValueWith = getReplaceMissingValueWith(args);
 
     validationHelperCheckArgIsLiteral(lookupExpr, "second argument");
     if (lookupExpr.getLiteralValue() == null) {
@@ -69,7 +70,7 @@ public Expr apply(final List<Expr> args)
         lookupExtractorFactoryContainerProvider,
         lookupName,
         false,
-        null,
+        replaceMissingValueWith,
         false,
         null
     );
@@ -104,6 +105,15 @@ public ExpressionType getOutputType(InputBindingInspector inspector)
       @Override
       public String stringify()
       {
+        if (replaceMissingValueWith != null) {
+          return StringUtils.format(
+              "%s(%s, %s, '%s')",
+              FN_NAME,
+              arg.stringify(),
+              lookupExpr.stringify(),
+              replaceMissingValueWith
+          );
+        }
         return StringUtils.format("%s(%s, %s)", FN_NAME, arg.stringify(), lookupExpr.stringify());
       }
 
@@ -116,4 +126,15 @@ public void decorateCacheKeyBuilder(CacheKeyBuilder builder)
 
     return new LookupExpr(arg);
   }
+
+  private String getReplaceMissingValueWith(final List<Expr> args)
+  {
+    if (args.size() > 2) {
+      final Expr missingValExpr = args.get(2);
+      if (missingValExpr.isLiteral() && missingValExpr.getLiteralValue() != null) {
+        return missingValExpr.getLiteralValue().toString();
+      }
+    }
+    return null;
+  }
 }

processing/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java

@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.expression;
+
+import com.google.common.collect.Lists;
+import org.apache.commons.compress.utils.Sets;
+import org.apache.druid.math.expr.Expr;
+import org.apache.druid.math.expr.ExprEval;
+import org.apache.druid.query.lookup.LookupExtractorFactoryContainer;
+import org.apache.druid.query.lookup.LookupExtractorFactoryContainerProvider;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+public class LookupExprMacroTest extends MacroTestBase
+{
+  public LookupExprMacroTest()
+  {
+    super(
+        new LookupExprMacro(new LookupExtractorFactoryContainerProvider()
+        {
+          @Override
+          public Set<String> getAllLookupNames()
+          {
+            return Sets.newHashSet("test_lookup");
+          }
+
+          @Override
+          public Optional<LookupExtractorFactoryContainer> get(String lookupName)
+          {
+            return Optional.empty();
+          }
+        })
+    );
+  }
+
+  @Test
+  public void testTooFewArgs()
+  {
+    expectException(IllegalArgumentException.class, "Function[lookup] requires at least 2 arguments");
+    apply(Collections.emptyList());
+  }
+
+  @Test
+  public void testNonLiteralLookupName()
+  {
+    expectException(
+        IllegalArgumentException.class,
+        "Function[lookup] second argument must be a registered lookup name"
+    );
+    apply(getArgs(Lists.newArrayList("1", new ArrayList<String>())));
+  }
+
+  @Test
+  public void testValidCalls()
+  {
+    Assert.assertNotNull(apply(getArgs(Lists.newArrayList("1", "test_lookup"))));
+    Assert.assertNotNull(apply(getArgs(Lists.newArrayList("null", "test_lookup"))));
+    Assert.assertNotNull(apply(getArgs(Lists.newArrayList("1", "test_lookup", null))));
+    Assert.assertNotNull(apply(getArgs(Lists.newArrayList("1", "test_lookup", "N/A"))));
+  }
+
+  private List<Expr> getArgs(List<Object> args)
+  {
+    return args.stream().map(a -> {
+      if (a != null && a instanceof String) {
+        return ExprEval.of(a.toString()).toExpr();
+      }
+      return ExprEval.bestEffortOf(null).toExpr();
+    }).collect(Collectors.toList());
+  }
+}

server/src/test/java/org/apache/druid/query/expression/LookupExprMacroTest.java

@@ -47,7 +47,12 @@ public void testLookup()
   {
     assertExpr("lookup(x, 'lookyloo')", "xfoo");
   }
-
+  @Test
+  public void testLookupMissingValue()
+  {
+    assertExpr("lookup(y, 'lookyloo', 'N/A')", "N/A");
+    assertExpr("lookup(y, 'lookyloo', null)", null);
+  }
   @Test
   public void testLookupNotFound()
   {

sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/QueryLookupOperatorConversion.java

@@ -30,16 +30,27 @@
 import org.apache.druid.query.lookup.LookupExtractorFactoryContainerProvider;
 import org.apache.druid.query.lookup.RegisteredLookupExtractionFn;
 import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.sql.calcite.expression.BasicOperandTypeChecker;
 import org.apache.druid.sql.calcite.expression.DruidExpression;
 import org.apache.druid.sql.calcite.expression.OperatorConversions;
 import org.apache.druid.sql.calcite.expression.SqlOperatorConversion;
 import org.apache.druid.sql.calcite.planner.PlannerContext;
 
+import java.util.List;
+
 public class QueryLookupOperatorConversion implements SqlOperatorConversion
 {
   private static final SqlFunction SQL_FUNCTION = OperatorConversions
       .operatorBuilder("LOOKUP")
-      .operandTypes(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER)
+      .operandTypeChecker(
+          BasicOperandTypeChecker.builder()
+                                 .operandTypes(
+                                     SqlTypeFamily.CHARACTER,
+                                     SqlTypeFamily.CHARACTER,
+                                     SqlTypeFamily.CHARACTER
+                                 )
+                                 .requiredOperandCount(2)
+                                 .build())
       .returnTypeNullable(SqlTypeName.VARCHAR)
       .functionCategory(SqlFunctionCategory.STRING)
       .build();
@@ -73,14 +84,15 @@ public DruidExpression toDruidExpression(
         inputExpressions -> {
           final DruidExpression arg = inputExpressions.get(0);
           final Expr lookupNameExpr = plannerContext.parseExpression(inputExpressions.get(1).getExpression());
+          final String replaceMissingValueWith = getReplaceMissingValueWith(inputExpressions, plannerContext);
 
           if (arg.isSimpleExtraction() && lookupNameExpr.isLiteral()) {
             return arg.getSimpleExtraction().cascade(
                 new RegisteredLookupExtractionFn(
                     lookupExtractorFactoryContainerProvider,
                     (String) lookupNameExpr.getLiteralValue(),
                     false,
-                    null,
+                    replaceMissingValueWith,
                     null,
                     true
                 )
@@ -91,4 +103,18 @@ public DruidExpression toDruidExpression(
         }
     );
   }
+
+  private String getReplaceMissingValueWith(
+      final List<DruidExpression> inputExpressions,
+      final PlannerContext plannerContext
+  )
+  {
+    if (inputExpressions.size() > 2) {
+      final Expr missingValExpr = plannerContext.parseExpression(inputExpressions.get(2).getExpression());
+      if (missingValExpr.isLiteral()) {
+        return missingValExpr.getLiteralValue().toString();
+      }
+    }
+    return null;
+  }
 }

sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java

@@ -8786,6 +8786,52 @@ public void testFilterAndGroupByLookup()
         )
     );
   }
+  @Test
+  public void testLookupReplaceMissingValueWith()
+  {
+    // Cannot vectorize due to extraction dimension specs.
+    cannotVectorize();
+
+    final RegisteredLookupExtractionFn extractionFn = new RegisteredLookupExtractionFn(
+        null,
+        "lookyloo",
+        false,
+        "Missing_Value",
+        null,
+        true
+    );
+
+    testQuery(
+        "SELECT LOOKUP(dim1, 'lookyloo', 'Missing_Value'), COUNT(*) FROM foo group by 1",
+        ImmutableList.of(
+            GroupByQuery.builder()
+                        .setDataSource(CalciteTests.DATASOURCE1)
+                        .setInterval(querySegmentSpec(Filtration.eternity()))
+                        .setGranularity(Granularities.ALL)
+                        .setDimensions(
+                            dimensions(
+                                new ExtractionDimensionSpec(
+                                    "dim1",
+                                    "d0",
+                                    ColumnType.STRING,
+                                    extractionFn
+                                )
+                            )
+                        )
+                        .setAggregatorSpecs(
+                            aggregators(
+                                new CountAggregatorFactory("a0")
+                            )
+                        )
+                        .setContext(QUERY_CONTEXT_DEFAULT)
+                        .build()
+        ),
+        ImmutableList.of(
+            new Object[]{"Missing_Value", 5L},
+            new Object[]{"xabc", 1L}
+        )
+    );
+  }
 
   @Test
   public void testCountDistinctOfLookup()