diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatch.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatch.scala index 84fd5865e76e7..4ad7cf649338e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatch.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatch.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.columnar.SimpleMetricsCachedBatch * The batch contains: * - `numRows`: Number of rows in this batch * - `arrowData`: One encapsulated Arrow RecordBatch message (with optional compression) - * - `stats`: Per-column statistics for partition pruning (upperBound, lowerBound, nullCount, etc.) + * - `stats`: Per-column statistics for partition pruning (lowerBound, upperBound, nullCount, etc.) * * This format enables: * - Zero-copy columnar reads when output is ColumnarBatch with ArrowColumnVector @@ -42,8 +42,9 @@ import org.apache.spark.sql.columnar.SimpleMetricsCachedBatch * * @param numRows Number of rows in this cached batch * @param arrowData One encapsulated Arrow RecordBatch message - * @param stats Per-column statistics as InternalRow (5 fields per column: - * upperBound, lowerBound, nullCount, rowCount, sizeInBytes) + * @param stats Per-column statistics as InternalRow (5 fields per column, in the + * `ColumnStats.collectedStatistics` order: + * lowerBound, upperBound, nullCount, count, sizeInBytes) */ case class ArrowCachedBatch( numRows: Int, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala index 9529e85698537..58b385071e775 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala @@ -166,6 +166,19 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { cacheAttributes: Seq[Attribute], selectedAttributes: Seq[Attribute], conf: SQLConf): RDD[InternalRow] = { + if (selectedAttributes.isEmpty) { + // Empty projection (e.g. a count aggregate over the cached relation): every cached batch + // already records its row count, so emit that many empty rows without touching the Arrow + // payload at all -- deserializing and decompressing it would be pure waste. The emitted + // row is a single reused 0-field UnsafeRow, matching the reuse contract of the regular + // path. + return input.mapPartitionsInternal { batchIterator => + val rowWriter = new UnsafeRowWriter(0) + rowWriter.reset() + val emptyRow = rowWriter.getRow + batchIterator.flatMap(batch => Iterator.fill(batch.numRows)(emptyRow)) + } + } val cacheSchema = DataTypeUtils.fromAttributes(cacheAttributes) val selectedSchema = DataTypeUtils.fromAttributes(selectedAttributes) val timeZoneId = conf.sessionLocalTimeZone diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala index 6576bc80b600e..f4119a306ce13 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala @@ -26,10 +26,11 @@ import org.apache.arrow.vector.{ TimeNanoVector, TimeStampMicroTZVector, TimeStampMicroVector, TinyIntVector, VarBinaryVector, VarCharVector, VectorSchemaRoot, VectorUnloader} -import org.apache.spark.{SparkConf, SparkUnsupportedOperationException} +import org.apache.spark.{SparkConf, SparkException, SparkUnsupportedOperationException} import org.apache.spark.sql.{QueryTest, Row} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{AttributeReference, GenericInternalRow} +import org.apache.spark.sql.columnar.CachedBatch import org.apache.spark.sql.execution.arrow.ArrowWriter import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.test.{ExamplePoint, ExamplePointUDT, SharedSparkSession} @@ -2401,6 +2402,51 @@ class ArrowCachedBatchSerializerSuite extends QueryTest with SharedSparkSession // If the produced root was not closed, this throws "Memory was leaked by query". alloc.close() } + + test("empty projection emits row counts without deserializing the Arrow payload") { + // A count-style read selects no columns, and the row count is already recorded on the cached + // batch, so the Arrow payload must not be deserialized at all. Prove it by handing the reader + // a batch whose payload is garbage: the empty projection succeeds purely from numRows, while + // a projection that actually needs the payload fails on the same batch. + val serializer = new ArrowCachedBatchSerializer + val attrs = Seq(AttributeReference("i", IntegerType)()) + val garbage = ArrowCachedBatch( + numRows = 5, + arrowData = Array[Byte](0x13, 0x37, 0x00, -1), + stats = InternalRow(null, null, 0, 5, 4L)) + val input = spark.sparkContext.parallelize(Seq[CachedBatch](garbage), 1) + val conf = spark.sessionState.conf + + val emptyProjection = + serializer.convertCachedBatchToInternalRow(input, attrs, Seq.empty, conf) + assert(emptyProjection.map(_.numFields).collect() === Array(0, 0, 0, 0, 0)) + + val fullProjection = + serializer.convertCachedBatchToInternalRow(input, attrs, attrs, conf) + intercept[SparkException] { + fullProjection.collect() + } + } + + test("count aggregate over the cached relation with the row-based reader") { + // End-to-end coverage of the empty-projection read: with the vectorized reader disabled the + // scan produces rows via convertCachedBatchToInternalRow, and a count aggregate selects no + // columns. Small Arrow batches make the count span many cached batches. + withSQLConf( + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "false", + SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "10") { + val df = (1 to 257).toList.toDF("i") + df.cache() + try { + assert(df.count() === 257) + checkAnswer(df.selectExpr("count(*)"), Seq(Row(257))) + // A projecting query over the same cached data still reads the payload correctly. + checkAnswer(df.selectExpr("sum(i)"), Seq(Row((1 to 257).sum.toLong))) + } finally { + df.unpersist() + } + } + } } /**