fix: output_bytes metric in hash aggregation

[ariel-miculas]

Which issue does this PR close?

• Closes Incorrect output_bytes metric in hash aggregation #23776.

Rationale for this change

Summing together the produced sliced RecordBatches in hash aggregation leads to inflated output_bytes metric

What changes are included in this PR?

Are these changes tested?

Yes

Are there any user-facing changes?

If the metrics are user-facing, then yes.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 95.74468% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.14%. Comparing base (7e015b7) to head (016f1a4).
⚠️ Report is 9 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/display.rs	0.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #23777      +/-   ##
==========================================
- Coverage   81.29%   81.14%   -0.15%     
==========================================
  Files        1110     1110              
  Lines      385197   386141     +944     
  Branches   385197   386141     +944     
==========================================
+ Hits       313132   313342     +210     
- Misses      53588    54332     +744     
+ Partials    18477    18467      -10     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Jithendra2608]

Hi @ariel-miculas! I was taking a look at the failing CI checks to see if I could help unblock this.
It looks like the typos job is failing because of a small spelling error in datafusion/physical-expr-common/src/metrics/baseline.rs on line 326 (accross instead of across).

Fixing that one word should turn the CI green. Hope this helps, and great work on the hash aggregation metrics!

[ariel-miculas]

@2010YOUY01 coud you please take a look?

[2010YOUY01]

Is it possible to implement it entirely within the record_output, and avoid large changes to different operators inner logic?

Also let's add some end-to-end tests, ideally we can do it with sqllogictest

[ariel-miculas]

Is it possible to implement it entirely within the record_output, and avoid large changes to different operators inner logic?

I don't see how, since RecordBatch::record_output doesn't keep any state related to previous RecordBatches that have been processed, and we need this information so that we can deduplicate the sliced RecordBatches to avoid summing up the memory for the same underlying original large RecordBatch.
This is a limitation that we need to deal with until we get #15591

Also let's add some end-to-end tests, ideally we can do it with sqllogictest

Added an sqllogictest

[ariel-miculas]

Closed and reopen to see if the test failure in an unrelated ordered_aggregate_spill.slt is transient:

External error: 1 errors in file /__w/datafusion/datafusion/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt

1. query failed: DataFusion error: Resources exhausted: Additional allocation failed for OrderedFinalAggregateStream[1] with top memory consumers (across reservations) as:
  OrderedFinalAggregateStream[0]#39115(can spill: true) consumed 406.8 KB, peak 496.2 KB,
  OrderedFinalAggregateStream[1]#39112(can spill: true) consumed 180.0 KB, peak 180.0 KB,
  OrderedFinalAggregateStream[1]#39235(can spill: false) consumed 5.4 KB, peak 11.5 KB.
Error: Failed to allocate additional 8.1 KB for OrderedFinalAggregateStream[1] with 5.4 KB already allocated for this reservation - 7.7 KB remain available for the total memory pool: greedy(used: 592.3 KB, pool_size: 600.0 KB)
[SQL] EXPLAIN ANALYZE
SELECT round(v1, -4), v1 % 5000,
  sum(v1 * 2), min(CAST(v1 % 2 AS VARCHAR))
FROM generate_series(20000) AS t1(v1)
GROUP BY round(v1, -4), v1 % 5000
at /__w/datafusion/datafusion/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt:166



Error: Execution("1 failures\n\nbacktrace:    0: sqllogictests::run_tests::{closure#0}\n   1: <tokio::runtime::park::CachedParkThread>::block_on::<sqllogictests::run_tests::{closure#0}>::{closure#0}\n   2: <tokio::runtime::park::CachedParkThread>::block_on::<sqllogictests::run_tests::{closure#0}>\n   3: <tokio::runtime::context::blocking::BlockingRegionGuard>::block_on::<sqllogictests::run_tests::{closure#0}>\n   4: <tokio::runtime::scheduler::multi_thread::MultiThread>::block_on::<sqllogictests::run_tests::{closure#0}>::{closure#0}\n   5: tokio::runtime::context::runtime::enter_runtime::<<tokio::runtime::scheduler::multi_thread::MultiThread>::block_on<sqllogictests::run_tests::{closure#0}>::{closure#0}, core::result::Result<(), datafusion_common::error::DataFusionError>>\n   6: <tokio::runtime::scheduler::multi_thread::MultiThread>::block_on::<sqllogictests::run_tests::{closure#0}>\n   7: <tokio::runtime::runtime::Runtime>::block_on_inner::<sqllogictests::run_tests::{closure#0}>\n   8: <tokio::runtime::runtime::Runtime>::block_on::<sqllogictests::run_tests::{closure#0}>\n   9: sqllogictests::main\n  10: <fn() -> core::result::Result<(), datafusion_common::error::DataFusionError> as core::ops::function::FnOnce<()>>::call_once\n  11: std::sys::backtrace::__rust_begin_short_backtrace::<fn() -> core::result::Result<(), datafusion_common::error::DataFusionError>, core::result::Result<(), datafusion_common::error::DataFusionError>>\n  12: std::rt::lang_start::<core::result::Result<(), datafusion_common::error::DataFusionError>>::{closure#0}\n  13: std::rt::lang_start_internal\n  14: std::rt::lang_start::<core::result::Result<(), datafusion_common::error::DataFusionError>>\n  15: main\n  16: <unknown>\n  17: __libc_start_main\n  18: _start\n")
error: test failed, to rerun pass `-p datafusion-sqllogictest --test sqllogictests`

Caused by:
  process didn't exit successfully: `/__w/datafusion/datafusion/target/llvm-cov-target/ci/deps/sqllogictests-917a8f0bbb363691` (exit status: 1)
error: process didn't exit successfully: `/usr/local/rustup/toolchains/1.97.0-x86_64-unknown-linux-gnu/bin/cargo test --profile ci --manifest-path /__w/datafusion/datafusion/Cargo.toml --target-dir /__w/datafusion/datafusion/target/llvm-cov-target --exclude datafusion-examples --exclude ffi_example_table_provider --exclude datafusion-cli --workspace --lib --tests --bins --features serde,avro,json,backtrace,integration-tests,parquet_encryption,substrait` (exit status: 1)

[ariel-miculas]

@2010YOUY01 let me know what you think

[2010YOUY01]

Is it possible to implement it entirely within the record_output, and avoid large changes to different operators inner logic?

I don't see how, since RecordBatch::record_output doesn't keep any state related to previous RecordBatches that have been processed, and we need this information so that we can deduplicate the sliced RecordBatches to avoid summing up the memory for the same underlying original large RecordBatch. This is a limitation that we need to deal with until we get #15591

Perhaps we can use something like RecordBatchMemoryCounter (with internal buffer tracking to avoid double-counting) to replace the below Counter:

datafusion/datafusion/physical-expr-common/src/metrics/baseline.rs

Line 69 in f27e50c

output_bytes: Count,

[ariel-miculas]

Perhaps we can use something like RecordBatchMemoryCounter (with internal buffer tracking to avoid double-counting) to replace the below Counter:

That's kind of what I'm doing, except I've wrapped RecordBatchMemoryCounter in a newtype struct:

#[derive(Debug, Default)]
pub struct RecordBatchMemoryMetrics(RecordBatchMemoryCounter);

impl RecordBatchMemoryMetrics {
    pub fn new() -> Self {
        Self::default()
    }

    /// Similar to RecordBatch.record_output, but deduplicating across batches to avoid
    /// output_size inflation due to shared memory
    pub fn record_output(&mut self, batch: &RecordBatch, bm: &BaselineMetrics) {
        bm.record_output(batch.num_rows());
        let n_bytes = self.0.count_batch(batch);
        bm.output_bytes.add(n_bytes);
        bm.output_batches.add(1);
    }
}

Unless I've missed something from your explanation.

[ariel-miculas]

Or you mean changing output_bytes itself from Count to RecordBatchMemoryCounter?
It would make sense if there are other operators which slice a large RecordBatch, otherwise is it worth doing this solely for the current aggregation implementation?

[2010YOUY01]

Or you mean changing output_bytes itself from Count to RecordBatchMemoryCounter? It would make sense if there are other operators which slice a large RecordBatch, otherwise is it worth doing this solely for the current aggregation implementation?

Yes, exactly.

Even if it's only for aggregation, I think this approach is still overall simpler. The existing approach seems to leak some complexity into the aggregation logic.

Also, the memory overcounting in the metrics is also likely in other operators, since now we're using batch concatenation quite commonly in joins/window executions.

[ariel-miculas]

Implemented the proposed approach, indeed it looks simpler overall.

[2010YOUY01]

Thank you, the implementation looks good to me.

There is a related inefficiency that was reported recently. I'll try to look into it soon and then figure out how we should proceed with this PR.

Even if we can't eliminate the extra overhead, we can still add a config option to enable or disable this hash-set-based deduplication, this shouldn't be a hard blocker.

#24310

[ariel-miculas]

Even if we can't eliminate the extra overhead, we can still add a config option to enable or disable this hash-set-based deduplication, this shouldn't be a hard blocker.

Instead of a config option, could we use a flag that gets set by select operators, such as hash aggregate? Because for hash aggregate we already know that sliced RecordBatches will be produced from a huge initial RecordBatch.

The issue is that we need to pass this information to the downstream consumer of these batches, since there's where the memory accounting is broken (in main).
Sorry, I was confusing the reservation mechanism with the output_bytes metrics. For output_bytes, we need to make the changes in the operator themselves, which should be easy to do.

With the config option my concern is that we'd be trading correctness for performance, since without the deduplication the hash aggregation shows heavily inflated output_bytes metrics.

[2010YOUY01]

Instead of a config option, could we use a flag that gets set by select operators, such as hash aggregate? Because for hash aggregate we already know that sliced RecordBatches will be produced from a huge initial RecordBatch.

Yes, I think starting with only aggregate is a good idea

With the config option my concern is that we'd be trading correctness for performance, since without the deduplication the hash aggregation shows heavily inflated output_bytes metrics.

I agree. That also adds extra complexity.

I imagine the performance impact would be minimal, since a hash set whose size is proportional to the number of batches is much smaller than the actual workload. We just need to verify that in #24310