Add IS [NOT] DISTINCT FROM to SQL and join matchers.

docs/querying/sql-operators.md

@@ -79,7 +79,9 @@ Also see the [CONCAT function](sql-scalar.md#string-functions).
 |Operator|Description|
 |--------|-----------|
 |`x = y` |Equal to|
+|`x IS NOT DISTINCT FROM y`|Equal to, considering `NULL` as a value. Never returns `NULL`.|
 |`x <> y`|Not equal to|
+|`x IS DISTINCT FROM y`|Not equal to, considering `NULL` as a value. Never returns `NULL`.|
 |`x > y` |Greater than|
 |`x >= y`|Greater than or equal to|
 |`x < y` |Less than|

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlan.java

@@ -218,7 +218,7 @@ private static JoinAlgorithm deduceJoinAlgorithm(JoinAlgorithm preferredJoinAlgo
     JoinAlgorithm deducedJoinAlgorithm;
     if (JoinAlgorithm.BROADCAST.equals(preferredJoinAlgorithm)) {
       deducedJoinAlgorithm = JoinAlgorithm.BROADCAST;
-    } else if (isConditionEqualityOnLeftAndRightColumns(joinDataSource.getConditionAnalysis())) {
+    } else if (canUseSortMergeJoin(joinDataSource.getConditionAnalysis())) {
       deducedJoinAlgorithm = JoinAlgorithm.SORT_MERGE;
     } else {
       deducedJoinAlgorithm = JoinAlgorithm.BROADCAST;
@@ -237,15 +237,21 @@ private static JoinAlgorithm deduceJoinAlgorithm(JoinAlgorithm preferredJoinAlgo
   }
 
   /**
-   * Checks if the join condition on two tables "table1" and "table2" is of the form
+   * Checks if the sortMerge algorithm can execute a particular join condition.
+   *
+   * Two checks:
+   * (1) join condition on two tables "table1" and "table2" is of the form
    * table1.columnA = table2.columnA && table1.columnB = table2.columnB && ....
-   * sortMerge algorithm can help these types of join conditions
+   *
+   * (2) join condition uses equals, not IS NOT DISTINCT FROM [sortMerge processor does not currently implement
+   * IS NOT DISTINCT FROM]
    */
-  private static boolean isConditionEqualityOnLeftAndRightColumns(JoinConditionAnalysis joinConditionAnalysis)
+  private static boolean canUseSortMergeJoin(JoinConditionAnalysis joinConditionAnalysis)
   {
-    return joinConditionAnalysis.getEquiConditions()
-                                .stream()
-                                .allMatch(equality -> equality.getLeftExpr().isIdentifier());
+    return joinConditionAnalysis
+        .getEquiConditions()
+        .stream()
+        .allMatch(equality -> equality.getLeftExpr().isIdentifier() && !equality.isIncludeNull());
   }
 
   /**

processing/src/main/java/org/apache/druid/math/expr/Exprs.java

@@ -19,11 +19,13 @@
 
 package org.apache.druid.math.expr;
 
-import org.apache.druid.java.util.common.Pair;
 import org.apache.druid.java.util.common.UOE;
+import org.apache.druid.segment.join.Equality;
+import org.apache.druid.segment.join.JoinPrefixUtils;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Stack;
 
@@ -79,16 +81,56 @@ public static List<Expr> decomposeAnd(final Expr expr)
   }
 
   /**
-   * Decomposes an equality expr into the left- and right-hand side.
+   * Decomposes an equality expr into an {@link Equality}. Used by join-related code to identify equi-joins.
    *
    * @return decomposed equality, or empty if the input expr was not an equality expr
    */
-  public static Optional<Pair<Expr, Expr>> decomposeEquals(final Expr expr)
+  public static Optional<Equality> decomposeEquals(final Expr expr, final String rightPrefix)
   {
+    final Expr lhs;
+    final Expr rhs;
+    final boolean includeNull;
+
     if (expr instanceof BinEqExpr) {
-      return Optional.of(Pair.of(((BinEqExpr) expr).left, ((BinEqExpr) expr).right));
+      lhs = ((BinEqExpr) expr).left;
+      rhs = ((BinEqExpr) expr).right;
+      includeNull = false;
+    } else if (expr instanceof FunctionExpr
+               && ((FunctionExpr) expr).function instanceof Function.IsNotDistinctFromFunc) {
+      final List<Expr> args = ((FunctionExpr) expr).args;
+      lhs = args.get(0);
+      rhs = args.get(1);
+      includeNull = true;
+    } else {
+      return Optional.empty();
+    }
+
+    if (isLeftExprAndRightColumn(lhs, rhs, rightPrefix)) {
+      // rhs is a right-hand column; lhs is an expression solely of the left-hand side.
+      return Optional.of(
+          new Equality(
+              lhs,
+              Objects.requireNonNull(rhs.getBindingIfIdentifier()).substring(rightPrefix.length()),
+              includeNull
+          )
+      );
+    } else if (isLeftExprAndRightColumn(rhs, lhs, rightPrefix)) {
+      return Optional.of(
+          new Equality(
+              rhs,
+              Objects.requireNonNull(lhs.getBindingIfIdentifier()).substring(rightPrefix.length()),
+              includeNull
+          )
+      );
     } else {
       return Optional.empty();
     }
   }
+
+  private static boolean isLeftExprAndRightColumn(final Expr a, final Expr b, final String rightPrefix)
+  {
+    return a.analyzeInputs().getRequiredBindings().stream().noneMatch(c -> JoinPrefixUtils.isPrefixedBy(c, rightPrefix))
+           && b.getBindingIfIdentifier() != null
+           && JoinPrefixUtils.isPrefixedBy(b.getBindingIfIdentifier(), rightPrefix);
+  }
 }

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

@@ -2224,6 +2224,89 @@ public <T> ExprVectorProcessor<T> asVectorProcessor(Expr.VectorInputBindingInspe
     }
   }
 
+  /**
+   * SQL function "x IS NOT DISTINCT FROM y". Very similar to "x = y", i.e. {@link BinEqExpr}, except this function
+   * never returns null, and this function considers NULL itself to be not-distinct-from NULL.
+   */
+  class IsNotDistinctFromFunc implements Function
+  {
+    @Override
+    public String name()
+    {
+      return "notdistinctfrom";
+    }
+
+    @Override
+    public ExprEval apply(List<Expr> args, Expr.ObjectBinding bindings)
+    {
+      final ExprEval leftVal = args.get(0).eval(bindings);
+      final ExprEval rightVal = args.get(1).eval(bindings);
+
+      if (leftVal.value() == null || rightVal.value() == null) {
+        return ExprEval.ofLongBoolean(leftVal.value() == null && rightVal.value() == null);
+      }
+
+      final ExpressionType comparisonType = ExpressionTypeConversion.autoDetect(leftVal, rightVal);
+      switch (comparisonType.getType()) {
+        case STRING:
+          return ExprEval.ofLongBoolean(Objects.equals(leftVal.asString(), rightVal.asString()));
+        case LONG:
+          return ExprEval.ofLongBoolean(leftVal.asLong() == rightVal.asLong());
+        case DOUBLE:
+        default:
+          if (leftVal.isNumericNull() || rightVal.isNumericNull()) {
+            return ExprEval.ofLongBoolean(leftVal.isNumericNull() && rightVal.isNumericNull());
+          } else {
+            return ExprEval.ofLongBoolean(leftVal.asDouble() == rightVal.asDouble());
+          }
+      }
+    }
+
+    @Override
+    public void validateArguments(List<Expr> args)
+    {
+      validationHelperCheckArgumentCount(args, 2);
+    }
+
+    @Nullable
+    @Override
+    public ExpressionType getOutputType(Expr.InputBindingInspector inspector, List<Expr> args)
+    {
+      return ExpressionType.LONG;
+    }
+  }
+
+  /**
+   * SQL function "x IS DISTINCT FROM y". Very similar to "x <> y", i.e. {@link BinNeqExpr}, except this function
+   * never returns null.
+   */
+  class IsDistinctFromFunc extends IsNotDistinctFromFunc
+  {
+    @Override
+    public String name()
+    {
+      return "isdistinctfrom";
+    }
+
+    @Override
+    public ExprEval apply(List<Expr> args, Expr.ObjectBinding bindings)
+    {
+      return ExprEval.ofLongBoolean(!super.apply(args, bindings).asBoolean());
+    }
+
+    @Override
+    public void validateArguments(List<Expr> args)
+    {
+      validationHelperCheckArgumentCount(args, 2);
+    }
+
+    @Nullable
+    @Override
+    public ExpressionType getOutputType(Expr.InputBindingInspector inspector, List<Expr> args)
+    {
+      return ExpressionType.LONG;
+    }
+  }
   class IsNullFunc implements Function
   {
     @Override

processing/src/main/java/org/apache/druid/segment/join/Equality.java

@@ -32,11 +32,13 @@ public class Equality
 {
   private final Expr leftExpr;
   private final String rightColumn;
+  private final boolean includeNull;
 
-  public Equality(final Expr leftExpr, final String rightColumn)
+  public Equality(final Expr leftExpr, final String rightColumn, final boolean includeNull)
   {
     this.leftExpr = leftExpr;
     this.rightColumn = rightColumn;
+    this.includeNull = includeNull;
   }
 
   public Expr getLeftExpr()
@@ -49,12 +51,22 @@ public String getRightColumn()
     return rightColumn;
   }
 
+  /**
+   * Whether null is treated as a value that can be equal to itself. True for conditions using "IS NOT DISTINCT FROM",
+   * false for conditions using regular equals.
+   */
+  public boolean isIncludeNull()
+  {
+    return includeNull;
+  }
+
   @Override
   public String toString()
   {
     return "Equality{" +
            "leftExpr=" + leftExpr +
            ", rightColumn='" + rightColumn + '\'' +
+           ", includeNull=" + includeNull +
            '}';
   }
 }

processing/src/main/java/org/apache/druid/segment/join/JoinConditionAnalysis.java

@@ -20,7 +20,6 @@
 package org.apache.druid.segment.join;
 
 import com.google.common.base.Preconditions;
-import org.apache.druid.java.util.common.Pair;
 import org.apache.druid.math.expr.Expr;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.math.expr.Exprs;
@@ -121,40 +120,18 @@ public static JoinConditionAnalysis forExpression(
 
     final List<Expr> exprs = Exprs.decomposeAnd(conditionExpr);
     for (Expr childExpr : exprs) {
-      final Optional<Pair<Expr, Expr>> maybeDecomposed = Exprs.decomposeEquals(childExpr);
+      final Optional<Equality> maybeEquality = Exprs.decomposeEquals(childExpr, rightPrefix);
 
-      if (!maybeDecomposed.isPresent()) {
+      if (!maybeEquality.isPresent()) {
         nonEquiConditions.add(childExpr);
       } else {
-        final Pair<Expr, Expr> decomposed = maybeDecomposed.get();
-        final Expr lhs = Objects.requireNonNull(decomposed.lhs);
-        final Expr rhs = Objects.requireNonNull(decomposed.rhs);
-
-        if (isLeftExprAndRightColumn(lhs, rhs, rightPrefix)) {
-          // rhs is a right-hand column; lhs is an expression solely of the left-hand side.
-          equiConditions.add(
-              new Equality(lhs, Objects.requireNonNull(rhs.getBindingIfIdentifier()).substring(rightPrefix.length()))
-          );
-        } else if (isLeftExprAndRightColumn(rhs, lhs, rightPrefix)) {
-          equiConditions.add(
-              new Equality(rhs, Objects.requireNonNull(lhs.getBindingIfIdentifier()).substring(rightPrefix.length()))
-          );
-        } else {
-          nonEquiConditions.add(childExpr);
-        }
+        equiConditions.add(maybeEquality.get());
       }
     }
 
     return new JoinConditionAnalysis(condition, rightPrefix, equiConditions, nonEquiConditions);
   }
 
-  private static boolean isLeftExprAndRightColumn(final Expr a, final Expr b, final String rightPrefix)
-  {
-    return a.analyzeInputs().getRequiredBindings().stream().noneMatch(c -> JoinPrefixUtils.isPrefixedBy(c, rightPrefix))
-           && b.getBindingIfIdentifier() != null
-           && JoinPrefixUtils.isPrefixedBy(b.getBindingIfIdentifier(), rightPrefix);
-  }
-
   /**
    * Return the condition expression.
    */

processing/src/main/java/org/apache/druid/segment/join/table/IndexedTable.java

@@ -95,6 +95,7 @@ default ColumnSelectorFactory makeColumnSelectorFactory(ReadableOffset offset, b
    * see {@link org.apache.druid.segment.join.JoinableFactory#computeJoinCacheKey}
    *
    * @return the byte array for cache key
+   *
    * @throws {@link IAE} if caching is not supported
    */
   default byte[] computeCacheKey()
@@ -125,23 +126,25 @@ interface Index
     /**
      * Returns whether keys are unique in this index. If this returns true, then {@link #find(Object)} will only ever
      * return a zero- or one-element list.
+     *
+     * @param includeNull whether null is considered a valid key
      */
-    boolean areKeysUnique();
+    boolean areKeysUnique(boolean includeNull);
 
     /**
      * Returns the list of row numbers corresponding to "key" in this index.
      *
      * If "key" is some type other than the natural type {@link #keyType()}, it will be converted before checking
      * the index.
      */
-    IntSortedSet find(Object key);
+    IntSortedSet find(@Nullable Object key);
 
     /**
      * Returns the row number corresponding to "key" in this index, or {@link #NOT_FOUND} if the key does not exist
      * in the index.
      *
-     * It is only valid to call this method if {@link #keyType()} is {@link ValueType#LONG} and {@link #areKeysUnique()}
-     * returns true.
+     * It is only valid to call this method if {@link #keyType()} is {@link ValueType#LONG} and
+     * {@link #areKeysUnique(boolean)} returns true.
      *
      * @throws UnsupportedOperationException if preconditions are not met
      */