[BUG] GPU range-partition boundary sampling crashes on rows-only (zero-column) batches

[amahussein]

Describe the bug

GPU range-partition boundary sampling crashes with ArrayIndexOutOfBoundsException when the input batch has no columns.

Every branch of SamplingUtils.selectWithoutReplacementFrom ends in GpuColumnVector.from(cb), which builds a cuDF Table, and a Table needs at least one column. GpuRangePartitioner.sketch feeds it the exchange input unchecked, so a column-less batch reaches it and the job aborts.

Steps/Code to reproduce bug

Range-partition on a key the optimizer folds to a constant, then run an aggregate so column pruning empties the scan.

from pyspark.sql import functions as F

path = "/tmp/zero-col-repro"
spark.range(0, 4000).select(
    F.col("id").alias("key"),
    F.sha2(F.col("id").cast("string"), 256).alias("payload"),
).repartition(4).write.mode("overwrite").parquet(path)

df = (spark.read.parquet(path)
      .select(F.lit(1).cast("long").alias("k"), F.col("payload"))
      .repartitionByRange(4, F.col("k")))

df.agg(F.count(F.lit(1))).collect()     # raises
df.collect()                            # succeeds: payload keeps the batch non-empty

The two calls differ only in column pruning. Executed plan for the failing one:

GpuHashAggregate (keys=[], functions=[gpucount(1, false)])
+- GpuShuffleCoalesce
   +- GpuColumnarExchange gpusinglepartitioning$(), ENSURE_REQUIREMENTS
      +- GpuHashAggregate (keys=[], functions=[partial_gpucount(1, false)])
         +- GpuShuffleCoalesce
            +- GpuColumnarExchange gpurangepartitioning(1 ASC NULLS FIRST, 4), REPARTITION_BY_NUM
               +- GpuFileGpuScan parquet [] ... ReadSchema: struct<>

1 ASC NULLS FIRST is the folded key; ReadSchema: struct<> is the column-less scan. Failure:

java.lang.ArrayIndexOutOfBoundsException: Index 0 out of bounds for length 0
  at ai.rapids.cudf.Table.<init>(Table.java:47)
  at com.nvidia.spark.rapids.GpuColumnVector.from(GpuColumnVector.java:656)
  at com.nvidia.spark.rapids.SamplingUtils$.selectWithoutReplacementFrom(SamplingUtils.scala:48)
  at com.nvidia.spark.rapids.SamplingUtils$.reservoirSampleAndCount(SamplingUtils.scala:175)
  at com.nvidia.spark.rapids.GpuRangePartitioner$.$anonfun$sketch$1(GpuRangePartitioner.scala:52)

Expected behavior

The query returns 4000. A batch with rows and no columns carries no key data, so sampling should count its rows and select nothing from it rather than building a table.

Environment details (please complete the following information)

• Environment location: Standalone, single host, one NVIDIA RTX A5000 (24 GB)
• Spark 3.5.3, buildver=353, plugin built from 4f46d0086c297170bc84845e25f6812772adafe4
• spark.plugins=com.nvidia.spark.SQLPlugin, spark.rapids.sql.enabled=true, spark.sql.adaptive.enabled=false, --master local[4]

Additional context

Same root cause as #15953, but a different caller. Both end at GpuColumnVector.from(ColumnarBatch), which builds a cuDF Table without checking the column count.

#15954, which fixes #15953, guards the aggregate call site rather than GpuColumnVector.from itself: it adds a dataTypes.nonEmpty branch in GpuAggregateIterator.concatenateBatchesWithRetry and carries the row count forward. That leaves this site unfixed, and it is not the same fix either. The aggregate can emit an equivalent zero-column batch; the sampler cannot, because a rows-only batch carries no key data, so the guard here is to count the rows and contribute no sample candidates.

Filing separately rather than asking #15954 to absorb it: different operator, different fix, and that PR already has its benchmarks and Databricks CI run against its current scope.

Found while reviewing #15925, which does not cause it. That PR's key-only boundary plan cannot reach this: GpuRangeBoundaryPlan.build returns None when the range ordering has no attribute references, which is exactly this case, and the failing plan contains no boundary node.

The trigger is narrow, so priority is low, but the failure mode is a hard abort rather than a fallback. A constant range key arises from generated SQL, from a view whose clustering column is a literal, or from any repartitionByRange whose key the optimizer folds.

[revans2]

We should fix this, but I would argue that it is a bug in Spark.

A batch with rows and no columns carries no key data, so sampling should count its rows and select nothing from it rather than building a table.

Why does Spark want us to sample NOTHING? So that we can partition on NOTHING? The range partitioning should have disappeared, and likely would have gone away if we had done an order by instead of an explicit range partition.

[amahussein]

We should fix this, but I would argue that it is a bug in Spark.

A batch with rows and no columns carries no key data, so sampling should count its rows and select nothing from it rather than building a table.

Why does Spark want us to sample NOTHING? So that we can partition on NOTHING? The range partitioning should have disappeared, and likely would have gone away if we had done an order by instead of an explicit range partition.

Agreed on the plan: a constant range key is pointless work either way.

But the crash is not specific to range partitioning. df.filter(expr("(rand() > 0.5) and (ceil(rand()) > 0)")) over a pruned scan hits the same Table.<init> through GpuAnd, with no repartition. I measured ten such sites, eight previously unknown. Umbrella issue to follow #15991