Repository navigation
fix: bound the batch under the operator that expands a large string scalar - #7316
zhangfengcdt wants to merge 4 commits into
Conversation
A string scalar from a literal or a scalar subquery is expanded to the batch length by aggregates, the hash partitioner and string kernels. A 256 KiB string over an 8192-row batch needs 2 GiB of value bytes, past Arrow's offsets, and native execution panics where Spark succeeds. Add up the string and binary scalars while planning and lower the session batch size so their expansion stays within 1 GiB. Native sources read the batch size when they run, the range source now included, and the JVM-fed leaves slice wider batches to it.
github-actions
Bot
added
bug
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Lowering the task's batch size from the planned scalar bytes reaches every operator that reads the batch size from the TaskContext at execute time, and the arithmetic in scalar_batch_rows is safe for zero and oversized scalars. Three comments, on an operator that regrows batches, on literals that are never expanded, and on slicing at default settings. Nothing here was run.
| let rows = scalar_batch_rows(planner.scalar_bytes(), batch_size); | ||
| if rows < batch_size { | ||
| let state = exec_context.session_ctx.state_ref(); | ||
| let mut state = state.write(); | ||
| let config = state.config_mut(); | ||
| *config = std::mem::take(config).with_batch_size(rows); |
There was a problem hiding this comment.
FilterExec keeps its own batch size, 8192 by default in DataFusion 55.2 (FILTER_EXEC_DEFAULT_BATCH_SIZE), and coalesces its output up to it through LimitedBatchCoalescer without reading the session value at execute time. Comet builds it with FilterExec::try_new, so a filter under the expanding projection turns two capped 4096-row batches back into one 8192-row batch. I expect SELECT concat((SELECT max(s) FROM wide), CAST(k AS STRING)) FROM ids WHERE k >= 0 to still hit the offset overflow. #7280 passes the session batch size into FilterExecBuilder, but at planning time, before this cap lowers it, so it would not cover this either. Could the cap reach filters, for example by computing the scalar bytes from the protobuf plan before the planner runs, and could the fixture add a query with a WHERE?
There was a problem hiding this comment.
I reproduced the filter case on e9b8b0c3 with Spark 4.1.3. Summing length(concat(repeat('x', 262144), cast(id % 2 AS string))) over a one-partition 8,192-row range returned 2147491840 in Spark and in Comet without the filter. Adding WHERE id % 7 >= 0 kept all 8,192 rows and the native CometFilter plan failed with byte array offset overflow.
There is also a case where setting the cap before planning would still be insufficient:
SELECT sum(length(concat(repeat('x', 262144), cast(tile AS string))))
FROM (SELECT ntile(2) OVER (ORDER BY id) AS tile FROM range(0, 8192, 1, 1))Spark returned 2147491840. The native window emitted its whole partition and the aggregate above it failed with the same overflow. The window contains only integers, so the failure happens after the window, separate from the large window payload tracked in #7270. Could the expanding operator enforce the cap on the input it actually receives?
| } | ||
| } | ||
| }; | ||
| self.note_scalar_bytes(&scalar_value); |
There was a problem hiding this comment.
Every item of an In list is planned through this arm, so its bytes count toward scalar_bytes, although list items are only compared and never expanded to the batch length. CometInSet turns a Spark InSet into this In message, so col IN (<25,000 UUID strings>) is about 900 KB of literals and lowers the batch size to about 1,165 rows for every operator in the task, with no overflow risk to justify it. Could the In arm plan its items without counting them?
| baseline_metrics, | ||
| cast_options: CastOptions::default(), | ||
| cast_time, | ||
| slicer: BatchSlicer::new(batch_size), |
There was a problem hiding this comment.
The description says that with the default configuration nothing is sliced, but the slicer is installed with the session batch size whether or not the cap lowered it. Utils.coalesceBroadcastBatches deliberately merges a broadcast into few large batches, so a broadcast build side can arrive here as one batch of millions of rows and is now cut into 8192-row slices. HashJoinExec then rebuilds the build side with concat_batches, which copies every slice, where a single input batch took arrow's zero-copy path. Could the slicer be installed only when the cap actually lowered the batch size? I have not confirmed at runtime that a broadcast reaches ScanExec as one batch.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
I ran the new SQL fixture on Spark 4.1.3 against e9b8b0c3. It passed in normal mode. I’m requesting changes because the scan can now pull an unneeded JVM batch and make a previously successful LIMIT fail. The review also covers the nullable-batch slicing cost, the grouped test’s strict-mode failure, and scalar expansions that still overflow. I added the filter and window results to the existing batch-size discussion.
| let mut timer = self.baseline_metrics.elapsed_compute().timer(); | ||
| let mut scan_batch = self.scan.batch.try_lock().unwrap(); | ||
| let this = self.get_mut(); | ||
| if let Some(batch) = this.slicer.pending() { |
There was a problem hiding this comment.
A pending slice can be skipped by an upstream error. I compared the base and head native libraries on Spark 4.1.3 with the same LIMIT 3 plan, a two-row batch size, and a JVM reader that returns ten rows on its first load and throws on its second. The base returned all three rows after one load. This head returned two rows, then opened the next JVM batch and threw before returning row three. With a ten-row batch size the head succeeds too.
After emitting the first slice the scan clears its shared input slot. On the next executePlan, pull_input_batches refills that slot before poll_next can return the pending slice here. Could both JVM-fed scans defer the next upstream read until the current batch’s slices are exhausted? This also keeps a short-circuiting consumer from paying for an unnecessary read.
| return batch; | ||
| } | ||
| let head = batch.slice(0, self.max_rows); | ||
| self.rest = Some(batch.slice(self.max_rows, batch.num_rows() - self.max_rows)); |
There was a problem hiding this comment.
Could we use DataFusion’s BatchSplitStream here, as ExplodeExec does? Making a slice of the entire remainder rescans its validity bitmap to compute its null count. Every subsequent poll scans most of those same bits again. In a release-mode probe with the pinned Arrow 59.3.1, splitting 4,194,304 alternating-null integers into 512-row batches took 164–181 ms with this implementation and about 1 ms when slicing the original batch at successive offsets. This measures just the slicer. It happens even when the scalar cap legitimately requires smaller batches.
| { | ||
| let no_rows = RecordBatch::new_empty(Arc::new(Schema::empty())); | ||
| if let ColumnarValue::Scalar(value) = subquery.evaluate(&no_rows)? { | ||
| self.note_scalar_bytes(&value); |
There was a problem hiding this comment.
The size recorded here is the size of the subquery result. A parent can produce a larger scalar before it is expanded. With the fixture’s 256 KiB s, I ran this on Spark 4.1.3:
SELECT sum(length(concat(repeat((SELECT max(s) FROM wide), 8), cast(id AS string))))
FROM range(0, 8192, 1, 1)Spark returned 17179900842. With a native range and aggregates, this head failed with byte array offset overflow. Changing only spark.comet.batchSize to 256 returned the Spark result. Here repeat creates a 2 MiB scalar before concat broadcasts it. Could the expansion use the resulting scalar’s width, or evaluate this case in bounded chunks?
| -- count of the scalar in a grouped aggregate. The plan runs natively and panicked before the | ||
| -- fix; the answer-only mode is because the harness's missingInput check trips on a grouped | ||
| -- count of a scalar subquery whatever the string's size, which is a separate plan issue. | ||
| query spark_answer_only |
There was a problem hiding this comment.
I ran this fixture with ENABLE_COMET_STRICT_TESTING=true on Spark 4.1.3. It failed here with “Plan was fully native in Comet. Call checkSparkAnswerAndOperator instead.” Both aggregate stages ran natively. In the default mode, this directive would also pass if the aggregate fell back to Spark and stopped exercising the original bug. Could we keep an explicit assertion that the grouped count runs in CometHashAggregate and make the case pass strict mode? The missingInput problem mentioned above can be handled separately.
| self.baseline_metrics.record_output(*num_rows); | ||
| let maybe_batch = self.build_record_batch(columns, *num_rows); | ||
| Poll::Ready(Some(maybe_batch)) | ||
| this.baseline_metrics.record_output(*num_rows); |
There was a problem hiding this comment.
Could ScanStream and ShuffleScanStream record output_rows for each slice they emit? With a 10-row input and a four-row cap, this records ten on the first poll. If a downstream limit stops after that four-row slice, the remaining six never leave the scan, but its output metric still reports ten. A test that stops polling after the first slice would catch this.
…calar Rework after review. Capping the input batch size does not bound what reaches the expanding operator: a filter coalesces its output to its own batch size, a window emits its whole partition, and a scalar built from a subquery result can be wider than the result. The scan slicers also dropped a pending slice under LIMIT and double counted output_rows. A new CometBatchSplitExec, over DataFusion's BatchSplitStream, is placed right under a projection, filter, hash aggregate or hash shuffle writer when that operator's string and binary scalars, expanded to the session batch size, would exceed 1 GiB. The planner measures each column-free subtree of those types by evaluating it once on a one-row batch, so that repeat(<subquery>, n) counts at its final width. IN list members are compared and never expanded, so they are skipped. The session batch size, the scan and shuffle scan operators and the range source are unchanged.
zhangfengcdt
changed the title
uuid(0) and randstr(10, 0) are column-free string expressions, so the planner evaluated them once at plan time to measure their width, which drew the first value from their seeded generator. Execution then produced Spark's second value onward and the bit-for-bit tests failed. A subtree with a volatile expression yields a value per row rather than a scalar, so it is never measured: Comet's stateful expressions and any DataFusion scalar function declared volatile.
|
Thanks @comphead and @sunchao ! Agreed that capping the input batch size cannot bound what reaches the expanding operator, and it is not a optimal solution and fix. So I reworked the fix to enforce the bound at the operator, with the following changes:
Appreciate if you both could take another look and let me know if you'd refine it or go with other directions. Thanks! |
Which issue does this PR close?
Closes #7247.
Rationale for this change
A string scalar, from a scalar subquery or a literal, is expanded to the batch length by aggregates, the hash partitioner and string kernels. A 256 KiB string over the default 8192-row batch needs 2 GiB of value bytes, past Arrow's 32-bit
Utf8offsets, so native execution panics withoffset overflowor fails withbyte array offset overflow. Spark evaluates such scalars row by row and succeeds.The expansion happens in DataFusion (aggregate argument evaluation,
concat) as well as in Comet's shuffle partitioner, and a plan-time fallback cannot see a subquery's value. So the fix bounds the rows per batch that reach the expanding operator.The first version of this PR lowered the session batch size and sliced batches at the scans. Review showed that does not bound what reaches the operator: a filter coalesces its output to its own batch size, a window emits its whole partition, and a scalar built from a subquery result (
repeat(<subquery>, 2)) is wider than the result the planner measured. This version enforces the bound at the operator instead.What changes are included in this PR?
CometBatchSplitExecin the operators crate, a thin operator over DataFusion'sBatchSplitStream, which emits an input batch in slices of at mostbatch_sizerows. Slices share buffers.CometBatchSplitExecwhen that operator's string and binary scalars, expanded to the session batch size, would exceed 1 GiB. For the issue's string that is 4096 rows. Other operators are left alone.Utf8orBinarytype. The planner evaluates it once on a one-row batch, so a scalar built by an expression counts at its final width. Members of anINlist are compared and never expanded, so they are skipped. A subtree that contains a volatile expression, such asuuid(0), or a JVM UDF is not evaluated at plan time.ScanExec,ShuffleScanExecand therangesource are unchanged frommain.How are these changes tested?
operators/large_scalar_string.sqlwith nine queries that fail onmainwithbyte array offset overflow: the grouped count, the hash partitioning key,concatwith a subquery and with a literal, the literal over a nativerange, a filter under the projection, a window under the aggregate, andrepeatof the subquery result. All run inquerymode, so the plan is checked to be native.CometBatchSplitExecand for the scalar measurement: literals and columns, a scalar built by an expression,INlist members, and the planner wrapping the input only when the scalars need a smaller batch.CometSqlFileTestSuiteforoperators/,scalar_subqueryandshuffleon Spark 4.1, and the core and operators crate test suites.