fix: return schema of ExtensionPlan instead of its children's (#5514)

* fix: return schema of ExtensionPlan instead of its children's Signed-off-by: Ruihang Xia <[email protected]> * add test case Signed-off-by: Ruihang Xia <[email protected]> --------- Signed-off-by: Ruihang Xia <[email protected]>
apache · Mar 9, 2023 · 6a8c4a6 · 6a8c4a6
1 parent 55c7c4d
commit 6a8c4a6
Showing 1 changed file with 68 additions and 1 deletion.
diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
@@ -201,13 +201,13 @@ impl LogicalPlan {
             | LogicalPlan::Values(_)
             | LogicalPlan::SubqueryAlias(_)
             | LogicalPlan::Union(_)
+            | LogicalPlan::Extension(_)
             | LogicalPlan::TableScan(_) => {
                 vec![self.schema()]
             }
             // return children schemas
             LogicalPlan::Limit(_)
             | LogicalPlan::Subquery(_)
-            | LogicalPlan::Extension(_)
             | LogicalPlan::Repartition(_)
             | LogicalPlan::Sort(_)
             | LogicalPlan::CreateMemoryTable(_)
@@ -2381,4 +2381,71 @@ mod tests {
             .build()
             .unwrap()
     }
+
+    /// Extension plan that panic when trying to access its input plan
+    #[derive(Debug)]
+    struct NoChildExtension {
+        empty_schema: DFSchemaRef,
+    }
+
+    impl NoChildExtension {
+        fn empty() -> Self {
+            Self {
+                empty_schema: Arc::new(DFSchema::empty()),
+            }
+        }
+    }
+
+    impl UserDefinedLogicalNode for NoChildExtension {
+        fn as_any(&self) -> &dyn std::any::Any {
+            unimplemented!()
+        }
+
+        fn name(&self) -> &str {
+            unimplemented!()
+        }
+
+        fn inputs(&self) -> Vec<&LogicalPlan> {
+            panic!("Should not be called")
+        }
+
+        fn schema(&self) -> &DFSchemaRef {
+            &self.empty_schema
+        }
+
+        fn expressions(&self) -> Vec<Expr> {
+            unimplemented!()
+        }
+
+        fn fmt_for_explain(&self, _: &mut fmt::Formatter) -> fmt::Result {
+            unimplemented!()
+        }
+
+        fn from_template(
+            &self,
+            _: &[Expr],
+            _: &[LogicalPlan],
+        ) -> Arc<dyn UserDefinedLogicalNode> {
+            unimplemented!()
+        }
+
+        fn dyn_hash(&self, _: &mut dyn Hasher) {
+            unimplemented!()
+        }
+
+        fn dyn_eq(&self, _: &dyn UserDefinedLogicalNode) -> bool {
+            unimplemented!()
+        }
+    }
+
+    #[test]
+    fn test_extension_all_schemas() {
+        let plan = LogicalPlan::Extension(Extension {
+            node: Arc::new(NoChildExtension::empty()),
+        });
+
+        let schemas = plan.all_schemas();
+        assert_eq!(1, schemas.len());
+        assert_eq!(0, schemas[0].fields().len());
+    }
 }