apache / apache/datafusion-comet

Support useLargeVarTypes in accelerated mapInArrow/mapInPandas

Offen
#5,555 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
enhancement
Vorherrschende Sprache
Scala
Sterne
1.3k
Forks
375
Ø Merge
2 T. 10 Std.
Gemergte PRs (30 T.)
231

Beschreibung

### 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.

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Beginne mit der Eligibility-Prüfung in spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala um Zeile 204 und verfolge CometMapInBatch durch die Grenze zwischen mapInArrow und mapInPandas. Führe die bereitgestellte Reproduktion mit Spark 4.0.4/PyArrow 25.0.1 aus und untersuche die für Worker sichtbaren Typen sowie die Explain-Ausgabe. Als abgeschlossen gilt die Arbeit, wenn berechtigte Pläne weiterhin beschleunigt ausgeführt werden, große String- und Binary-Typen über die aufgeführten Fälle hinweg erhalten bleiben, Allokationen freigegeben werden und die Einschränkung dokumentiert ist.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, scala, spark
Bereich
backend, data-engineering
Issue-Typ
Feature
Schwierigkeit
5/5
Geschätzter Aufwand
Über eine Woche
Aktivitätsstatus
Aktiv
Klarheit
Größtenteils klar
Anfängerfreundlichkeit
45/100

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.