Repository navigation
Conversation
|
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. |
| "because the cap is only reachable for very large per-partition group cardinalities; " + | ||
| "enable it when you see the offset-overflow error.") | ||
| .booleanConf | ||
| .createWithDefault(true) |
There was a problem hiding this comment.
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?
| super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "false") | ||
| super.sparkConf | ||
| .set(SQLConf.ANSI_ENABLED.key, "false") | ||
| .set("spark.memory.offHeap.enabled", "false") |
There was a problem hiding this comment.
These changes seem specific to the ignored test, but will impact all of the existing tests?
|
@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 |
| // The test to reproduce `offset overflow` for aggregation queries, when interim data | ||
| // get exploded 100x comparing to initial input size. | ||
| // It is not supposed to run on CI as the test requires significant RAM to succeed | ||
| ignore("CUBE(9) + COUNT(DISTINCT) wide Utf8 keys: useLargeDataTypes preserves correctness") { |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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
FWIW we are working on refactoring aggregates here Then I think we will be in position to update the allocation strategy internally (which currently uses large contiguous allocations for group values and aggregates |
Once we support non contiguous allocations, I think we'll be in the position to avoid blowing out the offsets for more than 2GB of string data |
|
(to be clear I doubt we'll have non contiguous allocations for DF 55 - maybe 56) |
Thanks @alamb for the input |
The problem is real:
(DataType::LargeUtf8, DataType::Utf8) | (DataType::LargeBinary, DataType::Binary) => {
Ok(Arc::clone(array))
}This is in The comment explains why the real cast is undesirable in the aggregate case (absolute offsets above Defaulting the config to
What does that cost on a normal aggregate? A Also, the description calls the config Shuffle wire format changes
Title and description "fix: experiments with large types for aggregated values" is not a merge-ready title, and the body starts with an How was the fix validated? The original repro needs more than 2 GiB of interned group keys, which is not something a unit test can do. |
|
FYI here is a ticket upstream in DataFusion that explains the root cause |
andygrove
added
enhancement
204a4dc to
7db0b70
Compare
github-actions
Bot
added
area:shuffle
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Summary
- Prior state and problem: String/binary aggregation keys can overflow DataFusion’s 32-bit offsets once accumulated values exceed 2 GiB.
- Design approach: Promote grouping expressions to
LargeUtf8/LargeBinary, then narrow aggregate output throughSchemaAlignExec. - Correctness / compatibility analysis: Found three introduced issues below. Bounded tests reproduce a final-aggregation spill failure for both types. Small-data null/empty-key cases pass. Checked relevant Spark sources across 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2 development.
- Key design decisions: Keeping external schemas unchanged preserves Spark’s type contract, but final aggregation’s internal spill schema must also match promoted keys.
- Implementation sketch: Scala serializes the new flag, the native planner adds promotion and output alignment, and conversion/shuffle/FFI helpers accept large-offset types.
- Behavioral changes worth calling out: The feature defaults to enabled. Affected aggregates lose their reported metrics and copy all key bytes during narrowing, including ordinary batches below the overflow limit.
- Suggested improvements: Align final-aggregate input with its spill schema, retain aggregate metrics through the wrapper, and use buffer-sharing casts for batches whose offsets fit.
Reviewed the entire diff from 64e98918ab11c41089416b18b262f56f83d41344 to 7db0b70d8f394e129dac613c39f3c2ab2cfd715d. The PR remains open and non-draft. Read existing discussion and threads. The previously discussed no-op cast is absent at this head.
Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-ffi-pr, review-comet-memory-pr, and review-comet-shuffle-pr.
Exact-head CI: All non-skipped checks passed, including Linux Rust tests, Spark 4.1 Comet suites, and TPC-H/TPC-DS. Spark SQL, Iceberg, macOS, and benchmark checks were skipped. The added large-data test was ignored.
Validation: A disposable harness used the exact SchemaAlignExec source with locked DataFusion 55.1.0 and Arrow 59.3.0. Its three existing tests passed. Additional bounded cases verified results, reproduced spill failures and missing metrics, and confirmed a projection-based correction. An optimized microbenchmark measured narrowing costs. No full Comet JVM/native rebuild, full Spark SQL/Iceberg run, or >2 GiB reproduction was performed. Benchmark timings cover conversion only. Project files remain unchanged.
| .map(|r| (r, format!("col_{idx}"))) | ||
| let raw = self.create_expr(expr, Arc::clone(&child_schema_ref))?; | ||
| let (wrapped, revert) = if use_large { | ||
| promote_byte_group_key(raw, child_schema)? |
There was a problem hiding this comment.
[P1] Keep the final aggregate’s input and spill schemas consistent with promoted keys. This promotion also runs in Final mode, while its child still produces Utf8/Binary. DataFusion 55.1.0 uses agg.input().schema() as the final aggregate’s state schema. When memory pressure triggers spilling, it stamps that small-offset schema onto promoted group arrays and fails with expected Utf8 but found LargeUtf8 (likewise for binary). The output alignment cannot fix an error inside the aggregate. With the default enabled, ordinary spilling GROUP BY/DISTINCT queries now fail. Promote the final input through a matching projection, or otherwise make its spill state schema agree with the promoted keys.
Evidence: Bounded reproduction in /tmp/comet-4791-validation/src/main.rs: native Final distinct aggregation over 80 batches of 256 rows, 10,000 distinct 128-byte keys, batch size 256, and a 1 MiB FairSpillPool. Without promotion, both Utf8 and Binary return 10,000 keys after five spills. With the PR’s grouping-expression promotion, both fail before returning rows with the corresponding small/large type mismatch. Promoting the input through ProjectionExec restores 10,000 results and five spills. DataFusion’s aggregate_hash_table/final_table.rs takes the state schema from agg.input().schema(), and common.rs::take_state_batch constructs the failing batch.
| }) | ||
| .collect(); | ||
| let target_schema: SchemaRef = Arc::new(Schema::new(target_fields)); | ||
| SchemaAlignExec::try_new_or_passthrough(aggregate, &target_schema) |
There was a problem hiding this comment.
[P2] Preserve aggregate metrics when installing this wrapper. SchemaAlignExec has no metrics() implementation, but it replaces AggregateExec as the native plan passed to SparkPlan::new. to_native_metric_node consequently reads no metrics and never visits the wrapped aggregate. Every promoted string/binary aggregation loses output-row, execution-time, and spill reporting, including the spill information propagated into Spark task metrics. Forward the aggregate’s metrics through the wrapper or explicitly retain it as the metric source.
Evidence: The exact-source harness executes string and binary aggregates containing duplicates, nulls, and empty keys. Both return four groups, and the underlying AggregateExec reports output_rows=4 and nonzero elapsed_compute, while the wrapper returns metrics=None. SparkPlan::new leaves additional_native_plans empty, and native/core/src/execution/metrics/utils.rs::to_native_metric_node only reads the root metrics in that case.
There was a problem hiding this comment.
I can confirm this end to end with a full build. SELECT k, sum(v) FROM t GROUP BY k on a string key shows the partial CometHashAggregateExec with output_rows=0 and elapsed_compute=0 when the flag is on, against 20000 rows and about 50 ms with it off. The final aggregate's elapsed_compute drops from about 10 ms to 17 µs. When fixing it, note that to_native_metric_node skips output_rows from additional_native_plans, so registering the AggregateExec through SparkPlan::new_with_additional also needs the wrapper to record its own output_rows. A test asserting these metrics on a string-keyed aggregate would catch this. The one in CometExecSuite uses a numeric key.
| // the underlying Vec never has to grow-and-memcpy while we replay rows. | ||
| let offsets = arr.value_offsets(); | ||
| let values_bytes = (offsets[arr.len()] - offsets[0]) as usize; | ||
| let mut builder = StringBuilder::with_capacity(arr.len(), values_bytes); |
There was a problem hiding this comment.
[P2] Avoid copying value buffers when the offsets already fit. This builder path runs for every promoted string key, with the equivalent copy for binary, even when no splitting or rebasing is needed. Since promotion defaults to enabled, each partial/final aggregation adds a full allocation and copy of its emitted key bytes. Arrow can narrow ordinary batches while sharing those bytes. Use that fast path when offsets fit, reserving rebuilding for overflow slices, or rebase offsets against a shared value-buffer slice.
Evidence: The exact SchemaAlignExec source copies a 4 MiB value buffer for 8,192 keys of 512 bytes, verified by buffer-pointer comparison. Arrow 59.3.0’s cast shares it and produces identical values. An optimized benchmark of the PR’s exact builder branch measured median conversion times of 168.1 µs versus 8.4 µs for that batch, and 7.81 ms versus 8.3 µs for 8,192 keys of 8 KiB. These are narrowing-only measurements, not whole-query timings. Reproduction: /tmp/comet-4791-cast-bench/src/main.rs.
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
Keeping the Large types inside the aggregate is a good call. Shuffle blocks and the FFI boundary stay Utf8, which settles the rolling-upgrade question from my earlier review, and the TPC-H and TPC-DS result checks pass with the flag on. The title and body still describe an earlier design, though. The body mentions spark.comet.exec.useLargeDataTypes, a ShuffleWriter.use_large_data_types flag, changes to align_shuffle_writer_input and CometNativeShuffleWriter.scala, and Large types carried through shuffle with no cast-back. None of that is in the diff. Could the body be rewritten for the current head, along with the title change I suggested earlier?
| } | ||
|
|
||
| #[derive(Debug, Clone)] | ||
| enum ColumnAction { |
There was a problem hiding this comment.
It looks like the rebase brought back the ColumnAction enum that #5138 removed from this file. #5138 routed SchemaAlignStream through cast_and_stamp_schema, so every batch is checked against the target schema and a failed cast names the operator and the column path. Here the choice is made once at plan time. Passthrough columns are no longer checked per batch, and the Cast arm returns a bare arrow error. This operator sits in front of every native shuffle writer, so that applies with the flag off too. Could the Large-to-small columns be handled first, with the result then handed to cast_and_stamp_schema?
| } | ||
| true | ||
| match (actual_field.data_type(), expected_field.data_type()) { | ||
| (DataType::LargeUtf8, DataType::Utf8) => ColumnAction::CastLargeStringToString, |
There was a problem hiding this comment.
With the flag on, reusing SchemaAlignExec above the aggregate hits the warning just above and logs ShuffleWriter input schema mismatch on col[0] 'col_0': child produced LargeUtf8, catalyst declared Utf8. Inserting a cast; please file the upstream function bug at .../issues/4515 on each executor. I saw it in my local runs. There is no shuffle writer or upstream bug involved here, and the module doc says this operator is enclosed by shuffle on purpose. Would a small operator dedicated to the cast-back be cleaner? It could record its own metrics, which would also help with the metrics problem sunchao raised in planner.rs.
| "because the cap is only reachable for very large per-partition group cardinalities; " + | ||
| "enable it when you see the offset-overflow error.") | ||
| .booleanConf | ||
| .createWithDefault(true) |
There was a problem hiding this comment.
I asked in July whether defaulting this to true was intentional, and the doc string still says it defaults to false. I measured it on a release build. SELECT count(*), sum(s) FROM (SELECT k, sum(v) s FROM t GROUP BY k) over 8M rows with 4M distinct 21-byte keys took 494 ms at the median with the flag on and 450 ms with it off, about 10% slower over five alternating runs. Low-cardinality and 200-byte keys were within noise. My laptop was busy, so treat the numbers as rough. The overflow needs more than 2 GiB of distinct key bytes in one task, and apache/datafusion#24704 tracks the real fix. Could this default to false, with the tuning guide pointing at it for the offset overflow error? Either way the doc needs updating. It describes the cast-back as a Projection and says the overhead is O(rows), but as sunchao pointed out, the cast-back copies every key byte.
| .createWithDefault(false) | ||
|
|
||
| val COMET_AGG_USE_LARGE_DATATYPES: ConfigEntry[Boolean] = | ||
| conf(s"$COMET_EXEC_CONFIG_PREFIX.aggregation.useLargeDataTypes") |
There was a problem hiding this comment.
The other aggregate keys use spark.comet.exec.aggregate.*, and config_conventions.md asks for feature flags to end in .enabled. This key will be covered by the versioning policy once it ships, so could we settle on something like spark.comet.exec.aggregate.largeGroupKeys.enabled now?
| // Native shuffle may dictionary-encode string/binary columns for efficiency, | ||
| // but downstream DataFusion operators expect the value types declared in the | ||
| // schema (e.g. Utf8, not Dictionary<Int32, Utf8>). | ||
| // Coerce each decoded column to the catalyst-declared type: |
There was a problem hiding this comment.
Since the aggregate now casts back before anything above it sees the batch, I don't think this feature can put LargeUtf8 into a shuffle block. If one did, the remote read path would reject it, because remote_schema.rs treats LargeUtf8 against Utf8 as an incompatible type. On the local path, ShuffleScanStream::poll_next already reconciles every column through cast_and_stamp_schema, so the extra cast here duplicates that. Could this go back to unpacking dictionaries only? Returning an error instead of the old expect is a good change and worth keeping.
| valueVector match { | ||
| case v if isSupportedFieldVector(v) => | ||
| v.asInstanceOf[FieldVector] | ||
| // Accepted here but left out of isSupportedFieldVector, which isArrowBacked uses to keep |
There was a problem hiding this comment.
Which path needs getFieldVector to accept LargeVarCharVector now that the aggregate casts back natively? As far as I can tell, the only Large vectors on the JVM side come from PyArrow UDFs returning large_string. So this changes native C2R export and broadcast serialization for that path without a test. It also makes the isSupportedFieldVector doc, the isArrowBacked comment and the UtilsSuite test comment wrong, since all three assume getFieldVector rejects these vectors. Could this be dropped here, or moved to its own PR with a PyArrow test? The same question applies to the LargeUtf8 passthrough at columnar_to_row.rs:1041.
| matches!(to_type, DataType::Binary | DataType::LargeUtf8) | ||
| } | ||
|
|
||
| pub(crate) fn is_df_cast_from_large_string_spark_compatible(to_type: &DataType) -> bool { |
There was a problem hiding this comment.
I can't find a way to reach these new arms, or the matching ones in cast.rs. types.proto has no Large type ids, so the serde never emits a Cast to or from one. The Parquet schema adapter's casts already go through the is_adapting_schema branch in cast_array. The comment also says SchemaAlignExec pre-splits arrays before they get here, but SchemaAlignExec builds its arrays directly and never calls this cast. Could these arms be removed from this PR?
| // The test to reproduce `offset overflow` for aggregation queries, when interim data | ||
| // get exploded 100x comparing to initial input size. | ||
| // It is not supposed to run on CI as the test requires significant RAM to succeed | ||
| ignore("CUBE(9) + COUNT(DISTINCT) wide Utf8 keys: useLargeDataTypes preserves correctness") { |
There was a problem hiding this comment.
This test is ignored, so CI never runs it. The sparkConf comment at line 60 still describes the off-heap and memory pool settings that 7252adb removed. This test now sets the pool through withSQLConf, which that comment says won't take effect. You mentioned in July that CI passes with the flag on. That covers the common path, but nothing exercises the split in compute_row_ranges or checks the metrics. Could the byte cap be a parameter so a Rust unit test can split a small batch? And could this test and the stale comment be replaced with small tests that run? If the default moves to false, those tests would also need to turn the flag on explicitly.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Summary
- Prior state and problem: High-cardinality string/binary grouping can exceed DataFusion’s 2 GiB group-key buffer limit.
- Design approach: Promote keys to
LargeUtf8/LargeBinaryinside aggregation, then narrow output throughSchemaAlignExec. - Correctness / compatibility analysis: Small null, empty-key, and duplicate-key cases pass. The existing P1 spill failure remains reproducible for both types with 10,000 keys and a 1 MiB pool. Compared relevant Spark sources across 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2 development.
- Key design decisions: Keeping external schemas unchanged preserves Spark’s type contract. Reusing the shuffle alignment wrapper introduces the already-reported metrics and abstraction concerns.
- Implementation sketch: Scala serializes the configuration flag. The native planner promotes grouping expressions and aligns aggregate output. Shuffle, cast, and FFI helpers gain large-offset handling.
- Behavioral changes worth calling out: Promotion defaults to enabled. Independently reproduced the existing missing-metrics concern and confirmed that narrowing copies a 4 MiB value buffer which Arrow’s ordinary cast shares.
- Suggested improvements: Resolve the existing spill-schema, aggregate-metrics, and unnecessary-copy threads before merging. No additional introduced P1/P2 issues found within this review.
Reviewed the entire 11-file diff from 64e98918ab11c41089416b18b262f56f83d41344 to 7db0b70d8f394e129dac613c39f3c2ab2cfd715d, including surrounding code and existing discussion. The PR remains non-draft. Existing findings are not duplicated below.
Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-ffi-pr, review-comet-memory-pr, and review-comet-shuffle-pr.
Exact-head CI: All executed checks passed, including Linux Rust tests, Spark 4.1 Comet suites, and TPC-H/TPC-DS result checks. Spark SQL, Iceberg, PyArrow UDF, macOS, and benchmark jobs were skipped.
Validation: Reran the disposable harness against the exact SchemaAlignExec source with DataFusion 55.1.0 and Arrow 59.3.0. Its 19 tests passed, including additional splitter tests using a reduced 32-byte cap to exercise multiple columns, sliced offsets, nulls, empty batches, and oversized-row rejection. No full Comet JVM/native rebuild, full Spark SQL/Iceberg run, or actual >2 GiB validation was performed. Project files and GitHub state were unchanged.
|
Final aggregates fail on spill with A job running this branch (rebased on Cause:
Suggested fix: for DataFusion Two planner tests cover it:
Both tests pass with the fix. With Workaround until then: Two more things on this branch:
|
27a9ba6 to
2f84cc0
Compare
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Summary
- Prior state and problem: High-cardinality string/binary grouping can exceed DataFusion’s 2 GiB group-key buffer limit.
- Design approach: Promote grouping keys to
LargeUtf8/LargeBinaryinside aggregation, then narrow output throughSchemaAlignExec. - Correctness / compatibility analysis: The previously reported final-aggregation spill failure is fixed by promoting input columns. Bounded tests passed for both key types, mixed keys, nulls, empty values, and sum/count states through spilling. Compared relevant Spark sources across 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. No additional introduced P1/P2 issues found within this review.
- Key design decisions: Keeping large-offset types inside aggregation preserves external schemas. Reusing the shuffle alignment wrapper still leaves the existing aggregate-metrics concern unresolved: the wrapped aggregate reports metrics, but
SchemaAlignExecreturnsNone. - Implementation sketch: Scala serializes the configuration flag. The native planner promotes partial grouping expressions or final input columns, then splits and narrows output. Cast, shuffle-read, and FFI helpers also gain large-offset handling.
- Behavioral changes worth calling out: Promotion defaults to enabled. The existing copying concern remains reproducible: narrowing an ordinary 4 MiB value buffer allocates and copies it, while Arrow’s cast shares the buffer and returns identical values.
- Suggested improvements: Resolve the existing metrics and unnecessary-copy findings before merging. Preserve aggregate metrics through the wrapper and share value buffers when narrowing permits it.
Reviewed the entire 11-file diff from ce455f32d948355e073638e81009ecc3e5dea349 to 2f84cc0c2bb9b5dea5ae5b97294bfeb1adcd9484, including surrounding code and existing discussion. The PR remains open and non-draft. Existing findings are not duplicated below.
Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-ffi-pr, review-comet-memory-pr, and review-comet-shuffle-pr.
Exact-head CI: All executed checks passed, including Linux Rust tests, both new planner regression tests, Spark 4.1 Comet suites, and TPC-H/TPC-DS result checks. Spark SQL, Iceberg, PyArrow UDF, macOS, and benchmark jobs were skipped.
Validation: The disposable harness at /tmp/comet-4791-2f84-validation-f9wv4e9l uses exact-head sources and planner helpers with locked DataFusion 55.1.0 and Arrow 59.3.0. Its 20 tests and an additional columnar-to-row equivalence test passed. Mixed-key aggregation returned all 9,970 expected groups after five spills with promotion enabled and disabled. Splitter tests used a reduced 32-byte cap. Separate checks reproduced the existing missing metrics and value-buffer copying.
Limits: No full local Comet JVM/native rebuild, Spark SQL/Iceberg run, or actual >2 GiB reproduction was performed. The large CUBE test remains ignored. Project files and GitHub state were unchanged.
|
Hi @comphead |
DataFusion's final hash aggregation spills against its input schema (`AggregateHashTable<FinalMarker>` takes `agg.input().schema()`), so casting a Final group-by expression to LargeUtf8/LargeBinary made the first spill fail with "column types must match schema types, expected Binary but found LargeBinary at column index 0". For Final aggregates, cast the key columns in a pass-through projection below the aggregate instead, so its input, group values and spill files share one type. Partial and PartialMerge keep the expression cast. Adds a plan-shape test and a spill test. Both fail with the original error when Final aggregates take the expression cast.
apache#6066 removed `spark.comet.exec.onHeap.memoryPool` (`COMET_ONHEAP_MEMORY_POOL_TYPE`). On-heap mode now always runs on an unbounded native memory pool, so the setting was redundant, and after rebasing onto main its reference no longer compiled.
2f84cc0 to
47ac89f
Compare
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
I compared this head with branch-1.2, which has the same DataFusion 55.2 aggregate planning, so the key promotion, the Final input projection and the revert are the whole difference a user upgrading from 1.2 would see. Moving the Final promotion to the aggregate's input is a clean fix for the spill failure, and the two planner tests pin it down well. With the flag on by default, though, sort aggregates now pick up an extra sort, and CI is red on the release's spill tests because the aggregate metrics are lost. The comments below cover those, plus the points from my earlier comments that are still open at this head.
1. native/core/src/execution/planner/aggregate.rs:463
With the flag on, every sort aggregate with a string or binary key gets a SortExec over its output. CometSortAggregateExec sets ordered_by_grouping_keys, and this block adds a sort unless the output reports the grouping order. The output is now SchemaAlignExec, which builds its properties from EquivalenceProperties::new(target_schema). It never reports an order, so the sort is always added. On branch-1.2, a sort aggregate over a native sort in the same plan streams with no sort above it. sort_aggregate_float_key_streams_over_native_sort checks that for a double key.
Binary keys lose more than that. DataFusion's CastExpr::check_bigger_cast treats (Utf8, LargeUtf8) as order-preserving but has no (Binary, LargeBinary) entry. So a binary-keyed sort aggregate also stops seeing its sorted input, and builds a hash table instead of streaming.
This path is common. Spark plans SortAggregateExec whenever an aggregate buffer holds a string or binary, since UnsafeRow.isMutable is false for them. So SELECT k, max(s) FROM t GROUP BY k runs both stages this way. Each stage pays for one more full sort, and one more spillable consumer shares the task's fair pool.
The hidden order also reaches aggregates stacked on a promoted one in the same native plan. The PartialMerge in the #6880 tests no longer switches to Final at line 353, and runs as SingleHashAggregateStream. The Final in the ordered #6254 test runs as FinalHashAggregateStream instead of OrderedFinalAggregateStream. Both still return the right rows, but those tests no longer cover the streams they were written for.
Could the revert keep its child's orderings? Narrowing the offsets doesn't change the order. with_new_schema rejects the changed types, so the orderings would need rebuilding, for example with EquivalenceProperties::new_with_orderings. maintains_input_order should also report true, so keeps_source_order still sees through the operator. Separately, a streaming aggregate over input sorted on all of its keys holds one group at a time and can't reach the cap, so the planner could skip the promotion there.
I haven't run this test. From the code, I'd expect it to pass on branch-1.2 without the use_large_data_types line, and to fail here with SortExec at the root:
#[test]
fn sort_aggregate_byte_key_streams_over_native_sort() {
// STRING and BYTES keys, in both stages of a sort aggregate.
for type_id in [7, 8] {
for mode in [
spark_operator::AggregateMode::Partial,
spark_operator::AggregateMode::Final,
] {
let key_type = spark_expression::DataType {
type_id,
type_info: None,
};
let key = Expr {
expr_struct: Some(Bound(spark_expression::BoundReference {
index: 0,
datatype: Some(key_type.clone()),
})),
..Default::default()
};
let sort = Operator {
children: vec![Operator {
op_struct: Some(OpStruct::Scan(spark_operator::Scan {
fields: vec![key_type],
source: String::new(),
})),
..Default::default()
}],
op_struct: Some(OpStruct::Sort(spark_operator::Sort {
sort_orders: vec![Expr {
expr_struct: Some(SortOrder(Box::new(spark_expression::SortOrder {
child: Some(Box::new(key.clone())),
direction: spark_expression::SortDirection::Ascending as i32,
null_ordering: spark_expression::NullOrdering::NullsFirst as i32,
}))),
..Default::default()
}],
..Default::default()
})),
..Default::default()
};
let op = Operator {
plan_id: 1,
children: vec![sort],
op_struct: Some(OpStruct::HashAgg(spark_operator::HashAggregate {
grouping_exprs: vec![key],
mode: mode as i32,
ordered_by_grouping_keys: true,
use_large_data_types: true,
..Default::default()
})),
..Default::default()
};
let (_, _, planned) = PhysicalPlanner::default()
.create_plan(&op, &mut vec![], 1)
.unwrap();
assert_eq!(
"AggregateExec",
planned.native_plan.name(),
"type_id {type_id}, {mode:?}"
);
}
}
}2. native/shuffle/src/schema_align.rs:199
CI on this head fails five tests in CometAggregateSuite on Spark 4.1 [exec]: both #6254 tests, the #6857 spill test, and both #6880 tests. Each one fails on spill_count > 0, with The Final aggregate did not spill or The PartialMerge aggregate did not spill, after its result check has passed. These are branch-1.2's spill regression tests. They landed on main in October, after the last green run here.
They fail because SchemaAlignExec is now the aggregate's native root and has no metrics(). So to_native_metric_node never reads the AggregateExec, as sunchao found. Users lose the operator's rows, time, spill count and peak memory in the SQL tab. The spilled bytes also stop reaching taskMetrics().incDiskBytesSpilled through CometMetricNode.reportSpillMetrics, so the stage's spill totals drop too.
Could the fix make these five tests pass with the flag on, rather than setting useLargeDataTypes=false in them? As I said in my earlier reply, to_native_metric_node skips output_rows from additional_native_plans. So the revert operator needs to report the aggregate's output rows itself. A test that checks these metrics on a string-keyed and a binary-keyed aggregate would keep this from regressing.
3. native/shuffle/src/schema_align.rs:326
sunchao's measurement of this copy still holds at this head. Every output batch of every promoted aggregate has its key bytes copied into a StringBuilder or BinaryBuilder, even when the batch is far below 2 GiB.
Could this rebase instead of copying? When offsets[len] - offsets[0] fits in an i32, the narrowed array can take o - offsets[0] as its offsets and share values.slice_with_length(offsets[0], total) as its values. That is O(rows). It also handles the slices DataFusion emits from its one large output batch, which is where arrow's cast fails, because it converts the whole offsets buffer behind the slice. Only a slice over 2 GiB would then need the split and a copy.
A micro-benchmark for this operator would show where this lands. So would the GROUP BY timing from my earlier comment, rerun with the flag on and off.
4. spark/src/main/scala/org/apache/comet/CometConf.scala:828
My earlier comments on this line still apply at this head. The default is true, but the doc says Defaults to false, describes the cast-back as a Projection, and puts the overhead at O(rows). With the default on, every string- or binary-keyed native aggregate pays for the promotion, and sort aggregates pick up an extra sort. That is all to avoid an overflow that needs more than 2 GiB of distinct key bytes in one task.
Unless the copy and the sort are fixed and measured, could this default to false for now? The tuning guide could then point at it for the offset overflow error.
The name is still open as well. The sibling key is spark.comet.exec.aggregate.skipPartial.enabled, so something like spark.comet.exec.aggregate.largeGroupKeys.enabled would fit. It is easier to settle that before the key ships.
5. native/core/src/execution/planner.rs:4819
The existing native aggregate tests build HashAggregate with ..Default::default(). That leaves use_large_data_types at the proto default of false, while CometConf defaults it to true. So apart from the two new tests, no Rust aggregate test runs the path users get.
Could the sort-aggregate, PartialMerge and spill tests in this module run with the flag both on and off? A string-keyed copy of sort_aggregate_float_key_streams_over_native_sort with the flag on would have caught the extra sort above sort aggregates.
6. spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:4865
This test is still ignored, and the comment above sparkConf (lines 60-65) still describes off-heap settings the suite no longer makes. So nothing that runs in CI reaches the split in compute_row_ranges, and the split has no unit test either.
Could the byte cap become a parameter, so a Rust test can split a small batch? Sliced input, nulls, two shrinking columns and the oversized-row error are the cases worth covering. A small Scala test could cover the rest. It would run string and binary keys with nulls and empty values through Partial, Final, PartialMerge and a sort aggregate, with the flag on, and check both results and metrics.
Since this changes every string-keyed aggregate by default, could the PR also get a run-spark-4.1-tests run? The Spark SQL suites were skipped at this head.
7. General
A few points from my earlier comments are still open at this head:
SchemaAlignExeclogsShuffleWriter input schema mismatch on col[0] 'col_0': child produced LargeUtf8, catalyst declared Utf8 ... please file the upstream function bugfor aggregate output, as the CI logs show. Its module doc still says the operator is enclosed by shuffle.- The cast in
shuffle_scan.rsrepeats whatcast_and_stamp_schemadoes inpoll_next. Its error also drops the column path thatcast_and_stamp_schemareports. - This feature still doesn't reach the
getFieldVectorchange, theLargeUtf8pass-through incolumnar_to_row.rs, or the new arms incast.rsandstring.rs. TheisArrowBackedcomment inUtils.scalastill saysgetFieldVectorrejects Large vectors, which is no longer true. - The title and body still describe the shuffle-level design.
Could these be addressed in this round? Two things from that round are resolved, which is great to see. Every column now goes through cast_and_stamp_schema in SchemaAlignExec, and the Final spill failure is fixed, with good planner tests.
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-taskByteGroupValueBuilder<i32>byte-buffer cap (i32::MAX = 2 147 483 647),surfacing as:
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(defaulttrue) that promotes Utf8/Binary group-by expressions to LargeUtf8/LargeBinary before the aggregate, routing DataFusion toByteGroupValueBuilder<i64>(i64 offsets, effectively unbounded buffer). The Large variant is preserved end-to-end through shuffle and mapped back to SparkStringTypeat the JVM boundary — no cast-backprojection, no lossy round-trip.
Changes
Rust
native/proto/src/proto/operator.proto— addedHashAggregate.use_large_data_typesandShuffleWriter.use_large_data_typesflags.native/core/src/execution/planner.rspromote_byte_group_keywraps each group-by expr inCastExpr(LargeUtf8|LargeBinary)when the flag is on.align_shuffle_writer_inputaccepts the flag and promotesUtf8/Binaryinexpected_output_schemato their Large variants before invokingSchemaAlignExec, so no down-cast is ever inserted.native/core/src/execution/columnar_to_row.rs—maybe_cast_to_schema_typepassesLargeUtf8/LargeBinarythrough unchanged (the row encoder'sTypedArray::LargeString/LargeBinaryvariants alreadyhandle both offset widths).
native/shuffle/src/schema_align.rs— newCastLargeStringToString/CastLargeBinaryToBinaryactions with a byte-aware row-range splitter that rebuilds each chunk viaStringBuilder/BinaryBuilder(arrow's
cast_byte_containerfails on sliced offsets, so we can't rely oncast_with_optionsalone).native/spark-expr/src/conversion_funcs/{cast,string}.rs— extendedis_datafusion_spark_compatibleto whitelist all four offset-width conversions (Utf8↔LargeUtf8,Binary↔LargeBinary), safety-net forany DF adapter that still constructs
spark_expr::Castfor these types.Scala / Java
spark/.../CometConf.scala— addedCOMET_AGG_USE_LARGE_DATATYPESwith explanatory doc.spark/.../operators.scala— wires the flag into bothHashAggregate.newBuilder()sites viaCometConf.COMET_AGG_USE_LARGE_DATATYPES.get(aggregate.conf).spark/.../CometNativeShuffleWriter.scala— wires the flag intoShuffleWriter.newBuilder().spark/.../comet/util/Utils.scala— mapsLargeUtf8 → StringType/LargeBinary → BinaryType; addsLargeVarCharVector/LargeVarBinaryVectorto the FFI export whitelist.Test coverage
CometAggregateSuitegains one test that runs the same CUBE(9) +COUNT(DISTINCT)shape twice:useLargeDataTypes=falseon 30K × 384B rows → asserts the"offset overflow"exception (existing behavior preserved).useLargeDataTypes=trueon 12K × 170B rows →checkSparkAnswerAndOperatorvalidates row-by-row equality against the Comet-disabled Spark baseline and re-assertsCometHashAggregateExecpresence.Residual limits
LargeUtf8 → Utf8casts elsewhere in the pipeline still have to fit a single arrow-Utf8 batch (i32::MAXbytes).SchemaAlignExecsplits by byte budget to stay under it; for extreme per-batch bytes thiscan still fail — mitigated by lowering
datafusion.execution.batch_sizefor 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<i64><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:#16a34aLegend: