Add RowIdSupplier to ColumnSelectorFactory. (#12577)

* Add RowIdSupplier to ColumnSelectorFactory.

This enables virtual columns to cache their outputs in case they are
called multiple times on the same underlying row. This is common for
numeric selectors, where the common pattern is to call isNull() and
then follow with getLong(), getFloat(), or getDouble(). Here, output
caching reduces the number of expression evals by half.

* Fix tests.

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

@@ -53,4 +53,15 @@ public interface ColumnSelectorFactory extends ColumnInspector
   @Override
   @Nullable
   ColumnCapabilities getColumnCapabilities(String column);
+
+  /**
+   * Returns a {@link RowIdSupplier} that allows callers to detect whether the selectors returned by this
+   * factory have moved or not. Useful for selectors that wrap other selectors, such as virtual columns,
+   * as it allows them to cache their outputs.
+   */
+  @Nullable
+  default RowIdSupplier getRowIdSupplier()
+  {
+    return null;
+  }
 }

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

@@ -39,7 +39,7 @@ import java.util.function.Function;
  * It's counterpart for incremental index is {@link
  * org.apache.druid.segment.incremental.IncrementalIndexColumnSelectorFactory}.
  */
-public class QueryableIndexColumnSelectorFactory implements ColumnSelectorFactory
+public class QueryableIndexColumnSelectorFactory implements ColumnSelectorFactory, RowIdSupplier
 {
   private final QueryableIndex index;
   private final VirtualColumns virtualColumns;
@@ -193,6 +193,19 @@ public class QueryableIndexColumnSelectorFactory implements ColumnSelectorFactor
     }
   }
 
+  @Nullable
+  @Override
+  public RowIdSupplier getRowIdSupplier()
+  {
+    return this;
+  }
+
+  @Override
+  public long getRowId()
+  {
+    return offset.getOffset();
+  }
+
   @Override
   @Nullable
   public ColumnCapabilities getColumnCapabilities(String columnName)

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

@@ -40,7 +40,6 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.function.Function;
-import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 import java.util.function.ToLongFunction;
 
@@ -49,12 +48,10 @@ import java.util.function.ToLongFunction;
  */
 public class RowBasedColumnSelectorFactory<T> implements ColumnSelectorFactory
 {
-  private static final long NO_ID = -1;
-
   private final Supplier<T> rowSupplier;
 
   @Nullable
-  private final LongSupplier rowIdSupplier;
+  private final RowIdSupplier rowIdSupplier;
   private final RowAdapter<T> adapter;
   private final ColumnInspector columnInspector;
   private final boolean throwParseExceptions;
@@ -65,7 +62,7 @@ public class RowBasedColumnSelectorFactory<T> implements ColumnSelectorFactory
    */
   RowBasedColumnSelectorFactory(
       final Supplier<T> rowSupplier,
-      @Nullable final LongSupplier rowIdSupplier,
+      @Nullable final RowIdSupplier rowIdSupplier,
       final RowAdapter<T> adapter,
       final ColumnInspector columnInspector,
       final boolean throwParseExceptions
@@ -167,7 +164,7 @@ public class RowBasedColumnSelectorFactory<T> implements ColumnSelectorFactory
 
       return new BaseSingleValueDimensionSelector()
       {
-        private long currentId = NO_ID;
+        private long currentId = RowIdSupplier.INIT;
         private String currentValue;
 
         @Override
@@ -186,11 +183,11 @@ public class RowBasedColumnSelectorFactory<T> implements ColumnSelectorFactory
 
         private void updateCurrentValue()
         {
-          if (rowIdSupplier == null || rowIdSupplier.getAsLong() != currentId) {
+          if (rowIdSupplier == null || rowIdSupplier.getRowId() != currentId) {
             currentValue = extractionFn.apply(timestampFunction.applyAsLong(rowSupplier.get()));
 
             if (rowIdSupplier != null) {
-              currentId = rowIdSupplier.getAsLong();
+              currentId = rowIdSupplier.getRowId();
             }
           }
         }
@@ -200,7 +197,7 @@ public class RowBasedColumnSelectorFactory<T> implements ColumnSelectorFactory
 
       return new DimensionSelector()
       {
-        private long currentId = NO_ID;
+        private long currentId = RowIdSupplier.INIT;
         private List<String> dimensionValues;
 
         private final RangeIndexedInts indexedInts = new RangeIndexedInts();
@@ -331,7 +328,7 @@ public class RowBasedColumnSelectorFactory<T> implements ColumnSelectorFactory
 
         private void updateCurrentValues()
         {
-          if (rowIdSupplier == null || rowIdSupplier.getAsLong() != currentId) {
+          if (rowIdSupplier == null || rowIdSupplier.getRowId() != currentId) {
             try {
               final Object rawValue = dimFunction.apply(rowSupplier.get());
 
@@ -377,12 +374,12 @@ public class RowBasedColumnSelectorFactory<T> implements ColumnSelectorFactory
               }
             }
             catch (Throwable e) {
-              currentId = NO_ID;
+              currentId = RowIdSupplier.INIT;
               throw e;
             }
 
             if (rowIdSupplier != null) {
-              currentId = rowIdSupplier.getAsLong();
+              currentId = rowIdSupplier.getRowId();
             }
           }
         }
@@ -491,6 +488,13 @@ public class RowBasedColumnSelectorFactory<T> implements ColumnSelectorFactory
     }
   }
 
+  @Nullable
+  @Override
+  public RowIdSupplier getRowIdSupplier()
+  {
+    return rowIdSupplier;
+  }
+
   @Nullable
   @Override
   public ColumnCapabilities getColumnCapabilities(String columnName)

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

@@ -0,0 +1,43 @@
+/*
+ * 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;
+
+/**
+ * Returned by {@link ColumnSelectorFactory#getRowIdSupplier()}. Allows users of {@link ColumnSelectorFactory}
+ * to cache objects returned by their selectors.
+ */
+public interface RowIdSupplier
+{
+  /**
+   * A number that will never be returned from {@link #getRowId()}. Useful for initialization.
+   */
+  long INIT = -1;
+
+  /**
+   * Returns a number that uniquely identifies the current position of some underlying cursor. This is useful for
+   * caching: it is safe to assume nothing has changed in the selector as long as the row ID stays the same.
+   *
+   * Row IDs do not need to be contiguous or monotonic. They need not have any meaning. In particular: they may not
+   * be row *numbers* (row number 0 may have any arbitrary row ID).
+   *
+   * Valid row IDs are always nonnegative.
+   */
+  long getRowId();
+}

processing/src/main/java/org/apache/druid/segment/incremental/IncrementalIndexColumnSelectorFactory.java

@@ -25,6 +25,7 @@ import org.apache.druid.segment.ColumnSelectorFactory;
 import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.DimensionIndexer;
 import org.apache.druid.segment.DimensionSelector;
+import org.apache.druid.segment.RowIdSupplier;
 import org.apache.druid.segment.SingleScanTimeDimensionSelector;
 import org.apache.druid.segment.VirtualColumns;
 import org.apache.druid.segment.column.ColumnCapabilities;
@@ -37,7 +38,7 @@ import javax.annotation.Nullable;
  * The basic implementation of {@link ColumnSelectorFactory} over an {@link IncrementalIndex}. It's counterpart for
  * historical segments is {@link org.apache.druid.segment.QueryableIndexColumnSelectorFactory}.
  */
-class IncrementalIndexColumnSelectorFactory implements ColumnSelectorFactory
+class IncrementalIndexColumnSelectorFactory implements ColumnSelectorFactory, RowIdSupplier
 {
   private final IncrementalIndexStorageAdapter adapter;
   private final IncrementalIndex index;
@@ -133,4 +134,17 @@ class IncrementalIndexColumnSelectorFactory implements ColumnSelectorFactory
     // Use adapter.getColumnCapabilities instead of index.getCapabilities (see note in IncrementalIndexStorageAdapater)
     return adapter.getColumnCapabilities(columnName);
   }
+
+  @Nullable
+  @Override
+  public RowIdSupplier getRowIdSupplier()
+  {
+    return this;
+  }
+
+  @Override
+  public long getRowId()
+  {
+    return rowHolder.get().getRowIndex();
+  }
 }

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

@@ -26,6 +26,7 @@ import org.apache.druid.segment.ColumnSelectorFactory;
 import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.Cursor;
 import org.apache.druid.segment.DimensionSelector;
+import org.apache.druid.segment.RowIdSupplier;
 import org.apache.druid.segment.column.ColumnCapabilities;
 import org.joda.time.DateTime;
 
@@ -69,8 +70,10 @@ public class HashJoinEngine
                                                       closer
                                                   );
 
-    class JoinColumnSelectorFactory implements ColumnSelectorFactory
+    class JoinColumnSelectorFactory implements ColumnSelectorFactory, RowIdSupplier
     {
+      private long rowId = 0;
+
       @Override
       public DimensionSelector makeDimensionSelector(DimensionSpec dimensionSpec)
       {
@@ -116,6 +119,29 @@ public class HashJoinEngine
           return leftColumnSelectorFactory.getColumnCapabilities(column);
         }
       }
+
+      @Nullable
+      @Override
+      public RowIdSupplier getRowIdSupplier()
+      {
+        return this;
+      }
+
+      @Override
+      public long getRowId()
+      {
+        return rowId;
+      }
+
+      void advanceRowId()
+      {
+        rowId++;
+      }
+
+      void resetRowId()
+      {
+        rowId = 0;
+      }
     }
 
     final JoinColumnSelectorFactory joinColumnSelectorFactory = new JoinColumnSelectorFactory();
@@ -171,6 +197,8 @@ public class HashJoinEngine
       @Override
       public void advanceUninterruptibly()
       {
+        joinColumnSelectorFactory.advanceRowId();
+
         if (joinMatcher.hasMatch()) {
           joinMatcher.nextMatch();
 
@@ -218,6 +246,7 @@ public class HashJoinEngine
       {
         leftCursor.reset();
         joinMatcher.reset();
+        joinColumnSelectorFactory.resetRowId();
       }
     }
 

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

@@ -0,0 +1,115 @@
+/*
+ * 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.virtual;
+
+import org.apache.druid.math.expr.ExprEval;
+import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
+import org.apache.druid.segment.ColumnValueSelector;
+import org.apache.druid.segment.RowIdSupplier;
+
+import javax.annotation.Nullable;
+
+/**
+ * Base class for many (although not all) {@code ColumnValueSelector<ExprEval>}.
+ */
+public abstract class BaseExpressionColumnValueSelector implements ColumnValueSelector<ExprEval>
+{
+  @Nullable
+  private final RowIdSupplier rowIdSupplier;
+  private long currentRowId = RowIdSupplier.INIT;
+  private ExprEval<?> currentEval;
+
+  protected BaseExpressionColumnValueSelector(@Nullable RowIdSupplier rowIdSupplier)
+  {
+    this.rowIdSupplier = rowIdSupplier;
+  }
+
+  @Override
+  public double getDouble()
+  {
+    // No assert for null handling as ExprEval already has it.
+    return computeCurrentEval().asDouble();
+  }
+
+  @Override
+  public float getFloat()
+  {
+    // No assert for null handling as ExprEval already has it.
+    return (float) computeCurrentEval().asDouble();
+  }
+
+  @Override
+  public long getLong()
+  {
+    // No assert for null handling as ExprEval already has it.
+    return computeCurrentEval().asLong();
+  }
+
+  @Override
+  public boolean isNull()
+  {
+    // It is possible for an expression to have a non-null String value, but be null when treated as long/float/double.
+    // Check specifically for numeric nulls, which matches the expected behavior of ColumnValueSelector.isNull.
+    return computeCurrentEval().isNumericNull();
+  }
+
+  @Nullable
+  @Override
+  public ExprEval<?> getObject()
+  {
+    return computeCurrentEval();
+  }
+
+  @Override
+  public Class<ExprEval> classOfObject()
+  {
+    return ExprEval.class;
+  }
+
+  @Override
+  public void inspectRuntimeShape(final RuntimeShapeInspector inspector)
+  {
+    inspector.visit("rowIdSupplier", this);
+  }
+
+  /**
+   * Implementations override this.
+   */
+  protected abstract ExprEval<?> eval();
+
+  /**
+   * Call {@link #eval()} or use {@code currentEval} as appropriate.
+   */
+  private ExprEval<?> computeCurrentEval()
+  {
+    if (rowIdSupplier == null) {
+      return eval();
+    } else {
+      final long rowId = rowIdSupplier.getRowId();
+
+      if (currentRowId != rowId) {
+        currentEval = eval();
+        currentRowId = rowId;
+      }
+
+      return currentEval;
+    }
+  }
+}

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

@@ -24,54 +24,32 @@ import org.apache.druid.math.expr.Expr;
 import org.apache.druid.math.expr.ExprEval;
 import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import org.apache.druid.segment.ColumnValueSelector;
+import org.apache.druid.segment.RowIdSupplier;
 
-import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 /**
  * Basic expression {@link ColumnValueSelector}. Evaluates {@link Expr} into {@link ExprEval} against
  * {@link Expr.ObjectBinding} which are backed by the underlying expression input {@link ColumnValueSelector}s
  */
-public class ExpressionColumnValueSelector implements ColumnValueSelector<ExprEval>
+public class ExpressionColumnValueSelector extends BaseExpressionColumnValueSelector
 {
-  final Expr.ObjectBinding bindings;
-  final Expr expression;
+  private final Expr.ObjectBinding bindings;
+  private final Expr expression;
 
-  public ExpressionColumnValueSelector(Expr expression, Expr.ObjectBinding bindings)
+  public ExpressionColumnValueSelector(
+      Expr expression,
+      Expr.ObjectBinding bindings,
+      @Nullable RowIdSupplier rowIdSupplier
+  )
   {
+    super(rowIdSupplier);
     this.bindings = Preconditions.checkNotNull(bindings, "bindings");
     this.expression = Preconditions.checkNotNull(expression, "expression");
   }
 
   @Override
-  public double getDouble()
-  {
-    // No Assert for null handling as ExprEval already have it.
-    return getObject().asDouble();
-  }
-
-  @Override
-  public float getFloat()
-  {
-    // No Assert for null handling as ExprEval already have it.
-    return (float) getObject().asDouble();
-  }
-
-  @Override
-  public long getLong()
-  {
-    // No Assert for null handling as ExprEval already have it.
-    return getObject().asLong();
-  }
-
-  @Override
-  public Class<ExprEval> classOfObject()
-  {
-    return ExprEval.class;
-  }
-
-  @Nonnull
-  @Override
-  public ExprEval getObject()
+  protected ExprEval<?> eval()
   {
     return expression.eval(bindings);
   }
@@ -79,16 +57,8 @@ public class ExpressionColumnValueSelector implements ColumnValueSelector<ExprEv
   @Override
   public void inspectRuntimeShape(RuntimeShapeInspector inspector)
   {
+    super.inspectRuntimeShape(inspector);
     inspector.visit("expression", expression);
     inspector.visit("bindings", bindings);
   }
-
-  @Override
-  public boolean isNull()
-  {
-    // It is possible for an expression to have a non-null String value but it can return null when parsed
-    // as a primitive long/float/double.
-    // ExprEval.isNumericNull checks whether the parsed primitive value is null or not.
-    return getObject().isNumericNull();
-  }
 }

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

@@ -40,6 +40,7 @@ import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.ConstantExprEvalSelector;
 import org.apache.druid.segment.DimensionSelector;
 import org.apache.druid.segment.NilColumnValueSelector;
+import org.apache.druid.segment.RowIdSupplier;
 import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.segment.column.ColumnType;
@@ -143,6 +144,8 @@ public class ExpressionSelectors
       ExpressionPlan plan
   )
   {
+    final RowIdSupplier rowIdSupplier = columnSelectorFactory.getRowIdSupplier();
+
     if (plan.is(ExpressionPlan.Trait.SINGLE_INPUT_SCALAR)) {
       final String column = plan.getSingleInputName();
       final ColumnType inputType = plan.getSingleInputType();
@@ -150,12 +153,14 @@ public class ExpressionSelectors
         return new SingleLongInputCachingExpressionColumnValueSelector(
             columnSelectorFactory.makeColumnValueSelector(column),
             plan.getExpression(),
-            !ColumnHolder.TIME_COLUMN_NAME.equals(column) // __time doesn't need an LRU cache since it is sorted.
+            !ColumnHolder.TIME_COLUMN_NAME.equals(column), // __time doesn't need an LRU cache since it is sorted.
+            rowIdSupplier
         );
       } else if (inputType.is(ValueType.STRING)) {
         return new SingleStringInputCachingExpressionColumnValueSelector(
             columnSelectorFactory.makeDimensionSelector(new DefaultDimensionSpec(column, column, ColumnType.STRING)),
-            plan.getExpression()
+            plan.getExpression(),
+            rowIdSupplier
         );
       }
     }
@@ -169,11 +174,11 @@ public class ExpressionSelectors
     // if any unknown column input types, fall back to an expression selector that examines input bindings on a
     // per row basis
     if (plan.any(ExpressionPlan.Trait.UNKNOWN_INPUTS, ExpressionPlan.Trait.INCOMPLETE_INPUTS)) {
-      return new RowBasedExpressionColumnValueSelector(plan, bindings);
+      return new RowBasedExpressionColumnValueSelector(plan, bindings, rowIdSupplier);
     }
 
     // generic expression value selector for fully known input types
-    return new ExpressionColumnValueSelector(plan.getAppliedExpression(), bindings);
+    return new ExpressionColumnValueSelector(plan.getAppliedExpression(), bindings, rowIdSupplier);
   }
 
   /**

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

@@ -24,7 +24,10 @@ import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import org.apache.druid.math.expr.Expr;
 import org.apache.druid.math.expr.ExprEval;
 import org.apache.druid.math.expr.Parser;
+import org.apache.druid.segment.RowIdSupplier;
 
+import javax.annotation.Nullable;
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -37,8 +40,10 @@ import java.util.stream.Collectors;
  * Currently, string dimensions are the only bindings which might appear as a {@link String} or a {@link Object[]}, so
  * numbers are eliminated from the set of 'unknown' bindings to check as they are encountered.
  */
-public class RowBasedExpressionColumnValueSelector extends ExpressionColumnValueSelector
+public class RowBasedExpressionColumnValueSelector extends BaseExpressionColumnValueSelector
 {
+  private final Expr.ObjectBinding bindings;
+  private final Expr expression;
   private final List<String> unknownColumns;
   private final Expr.BindingAnalysis baseBindingAnalysis;
   private final Set<String> ignoredColumns;
@@ -46,10 +51,13 @@ public class RowBasedExpressionColumnValueSelector extends ExpressionColumnValue
 
   public RowBasedExpressionColumnValueSelector(
       ExpressionPlan plan,
-      Expr.ObjectBinding bindings
+      Expr.ObjectBinding bindings,
+      @Nullable RowIdSupplier rowIdSupplier
   )
   {
-    super(plan.getAppliedExpression(), bindings);
+    super(rowIdSupplier);
+    this.bindings = bindings;
+    this.expression = plan.getAppliedExpression();
     this.unknownColumns = plan.getUnknownInputs()
                               .stream()
                               .filter(x -> !plan.getAnalysis().getArrayBindings().contains(x))
@@ -60,10 +68,16 @@ public class RowBasedExpressionColumnValueSelector extends ExpressionColumnValue
   }
 
   @Override
-  public ExprEval getObject()
+  protected ExprEval<?> eval()
   {
     // check to find any arrays for this row
-    List<String> arrayBindings = unknownColumns.stream().filter(this::isBindingArray).collect(Collectors.toList());
+    List<String> arrayBindings = new ArrayList<>();
+
+    for (String unknownColumn : unknownColumns) {
+      if (isBindingArray(unknownColumn)) {
+        arrayBindings.add(unknownColumn);
+      }
+    }
 
     // eliminate anything that will never be an array
     if (ignoredColumns.size() > 0) {

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

@@ -27,6 +27,7 @@ import org.apache.druid.math.expr.ExprEval;
 import org.apache.druid.math.expr.ExpressionType;
 import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import org.apache.druid.segment.ColumnValueSelector;
+import org.apache.druid.segment.RowIdSupplier;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 
 import javax.annotation.Nonnull;
@@ -36,7 +37,7 @@ import javax.annotation.Nullable;
  * Like {@link ExpressionColumnValueSelector}, but caches the most recently computed value and re-uses it in the case
  * of runs in the underlying column. This is especially useful for the __time column, where we expect runs.
  */
-public class SingleLongInputCachingExpressionColumnValueSelector implements ColumnValueSelector<ExprEval>
+public class SingleLongInputCachingExpressionColumnValueSelector extends BaseExpressionColumnValueSelector
 {
   private static final int CACHE_SIZE = 1000;
 
@@ -61,9 +62,12 @@ public class SingleLongInputCachingExpressionColumnValueSelector implements Colu
   public SingleLongInputCachingExpressionColumnValueSelector(
       final ColumnValueSelector selector,
       final Expr expression,
-      final boolean useLruCache
+      final boolean useLruCache,
+      @Nullable RowIdSupplier rowIdSupplier
   )
   {
+    super(rowIdSupplier);
+
     // Verify expression has just one binding.
     if (expression.analyzeInputs().getRequiredBindings().size() != 1) {
       throw new ISE("Expected expression with just one binding");
@@ -77,31 +81,14 @@ public class SingleLongInputCachingExpressionColumnValueSelector implements Colu
   @Override
   public void inspectRuntimeShape(final RuntimeShapeInspector inspector)
   {
+    super.inspectRuntimeShape(inspector);
     inspector.visit("selector", selector);
     inspector.visit("expression", expression);
   }
 
-  @Override
-  public double getDouble()
-  {
-    return getObject().asDouble();
-  }
-
-  @Override
-  public float getFloat()
-  {
-    return (float) getObject().asDouble();
-  }
-
-  @Override
-  public long getLong()
-  {
-    return getObject().asLong();
-  }
-
   @Nonnull
   @Override
-  public ExprEval getObject()
+  protected ExprEval<?> eval()
   {
     // things can still call this even when underlying selector is null (e.g. ExpressionColumnValueSelector#isNull)
     if (selector.isNull()) {
@@ -129,21 +116,6 @@ public class SingleLongInputCachingExpressionColumnValueSelector implements Colu
     return lastOutput;
   }
 
-  @Override
-  public Class<ExprEval> classOfObject()
-  {
-    return ExprEval.class;
-  }
-
-  @Override
-  public boolean isNull()
-  {
-    // It is possible for an expression to have a non-null String value but it can return null when parsed
-    // as a primitive long/float/double.
-    // ExprEval.isNumericNull checks whether the parsed primitive value is null or not.
-    return getObject().isNumericNull();
-  }
-
   public class LruEvalCache
   {
     private final Long2ObjectLinkedOpenHashMap<ExprEval> m = new Long2ObjectLinkedOpenHashMap<>();

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

@@ -28,9 +28,9 @@ import org.apache.druid.math.expr.ExprEval;
 import org.apache.druid.math.expr.ExpressionType;
 import org.apache.druid.math.expr.InputBindings;
 import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
-import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.DimensionDictionarySelector;
 import org.apache.druid.segment.DimensionSelector;
+import org.apache.druid.segment.RowIdSupplier;
 import org.apache.druid.segment.data.IndexedInts;
 
 import javax.annotation.Nullable;
@@ -39,7 +39,7 @@ import javax.annotation.Nullable;
  * Like {@link ExpressionColumnValueSelector}, but caches results for the first CACHE_SIZE dictionary IDs of
  * a string column. Must only be used on selectors with dictionaries.
  */
-public class SingleStringInputCachingExpressionColumnValueSelector implements ColumnValueSelector<ExprEval>
+public class SingleStringInputCachingExpressionColumnValueSelector extends BaseExpressionColumnValueSelector
 {
   private static final int CACHE_SIZE = 1000;
 
@@ -53,9 +53,12 @@ public class SingleStringInputCachingExpressionColumnValueSelector implements Co
 
   public SingleStringInputCachingExpressionColumnValueSelector(
       final DimensionSelector selector,
-      final Expr expression
+      final Expr expression,
+      @Nullable final RowIdSupplier rowIdSupplier
   )
   {
+    super(rowIdSupplier);
+
     // Verify expression has just one binding.
     if (expression.analyzeInputs().getRequiredBindings().size() != 1) {
       throw new ISE("Expected expression with just one binding");
@@ -81,45 +84,13 @@ public class SingleStringInputCachingExpressionColumnValueSelector implements Co
   @Override
   public void inspectRuntimeShape(final RuntimeShapeInspector inspector)
   {
+    super.inspectRuntimeShape(inspector);
     inspector.visit("selector", selector);
     inspector.visit("expression", expression);
   }
 
   @Override
-  public double getDouble()
-  {
-    // No Assert for null handling as ExprEval already have it.
-    return eval().asDouble();
-  }
-
-  @Override
-  public float getFloat()
-  {
-    // No Assert for null handling as ExprEval already have it.
-    return (float) eval().asDouble();
-  }
-
-  @Override
-  public long getLong()
-  {
-    // No Assert for null handling as ExprEval already have it.
-    return eval().asLong();
-  }
-
-  @Nullable
-  @Override
-  public ExprEval getObject()
-  {
-    return eval();
-  }
-
-  @Override
-  public Class<ExprEval> classOfObject()
-  {
-    return ExprEval.class;
-  }
-
-  private ExprEval eval()
+  protected ExprEval<?> eval()
   {
     final IndexedInts row = selector.getRow();
 

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

@@ -24,8 +24,11 @@ import org.apache.druid.query.dimension.DimensionSpec;
 import org.apache.druid.segment.ColumnSelectorFactory;
 import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.DimensionSelector;
+import org.apache.druid.segment.RowIdSupplier;
 import org.apache.druid.segment.VirtualColumns;
 
+import javax.annotation.Nullable;
+
 /**
  * {@link ColumnSelectorFactory} which can create selectors for both virtual and non-virtual columns
  */
@@ -61,4 +64,11 @@ public class VirtualizedColumnSelectorFactory extends VirtualizedColumnInspector
       return baseFactory.makeColumnValueSelector(columnName);
     }
   }
+
+  @Nullable
+  @Override
+  public RowIdSupplier getRowIdSupplier()
+  {
+    return baseFactory.getRowIdSupplier();
+  }
 }

processing/src/test/java/org/apache/druid/segment/virtual/VirtualColumnsTest.java

@@ -56,6 +56,11 @@ import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+import org.mockito.quality.Strictness;
 
 import javax.annotation.Nullable;
 import java.util.Arrays;
@@ -68,6 +73,12 @@ public class VirtualColumnsTest extends InitializedNullHandlingTest
   @Rule
   public ExpectedException expectedException = ExpectedException.none();
 
+  @Rule
+  public MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS);
+
+  @Mock
+  public ColumnSelectorFactory baseColumnSelectorFactory;
+
   @Test
   public void testExists()
   {
@@ -197,24 +208,35 @@ public class VirtualColumnsTest extends InitializedNullHandlingTest
     expectedException.expect(IllegalArgumentException.class);
     expectedException.expectMessage("No such virtual column[bar]");
 
-    virtualColumns.makeColumnValueSelector("bar", null);
+    virtualColumns.makeColumnValueSelector("bar", baseColumnSelectorFactory);
   }
 
   @Test
   public void testMakeSelectors()
   {
+    Mockito.when(baseColumnSelectorFactory.getRowIdSupplier()).thenReturn(null);
+
     final VirtualColumns virtualColumns = makeVirtualColumns();
-    final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector("expr", null);
+    final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector(
+        "expr",
+        baseColumnSelectorFactory
+    );
     final DimensionSelector dimensionSelector = virtualColumns.makeDimensionSelector(
         new DefaultDimensionSpec("expr", "x"),
-        null
+        baseColumnSelectorFactory
     );
     final DimensionSelector extractionDimensionSelector = virtualColumns.makeDimensionSelector(
         new ExtractionDimensionSpec("expr", "x", new BucketExtractionFn(1.0, 0.5)),
-        null
+        baseColumnSelectorFactory
+    );
+    final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector(
+        "expr",
+        baseColumnSelectorFactory
+    );
+    final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector(
+        "expr",
+        baseColumnSelectorFactory
     );
-    final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector("expr", null);
-    final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector("expr", null);
 
     Assert.assertEquals(1L, objectSelector.getObject());
     Assert.assertEquals("1", dimensionSelector.lookupName(dimensionSelector.getRow().get(0)));
@@ -227,13 +249,22 @@ public class VirtualColumnsTest extends InitializedNullHandlingTest
   public void testMakeSelectorsWithDotSupport()
   {
     final VirtualColumns virtualColumns = makeVirtualColumns();
-    final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector("foo.5", null);
+    final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector(
+        "foo.5",
+        baseColumnSelectorFactory
+    );
     final DimensionSelector dimensionSelector = virtualColumns.makeDimensionSelector(
         new DefaultDimensionSpec("foo.5", "x"),
-        null
+        baseColumnSelectorFactory
+    );
+    final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector(
+        "foo.5",
+        baseColumnSelectorFactory
+    );
+    final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector(
+        "foo.5",
+        baseColumnSelectorFactory
     );
-    final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector("foo.5", null);
-    final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector("foo.5", null);
 
     Assert.assertEquals(5L, objectSelector.getObject());
     Assert.assertEquals("5", dimensionSelector.lookupName(dimensionSelector.getRow().get(0)));
@@ -245,13 +276,22 @@ public class VirtualColumnsTest extends InitializedNullHandlingTest
   public void testMakeSelectorsWithDotSupportBaseNameOnly()
   {
     final VirtualColumns virtualColumns = makeVirtualColumns();
-    final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector("foo", null);
+    final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector(
+        "foo",
+        baseColumnSelectorFactory
+    );
     final DimensionSelector dimensionSelector = virtualColumns.makeDimensionSelector(
         new DefaultDimensionSpec("foo", "x"),
-        null
+        baseColumnSelectorFactory
+    );
+    final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector(
+        "foo",
+        baseColumnSelectorFactory
+    );
+    final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector(
+        "foo",
+        baseColumnSelectorFactory
     );
-    final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector("foo", null);
-    final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector("foo", null);
 
     Assert.assertEquals(-1L, objectSelector.getObject());
     Assert.assertEquals("-1", dimensionSelector.lookupName(dimensionSelector.getRow().get(0)));