This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new ca6a8690aae6 [SPARK-58918][SQL] Push column pruning into the Arrow
cache read path
ca6a8690aae6 is described below
commit ca6a8690aae663af547bcb0155e6abfa69332656
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Fri Aug 21 00:31:10 2026 -0700
[SPARK-58918][SQL] Push column pruning into the Arrow cache read path
### What changes were proposed in this pull request?
When reading a projection of an Arrow-cached relation,
`ArrowCachedBatchSerializer`'s read path deserialized **every** cached column
-- decompressing and loading all of them off-heap -- and then discarded the
unselected ones. This change reads only the selected columns out of the cached
bytes.
A new helper `ArrowCachedBatchSerializer.readProjectedRecordBatch` reads
the encapsulated IPC RecordBatch message's metadata (a small flatbuffer that
lists every buffer's offset and length within the body) and copies just the
byte ranges belonging to the selected columns straight out of the in-memory
cached `Array[Byte]`, so the unselected columns are never copied off-heap,
loaded, or decompressed. The selected buffers become windows into a single
off-heap allocation, mirroring how th [...]
It is wired into both read paths -- `convertCachedBatchToColumnarBatch`
(columnar) and `convertCachedBatchToInternalRow` (row). A projection whose
selected attribute is absent from the cache schema (index `-1`) falls back to
reading the full batch.
### Why are the changes needed?
The wasted work is proportional to the pruned columns and dominates
wide-relation scans, especially with compression. An e2e SQL benchmark
(`sum(col0)` over a cached relation, cache materialized outside the timed
region, 1M rows x 50 long columns) shows:
| | before | after |
|---|---|---|
| Arrow cache -- uncompressed | 31 ms | 16 ms (~1.9x) |
| Arrow cache -- zstd level 1 | 326 ms | 22 ms (~15x) |
The compressed case gains most, since the 49 pruned columns are no longer
decompressed.
### Does this PR introduce _any_ user-facing change?
No. The Arrow cache serializer (SPARK-57268) is unreleased, and this is a
read-path performance improvement with identical results.
### How was this patch tested?
- A new `ArrowCachedBatchSerializerSuite` test covers projection shapes
(reordering, single column at each position, complex-after-var-width, duplicate
selection, full projection) on both read paths, and is negative-validated (an
off-by-one in the buffer-span arithmetic fails it).
- Full `ArrowCachedBatchSerializerSuite` and
`ArrowCachedBatchKryoRegistrationSuite` pass (77 tests).
- A new `columnPruningWideTable` case in `ArrowCacheBenchmark` measures the
read path with the cache built outside the timed region; the committed
benchmark result files are regenerated by the benchmark CI job.
### Was this patch authored or co-authored using generative AI tooling?
Yes, this pull request and its description were written by Claude Code.
Closes #58177 from viirya/arrow-cache-projection-pushdown.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
.../columnar/ArrowCachedBatchSerializer.scala | 215 +++++++++++++++++++--
.../execution/benchmark/ArrowCacheBenchmark.scala | 56 ++++++
.../columnar/ArrowCachedBatchSerializerSuite.scala | 40 ++++
3 files changed, 295 insertions(+), 16 deletions(-)
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 7f8ddb177788..baba2c1365c8 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
@@ -23,10 +23,14 @@ import java.nio.channels.Channels
import scala.jdk.CollectionConverters._
import org.apache.arrow.compression.{Lz4CompressionCodec, ZstdCompressionCodec}
-import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot, VectorUnloader}
+import org.apache.arrow.flatbuf.{RecordBatch => FlatBufRecordBatch}
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.{TypeLayout, VectorLoader, VectorSchemaRoot,
VectorUnloader}
import org.apache.arrow.vector.compression.{CompressionCodec,
NoCompressionCodec}
import org.apache.arrow.vector.ipc.{ReadChannel, WriteChannel}
+import org.apache.arrow.vector.ipc.message.{ArrowBodyCompression,
ArrowFieldNode}
import org.apache.arrow.vector.ipc.message.{ArrowRecordBatch,
MessageSerializer}
+import org.apache.arrow.vector.types.pojo.Field
import org.apache.spark.{SparkException, TaskContext}
import org.apache.spark.rdd.RDD
@@ -285,6 +289,135 @@ private object ArrowCachedBatchSerializer {
out.toByteArray
}
+ /**
+ * Number of Arrow buffers a field occupies in a RecordBatch body, including
all of its
+ * descendants, in the depth-first order `VectorLoader` consumes them. The
type's own buffer
+ * count comes from `TypeLayout` (validity + offset/data buffers), then each
child contributes
+ * its whole subtree recursively. Used to map each top-level column to its
run of buffers.
+ */
+ private def fieldBufferCount(field: Field): Int =
+ TypeLayout.getTypeBufferCount(field.getType) +
+ field.getChildren.asScala.map(fieldBufferCount).sum
+
+ /** Number of Arrow field nodes a field occupies (itself plus every
descendant). */
+ private def fieldNodeCount(field: Field): Int =
+ 1 + field.getChildren.asScala.map(fieldNodeCount).sum
+
+ /** Number of variadic buffer counts a field contributes (one per view-type
buffer, recursive). */
+ private def fieldVariadicCount(field: Field): Int = {
+ val own = field.getType match {
+ // View types (Utf8View/BinaryView) carry a variadic-buffer count in the
RecordBatch;
+ // no other type does. The cache never writes view vectors today, but
account for them so
+ // the span arithmetic stays correct if that changes.
+ case _: org.apache.arrow.vector.types.pojo.ArrowType.Utf8View |
+ _: org.apache.arrow.vector.types.pojo.ArrowType.BinaryView => 1
+ case _ => 0
+ }
+ own + field.getChildren.asScala.map(fieldVariadicCount).sum
+ }
+
+ /**
+ * Read an encapsulated IPC RecordBatch message from `data`, materializing
off-heap only the
+ * buffers of the requested top-level columns. This is the
projection-pushdown read path: the
+ * message metadata (a small flatbuffer) lists every buffer's (offset,
length) within the body,
+ * so we copy just the byte ranges belonging to the selected columns
straight out of the
+ * in-memory `data` array, never touching (or allocating off-heap for) the
unselected columns.
+ *
+ * The body is a flat, depth-first sequence of buffers in schema order, so
each top-level column
+ * owns a contiguous run of buffers whose span is `fieldBufferCount`; field
nodes and variadic
+ * counts run in the same order. The selected columns' bytes are copied into
one off-heap buffer
+ * (each buffer 8-byte aligned, matching Arrow's IPC body layout) and the
returned batch's
+ * buffers are windows into it, exactly like the standard reader slices one
body buffer -- so the
+ * batch has a single underlying allocation and no per-buffer bookkeeping.
The returned batch
+ * owns its buffers (the constructor retains each), so the caller closes it
as usual.
+ *
+ * Compression is preserved unchanged: buffer (offset, length) spans cover
the on-body bytes
+ * including any per-buffer uncompressed-length prefix, so the copied
windows are still compressed
+ * as written; `VectorLoader.load` decompresses only the selected ones later.
+ */
+ def readProjectedRecordBatch(
+ data: Array[Byte],
+ schemaFields: Seq[Field],
+ selectedIndices: Array[Int],
+ allocator: BufferAllocator): ArrowRecordBatch = {
+ val in = new ByteArrayInputStream(data)
+ val readChannel = new ReadChannel(Channels.newChannel(in))
+ // Read only the message metadata; the body bytes stay in `data` and are
copied selectively.
+ val metadata = MessageSerializer.readMessage(readChannel)
+ require(metadata != null, "Unexpected end of input reading cached batch
message")
+ val batch =
+ metadata.getMessage.header(new
FlatBufRecordBatch()).asInstanceOf[FlatBufRecordBatch]
+ // serializeBatch writes exactly [encapsulated message][body] with no
end-of-stream marker, so
+ // the body is the tail of `data`: it starts at data.length minus the
declared body length.
+ val bodyStart = data.length - metadata.getMessageBodyLength().toInt
+
+ val compression: ArrowBodyCompression =
+ if (batch.compression() == null)
NoCompressionCodec.DEFAULT_BODY_COMPRESSION
+ else new ArrowBodyCompression(batch.compression().codec(),
batch.compression().method())
+
+ val nodeStarts = schemaFields.scanLeft(0)(_ + fieldNodeCount(_)).toArray
+ val bufferStarts = schemaFields.scanLeft(0)(_ +
fieldBufferCount(_)).toArray
+ val variadicStarts = schemaFields.scanLeft(0)(_ +
fieldVariadicCount(_)).toArray
+ val hasVariadic = batch.variadicBufferCountsLength() > 0
+
+ // Enumerate the selected columns' nodes, buffer indices and variadic
counts, in output order.
+ val selectedNodes = new java.util.ArrayList[ArrowFieldNode]()
+ val selectedBufferIdx = new scala.collection.mutable.ArrayBuffer[Int]()
+ val selectedVariadic = new java.util.ArrayList[java.lang.Long]()
+ selectedIndices.foreach { i =>
+ val field = schemaFields(i)
+ val nStart = nodeStarts(i)
+ (nStart until nStart + fieldNodeCount(field)).foreach { j =>
+ val node = batch.nodes(j)
+ selectedNodes.add(new ArrowFieldNode(node.length(), node.nullCount()))
+ }
+ val bStart = bufferStarts(i)
+ (bStart until bStart +
fieldBufferCount(field)).foreach(selectedBufferIdx += _)
+ if (hasVariadic) {
+ val vStart = variadicStarts(i)
+ (vStart until vStart + fieldVariadicCount(field)).foreach(j =>
+ selectedVariadic.add(batch.variadicBufferCounts(j)))
+ }
+ }
+
+ val layout = selectedBufferIdx.map { j =>
+ val buf = batch.buffers(j)
+ (buf.offset(), buf.length())
+ }
+ val alignedSizes = layout.map { case (_, len) => ((len + 7) / 8) * 8 }
+ val body = allocator.buffer(math.max(alignedSizes.sum, 1))
+ try {
+ val selectedBuffers = new
java.util.ArrayList[org.apache.arrow.memory.ArrowBuf]()
+ var pos = 0L
+ layout.indices.foreach { k =>
+ val (srcOffset, len) = layout(k)
+ if (len > 0) {
+ body.setBytes(pos, data, bodyStart + srcOffset.toInt, len.toInt)
+ }
+ val window = body.slice(pos, len)
+ window.writerIndex(len)
+ selectedBuffers.add(window)
+ pos += alignedSizes(k)
+ }
+ val recordBatch = new ArrowRecordBatch(
+ batch.length().toInt,
+ selectedNodes,
+ selectedBuffers,
+ compression,
+ selectedVariadic,
+ false)
+ // The constructor retained each window (slice() itself does not), so
the batch now holds one
+ // reference per window into `body`. Drop `body`'s original allocation
reference; the batch is
+ // then the sole owner and the caller's recordBatch.close() frees the
single allocation.
+ body.close()
+ recordBatch
+ } catch {
+ case t: Throwable =>
+ body.close()
+ throw t
+ }
+ }
+
/**
* Byte offset of the unscaled low-order word within a 16-byte Arrow
Decimal128 slot, for the
* given native byte order. Arrow Java writes decimal values in native byte
order
@@ -1143,6 +1276,21 @@ private class ArrowCachedBatchToColumnarBatchIterator(
private val arrowSchema = ArrowUtils.toArrowSchema(
cacheSchema, timeZoneId, false, false, losslessInternalTypes = true)
+ // Projection pushdown: the cached batch stores all cache columns, but only
the selected ones
+ // are needed. When every selected column maps to a distinct cached column,
read a batch holding
+ // only the selected columns' buffers (in columnIndices order) so unselected
columns are never
+ // copied off-heap, loaded, or decompressed. The projected schema's field
order matches
+ // columnIndices, so the loaded root's vectors are already in output order.
If any selected
+ // attribute is absent from the cache schema (index -1), fall back to
reading the full batch.
+ private val cacheFields = arrowSchema.getFields.asScala.toSeq
+ private val canProjectOnLoad = columnIndices.forall(_ >= 0)
+ private val projectedSchema =
+ if (canProjectOnLoad) {
+ new
org.apache.arrow.vector.types.pojo.Schema(columnIndices.map(cacheFields).toList.asJava)
+ } else {
+ arrowSchema
+ }
+
// Track only the previous root to close it when next batch is produced
private var previousRoot: VectorSchemaRoot = null
@@ -1202,11 +1350,16 @@ private class ArrowCachedBatchToColumnarBatchIterator(
previousRoot = root
- // Wrap vectors in ArrowColumnVector and project to selected columns.
- val allColumns = root.getFieldVectors.asScala.map { vector =>
- new ArrowColumnVector(vector)
- }.toArray[ColumnVector]
- val selectedColumns = columnIndices.map(allColumns(_))
+ // When projected on load, the root already holds only the selected
columns in output order,
+ // so wrap its vectors directly. Otherwise it holds all cache columns and
must be selected.
+ val selectedColumns = if (canProjectOnLoad) {
+ root.getFieldVectors.asScala.map(v => new
ArrowColumnVector(v)).toArray[ColumnVector]
+ } else {
+ val allColumns = root.getFieldVectors.asScala.map { vector =>
+ new ArrowColumnVector(vector)
+ }.toArray[ColumnVector]
+ columnIndices.map(allColumns(_))
+ }
val batch = new ColumnarBatch(selectedColumns, root.getRowCount)
// Start prefetching the next batch while this one is being consumed.
@@ -1217,11 +1370,18 @@ private class ArrowCachedBatchToColumnarBatchIterator(
/** Deserialize a cached batch into its own freshly-created root. Does not
touch other roots. */
private def deserializeToRoot(cachedBatch: ArrowCachedBatch):
VectorSchemaRoot = {
- val in = new ByteArrayInputStream(cachedBatch.arrowData)
- val readChannel = new ReadChannel(Channels.newChannel(in))
- val recordBatch = MessageSerializer.deserializeRecordBatch(readChannel,
allocator)
+ // Projection pushdown: read only the selected columns' buffers out of the
cached bytes, so
+ // unselected columns are never copied off-heap, loaded, or decompressed.
+ val recordBatch = if (canProjectOnLoad) {
+ ArrowCachedBatchSerializer.readProjectedRecordBatch(
+ cachedBatch.arrowData, cacheFields, columnIndices, allocator)
+ } else {
+ val in = new ByteArrayInputStream(cachedBatch.arrowData)
+ val readChannel = new ReadChannel(Channels.newChannel(in))
+ MessageSerializer.deserializeRecordBatch(readChannel, allocator)
+ }
Utils.tryWithSafeFinally {
- val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ val root = VectorSchemaRoot.create(projectedSchema, allocator)
// VectorLoader.load fills vectors incrementally, so a failure
(malformed data, decompression
// error, OOM) can occur after earlier vectors have allocated buffers.
Close the partially
// loaded root on failure, otherwise it becomes unreachable and the
later allocator.close()
@@ -1424,6 +1584,20 @@ private class ArrowCachedBatchToInternalRowIterator(
private val arrowSchema = ArrowUtils.toArrowSchema(
cacheSchema, timeZoneId, false, false, losslessInternalTypes = true)
+ // Projection pushdown: see ArrowCachedBatchToColumnarBatchIterator. When
every selected column
+ // maps to a distinct cached column, read a batch holding only the selected
columns' buffers so
+ // unselected columns are never copied off-heap, loaded, or decompressed,
and readers bind
+ // positionally. If any selected attribute is absent from the cache (index
-1), fall back to the
+ // full batch and bind readers via columnIndices.
+ private val cacheFields = arrowSchema.getFields.asScala.toSeq
+ private val canProjectOnLoad = columnIndices.forall(_ >= 0)
+ private val projectedSchema =
+ if (canProjectOnLoad) {
+ new
org.apache.arrow.vector.types.pojo.Schema(columnIndices.map(cacheFields).toList.asJava)
+ } else {
+ arrowSchema
+ }
+
// Pre-build typed readers per column at init time -- no per-row pattern
match
private val columnReaders: Array[ArrowColumnReader] =
selectedSchema.fields.map(f => ArrowColumnReader.create(f.dataType))
@@ -1504,11 +1678,17 @@ private class ArrowCachedBatchToInternalRowIterator(
/** Deserialize a cached batch into a VectorSchemaRoot. */
private def deserializeBatch(cachedBatch: ArrowCachedBatch):
VectorSchemaRoot = {
- val in = new ByteArrayInputStream(cachedBatch.arrowData)
- val readChannel = new ReadChannel(Channels.newChannel(in))
- val recordBatch = MessageSerializer.deserializeRecordBatch(readChannel,
allocator)
+ // Projection pushdown: read only the selected columns' buffers out of the
cached bytes.
+ val recordBatch = if (canProjectOnLoad) {
+ ArrowCachedBatchSerializer.readProjectedRecordBatch(
+ cachedBatch.arrowData, cacheFields, columnIndices, allocator)
+ } else {
+ val in = new ByteArrayInputStream(cachedBatch.arrowData)
+ val readChannel = new ReadChannel(Channels.newChannel(in))
+ MessageSerializer.deserializeRecordBatch(readChannel, allocator)
+ }
try {
- val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ val root = VectorSchemaRoot.create(projectedSchema, allocator)
// VectorLoader.load fills vectors incrementally, so a failure
(malformed data, decompression
// error, OOM) can occur after earlier vectors have allocated buffers.
Close the partially
// loaded root on failure, otherwise it becomes unreachable and the
later allocator.close()
@@ -1560,10 +1740,13 @@ private class ArrowCachedBatchToInternalRowIterator(
currentRoot = root
- // Update pre-built readers with new vectors
+ // Update pre-built readers with new vectors. When projected on load, the
root holds the
+ // selected columns positionally; otherwise it holds all cache columns,
selected via
+ // columnIndices.
var i = 0
while (i < numFields) {
- columnReaders(i).setVector(root.getVector(columnIndices(i)))
+ val vectorIndex = if (canProjectOnLoad) i else columnIndices(i)
+ columnReaders(i).setVector(root.getVector(vectorIndex))
i += 1
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowCacheBenchmark.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowCacheBenchmark.scala
index afb732a00d53..7df7e97ecb25 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowCacheBenchmark.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowCacheBenchmark.scala
@@ -793,6 +793,61 @@ object ArrowCacheBenchmark extends SqlBasedBenchmark {
}
}
+ private def columnPruningWideTable(): Unit = {
+ // 1M rows x 50 long columns is a ~400MB uncompressed cache -- wide enough
that pruning 49 of
+ // 50 columns matters, but small enough to sit comfortably in the
benchmark heap so the timing
+ // reflects the read, not GC pressure from a cache that nearly fills the
heap.
+ val numRows = 1000000
+ val numCols = 50
+ val cols = (0 until numCols).map(i => s"id + $i as col$i")
+
+ // Measure only the read: summing one column while pruning the other 49.
Each case builds and
+ // fully materializes its cache before starting the timer (via
addTimerCase), so the numbers
+ // reflect the cached-scan read path -- which column pruning speeds up --
not the one-time cache
+ // materialization. An aggregate is used rather than a projection to a
noop sink because the
+ // latter does not pull the selected column's data through the read path.
The cache serializer
+ // is a JVM-wide static, so each case creates its own fresh session
(excluded from timing),
+ // which also lets the Default and Arrow serializers share one comparison
table.
+ val default =
"org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer"
+ val arrow =
"org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer"
+ runBenchmark(s"Cache with column pruning (sum 1 of $numCols columns)") {
+ val benchmark = new Benchmark(
+ s"Sum 1 of $numCols columns, $numRows rows", numRows, output = output)
+
+ def addPruningCase(name: String, serializer: String)(
+ configure: SparkSession => Unit): Unit = {
+ benchmark.addTimerCase(name) { timer =>
+ val spark = createFreshSession(serializer)
+ try {
+ configure(spark)
+ val df = spark.range(numRows).selectExpr(cols: _*)
+ df.cache()
+ df.count() // materialize the cache before timing
+ timer.startTiming()
+ df.selectExpr("sum(col0)").collect()
+ timer.stopTiming()
+ df.unpersist(blocking = true)
+ } finally {
+ spark.stop()
+ }
+ }
+ }
+
+ addPruningCase("Default cache", default)(_ => ())
+ addPruningCase("Default cache (uncompressed)", default) { spark =>
+ spark.conf.set("spark.sql.inMemoryColumnarStorage.compressed", "false")
+ }
+ addPruningCase("Arrow cache", arrow)(_ => ())
+ Seq("-1", "1", "3").foreach { level =>
+ addPruningCase(s"Arrow cache (zstd level $level)", arrow) { spark =>
+ spark.conf.set(SQLConf.ARROW_EXECUTION_COMPRESSION_CODEC.key, "zstd")
+ spark.conf.set(SQLConf.ARROW_EXECUTION_ZSTD_COMPRESSION_LEVEL.key,
level)
+ }
+ }
+ benchmark.run()
+ }
+ }
+
override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
runBenchmark("Arrow Cache vs Default Cache") {
cachePrimitiveTypes()
@@ -800,6 +855,7 @@ object ArrowCacheBenchmark extends SqlBasedBenchmark {
cacheColumnarInput()
recacheArrowData()
columnPruning()
+ columnPruningWideTable()
}
}
}
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 73c2b15daf2f..9fe540388e37 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
@@ -264,6 +264,46 @@ class ArrowCachedBatchSerializerSuite extends QueryTest
with SharedSparkSession
assert(projected.queryExecution.executedPlan.toString.contains("InMemoryTableScan"))
}
+ test("column projection prunes columns on load for every projection shape") {
+ // The read path reads only the selected columns' buffers out of the
cached bytes. Exercise a
+ // spread of projection shapes -- reordering, single column at each
position, complex columns
+ // mixed with primitives, and duplicate selection -- under both the row
and vectorized read
+ // paths, so the buffer-span arithmetic (which must skip the exact
node/buffer runs of
+ // unselected columns, including the child buffers of complex columns) is
covered end to end.
+ Seq(false, true).foreach { vectorized =>
+ withSQLConf(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key ->
vectorized.toString) {
+ val df = (1 to 50).map { i =>
+ (i, s"str$i", Seq(i, i + 1), (i.toLong, s"n$i"))
+ }.toDF("a", "b", "arr", "st")
+ df.cache()
+ try {
+ def expected(cols: Seq[String]): Seq[Row] = (1 to 50).map { i =>
+ Row.fromSeq(cols.map {
+ case "a" => i
+ case "b" => s"str$i"
+ case "arr" => Seq(i, i + 1)
+ case "st" => Row(i.toLong, s"n$i")
+ })
+ }
+ // Selecting a complex column after skipping a var-width one
exercises skipping the
+ // multi-buffer runs (offset + data) of the unselected string.
+ checkAnswer(df.select("arr"), expected(Seq("arr")))
+ checkAnswer(df.select("st"), expected(Seq("st")))
+ checkAnswer(df.select("b"), expected(Seq("b")))
+ // Reordered projection: the loaded root must be in output
(columnIndices) order.
+ checkAnswer(df.select("st", "a"), expected(Seq("st", "a")))
+ checkAnswer(df.select("arr", "b", "a"), expected(Seq("arr", "b",
"a")))
+ // Duplicate selection maps two output columns to one cached column.
+ checkAnswer(df.select("a", "a"), (1 to 50).map(i => Row(i, i)))
+ // Full projection (no pruning) still round-trips.
+ checkAnswer(df.select("a", "b", "arr", "st"), expected(Seq("a", "b",
"arr", "st")))
+ } finally {
+ df.unpersist()
+ }
+ }
+ }
+ }
+
test("caching with multiple batches") {
withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "10") {
val df = (1 to 50).map(i => (i, s"str$i")).toDF("a", "b")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
