[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #6124: Implement Streaming Aggregation: Do not break pipeline in aggregation if group by columns are ordered (V2)

["alamb (via GitHub)" <gi...@apache.org>]

alamb commented on code in PR #6124:
URL: https://github.com/apache/arrow-datafusion/pull/6124#discussion_r1178097828


##########
datafusion/core/tests/sql/group_by.rs:
##########
@@ -160,3 +161,99 @@ async fn group_by_dictionary() {
     run_test_case::<UInt32Type>().await;
     run_test_case::<UInt64Type>().await;
 }
+
+#[tokio::test]
+async fn test_source_sorted_groupby() -> Result<()> {

Review Comment:
   I wonder if there is a reason to use `rs` here and not sqllogictest ? I don't see it as a blocker for this PR it just makes that much more work to port these tests over eventually



##########
datafusion/core/tests/sql/group_by.rs:
##########
@@ -160,3 +161,99 @@ async fn group_by_dictionary() {
     run_test_case::<UInt32Type>().await;
     run_test_case::<UInt64Type>().await;
 }
+
+#[tokio::test]
+async fn test_source_sorted_groupby() -> Result<()> {
+    let tmpdir = TempDir::new().unwrap();
+    let session_config = SessionConfig::new().with_target_partitions(1);
+    let ctx = get_test_context2(&tmpdir, true, session_config).await?;
+
+    let sql = "SELECT a, b,
+           SUM(c) as summation1
+           FROM annotated_data
+           GROUP BY b, a";
+
+    let msg = format!("Creating logical plan for '{sql}'");
+    let dataframe = ctx.sql(sql).await.expect(&msg);
+    let physical_plan = dataframe.create_physical_plan().await?;
+    let formatted = displayable(physical_plan.as_ref()).indent().to_string();
+    let expected = {
+        vec![
+            "ProjectionExec: expr=[a@1 as a, b@0 as b, SUM(annotated_data.c)@2 as summation1]",
+            "  AggregateExec: mode=Single, gby=[b@1 as b, a@0 as a], aggr=[SUM(annotated_data.c)]",
+        ]
+    };
+
+    let actual: Vec<&str> = formatted.trim().lines().collect();
+    let actual_len = actual.len();
+    let actual_trim_last = &actual[..actual_len - 1];
+    assert_eq!(
+        expected, actual_trim_last,
+        "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
+    );
+
+    let actual = execute_to_batches(&ctx, sql).await;
+    let expected = vec![
+        "+---+---+------------+",
+        "| a | b | summation1 |",
+        "+---+---+------------+",
+        "| 0 | 0 | 300        |",
+        "| 0 | 1 | 925        |",
+        "| 1 | 2 | 1550       |",
+        "| 1 | 3 | 2175       |",
+        "+---+---+------------+",
+    ];
+    assert_batches_eq!(expected, &actual);
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_source_sorted_groupby2() -> Result<()> {
+    let tmpdir = TempDir::new().unwrap();
+    let session_config = SessionConfig::new().with_target_partitions(1);
+    let ctx = get_test_context2(&tmpdir, true, session_config).await?;
+
+    let sql = "SELECT a, d,
+           SUM(c) as summation1
+           FROM annotated_data
+           GROUP BY d, a";
+
+    let msg = format!("Creating logical plan for '{sql}'");
+    let dataframe = ctx.sql(sql).await.expect(&msg);
+    let physical_plan = dataframe.create_physical_plan().await?;
+    let formatted = displayable(physical_plan.as_ref()).indent().to_string();
+    let expected = {
+        vec![
+            "ProjectionExec: expr=[a@1 as a, d@0 as d, SUM(annotated_data.c)@2 as summation1]",
+            "  AggregateExec: mode=Single, gby=[d@2 as d, a@0 as a], aggr=[SUM(annotated_data.c)]",

Review Comment:
   I think it would be good to add  an annotation of the `GroupByOrderMode` in the explain plan so it is clear which stream type would be used.
   
   I think we could do so in a follow on PR



##########
datafusion/core/tests/aggregate_fuzz.rs:
##########
@@ -0,0 +1,221 @@
+// 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.
+
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array};
+use arrow::compute::{concat_batches, SortOptions};
+use arrow::datatypes::DataType;
+use arrow::record_batch::RecordBatch;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::physical_plan::aggregates::{
+    AggregateExec, AggregateMode, PhysicalGroupBy,
+};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+
+use datafusion::physical_plan::collect;
+use datafusion::physical_plan::memory::MemoryExec;
+use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion_physical_expr::expressions::{col, Sum};
+use datafusion_physical_expr::{AggregateExpr, PhysicalSortExpr};
+use test_utils::add_empty_batches;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
+    async fn aggregate_test() {
+        let test_cases = vec![
+            vec!["a"],
+            vec!["b", "a"],
+            vec!["c", "a"],
+            vec!["c", "b", "a"],
+            vec!["d", "a"],
+            vec!["d", "b", "a"],
+            vec!["d", "c", "a"],
+            vec!["d", "c", "b", "a"],
+        ];
+        let n = 300;
+        let distincts = vec![10, 20];
+        for distinct in distincts {
+            let mut handles = Vec::new();
+            for i in 0..n {
+                let test_idx = i % test_cases.len();
+                let group_by_columns = test_cases[test_idx].clone();
+                let job = tokio::spawn(run_aggregate_test(
+                    make_staggered_batches::<true>(1000, distinct, i as u64),
+                    group_by_columns,
+                ));
+                handles.push(job);
+            }
+            for job in handles {
+                job.await.unwrap();
+            }
+        }
+    }
+}
+
+/// Perform batch and running window same input
+/// and verify outputs of `WindowAggExec` and `BoundedWindowAggExec` are equal

Review Comment:
   I think this comment needs to be updated to refer to GroupBys that have sorted and non sorted inputs



##########
datafusion/core/tests/sqllogictests/test_files/window.slt:
##########
@@ -356,23 +356,22 @@ Sort: d.b ASC NULLS LAST
                       EmptyRelation
 physical_plan
 SortPreservingMergeExec: [b@0 ASC NULLS LAST]
-  SortExec: expr=[b@0 ASC NULLS LAST]

Review Comment:
   this is nice that it no longer has to resort the input by b (as the group by preserves the sort)



##########
datafusion/core/src/physical_plan/aggregates/bounded_aggregate_stream.rs:
##########
@@ -0,0 +1,1041 @@
+// 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.
+
+//! Hash aggregation through row format

Review Comment:
   I think this comment could be updated to refer to operating on sorted inputs



##########
datafusion/core/tests/aggregate_fuzz.rs:
##########


Review Comment:
   thank you for this test / file



##########
datafusion/core/tests/window_fuzz.rs:
##########
@@ -480,6 +480,10 @@ async fn run_window_test(
     let collected_running = collect(running_window_exec, task_ctx.clone())
         .await
         .unwrap();
+

Review Comment:
   👍 



##########
datafusion/core/src/physical_plan/aggregates/mod.rs:
##########
@@ -72,6 +76,15 @@ pub enum AggregateMode {
     Single,
 }
 
+/// Group By expression modes
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum GroupByOrderMode {
+    /// Some of the expressions in the GROUP BY clause have an ordering.

Review Comment:
   I wonder if it we could add a more precise definition of "partial" here
   
   perhaps something like the "prefix of the group keys is some permutation of the sort key prefix "?
   
   I think your inline comment below explains it nicely with an example:
   ```
       // Find out how many expressions of the existing ordering define ordering
       // for expressions in the GROUP BY clause. For example, if the input is
       // ordered by a, b, c, d and we group by b, a, d; the result below would be.
       // 2, meaning 2 elements (a, b) among the GROUP BY columns define ordering.
   ```
   
   Maybe we can replicate that comment here in the documentation for GroupByOrderMode
   
   



-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org