feat: remap query virtual columns to clustered base table columns when equivalent

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

@@ -25,6 +25,7 @@
 
 import javax.annotation.Nullable;
 import java.util.List;
+import java.util.Map;
 
 /**
  * {@link Cursor} that walks a sequence of per-group cursors back-to-back, presenting them to the caller as a single
@@ -40,6 +41,7 @@ public final class ConcatenatingCursor implements Cursor
   private final List<Supplier<CursorHolder>> holderSuppliers;
   private final List<Object[]> clusteringValuesByGroup;
   private final ClusteringColumnSelectorFactory wrapperFactory;
+  private final ColumnSelectorFactory exposedFactory;
 
   private int currentIdx;
   @Nullable
@@ -49,7 +51,8 @@ public final class ConcatenatingCursor implements Cursor
   public ConcatenatingCursor(
       List<Supplier<CursorHolder>> holderSuppliers,
       List<Object[]> clusteringValuesByGroup,
-      ClusteringColumnSelectorFactory wrapperFactory
+      ClusteringColumnSelectorFactory wrapperFactory,
+      Map<String, String> virtualColumnRemap
   )
   {
     if (holderSuppliers.size() != clusteringValuesByGroup.size()) {
@@ -65,6 +68,9 @@ public ConcatenatingCursor(
     this.holderSuppliers = holderSuppliers;
     this.clusteringValuesByGroup = clusteringValuesByGroup;
     this.wrapperFactory = wrapperFactory;
+    this.exposedFactory = virtualColumnRemap.isEmpty()
+                          ? wrapperFactory
+                          : new RemapColumnSelectorFactory(wrapperFactory, virtualColumnRemap);
     this.currentIdx = -1;
   }
 
@@ -101,7 +107,7 @@ private void advanceToNextNonEmptyGroup()
   public ColumnSelectorFactory getColumnSelectorFactory()
   {
     initializeIfNeeded();
-    return wrapperFactory;
+    return exposedFactory;
   }
 
   @Override

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

@@ -44,6 +44,7 @@
 import org.apache.druid.segment.vector.ConcatenatingVectorCursor;
 import org.apache.druid.segment.vector.MultiValueDimensionVectorSelector;
 import org.apache.druid.segment.vector.ReadableVectorInspector;
+import org.apache.druid.segment.vector.RemapVectorColumnSelectorFactory;
 import org.apache.druid.segment.vector.SingleValueDimensionVectorSelector;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorCursor;
@@ -215,11 +216,15 @@ protected ColumnSelectorFactory makeColumnSelectorFactoryForOffset(
           Offset baseOffset
       )
       {
-        return new SingleGroupClusteringColumnSelectorFactory(
+        final ColumnSelectorFactory clusteringFactory = new SingleGroupClusteringColumnSelectorFactory(
             super.makeColumnSelectorFactoryForOffset(columnCache, baseOffset),
             valueGroup.getSummary().getClusteringColumns(),
             valueGroup.lookupClusteringValues()
         );
+        if (plan.virtualColumnRemap().isEmpty()) {
+          return clusteringFactory;
+        }
+        return new RemapColumnSelectorFactory(clusteringFactory, plan.virtualColumnRemap());
       }
 
       @Override
@@ -228,11 +233,15 @@ protected VectorColumnSelectorFactory makeVectorColumnSelectorFactoryForOffset(
           VectorOffset baseOffset
       )
       {
-        return new SingleGroupClusteringVectorColumnSelectorFactory(
+        final VectorColumnSelectorFactory clusteringVectorFactory = new SingleGroupClusteringVectorColumnSelectorFactory(
             super.makeVectorColumnSelectorFactoryForOffset(columnCache, baseOffset),
             valueGroup.getSummary().getClusteringColumns(),
             valueGroup.lookupClusteringValues()
         );
+        if (plan.virtualColumnRemap().isEmpty()) {
+          return clusteringVectorFactory;
+        }
+        return new RemapVectorColumnSelectorFactory(clusteringVectorFactory, plan.virtualColumnRemap());
       }
     };
   }
@@ -299,12 +308,14 @@ private CursorHolder makeMultiGroupClusteredCursorHolder(
     final ConcatenatingCursor cursor = new ConcatenatingCursor(
         holderSuppliers,
         clusteringValuesByGroup,
-        wrapperFactory
+        wrapperFactory,
+        plan.virtualColumnRemap()
     );
     final ConcatenatingVectorCursor vectorCursor = new ConcatenatingVectorCursor(
         holderSuppliers,
         clusteringValuesByGroup,
-        vectorWrapperFactory
+        vectorWrapperFactory,
+        plan.virtualColumnRemap()
     );
 
     // each group gets a different rewritten filter, so the conservative thing to do here is require the original query
@@ -313,8 +324,8 @@ private CursorHolder makeMultiGroupClusteredCursorHolder(
     final Filter queryFilter = spec.getFilter();
     final boolean filterCanVectorize =
         queryFilter == null || queryFilter.canVectorizeMatcher(spec.getVirtualColumns().wrapInspector(this));
-    // we still check that the first holder is vectorizable to make sure all the non-filter parts can be vectorized
-    final boolean canVectorize = filterCanVectorize && holderSuppliers.get(0).get().canVectorize();
+    // we still check that the first holder is vectorizable to make sure all the non-filter parts can be vectorized.
+    final boolean canVectorize = filterCanVectorize && holderSuppliers.getFirst().get().canVectorize();
 
     return new CursorHolder()
     {

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

@@ -175,9 +175,14 @@ private CursorHolder makeClusteredCursorHolder(CursorBuildSpec spec, OnHeapClust
     final ClusteringColumnSelectorFactory wrapperFactory = new ClusteringColumnSelectorFactory(
         ClusteringColumnSelectorFactory.UNINITIALIZED_DELEGATE,
         clusteringColumns,
-        clusteringValuesByGroup.get(0)
+        clusteringValuesByGroup.getFirst()
+    );
+    final ConcatenatingCursor cursor = new ConcatenatingCursor(
+        holderSuppliers,
+        clusteringValuesByGroup,
+        wrapperFactory,
+        plan.virtualColumnRemap()
     );
-    final ConcatenatingCursor cursor = new ConcatenatingCursor(holderSuppliers, clusteringValuesByGroup, wrapperFactory);
 
     return new CursorHolder()
     {

processing/src/main/java/org/apache/druid/segment/projections/ClusterGroupQueryPlan.java

@@ -21,11 +21,15 @@
 
 import org.apache.druid.query.filter.Filter;
 import org.apache.druid.segment.CursorBuildSpec;
+import org.apache.druid.segment.VirtualColumn;
+import org.apache.druid.segment.VirtualColumns;
 import org.apache.druid.segment.filter.FalseFilter;
 import org.apache.druid.segment.filter.TrueFilter;
 
 import javax.annotation.Nullable;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Function;
 
 /**
@@ -39,14 +43,17 @@ public final class ClusterGroupQueryPlan
 {
   private final List<TableClusterGroupSpec> survivingGroups;
   private final Function<TableClusterGroupSpec, Filter> rewriter;
+  private final Map<String, String> virtualColumnRemap;
 
   ClusterGroupQueryPlan(
       List<TableClusterGroupSpec> survivingGroups,
-      Function<TableClusterGroupSpec, Filter> rewriter
+      Function<TableClusterGroupSpec, Filter> rewriter,
+      Map<String, String> virtualColumnRemap
   )
   {
     this.survivingGroups = survivingGroups;
     this.rewriter = rewriter;
+    this.virtualColumnRemap = virtualColumnRemap;
   }
 
   /**
@@ -58,6 +65,16 @@ public List<TableClusterGroupSpec> survivingGroups()
     return survivingGroups;
   }
 
+  /**
+   * Query-level mapping of {@code queryVirtualColumnOutputName -> materializedColumnName} for query virtual columns
+   * that are equivalent to a clustering column or a non-clustering materialized column in the clustered base table.
+   * Empty when no query virtual column has a materialized equivalent (or there are no query virtual columns).
+   */
+  public Map<String, String> virtualColumnRemap()
+  {
+    return virtualColumnRemap;
+  }
+
   /**
    * Per-group rewrite of the query's filter against {@code group}'s constant clustering tuple: clustering-column
    * leaves collapse to {@link TrueFilter} / {@link FalseFilter} and fold through AND / OR / NOT, so the rewritten
@@ -76,21 +93,51 @@ public Filter rewriteFor(TableClusterGroupSpec group)
 
   /**
    * Rebuild {@code spec} for {@code group}'s per-group cursor by swapping in this plan's per-group filter rewrite
-   * (see {@link #rewriteFor}). Returns {@code spec} unchanged when there is no filter, or when the rewrite is
-   * identical to the original (no clustering leaves folded), so the common no-op case allocates nothing. Shared by
-   * the {@link org.apache.druid.segment.QueryableIndexCursorFactory} (historical) and
-   * {@link org.apache.druid.segment.incremental.IncrementalIndexCursorFactory} (realtime) clustered dispatch.
+   * (see {@link #rewriteFor}) and, when {@link #virtualColumnRemap()} is non-empty, substituting any query virtual
+   * columns that are equivalent to a materialized column.
+   * <p>
+   * When there is a remap, the matched query virtual columns (the remap keys) are dropped from the spec's virtual
+   * columns, and the per-group filter's required columns are rewritten via the remap so that non-clustering-VC filter
+   * leaves point at the materialized physical column (clustering-VC leaves were already folded to TRUE / FALSE by the
+   * per-group rewrite walk, so they won't reference the dropped virtual columns). The grouping / select / aggregator
+   * references are served instead by the {@link org.apache.druid.segment.RemapColumnSelectorFactory} the cursor
+   * factory wraps around the {@link ClusteringColumnSelectorFactory}.
    */
   public CursorBuildSpec rebuildCursorBuildSpec(CursorBuildSpec spec, TableClusterGroupSpec group)
   {
-    if (spec.getFilter() == null) {
-      return spec;
+    if (virtualColumnRemap.isEmpty()) {
+      if (spec.getFilter() == null) {
+        return spec;
+      }
+      final Filter rewritten = rewriteFor(group);
+      if (rewritten == spec.getFilter()) {
+        return spec;
+      }
+      return CursorBuildSpec.builder(spec).setFilter(rewritten).build();
     }
-    final Filter rewritten = rewriteFor(group);
-    if (rewritten == spec.getFilter()) {
-      return spec;
+
+    // Drop the remapped query virtual columns from the per-group spec; the materialized columns they map to are
+    // either the per-group constant clustering column or a per-group physical column, both served directly by the
+    // ClusteringColumnSelectorFactory (the cursor factory additionally wraps it with a RemapColumnSelectorFactory so
+    // the original query VC names resolve to the materialized columns).
+    final List<VirtualColumn> prunedVcs = new ArrayList<>();
+    for (VirtualColumn vc : spec.getVirtualColumns().getVirtualColumns()) {
+      if (!virtualColumnRemap.containsKey(vc.getOutputName())) {
+        prunedVcs.add(vc);
+      }
     }
-    return CursorBuildSpec.builder(spec).setFilter(rewritten).build();
+
+    // Per-group filter rewrite first (folds clustering leaves), then remap any surviving non-clustering-VC leaves to
+    // the materialized physical column.
+    Filter rewritten = spec.getFilter() == null ? null : rewriteFor(group);
+    if (rewritten != null) {
+      rewritten = rewritten.rewriteRequiredColumns(virtualColumnRemap);
+    }
+
+    return CursorBuildSpec.builder(spec)
+                          .setFilter(rewritten)
+                          .setVirtualColumns(VirtualColumns.create(prunedVcs))
+                          .build();
   }
 }
 

processing/src/main/java/org/apache/druid/segment/projections/Projections.java

@@ -56,8 +56,11 @@
 import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.Deque;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
@@ -652,9 +655,9 @@ public static ClusterGroupQueryPlan planClusterGroupQuery(
   {
     final Filter queryFilter = cursorBuildSpec.getFilter();
     final VirtualColumns queryVcs = cursorBuildSpec.getVirtualColumns();
-    if (groups.isEmpty() || queryFilter == null) {
-      // No filter (or no groups): every group survives, per-group rewrite is a no-op (null filter).
-      return new ClusterGroupQueryPlan(groups, group -> null);
+    if (groups.isEmpty()) {
+      // No groups: nothing to plan, nothing to remap.
+      return new ClusterGroupQueryPlan(groups, group -> null, Map.of());
     }
 
     // Every spec in the list shares one summary by construction (set once in the schema constructor), so
@@ -663,6 +666,14 @@ public static ClusterGroupQueryPlan planClusterGroupQuery(
     final RowSignature clusteringColumns = summary.getClusteringColumns();
     final VirtualColumns groupVcs = summary.getVirtualColumns();
 
+    final Map<String, String> virtualColumnRemap = buildClusterVirtualColumnRemap(queryVcs, groupVcs);
+
+    if (queryFilter == null) {
+      // No filter: every group survives, per-group filter rewrite is a no-op (null filter), but the VC remap (if any)
+      // still drives grouping / select substitution at cursor build time.
+      return new ClusterGroupQueryPlan(groups, group -> null, virtualColumnRemap);
+    }
+
     // Single walk per group: produces the rewritten filter, and a top-level FalseFilter means the group prunes.
     // Cache the rewrite for every group (including pruned ones, where it's FalseFilter) so rewriteFor doesn't
     // re-walk for either the cursor factory or callers that want to inspect a pruned group's outcome directly.
@@ -681,7 +692,74 @@ public static ClusterGroupQueryPlan planClusterGroupQuery(
         kept.add(group);
       }
     }
-    return new ClusterGroupQueryPlan(kept, rewriteCache::get);
+    return new ClusterGroupQueryPlan(kept, rewriteCache::get, virtualColumnRemap);
+  }
+
+  /**
+   * Build a query-level remap of {@code queryVirtualColumnOutputName -> materializedColumnName} for each query virtual
+   * column that has an equivalent materialized column in the clustered base table (a clustering column produced by a
+   * group virtual column, or a non-clustering materialized virtual-column output).
+   * <p>
+   * A substituted (dropped) query virtual column is read from its materialized column and never recomputes, so it
+   * imposes no requirement on its own inputs. A query virtual column must therefore be kept (recomputed, not
+   * substituted) only when it is transitively required by a kept query virtual column (because query VCs are computed
+   * in the per-group cursor, below the concat-level remap, so a dropped input a kept VC still references would
+   * incorrectly resolve to null.)
+   */
+  private static Map<String, String> buildClusterVirtualColumnRemap(VirtualColumns queryVcs, VirtualColumns groupVcs)
+  {
+    final VirtualColumn[] all = queryVcs.getVirtualColumns();
+    if (all.length == 0) {
+      return Map.of();
+    }
+    // Candidate substitutions: query VCs that have a differently-named equivalent materialized column.
+    final Map<String, String> candidates = new HashMap<>();
+    final Map<String, VirtualColumn> byName = new HashMap<>();
+    for (VirtualColumn vc : all) {
+      final String outputName = vc.getOutputName();
+      byName.put(outputName, vc);
+      final VirtualColumns.Node queryNode = queryVcs.getNode(outputName);
+      if (queryNode == null) {
+        continue;
+      }
+      final VirtualColumn equivalent = groupVcs.findEquivalent(queryNode);
+      if (equivalent != null && !outputName.equals(equivalent.getOutputName())) {
+        candidates.put(outputName, equivalent.getOutputName());
+      }
+    }
+    if (candidates.isEmpty()) {
+      return Map.of();
+    }
+
+    // Transitive "must keep" closure: seed with the non-substituted query VCs (those recompute), then follow
+    // requiredColumns() to every query VC they depend on.
+    final Set<String> keep = new HashSet<>();
+    final Deque<String> toVisit = new ArrayDeque<>();
+    for (VirtualColumn vc : all) {
+      if (!candidates.containsKey(vc.getOutputName())) {
+        toVisit.add(vc.getOutputName());
+      }
+    }
+    while (!toVisit.isEmpty()) {
+      final VirtualColumn vc = byName.get(toVisit.poll());
+      if (vc == null) {
+        continue;
+      }
+      for (String dep : vc.requiredColumns()) {
+        // Only query VCs matter here; a dep that a kept VC recomputes from must itself be kept (and its deps)
+        if (byName.containsKey(dep) && keep.add(dep)) {
+          toVisit.add(dep);
+        }
+      }
+    }
+
+    final Map<String, String> remap = new HashMap<>();
+    for (Map.Entry<String, String> e : candidates.entrySet()) {
+      if (!keep.contains(e.getKey())) {
+        remap.put(e.getKey(), e.getValue());
+      }
+    }
+    return remap;
   }
 
   /**