[HyukjinKwon]

What changes were proposed in this pull request?

This PR disables spark.sql.execution.arrow.pythonUDF.columnarInput.enabled by default, and makes ArrowEvalPythonExec.doExecute check it too.

With it enabled, the Arrow columnar path of ColumnarArrowEvalPythonEvaluatorFactory, added by SPARK-56350, keeps the pass-through ColumnVectors of each input batch until the Python result of that batch comes back, which can be many batches later. That is only correct if the batch stays valid until then, which a columnar PartitionReader does not promise: get() only returns the same value until next() is called. A reader may refill the same vectors for every batch, or release them in close(), which DataSourceRDD calls as soon as the reader is exhausted, since SPARK-55715. With columnar input disabled, ColumnarToRow converts the batches to rows before the UDFs, as before SPARK-56350. But ArrowEvalPythonExec.doExecute took the columnar input path whenever its child supported columnar, whatever the configuration: a child that also returns rows, such as the scan of the Arrow cache, gets no ColumnarToRowExec, so the Arrow cache failed even with columnar input disabled. doExecute now checks supportsColumnar, which includes the configuration, so that such a child returns rows.

As discussed in the review, the Arrow columnar path will instead respect the vectors that their producer reuses, with a marker interface that @viirya proposed, which can enable columnar input again on master. This PR is the fix for the maintenance branches, and for master until then.

It also adds two options to the test source ArrowBackedDataSourceV2: reuseVectors, under which its reader refills the same vectors for every batch and closes them when it is closed, and batchSize.

Why are the changes needed?

With columnar input enabled, a scalar pandas or Arrow UDF on a column of a source whose Arrow vectors are refilled or released fails with IndexOutOfBoundsException: index: 0, length: 1 (expected: range(0, 0)), or pairs the pass-through values of a later batch with the UDF results of an earlier one. It was reported by @FelixYBW and reproduced by @sarutak in #59424 with native data sources, whose reader imports every batch into the same vectors, and also without them, with ArrowCachedBatchSerializer: spark.range(0, 2000000, 1, 1).cache() and the same UDF, over 200 batches of 10000 rows (spark.sql.execution.arrow.maxRecordsPerBatch), fails with the same exception, also with columnar input disabled before this PR. As @sarutak found, the Arrow columnar case of the benchmark of SPARK-56350, bench_arrow_columnar_udf.py, fails the same way on master and on v4.2.0, since SPARK-55715: its source frees its vectors only when it is closed, which DataSourceRDD now does before the last results come back from Python.

Does this PR introduce any user-facing change?

Yes. By default, Python UDFs on the columns of an Arrow-backed source no longer fail or return wrong results: Spark converts the batches to rows before the UDFs, as before 4.2.0. Setting spark.sql.execution.arrow.pythonUDF.columnarInput.enabled to true enables columnar input again, for upstream operators whose batches stay valid until their results come back. The migration guide says so, for 4.2.1 and 4.4.

How was this patch tested?

• A new test in ArrowColumnarPythonUDFSuite, Arrow-backed source that reuses its vectors: correct with the default configuration, runs a UDF on a string column, over 2 batches and over 2000 batches of 100 rows of a source that refills its vectors. It checks that ArrowEvalPythonExec does not take columnar input, and the pass-through columns of every row against its UDF result. With columnar input enabled, it fails with the IndexOutOfBoundsException above in both cases.
• A new suite, ArrowCachedPythonUDFSuite, runs a UDF over 200 batches of 1000 rows of the Arrow cache with the default configuration, and checks that ArrowEvalPythonExec reads rows although its child supports columnar, and the pass-through column of every row against its UDF result. Without the change to doExecute, it fails with the IndexOutOfBoundsException above.
• The other tests of ArrowColumnarPythonUDFSuite, which test the columnar input path, enable it, as does ArrowPythonUDFCachedInputTests in PySpark, which tests it over the Arrow cache. Both suites pass (8 tests), as do scalastyle and ruff.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 5.5)

This pull request and its description were written by Isaac.

[sarutak]

Thanks for the fix, @HyukjinKwon. I built 9c9a900 with Java 22 and ran the following:

• The new test passes, and it fails with the same IndexOutOfBoundsException when only ColumnarArrowEvalPythonEvaluatorFactory.scala is reverted to the merge base. So it goes through the Arrow columnar path, as intended.
• The same test with spark.python.udf.pipelined.enabled=true (a suite that sets it in sparkConf): correct for 200 and 200000 rows.
• The ArrowCachedBatchSerializer case from #59424, with an arrow_udf and a pandas_udf stacked on top of each other (two ArrowEvalPythonExecs), 2000000 rows in batches of 1000: no mismatch, in both the default and the pipelined mode.

[sarutak]

VectorAppender does not support ListViewVector, which ArrowColumnVector reads since SPARK-57929. With Arrow 19.0.0, I copied a vector of each type that ArrowColumnVector accepts this way: all of them work (including ViewVarCharVector, ViewVarBinaryVector, map, list and struct), except ListViewVector, which fails with UnsupportedOperationException: VectorVisitor for ListViewVector is not supported.

Since every pass-through column is copied, a query over a source that returns a ListViewVector in any column now fails on this path, even when the UDF does not read that column. Before this PR, the column was only referenced, so it worked. Could the check at L157 also send such batches to the row queue path (for example, by requiring that no column is a ListViewVector), and could a test cover it?

[HyukjinKwon]

Thanks, fixed in c6bcd01: the Arrow columnar path is now only taken when no column has a list view (ListViewVector or LargeListViewVector) at any level, so these batches take the row queue path. Arrow-backed source with columns that the Arrow columnar path does not copy reads a list view column, and fails with UnsupportedOperationException: VectorVisitor for ListViewVector is not supported. without the check.

[sarutak]

In pipelined mode, this listener can run while the writer thread is still in copyColumns or before its passThruQueue.add at L216. Then the copy is not closed, and allocator.close() fails with "Memory was leaked".

The root cause is not in this PR but in the listener of BasePythonRunner.createPipelinedDataIn, which runs before this one. It calls writerFuture.cancel(true) and then writerFuture.get(), but get() on a cancelled FutureTask throws CancellationException right away, without waiting for the thread. A small program confirmed this: get() returned while the task was still running, and the task finished 500 ms later. So the comment at PythonRunner.scala L545 ("Wait for the writer to actually exit ...") does not hold. This PR adds a resource that the writer thread allocates and that a listener frees.

To check it, I forced this order on 9c9a900 without changing its non-test code, in pipelined mode with the reuseVectors source:

1. The task thread reads 10 rows and then stops reading (executedPlan.execute().mapPartitions(it => Iterator(it.take(10).size)), since limit is pushed below the UDF).
2. The writer thread, in copyColumns of the next batch, blocks without being interruptible in getValueVector of the last column. I did this with a test-only subclass of ArrowColumnVector, to stand in for the copy of a wide column. At that point, the copies of the other three columns are allocated.
3. The task ends, and this listener runs while the writer is blocked there. (A listener that the test reader registers, and that runs after this one, confirms the order.)

All 10 runs failed with TaskCompletionListenerException: Memory was leaked by query. Memory leaked: (3072) from Allocator(pass-through columns of Python UDFs in task ...). Also, the memory allocated in ArrowUtils.rootAllocator grew by 3072 bytes with each run, and was still allocated 2 seconds after the job ended, so the copies were not freed. Without forcing the order, I did not see it in 100 runs on the narrow columns of the test source. A task stops early this way when it is killed or fails, or when an operator above it stops reading.

Could the copy and add be guarded against the listener here, for example with a lock and a closed flag that both check, so that a copy made after the listener ran is closed right away? Fixing the runner so that its listener really waits for the writer (for example, with a latch released at the end of PipelinedWriterRunnable.run) would also cover this, and the existing cases that the comment there describes, but that seems better as a separate JIRA.

[viirya]

The root cause here, get() returning right after cancel(true) without waiting for the writer, is tracked in SPARK-60113, which makes the runner's listener wait for the writer to exit. Once that lands, this path is covered without its own lock. This listener is registered before pyRunner.compute, so it runs after the runner's listener, by which time the writer has finished copyColumns and add. The same holds for stacked ArrowEvalPythonExecs in pipelined mode, where the outer writer thread drives the inner resultIter.map and the assignment of outputPassThruCols.

@HyukjinKwon, could you add a comment here saying that this listener must be registered before compute for that reason? Nothing else enforces the order.

[HyukjinKwon]

Thanks for forcing the order to show it. Fixed in c6bcd01 with a lock, as you suggested. The listener sets completed and releases the queued copies under it. The writer copies outside of it, then adds the copies, or releases them right away if the task completed meanwhile. The allocator is closed by whichever of the two ends last, so a copy that blocks does not block the listener, and nothing leaks. As @viirya pointed out, the runner itself is SPARK-60113 (#59332), and the comment here now also says that this listener must be registered before compute.

[HyukjinKwon]

Thanks, I added the comment in c6bcd01: this listener must be registered before compute, so that it runs after the listener of the runner. Until SPARK-60113 makes that listener wait for the writer, the copies are guarded against it, as @sarutak suggested: a copy made once the task completed is released right away, and the allocator is closed by whichever of the two ends last. Once SPARK-60113 lands, the guard can be removed.

[sarutak]

Two suggestions for the test:

• child.supportsColumnar is also true for the other two paths. The test takes the Arrow columnar path only because the cast of the test UDF is removed, so if that changes, it would still pass without checking anything. Could it also assert that all the inputs of the UDF are column references, for example arrowExec.head.udfs.flatMap(_.children).forall(_.isInstanceOf[AttributeReference])?
• No test sets spark.python.udf.pipelined.enabled=true, so the change to ConcurrentLinkedQueue is not covered, and no test has a source with both Arrow and non-Arrow columns, which the new check at L157 is for. As the pipelined mode is a SparkConf, the first one would need a suite that sets it in sparkConf.

[HyukjinKwon]

Thanks, done in c6bcd01:

• The tests now also check that the inputs of the UDF are columns, besides child.supportsColumnar.
• ArrowColumnarPythonUDFPipelinedSuite runs the tests of the suite with spark.python.udf.pipelined.enabled=true, and Arrow-backed source with columns that the Arrow columnar path does not copy reads a column that is not Arrow-backed, and one that is a list view. Without the check at L157, it fails with an internal error for the first one.

[sarutak]

The copy takes source.getValueCount rows, while the output batch has batch.numRows() rows. Before this PR, a reader that set the number of rows of the batch but not the value count of its vectors worked on this path, since ArrowColumnVector does not read the value count and writeArrowDirect sets it on the inputs of the UDF. Now the copy is shorter, and reading the missing rows throws IndexOutOfBoundsException (I checked this with a value count of 0) or, within the allocated capacity, may read as nulls. The readers in Spark set it, so this is minor, but a check that source.getValueCount >= batch.numRows(), failing with a clear error, would make this case explicit.

[HyukjinKwon]

Thanks, done in c6bcd01: a vector with fewer values than the rows of its batch fails with an internal error that says so, such as The Arrow vector of the column name has 0 values, but its batch has 100 rows.

[sarutak]

Two more points, on the cost of the copy and on memory.

The benchmark of SPARK-56350

I ran bench_arrow_columnar_udf.py (one partition, 100-character strings in data). On master without this PR, and on v4.2.0, its Arrow columnar case fails with the same IndexOutOfBoundsException for 10000, 500000 and 2000000 rows. It worked when SPARK-56350 was merged (2d5c430), and it stopped working with SPARK-55715 (f8f4813, five days later; both are in 4.2.0): I ran it on f8f4813 and on its parent. Before SPARK-55715, DataSourceRDD closed the last reader of a task only at the end of the task; since then, PartitionIterator.hasNext closes a reader as soon as it is exhausted, before the last results come back from Python. The source of the benchmark frees its vectors only in close(), so it became one more case of the bug that this PR fixes. A source that refills or closes its vectors for every batch, such as ArrowCachedBatchSerializer, fails whatever DataSourceRDD does. The existing tests of ArrowColumnarPythonUDFSuite did not catch this, since their UDFs on id keep the cast and take the third path.

With this PR, it runs again. Three rounds with --rows 2000000 on a laptop (so only the ratio is meaningful):

• columnar input enabled: 2697, 2643 and 2654 ms on average;
• disabled (ColumnarToRow): 2887, 2889 and 2899 ms;
• so 1.07x to 1.09x faster. SPARK-56350 reported 1.10x for 2M rows of 200-character strings and 1.17x for 10M rows of 100-character strings.

With timers added around copyColumns, the copy took about 62 ms of a run of about 2.7 s (1954 batches, peak of the allocator 11.6 to 12.1 MB), so it takes a part of the gain but does not remove it. Could you add these numbers, or your own, to the description?

Memory

The copies are allocated from a child of ArrowUtils.rootAllocator with no limit, so they are not counted by the TaskMemoryManager and cannot spill, while the row queue path keeps the same rows in a HybridRowQueue, which is counted and can spill. How many batches are in flight is set by how far the input runs ahead of the results, not by the size of the pass-through columns. On the test source (4 narrow columns), the peak of the allocator was 5.9 MB for 200000 rows in batches of 100 and 20 MB for 2000000 rows in batches of 1000, so a few hundred batches were in flight, and this grows with the width of the pass-through columns. Before this PR, a source that allocates new vectors for every batch and does not free them early kept about the same memory alive, so this is not necessarily worse in practice, but it may be worth a sentence in the comment of evalArrowColumnar.

[viirya]

Thanks for the fix, @HyukjinKwon. The copy is needed: retaining the source buffers would not survive a reader that refills its vectors in place, and transferring them would take away the UDF inputs before the runner writes them. The new test fails on the merge base with the reported IndexOutOfBoundsException and passes with this change.

The ListViewVector case that @sarutak raised looks like the one to address before merging, since the copy now covers every pass-through column, including those the UDF does not read.

I have one design question. Could Arrow-backed batches simply take path 2 instead? Path 2 already writes Arrow-backed UDF inputs directly, since ColumnarArrowPythonInput.isArrowBacked is checked for every batch, and it keeps the pass-through rows in a HybridRowQueue, which the TaskMemoryManager counts and which can spill. Sending these batches to path 2 would also fix the bug, without a new allocator outside memory accounting, the ListViewVector and value count cases, or the lifecycle code. What path 1 still adds is avoiding the row copy of the pass-through columns and the RowToColumnConverter on output. The numbers above compare against disabling columnar input, not against path 2. Could you also measure path 2 on the same benchmark? If the gap is small, path 2 may be the simpler fix, especially for the backports.

[viirya]

Now that the pass-through columns are copied, the input batch is no longer needed once it has been written to Python, so this path could also release it, as ColumnarToRow does with closeIfFreeable (SPARK-50235). Today nothing does: a row-based parent calls ArrowEvalPythonExec.doExecute, since it is also supportsRowBased, and that flattens the batches without closeIfFreeable. So a source that allocates new vectors for every batch and relies on the consumer to release them keeps all of them until the end of the task. I checked it with such a test source over 50000 rows in batches of 100: none of the 500 batches was released, both on 9c9a900 and on the merge base. So this is not a regression of this PR. With the change below, all 500 were released as they were consumed, the peak of the source's allocator dropped from 73.7 MB to 147 KB, and the new reuseVectors test still passes:

val bufferedIter = new Iterator[ColumnarBatch] {
  private var previous: ColumnarBatch = null
  override def hasNext: Boolean = {
    if (previous != null) {
      previous.closeIfFreeable()
      previous = null
    }
    inputIter.hasNext
  }
  override def next(): ColumnarBatch = {
    val batch = inputIter.next()
    passThruQueue.add((copyColumns(batch, allocator), batch.numRows()))
    previous = batch
    batch
  }
}

The previous batch is closed before the producer is asked for the next one, so a reader that refills the same vectors is not affected. Paths 2 and 3 have the same gap, so a separate JIRA for all three paths would also work.

[HyukjinKwon]

Thanks for checking it. As you said, it is not a regression of this PR, and paths 2 and 3 have the same gap, so I opened SPARK-60306 to release the input batches of all three paths once they are written to Python, with your change for this path.

[viirya]

nit: when the parent consumes columnar batches, such as the Arrow cache, it may already have closed these columns with closeIfFreeable. Closing them again is harmless, but saying so here would save the next reader from checking.

[HyukjinKwon]

Done in c6bcd01: the comment says that a columnar parent, such as the Arrow cache, may already have closed the columns with closeIfFreeable, and that closing them again has no effect.

[viirya]

nit: batchSize only applies when reuseVectors is true. ArrowBackedPartitionReader still uses BATCH_SIZE, so readArrowSource(batchSize = ...) without reuseVectors ignores it without saying so. Could the existing reader take it too, or could the option say that it only applies to the reusing reader?

[HyukjinKwon]

Done in c6bcd01: ArrowBackedPartitionReader takes batchSize too.

[HyukjinKwon]

@sarutak Thank you for running the benchmark, and for finding that it fails on master and v4.2.0 since SPARK-55715. I added your numbers to the description, with the cause.

@viirya Thank you for the review. On sending Arrow-backed batches to the row queue path instead: I measured it on 2,000,000 rows in one partition, with 100-character strings in data, a UDF on data and the noop sink, in 5 runs after a warm-up:

Path	Mean
Arrow columnar path, with the copies	1458 ms
Row queue path, forced	1862 ms
Columnar input disabled	1661 ms

The row queue path is slower than without columnar input, since it converts the joined rows back to columns, which the row-based parent then flattens again, while the Arrow columnar path with the copies stays 1.14x faster than without columnar input. So I kept the copy, and fixed the list views and the value counts as you and @sarutak suggested. For the backports, branch-4.2 has no pipelined mode, so the guard against the listener is not needed there.

On memory, the comment of evalArrowColumnar now says that the copies are not counted by the task memory manager and cannot spill, and that their memory grows with how far the input runs ahead of the results, as for a source that allocates new vectors for every batch. The description has the peaks that @sarutak measured. Releasing the input batches once they are written, as you suggested, applies to all three paths, so I opened SPARK-60306 for it.

[viirya]

Following up on my review: I wrote there that the copy is needed, and I'd like to take that back. I don't think path 1 should copy every pass-through column unconditionally. Path 1 was designed to keep the pass-through columns as they are, and copying them on every batch turns it into a different path, with an allocator outside memory accounting, the ListViewVector gap, and the lifecycle code that comes with it.

[HyukjinKwon]

@viirya Thanks for following up. I agree that the copy changes what path 1 is. The difficulty is that keeping the pass-through columns as they are is only correct if a batch stays valid until its result comes back, which ColumnarBatch does not promise, and which no source can promise today: since SPARK-55715, DataSourceRDD closes a reader as soon as it is exhausted, so even the source of the benchmark, which allocates new vectors for every batch, frees them before the last results come back. So I see these options, with the numbers above:

1. Copy the pass-through columns, as this PR does: 1458 ms, 1.14x faster than without columnar input, but with an allocator outside memory accounting, the list views, and the lifecycle code.
2. Send these batches to path 2: correct and accounted for, but 1862 ms, slower than without columnar input.
3. Keep path 1 as it is, only for sources that declare that their batches stay valid until the end of the task, which would be a new opt-in, and otherwise not take columnar input: 1661 ms for the others.
4. Disable spark.sql.execution.arrow.pythonUDF.columnarInput.enabled by default as the fix for 4.2, and settle the design of path 1 separately.

Which would you prefer? For the backports, 4 is the smallest change, and 3 needs a new API. I'm happy to switch this PR to either.

[viirya]

@HyukjinKwon Thanks for laying out the options. I'd go with 4 for the backports, and for master I'd rather settle path 1 as below than copy every pass-through column.

DataSourceRDD closes the reader, not the batches it has returned; whether those are released at that point is up to the reader's own close(). Nothing in the PartitionReader contract lets a reader release, in close(), batches it has already returned: get() only says it returns the same value until next() is called. A producer that releases or rewrites batches it has returned cannot know whether its consumer still holds them. That is the case here: DataSourceRDD closes the reader as soon as next() returns false, inside the consumer's hasNext, and path 1 calls hasNext right after writing a batch to Python, while it still holds that batch's pass-through columns.

What a consumer can rely on is the batch itself. closeIfFreeable (SPARK-50235) already works that way: a writable or constant vector is reused by its producer, and a consumer must not release it, while any other vector is the consumer's to release once it is done with it. NativePartitionReader in #59424 follows this: it imports every batch into the same vectors, and wraps them in ReusedArrowColumnVector, whose closeIfFreeable does nothing, for that reason. Path 1 ignores that signal, since it only checks for ArrowColumnVector.

So instead of copying, path 1 could respect it. For example, an internal marker interface in sql/core, implemented by WritableColumnVector and ReusedArrowColumnVector, which keeps extending ArrowColumnVector, would say that a vector is reused by its producer. Path 1 would send batches with such a column to path 2, and hold the other columns without copying them, leaving them to the consumer of its output batch to release with closeIfFreeable. Batches of a native data source would then take path 2, and how fast path 2 is is a separate question from this bug. We'd be happy to open a PR that adds this marker interface and closes the gaps in path 1 and in the producers and consumers described below.

That leaves the producers that reuse or release Arrow vectors they have returned without saying so:

• ArrowCachedBatchToColumnarBatchIterator closes the previous root in next(), on the premise that a consumer which asks for the next batch is done with the previous one. It should only close a root once none of its vectors is still in use, or mark its vectors as reused.
• ArrowOutputProcessorImpl loads every batch of the Python runner into the same ArrowColumnVectors, which an outer ArrowEvalPythonExec would hold on path 1. Its data buffers are already new for every batch, as the reader deserializes each record batch into new buffers, so loading into a new root for every batch would cost little; marking its vectors as reused also works.

On the consumer side, closeIfFreeable is not called by ArrowEvalPythonExec.doExecute and the non-codegen ColumnarToRowEvaluatorFactory, which flatten the batches without it, so a source that relies on its consumer keeps its vectors until the end of the task. SPARK-60306 could cover these consumers, rather than releasing the input batches of path 1 once they are written.

What do you think?

[HyukjinKwon]

@viirya Thanks, that design makes sense to me: respecting the signal of the batch itself, as closeIfFreeable does, is better than copying every pass-through column, and the producers that reuse or release Arrow vectors without saying so are the other half of it.

I switched this PR to 4 in 91aa31b: it disables spark.sql.execution.arrow.pythonUDF.columnarInput.enabled by default, and the evaluator is back to what it is on master. I kept it disabled on master too, so that master and the maintenance branches are correct with the default configuration until path 1 respects the marker, and your PR can enable it again there; tell me if you'd rather leave master as it is. The tests of the columnar input path enable it, and a new test reads the source that refills its vectors with the default configuration, which fails with columnar input enabled. I updated the title and the description.

Please go ahead with the marker interface; I'm happy to review it. ReusedArrowColumnVector in #59424 can implement the marker, whichever of the two lands first. I also re-scoped SPARK-60306 to the consumers that do not call closeIfFreeable, as you suggested.

@sarutak The copy, the guard and the checks that you reviewed are gone with this change. Thanks for reviewing them so carefully: the list views and the pipelined mode were part of why we moved away from the copy.

[sarutak]

Thanks for switching to this, @HyukjinKwon. I built 21eb5b5: ArrowColumnarPythonUDFSuite passes, and with the old default true restored, the new test fails (at the !supportsColumnar assertion; without that assertion, with the IndexOutOfBoundsException).

The new default does not fix the Arrow cache case

With spark.sql.cache.serializer=org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer, spark.range(0, 2000000, 1, 1).cache() and a pandas UDF that returns its input, select("id", f(col("id"))) still fails with IndexOutOfBoundsException: index: 0, length: 1 (expected: range(0, 0)) on 21eb5b5, with the default configuration and with columnarInput.enabled=false set explicitly. It fails the same way on the merge base with false, so the configuration was never a workaround for this case. (The same happens with 200000 rows and spark.sql.execution.arrow.maxRecordsPerBatch=1000.)

The reason is that ArrowEvalPythonExec.doExecute decides on child.supportsColumnar, not on the configuration:

override def supportsColumnar: Boolean =
  child.supportsColumnar && conf.arrowPySparkUDFColumnarInputEnabled
override def supportsRowBased: Boolean = true

override protected def doExecute(): RDD[InternalRow] = {
  if (child.supportsColumnar) {
    // Columnar path: delegate to doExecuteColumnar, ...

For a DSv2 scan, ApplyColumnarRulesAndInsertTransitions puts a ColumnarToRowExec below ArrowEvalPythonExec once it does not support columnar, so its child no longer supports columnar, and the configuration works, as the new test shows. But InMemoryTableScanExec supports both rows and batches, so ensureOutputsRowBased adds no transition, and the plan is ArrowEvalPython <- InMemoryTableScan. The child still supports columnar, so doExecute goes through doExecuteColumnar and the Arrow columnar path, whatever the configuration says.

Checking the configuration there too, with if (supportsColumnar), fixes it: the same queries have no mismatch with the default and with false, and still fail with true, and ArrowColumnarPythonUDFSuite passes. Then InMemoryTableScanExec returns rows on this path. Could this PR include that change, and a test over the Arrow cache with the default configuration and more than one batch? ArrowPythonUDFCachedInputTests now enables columnar input, so nothing tests the default there.

Smaller points

• The Scaladoc of ColumnarArrowEvalPythonEvaluatorFactory (L47-49) and of evalArrowColumnar (L178-179) still say that the references to the pass-through columns are safe "because Arrow vectors are independently allocated per batch", which the new doc of the configuration contradicts.
• Since this changes the default of a configuration released in 4.2.0, could docs/sql-migration-guide.md have an entry, in "Upgrading from Spark SQL 4.2.0 to 4.2.1" and in the section of the release it lands in, with how to enable it again?
• The first version of this PR replaced the ArrayDeque of evalArrowColumnar with a ConcurrentLinkedQueue, since in pipelined mode the writer thread adds to it and the task thread polls it. With that change dropped, the race is still there when both columnarInput.enabled and spark.python.udf.pipelined.enabled are set. Both are opt-in now, so it does not block this PR, but it would be worth tracking with the follow-up that enables the Arrow columnar path again. (I did not reproduce this race; it follows from ArrayDeque not being thread-safe.)
• In the description, spark.sql.inMemoryColumnarStorage.batchSize=1000 has no effect with ArrowCachedBatchSerializer, which sizes its batches by spark.sql.execution.arrow.maxRecordsPerBatch (10000 by default), so that repro has 200 batches of 10000 rows. I wrote the same setting in #59424, sorry for that.

[sarutak]

!arrowExec.head.supportsColumnar is also true in the Arrow cache case above, where the Arrow columnar path still runs, since doExecute decides on child.supportsColumnar. Could this assert that the child does not support columnar (!arrowExec.head.child.supportsColumnar), as the first test of the suite does for the opposite case? Then the assertion checks the path that runs.

[HyukjinKwon]

Thanks, done in a37bb49: the test now checks that the child of ArrowEvalPythonExec does not support columnar. The new ArrowCachedPythonUDFSuite covers the Arrow cache, whose scan supports columnar while ArrowEvalPythonExec does not.

[HyukjinKwon]

@sarutak Thank you, you're right: the configuration did not apply to the Arrow cache. Fixed in a37bb49:

• ArrowEvalPythonExec.doExecute now checks supportsColumnar, which includes the configuration, instead of child.supportsColumnar, so that a child that also returns rows, such as the scan of the Arrow cache, returns rows when columnar input is disabled.
• A new suite, ArrowCachedPythonUDFSuite, runs a UDF over 200 batches of 1000 rows of the Arrow cache with the default configuration. It fails with the IndexOutOfBoundsException without the change to doExecute, and passes with it. It is a suite of its own, since the serializer of the cache is kept for the JVM.
• The Scaladoc of the evaluator and of evalArrowColumnar no longer says that the references are safe.
• The migration guide has the new default in "Upgrading from Spark SQL 4.2.0 to 4.2.1" and in "Upgrading from Spark SQL 4.3 to 4.4", where a PR against master first ships. As 4.2.1 and 4.3.0 are at their release candidates, the sections may need to follow the branches that it is backported to.
• I fixed the repro in the description: 200 batches of 10000 rows, by spark.sql.execution.arrow.maxRecordsPerBatch.
• The ArrayDeque race in pipelined mode: agreed, both settings are opt-in now. @viirya, could the follow-up that enables the Arrow columnar path again cover it too?

[viirya]

@HyukjinKwon Yes, the follow-up will cover the ArrayDeque race: it replaces the queue of evalArrowColumnar with a ConcurrentLinkedQueue, since the Arrow columnar path will hold the pass-through columns again, and in pipelined mode the writer thread adds to the queue while the task thread polls it.

Keeping columnar input disabled on master too until then is fine with me. The follow-up will enable it again, and rewrite the doc of the configuration and the migration guide entry for 4.4 at the same time.

Thanks for checking the configuration in doExecute as well, and thanks @sarutak for catching the Arrow cache case.

[sarutak]

Thanks for the update, @HyukjinKwon. I built a37bb49 with Java 22 and checked the following:

• The Arrow cache repros from my last comment (spark.range(0, 2000000, 1, 1).cache(), and 200000 rows with spark.sql.execution.arrow.maxRecordsPerBatch=1000) have no mismatch with the default configuration and with columnarInput.enabled=false, and still fail with true.
• ArrowCachedPythonUDFSuite and ArrowColumnarPythonUDFSuite pass. With only ArrowEvalPythonExec.scala reverted to 21eb5b5, ArrowCachedPythonUDFSuite fails with the IndexOutOfBoundsException, so it covers the change to doExecute.
• The Scaladoc, the migration guide entries and the description look good to me.

LGTM.

[HyukjinKwon]

Merge Summary:

• merged into master 2487c98
• merged into branch-4.x d5a46c9

Posted by merge_spark_pr.py