perf: Reduce peak memory usage of NLJ by 2X using segmented batch layout - #25371
2010YOUY01 wants to merge 3 commits into
Conversation
Add one NLJ query at the existing 10-million-row memory-validation scale. Keep the 80 MB build input on the left and cap baseline-adjusted RSS at 90 MB, allowing overhead without a second build-side copy. The test fails with concatenation and passes with segmented storage.
github-actions
Bot
added
core
|
|
||
| /// One consecutive group of output rows gathered from the same segment. | ||
| #[derive(Debug)] | ||
| struct Single { |
There was a problem hiding this comment.
This inner encoding seems inefficient, but end-to-end NLJ microbench execution time is not influenced (see PR description)
And the modular design make it easy to optimize it in the future if necessary.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25371 +/- ##
==========================================
+ Coverage 81.80% 82.29% +0.48%
==========================================
Files 1130 1138 +8
Lines 417754 430594 +12840
Branches 417754 430594 +12840
==========================================
+ Hits 341754 354346 +12592
+ Misses 55875 54799 -1076
- Partials 20125 21449 +1324 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
jayzhan211
left a comment
jayzhan211
left a comment
There was a problem hiding this comment.
Thanks @2010YOUY01 , here are some suggestions
| batch.num_columns() | ||
| ) | ||
| })?; | ||
| take( |
There was a problem hiding this comment.
Many small build-side batches get 1.2–1.4x slower
When the build side is many small batches and the probe side is small, one range join covers up to batch_size segments. That becomes one take per segment per column plus a concat, where main did a single take. Measured with a build-side MemTable of 1-row batches and a 1-row probe side: 200K rows count(*) 48–61 ms → 66–73 ms; 1M rows projecting a Utf8 column 508–559 ms → 622–788 ms. The nlj bench suite is unchanged.
Fine to handle in a follow-up. Merging small segments while buffering (e.g. BatchCoalescer once a batch is under batch_size / 2) keeps the memory saving and bounds the extra copy to about one batch.
There was a problem hiding this comment.
The existing convention is: each operator promise to output at batch_size, so they can assume input don't have many small batches, so they don't have to handle input coalescing repeatedly.
I find it causes confusion very often, I'll try to get it documented somewhere.
There was a problem hiding this comment.
Is this based on the fact that we usually use 8192 as the default batch size?
There was a problem hiding this comment.
Yes, even if it's tuned for different workloads, batch_size is always reasonably large for better vectorization.
There was a problem hiding this comment.
Let's add documentation for this!
| /// Returns an execution error if a batch has a different schema or the | ||
| /// total row count overflows. | ||
| pub(crate) fn new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Result<Self> { | ||
| if batches.iter().any(|batch| batch.schema() != schema) { |
There was a problem hiding this comment.
Schema equality check rejects batches that main accepts
batch.schema() != schema also compares field metadata, nullability, and names; concat_batches on main only needed matching column types. A build-side child whose batches carry extra field metadata now fails with LogicalBatch input batches must have the same schema, while main returns the join result. I reproduced this with a unit test: a TestMemoryExec with a plain schema whose batch has field metadata, joined via multi_partitioned_join_collect. It passes on main and errors on this branch.
Fix: check only what take/concat need:
if batches.iter().any(|batch| {
batch.num_columns() != schema.fields().len()
|| batch
.columns()
.iter()
.zip(schema.fields())
.any(|(array, field)| array.data_type() != field.data_type())
}) {
return exec_err!("LogicalBatch input batches must have the same column types");
}There was a problem hiding this comment.
I'm not sure about this, currently on Schema or RecordBatch doc, it doesn't specify the invariants for metadata. So I think now it's better to do it more conservatively (include the metadata for = comparison), and later if we find it too strict, we can always relax it.
Additionally, perhaps it should be better to deprecate derived = comparisons on Schema, and add explicit APIs for equal including/excluding metadata, so it's easier to use it correctly.
There was a problem hiding this comment.
A test that failed with current check
/// Batches from an arbitrary `ExecutionPlan` may carry a schema that is
/// stricter than the plan's advertised schema (non-nullable fields, extra
/// field metadata); DataFusion does not adapt them and the previous
/// `concat_batches` path never required an exact match. Buffering the
/// build side must keep accepting such batches.
#[tokio::test]
async fn build_side_batches_with_stricter_schema() -> Result<()> {
let plan_schema =
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
let batch_schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false).with_metadata(HashMap::from([(
"PARQUET:field_id".to_string(),
"1".to_string(),
)])),
]));
let batch = RecordBatch::try_new(
batch_schema,
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)?;
let left: Arc<dyn ExecutionPlan> =
Arc::new(MockExec::new(vec![Ok(batch)], plan_schema));
let right = build_table(
("b", &vec![10, 20]),
("c", &vec![1, 2]),
("d", &vec![0, 0]),
None,
vec![],
);
let (_, batches, _) = join_collect(
left,
right,
&JoinType::Inner,
None,
Arc::new(TaskContext::default()),
)
.await?;
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(rows, 6);
Ok(())
}This is only a mock test, so whether we need it comes down to whether the case can actually happen in production. If it can, relaxing the check seems worth doing.
There was a problem hiding this comment.
It's reasonable to allow not 100% equivalent metadata for different batches, the issue is we don't have a project-level spec for how to handle batch-level metadata, and don't know how to test that e2e (this requires use cases that depend on metadata)
For instance, what should we do if we want to concat batches with conflicting metadata keys:
-- Conflicting metadata keys across batch
batch1 metadata:
key: config, value: foo
batch2 metadata:
key: config, value: bar
or after left_batch join right_batch, how to propagate metadata to output batch metadata.
The existing implementation is likely quite random across the codebase, so I'd prefer to keep it stricter. If someone have application with such metadata, they must find the existing random behavior, and next we can carry out the plan to agree on spec, and figure out how to get it tested, etc.
There was a problem hiding this comment.
Makes sense — let's keep it strict for now and revisit once there's a real use case that depends on batch-level metadata.
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
Thanks for working on this. The segmented LogicalBatch approach looks like a solid way to avoid the build-side concat_batches memory amplification while keeping the NLJ implementation reasonably clean.
I only have one non-blocking testing suggestion below.
| @@ -3055,8 +3057,7 @@ impl NestedLoopJoinStream { | |||
| Vec::with_capacity(filter.column_indices().len()); | |||
There was a problem hiding this comment.
Could we add an execution-level regression test that combines a multi-batch left input, a join filter, and a reordered NestedLoopJoinExec projection? The new LogicalBatch::take_column path is now used for both filter columns and output columns, but the current runtime join tests appear to construct the executor with projection: None, while the existing projection coverage only checks proto state.
I don't think this needs to block the PR since the unprojected multi-segment paths and the LogicalBatch gather behavior are already well covered, but having this combination exercised end to end would help protect the new path.
There was a problem hiding this comment.
I don't fully understand what coverage is missing. At least the coverage report looks okay: #25371 (comment)
If you can come up with a concrete test case, I'm happy to add it. More test cases are always a win.
|
I plan to merge it after 1 day, unless someone needs more time to review or have additional feedbacks. |
4 participants
Which issue does this PR close?
concat_batchesin joins #23076Rationale for this change
NestedLoopJoinExecbuffers the build side as multipleRecordBatches and then concatenates them into a single batch. During concatenation, the original batches and the concatenated output coexist, temporarily requiring roughly 2x the memory for the buffered data. See below issue for more details:concat_batchesin joins #23076This PR avoids that memory amplification by keeping the build side in a segmented layout:
The segmented batches are wrapped in a reusable
LogicalBatchabstraction, which presents them as one logically contiguous batch while hiding the underlying physical layout from operators.This has two main benefits:
concat_batchesmemory amplification.LogicalBatchas a single logical batch, to simplify their inner logic.This PR supersedes #24820. The main difference is that this PR extracts the segmented batch representation into the reusable
LogicalBatchabstraction, while #24820 handles theVec<RecordBatch>layout directly insideNestedLoopJoinExec.Reproducer
The query below has approximately 800 MB of raw build-side
Int64values. Avoiding concatenation eliminates the additional full-size copy created byconcat_batches.Measure peak RSS on both the PR branch and
main:Peak RSS:
mainWhat changes are included in this PR?
LogicalBatchmodule, with single contiguous batch abstraction, but internally use segmented physical layoutLogicalBatchin NLJWhat is the testing strategy for this PR?
LogicalBatchmoduleAre there any user-facing changes?
No