vectorize constant expressions with optimized selectors

processing/src/main/java/org/apache/druid/segment/vector/ConstantVectorSelectors.java

@@ -0,0 +1,172 @@
+/*
+ * 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.segment.vector;
+
+import org.apache.druid.segment.IdLookup;
+
+import javax.annotation.Nullable;
+import java.util.Arrays;
+
+public class ConstantVectorSelectors
+{
+  public static VectorValueSelector vectorValueSelector(VectorSizeInspector inspector, @Nullable Number constant)
+  {
+    if (constant == null) {
+      return NilVectorSelector.create(inspector);
+    }
+    final long[] longVector = new long[inspector.getMaxVectorSize()];
+    final float[] floatVector = new float[inspector.getMaxVectorSize()];
+    final double[] doubleVector = new double[inspector.getMaxVectorSize()];
+    Arrays.fill(longVector, constant.longValue());
+    Arrays.fill(floatVector, constant.floatValue());
+    Arrays.fill(doubleVector, constant.doubleValue());
+    return new VectorValueSelector()
+    {
+      @Override
+      public long[] getLongVector()
+      {
+        return longVector;
+      }
+
+      @Override
+      public float[] getFloatVector()
+      {
+        return floatVector;
+      }
+
+      @Override
+      public double[] getDoubleVector()
+      {
+        return doubleVector;
+      }
+
+      @Nullable
+      @Override
+      public boolean[] getNullVector()
+      {
+        return null;
+      }
+
+      @Override
+      public int getMaxVectorSize()
+      {
+        return inspector.getMaxVectorSize();
+      }
+
+      @Override
+      public int getCurrentVectorSize()
+      {
+        return inspector.getCurrentVectorSize();
+      }
+    };
+  }
+
+  public static VectorObjectSelector vectorObjectSelector(
+      VectorSizeInspector inspector,
+      @Nullable Object object
+  )
+  {
+    if (object == null) {
+      return NilVectorSelector.create(inspector);
+    }
+
+    final Object[] objects = new Object[inspector.getMaxVectorSize()];
+    Arrays.fill(objects, object);
+
+    return new VectorObjectSelector()
+    {
+      @Override
+      public Object[] getObjectVector()
+      {
+        return objects;
+      }
+
+      @Override
+      public int getMaxVectorSize()
+      {
+        return inspector.getMaxVectorSize();
+      }
+
+      @Override
+      public int getCurrentVectorSize()
+      {
+        return inspector.getCurrentVectorSize();
+      }
+    };
+  }
+
+  public static SingleValueDimensionVectorSelector singleValueDimensionVectorSelector(
+      VectorSizeInspector inspector,
+      @Nullable String value
+  )
+  {
+    if (value == null) {
+      return NilVectorSelector.create(inspector);
+    }
+
+    final int[] row = new int[inspector.getMaxVectorSize()];
+    return new SingleValueDimensionVectorSelector()
+    {
+      @Override
+      public int[] getRowVector()
+      {
+        return row;
+      }
+
+      @Override
+      public int getValueCardinality()
+      {
+        return 1;
+      }
+
+      @Nullable
+      @Override
+      public String lookupName(int id)
+      {
+        return value;
+      }
+
+      @Override
+      public boolean nameLookupPossibleInAdvance()
+      {
+        return true;
+      }
+
+      @Nullable
+      @Override
+      public IdLookup idLookup()
+      {
+        return null;
+      }
+
+      @Override
+      public int getMaxVectorSize()
+      {
+        return inspector.getMaxVectorSize();
+      }
+
+      @Override
+      public int getCurrentVectorSize()
+      {
+        return inspector.getCurrentVectorSize();
+      }
+    };
+  }
+}

processing/src/main/java/org/apache/druid/segment/virtual/ExpressionPlan.java

@@ -100,6 +100,11 @@ public enum Trait
     this.unappliedInputs = unappliedInputs;
   }
 
+  public boolean isConstant()
+  {
+    return analysis.getRequiredBindings().isEmpty();
+  }
+
   public Expr getExpression()
   {
     return expression;

processing/src/main/java/org/apache/druid/segment/virtual/ExpressionVectorSelectors.java

@@ -23,8 +23,11 @@
 import org.apache.druid.math.expr.Expr;
 import org.apache.druid.math.expr.ExprType;
 import org.apache.druid.math.expr.vector.ExprVectorProcessor;
+import org.apache.druid.query.expression.ExprUtils;
 import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ValueType;
+import org.apache.druid.segment.vector.ConstantVectorSelectors;
+import org.apache.druid.segment.vector.SingleValueDimensionVectorSelector;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorObjectSelector;
 import org.apache.druid.segment.vector.VectorValueSelector;
@@ -38,13 +41,36 @@ private ExpressionVectorSelectors()
     // No instantiation.
   }
 
+  public static SingleValueDimensionVectorSelector makeSingleValueDimensionVectorSelector(
+      VectorColumnSelectorFactory factory,
+      Expr expression
+  )
+  {
+    final ExpressionPlan plan = ExpressionPlanner.plan(factory, expression);
+    Preconditions.checkArgument(plan.is(ExpressionPlan.Trait.VECTORIZABLE));
+    // only constant expressions are currently supported, nothing else should get here
+
+    if (plan.isConstant()) {
+      String constant = plan.getExpression().eval(ExprUtils.nilBindings()).asString();
+      return ConstantVectorSelectors.singleValueDimensionVectorSelector(factory.getVectorSizeInspector(), constant);
+    }
+    throw new IllegalStateException("Only constant expressions currently support dimension selectors");
+  }
+
   public static VectorValueSelector makeVectorValueSelector(
       VectorColumnSelectorFactory factory,
       Expr expression
   )
   {
     final ExpressionPlan plan = ExpressionPlanner.plan(factory, expression);
     Preconditions.checkArgument(plan.is(ExpressionPlan.Trait.VECTORIZABLE));
+
+    if (plan.isConstant()) {
+      return ConstantVectorSelectors.vectorValueSelector(
+          factory.getVectorSizeInspector(),
+          (Number) plan.getExpression().eval(ExprUtils.nilBindings()).value()
+      );
+    }
     final Expr.VectorInputBinding bindings = createVectorBindings(plan.getAnalysis(), factory);
     final ExprVectorProcessor<?> processor = plan.getExpression().buildVectorized(bindings);
     return new ExpressionVectorValueSelector(processor, bindings);
@@ -57,6 +83,14 @@ public static VectorObjectSelector makeVectorObjectSelector(
   {
     final ExpressionPlan plan = ExpressionPlanner.plan(factory, expression);
     Preconditions.checkArgument(plan.is(ExpressionPlan.Trait.VECTORIZABLE));
+
+    if (plan.isConstant()) {
+      return ConstantVectorSelectors.vectorObjectSelector(
+          factory.getVectorSizeInspector(),
+          plan.getExpression().eval(ExprUtils.nilBindings()).value()
+      );
+    }
+
     final Expr.VectorInputBinding bindings = createVectorBindings(plan.getAnalysis(), factory);
     final ExprVectorProcessor<?> processor = plan.getExpression().buildVectorized(bindings);
     return new ExpressionVectorObjectSelector(processor, bindings);

processing/src/main/java/org/apache/druid/segment/virtual/ExpressionVirtualColumn.java

@@ -42,6 +42,7 @@
 import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ColumnCapabilitiesImpl;
 import org.apache.druid.segment.column.ValueType;
+import org.apache.druid.segment.vector.SingleValueDimensionVectorSelector;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorObjectSelector;
 import org.apache.druid.segment.vector.VectorValueSelector;
@@ -146,6 +147,15 @@ public boolean canVectorize(ColumnInspector inspector)
     return plan.is(ExpressionPlan.Trait.VECTORIZABLE);
   }
 
+  @Override
+  public SingleValueDimensionVectorSelector makeSingleValueVectorDimensionSelector(
+      DimensionSpec dimensionSpec,
+      VectorColumnSelectorFactory factory
+  )
+  {
+    return ExpressionVectorSelectors.makeSingleValueDimensionVectorSelector(factory, parsedExpression.get());
+  }
+
   @Override
   public VectorValueSelector makeVectorValueSelector(String columnName, VectorColumnSelectorFactory factory)
   {
@@ -200,14 +210,26 @@ public ColumnCapabilities capabilities(ColumnInspector inspector, String columnN
         return ColumnCapabilitiesImpl.createSimpleNumericColumnCapabilities(outputType);
       }
 
+      // array types shouldn't escape the expression system currently, so coerce anything past this point into some
+      // style of string
+
       // we don't have to check for unknown input here because output type is unable to be inferred if we don't know
       // the complete set of input types
       if (plan.any(ExpressionPlan.Trait.NON_SCALAR_OUTPUT, ExpressionPlan.Trait.NEEDS_APPLIED)) {
         // always a multi-value string since wider engine does not yet support array types
         return new ColumnCapabilitiesImpl().setType(ValueType.STRING).setHasMultipleValues(true);
       }
 
-      // if we got here, lets call it single value string output
+      // constant strings are supported as dimension selectors, set them as dictionary encoded and unique
+      if (plan.isConstant()) {
+        return new ColumnCapabilitiesImpl().setType(ValueType.STRING)
+                                           .setDictionaryEncoded(true)
+                                           .setDictionaryValuesUnique(true)
+                                           .setDictionaryValuesSorted(true)
+                                           .setHasMultipleValues(false);
+      }
+
+      // if we got here, lets call it single value string output, non-dictionary encoded
       return new ColumnCapabilitiesImpl().setType(ValueType.STRING)
                                          .setHasMultipleValues(false)
                                          .setDictionaryEncoded(false);