[arrow-datafusion] branch main updated: minor: do not fail analyzer if subquery plan contains extension (#7455)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 4e33be5e12 minor: do not fail analyzer if subquery plan contains extension (#7455)
4e33be5e12 is described below

commit 4e33be5e12726cd127405e586246f2c73007e02f
Author: Ruihang Xia <wa...@gmail.com>
AuthorDate: Wed Sep 6 12:46:36 2023 -0500

    minor: do not fail analyzer if subquery plan contains extension (#7455)
    
    * do not fail analyzer if subquery plan contains extension
    
    Signed-off-by: Ruihang Xia <wa...@gmail.com>
    
    * add a simple UT case
    
    Signed-off-by: Ruihang Xia <wa...@gmail.com>
    
    ---------
    
    Signed-off-by: Ruihang Xia <wa...@gmail.com>
---
 datafusion/optimizer/src/analyzer/subquery.rs | 55 +++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/datafusion/optimizer/src/analyzer/subquery.rs b/datafusion/optimizer/src/analyzer/subquery.rs
index fee6377114..6b8b1020cd 100644
--- a/datafusion/optimizer/src/analyzer/subquery.rs
+++ b/datafusion/optimizer/src/analyzer/subquery.rs
@@ -225,6 +225,7 @@ fn check_inner_plan(
                 Ok(())
             }
         },
+        LogicalPlan::Extension(_) => Ok(()),
         _ => plan_err!("Unsupported operator in the subquery plan."),
     }
 }
@@ -338,3 +339,57 @@ fn check_mixed_out_refer_in_window(window: &Window) -> Result<()> {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod test {
+    use std::sync::Arc;
+
+    use datafusion_common::{DFSchema, DFSchemaRef};
+    use datafusion_expr::{Extension, UserDefinedLogicalNodeCore};
+
+    use super::*;
+
+    #[derive(Debug, PartialEq, Eq, Hash)]
+    struct MockUserDefinedLogicalPlan {
+        empty_schema: DFSchemaRef,
+    }
+
+    impl UserDefinedLogicalNodeCore for MockUserDefinedLogicalPlan {
+        fn name(&self) -> &str {
+            "MockUserDefinedLogicalPlan"
+        }
+
+        fn inputs(&self) -> Vec<&LogicalPlan> {
+            vec![]
+        }
+
+        fn schema(&self) -> &datafusion_common::DFSchemaRef {
+            &self.empty_schema
+        }
+
+        fn expressions(&self) -> Vec<Expr> {
+            vec![]
+        }
+
+        fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+            write!(f, "MockUserDefinedLogicalPlan")
+        }
+
+        fn from_template(&self, _exprs: &[Expr], _inputs: &[LogicalPlan]) -> Self {
+            Self {
+                empty_schema: self.empty_schema.clone(),
+            }
+        }
+    }
+
+    #[test]
+    fn wont_fail_extension_plan() {
+        let plan = LogicalPlan::Extension(Extension {
+            node: Arc::new(MockUserDefinedLogicalPlan {
+                empty_schema: DFSchemaRef::new(DFSchema::empty()),
+            }),
+        });
+
+        check_inner_plan(&plan, false, false, true).unwrap();
+    }
+}