[BUG] Enable groupby with alias for native executor (#2917)

Fixed this already for the current translation path back in #2790, doing the same for swordfish --------- Co-authored-by: Colin Ho <[email protected]> Co-authored-by: Colin Ho <[email protected]>
Eventual-Inc · Sep 26, 2024 · e52d965 · e52d965
1 parent c84940e
commit e52d965
Show file tree

Hide file tree

Showing 6 changed files with 19 additions and 41 deletions.
diff --git a/src/daft-local-execution/src/pipeline.rs b/src/daft-local-execution/src/pipeline.rs
@@ -7,7 +7,7 @@ use daft_core::{
     prelude::{Schema, SchemaRef},
     utils::supertype,
 };
-use daft_dsl::{join::get_common_join_keys, Expr};
+use daft_dsl::{col, join::get_common_join_keys, Expr};
 use daft_micropartition::MicroPartition;
 use daft_physical_plan::{
     Filter, HashAggregate, HashJoin, InMemoryScan, Limit, LocalPhysicalPlan, Project, Sort,
@@ -186,17 +186,23 @@ pub fn physical_plan_to_pipeline(
         }) => {
             let (first_stage_aggs, second_stage_aggs, final_exprs) =
                 populate_aggregation_stages(aggregations, schema, group_by);
-            let first_stage_agg_op = AggregateOperator::new(
-                first_stage_aggs
-                    .values()
-                    .cloned()
-                    .map(|e| Arc::new(Expr::Agg(e.clone())))
-                    .collect(),
-                group_by.clone(),
-            );
             let child_node = physical_plan_to_pipeline(input, psets)?;
-            let post_first_agg_node =
-                IntermediateNode::new(Arc::new(first_stage_agg_op), vec![child_node]).boxed();
+            let (post_first_agg_node, group_by) = if !first_stage_aggs.is_empty() {
+                let agg_op = AggregateOperator::new(
+                    first_stage_aggs
+                        .values()
+                        .cloned()
+                        .map(|e| Arc::new(Expr::Agg(e.clone())))
+                        .collect(),
+                    group_by.clone(),
+                );
+                (
+                    IntermediateNode::new(Arc::new(agg_op), vec![child_node]).boxed(),
+                    &group_by.iter().map(|e| col(e.name())).collect(),
+                )
+            } else {
+                (child_node, group_by)
+            };
 
             let second_stage_agg_sink = AggregateSink::new(
                 second_stage_aggs

diff --git a/tests/cookbook/test_aggregations.py b/tests/cookbook/test_aggregations.py
@@ -6,17 +6,11 @@
 import pandas as pd
 import pytest
 
-from daft import context
 from daft.datatype import DataType
 from daft.expressions import col
 from daft.udf import udf
 from tests.conftest import assert_df_equals
 
-pytestmark = pytest.mark.skipif(
-    context.get_context().daft_execution_config.enable_native_executor is True,
-    reason="Native executor fails for these tests",
-)
-
 
 def test_sum(daft_df, service_requests_csv_pd_df, repartition_nparts):
     """Sums across an entire column for the entire table"""

diff --git a/tests/dataframe/test_aggregations.py b/tests/dataframe/test_aggregations.py
@@ -7,18 +7,13 @@
 import pytest
 
 import daft
-from daft import col, context
+from daft import col
 from daft.context import get_context
 from daft.datatype import DataType
 from daft.errors import ExpressionTypeError
 from daft.utils import freeze
 from tests.utils import sort_arrow_table
 
-pytestmark = pytest.mark.skipif(
-    context.get_context().daft_execution_config.enable_native_executor is True,
-    reason="Native executor fails for these tests",
-)
-
 
 @pytest.mark.parametrize("repartition_nparts", [1, 2, 4])
 def test_agg_global(make_df, repartition_nparts):

diff --git a/tests/dataframe/test_approx_percentiles_aggregations.py b/tests/dataframe/test_approx_percentiles_aggregations.py
@@ -4,12 +4,7 @@
 import pyarrow as pa
 import pytest
 
-from daft import col, context
-
-pytestmark = pytest.mark.skipif(
-    context.get_context().daft_execution_config.enable_native_executor is True,
-    reason="Native executor fails for these tests",
-)
+from daft import col
 
 
 @pytest.mark.parametrize("repartition_nparts", [1, 2, 4])

diff --git a/tests/dataframe/test_distinct.py b/tests/dataframe/test_distinct.py
@@ -3,15 +3,9 @@
 import pyarrow as pa
 import pytest
 
-from daft import context
 from daft.datatype import DataType
 from tests.utils import sort_arrow_table
 
-pytestmark = pytest.mark.skipif(
-    context.get_context().daft_execution_config.enable_native_executor is True,
-    reason="Native executor fails for these tests",
-)
-
 
 @pytest.mark.parametrize("repartition_nparts", [1, 2, 5])
 def test_distinct_with_nulls(make_df, repartition_nparts):

diff --git a/tests/dataframe/test_map_groups.py b/tests/dataframe/test_map_groups.py
@@ -3,12 +3,6 @@
 import pytest
 
 import daft
-from daft import context
-
-pytestmark = pytest.mark.skipif(
-    context.get_context().daft_execution_config.enable_native_executor is True,
-    reason="Native executor fails for these tests",
-)
 
 
 @pytest.mark.parametrize("repartition_nparts", [1, 2, 4])