Liquid-clustering boundary sampling [databricks]

[sdrp713]

Fixes #15903

Description

Problem

GPU range-boundary collection currently samples the full input row. For wide tables, this means decoding and retaining hundreds of payload columns even though boundary selection only needs the range key dependencies. This can cause excessive GPU memory pressure, spill, and semaphore wait.

Change

Build a narrow GPU plan for GpuRangePartitioner.createRangeBounds containing only:

• Range key dependencies.
• Columns required by deterministic filters.
• Required partition columns.

The sampling algorithm, weighting, ordering, and boundary selection are unchanged. The subsequent range shuffle still processes the complete rows.

Pruning is allowlist based. Unsupported or nondeterministic plans use the original full width GPU path. The optimization can also be disabled with:

spark.rapids.sql.rangePartitioning.sampleKeysOnly=false

Different valid samples may produce different boundaries and therefore different skew or spill, but do not change result rows.

This applies to plans containing a GPU range exchange, such as:

GpuColumnarExchange gpurangepartitioning(<rpKey1> ASC NULLS FIRST, <rpKey2> ASC NULLS FIRST, 58), REPARTITION_BY_NUM

It does not affect paths that reuse range bounds already computed by Delta through ZOrderRules.partExprRule. The benefit approaches zero for narrow inputs or when range key dependencies include most columns.

Performance Results

The optimization was evaluated using a Delta Lake liquid-clustering workload consisting of:

1. A plain append with optimized writes and automatic compaction disabled.
2. An OPTIMIZE FULL operation on a table clustered by a column.

The measured physical plan contained this GPU range exchange:

GpuColumnarExchange gpurangepartitioning(<rpKey1> ASC NULLS FIRST, <rpKey2> ASC NULLS FIRST, 58), REPARTITION_BY_NUM

The two rpKey expressions are generated ordering keys, and the exchange creates 58 range partitions.

The input contained:

• 5,887,163 rows
• 585 columns
• 30 input tasks
• Approximately 58.32 GiB read by the original full-row boundary scan

The following configuration was used for both runs:

• Spark 3.5.3
• Two g4dn.8xlarge GPU executors
• 4 executor cores
• 60 GiB executor memory
• spark.sql.files.maxPartitionBytes=2g
• 512 MiB target GPU batch size
• RAPIDS shuffle manager
• Optimized writes disabled
• Automatic compaction disabled

The following is a single-run comparison of the GPU range-boundary sampling stage executed by GpuRangePartitioner.createRangeBounds. These are not end-to-end results or medians.

Range-boundary sampling metrics	Full-row sampling	Key-only sampling	Improvement
Wall-clock duration	363.478 s	3.399 s	106.9× faster
Aggregate executor runtime	2,752.704 s	23.303 s	118.1× lower
Input rows	5,887,163	5,887,163	Unchanged
Input read	58.32 GiB	4.10 MiB	Payload columns removed
Input batches	938	30	96.8% fewer
GPU decode time	272.929 s	0.331 s	824× lower
Host spill	83.78 GiB	0	Eliminated
Disk spill	15.48 GiB	0	Eliminated
Maximum task duration	134.815 s	1.111 s	121.3× lower

Checklists

Documentation

• Updated for new or modified user-facing features or behaviors
• No user-facing change

Testing

• Added or modified tests to cover new code paths
• Covered by existing tests
  (Please provide the names of the existing tests in the PR description.)
• Not required

Performance

• Tests ran and results are added in the PR description
• Issue filed with a link in the PR description
• Not required

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge with no outstanding actionable findings.

Summary

This PR reduces GPU range-boundary sampling width by constructing an auxiliary plan containing only range-key dependencies, deterministic filter dependencies, and required partition columns.

• Falls back to the existing full-width GPU sampling path for unsupported or nondeterministic plans.
• Adds an internal configuration switch to disable key-only sampling.
• Uses the narrow plan to calculate boundaries while retaining complete rows for the subsequent shuffle.
• Adds coverage for computed keys, partition-only keys, disabled optimization, unsupported plans, and nondeterministic keys.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[GPU range exchange] --> B{Key-only sampling enabled?}
    B -- No --> F[Sample original full-width GPU RDD]
    B -- Yes --> C{Boundary plan safely prunable?}
    C -- No --> F
    C -- Yes --> D[Build narrow GPU scan and projection]
    D --> E[Collect range boundaries]
    F --> E
    E --> G[Create range partitioner]
    G --> H[Shuffle complete input rows]

Loading

_{Reviews (5) · Last reviewed commit: "Test successful range boundary pruning p..."}

[amahussein]

The design is the conservative one. prune is an allowlist ending in case _ => None, so an unknown node falls back rather than opting in, and the boundary pass replaces the sampling job createRangeBounds already ran over the full input rather than adding a third read. Worth saying plainly that a bad boundary costs skew and spill here, not wrong results.

Nothing in CI can detect a wrong partition assignment, which is exactly what a boundary bug produces; details inline. And Fixes #15444 does not hold: that reproducer has no clustering key, and optimized write here partitions round-robin or hash, never by range, so getPartitioner's range case is not on its path. Please drop it or say which part you address.

other questions related to the description

• Which stage did the 107x come from, and was it one run per arm or a median of several? The table is headed "Stage metrics", so it reads as the boundary stage rather than end-to-end. Both numbers are interesting; they are different claims.
• Which plan produced it? The plugin's OSS z-order and liquid-clustering acceleration (ZOrderRules.partExprRule) reuses the rangeBounds Delta's CPU RangePartitioner already computed and never builds a GpuRangePartitioning exchange, so that path is untouched by this change. A one-line plan fragment showing GpuColumnarExchange gpurangepartitioning(...) would settle which path the measurement covers, and would also tell the next reader which workloads benefit.
• What does a narrow table do? 585 columns with a one-column clustering key is the best case. When the key covers most columns, build still succeeds and the boundary pass re-reads at nearly full width, adding a plan node and a per-batch projection. I see no mechanism for a meaningful regression, but the description should say where the benefit goes to zero.

[amahussein]

Non-blocking: unconditional for every GPU range exchange the pruner accepts, with no off switch, so a field regression needs a patched release to mitigate.

Suggest an internal boolean defaulting to true, shaped like spark.rapids.sql.shuffledHashJoin.optimizeShuffle. It also lets a test select between the two paths, which the suite comment needs.

[sdrp713]

Fixed. I added the default on internal setting spark.rapids.sql.rangePartitioning.sampleKeysOnly. Setting it to false restores the original full width GPU sampling path, and the disabled path is covered by a correctness test.

[sdrp713]

• Which stage did the 107x come from, and was it one run per arm or a median of several? The table is headed "Stage metrics", so it reads as the boundary stage rather than end-to-end. Both numbers are interesting; they are different claims.

The 107× result is boundary-collection time, comparing one baseline run with one optimized run. It is not a median or an end-to-end result. I updated the description to make that explicit.

[sdrp713]

• Which plan produced it? The plugin's OSS z-order and liquid-clustering acceleration (ZOrderRules.partExprRule) reuses the rangeBounds Delta's CPU RangePartitioner already computed and never builds a GpuRangePartitioning exchange, so that path is untouched by this change. A one-line plan fragment showing GpuColumnarExchange gpurangepartitioning(...) would settle which path the measurement covers, and would also tell the next reader which workloads benefit.

The measured plan contained GpuColumnarExchange gpurangepartitioning( ASC, ASC, 58), so it exercised GpuRangePartitioner.createRangeBounds. I added this fragment and clarified that paths reusing Delta’s CPU-computed bounds are not affected.

[sdrp713]

• What does a narrow table do? 585 columns with a one-column clustering key is the best case. When the key covers most columns, build still succeeds and the boundary pass re-reads at nearly full width, adding a plan node and a per-batch projection. I see no mechanism for a meaningful regression, but the description should say where the benefit goes to zero.

Agreed. I clarified that the benefit approaches zero when the input is narrow or the ordering keys require most columns. I also added a default-on internal kill switch so the original full-width path can be restored if necessary.

[amahussein]

Thanks @sdrp713
It is recommended to avoid rebasing for PRs with ongoing reviews because it moves the commits. Instead, use merge.

[amahussein]

Thanks, this addresses everything from the last round.

The boundary test now asserts the range invariant itself, which is the property a mis-mapped projection or a wrongly bound sorter would break, and the fiddly parts are right: empty partitions dropped before comparison, and <= rather than <. The conf is .internal() and reads through the same idiom the file already uses, so no generated docs are owed.

One blocking item, and it is a title edit: please add [databricks]. The diff touches no ...db/ path, so DBR premerge will not run on its own, and this is the only part of the support matrix with no evidence behind it. It is reachable there, since the DBR scan metas convert to the same shared GpuFileSourceScanExec, and the exchange's DBR subclasses inherit the new subqueries override.

Premerge still has not run on this head.

[amahussein]

• Which plan produced it? The plugin's OSS z-order and liquid-clustering acceleration (ZOrderRules.partExprRule) reuses the rangeBounds Delta's CPU RangePartitioner already computed and never builds a GpuRangePartitioning exchange, so that path is untouched by this change. A one-line plan fragment showing GpuColumnarExchange gpurangepartitioning(...) would settle which path the measurement covers, and would also tell the next reader which workloads benefit.

The measured plan contained GpuColumnarExchange gpurangepartitioning( ASC, ASC, 58), so it exercised GpuRangePartitioner.createRangeBounds. I added this fragment and clarified that paths reusing Delta’s CPU-computed bounds are not affected.

That settles which path the numbers cover, thanks.

One inconsistency: the fragment is two keys and 58 ranges, but the Performance section says 1,000 ranges and a table clustered by a column. With the table declared as one run per arm, the plan is what makes it checkable, so worth saying which run it came from.

[gerashegalov]

Could we add coverage for the two successful pruning paths that are currently untested?

• A deterministic computed range key through GpuProjectExec, checking that the scan retains all key dependencies and that the resulting partitions satisfy the range-ordering invariant.
• A partition-column-only range key, checking the empty data schema, pruned partition schema, row preservation, and the same ordering invariant.

[gerashegalov]

build

[sdrp713]

build

[sdrp713]

builds are failing due to unrelated cudf dependency