You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "vbarua (via GitHub)" <gi...@apache.org> on 2024/09/16 22:06:47 UTC

[PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

vbarua opened a new pull request, #12495:
URL: https://github.com/apache/datafusion/pull/12495

   ## Which issue does this PR close?
   Continues #12347 
   
   This PR sets the [output_mapping](https://github.com/substrait-io/substrait/blob/bc4d6fb9bc0435c3db24172566c343e119fc50a9/proto/substrait/algebra.proto#L30-L33) field on ProjectRels for plans _produced_ by DataFusion. 
   
   ## Rationale for this change
   Given a plan
   ```
   Projection: data.b, data.a + data.a, data.a
     TableScan: data projection=[a, b]",
   ```
   
   When creating the Substrait ProjectRel, DataFusion will now set the output mapping to [2,3,4] so that the Substrait ProjectRel will match the DataFusion behaviour of _only_ outputting expressions.
   
   By default, if the `emit_kind` field is not set with an `output_mapping`, [a Substrait ProjectRel emits all input columns followed by all expressions](https://substrait.io/relations/logical_relations/#project-operation). This does not match the DataFusion behaviour.
   
   This change only includes _producer_ changes in order to ease the migration for users. Because DataFusion ignores the `emit_kind` field when consuming plans, DataFusion will be able to consume plans that set the output mapping correctly and plans that don't.
   
   Once the _consumer_ is updated to respect `emit_kind`, plans that don't set the `output_mapping` will not work as before. By shipping the producer change by itself, we allow users who serialize plans between versions the opportunity to update all of their plans before making a consumer change that would break them.
   
   ## Are these changes tested?
   Yes
   
   ## Are there any user-facing changes?
   Substrait plans with ProjectRels will now set the output_mapping field. 
   


-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #12495:
URL: https://github.com/apache/datafusion/pull/12495#discussion_r1766422065


##########
datafusion/substrait/src/logical_plan/producer.rs:
##########
@@ -432,38 +445,39 @@ pub fn to_substrait_rel(
         }
         LogicalPlan::Window(window) => {
             let input = to_substrait_rel(window.input.as_ref(), ctx, extensions)?;
-            // If the input is a Project relation, we can just append the WindowFunction expressions
-            // before returning
-            // Otherwise, wrap the input in a Project relation before appending the WindowFunction
-            // expressions
-            let mut project_rel: Box<ProjectRel> = match &input.as_ref().rel_type {
-                Some(RelType::Project(p)) => Box::new(*p.clone()),
-                _ => {
-                    // Create Projection with field referencing all output fields in the input relation
-                    let expressions = (0..window.input.schema().fields().len())
-                        .map(substrait_field_ref)
-                        .collect::<Result<Vec<_>>>()?;
-                    Box::new(ProjectRel {
-                        common: None,
-                        input: Some(input),
-                        expressions,
-                        advanced_extension: None,
-                    })
-                }
-            };
-            // Parse WindowFunction expression
-            let mut window_exprs = vec![];
+
+            // create a field reference for each input field
+            let mut expressions = (0..window.input.schema().fields().len())

Review Comment:
   I guess you could not add these ref expressions, and then use a direct emit (or remap starting from 0) instead, would that not be the same? That said it'd break roundtrip until you've gotten the consumer part done as well 😅 



-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #12495:
URL: https://github.com/apache/datafusion/pull/12495#discussion_r1762024765


##########
datafusion/substrait/src/logical_plan/producer.rs:
##########
@@ -432,38 +445,39 @@ pub fn to_substrait_rel(
         }
         LogicalPlan::Window(window) => {
             let input = to_substrait_rel(window.input.as_ref(), ctx, extensions)?;
-            // If the input is a Project relation, we can just append the WindowFunction expressions
-            // before returning
-            // Otherwise, wrap the input in a Project relation before appending the WindowFunction
-            // expressions
-            let mut project_rel: Box<ProjectRel> = match &input.as_ref().rel_type {
-                Some(RelType::Project(p)) => Box::new(*p.clone()),
-                _ => {
-                    // Create Projection with field referencing all output fields in the input relation
-                    let expressions = (0..window.input.schema().fields().len())
-                        .map(substrait_field_ref)
-                        .collect::<Result<Vec<_>>>()?;
-                    Box::new(ProjectRel {
-                        common: None,
-                        input: Some(input),
-                        expressions,
-                        advanced_extension: None,
-                    })
-                }
-            };

Review Comment:
   This is potentially a useful optimization, however it becomes a bit more complicated with the introduction of the `output_mapping` because you a need to modify it along with the expressions. I've opted to simplify this for now to favour simplicity and correctness. 
   
   As well, I think this is better handled in when consuming plans and/or by the optimizer.



-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #12495:
URL: https://github.com/apache/datafusion/pull/12495#discussion_r1765310035


##########
datafusion/substrait/tests/cases/serialize.rs:
##########
@@ -63,6 +67,103 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn include_remaps_for_projects() -> Result<()> {
+        let ctx = create_context().await?;
+        let df = ctx.sql("SELECT b, a + a, a FROM data").await?;
+        let datafusion_plan = df.into_optimized_plan()?;
+
+        assert_eq!(
+            format!("{}", datafusion_plan),
+            "Projection: data.b, data.a + data.a, data.a\
+            \n  TableScan: data projection=[a, b]",
+        );
+
+        let plan = to_substrait_plan(&datafusion_plan, &ctx)?.as_ref().clone();
+
+        let relation = plan.relations.first().unwrap().rel_type.as_ref();
+        let root_rel = match relation {
+            Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
+            _ => panic!("expected Root"),
+        };
+        if let Some(rel::RelType::Project(p)) = root_rel.rel_type.as_ref() {
+            // The input has 2 columns [a, b], the Projection has 3 expressions [b, a + a, a]
+            // The required output mapping is [2,3,4], which skips the 2 input columns.
+            assert_emit(p.common.as_ref(), vec![2, 3, 4]);
+
+            if let Some(rel::RelType::Read(r)) =
+                p.input.as_ref().unwrap().rel_type.as_ref()
+            {
+                let mask_expression = r.projection.as_ref().unwrap();
+                let select = mask_expression.select.as_ref().unwrap();
+                assert_eq!(
+                    2,
+                    select.struct_items.len(),
+                    "Read outputs two columns: a, b"
+                );
+                return Ok(());
+            }
+        }
+        panic!("plan did not match expected structure")
+    }
+
+    #[tokio::test]
+    async fn include_remaps_for_windows() -> Result<()> {
+        let ctx = create_context().await?;
+        // let df = ctx.sql("SELECT a, b, lead(b) OVER (PARTITION BY a) FROM data").await?;
+        let df = ctx
+            .sql("SELECT b, RANK() OVER (PARTITION BY a), c FROM data;")
+            .await?;
+        let datafusion_plan = df.into_optimized_plan()?;
+        assert_eq!(
+            format!("{}", datafusion_plan),
+            "Projection: data.b, RANK() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING, data.c\
+            \n  WindowAggr: windowExpr=[[RANK() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]\
+            \n    TableScan: data projection=[a, b, c]",
+        );
+
+        let plan = to_substrait_plan(&datafusion_plan, &ctx)?.as_ref().clone();
+
+        let relation = plan.relations.first().unwrap().rel_type.as_ref();
+        let root_rel = match relation {
+            Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
+            _ => panic!("expected Root"),
+        };
+
+        if let Some(rel::RelType::Project(p1)) = root_rel.rel_type.as_ref() {
+            // The WindowAggr outputs 4 columns, the Projection has 4 columns

Review Comment:
   ~I think both are 3 columns?~
   
   ~Or well the window outputs 4 but the project outputs only 3.~
   
   I guess that's what you're saying here anyways 😅 
   



-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #12495:
URL: https://github.com/apache/datafusion/pull/12495#issuecomment-2356542050

   FYI @Blizzara 


-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #12495:
URL: https://github.com/apache/datafusion/pull/12495#discussion_r1763640110


##########
datafusion/substrait/src/logical_plan/producer.rs:
##########
@@ -432,38 +445,39 @@ pub fn to_substrait_rel(
         }
         LogicalPlan::Window(window) => {
             let input = to_substrait_rel(window.input.as_ref(), ctx, extensions)?;
-            // If the input is a Project relation, we can just append the WindowFunction expressions
-            // before returning
-            // Otherwise, wrap the input in a Project relation before appending the WindowFunction
-            // expressions
-            let mut project_rel: Box<ProjectRel> = match &input.as_ref().rel_type {
-                Some(RelType::Project(p)) => Box::new(*p.clone()),
-                _ => {
-                    // Create Projection with field referencing all output fields in the input relation
-                    let expressions = (0..window.input.schema().fields().len())
-                        .map(substrait_field_ref)
-                        .collect::<Result<Vec<_>>>()?;
-                    Box::new(ProjectRel {
-                        common: None,
-                        input: Some(input),
-                        expressions,
-                        advanced_extension: None,
-                    })
-                }
-            };

Review Comment:
   Yes, the datafusion optimizer already handles projection pushdown quite well -- so keeping the substrait producer simpler makes the most sense to me



-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #12495:
URL: https://github.com/apache/datafusion/pull/12495#discussion_r1765310035


##########
datafusion/substrait/tests/cases/serialize.rs:
##########
@@ -63,6 +67,103 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn include_remaps_for_projects() -> Result<()> {
+        let ctx = create_context().await?;
+        let df = ctx.sql("SELECT b, a + a, a FROM data").await?;
+        let datafusion_plan = df.into_optimized_plan()?;
+
+        assert_eq!(
+            format!("{}", datafusion_plan),
+            "Projection: data.b, data.a + data.a, data.a\
+            \n  TableScan: data projection=[a, b]",
+        );
+
+        let plan = to_substrait_plan(&datafusion_plan, &ctx)?.as_ref().clone();
+
+        let relation = plan.relations.first().unwrap().rel_type.as_ref();
+        let root_rel = match relation {
+            Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
+            _ => panic!("expected Root"),
+        };
+        if let Some(rel::RelType::Project(p)) = root_rel.rel_type.as_ref() {
+            // The input has 2 columns [a, b], the Projection has 3 expressions [b, a + a, a]
+            // The required output mapping is [2,3,4], which skips the 2 input columns.
+            assert_emit(p.common.as_ref(), vec![2, 3, 4]);
+
+            if let Some(rel::RelType::Read(r)) =
+                p.input.as_ref().unwrap().rel_type.as_ref()
+            {
+                let mask_expression = r.projection.as_ref().unwrap();
+                let select = mask_expression.select.as_ref().unwrap();
+                assert_eq!(
+                    2,
+                    select.struct_items.len(),
+                    "Read outputs two columns: a, b"
+                );
+                return Ok(());
+            }
+        }
+        panic!("plan did not match expected structure")
+    }
+
+    #[tokio::test]
+    async fn include_remaps_for_windows() -> Result<()> {
+        let ctx = create_context().await?;
+        // let df = ctx.sql("SELECT a, b, lead(b) OVER (PARTITION BY a) FROM data").await?;
+        let df = ctx
+            .sql("SELECT b, RANK() OVER (PARTITION BY a), c FROM data;")
+            .await?;
+        let datafusion_plan = df.into_optimized_plan()?;
+        assert_eq!(
+            format!("{}", datafusion_plan),
+            "Projection: data.b, RANK() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING, data.c\
+            \n  WindowAggr: windowExpr=[[RANK() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]\
+            \n    TableScan: data projection=[a, b, c]",
+        );
+
+        let plan = to_substrait_plan(&datafusion_plan, &ctx)?.as_ref().clone();
+
+        let relation = plan.relations.first().unwrap().rel_type.as_ref();
+        let root_rel = match relation {
+            Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
+            _ => panic!("expected Root"),
+        };
+
+        if let Some(rel::RelType::Project(p1)) = root_rel.rel_type.as_ref() {
+            // The WindowAggr outputs 4 columns, the Projection has 4 columns

Review Comment:
   ~I think both are 3 columns?~
   
   Or well the window outputs 4 but the project outputs only 3.



-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "Blizzara (via GitHub)" <gi...@apache.org>.
Blizzara commented on code in PR #12495:
URL: https://github.com/apache/datafusion/pull/12495#discussion_r1765310035


##########
datafusion/substrait/tests/cases/serialize.rs:
##########
@@ -63,6 +67,103 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn include_remaps_for_projects() -> Result<()> {
+        let ctx = create_context().await?;
+        let df = ctx.sql("SELECT b, a + a, a FROM data").await?;
+        let datafusion_plan = df.into_optimized_plan()?;
+
+        assert_eq!(
+            format!("{}", datafusion_plan),
+            "Projection: data.b, data.a + data.a, data.a\
+            \n  TableScan: data projection=[a, b]",
+        );
+
+        let plan = to_substrait_plan(&datafusion_plan, &ctx)?.as_ref().clone();
+
+        let relation = plan.relations.first().unwrap().rel_type.as_ref();
+        let root_rel = match relation {
+            Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
+            _ => panic!("expected Root"),
+        };
+        if let Some(rel::RelType::Project(p)) = root_rel.rel_type.as_ref() {
+            // The input has 2 columns [a, b], the Projection has 3 expressions [b, a + a, a]
+            // The required output mapping is [2,3,4], which skips the 2 input columns.
+            assert_emit(p.common.as_ref(), vec![2, 3, 4]);
+
+            if let Some(rel::RelType::Read(r)) =
+                p.input.as_ref().unwrap().rel_type.as_ref()
+            {
+                let mask_expression = r.projection.as_ref().unwrap();
+                let select = mask_expression.select.as_ref().unwrap();
+                assert_eq!(
+                    2,
+                    select.struct_items.len(),
+                    "Read outputs two columns: a, b"
+                );
+                return Ok(());
+            }
+        }
+        panic!("plan did not match expected structure")
+    }
+
+    #[tokio::test]
+    async fn include_remaps_for_windows() -> Result<()> {
+        let ctx = create_context().await?;
+        // let df = ctx.sql("SELECT a, b, lead(b) OVER (PARTITION BY a) FROM data").await?;
+        let df = ctx
+            .sql("SELECT b, RANK() OVER (PARTITION BY a), c FROM data;")
+            .await?;
+        let datafusion_plan = df.into_optimized_plan()?;
+        assert_eq!(
+            format!("{}", datafusion_plan),
+            "Projection: data.b, RANK() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING, data.c\
+            \n  WindowAggr: windowExpr=[[RANK() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]\
+            \n    TableScan: data projection=[a, b, c]",
+        );
+
+        let plan = to_substrait_plan(&datafusion_plan, &ctx)?.as_ref().clone();
+
+        let relation = plan.relations.first().unwrap().rel_type.as_ref();
+        let root_rel = match relation {
+            Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
+            _ => panic!("expected Root"),
+        };
+
+        if let Some(rel::RelType::Project(p1)) = root_rel.rel_type.as_ref() {
+            // The WindowAggr outputs 4 columns, the Projection has 4 columns

Review Comment:
   I think both are 3 columns?



-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #12495:
URL: https://github.com/apache/datafusion/pull/12495#issuecomment-2358205561

   Thanks again @vbarua 


-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #12495:
URL: https://github.com/apache/datafusion/pull/12495


-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat(substrait): set ProjectRel output_mapping in producer [datafusion]

Posted by "vbarua (via GitHub)" <gi...@apache.org>.
vbarua commented on code in PR #12495:
URL: https://github.com/apache/datafusion/pull/12495#discussion_r1765349366


##########
datafusion/substrait/tests/cases/serialize.rs:
##########
@@ -63,6 +67,103 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn include_remaps_for_projects() -> Result<()> {
+        let ctx = create_context().await?;
+        let df = ctx.sql("SELECT b, a + a, a FROM data").await?;
+        let datafusion_plan = df.into_optimized_plan()?;
+
+        assert_eq!(
+            format!("{}", datafusion_plan),
+            "Projection: data.b, data.a + data.a, data.a\
+            \n  TableScan: data projection=[a, b]",
+        );
+
+        let plan = to_substrait_plan(&datafusion_plan, &ctx)?.as_ref().clone();
+
+        let relation = plan.relations.first().unwrap().rel_type.as_ref();
+        let root_rel = match relation {
+            Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
+            _ => panic!("expected Root"),
+        };
+        if let Some(rel::RelType::Project(p)) = root_rel.rel_type.as_ref() {
+            // The input has 2 columns [a, b], the Projection has 3 expressions [b, a + a, a]
+            // The required output mapping is [2,3,4], which skips the 2 input columns.
+            assert_emit(p.common.as_ref(), vec![2, 3, 4]);
+
+            if let Some(rel::RelType::Read(r)) =
+                p.input.as_ref().unwrap().rel_type.as_ref()
+            {
+                let mask_expression = r.projection.as_ref().unwrap();
+                let select = mask_expression.select.as_ref().unwrap();
+                assert_eq!(
+                    2,
+                    select.struct_items.len(),
+                    "Read outputs two columns: a, b"
+                );
+                return Ok(());
+            }
+        }
+        panic!("plan did not match expected structure")
+    }
+
+    #[tokio::test]
+    async fn include_remaps_for_windows() -> Result<()> {
+        let ctx = create_context().await?;
+        // let df = ctx.sql("SELECT a, b, lead(b) OVER (PARTITION BY a) FROM data").await?;
+        let df = ctx
+            .sql("SELECT b, RANK() OVER (PARTITION BY a), c FROM data;")
+            .await?;
+        let datafusion_plan = df.into_optimized_plan()?;
+        assert_eq!(
+            format!("{}", datafusion_plan),
+            "Projection: data.b, RANK() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING, data.c\
+            \n  WindowAggr: windowExpr=[[RANK() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]\
+            \n    TableScan: data projection=[a, b, c]",
+        );
+
+        let plan = to_substrait_plan(&datafusion_plan, &ctx)?.as_ref().clone();
+
+        let relation = plan.relations.first().unwrap().rel_type.as_ref();
+        let root_rel = match relation {
+            Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
+            _ => panic!("expected Root"),
+        };
+
+        if let Some(rel::RelType::Project(p1)) = root_rel.rel_type.as_ref() {
+            // The WindowAggr outputs 4 columns, the Projection has 4 columns

Review Comment:
   Ah, that is a typo. It should have been
   ```
   The WindowAggr outputs 4 columns, the Projection has 3 columns
                                                       ^^^
   ```
   which is what the assert below is actually (correctly) checking.
   
   I'll make a note to myself to update this when I make the consumer changes.



-- 
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@datafusion.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org