apache / apache/datafusion-comet
Support useLargeVarTypes in accelerated mapInArrow/mapInPandas
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
### What is the problem the feature request solves?
Spark's `spark.sql.execution.arrow.useLargeVarTypes=true` requests Arrow `large_string` and `large_binary` input columns with 64-bit offsets. Comet's accelerated `mapInArrow` / `mapInPandas` path currently falls back whenever this setting is enabled, even when `spark.comet.exec.pyarrowUDF.enabled=true` and the input is already a Comet columnar batch.
The [current eligibility check on upstream main](https://github.com/apache/datafusion-comet/blob/e13f9d265af0addb81385305d9c6553cdb6f9dfc/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala#L204) returns `None` unconditionally for this configuration. Users therefore lose the columnar Python execution path and pay the Arrow-to-row-to-Arrow conversion for that stage. Other operators may still use Comet. The fallback applies even to small batches; it is not conditional on a column actually exceeding the ordinary Arrow size limit.
Native Comet input vectors use ordinary string/binary layouts with 32-bit offsets. Serializing those buffers under their matching schema produces valid IPC, but does not supply the large input types requested by Spark. Removing the fallback or changing only the advertised field types would not implement support: the offset buffers must actually match the large types.
A small reproduction was verified on Spark 4.0.4 / JDK 17 / PyArrow 25.0.1 using the Comet jar from PR #5368 at `5fd48b3e`. With the Python acceleration flag enabled and three Parquet rows containing strings, binary values, nulls, and empty values:
| `useLargeVarTypes` | Python operator | Types observed inside the worker |
| --- | --- | --- |
| `false` | `CometMapInBatch` | `string`, `binary` |
| `true` | Spark `MapInArrow`, preceded by `CometColumnarToRow` | `large_string`, `large_binary` |
Both executions returned the same values. This is a missing acceleration capability, not a query correctness failure.
For an existing Comet-enabled Spark 4.x session with a working native Parquet scan:
```python
import tempfile
import pyarrow as pa
spark.conf.set("spark.comet.exec.pyarrowUDF.enabled", "true")
spark.conf.set("spark.sql.execution.arrow.useLargeVarTypes", "true")
def passthrough(batches):
for batch in batches:
assert pa.types.is_large_string(batch.schema.field("s").type)
assert pa.types.is_large_binary(batch.schema.field("b").type)
yield batch
with tempfile.TemporaryDirectory() as directory:
path = directory + "/input"
spark.createDataFrame(
[(1, "a", bytearray(b"x")), (2, None, None), (3, "", bytearray())],
"id int, s string, b binary",
).coalesce(1).write.parquet(path)
source = spark.read.parquet(path)
result = source.mapInArrow(passthrough, source.schema)
result.explain()
print(result.collect())
```
### Describe the potential solution
Allow eligible `mapInArrow` and `mapInPandas` operations to use `CometMapInBatch` while preserving Spark's requested large Arrow input types. This could widen offsets at the Comet-to-Python boundary or produce suitable vectors upstream; the implementation should preserve source ownership and reuse payload buffers where practical.
Acceptance criteria:
- With both settings enabled, an otherwise eligible plan remains accelerated and Python receives the same large string/binary input types as vanilla Spark.
- Preserve values, nulls, nested strings/binary values, empty batches, and behavior across multiple batches and chained UDFs for both APIs.
- Account for and release conversion allocations on success and failure, without changing the ordinary `useLargeVarTypes=false` fast path.
- Add worker-visible type assertions and plan assertions, measure the conversion cost, and update the documented limitation. Document any remaining size limits in native producers; small-batch type tests alone do not demonstrate support beyond 2 GiB.
### Additional context
This limitation predates #5368, which removes the intermediate input-buffer copy but deliberately retains the fallback.
Related issues have different scopes:
- #5488 concerns serializing large-offset vectors returned by Python for broadcast or collect. That output-side issue does not cover supplying large input vectors to Python.
- #5386 extends acceleration to scalar Python UDFs and mentions this configuration as an existing limitation; it does not track implementing large input types.
Contributor guide
Research direction
Start with the eligibility check in spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala around line 204, then trace CometMapInBatch through the mapInArrow and mapInPandas boundary. Run the provided Spark 4.0.4/PyArrow 25.0.1 reproduction and inspect worker-visible types and explain output. Done means eligible plans stay accelerated, large string and binary types are preserved across the listed cases, allocations are released, and the limitation is documented.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, scala, spark
- Domain
- backend, data-engineering
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100