Reduce amount of expression objects created during evaluations

benchmarks/src/test/java/org/apache/druid/benchmark/ExpressionFilterBenchmark.java

@@ -25,6 +25,7 @@
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.math.expr.ExpressionProcessing;
 import org.apache.druid.query.expression.TestExprMacroTable;
 import org.apache.druid.query.filter.AndDimFilter;
 import org.apache.druid.query.filter.DimFilter;
@@ -70,6 +71,7 @@ public class ExpressionFilterBenchmark
 {
   static {
     NullHandling.initializeForTests();
+    ExpressionProcessing.initializeForTests();
   }
 
   @Param({"1000000"})

benchmarks/src/test/java/org/apache/druid/benchmark/ExpressionSelectorBenchmark.java

@@ -25,8 +25,10 @@
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.math.expr.ExpressionProcessing;
 import org.apache.druid.query.dimension.DefaultDimensionSpec;
 import org.apache.druid.query.dimension.ExtractionDimensionSpec;
+import org.apache.druid.query.expression.LookupEnabledTestExprMacroTable;
 import org.apache.druid.query.expression.TestExprMacroTable;
 import org.apache.druid.query.extraction.StrlenExtractionFn;
 import org.apache.druid.query.extraction.TimeFormatExtractionFn;
@@ -59,21 +61,21 @@
 import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.annotations.Warmup;
 import org.openjdk.jmh.infra.Blackhole;
-
 import java.util.BitSet;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
 @State(Scope.Benchmark)
 @Fork(value = 1)
-@Warmup(iterations = 15)
-@Measurement(iterations = 30)
+@Warmup(iterations = 3, time = 3)
+@Measurement(iterations = 10, time = 3)
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
 public class ExpressionSelectorBenchmark
 {
   static {
     NullHandling.initializeForTests();
+    ExpressionProcessing.initializeForTests();
   }
 
   @Param({"1000000"})
@@ -451,6 +453,154 @@ public void stringConcatAndCompareOnLong(Blackhole blackhole)
     blackhole.consume(results);
   }
 
+  @Benchmark
+  public void caseSearched1(Blackhole blackhole)
+  {
+    final Sequence<Cursor> cursors = new QueryableIndexStorageAdapter(index).makeCursors(
+        null,
+        index.getDataInterval(),
+        VirtualColumns.create(
+            ImmutableList.of(
+                new ExpressionVirtualColumn(
+                    "v",
+                    "case_searched(s == 'asd' || isnull(s) || s == 'xxx', 1, s == 'foo' || s == 'bar', 2, 3)",
+                    ColumnType.LONG,
+                    TestExprMacroTable.INSTANCE
+                )
+            )
+        ),
+        Granularities.ALL,
+        false,
+        null
+    );
+
+    final List<?> results = cursors
+        .map(cursor -> {
+          final ColumnValueSelector selector = cursor.getColumnSelectorFactory().makeColumnValueSelector("v");
+          consumeLong(cursor, selector, blackhole);
+          return null;
+        })
+        .toList();
+
+    blackhole.consume(results);
+  }
+
+  @Benchmark
+  public void caseSearched2(Blackhole blackhole)
+  {
+    final Sequence<Cursor> cursors = new QueryableIndexStorageAdapter(index).makeCursors(
+        null,
+        index.getDataInterval(),
+        VirtualColumns.create(
+            ImmutableList.of(
+                new ExpressionVirtualColumn(
+                    "v",
+                    "case_searched(s == 'asd' || isnull(s) || n == 1, 1, n == 2, 2, 3)",
+                    ColumnType.LONG,
+                    TestExprMacroTable.INSTANCE
+                )
+            )
+        ),
+        Granularities.ALL,
+        false,
+        null
+    );
+
+    final List<?> results = cursors
+        .map(cursor -> {
+          final ColumnValueSelector selector = cursor.getColumnSelectorFactory().makeColumnValueSelector("v");
+          consumeLong(cursor, selector, blackhole);
+          return null;
+        })
+        .toList();
+
+    blackhole.consume(results);
+  }
+
+  @Benchmark
+  public void caseSearchedWithLookup(Blackhole blackhole)
+  {
+    final Sequence<Cursor> cursors = new QueryableIndexStorageAdapter(index).makeCursors(
+        null,
+        index.getDataInterval(),
+        VirtualColumns.create(
+            ImmutableList.of(
+                new ExpressionVirtualColumn(
+                    "v",
+                    "case_searched(n == 1001, -1, "
+                        + "lookup(s, 'lookyloo') == 'asd1', 1, "
+                        + "lookup(s, 'lookyloo') == 'asd2', 2, "
+                        + "lookup(s, 'lookyloo') == 'asd3', 3, "
+                        + "lookup(s, 'lookyloo') == 'asd4', 4, "
+                        + "lookup(s, 'lookyloo') == 'asd5', 5, "
+                        + "-2)",
+                    ColumnType.LONG,
+                    LookupEnabledTestExprMacroTable.INSTANCE
+                )
+            )
+        ),
+        Granularities.ALL,
+        false,
+        null
+    );
+
+    final List<?> results = cursors
+        .map(cursor -> {
+          final ColumnValueSelector selector = cursor.getColumnSelectorFactory().makeColumnValueSelector("v");
+          consumeLong(cursor, selector, blackhole);
+          return null;
+        })
+        .toList();
+
+    blackhole.consume(results);
+  }
+
+  @Benchmark
+  public void caseSearchedWithLookup2(Blackhole blackhole)
+  {
+    final Sequence<Cursor> cursors = new QueryableIndexStorageAdapter(index).makeCursors(
+        null,
+        index.getDataInterval(),
+        VirtualColumns.create(
+            ImmutableList.of(
+                new ExpressionVirtualColumn(
+                    "ll",
+                    "lookup(s, 'lookyloo')",
+                    ColumnType.STRING,
+                    LookupEnabledTestExprMacroTable.INSTANCE
+                ),
+                new ExpressionVirtualColumn(
+                    "v",
+                    "case_searched(n == 1001, -1, "
+                        + "ll == 'asd1', 1, "
+                        + "ll == 'asd2', 2, "
+                        + "ll == 'asd3', 3, "
+                        + "ll == 'asd4', 4, "
+                        + "ll == 'asd5', 5, "
+                        + "-2)",
+                    ColumnType.LONG,
+                    LookupEnabledTestExprMacroTable.INSTANCE
+                )
+            )
+        ),
+        Granularities.ALL,
+        false,
+        null
+    );
+
+    final List<?> results = cursors
+        .map(cursor -> {
+          final ColumnValueSelector selector = cursor.getColumnSelectorFactory().makeColumnValueSelector("v");
+          consumeLong(cursor, selector, blackhole);
+          return null;
+        })
+        .toList();
+
+    blackhole.consume(results);
+  }
+
+
+
   private void consumeDimension(final Cursor cursor, final DimensionSelector selector, final Blackhole blackhole)
   {
     if (selector.getValueCardinality() >= 0) {

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

@@ -157,9 +157,12 @@ public int hashCode()
 
 class LongExpr extends ConstantExpr<Long>
 {
+  private final ExprEval expr;
+
   LongExpr(Long value)
   {
     super(ExpressionType.LONG, Preconditions.checkNotNull(value, "value"));
+    expr = ExprEval.ofLong(value);
   }
 
   @Override
@@ -171,7 +174,7 @@ public String toString()
   @Override
   public ExprEval eval(ObjectBinding bindings)
   {
-    return ExprEval.ofLong(value);
+    return expr;
   }
 
   @Override
@@ -240,9 +243,12 @@ public String toString()
 
 class DoubleExpr extends ConstantExpr<Double>
 {
+  private final ExprEval expr;
+
   DoubleExpr(Double value)
   {
     super(ExpressionType.DOUBLE, Preconditions.checkNotNull(value, "value"));
+    expr = ExprEval.ofDouble(value);
   }
 
   @Override
@@ -254,7 +260,7 @@ public String toString()
   @Override
   public ExprEval eval(ObjectBinding bindings)
   {
-    return ExprEval.ofDouble(value);
+    return expr;
   }
 
   @Override
@@ -323,9 +329,12 @@ public String toString()
 
 class StringExpr extends ConstantExpr<String>
 {
+  private final ExprEval expr;
+
   StringExpr(@Nullable String value)
   {
     super(ExpressionType.STRING, NullHandling.emptyToNullIfNeeded(value));
+    expr = ExprEval.of(value);
   }
 
   @Override
@@ -337,7 +346,7 @@ public String toString()
   @Override
   public ExprEval eval(ObjectBinding bindings)
   {
-    return ExprEval.of(value);
+    return expr;
   }
 
   @Override
@@ -464,15 +473,18 @@ public String toString()
 
 class ComplexExpr extends ConstantExpr<Object>
 {
+  private final ExprEval expr;
+
   protected ComplexExpr(ExpressionType outputType, @Nullable Object value)
   {
     super(outputType, value);
+    expr = ExprEval.ofComplex(outputType, value);
   }
 
   @Override
   public ExprEval eval(ObjectBinding bindings)
   {
-    return ExprEval.ofComplex(outputType, value);
+    return expr;
   }
 
   @Override

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

@@ -363,7 +363,7 @@ public static ExprEval ofBoolean(boolean value, ExprType type)
       case DOUBLE:
         return ExprEval.of(Evals.asDouble(value));
       case LONG:
-        return ExprEval.of(Evals.asLong(value));
+        return ofLongBoolean(value);
       case STRING:
         return ExprEval.of(String.valueOf(value));
       default:
@@ -376,7 +376,7 @@ public static ExprEval ofBoolean(boolean value, ExprType type)
    */
   public static ExprEval ofLongBoolean(boolean value)
   {
-    return ExprEval.of(Evals.asLong(value));
+    return value ? LongExprEval.TRUE : LongExprEval.FALSE;
   }
 
   public static ExprEval ofComplex(ExpressionType outputType, @Nullable Object value)
@@ -922,6 +922,8 @@ public Expr toExpr()
 
   private static class LongExprEval extends NumericExprEval
   {
+    private static final LongExprEval TRUE = new LongExprEval(Evals.asLong(true));
+    private static final LongExprEval FALSE = new LongExprEval(Evals.asLong(false));
     private static final LongExprEval OF_NULL = new LongExprEval(null);
 
     private LongExprEval(@Nullable Number value)

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

@@ -55,6 +55,9 @@ public static ExpressionType autoDetect(ExprEval eval, ExprEval otherEval)
   {
     ExpressionType type = eval.type();
     ExpressionType otherType = otherEval.type();
+    if (type == otherType && type.getType().isPrimitive()) {
+      return type;
+    }
     if (Types.is(type, ExprType.STRING) && Types.is(otherType, ExprType.STRING)) {
       return ExpressionType.STRING;
     }

processing/src/test/java/org/apache/druid/math/expr/ExprTest.java

@@ -139,7 +139,7 @@ public void testEqualsContractForUnaryMinusExpr()
   public void testEqualsContractForStringExpr()
   {
     EqualsVerifier.forClass(StringExpr.class)
-                  .withIgnoredFields("outputType")
+                  .withIgnoredFields("outputType", "expr")
                   .withPrefabValues(ExpressionType.class, ExpressionType.STRING, ExpressionType.DOUBLE)
                   .usingGetClass()
                   .verify();
@@ -149,7 +149,7 @@ public void testEqualsContractForStringExpr()
   public void testEqualsContractForDoubleExpr()
   {
     EqualsVerifier.forClass(DoubleExpr.class)
-                  .withIgnoredFields("outputType")
+                  .withIgnoredFields("outputType", "expr")
                   .withPrefabValues(ExpressionType.class, ExpressionType.DOUBLE, ExpressionType.LONG)
                   .usingGetClass()
                   .verify();
@@ -159,7 +159,7 @@ public void testEqualsContractForDoubleExpr()
   public void testEqualsContractForLongExpr()
   {
     EqualsVerifier.forClass(LongExpr.class)
-                  .withIgnoredFields("outputType")
+                  .withIgnoredFields("outputType", "expr")
                   .withPrefabValues(ExpressionType.class, ExpressionType.LONG, ExpressionType.STRING)
                   .usingGetClass()
                   .verify();
@@ -187,6 +187,7 @@ public void testEqualsContractForComplexExpr()
                       ExpressionTypeFactory.getInstance().ofComplex("bar")
                   )
                   .withNonnullFields("outputType")
+                  .withIgnoredFields("expr")
                   .usingGetClass()
                   .verify();
   }

processing/src/test/java/org/apache/druid/math/expr/OutputTypeTest.java

@@ -453,6 +453,7 @@ public void testEvalAutoConversion()
     final ExprEval<?> longEval = ExprEval.of(1L);
     final ExprEval<?> doubleEval = ExprEval.of(1.0);
     final ExprEval<?> arrayEval = ExprEval.ofLongArray(new Long[]{1L, 2L, 3L});
+    final ExprEval<?> complexEval = ExprEval.ofComplex(ExpressionType.UNKNOWN_COMPLEX, new Object());
 
     // only long stays long
     Assert.assertEquals(ExpressionType.LONG, ExpressionTypeConversion.autoDetect(longEval, longEval));
@@ -479,6 +480,11 @@ public void testEvalAutoConversion()
     Assert.assertEquals(ExpressionType.DOUBLE, ExpressionTypeConversion.autoDetect(nullStringEval, arrayEval));
     Assert.assertEquals(ExpressionType.DOUBLE, ExpressionTypeConversion.autoDetect(doubleEval, arrayEval));
     Assert.assertEquals(ExpressionType.DOUBLE, ExpressionTypeConversion.autoDetect(longEval, arrayEval));
+
+    Assert.assertEquals(ExpressionType.DOUBLE, ExpressionTypeConversion.autoDetect(longEval, complexEval));
+    Assert.assertEquals(ExpressionType.DOUBLE, ExpressionTypeConversion.autoDetect(doubleEval, complexEval));
+    Assert.assertEquals(ExpressionType.DOUBLE, ExpressionTypeConversion.autoDetect(arrayEval, complexEval));
+    Assert.assertEquals(ExpressionType.DOUBLE, ExpressionTypeConversion.autoDetect(complexEval, complexEval));
   }
 
   @Test