Add aggregator factory check to determine if vectorized aggregators are supported for the inputs

processing/src/main/java/org/apache/druid/query/aggregation/AggregatorFactory.java

@@ -24,6 +24,7 @@
 import org.apache.druid.java.util.common.UOE;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.query.PerSegmentQueryOptimizationContext;
+import org.apache.druid.segment.ColumnInspector;
 import org.apache.druid.segment.ColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 
@@ -68,7 +69,7 @@ public VectorAggregator factorizeVector(VectorColumnSelectorFactory selectorFact
   /**
    * Returns whether or not this aggregation class supports vectorization. The default implementation returns false.
    */
-  public boolean canVectorize()
+  public boolean canVectorize(ColumnInspector columnInspector)
   {
     return false;
   }

processing/src/main/java/org/apache/druid/query/aggregation/CountAggregatorFactory.java

@@ -23,6 +23,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import org.apache.druid.segment.ColumnInspector;
 import org.apache.druid.segment.ColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 
@@ -72,7 +73,7 @@ public Comparator getComparator()
   }
 
   @Override
-  public boolean canVectorize()
+  public boolean canVectorize(ColumnInspector columnInspector)
   {
     return true;
   }

processing/src/main/java/org/apache/druid/query/aggregation/DoubleSumAggregatorFactory.java

@@ -25,6 +25,9 @@
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.segment.BaseDoubleColumnValueSelector;
+import org.apache.druid.segment.ColumnInspector;
+import org.apache.druid.segment.column.ColumnCapabilities;
+import org.apache.druid.segment.column.ValueType;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorValueSelector;
 
@@ -79,8 +82,12 @@ protected VectorValueSelector vectorSelector(VectorColumnSelectorFactory columnS
   }
 
   @Override
-  public boolean canVectorize()
+  public boolean canVectorize(ColumnInspector columnInspector)
   {
+    if (fieldName != null) {
+      final ColumnCapabilities capabilities = columnInspector.getColumnCapabilities(fieldName);
+      return expression == null && (capabilities == null || ValueType.isNumeric(capabilities.getType()));
+    }
     return expression == null;
   }
 

processing/src/main/java/org/apache/druid/query/aggregation/FilteredAggregatorFactory.java

@@ -29,6 +29,7 @@
 import org.apache.druid.query.filter.IntervalDimFilter;
 import org.apache.druid.query.filter.ValueMatcher;
 import org.apache.druid.query.filter.vector.VectorValueMatcher;
+import org.apache.druid.segment.ColumnInspector;
 import org.apache.druid.segment.ColumnSelectorFactory;
 import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
@@ -98,7 +99,7 @@ public BufferAggregator factorizeBuffered(ColumnSelectorFactory columnSelectorFa
   @Override
   public VectorAggregator factorizeVector(VectorColumnSelectorFactory columnSelectorFactory)
   {
-    Preconditions.checkState(canVectorize(), "Cannot vectorize");
+    Preconditions.checkState(canVectorize(columnSelectorFactory), "Cannot vectorize");
     final VectorValueMatcher valueMatcher = filter.makeVectorMatcher(columnSelectorFactory);
     return new FilteredVectorAggregator(
         valueMatcher,
@@ -107,9 +108,9 @@ public VectorAggregator factorizeVector(VectorColumnSelectorFactory columnSelect
   }
 
   @Override
-  public boolean canVectorize()
+  public boolean canVectorize(ColumnInspector columnInspector)
   {
-    return delegate.canVectorize() && filter.canVectorizeMatcher();
+    return delegate.canVectorize(columnInspector) && filter.canVectorizeMatcher();
   }
 
   @Override

processing/src/main/java/org/apache/druid/query/aggregation/FloatSumAggregatorFactory.java

@@ -25,6 +25,8 @@
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.segment.BaseFloatColumnValueSelector;
+import org.apache.druid.segment.ColumnInspector;
+import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorValueSelector;
 
@@ -78,12 +80,15 @@ protected VectorValueSelector vectorSelector(VectorColumnSelectorFactory columnS
   }
 
   @Override
-  public boolean canVectorize()
+  public boolean canVectorize(ColumnInspector columnInspector)
   {
+    if (fieldName != null) {
+      final ColumnCapabilities capabilities = columnInspector.getColumnCapabilities(fieldName);
+      return expression == null && (capabilities == null || capabilities.getType().isNumeric());
+    }
     return expression == null;
   }
 
-
   @Override
   protected VectorAggregator factorizeVector(
       VectorColumnSelectorFactory columnSelectorFactory,

processing/src/main/java/org/apache/druid/query/aggregation/LongSumAggregatorFactory.java

@@ -25,6 +25,8 @@
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.segment.BaseLongColumnValueSelector;
+import org.apache.druid.segment.ColumnInspector;
+import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorValueSelector;
 
@@ -87,8 +89,12 @@ protected VectorAggregator factorizeVector(
   }
 
   @Override
-  public boolean canVectorize()
+  public boolean canVectorize(ColumnInspector columnInspector)
   {
+    if (fieldName != null) {
+      final ColumnCapabilities capabilities = columnInspector.getColumnCapabilities(fieldName);
+      return expression == null && (capabilities == null || capabilities.getType().isNumeric());
+    }
     return expression == null;
   }
 

processing/src/main/java/org/apache/druid/query/aggregation/NullableNumericAggregatorFactory.java

@@ -65,7 +65,7 @@ public final BufferAggregator factorizeBuffered(ColumnSelectorFactory columnSele
   @Override
   public final VectorAggregator factorizeVector(VectorColumnSelectorFactory columnSelectorFactory)
   {
-    Preconditions.checkState(canVectorize(), "Cannot vectorize");
+    Preconditions.checkState(canVectorize(columnSelectorFactory), "Cannot vectorize");
     VectorValueSelector selector = vectorSelector(columnSelectorFactory);
     VectorAggregator aggregator = factorizeVector(columnSelectorFactory, selector);
     return NullHandling.replaceWithDefault() ? aggregator : new NullableNumericVectorAggregator(aggregator, selector);
@@ -135,12 +135,11 @@ protected abstract BufferAggregator factorizeBuffered(
    * @see BufferAggregator
    */
   protected VectorAggregator factorizeVector(
-      // Not used by current aggregators, but here for parity with "factorizeBuffered".
-      @SuppressWarnings("unused") VectorColumnSelectorFactory columnSelectorFactory,
+      VectorColumnSelectorFactory columnSelectorFactory,
       VectorValueSelector selector
   )
   {
-    if (!canVectorize()) {
+    if (!canVectorize(columnSelectorFactory)) {
       throw new UnsupportedOperationException("Cannot vectorize");
     } else {
       throw new UnsupportedOperationException("canVectorize returned true but 'factorizeVector' is not implemented");

processing/src/main/java/org/apache/druid/query/aggregation/SuppressedAggregatorFactory.java

@@ -22,6 +22,7 @@
 import org.apache.druid.query.PerSegmentQueryOptimizationContext;
 import org.apache.druid.query.cache.CacheKeyBuilder;
 import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
+import org.apache.druid.segment.ColumnInspector;
 import org.apache.druid.segment.ColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 
@@ -70,9 +71,9 @@ public VectorAggregator factorizeVector(VectorColumnSelectorFactory columnSelect
   }
 
   @Override
-  public boolean canVectorize()
+  public boolean canVectorize(ColumnInspector columnInspector)
   {
-    return delegate.canVectorize();
+    return delegate.canVectorize(columnInspector);
   }
 
   @Override

processing/src/main/java/org/apache/druid/query/aggregation/hyperloglog/HyperUniquesAggregatorFactory.java

@@ -38,6 +38,7 @@
 import org.apache.druid.query.aggregation.cardinality.HyperLogLogCollectorAggregateCombiner;
 import org.apache.druid.query.cache.CacheKeyBuilder;
 import org.apache.druid.segment.BaseObjectColumnValueSelector;
+import org.apache.druid.segment.ColumnInspector;
 import org.apache.druid.segment.ColumnSelectorFactory;
 import org.apache.druid.segment.NilColumnValueSelector;
 import org.apache.druid.segment.column.ColumnCapabilities;
@@ -140,7 +141,7 @@ public VectorAggregator factorizeVector(final VectorColumnSelectorFactory select
   }
 
   @Override
-  public boolean canVectorize()
+  public boolean canVectorize(ColumnInspector columnInspector)
   {
     return true;
   }

processing/src/main/java/org/apache/druid/query/aggregation/mean/DoubleMeanAggregatorFactory.java

@@ -31,7 +31,9 @@
 import org.apache.druid.query.aggregation.BufferAggregator;
 import org.apache.druid.query.aggregation.VectorAggregator;
 import org.apache.druid.query.cache.CacheKeyBuilder;
+import org.apache.druid.segment.ColumnInspector;
 import org.apache.druid.segment.ColumnSelectorFactory;
+import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 
 import javax.annotation.Nullable;
@@ -106,9 +108,10 @@ public VectorAggregator factorizeVector(final VectorColumnSelectorFactory select
   }
 
   @Override
-  public boolean canVectorize()
+  public boolean canVectorize(ColumnInspector columnInspector)
   {
-    return true;
+    final ColumnCapabilities capabilities = columnInspector.getColumnCapabilities(fieldName);
+    return capabilities == null || capabilities.getType().isNumeric();
   }
 
   @Override

processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/vector/VectorGroupByEngine.java

@@ -28,7 +28,6 @@
 import org.apache.druid.java.util.common.parsers.CloseableIterator;
 import org.apache.druid.query.QueryConfig;
 import org.apache.druid.query.aggregation.AggregatorAdapters;
-import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.dimension.DimensionSpec;
 import org.apache.druid.query.filter.Filter;
 import org.apache.druid.query.groupby.GroupByQuery;
@@ -85,7 +84,7 @@ public static boolean canVectorize(
 
     return GroupByQueryEngineV2.isAllSingleValueDims(adapter::getColumnCapabilities, query.getDimensions(), true)
            && query.getDimensions().stream().allMatch(DimensionSpec::canVectorize)
-           && query.getAggregatorSpecs().stream().allMatch(AggregatorFactory::canVectorize)
+           && query.getAggregatorSpecs().stream().allMatch(aggregatorFactory -> aggregatorFactory.canVectorize(adapter))
            && adapter.canVectorize(filter, query.getVirtualColumns(), false);
   }
 

processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQueryEngine.java

@@ -102,7 +102,7 @@ public Sequence<Result<TimeseriesResultValue>> process(final TimeseriesQuery que
 
     final boolean doVectorize = queryConfigToUse.getVectorize().shouldVectorize(
         adapter.canVectorize(filter, query.getVirtualColumns(), descending)
-        && query.getAggregatorSpecs().stream().allMatch(AggregatorFactory::canVectorize)
+        && query.getAggregatorSpecs().stream().allMatch(aggregatorFactory -> aggregatorFactory.canVectorize(adapter))
     );
 
     final Sequence<Result<TimeseriesResultValue>> result;

processing/src/main/java/org/apache/druid/segment/ColumnInspector.java

@@ -0,0 +1,37 @@
+/*
+ * 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;
+
+import org.apache.druid.segment.column.ColumnCapabilities;
+
+import javax.annotation.Nullable;
+
+public interface ColumnInspector
+{
+  /**
+   * Returns capabilities of a particular column.
+   *
+   * @param column column name
+   *
+   * @return capabilities, or null
+   */
+  @Nullable
+  ColumnCapabilities getColumnCapabilities(String column);
+}

processing/src/main/java/org/apache/druid/segment/ColumnSelectorFactory.java

@@ -31,7 +31,7 @@
  * @see org.apache.druid.segment.vector.VectorColumnSelectorFactory, the vectorized version
  */
 @PublicApi
-public interface ColumnSelectorFactory
+public interface ColumnSelectorFactory extends ColumnInspector
 {
   DimensionSelector makeDimensionSelector(DimensionSpec dimensionSpec);
 
@@ -50,6 +50,7 @@ public interface ColumnSelectorFactory
    *
    * @return capabilities, or null
    */
+  @Override
   @Nullable
   ColumnCapabilities getColumnCapabilities(String column);
 }

processing/src/main/java/org/apache/druid/segment/StorageAdapter.java

@@ -30,7 +30,7 @@
 /**
  */
 @PublicApi
-public interface StorageAdapter extends CursorFactory
+public interface StorageAdapter extends CursorFactory, ColumnInspector
 {
   Interval getInterval();
   Indexed<String> getAvailableDimensions();
@@ -62,6 +62,7 @@ public interface StorageAdapter extends CursorFactory
    *
    * @return capabilities, or null
    */
+  @Override
   @Nullable
   ColumnCapabilities getColumnCapabilities(String column);
 

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

@@ -20,6 +20,7 @@
 package org.apache.druid.segment.vector;
 
 import org.apache.druid.query.dimension.DimensionSpec;
+import org.apache.druid.segment.ColumnInspector;
 import org.apache.druid.segment.column.ColumnCapabilities;
 
 import javax.annotation.Nullable;
@@ -29,7 +30,7 @@
  *
  * @see org.apache.druid.segment.ColumnSelectorFactory, the non-vectorized version.
  */
-public interface VectorColumnSelectorFactory
+public interface VectorColumnSelectorFactory extends ColumnInspector
 {
   /**
    * Returns a {@link VectorSizeInspector} for the {@link VectorCursor} that generated this object.
@@ -72,6 +73,7 @@ default int getMaxVectorSize()
    *
    * @return capabilities, or null if the column doesn't exist.
    */
+  @Override
   @Nullable
   ColumnCapabilities getColumnCapabilities(String column);
 }