feat: restrict virtual column usage in clustered segments, allow remapping them at query time

processing/src/main/java/org/apache/druid/data/input/impl/ClusteredValueGroupsBaseTableProjectionSpec.java

@@ -33,24 +33,29 @@
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.segment.VirtualColumn;
 import org.apache.druid.segment.VirtualColumns;
+import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ColumnHolder;
+import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.segment.projections.Projections;
 import org.apache.druid.utils.CollectionUtils;
 
 import javax.annotation.Nullable;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
 
 /**
  * {@link BaseTableProjectionSpec} for the clustered-value-groups base table mode: rows are partitioned into per-tuple
  * "cluster groups" keyed by one or more typed clustering columns, optionally derived from {@link #virtualColumns}.
- * Essentially, each group is stored as its own internal sub-segment..
+ * Essentially, each group is stored as its own internal sub-segment. Logically, a clustered base table spec is a
+ * projection of an (imaginary/unstored) base table into the stored clustered table.
  * <p>
- * The operator declares a single ordered {@link #columns} list — the full set of columns in segment order, plus a
+ * The operator declares a single ordered {@link #columns} list, the full set of columns in segment order, plus a
  * {@link #clusteringColumns} list of NAMES designating the leading prefix of {@link #columns} that rows are
  * clustered by. The time position is an explicit positional entry in {@link #columns} named {@code __time}; clustering
  * by the time column is not yet supported, so {@code __time} must be a non-clustering column. A clustered base table
@@ -94,6 +99,7 @@ public ClusteredValueGroupsBaseTableProjectionSpec(
   {
     validate(columns, clusteringColumns);
     this.virtualColumns = virtualColumns == null ? VirtualColumns.EMPTY : virtualColumns;
+    validateVirtualColumns(this.virtualColumns, columns);
     this.columns = Collections.unmodifiableList(new ArrayList<>(columns));
     this.clusteringColumns = Collections.unmodifiableList(new ArrayList<>(clusteringColumns));
     this.clusteringColumnSchemas = this.columns.subList(0, this.clusteringColumns.size());
@@ -311,6 +317,98 @@ private static void validate(List<DimensionSchema> columns, List<String> cluster
     }
   }
 
+  /**
+   * Rules to keep virtual columns definitions reasonable:
+   * <ul>
+   *   <li><b>dot notation</b>: a virtual column that {@link VirtualColumn#usesDotNotation()} is referenced as
+   *   {@code name.subfield} rather than by a fixed output name, so it has no stable identity to materialize or cluster
+   *   on; it is rejected outright.</li>
+   *   <li><b>inputs</b>: every input of a virtual column must be a stored column (declared in {@code columns}) or
+   *   another virtual column in the spec.</li>
+   *   <li><b>outputs</b>: every virtual column must either be materialized (its output declared in {@code columns}) or
+   *   be an intermediary that feeds another virtual column. A virtual column that is neither materializes nothing and
+   *   is used by nothing and dead metadata and so it is rejected.</li>
+   *   <li><b>output type</b>: a materialized virtual column must produce the {@link ColumnType} declared for its column
+   *   in {@code columns}. This is intentionally strict to avoid a virtual column silently filling a column with a
+   *   coerced/mismatched type.</li>
+   * </ul>
+   * The query-granularity carrier ({@link Granularities#GRANULARITY_VIRTUAL_COLUMN_NAME}) is special handled to
+   * capture how __time is computed, so it is exempt from the output rule.
+   */
+  private static void validateVirtualColumns(VirtualColumns virtualColumns, List<DimensionSchema> columns)
+  {
+    final VirtualColumn[] all = virtualColumns.getVirtualColumns();
+    if (all.length == 0) {
+      return;
+    }
+    // Declared column types, doubling as the ColumnInspector used to infer each materialized virtual column's output
+    // type (an expression's output type can depend on its input column types).
+    final RowSignature.Builder signatureBuilder = RowSignature.builder();
+    for (DimensionSchema column : columns) {
+      signatureBuilder.add(column.getName(), column.getColumnType());
+    }
+    final RowSignature rowSignature = signatureBuilder.build();
+    // The output rule below lets a virtual column go unstored when it is exempt: an intermediary that feeds another
+    // virtual column (collected during the input pass), or the metadata-only query-granularity carrier (seeded here).
+    final Set<String> outputExempt = new HashSet<>();
+    outputExempt.add(Granularities.GRANULARITY_VIRTUAL_COLUMN_NAME);
+    // input rule: every input must be a stored column or another virtual column; an input that is a virtual column
+    // makes that virtual column an intermediary, exempt from having to be materialized itself.
+    for (VirtualColumn virtualColumn : all) {
+      if (virtualColumn.usesDotNotation()) {
+        throw InvalidInput.exception(
+            "virtual column [%s] uses dot notation, which is not supported for clusteredValueGroups base table specs",
+            virtualColumn.getOutputName()
+        );
+      }
+      for (String input : virtualColumn.requiredColumns()) {
+        final boolean isStored = rowSignature.contains(input);
+        final boolean isVirtual = virtualColumns.exists(input);
+        if (!isStored && !isVirtual) {
+          throw InvalidInput.exception(
+              "virtual column [%s] reads column [%s], which is neither a stored column nor another virtual column;"
+              + " clustered base table virtual columns must be computable from stored columns (retain [%s] in"
+              + " 'columns', or use a transformSpec for in-place transforms)",
+              virtualColumn.getOutputName(),
+              input,
+              input
+          );
+        }
+        if (isVirtual) {
+          outputExempt.add(input);
+        }
+      }
+    }
+    // output rule: every virtual column must be materialized (stored) or exempt (intermediary / granularity carrier);
+    // a materialized virtual column must additionally produce the declared type of the column it fills.
+    for (VirtualColumn virtualColumn : all) {
+      final String outputName = virtualColumn.getOutputName();
+      final boolean stored = rowSignature.contains(outputName);
+      if (!stored && !outputExempt.contains(outputName)) {
+        throw InvalidInput.exception(
+            "virtual column [%s] is not stored (not declared in 'columns') and does not feed another virtual column;"
+            + " clustered base table virtual columns must materialize a stored column or be an input to one that does",
+            outputName
+        );
+      }
+      if (stored) {
+        final ColumnCapabilities outputCapabilities =
+            virtualColumns.getColumnCapabilitiesWithFallback(rowSignature, outputName);
+        final ColumnType outputType = ColumnType.fromCapabilities(outputCapabilities);
+        final ColumnType declaredType = rowSignature.getColumnType(outputName).orElse(null);
+        if (outputType != null && !outputType.equals(declaredType)) {
+          throw InvalidInput.exception(
+              "virtual column [%s] produces type [%s] but the column it materializes is declared as [%s] in 'columns';"
+              + " a clustered base table virtual column must match the declared type of the column it fills",
+              outputName,
+              outputType,
+              declaredType
+          );
+        }
+      }
+    }
+  }
+
   private static DruidException clusteringPrefixException(
       List<DimensionSchema> columns,
       List<String> clusteringColumns

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/MergingClusterGroupCursor.java

@@ -26,6 +26,7 @@
 import org.apache.druid.segment.projections.MergingColumnSelectorFactory;
 
 import java.util.List;
+import java.util.Map;
 
 /**
  * {@link Cursor} that presents a set of individually {@code __time}-sorted per-cluster-group cursors as a single
@@ -38,7 +39,11 @@
  * <em>all</em> surviving group cursors up front (so unlike the concatenating path it does not benefit from early-exit
  * laziness) and, on each {@link #advance()}, emits the row with the smallest (or largest, when descending)
  * {@code __time} across the groups. Each per-group sub-index already exposes the clustering columns as constants, so
- * the {@link MergingColumnSelectorFactory} simply dispatches every column to the winning group.
+ * the {@link MergingColumnSelectorFactory} simply dispatches every column to the winning group. When the plan carries a
+ * query-virtual-column-to-materialized-column remap, the merging factory is wrapped in a
+ * {@link RemapColumnSelectorFactory} so reads of a remapped query virtual column's output name resolve to its
+ * materialized column (mirroring {@link ConcatenatingCursor}); the per-group specs already dropped those virtual
+ * columns and added their materialized targets as physical columns.
  * <p>
  * An {@link IntHeapPriorityQueue} of group indices, keyed on each group's current {@code __time} via
  * {@link #compareGroups}, drives the merge: the queue head is the winning group, {@link IntHeapPriorityQueue#changed()}
@@ -50,6 +55,7 @@ public final class MergingClusterGroupCursor implements Cursor
 {
   private final List<Supplier<CursorHolder>> holderSuppliers;
   private final boolean descending;
+  private final Map<String, String> virtualColumnRemap;
 
   private boolean initialized;
   // Indexed by group. groupCursors[i] is null if that group's holder produced no cursor.
@@ -62,18 +68,22 @@ public final class MergingClusterGroupCursor implements Cursor
   private int currentGroup;
   // Monotonically increasing output-row id, one per emitted row; the merge emits exactly one row per advance.
   private long outputRowId;
-  private MergingColumnSelectorFactory factory;
+  // The factory exposed to callers: the MergingColumnSelectorFactory, wrapped in a RemapColumnSelectorFactory when
+  // virtualColumnRemap is non-empty.
+  private ColumnSelectorFactory exposedFactory;
 
   public MergingClusterGroupCursor(
       List<Supplier<CursorHolder>> holderSuppliers,
-      boolean descending
+      boolean descending,
+      Map<String, String> virtualColumnRemap
   )
   {
     if (holderSuppliers.isEmpty()) {
       throw DruidException.defensive("MergingClusterGroupCursor requires at least one cluster group");
     }
     this.holderSuppliers = holderSuppliers;
     this.descending = descending;
+    this.virtualColumnRemap = virtualColumnRemap;
   }
 
   private void initializeIfNeeded()
@@ -106,14 +116,18 @@ private void initializeIfNeeded()
       }
     }
     currentGroup = heap.isEmpty() ? 0 : heap.firstInt();
-    factory = new MergingColumnSelectorFactory(groupFactories, () -> currentGroup, () -> outputRowId);
+    final MergingColumnSelectorFactory mergingFactory =
+        new MergingColumnSelectorFactory(groupFactories, () -> currentGroup, () -> outputRowId);
+    exposedFactory = virtualColumnRemap.isEmpty()
+                     ? mergingFactory
+                     : new RemapColumnSelectorFactory(mergingFactory, virtualColumnRemap);
   }
 
   @Override
   public ColumnSelectorFactory getColumnSelectorFactory()
   {
     initializeIfNeeded();
-    return factory;
+    return exposedFactory;
   }
 
   @Override