From f7de9b6d9e6054b5804f6571f8765a305b3af505 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Thu, 2 Jul 2026 11:21:09 -0700 Subject: [PATCH 1/2] feat: remap query virtual columns to clustered base table columns when equivalent --- .../druid/segment/ConcatenatingCursor.java | 10 +- .../segment/QueryableIndexCursorFactory.java | 23 +- .../IncrementalIndexCursorFactory.java | 9 +- .../projections/ClusterGroupQueryPlan.java | 69 ++++- .../segment/projections/Projections.java | 86 ++++++- .../vector/ConcatenatingVectorCursor.java | 10 +- .../segment/ConcatenatingCursorTest.java | 25 +- ...ryableIndexCursorFactoryClusteredTest.java | 235 ++++++++++++++++++ ...mentalIndexCursorFactoryClusteredTest.java | 119 +++++++++ .../ProjectionsPlanClusterGroupQueryTest.java | 177 +++++++++++++ .../vector/ConcatenatingVectorCursorTest.java | 28 ++- 11 files changed, 747 insertions(+), 44 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/segment/ConcatenatingCursor.java b/processing/src/main/java/org/apache/druid/segment/ConcatenatingCursor.java index 1208c3df2021..d23069c3137a 100644 --- a/processing/src/main/java/org/apache/druid/segment/ConcatenatingCursor.java +++ b/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> holderSuppliers; private final List 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> holderSuppliers, List clusteringValuesByGroup, - ClusteringColumnSelectorFactory wrapperFactory + ClusteringColumnSelectorFactory wrapperFactory, + Map 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 diff --git a/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorFactory.java b/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorFactory.java index 1803554dc56e..de4e57ee6af7 100644 --- a/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorFactory.java +++ b/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() { diff --git a/processing/src/main/java/org/apache/druid/segment/incremental/IncrementalIndexCursorFactory.java b/processing/src/main/java/org/apache/druid/segment/incremental/IncrementalIndexCursorFactory.java index 07cebffa2c96..3c3c04970781 100644 --- a/processing/src/main/java/org/apache/druid/segment/incremental/IncrementalIndexCursorFactory.java +++ b/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() { diff --git a/processing/src/main/java/org/apache/druid/segment/projections/ClusterGroupQueryPlan.java b/processing/src/main/java/org/apache/druid/segment/projections/ClusterGroupQueryPlan.java index 6f87c1984df6..629b3b408385 100644 --- a/processing/src/main/java/org/apache/druid/segment/projections/ClusterGroupQueryPlan.java +++ b/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 survivingGroups; private final Function rewriter; + private final Map virtualColumnRemap; ClusterGroupQueryPlan( List survivingGroups, - Function rewriter + Function rewriter, + Map virtualColumnRemap ) { this.survivingGroups = survivingGroups; this.rewriter = rewriter; + this.virtualColumnRemap = virtualColumnRemap; } /** @@ -58,6 +65,16 @@ public List 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 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. + *

+ * 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 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(); } } diff --git a/processing/src/main/java/org/apache/druid/segment/projections/Projections.java b/processing/src/main/java/org/apache/druid/segment/projections/Projections.java index 54c214f4b37d..26ff0e034c11 100644 --- a/processing/src/main/java/org/apache/druid/segment/projections/Projections.java +++ b/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 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). + *

+ * 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 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 candidates = new HashMap<>(); + final Map 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 keep = new HashSet<>(); + final Deque 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 remap = new HashMap<>(); + for (Map.Entry e : candidates.entrySet()) { + if (!keep.contains(e.getKey())) { + remap.put(e.getKey(), e.getValue()); + } + } + return remap; } /** diff --git a/processing/src/main/java/org/apache/druid/segment/vector/ConcatenatingVectorCursor.java b/processing/src/main/java/org/apache/druid/segment/vector/ConcatenatingVectorCursor.java index c2cfe44efd7d..73328d4d21b3 100644 --- a/processing/src/main/java/org/apache/druid/segment/vector/ConcatenatingVectorCursor.java +++ b/processing/src/main/java/org/apache/druid/segment/vector/ConcatenatingVectorCursor.java @@ -25,6 +25,7 @@ import org.apache.druid.segment.projections.ClusteringVectorColumnSelectorFactory; import java.util.List; +import java.util.Map; /** * Vector-cursor counterpart of {@link org.apache.druid.segment.ConcatenatingCursor}. Walks a sequence of per-group @@ -45,6 +46,7 @@ public final class ConcatenatingVectorCursor implements VectorCursor private final List> holderSuppliers; private final List clusteringValuesByGroup; private final ClusteringVectorColumnSelectorFactory wrapperFactory; + private final VectorColumnSelectorFactory exposedFactory; private int currentIdx; private VectorCursor currentCursor; @@ -53,7 +55,8 @@ public final class ConcatenatingVectorCursor implements VectorCursor public ConcatenatingVectorCursor( List> holderSuppliers, List clusteringValuesByGroup, - ClusteringVectorColumnSelectorFactory wrapperFactory + ClusteringVectorColumnSelectorFactory wrapperFactory, + Map virtualColumnRemap ) { if (holderSuppliers.size() != clusteringValuesByGroup.size()) { @@ -69,6 +72,9 @@ public ConcatenatingVectorCursor( this.holderSuppliers = holderSuppliers; this.clusteringValuesByGroup = clusteringValuesByGroup; this.wrapperFactory = wrapperFactory; + this.exposedFactory = virtualColumnRemap.isEmpty() + ? wrapperFactory + : new RemapVectorColumnSelectorFactory(wrapperFactory, virtualColumnRemap); this.currentIdx = -1; } @@ -106,7 +112,7 @@ private void advanceToNextNonEmptyGroup() public VectorColumnSelectorFactory getColumnSelectorFactory() { initializeIfNeeded(); - return wrapperFactory; + return exposedFactory; } @Override diff --git a/processing/src/test/java/org/apache/druid/segment/ConcatenatingCursorTest.java b/processing/src/test/java/org/apache/druid/segment/ConcatenatingCursorTest.java index 36d2f8c547fc..29a6cc54c260 100644 --- a/processing/src/test/java/org/apache/druid/segment/ConcatenatingCursorTest.java +++ b/processing/src/test/java/org/apache/druid/segment/ConcatenatingCursorTest.java @@ -34,6 +34,7 @@ import javax.annotation.Nullable; import java.util.ArrayList; import java.util.List; +import java.util.Map; class ConcatenatingCursorTest { @@ -56,7 +57,8 @@ void testWalksTwoNonEmptyGroupsBackToBack() ConcatenatingCursor c = new ConcatenatingCursor( List.of(holderSupplier(a), holderSupplier(b)), List.of(new Object[]{"acme"}, new Object[]{"globex"}), - wrapper + wrapper, + Map.of() ); ColumnValueSelector tenant = c.getColumnSelectorFactory().makeColumnValueSelector("tenant"); @@ -99,7 +101,8 @@ void testSkipsLeadingEmptyGroup() ConcatenatingCursor c = new ConcatenatingCursor( List.of(holderSupplier(empty), holderSupplier(full)), List.of(new Object[]{"a"}, new Object[]{"b"}), - wrapper + wrapper, + Map.of() ); ColumnValueSelector tenant = c.getColumnSelectorFactory().makeColumnValueSelector("tenant"); @@ -127,7 +130,8 @@ void testSkipsTrailingEmptyGroup() ConcatenatingCursor c = new ConcatenatingCursor( List.of(holderSupplier(full), holderSupplier(empty)), List.of(new Object[]{"a"}, new Object[]{"b"}), - wrapper + wrapper, + Map.of() ); ColumnValueSelector tenant = c.getColumnSelectorFactory().makeColumnValueSelector("tenant"); @@ -152,7 +156,8 @@ void testAllEmptyGroups() ConcatenatingCursor c = new ConcatenatingCursor( List.of(holderSupplier(e1), holderSupplier(e2)), List.of(new Object[]{"a"}, new Object[]{"b"}), - wrapper + wrapper, + Map.of() ); Assertions.assertTrue(c.isDone()); @@ -172,7 +177,8 @@ void testSingleGroupDegenerateCase() ConcatenatingCursor c = new ConcatenatingCursor( List.of(holderSupplier(only)), List.of(new Object[]{"a"}), - wrapper + wrapper, + Map.of() ); ColumnValueSelector tenant = c.getColumnSelectorFactory().makeColumnValueSelector("tenant"); @@ -202,7 +208,8 @@ void testCloseClosesAllOpenedHolders() ConcatenatingCursor c = new ConcatenatingCursor( List.of(holderSupplier(a), holderSupplier(b)), List.of(new Object[]{"x"}, new Object[]{"y"}), - wrapper + wrapper, + Map.of() ); // Walk through, opens both holders. @@ -236,7 +243,8 @@ void testGroupsAreOpenedLazilyOnTransitionNotEagerly() ConcatenatingCursor c = new ConcatenatingCursor( suppliers, List.of(new Object[]{"a"}, new Object[]{"b"}), - wrapper + wrapper, + Map.of() ); // Init opens first group only. Second is still untouched. @@ -265,7 +273,8 @@ void testResetReIteratesWithoutClosingHolders() ConcatenatingCursor c = new ConcatenatingCursor( List.of(holderSupplier(a), holderSupplier(b)), List.of(new Object[]{"acme"}, new Object[]{"globex"}), - wrapper + wrapper, + Map.of() ); ColumnValueSelector tenant = c.getColumnSelectorFactory().makeColumnValueSelector("tenant"); diff --git a/processing/src/test/java/org/apache/druid/segment/QueryableIndexCursorFactoryClusteredTest.java b/processing/src/test/java/org/apache/druid/segment/QueryableIndexCursorFactoryClusteredTest.java index 8f58fac90443..910fb5974c98 100644 --- a/processing/src/test/java/org/apache/druid/segment/QueryableIndexCursorFactoryClusteredTest.java +++ b/processing/src/test/java/org/apache/druid/segment/QueryableIndexCursorFactoryClusteredTest.java @@ -32,6 +32,7 @@ import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.granularity.Granularities; import org.apache.druid.java.util.common.io.Closer; +import org.apache.druid.math.expr.ExpressionProcessing; import org.apache.druid.query.DruidProcessingConfig; import org.apache.druid.query.Druids; import org.apache.druid.query.OrderBy; @@ -39,6 +40,7 @@ import org.apache.druid.query.Result; import org.apache.druid.query.aggregation.CountAggregatorFactory; import org.apache.druid.query.dimension.DefaultDimensionSpec; +import org.apache.druid.query.expression.TestExprMacroTable; import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.query.filter.EqualityFilter; import org.apache.druid.query.filter.Filter; @@ -62,6 +64,9 @@ import org.apache.druid.segment.filter.OrFilter; import org.apache.druid.segment.incremental.IncrementalIndexSchema; import org.apache.druid.segment.index.BitmapColumnIndex; +import org.apache.druid.segment.vector.VectorCursor; +import org.apache.druid.segment.vector.VectorObjectSelector; +import org.apache.druid.segment.virtual.ExpressionVirtualColumn; import org.joda.time.Interval; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; @@ -106,6 +111,9 @@ class QueryableIndexCursorFactoryClusteredTest @BeforeAll static void setUpEngines() { + // Some tests build segments with ExpressionVirtualColumn clustering / materialized columns, which require the + // expression processing module to be initialized. + ExpressionProcessing.initializeForTests(); engineCloser = Closer.create(); nonBlockingPool = engineCloser.register( new CloseableStupidPool<>("ClusteredCursorFactoryTest-bufferPool", () -> ByteBuffer.allocate(50000)) @@ -182,6 +190,233 @@ private QueryableIndex standardTwoGroup() )); } + /** + * Cluster spec for a segment clustered on {@code tenant_lower := lower(tenant)} (a clustering column produced by a + * group VC; raw {@code tenant} is NOT a stored column) plus a non-clustering materialized + * {@code region_upper := upper(region)} column. Columns: {@code [tenant_lower (clustering), region, region_upper, + * __time]}. + */ + private static final ClusteredValueGroupsBaseTableProjectionSpec VIRTUAL_CLUSTER_SPEC = + ClusteredValueGroupsBaseTableProjectionSpec.builder() + .virtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("tenant_lower", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("region_upper", "upper(region)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .columns( + new StringDimensionSchema("tenant_lower"), + StringDimensionSchema.create("region"), + StringDimensionSchema.create("region_upper"), + new LongDimensionSchema("__time") + ) + .clusteringColumns("tenant_lower") + .build(); + + private QueryableIndex buildVirtualClusteringSegment() + { + final IncrementalIndexSchema schema = + IncrementalIndexSchema.builder() + .withMinTimestamp(INTERVAL.getStartMillis()) + .withTimestampSpec(new TimestampSpec("__time", "auto", null)) + .withQueryGranularity(Granularities.NONE) + .withDimensionsSpec(VIRTUAL_CLUSTER_SPEC.getDimensionsSpec()) + .withRollup(false) + .withClusterSpec(VIRTUAL_CLUSTER_SPEC) + .build(); + return IndexBuilder.create() + .useV10() + .tmpDir(tmpDir) + .schema(schema) + .rows(List.of( + row("Acme", "2025-01-01T00:00:00", "us-east-1"), + row("Acme", "2025-01-01T01:00:00", "us-west-2"), + row("Globex", "2025-01-01T00:30:00", "eu-west-1") + )) + .buildMMappedIndex(INTERVAL); + } + + @Test + void testQueryVcEquivalentToClusteringColumnReadsMaterializedColumnViaAsCursor() + { + // Raw `tenant` is NOT stored; query VC v0 := lower(tenant) is equivalent to the clustering column tenant_lower. + // The scalar selector remap must substitute the materialized tenant_lower clustering constant — recompute would + // be null. Read via asCursor() (non-vector) to exercise the scalar ClusteringColumnSelectorFactory path. + segmentIndex = buildVirtualClusteringSegment(); + final QueryableIndexCursorFactory factory = new QueryableIndexCursorFactory( + segmentIndex, + QueryableIndexTimeBoundaryInspector.create(segmentIndex) + ); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + Assertions.assertEquals(List.of("acme", "acme", "globex"), collectDimension(holder.asCursor(), "v0")); + } + } + + @Test + void testQueryVcEquivalentToNonClusteringMaterializedColumnReadsMaterializedColumnViaAsCursor() + { + // Query VC v1 := upper(region) is equivalent to the non-clustering materialized column region_upper. The remap + // makes makeDimensionSelector("v1") read the per-group physical region_upper column. + segmentIndex = buildVirtualClusteringSegment(); + final QueryableIndexCursorFactory factory = new QueryableIndexCursorFactory( + segmentIndex, + QueryableIndexTimeBoundaryInspector.create(segmentIndex) + ); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v1", "upper(region)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + Assertions.assertEquals( + List.of("US-EAST-1", "US-WEST-2", "EU-WEST-1"), + collectDimension(holder.asCursor(), "v1") + ); + } + } + + @Test + void testQueryVcWithNoEquivalentStillRecomputesViaAsCursor() + { + // No-regression: query VC with no materialized equivalent (lower(region_upper)) is recomputed from the stored + // region_upper column, not remapped. + segmentIndex = buildVirtualClusteringSegment(); + final QueryableIndexCursorFactory factory = new QueryableIndexCursorFactory( + segmentIndex, + QueryableIndexTimeBoundaryInspector.create(segmentIndex) + ); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v2", "lower(region_upper)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + Assertions.assertEquals( + List.of("us-east-1", "us-west-2", "eu-west-1"), + collectDimension(holder.asCursor(), "v2") + ); + } + } + + @Test + void testQueryVcEquivalentToClusteringColumnSingleGroupReadsMaterializedColumn() + { + // Single surviving group (filter on the equivalent VC prunes to one group), exercising the single-group cursor + // holder path's scalar remap wrap. Raw `tenant` is not stored, so the materialized clustering constant is the + // only correct source. + segmentIndex = buildVirtualClusteringSegment(); + final QueryableIndexCursorFactory factory = new QueryableIndexCursorFactory( + segmentIndex, + QueryableIndexTimeBoundaryInspector.create(segmentIndex) + ); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .setFilter(new EqualityFilter("v0", ColumnType.STRING, "globex", null)) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + Assertions.assertEquals(List.of("globex"), collectDimension(holder.asCursor(), "v0")); + } + } + + @Test + void testQueryVcEquivalentToClusteringColumnReadsMaterializedColumnViaVectorCursor() + { + // the query-VC -> materialized-column remap is applied on the vector factory too, so a + // vectorized read of v0 := lower(tenant) resolves to the materialized clustering column tenant_lower (raw tenant + // not stored -> recompute would be null), and the remap no longer forces the scalar path. + segmentIndex = buildVirtualClusteringSegment(); + final QueryableIndexCursorFactory factory = new QueryableIndexCursorFactory( + segmentIndex, + QueryableIndexTimeBoundaryInspector.create(segmentIndex) + ); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + Assertions.assertTrue(holder.canVectorize()); + Assertions.assertEquals(List.of("acme", "acme", "globex"), collectObjectVector(holder.asVectorCursor(), "v0")); + } + } + + @Test + void testQueryVcEquivalentToNonClusteringMaterializedColumnReadsMaterializedColumnViaVectorCursor() + { + // non-clustering: v1 := upper(region) resolves to the materialized region_upper physical column. + segmentIndex = buildVirtualClusteringSegment(); + final QueryableIndexCursorFactory factory = new QueryableIndexCursorFactory( + segmentIndex, + QueryableIndexTimeBoundaryInspector.create(segmentIndex) + ); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v1", "upper(region)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + Assertions.assertTrue(holder.canVectorize()); + Assertions.assertEquals( + List.of("US-EAST-1", "US-WEST-2", "EU-WEST-1"), + collectObjectVector(holder.asVectorCursor(), "v1") + ); + } + } + + @Test + void testQueryVcEquivalentToClusteringColumnSingleGroupVectorCursor() + { + // single surviving group (filter on the equivalent VC prunes to one group): exercises the single-group holder's + // vector remap wrap. + segmentIndex = buildVirtualClusteringSegment(); + final QueryableIndexCursorFactory factory = new QueryableIndexCursorFactory( + segmentIndex, + QueryableIndexTimeBoundaryInspector.create(segmentIndex) + ); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .setFilter(new EqualityFilter("v0", ColumnType.STRING, "globex", null)) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + Assertions.assertTrue(holder.canVectorize()); + Assertions.assertEquals(List.of("globex"), collectObjectVector(holder.asVectorCursor(), "v0")); + } + } + + private static List collectDimension(Cursor cursor, String column) + { + final DimensionSelector sel = + cursor.getColumnSelectorFactory().makeDimensionSelector(DefaultDimensionSpec.of(column)); + final List out = new ArrayList<>(); + while (!cursor.isDone()) { + out.add(sel.getRow().size() == 0 ? null : sel.lookupName(sel.getRow().get(0))); + cursor.advance(); + } + return out; + } + + private static List collectObjectVector(VectorCursor cursor, String column) + { + final VectorObjectSelector sel = cursor.getColumnSelectorFactory().makeObjectSelector(column); + final List out = new ArrayList<>(); + while (!cursor.isDone()) { + final Object[] vector = sel.getObjectVector(); + final int size = cursor.getCurrentVectorSize(); + for (int i = 0; i < size; i++) { + out.add((String) vector[i]); + } + cursor.advance(); + } + return out; + } + @Test void testGetRowSignatureCombinesClusteringFromSummaryAndDataFromGroup() { diff --git a/processing/src/test/java/org/apache/druid/segment/incremental/IncrementalIndexCursorFactoryClusteredTest.java b/processing/src/test/java/org/apache/druid/segment/incremental/IncrementalIndexCursorFactoryClusteredTest.java index 36b663d45051..3ff7d55c61a7 100644 --- a/processing/src/test/java/org/apache/druid/segment/incremental/IncrementalIndexCursorFactoryClusteredTest.java +++ b/processing/src/test/java/org/apache/druid/segment/incremental/IncrementalIndexCursorFactoryClusteredTest.java @@ -176,6 +176,125 @@ void testNonClusteringVirtualColumnDimensionIsMaterialized() } } + /** + * Build an index clustered on {@code tenant_lower := lower(tenant)} (a clustering column produced by a group VC; + * raw {@code tenant} is NOT a stored column) with a non-clustering materialized {@code region_upper := upper(region)} + * column. Columns are {@code [tenant_lower (clustering), region, region_upper, __time]}. + */ + private static OnheapIncrementalIndex virtualClusteringIndex() + { + final ClusteredValueGroupsBaseTableProjectionSpec spec = ClusteredValueGroupsBaseTableProjectionSpec.builder() + .virtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("tenant_lower", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("region_upper", "upper(region)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .columns( + new StringDimensionSchema("tenant_lower"), + new StringDimensionSchema("region"), + new StringDimensionSchema("region_upper"), + new LongDimensionSchema("__time") + ) + .clusteringColumns("tenant_lower") + .build(); + final IncrementalIndexSchema schema = IncrementalIndexSchema.builder() + .withMinTimestamp(T0) + .withTimestampSpec(TIMESTAMP_SPEC) + .withQueryGranularity(Granularities.NONE) + .withDimensionsSpec(spec.getDimensionsSpec()) + .withRollup(false) + .withClusterSpec(spec) + .build(); + final OnheapIncrementalIndex index = (OnheapIncrementalIndex) new OnheapIncrementalIndex.Builder() + .setIndexSchema(schema) + .setMaxRowCount(10_000) + .build(); + index.add(row(T0, "Acme", "us-east-1")); + index.add(row(T0 + 1, "Acme", "us-west-2")); + index.add(row(T0 + 2, "Globex", "eu-west-1")); + return index; + } + + @Test + void testQueryVirtualColumnEquivalentToClusteringColumnReadsMaterializedColumn() + { + // Query VC v0 := lower(tenant) is equivalent to the clustering column tenant_lower (also lower(tenant)). Since + // raw `tenant` is NOT stored, recomputing the expression per-group would yield null; the remap substitutes the + // materialized tenant_lower clustering constant so makeDimensionSelector("v0") returns the per-group value. + try (OnheapIncrementalIndex index = virtualClusteringIndex()) { + final IncrementalIndexCursorFactory factory = new IncrementalIndexCursorFactory(index); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + final Cursor cursor = holder.asCursor(); + final DimensionSelector v0Sel = + cursor.getColumnSelectorFactory().makeDimensionSelector(DefaultDimensionSpec.of("v0")); + final List out = new ArrayList<>(); + while (!cursor.isDone()) { + out.add(v0Sel.getRow().size() == 0 ? null : v0Sel.lookupName(v0Sel.getRow().get(0))); + cursor.advance(); + } + // acme group (tenant_lower=acme) first, then globex; both materialized from the clustering constant. + Assertions.assertEquals(List.of("acme", "acme", "globex"), out); + } + } + } + + @Test + void testQueryVirtualColumnEquivalentToNonClusteringMaterializedColumnReadsMaterializedColumn() + { + // Query VC v1 := upper(region) is equivalent to the non-clustering materialized column region_upper. The remap + // makes makeDimensionSelector("v1") read the per-group physical region_upper column. + try (OnheapIncrementalIndex index = virtualClusteringIndex()) { + final IncrementalIndexCursorFactory factory = new IncrementalIndexCursorFactory(index); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v1", "upper(region)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + final Cursor cursor = holder.asCursor(); + final DimensionSelector v1Sel = + cursor.getColumnSelectorFactory().makeDimensionSelector(DefaultDimensionSpec.of("v1")); + final List out = new ArrayList<>(); + while (!cursor.isDone()) { + out.add(v1Sel.getRow().size() == 0 ? null : v1Sel.lookupName(v1Sel.getRow().get(0))); + cursor.advance(); + } + Assertions.assertEquals(List.of("US-EAST-1", "US-WEST-2", "EU-WEST-1"), out); + } + } + } + + @Test + void testQueryVirtualColumnWithoutEquivalentStillRecomputes() + { + // a query VC with NO materialized equivalent (upper(region_upper) -> a different expression) is not remapped and + // is recomputed normally from the stored region_upper column. + try (OnheapIncrementalIndex index = virtualClusteringIndex()) { + final IncrementalIndexCursorFactory factory = new IncrementalIndexCursorFactory(index); + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(VirtualColumns.create( + new ExpressionVirtualColumn("v2", "lower(region_upper)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + )) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + final Cursor cursor = holder.asCursor(); + final DimensionSelector v2Sel = + cursor.getColumnSelectorFactory().makeDimensionSelector(DefaultDimensionSpec.of("v2")); + final List out = new ArrayList<>(); + while (!cursor.isDone()) { + out.add(v2Sel.getRow().size() == 0 ? null : v2Sel.lookupName(v2Sel.getRow().get(0))); + cursor.advance(); + } + // lower(upper(region)) recomputed from the materialized region_upper column. + Assertions.assertEquals(List.of("us-east-1", "us-west-2", "eu-west-1"), out); + } + } + } + @Test void testRowSignatureExposesClusteringAndNonClusteringColumns() { diff --git a/processing/src/test/java/org/apache/druid/segment/projections/ProjectionsPlanClusterGroupQueryTest.java b/processing/src/test/java/org/apache/druid/segment/projections/ProjectionsPlanClusterGroupQueryTest.java index b013feb333f7..e4937c9c5f0b 100644 --- a/processing/src/test/java/org/apache/druid/segment/projections/ProjectionsPlanClusterGroupQueryTest.java +++ b/processing/src/test/java/org/apache/druid/segment/projections/ProjectionsPlanClusterGroupQueryTest.java @@ -45,6 +45,7 @@ import java.util.Arrays; import java.util.LinkedHashSet; import java.util.List; +import java.util.Map; /** * Coverage for {@link Projections#planClusterGroupQuery}, exercised through both facets of the returned @@ -194,6 +195,60 @@ private static TableClusterGroupSpec virtualClusteringGroup(String loweredTenant return built.specs().get(0); } + /** + * Group clustered on {@code tenant_lower := lower(tenant)} (a clustering column produced by a group VC; raw + * {@code tenant} is NOT stored) with a non-clustering materialized column {@code region_upper := upper(region)}. + * Both equivalences exercise the query-VC -> materialized-column remap. + */ + private static TableClusterGroupSpec materializedVcGroup(String loweredTenant) + { + final RowSignature clustering = RowSignature.builder().add("tenant_lower", ColumnType.STRING).build(); + final ClusterGroupSchemaTestHelpers.Built built = ClusterGroupSchemaTestHelpers.buildClusterGroups( + clustering, + List.of(Arrays.asList(loweredTenant)) + ); + new ClusteredValueGroupsBaseTableSchema( + VirtualColumns.create( + new ExpressionVirtualColumn("tenant_lower", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("region_upper", "upper(region)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ), + List.of("tenant_lower", "region_upper", ColumnHolder.TIME_COLUMN_NAME, "metric"), + List.of(OrderBy.ascending("tenant_lower"), OrderBy.ascending(ColumnHolder.TIME_COLUMN_NAME)), + clustering, + null, + built.dictionaries(), + built.specs() + ); + return built.specs().get(0); + } + + /** + * Like {@link #materializedVcGroup} but also materializes a NESTED virtual column + * {@code tenant_upper_lower := upper(tenant_lower)} (derived from the clustering column), so a query VC chain + * {@code v0 := lower(tenant)} -> {@code v1 := upper(v0)} can be materialized end-to-end. + */ + private static TableClusterGroupSpec nestedMaterializedVcGroup(String loweredTenant) + { + final RowSignature clustering = RowSignature.builder().add("tenant_lower", ColumnType.STRING).build(); + final ClusterGroupSchemaTestHelpers.Built built = ClusterGroupSchemaTestHelpers.buildClusterGroups( + clustering, + List.of(Arrays.asList(loweredTenant)) + ); + new ClusteredValueGroupsBaseTableSchema( + VirtualColumns.create( + new ExpressionVirtualColumn("tenant_lower", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("tenant_upper_lower", "upper(tenant_lower)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ), + List.of("tenant_lower", "tenant_upper_lower", ColumnHolder.TIME_COLUMN_NAME, "metric"), + List.of(OrderBy.ascending("tenant_lower"), OrderBy.ascending(ColumnHolder.TIME_COLUMN_NAME)), + clustering, + null, + built.dictionaries(), + built.specs() + ); + return built.specs().get(0); + } + private static CursorBuildSpec buildSpec(@Nullable Filter filter) { return buildSpec(filter, VirtualColumns.EMPTY); @@ -900,4 +955,126 @@ void testQueryVcShadowingClusteringNameWithoutEquivalenceLeavesLeafUnchanged() Projections.planClusterGroupQuery(List.of(group), buildSpec(f, queryVcs)).rewriteFor(group) ); } + + @Test + void testVirtualColumnRemapMapsQueryVcEquivalentToClusteringColumn() + { + // Query VC v0 := lower(tenant) is equivalent to the group's clustering column tenant_lower (also lower(tenant)). + // The remap should map v0 -> tenant_lower so grouping/select reads the materialized clustering column. + final VirtualColumns queryVcs = VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final ClusterGroupQueryPlan plan = Projections.planClusterGroupQuery( + List.of(materializedVcGroup("acme")), + buildSpec(null, queryVcs) + ); + Assertions.assertEquals(Map.of("v0", "tenant_lower"), plan.virtualColumnRemap()); + } + + @Test + void testVirtualColumnRemapMapsQueryVcEquivalentToNonClusteringMaterializedColumn() + { + // Query VC v1 := upper(region) is equivalent to the group's non-clustering materialized column region_upper. + final VirtualColumns queryVcs = VirtualColumns.create( + new ExpressionVirtualColumn("v1", "upper(region)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final ClusterGroupQueryPlan plan = Projections.planClusterGroupQuery( + List.of(materializedVcGroup("acme")), + buildSpec(null, queryVcs) + ); + Assertions.assertEquals(Map.of("v1", "region_upper"), plan.virtualColumnRemap()); + } + + @Test + void testVirtualColumnRemapEmptyWhenNoEquivalent() + { + // Query VC has no materialized equivalent in the group (no length(region) column / VC) → no remap. + final VirtualColumns queryVcs = VirtualColumns.create( + new ExpressionVirtualColumn("v0", "strlen(region)", ColumnType.LONG, TestExprMacroTable.INSTANCE) + ); + final ClusterGroupQueryPlan plan = Projections.planClusterGroupQuery( + List.of(materializedVcGroup("acme")), + buildSpec(null, queryVcs) + ); + Assertions.assertTrue(plan.virtualColumnRemap().isEmpty()); + } + + @Test + void testVirtualColumnRemapEmptyWhenNoQueryVirtualColumns() + { + final ClusterGroupQueryPlan plan = Projections.planClusterGroupQuery( + List.of(materializedVcGroup("acme")), + buildSpec(null) + ); + Assertions.assertTrue(plan.virtualColumnRemap().isEmpty()); + } + + @Test + void testVirtualColumnRemapBuiltEvenWithFilter() + { + // The remap must be computed even when the query has a filter (grouping/select substitution applies regardless). + final VirtualColumns queryVcs = VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final Filter f = new EqualityFilter("v0", ColumnType.STRING, "acme", null); + final ClusterGroupQueryPlan plan = Projections.planClusterGroupQuery( + List.of(materializedVcGroup("acme")), + buildSpec(f, queryVcs) + ); + Assertions.assertEquals(Map.of("v0", "tenant_lower"), plan.virtualColumnRemap()); + } + + @Test + void testVirtualColumnRemapRespectsDependencyGuard() + { + // v0 := lower(tenant) (equivalent to clustering tenant_lower) AND v1 := concat(v0, 'x') which DEPENDS on v0. + // v0 must NOT be remapped/dropped because v1 still needs it as an input; v1 itself has no equivalent → no remap. + final VirtualColumns queryVcs = VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("v1", "concat(v0, 'x')", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final ClusterGroupQueryPlan plan = Projections.planClusterGroupQuery( + List.of(materializedVcGroup("acme")), + buildSpec(null, queryVcs) + ); + Assertions.assertTrue(plan.virtualColumnRemap().isEmpty()); + } + + @Test + void testVirtualColumnRemapSubstitutesFullyMaterializedDependencyChain() + { + // v0 := lower(tenant) (equivalent to clustering tenant_lower) AND v1 := upper(v0) which DEPENDS on v0 but is + // ALSO materialized (equivalent to tenant_upper_lower := upper(tenant_lower)). Because v1 is substituted (dropped, + // read from its materialized column) it never recomputes, so it does not force v0 to be kept -> BOTH substitute. + final VirtualColumns queryVcs = VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("v1", "upper(v0)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final ClusterGroupQueryPlan plan = Projections.planClusterGroupQuery( + List.of(nestedMaterializedVcGroup("acme")), + buildSpec(null, queryVcs) + ); + Assertions.assertEquals( + Map.of("v0", "tenant_lower", "v1", "tenant_upper_lower"), + plan.virtualColumnRemap() + ); + } + + @Test + void testVirtualColumnRemapKeepsMaterializedChainNeededByNonMaterializedDependent() + { + // Transitive keep: v0 := lower(tenant) and v1 := upper(v0) are both materialized, but v2 := strlen(v1) is NOT + // (no strlen column). v2 is kept and recomputes -> needs v1 -> v1 kept -> v1 recomputes -> needs v0 -> v0 kept. + // So nothing substitutes, even though v0/v1's only direct dependents are themselves materialized. + final VirtualColumns queryVcs = VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("v1", "upper(v0)", ColumnType.STRING, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("v2", "strlen(v1)", ColumnType.LONG, TestExprMacroTable.INSTANCE) + ); + final ClusterGroupQueryPlan plan = Projections.planClusterGroupQuery( + List.of(nestedMaterializedVcGroup("acme")), + buildSpec(null, queryVcs) + ); + Assertions.assertTrue(plan.virtualColumnRemap().isEmpty()); + } } diff --git a/processing/src/test/java/org/apache/druid/segment/vector/ConcatenatingVectorCursorTest.java b/processing/src/test/java/org/apache/druid/segment/vector/ConcatenatingVectorCursorTest.java index 5a9ebf1ba2c2..80fa379a849c 100644 --- a/processing/src/test/java/org/apache/druid/segment/vector/ConcatenatingVectorCursorTest.java +++ b/processing/src/test/java/org/apache/druid/segment/vector/ConcatenatingVectorCursorTest.java @@ -36,6 +36,7 @@ import javax.annotation.Nullable; import java.util.ArrayList; import java.util.List; +import java.util.Map; class ConcatenatingVectorCursorTest { @@ -58,7 +59,8 @@ void testWalksTwoNonEmptyGroupsBackToBack() ConcatenatingVectorCursor c = new ConcatenatingVectorCursor( List.of(holderSupplier(a), holderSupplier(b)), List.of(new Object[]{"acme"}, new Object[]{"globex"}), - wrapper + wrapper, + Map.of() ); VectorObjectSelector tenant = c.getColumnSelectorFactory().makeObjectSelector("tenant"); @@ -105,7 +107,8 @@ void testSkipsLeadingEmptyGroup() ConcatenatingVectorCursor c = new ConcatenatingVectorCursor( List.of(holderSupplier(empty), holderSupplier(full)), List.of(new Object[]{"a"}, new Object[]{"b"}), - wrapper + wrapper, + Map.of() ); VectorObjectSelector tenant = c.getColumnSelectorFactory().makeObjectSelector("tenant"); @@ -133,7 +136,8 @@ void testSkipsTrailingEmptyGroup() ConcatenatingVectorCursor c = new ConcatenatingVectorCursor( List.of(holderSupplier(full), holderSupplier(empty)), List.of(new Object[]{"a"}, new Object[]{"b"}), - wrapper + wrapper, + Map.of() ); VectorObjectSelector tenant = c.getColumnSelectorFactory().makeObjectSelector("tenant"); @@ -158,7 +162,8 @@ void testAllEmptyGroups() ConcatenatingVectorCursor c = new ConcatenatingVectorCursor( List.of(holderSupplier(e1), holderSupplier(e2)), List.of(new Object[]{"a"}, new Object[]{"b"}), - wrapper + wrapper, + Map.of() ); Assertions.assertTrue(c.isDone()); @@ -179,7 +184,8 @@ void testSingleGroupDegenerateCase() ConcatenatingVectorCursor c = new ConcatenatingVectorCursor( List.of(holderSupplier(only)), List.of(new Object[]{"a"}), - wrapper + wrapper, + Map.of() ); VectorObjectSelector tenant = c.getColumnSelectorFactory().makeObjectSelector("tenant"); @@ -210,7 +216,8 @@ void testPartialVectorAtGroupBoundary() ConcatenatingVectorCursor c = new ConcatenatingVectorCursor( List.of(holderSupplier(a), holderSupplier(b)), List.of(new Object[]{"acme"}, new Object[]{"globex"}), - wrapper + wrapper, + Map.of() ); VectorObjectSelector tenant = c.getColumnSelectorFactory().makeObjectSelector("tenant"); @@ -245,7 +252,8 @@ void testCloserClosesAllOpenedHolders() ConcatenatingVectorCursor c = new ConcatenatingVectorCursor( List.of(holderSupplier(a), holderSupplier(b)), List.of(new Object[]{"x"}, new Object[]{"y"}), - wrapper + wrapper, + Map.of() ); c.getColumnSelectorFactory(); @@ -279,7 +287,8 @@ void testGroupsAreOpenedLazilyOnTransitionNotEagerly() ConcatenatingVectorCursor c = new ConcatenatingVectorCursor( suppliers, List.of(new Object[]{"a"}, new Object[]{"b"}), - wrapper + wrapper, + Map.of() ); c.getColumnSelectorFactory(); @@ -307,7 +316,8 @@ void testMaxVectorSizeIsConfiguredValueAcrossAllStates() ConcatenatingVectorCursor c = new ConcatenatingVectorCursor( List.of(holderSupplier(a)), List.of(new Object[]{"acme"}), - wrapper + wrapper, + Map.of() ); // Pre-init. From b76b28a73eec86ba83b4014d92d31e93f42a76c1 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Thu, 2 Jul 2026 12:10:20 -0700 Subject: [PATCH 2/2] more test --- ...ryableIndexCursorFactoryClusteredTest.java | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/processing/src/test/java/org/apache/druid/segment/QueryableIndexCursorFactoryClusteredTest.java b/processing/src/test/java/org/apache/druid/segment/QueryableIndexCursorFactoryClusteredTest.java index 910fb5974c98..b86c18bf6878 100644 --- a/processing/src/test/java/org/apache/druid/segment/QueryableIndexCursorFactoryClusteredTest.java +++ b/processing/src/test/java/org/apache/druid/segment/QueryableIndexCursorFactoryClusteredTest.java @@ -323,6 +323,40 @@ void testQueryVcEquivalentToClusteringColumnSingleGroupReadsMaterializedColumn() } } + @Test + void testEquivalentVcFilterPrunesToSingleGroupNotConcatenated() + { + // A filter on the ORIGINAL expression used to build the clustering column (lower(tenant) = 'acme', planned as a + // query VC) must prune to the single matching cluster group and take the single-group cursor path -- NOT survive + // every group and post-filter through the ConcatenatingCursor. + segmentIndex = buildVirtualClusteringSegment(); + final QueryableIndexCursorFactory factory = new QueryableIndexCursorFactory( + segmentIndex, + QueryableIndexTimeBoundaryInspector.create(segmentIndex) + ); + final VirtualColumns queryVcs = VirtualColumns.create( + new ExpressionVirtualColumn("v0", "lower(tenant)", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + + // Control: unfiltered -> both groups survive -> multi-group ConcatenatingCursor. + try (CursorHolder holder = + factory.makeCursorHolder(CursorBuildSpec.builder().setVirtualColumns(queryVcs).build())) { + Assertions.assertInstanceOf(ConcatenatingCursor.class, holder.asCursor()); + } + + // Filtered on the equivalent VC -> the clustering leaf resolves to tenant_lower and folds, pruning to the single + // acme group -> single-group cursor (not a ConcatenatingCursor), and no residual filter on the missing raw column. + final CursorBuildSpec buildSpec = CursorBuildSpec.builder() + .setVirtualColumns(queryVcs) + .setFilter(new EqualityFilter("v0", ColumnType.STRING, "acme", null)) + .build(); + try (CursorHolder holder = factory.makeCursorHolder(buildSpec)) { + final Cursor cursor = holder.asCursor(); + Assertions.assertFalse(cursor instanceof ConcatenatingCursor); + Assertions.assertEquals(List.of("us-east-1", "us-west-2"), collectDimension(cursor, "region")); + } + } + @Test void testQueryVcEquivalentToClusteringColumnReadsMaterializedColumnViaVectorCursor() {