kosiew commented on code in PR #25086: URL: https://github.com/apache/datafusion/pull/25086#discussion_r4046612859
########## datafusion/sqllogictest/test_files/aggregate_output_bytes.slt: ########## @@ -0,0 +1,91 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# -------------------------------------------- +# Regression test for the hash aggregation `output_bytes` metric. +# +# The hash aggregate's grouping table materializes its emitted output as a +# single RecordBatch and then slices it into `batch_size`-sized chunks across +# successive `poll_next` calls. Those slices share the same underlying +# buffers (Arrow `.slice()` is zero-copy), so each buffer must only be +# counted once in `output_bytes`, no matter how many slices reference it. +# +# Using `RecordBatch::record_output` (as the streams used to) recomputes +# memory usage from scratch on every slice, with no memory of buffers +# already counted -- so all N slices of one materialization each count the +# shared buffers again, inflating `output_bytes` by ~Nx. `RecordBatchMemoryMetrics` Review Comment: Small documentation nit: the new implementation doesn't use `RecordBatchMemoryMetrics`, and I don't think that type exists. The fix records the unsliced materialized batch once via `BaselineMetrics::record_output_bytes`. Could you update this test header to describe the mechanism that's actually being tested? ########## datafusion/physical-plan/src/aggregates/mod.rs: ########## @@ -8467,4 +8467,160 @@ mod tests { assert!(agg.set_dynamic_filter(df).is_err()); Ok(()) } + + /// Regression test for the aggregate's `output_bytes` metric. + /// + /// The aggregate materializes each emit as one batch and hands it + /// downstream in `batch_size` slices that share the materialized buffers. + /// Two failure modes are guarded against: + /// + /// * Recording every slice with `RecordBatch::record_output` charges the + /// shared buffers once per slice, inflating `output_bytes`. + /// * Deduplicating buffers by address for the life of the stream + /// undercounts, because a downstream consumer drops each batch and the + /// allocator recycles its address for the next emit. + /// + /// The fix records rows and bytes once per materialization and only the + /// batch count per slice, so the result must not depend on whether the + /// consumer retains or drops batches, nor on how many slices an emit is + /// cut into. + /// + /// Sorted partial aggregation emits one materialization per completed key + /// range, giving many independently allocated emits. `batch_size` 8192 + /// leaves them whole; `batch_size` 100 slices each 512-row emit into 6. + #[tokio::test] + async fn output_bytes_metric_is_exact_for_sliced_and_dropped_batches() -> Result<()> { Review Comment: Could we add some focused metric regression coverage for the migrated hash streams changed by this PR: `PartialHashAggregateStream`, `FinalHashAggregateStream`, `SingleHashAggregateStream`, and `PartialReduceHashAggregateStream`? The migration flag already defaults to `true`, but this test provides sorted input, so it selects `OrderedPartialAggregateStream` rather than `PartialHashAggregateStream`. It would be useful to use linear input and assert which stream path is selected, or construct the streams directly, while still exercising sliced output. Similar legacy fallback coverage would also be useful while `GroupedHashAggregateStream` remains supported. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
