fix: experiments with large types for aggregated values

[comphead]

Which issue does this PR close?

Closes #4718 .

Support LargeUtf8/LargeBinary group keys in HashAggregate to bypass the 2 GiB offset cap

Problem

CUBE / GROUPING SETS + COUNT(DISTINCT wide_string) workloads with high group-key cardinality trip DataFusion's per-task ByteGroupValueBuilder<i32> byte-buffer cap (i32::MAX = 2 147 483 647),
surfacing as:

org.apache.comet.CometNativeException: offset overflow, buffer size > 2147483647
  at datafusion physical-plan/.../multi_group_by/bytes.rs:202-206

The cap is per-column per-task on the group-key accumulator, hit when the cumulative bytes of one Utf8 column across all interned distinct group tuples exceed 2 GiB.

Fix

Add a new config spark.comet.exec.useLargeDataTypes (default true) that promotes Utf8/Binary group-by expressions to LargeUtf8/LargeBinary before the aggregate, routing DataFusion to
ByteGroupValueBuilder<i64> (i64 offsets, effectively unbounded buffer). The Large variant is preserved end-to-end through shuffle and mapped back to Spark StringType at the JVM boundary — no cast-back
projection, no lossy round-trip.

Changes

Rust

• native/proto/src/proto/operator.proto — added HashAggregate.use_large_data_types and ShuffleWriter.use_large_data_types flags.
• native/core/src/execution/planner.rs
  • promote_byte_group_key wraps each group-by expr in CastExpr(LargeUtf8|LargeBinary) when the flag is on.
  • align_shuffle_writer_input accepts the flag and promotes Utf8/Binary in expected_output_schema to their Large variants before invoking SchemaAlignExec, so no down-cast is ever inserted.
• native/core/src/execution/columnar_to_row.rs — maybe_cast_to_schema_type passes LargeUtf8/LargeBinary through unchanged (the row encoder's TypedArray::LargeString/LargeBinary variants already
  handle both offset widths).
• native/shuffle/src/schema_align.rs — new CastLargeStringToString/CastLargeBinaryToBinary actions with a byte-aware row-range splitter that rebuilds each chunk via StringBuilder/BinaryBuilder
  (arrow's cast_byte_container fails on sliced offsets, so we can't rely on cast_with_options alone).
• native/spark-expr/src/conversion_funcs/{cast,string}.rs — extended is_datafusion_spark_compatible to whitelist all four offset-width conversions (Utf8↔LargeUtf8, Binary↔LargeBinary), safety-net for
  any DF adapter that still constructs spark_expr::Cast for these types.

Scala / Java

• spark/.../CometConf.scala — added COMET_AGG_USE_LARGE_DATATYPES with explanatory doc.
• spark/.../operators.scala — wires the flag into both HashAggregate.newBuilder() sites via CometConf.COMET_AGG_USE_LARGE_DATATYPES.get(aggregate.conf).
• spark/.../CometNativeShuffleWriter.scala — wires the flag into ShuffleWriter.newBuilder().
• spark/.../comet/util/Utils.scala — maps LargeUtf8 → StringType / LargeBinary → BinaryType; adds LargeVarCharVector / LargeVarBinaryVector to the FFI export whitelist.

Test coverage

CometAggregateSuite gains one test that runs the same CUBE(9) + COUNT(DISTINCT) shape twice:

• useLargeDataTypes=false on 30K × 384B rows → asserts the "offset overflow" exception (existing behavior preserved).
• useLargeDataTypes=true on 12K × 170B rows → checkSparkAnswerAndOperator validates row-by-row equality against the Comet-disabled Spark baseline and re-asserts CometHashAggregateExec presence.

Residual limits

LargeUtf8 → Utf8 casts elsewhere in the pipeline still have to fit a single arrow-Utf8 batch (i32::MAX bytes). SchemaAlignExec splits by byte budget to stay under it; for extreme per-batch bytes this
can still fail — mitigated by lowering datafusion.execution.batch_size for the aggregate subtree if needed.

flowchart TD
    Scan["CometNativeScan parquet<br/><b>Utf8</b> (i32 offsets, ≤ 2 GiB per batch)"]
    Filter["CometFilter · CometProject · CometExpand<br/><b>Utf8</b> passthrough"]

    subgraph Agg["CometHashAggregate <i>(useLargeDataTypes=true)</i>"]
        direction TB
        Cast["CastExpr(Utf8 → LargeUtf8)<br/><i>promote_byte_group_key</i><br/><b>widen offsets i32 → i64</b>"]
        AggCore["AggregateExec (Partial / PartialMerge / Final)<br/>PhysicalGroupBy sees <b>LargeUtf8</b> keys<br/>→ dispatch ByteGroupValueBuilder&lt;i64&gt;<br/>buffer cap = i64::MAX (unbounded)"]
        Cast --> AggCore
    end

    subgraph Shuffle["CometExchange / CometNativeShuffleWriter"]
        direction TB
        Align["align_shuffle_writer_input<br/>promote expected_output_schema<br/><b>Utf8 → LargeUtf8</b> per <i>use_large_data_types</i>"]
        SchemaAlign["SchemaAlignExec<br/><b>passthrough</b> (no cast)"]
        IPC["Shuffle blocks encoded as <b>LargeUtf8</b>"]
        Align --> SchemaAlign --> IPC
    end

    subgraph JVMSide["JVM boundary"]
        direction TB
        Import["NativeUtil.importVector →<br/><b>LargeVarCharVector</b>"]
        TypeMap["Utils.fromArrowType<br/><b>LargeUtf8 → StringType</b>"]
        Import --> TypeMap
    end

    Downstream["Downstream CometHashAggregate<br/>re-promotion is a no-op<br/>(child schema already LargeUtf8)"]
    C2R["CometNativeColumnarToRow<br/>maybe_cast_to_schema_type:<br/><b>(LargeUtf8, Utf8) → passthrough</b><br/>TypedArray::LargeString → UnsafeRow"]
    Spark["Spark UnsafeRow (byte-oriented,<br/>offset width irrelevant)"]

    Scan --> Filter --> Agg --> Shuffle --> JVMSide --> Downstream --> C2R --> Spark

    style Cast fill:#fef3c7,stroke:#d97706
    style AggCore fill:#dbeafe,stroke:#2563eb
    style Align fill:#fef3c7,stroke:#d97706
    style SchemaAlign fill:#dcfce7,stroke:#16a34a
    style TypeMap fill:#fef3c7,stroke:#d97706
    style C2R fill:#dcfce7,stroke:#16a34a

Loading

Legend:

• 🟨 amber = new/modified code path (this PR)
• 🟦 blue = existing DF dispatch (already correct for LargeUtf8)
• 🟩 green = passthrough / no-op (LargeUtf8 preserved)

[mbutrovich]

We should check with upstream datafusion. @alamb mentioned that the group by stuff is being rewritten to be more efficient and avoid this giant interim batch in the first place. I'd like to make sure if we tackle this in Comet it's a generalizable solution and not special-casing a hack. Also at a glance the comments seem huge and redundant.

[andygrove]

The PR description says that this is an experimental feature and is disabled by default. Is it intentional that it is enabled by default here?

[andygrove]

These changes seem specific to the ignored test, but will impact all of the existing tests?

[comphead]

@mbutrovich this is more like a hack right now to support aggregation explosion cases, as a more long term solution we can consider StringView usage later on, or if DataFusion GROUP BY being rewritten, then the hack can be removed.

[andygrove]

I don't see any new tests for the functional changes in this PR. Could you add functional tests that use small amounts of data, just to test for correctness?

[comphead]

The entire CI is passed with useLargeDataTypes enabled.
on CI grade machines we cannot reproduce issue as we run out of memory earlier than hit the offset limit