[HSTACK] - inlinetablescan top-down instead of bottom-up

Author: aditanase

datafusion/optimizer/src/analyzer/inline_table_scan.rs

@@ -38,7 +38,7 @@ impl InlineTableScan {
 
 impl AnalyzerRule for InlineTableScan {
     fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) -> Result<LogicalPlan> {
-        plan.transform_up(analyze_internal).data()
+        plan.transform_down_with_subqueries(analyze_internal).data()
     }
 
     fn name(&self) -> &str {
@@ -47,11 +47,6 @@ impl AnalyzerRule for InlineTableScan {
 }
 
 fn analyze_internal(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
-    // rewrite any subqueries in the plan first
-    let transformed_plan =
-        plan.map_subqueries(|plan| plan.transform_up(analyze_internal))?;
-
-    let transformed_plan = transformed_plan.transform_data(|plan| {
     match plan {
         // Match only on scans without filter / projection / fetch
         // Views and DataFrames won't have those added
@@ -61,6 +56,7 @@ fn analyze_internal(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
                 let sub_plan = sub_plan.into_owned();
                 let projection_exprs =
                     generate_projection_expr(&table_scan.projection, &sub_plan)?;
+
                 LogicalPlanBuilder::from(sub_plan)
                     .project(projection_exprs)?
                     // Ensures that the reference to the inlined table remains the
@@ -75,9 +71,6 @@ fn analyze_internal(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
         }
         _ => Ok(Transformed::no(plan)),
     }
-    })?;
-
-    Ok(transformed_plan)
 }
 
 fn generate_projection_expr(
@@ -202,4 +195,94 @@ mod tests {
 
         assert_analyzed_plan_eq(Arc::new(InlineTableScan::new()), plan, expected)
     }
+
+    #[derive(Debug)]
+    // stand-in for DataFrameTableProvider which we can't access here
+    struct WrappingSource {
+        plan: LogicalPlan,
+    }
+
+    impl TableSource for WrappingSource {
+        fn as_any(&self) -> &dyn std::any::Any {
+            self
+        }
+
+        fn schema(&self) -> arrow::datatypes::SchemaRef {
+            let schema: Schema = self.plan.schema().as_ref().into();
+            Arc::new(schema)
+        }
+
+        fn supports_filters_pushdown(
+            &self,
+            filters: &[&Expr],
+        ) -> datafusion_common::Result<Vec<datafusion_expr::TableProviderFilterPushDown>>
+        {
+            // A filter is added on the DataFrame when given
+            Ok(vec![
+                datafusion_expr::TableProviderFilterPushDown::Exact;
+                filters.len()
+            ])
+        }
+
+        fn get_logical_plan(&self) -> Option<Cow<LogicalPlan>> {
+            Some(Cow::Borrowed(&self.plan))
+        }
+    }
+
+    #[test]
+    fn inline_table_scan_wrapped_once() -> datafusion_common::Result<()> {
+        let custom_source = CustomSource::new();
+        let wrapped_source = WrappingSource {
+            plan: LogicalPlanBuilder::scan("wrapped", Arc::new(custom_source), None)?
+                .build()?,
+        };
+
+        let scan =
+            LogicalPlanBuilder::scan("x".to_string(), Arc::new(wrapped_source), None)?;
+        let plan = scan.build()?;
+        let expected = "SubqueryAlias: x\
+        \n  Projection: *\
+        \n    SubqueryAlias: wrapped\
+        \n      Projection: *\
+        \n        TableScan: y";
+
+        assert_analyzed_plan_eq(Arc::new(InlineTableScan::new()), plan, expected)
+    }
+
+    #[test]
+    fn inline_table_scan_wrapped_twice_with_projection() -> datafusion_common::Result<()> {
+        let custom_source = CustomSource::new();
+        let wrapped_source_once = WrappingSource {
+            plan: LogicalPlanBuilder::scan(
+                "wrapped_once",
+                Arc::new(custom_source),
+                Some(vec![0]),
+            )?
+            .build()?,
+        };
+        let wrapped_source_twice = WrappingSource {
+            plan: LogicalPlanBuilder::scan(
+                "wrapped_twice",
+                Arc::new(wrapped_source_once),
+                None,
+            )?
+            .build()?,
+        };
+
+        let scan = LogicalPlanBuilder::scan(
+            "x".to_string(),
+            Arc::new(wrapped_source_twice),
+            None,
+        )?;
+        let plan = scan.build()?;
+        let expected = "SubqueryAlias: x\
+        \n  Projection: *\
+        \n    SubqueryAlias: wrapped_twice\
+        \n      Projection: *\
+        \n        SubqueryAlias: wrapped_once\
+        \n          Projection: y.a\
+        \n            TableScan: y";
+
+        assert_analyzed_plan_eq(Arc::new(InlineTableScan::new()), plan, expected)
+    }
 }