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: [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]
