Add native filter conversion for SCALAR_IN_ARRAY. (#16312)

* Add native filter conversion for SCALAR_IN_ARRAY. Main changes: 1) Add an implementation of "toDruidFilter" in ScalarInArrayOperatorConversion. 2) Split up Expressions.literalToDruidExpression into two functions, so the first half (literalToExprEval) can be used by ScalarInArrayOperatorConversion to more efficiently create the list of match values. * Fix type in time arithmetic conversion. * Test updates. * Update test cases to use null instead of '' in default-value mode. * Switch test from msqIncompatible to compatible with a different result. * Update one more test. * Fix test. * Update tests. * Use ExprEvalWrapper to differentiate between empty string and null. * Fix tests some more. * Fix test. * Additional comment. * Style adjustment. * Fix tests. * trueValue -> actualValue. * Use different approach, DruidLiteral instead of ExprEvalWrapper. * Revert changes in ArrayOfDoublesSketchSqlAggregatorTest.
apache · May 3, 2024 · 588d442 · 588d442
1 parent 1b107ff
commit 588d442
Show file tree

Hide file tree

Showing 10 changed files with 723 additions and 149 deletions.
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/expression/DruidExpression.java b/sql/src/main/java/org/apache/druid/sql/calcite/expression/DruidExpression.java
@@ -26,6 +26,9 @@
 import com.google.common.primitives.Chars;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.druid.math.expr.Expr;
+import org.apache.druid.math.expr.ExprEval;
+import org.apache.druid.math.expr.ExprType;
+import org.apache.druid.math.expr.ExpressionType;
 import org.apache.druid.segment.VirtualColumn;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.virtual.ExpressionVirtualColumn;
@@ -180,6 +183,19 @@ public static DruidExpression ofLiteral(
     );
   }
 
+  /**
+   * Create a literal expression from an {@link ExprEval}.
+   */
+  public static DruidExpression ofLiteral(final DruidLiteral literal)
+  {
+    if (literal.type() != null && literal.type().is(ExprType.STRING)) {
+      return ofStringLiteral((String) literal.value());
+    } else {
+      final ColumnType evalColumnType = literal.type() != null ? ExpressionType.toColumnType(literal.type()) : null;
+      return ofLiteral(evalColumnType, ExprEval.ofType(literal.type(), literal.value()).toExpr().stringify());
+    }
+  }
+
   public static DruidExpression ofStringLiteral(final String s)
   {
     return ofLiteral(ColumnType.STRING, stringLiteral(s));

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/expression/DruidLiteral.java b/sql/src/main/java/org/apache/druid/sql/calcite/expression/DruidLiteral.java
@@ -0,0 +1,70 @@
+/*
+ * 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.sql.calcite.expression;
+
+import org.apache.druid.common.config.NullHandling;
+import org.apache.druid.math.expr.ExprEval;
+import org.apache.druid.math.expr.ExpressionType;
+
+import javax.annotation.Nullable;
+
+/**
+ * Literal value, plus a {@link ExpressionType} that represents how to interpret the literal value.
+ *
+ * These are similar to {@link ExprEval}, but not identical: unlike {@link ExprEval}, string values in this class
+ * are not normalized through {@link NullHandling#emptyToNullIfNeeded(String)}. This allows us to differentiate
+ * between null and empty-string literals even when {@link NullHandling#replaceWithDefault()}.
+ */
+public class DruidLiteral
+{
+  @Nullable
+  private final ExpressionType type;
+
+  @Nullable
+  private final Object value;
+
+  DruidLiteral(final ExpressionType type, @Nullable final Object value)
+  {
+    this.type = type;
+    this.value = value;
+  }
+
+  @Nullable
+  public ExpressionType type()
+  {
+    return type;
+  }
+
+  @Nullable
+  public Object value()
+  {
+    return value;
+  }
+
+  public DruidLiteral castTo(final ExpressionType toType)
+  {
+    if (type.equals(toType)) {
+      return this;
+    }
+
+    final ExprEval<?> castEval = ExprEval.ofType(type, value).castTo(toType);
+    return new DruidLiteral(castEval.type(), castEval.value());
+  }
+}
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/expression/Expressions.java b/sql/src/main/java/org/apache/druid/sql/calcite/expression/Expressions.java
@@ -38,6 +38,7 @@
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.granularity.Granularity;
 import org.apache.druid.math.expr.Expr;
+import org.apache.druid.math.expr.ExpressionType;
 import org.apache.druid.query.aggregation.PostAggregator;
 import org.apache.druid.query.expression.TimestampFloorExprMacro;
 import org.apache.druid.query.extraction.ExtractionFn;
@@ -240,7 +241,8 @@ public static DruidExpression toDruidExpressionWithPostAggOperands(
     } else if (rexNode instanceof RexCall) {
       return rexCallToDruidExpression(plannerContext, rowSignature, rexNode, postAggregatorVisitor);
     } else if (kind == SqlKind.LITERAL) {
-      return literalToDruidExpression(plannerContext, rexNode);
+      final DruidLiteral eval = calciteLiteralToDruidLiteral(plannerContext, rexNode);
+      return eval != null ? DruidExpression.ofLiteral(eval) : null;
     } else {
       // Can't translate.
       return null;
@@ -306,61 +308,85 @@ private static DruidExpression rexCallToDruidExpression(
     }
   }
 
+  /**
+   * Create a {@link DruidLiteral} from a literal {@link RexNode}. Necessary because Calcite represents literals using
+   * different Java classes than Druid does.
+   *
+   * @param plannerContext planner context
+   * @param rexNode        Calcite literal
+   *
+   * @return converted literal, or null if the literal cannot be converted
+   */
   @Nullable
-  static DruidExpression literalToDruidExpression(
+  public static DruidLiteral calciteLiteralToDruidLiteral(
       final PlannerContext plannerContext,
       final RexNode rexNode
   )
   {
-    final SqlTypeName sqlTypeName = rexNode.getType().getSqlTypeName();
+    if (rexNode.isA(SqlKind.CAST)) {
+      if (SqlTypeFamily.DATE.contains(rexNode.getType())) {
+        // Cast to DATE suggests some timestamp flooring. We don't deal with that here, so return null.
+        return null;
+      }
+
+      final DruidLiteral innerLiteral =
+          calciteLiteralToDruidLiteral(plannerContext, ((RexCall) rexNode).getOperands().get(0));
+      if (innerLiteral == null) {
+        return null;
+      }
+
+      final ColumnType castToColumnType = Calcites.getColumnTypeForRelDataType(rexNode.getType());
+      if (castToColumnType == null) {
+        return null;
+      }
+
+      final ExpressionType castToExprType = ExpressionType.fromColumnType(castToColumnType);
+      if (castToExprType == null) {
+        return null;
+      }
+
+      return innerLiteral.castTo(castToExprType);
+    }
 
     // Translate literal.
-    final ColumnType columnType = Calcites.getColumnTypeForRelDataType(rexNode.getType());
+    final SqlTypeName sqlTypeName = rexNode.getType().getSqlTypeName();
+    final DruidLiteral retVal;
+
     if (RexLiteral.isNullLiteral(rexNode)) {
-      return DruidExpression.ofLiteral(columnType, DruidExpression.nullLiteral());
+      final ColumnType columnType = Calcites.getColumnTypeForRelDataType(rexNode.getType());
+      final ExpressionType expressionType = columnType == null ? null : ExpressionType.fromColumnTypeStrict(columnType);
+      retVal = new DruidLiteral(expressionType, null);
     } else if (SqlTypeName.INT_TYPES.contains(sqlTypeName)) {
       final Number number = (Number) RexLiteral.value(rexNode);
-      return DruidExpression.ofLiteral(
-          columnType,
-          number == null ? DruidExpression.nullLiteral() : DruidExpression.longLiteral(number.longValue())
-      );
+      retVal = new DruidLiteral(ExpressionType.LONG, number == null ? null : number.longValue());
     } else if (SqlTypeName.NUMERIC_TYPES.contains(sqlTypeName)) {
       // Numeric, non-INT, means we represent it as a double.
       final Number number = (Number) RexLiteral.value(rexNode);
-      return DruidExpression.ofLiteral(
-          columnType,
-          number == null ? DruidExpression.nullLiteral() : DruidExpression.doubleLiteral(number.doubleValue())
-      );
+      retVal = new DruidLiteral(ExpressionType.DOUBLE, number == null ? null : number.doubleValue());
     } else if (SqlTypeFamily.INTERVAL_DAY_TIME == sqlTypeName.getFamily()) {
       // Calcite represents DAY-TIME intervals in milliseconds.
       final long milliseconds = ((Number) RexLiteral.value(rexNode)).longValue();
-      return DruidExpression.ofLiteral(columnType, DruidExpression.longLiteral(milliseconds));
+      retVal = new DruidLiteral(ExpressionType.LONG, milliseconds);
     } else if (SqlTypeFamily.INTERVAL_YEAR_MONTH == sqlTypeName.getFamily()) {
       // Calcite represents YEAR-MONTH intervals in months.
       final long months = ((Number) RexLiteral.value(rexNode)).longValue();
-      return DruidExpression.ofLiteral(columnType, DruidExpression.longLiteral(months));
+      retVal = new DruidLiteral(ExpressionType.LONG, months);
     } else if (SqlTypeName.STRING_TYPES.contains(sqlTypeName)) {
-      return DruidExpression.ofStringLiteral(RexLiteral.stringValue(rexNode));
+      final String s = RexLiteral.stringValue(rexNode);
+      retVal = new DruidLiteral(ExpressionType.STRING, s);
     } else if (SqlTypeName.TIMESTAMP == sqlTypeName || SqlTypeName.DATE == sqlTypeName) {
-      if (RexLiteral.isNullLiteral(rexNode)) {
-        return DruidExpression.ofLiteral(columnType, DruidExpression.nullLiteral());
-      } else {
-        return DruidExpression.ofLiteral(
-            columnType,
-            DruidExpression.longLiteral(
-                Calcites.calciteDateTimeLiteralToJoda(rexNode, plannerContext.getTimeZone()).getMillis()
-            )
-        );
-      }
-    } else if (SqlTypeName.BOOLEAN == sqlTypeName) {
-      return DruidExpression.ofLiteral(
-          columnType,
-          DruidExpression.longLiteral(RexLiteral.booleanValue(rexNode) ? 1 : 0)
+      retVal = new DruidLiteral(
+          ExpressionType.LONG,
+          Calcites.calciteDateTimeLiteralToJoda(rexNode, plannerContext.getTimeZone()).getMillis()
       );
+    } else if (SqlTypeName.BOOLEAN == sqlTypeName) {
+      retVal = new DruidLiteral(ExpressionType.LONG, RexLiteral.booleanValue(rexNode) ? 1L : 0L);
     } else {
       // Can't translate other literals.
       return null;
     }
+
+    return retVal;
   }
 
   /**
@@ -647,8 +673,8 @@ private static DimFilter toSimpleLeafFilter(
 
       final DruidExpression rhsExpression = toDruidExpression(plannerContext, rowSignature, rhs);
       final Expr rhsParsed = rhsExpression != null
-                                       ? plannerContext.parseExpression(rhsExpression.getExpression())
-                                       : null;
+                             ? plannerContext.parseExpression(rhsExpression.getExpression())
+                             : null;
       // rhs must be a literal
       if (rhsParsed == null || !rhsParsed.isLiteral()) {
         return null;
@@ -815,7 +841,9 @@ private static DimFilter toSimpleLeafFilter(
       }
     } else if (rexNode instanceof RexCall) {
       final SqlOperator operator = ((RexCall) rexNode).getOperator();
-      final SqlOperatorConversion conversion = plannerContext.getPlannerToolbox().operatorTable().lookupOperatorConversion(operator);
+      final SqlOperatorConversion conversion = plannerContext.getPlannerToolbox()
+                                                             .operatorTable()
+                                                             .lookupOperatorConversion(operator);
 
       if (conversion == null) {
         return null;

diff --git a/.../java/org/apache/druid/sql/calcite/expression/builtin/ArrayOverlapOperatorConversion.java b/.../java/org/apache/druid/sql/calcite/expression/builtin/ArrayOverlapOperatorConversion.java
@@ -25,7 +25,6 @@
 import org.apache.calcite.sql.type.OperandTypes;
 import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeFamily;
-import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.math.expr.Evals;
 import org.apache.druid.math.expr.Expr;
 import org.apache.druid.math.expr.ExprEval;
@@ -34,10 +33,8 @@
 import org.apache.druid.query.filter.ArrayContainsElementFilter;
 import org.apache.druid.query.filter.DimFilter;
 import org.apache.druid.query.filter.EqualityFilter;
-import org.apache.druid.query.filter.InDimFilter;
 import org.apache.druid.query.filter.NullFilter;
 import org.apache.druid.query.filter.OrDimFilter;
-import org.apache.druid.query.filter.TypedInFilter;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.sql.calcite.expression.DruidExpression;
@@ -158,27 +155,13 @@ public DimFilter toDruidFilter(
           );
         }
       } else {
-        if (plannerContext.isUseBoundsAndSelectors() || NullHandling.replaceWithDefault() || !simpleExtractionExpr.isDirectColumnAccess()) {
-          final InDimFilter.ValuesSet valuesSet = InDimFilter.ValuesSet.create();
-          for (final Object arrayElement : arrayElements) {
-            valuesSet.add(Evals.asString(arrayElement));
-          }
-
-          return new InDimFilter(
-              simpleExtractionExpr.getSimpleExtraction().getColumn(),
-              valuesSet,
-              simpleExtractionExpr.getSimpleExtraction().getExtractionFn(),
-              null
-          );
-        } else {
-          return new TypedInFilter(
-             simpleExtractionExpr.getSimpleExtraction().getColumn(),
-             ExpressionType.toColumnType((ExpressionType) exprEval.type().getElementType()),
-             Arrays.asList(arrayElements),
-             null,
-             null
-          );
-        }
+        return ScalarInArrayOperatorConversion.makeInFilter(
+            plannerContext,
+            simpleExtractionExpr.getSimpleExtraction().getColumn(),
+            simpleExtractionExpr.getSimpleExtraction().getExtractionFn(),
+            Arrays.asList(arrayElements),
+            ExpressionType.toColumnType((ExpressionType) exprEval.type().getElementType())
+        );
       }
     }