refactor(optimizer): add .with_schema for defining test tables

Added `tables: HashMap<String, Arc<dyn TableSource>>` and `MyContextProvider::with_schema` method for dynamically defining tables for optimizer integration tests.

Author: atahanyorganci

datafusion/optimizer/tests/optimizer_integration.rs

@@ -492,7 +492,32 @@ fn test_sql(sql: &str) -> Result<LogicalPlan> {
         .with_expr_planners(vec![
             Arc::new(AggregateFunctionPlanner),
             Arc::new(WindowFunctionPlanner),
-        ]);
+        ])
+        .with_schema(
+            "test",
+            Schema::new_with_metadata(
+                vec![
+                    Field::new("col_int32", DataType::Int32, true),
+                    Field::new("col_uint32", DataType::UInt32, true),
+                    Field::new("col_utf8", DataType::Utf8, true),
+                    Field::new("col_date32", DataType::Date32, true),
+                    Field::new("col_date64", DataType::Date64, true),
+                    // timestamp with no timezone
+                    Field::new(
+                        "col_ts_nano_none",
+                        DataType::Timestamp(TimeUnit::Nanosecond, None),
+                        true,
+                    ),
+                    // timestamp with UTC timezone
+                    Field::new(
+                        "col_ts_nano_utc",
+                        DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
+                        true,
+                    ),
+                ],
+                HashMap::new(),
+            ),
+        );
 
     let sql_to_rel = SqlToRel::new(&context_provider);
     let plan = sql_to_rel.sql_statement_to_plan(statement.clone())?;
@@ -510,6 +535,7 @@ fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
 #[derive(Default)]
 struct MyContextProvider {
     options: ConfigOptions,
+    tables: HashMap<String, Arc<dyn TableSource>>,
     udafs: HashMap<String, Arc<AggregateUDF>>,
     expr_planners: Vec<Arc<dyn ExprPlanner>>,
 }
@@ -525,38 +551,23 @@ impl MyContextProvider {
         self.expr_planners = expr_planners;
         self
     }
+
+    fn with_schema(mut self, name: impl Into<String>, schema: Schema) -> Self {
+        self.tables.insert(
+            name.into(),
+            Arc::new(MyTableSource {
+                schema: Arc::new(schema),
+            }),
+        );
+        self
+    }
 }
 
 impl ContextProvider for MyContextProvider {
     fn get_table_source(&self, name: TableReference) -> Result<Arc<dyn TableSource>> {
         let table_name = name.table();
-        if table_name.starts_with("test") {
-            let schema = Schema::new_with_metadata(
-                vec![
-                    Field::new("col_int32", DataType::Int32, true),
-                    Field::new("col_uint32", DataType::UInt32, true),
-                    Field::new("col_utf8", DataType::Utf8, true),
-                    Field::new("col_date32", DataType::Date32, true),
-                    Field::new("col_date64", DataType::Date64, true),
-                    // timestamp with no timezone
-                    Field::new(
-                        "col_ts_nano_none",
-                        DataType::Timestamp(TimeUnit::Nanosecond, None),
-                        true,
-                    ),
-                    // timestamp with UTC timezone
-                    Field::new(
-                        "col_ts_nano_utc",
-                        DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
-                        true,
-                    ),
-                ],
-                HashMap::new(),
-            );
-
-            Ok(Arc::new(MyTableSource {
-                schema: Arc::new(schema),
-            }))
+        if let Some(table) = self.tables.get(table_name) {
+            Ok(table.clone())
         } else {
             plan_err!("table does not exist")
         }