Support useLargeVarTypes in accelerated mapInArrow/mapInPandas

[sunchao]

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 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:

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:

• Large-offset Arrow vectors from PyArrow UDFs cannot be serialized for broadcast or collect #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.
• Extend native Arrow UDF path to scalar Python UDFs (ArrowEvalPythonExec) #5386 extends acceleration to scalar Python UDFs and mentions this configuration as an existing limitation; it does not track implementing large input types.