add test

Author: niebayes

datafusion/core/src/execution/session_state.rs

@@ -1955,6 +1955,16 @@ mod tests {
     use super::{SessionContextProvider, SessionStateBuilder};
     use crate::datasource::MemTable;
     use crate::execution::context::SessionState;
+    use crate::{
+        common::assert_contains,
+        config::ConfigOptions,
+        datasource::{empty::EmptyTable, provider_as_source},
+        logical_expr::{
+            planner::ExprPlanner, AggregateUDF, ScalarUDF, TableSource, WindowUDF,
+        },
+        physical_plan::ExecutionPlan,
+        sql::{planner::ContextProvider, ResolvedTableReference, TableReference},
+    };
     use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray};
     use arrow::datatypes::{DataType, Field, Schema};
     use datafusion_catalog::MemoryCatalogProviderList;
@@ -1964,6 +1974,7 @@ mod tests {
     use datafusion_expr::Expr;
     use datafusion_optimizer::optimizer::OptimizerRule;
     use datafusion_optimizer::Optimizer;
+    use datafusion_physical_plan::display::DisplayableExecutionPlan;
     use datafusion_sql::planner::{PlannerContext, SqlToRel};
     use std::collections::HashMap;
     use std::sync::Arc;
@@ -2121,4 +2132,148 @@ mod tests {
 
         Ok(())
     }
+
+    /// This test demonstrates why it's more convenient and somewhat necessary to provide
+    /// an `expr_planners` method for `SessionState`.
+    #[tokio::test]
+    async fn test_with_expr_planners() -> Result<()> {
+        // A helper method for planning count wildcard with or without expr planners.
+        async fn plan_count_wildcard(
+            with_expr_planners: bool,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            let mut context_provider = MyContextProvider::new().with_table(
+                "t",
+                provider_as_source(Arc::new(EmptyTable::new(Schema::empty().into()))),
+            );
+            if with_expr_planners {
+                context_provider = context_provider.with_expr_planners();
+            }
+
+            let state = &context_provider.state;
+            let statement = state.sql_to_statement("select count(*) from t", "mysql")?;
+            let plan = SqlToRel::new(&context_provider).statement_to_plan(statement)?;
+            state.create_physical_plan(&plan).await
+        }
+
+        // Planning count wildcard without expr planners should fail.
+        let got = plan_count_wildcard(false).await;
+        assert_contains!(
+            got.unwrap_err().to_string(),
+            "Physical plan does not support logical expression Wildcard"
+        );
+
+        // Planning count wildcard with expr planners should succeed.
+        let got = plan_count_wildcard(true).await?;
+        let displayable = DisplayableExecutionPlan::new(got.as_ref());
+        assert_eq!(
+            displayable.indent(false).to_string(),
+            "ProjectionExec: expr=[0 as count(*)]\n  PlaceholderRowExec\n"
+        );
+
+        Ok(())
+    }
+
+    /// A `ContextProvider` based on `SessionState`.
+    ///
+    /// Almost all planning context are retrieved from the `SessionState`.
+    struct MyContextProvider {
+        /// The session state.
+        state: SessionState,
+        /// Registered tables.
+        tables: HashMap<ResolvedTableReference, Arc<dyn TableSource>>,
+        /// Controls whether to return expression planners when called `ContextProvider::expr_planners`.
+        return_expr_planners: bool,
+    }
+
+    impl MyContextProvider {
+        /// Creates a new `SessionContextProvider`.
+        pub fn new() -> Self {
+            Self {
+                state: SessionStateBuilder::default()
+                    .with_default_features()
+                    .build(),
+                tables: HashMap::new(),
+                return_expr_planners: false,
+            }
+        }
+
+        /// Registers a table.
+        ///
+        /// The catalog and schema are provided by default.
+        pub fn with_table(mut self, table: &str, source: Arc<dyn TableSource>) -> Self {
+            self.tables.insert(
+                ResolvedTableReference {
+                    catalog: "default".to_string().into(),
+                    schema: "public".to_string().into(),
+                    table: table.to_string().into(),
+                },
+                source,
+            );
+            self
+        }
+
+        /// Sets the `return_expr_planners` flag to true.
+        pub fn with_expr_planners(self) -> Self {
+            Self {
+                return_expr_planners: true,
+                ..self
+            }
+        }
+    }
+
+    impl ContextProvider for MyContextProvider {
+        fn get_table_source(&self, name: TableReference) -> Result<Arc<dyn TableSource>> {
+            let resolved_table_ref = ResolvedTableReference {
+                catalog: "default".to_string().into(),
+                schema: "public".to_string().into(),
+                table: name.table().to_string().into(),
+            };
+            let source = self.tables.get(&resolved_table_ref).cloned().unwrap();
+            Ok(source)
+        }
+
+        /// We use a `return_expr_planners` flag to demonstrate why it's necessary to
+        /// return the expression planners in the `SessionState`.
+        ///
+        /// Note, the default implementation returns an empty slice.
+        fn get_expr_planners(&self) -> &[Arc<dyn ExprPlanner>] {
+            if self.return_expr_planners {
+                self.state.expr_planners()
+            } else {
+                &[]
+            }
+        }
+
+        fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
+            self.state.scalar_functions().get(name).cloned()
+        }
+
+        fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
+            self.state.aggregate_functions().get(name).cloned()
+        }
+
+        fn get_window_meta(&self, name: &str) -> Option<Arc<WindowUDF>> {
+            self.state.window_functions().get(name).cloned()
+        }
+
+        fn get_variable_type(&self, _variable_names: &[String]) -> Option<DataType> {
+            None
+        }
+
+        fn options(&self) -> &ConfigOptions {
+            self.state.config_options()
+        }
+
+        fn udf_names(&self) -> Vec<String> {
+            self.state.scalar_functions().keys().cloned().collect()
+        }
+
+        fn udaf_names(&self) -> Vec<String> {
+            self.state.aggregate_functions().keys().cloned().collect()
+        }
+
+        fn udwf_names(&self) -> Vec<String> {
+            self.state.window_functions().keys().cloned().collect()
+        }
+    }
 }