sunchao opened a new issue, #5555: URL: https://github.com/apache/datafusion-comet/issues/5555
### 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
