[Enhancement] Pre-split base and rollup for INSERT-from-table into a range table

Sample-Based Tablet Pre-Split previously skipped INSERT ... SELECT ... FROM a
single OLAP table when the range-distribution target carried a rollup, because
the sampler could only project the base sort key and only by target column name.

This makes INSERT-from-table split every visible index (base plus each rollup):

- InsertFromTableScanContext carries the full target->source column-name map
  instead of a base-only source-column list, and the data-tier executor maps
  whatever sort key it is handed (request.getSortKey() for the base; each rollup
  spec's sort key for the secondary projection) through that map.
- The shared secondary-index projection becomes an overridable executor seam
  whose default projects by target column name (FILES / Broker byte-identical);
  INSERT-from-table overrides it to remap to source names.
- TablePreSplitSource builds one secondary index spec per visible rollup.
- The two guards that skipped multi-index INSERT-from-table are removed.

Single-index INSERT-from-table is byte-identical (the base sort key maps to the
same source names). FILES / Broker and the ALTER-side pre-split path are
untouched. Any unmapped sort-key column fails the sample so the load proceeds
without pre-split. Data tier only (an OLAP source has no Parquet/ORC footer).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: xiangguangyxg <xiangguangyxg@gmail.com>

Author: xiangguangyxg

fe/fe-core/src/main/java/com/starrocks/alter/reshard/presplit/AbstractSqlSampleSubqueryExecutor.java

@@ -45,7 +45,9 @@
  *
  * <p>Subclasses implement {@link #resolveSampleSpec} to supply the FROM clause
  * SQL, an optional WHERE predicate, the total input byte count, the compute
- * resource, and the projected column identifier lists. Everything else —
+ * resource, and the base sort-key / partition projection identifier lists; a
+ * subclass may also override {@link #secondaryIndexProjectionIdents} to control
+ * how each secondary index's sort key is projected. Everything else —
  * sampling-rate math, SQL synthesis, BE invocation, JSON row decode — is shared
  * here.
  */
@@ -168,7 +170,7 @@ public final SampleExecution execute(SampleRequest request) throws StarRocksExce
         List<SecondaryIndexSpec> secondaryIndexSortKeys = request.getSecondaryIndexSortKeys();
         String sampleSql = buildSampleSql(
                 spec.fromClauseSql(), spec.whereClauseSqlOrNull(),
-                spec.sortKeyProjectionIdents(), secondaryProjectionIdents(secondaryIndexSortKeys),
+                spec.sortKeyProjectionIdents(), secondaryProjectionIdents(request),
                 spec.partitionProjectionIdents(),
                 samplingRate, rowLimit, request.getSeed());
         List<TResultBatch> resultBatches = runSampleQuery(
@@ -180,20 +182,33 @@ public final SampleExecution execute(SampleRequest request) throws StarRocksExce
     }
 
     /**
-     * Flattens each {@link SecondaryIndexSpec}'s sort-key column names into
-     * backtick-quoted SQL identifiers, in spec order then per-spec column
-     * order. Empty when {@code secondaryIndexSortKeys} is empty, so the
-     * projection is byte-identical to the pre-multi-index SQL.
+     * Flattens each {@link SecondaryIndexSpec}'s projection idents, in spec order
+     * then per-spec column order. Empty when the request carries no secondary
+     * indexes, so the projection is byte-identical to the pre-multi-index SQL.
      */
-    private static List<String> secondaryProjectionIdents(List<SecondaryIndexSpec> secondaryIndexSortKeys) {
+    private List<String> secondaryProjectionIdents(SampleRequest request) throws StarRocksException {
+        List<SecondaryIndexSpec> secondaryIndexSortKeys = request.getSecondaryIndexSortKeys();
         if (secondaryIndexSortKeys.isEmpty()) {
             return List.of();
         }
         List<String> idents = new ArrayList<>();
         for (SecondaryIndexSpec secondaryIndexSpec : secondaryIndexSortKeys) {
-            for (Column column : secondaryIndexSpec.sortKey()) {
-                idents.add(SqlUtils.getIdentSql(column.getName()));
-            }
+            idents.addAll(secondaryIndexProjectionIdents(request, secondaryIndexSpec));
+        }
+        return idents;
+    }
+
+    /**
+     * Projection idents to SELECT for one secondary index's sort key. The default
+     * projects each column by its own name, correct when the source is name-aligned
+     * to the target. A subclass that remaps which source columns back a target column
+     * overrides this to project by the source-table column name instead.
+     */
+    protected List<String> secondaryIndexProjectionIdents(SampleRequest request, SecondaryIndexSpec spec)
+            throws StarRocksException {
+        List<String> idents = new ArrayList<>(spec.sortKey().size());
+        for (Column column : spec.sortKey()) {
+            idents.add(SqlUtils.getIdentSql(column.getName()));
         }
         return idents;
     }

fe/fe-core/src/main/java/com/starrocks/alter/reshard/presplit/InsertFromTableSampleSubqueryExecutor.java

@@ -15,8 +15,13 @@
 package com.starrocks.alter.reshard.presplit;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.starrocks.catalog.Column;
 import com.starrocks.common.StarRocksException;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
 /**
  * Data-tier {@link SampleSubqueryExecutor} for the INSERT-from-OLAP-table path.
  * Synthesizes a {@code SELECT <sort_key_cols>[, <partition_source_cols>] FROM
@@ -45,20 +50,58 @@ final class InsertFromTableSampleSubqueryExecutor extends AbstractSqlSampleSubqu
 
     @Override
     protected SampleSpec resolveSampleSpec(SampleRequest request) throws StarRocksException {
-        ScanContext scanContext = request.getScanContext();
-        if (!(scanContext instanceof InsertFromTableScanContext context)) {
-            throw new StarRocksException(ERROR_PREFIX + "received a "
-                    + scanContext.getClass().getSimpleName()
-                    + " — wire only the INSERT-from-table load kind here");
-        }
+        InsertFromTableScanContext context = contextOf(request);
         return new SampleSpec(
                 context.sourceFromSql(),
                 context.wherePredicateSql(),
                 Math.max(0L, context.sourceTable().getDataSize()),
                 context.computeResource(),
-                identsOf(context.sortKeySourceColumnNames()),
+                identsOf(mapToSource(request.getSortKey(), context.targetToSourceColumnNames())),
                 identsOf(context.partitionSourceColumnNames()),
                 request.getSortKey(),
                 request.getPartitionSourceColumns());
     }
+
+    /**
+     * Projects one rollup's sort key by remapping each target column to its source-table
+     * name. The base sort key is projected the same way in {@link #resolveSampleSpec}; both
+     * go through {@link #mapToSource}, so a rollup whose sort key reorders or subsets the base
+     * is projected correctly regardless of source column naming.
+     */
+    @Override
+    protected List<String> secondaryIndexProjectionIdents(SampleRequest request, SecondaryIndexSpec spec)
+            throws StarRocksException {
+        InsertFromTableScanContext context = contextOf(request);
+        return identsOf(mapToSource(spec.sortKey(), context.targetToSourceColumnNames()));
+    }
+
+    private static InsertFromTableScanContext contextOf(SampleRequest request) throws StarRocksException {
+        ScanContext scanContext = request.getScanContext();
+        if (!(scanContext instanceof InsertFromTableScanContext context)) {
+            throw new StarRocksException(ERROR_PREFIX + "received a "
+                    + scanContext.getClass().getSimpleName()
+                    + " -- wire only the INSERT-from-table load kind here");
+        }
+        return context;
+    }
+
+    /**
+     * Remaps target sort-key columns to their source-table column names via the scan
+     * context's target-&gt;source map. Throws (-&gt; the sample fails -&gt; the load proceeds
+     * without pre-split) if any column is unmapped, so a boundary is never computed against
+     * the wrong source column.
+     */
+    private static List<String> mapToSource(
+            List<Column> targetColumns, Map<String, String> targetToSourceColumnNames) throws StarRocksException {
+        List<String> sourceNames = new ArrayList<>(targetColumns.size());
+        for (Column column : targetColumns) {
+            String sourceName = targetToSourceColumnNames.get(column.getName().toLowerCase());
+            if (sourceName == null) {
+                throw new StarRocksException(ERROR_PREFIX
+                        + "no source column mapped for target sort-key column " + column.getName());
+            }
+            sourceNames.add(sourceName);
+        }
+        return sourceNames;
+    }
 }

fe/fe-core/src/main/java/com/starrocks/alter/reshard/presplit/InsertFromTableScanContext.java

@@ -18,29 +18,31 @@
 import com.starrocks.warehouse.cngroup.ComputeResource;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
 /**
  * {@link ScanContext} concrete for the INSERT-from-OLAP-table integration.
  * Carries the source {@link OlapTable} reference so the sampler can obtain its
- * data-size estimate and the pre-quoted FROM clause SQL, plus the source-column
- * name lists that map the target sort key and partition columns back to their
- * source equivalents. The optional WHERE predicate SQL is threaded through
- * verbatim from the INSERT-SELECT statement so the sample covers only the rows
- * the load will actually write.
+ * data-size estimate and the pre-quoted FROM clause SQL, the full target->source
+ * column-name map the sampler uses to project any index's sort key (base or rollup)
+ * by its source-table column name, and the source-column names of the target
+ * partition columns. The optional WHERE predicate SQL is threaded through verbatim
+ * from the INSERT-SELECT statement so the sample covers only the rows the load will
+ * actually write.
  */
 public record InsertFromTableScanContext(
         OlapTable sourceTable,
         String sourceFromSql,                       // "`db`.`tbl` `alias`" or "`db`.`tbl`"
-        List<String> sortKeySourceColumnNames,
+        Map<String, String> targetToSourceColumnNames,   // lower-cased target name -> source column name
         List<String> partitionSourceColumnNames,
         String wherePredicateSql,                   // nullable
         ComputeResource computeResource) implements ScanContext {
 
     public InsertFromTableScanContext {
         Objects.requireNonNull(sourceTable, "sourceTable");
         Objects.requireNonNull(sourceFromSql, "sourceFromSql");
-        Objects.requireNonNull(sortKeySourceColumnNames, "sortKeySourceColumnNames");
+        Objects.requireNonNull(targetToSourceColumnNames, "targetToSourceColumnNames");
         Objects.requireNonNull(partitionSourceColumnNames, "partitionSourceColumnNames");
         Objects.requireNonNull(computeResource, "computeResource");
     }

fe/fe-core/src/main/java/com/starrocks/alter/reshard/presplit/InsertSelectSourceColumns.java

@@ -37,23 +37,28 @@
  * (caller then silently skips pre-split).
  */
 final class InsertSelectSourceColumns {
-    private final List<String> sortKeySourceColumnNames;
     private final List<String> partitionSourceColumnNames;
+    private final Map<String, String> targetToSourceColumnNames;
 
-    private InsertSelectSourceColumns(List<String> sortKeySourceColumnNames,
-                                      List<String> partitionSourceColumnNames) {
-        this.sortKeySourceColumnNames = sortKeySourceColumnNames;
+    private InsertSelectSourceColumns(List<String> partitionSourceColumnNames,
+                                      Map<String, String> targetToSourceColumnNames) {
         this.partitionSourceColumnNames = partitionSourceColumnNames;
-    }
-
-    List<String> sortKeySourceColumnNames() {
-        return sortKeySourceColumnNames;
+        this.targetToSourceColumnNames = targetToSourceColumnNames;
     }
 
     List<String> partitionSourceColumnNames() {
         return partitionSourceColumnNames;
     }
 
+    /**
+     * Full target-&gt;source column-name map (lower-cased target name -&gt; source column name),
+     * total over the target's non-generated base columns. The sampler uses it to project any
+     * index's sort key (base or rollup) by its source-table column name.
+     */
+    Map<String, String> targetToSourceColumnNames() {
+        return targetToSourceColumnNames;
+    }
+
     /**
      * Resolves the source-column names for each sort-key and partition column of the target table.
      *
@@ -148,15 +153,15 @@ static InsertSelectSourceColumns resolve(
             }
         }
 
-        List<String> sortKeyNames = lookup(sortKeyColumns, targetToSource);
-        if (sortKeyNames == null) {
+        // A sort-key column with no source mapping cannot be sampled -> skip pre-split.
+        if (lookup(sortKeyColumns, targetToSource) == null) {
             return null;
         }
         List<String> partitionNames = lookup(partitionColumns, targetToSource);
         if (partitionNames == null) {
             return null;
         }
-        return new InsertSelectSourceColumns(sortKeyNames, partitionNames);
+        return new InsertSelectSourceColumns(partitionNames, Map.copyOf(targetToSource));
     }
 
     private static Set<String> targetNames(List<Column> targetCols) {

fe/fe-core/src/main/java/com/starrocks/alter/reshard/presplit/PreSplitFlow.java

@@ -42,9 +42,7 @@
  *   <li>{@link #dispatch} — partitioned-vs-unpartitioned routing, including the
  *       automatic-partition gate that conservatively skips manually
  *       list/range-partitioned targets (those cannot pre-create partitions from
- *       sampled values), and the INSERT-from-table rollup descope (a rollup's sort
- *       key cannot be remapped to source columns, so a multi-index target skips
- *       before sampling for that load kind only).</li>
+ *       sampled values).</li>
  *   <li>{@link #runSinglePartitionFlow} — resolve the unique partition + base
  *       tablet, build a {@link DefaultPreSplitPipeline}, submit via
  *       {@link TabletPreSplitCoordinator#submitAsynchronously}, then sync-await
@@ -76,7 +74,7 @@ private PreSplitFlow() {
      * requested tablet count; computeResource sizes the active CN count; scanContext carries
      * the source-specific scan inputs; secondaryIndexSpecs names every OTHER visible index
      * (rollup) whose sort key the multi-partition data-tier sampler should project alongside
-     * the base sort key -- empty for single-index targets and for INSERT-from-table (descoped).
+     * the base sort key -- empty for single-index targets.
      */
     record Prepared(ScanContext scanContext, List<Column> sortKeyColumns,
                     List<Column> partitionColumns, long estimatedBytes,
@@ -94,11 +92,6 @@ record Prepared(ScanContext scanContext, List<Column> sortKeyColumns,
 
     static void dispatch(Database database, OlapTable target, Prepared prepared,
                          LoadKind loadKind, BooleanSupplier shouldAbort, ConnectContext context) {
-        if (loadKind == LoadKind.INSERT_FROM_TABLE && target.getVisibleIndexMetas().size() > 1) {
-            // INSERT-from-table cannot remap a divergent rollup sort key to source columns (descoped).
-            PreSplitMetrics.recordEligibilitySkip(SkipReason.HAS_MATERIALIZED_VIEW_OR_ROLLUP);
-            return;
-        }
         if (target.getPartitionInfo().isPartitioned()) {
             // Manually list/range-partitioned targets do not support pre-creating partitions
             // from sampled values; skip conservatively, let the load proceed.
@@ -117,13 +110,6 @@ static void runSinglePartitionFlow(Database database, OlapTable table, Prepared
         if (target == null) {
             return;
         }
-        if (loadKind == LoadKind.INSERT_FROM_TABLE && target.indexTargets().size() > 1) {
-            // Authoritative re-check on the resolved index set: a rollup that became visible after the
-            // dispatch-time descope gate must not be sampled with the base-only INSERT-from-table source
-            // mapping. Skip pre-split (load proceeds) -- INSERT-from-table with a rollup is out of scope.
-            PreSplitMetrics.recordEligibilitySkip(SkipReason.HAS_MATERIALIZED_VIEW_OR_ROLLUP);
-            return;
-        }
         int activeComputeNodeCount = TabletReshardUtils.computeNodeCount(prepared.computeResource());
         DefaultPreSplitPipeline pipeline = DefaultPreSplitPipeline.forLoadKind(
                 target.database(), target.olapTable(), target.indexTargets(), prepared.estimatedBytes(), loadKind,

fe/fe-core/src/main/java/com/starrocks/alter/reshard/presplit/TablePreSplitSource.java

@@ -18,6 +18,7 @@
 import com.starrocks.authorization.PrivilegeType;
 import com.starrocks.catalog.Column;
 import com.starrocks.catalog.Database;
+import com.starrocks.catalog.MaterializedIndexMeta;
 import com.starrocks.catalog.OlapTable;
 import com.starrocks.catalog.Table;
 import com.starrocks.catalog.TableName;
@@ -37,6 +38,7 @@
 import com.starrocks.sql.ast.expression.SlotRef;
 import com.starrocks.sql.common.MetaUtils;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
@@ -109,14 +111,30 @@ public PreSplitFlow.Prepared prepare(InsertStmt insertStmt, SelectRelation selec
         }
         InsertFromTableScanContext scanContext = new InsertFromTableScanContext(
                 resolvedSource.sourceTable(), resolvedSource.sourceFromSql(),
-                mapping.sortKeySourceColumnNames(), mapping.partitionSourceColumnNames(),
+                mapping.targetToSourceColumnNames(), mapping.partitionSourceColumnNames(),
                 wherePredicateSql, context.getCurrentComputeResource());
         long estimatedBytes = Math.max(0L, resolvedSource.sourceTable().getDataSize());
-        // INSERT-from-table cannot remap a rollup's sort key to source columns, so the
-        // dispatch gate already skips multi-index targets for this load kind; secondaryIndexSpecs
-        // stays empty here via the backward-compatible Prepared constructor.
         return new PreSplitFlow.Prepared(scanContext, sortKeyColumns, partitionColumns,
-                estimatedBytes, context.getCurrentComputeResource());
+                estimatedBytes, context.getCurrentComputeResource(), buildSecondaryIndexSpecs(target));
+    }
+
+    /**
+     * Builds one {@link SecondaryIndexSpec} per visible rollup index of {@code target}
+     * (every visible index except the base), each carrying its own TARGET sort-key columns.
+     * The data-tier executor remaps every column to its source-table name via the scan
+     * context's target->source map at sample time, so the specs carry TARGET columns here --
+     * identical in shape to the FILES / Broker builders.
+     */
+    private static List<SecondaryIndexSpec> buildSecondaryIndexSpecs(OlapTable target) {
+        List<SecondaryIndexSpec> specs = new ArrayList<>();
+        for (MaterializedIndexMeta meta : target.getVisibleIndexMetas()) {
+            if (meta.getIndexMetaId() == target.getBaseIndexMetaId()) {
+                continue;
+            }
+            specs.add(new SecondaryIndexSpec(meta.getIndexMetaId(),
+                    MetaUtils.getRangeDistributionColumns(target, meta.getIndexMetaId())));
+        }
+        return specs;
     }
 
     /**