query/datafusion/
planner.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::collections::hash_map::Entry;
17use std::sync::Arc;
18
19use arrow_schema::DataType;
20use catalog::table_source::DfTableSourceProvider;
21use common_function::function::FunctionContext;
22use datafusion::catalog::TableFunctionArgs;
23use datafusion::common::{DFSchema, TableReference};
24use datafusion::datasource::cte_worktable::CteWorkTable;
25use datafusion::datasource::file_format::{FileFormatFactory, format_as_file_type};
26use datafusion::datasource::provider_as_source;
27use datafusion::error::Result as DfResult;
28use datafusion::execution::SessionStateDefaults;
29use datafusion::execution::context::SessionState;
30use datafusion::sql::planner::ContextProvider;
31use datafusion::variable::VarType;
32use datafusion_common::DataFusionError;
33use datafusion_common::config::ConfigOptions;
34use datafusion_common::file_options::file_type::FileType;
35use datafusion_expr::planner::{ExprPlanner, TypePlanner};
36use datafusion_expr::var_provider::is_system_variables;
37use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF};
38use datafusion_sql::parser::Statement as DfStatement;
39use session::context::QueryContextRef;
40use snafu::{Location, ResultExt};
41
42use crate::datafusion::json_expr_planner::JsonExprPlanner;
43use crate::datafusion::pg_oid_alias_expr_planner::PgOidAliasExprPlanner;
44use crate::error::{CatalogSnafu, Result};
45use crate::query_engine::{DefaultPlanDecoder, QueryEngineState};
46
47pub struct DfContextProviderAdapter {
48    engine_state: Arc<QueryEngineState>,
49    session_state: SessionState,
50    tables: HashMap<String, Arc<dyn TableSource>>,
51    table_provider: DfTableSourceProvider,
52    query_ctx: QueryContextRef,
53
54    // Fields from session state defaults:
55    /// Holds registered external FileFormat implementations
56    /// DataFusion doesn't pub this field, so we need to store it here.
57    file_formats: HashMap<String, Arc<dyn FileFormatFactory>>,
58    /// Provides support for customising the SQL planner, e.g. to add support for custom operators like `->>` or `?`
59    /// DataFusion doesn't pub this field, so we need to store it here.
60    expr_planners: Vec<Arc<dyn ExprPlanner>>,
61}
62
63impl DfContextProviderAdapter {
64    pub(crate) async fn try_new(
65        engine_state: Arc<QueryEngineState>,
66        session_state: SessionState,
67        df_stmt: Option<&DfStatement>,
68        query_ctx: QueryContextRef,
69    ) -> Result<Self> {
70        let table_names = if let Some(df_stmt) = df_stmt {
71            session_state.resolve_table_references(df_stmt)?
72        } else {
73            vec![]
74        };
75
76        let mut table_provider = DfTableSourceProvider::new(
77            engine_state.catalog_manager().clone(),
78            engine_state.disallow_cross_catalog_query(),
79            query_ctx.clone(),
80            Arc::new(DefaultPlanDecoder::new(session_state.clone(), &query_ctx)?),
81            session_state
82                .config_options()
83                .sql_parser
84                .enable_ident_normalization,
85        );
86
87        let tables = resolve_tables(table_names, &mut table_provider).await?;
88        let file_formats = SessionStateDefaults::default_file_formats()
89            .into_iter()
90            .map(|format| (format.get_ext().to_lowercase(), format))
91            .collect();
92
93        let mut expr_planners = SessionStateDefaults::default_expr_planners();
94        expr_planners.insert(0, Arc::new(JsonExprPlanner));
95        expr_planners.insert(0, Arc::new(PgOidAliasExprPlanner));
96
97        Ok(Self {
98            engine_state,
99            session_state,
100            tables,
101            table_provider,
102            query_ctx,
103            file_formats,
104            expr_planners,
105        })
106    }
107}
108
109async fn resolve_tables(
110    table_names: Vec<TableReference>,
111    table_provider: &mut DfTableSourceProvider,
112) -> Result<HashMap<String, Arc<dyn TableSource>>> {
113    let mut tables = HashMap::with_capacity(table_names.len());
114
115    for table_name in table_names {
116        let resolved_name = table_provider
117            .resolve_table_ref(table_name.clone())
118            .context(CatalogSnafu)?;
119
120        if let Entry::Vacant(v) = tables.entry(resolved_name.to_string()) {
121            // Try our best to resolve the tables here, but we don't return an error if table is not found,
122            // because the table name may be a temporary name of CTE, they can't be found until plan
123            // execution.
124            match table_provider.resolve_table(table_name).await {
125                Ok(table) => {
126                    let _ = v.insert(table);
127                }
128                Err(e) if e.should_fail() => {
129                    return Err(e).context(CatalogSnafu);
130                }
131                _ => {
132                    // ignore
133                }
134            }
135        }
136    }
137    Ok(tables)
138}
139
140impl ContextProvider for DfContextProviderAdapter {
141    fn get_table_source(&self, name: TableReference) -> DfResult<Arc<dyn TableSource>> {
142        let table_ref = self.table_provider.resolve_table_ref(name)?;
143        self.tables
144            .get(&table_ref.to_string())
145            .cloned()
146            .ok_or_else(|| {
147                crate::error::Error::TableNotFound {
148                    table: table_ref.to_string(),
149                    location: Location::default(),
150                }
151                .into()
152            })
153    }
154
155    fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
156        self.engine_state.scalar_function(name).map_or_else(
157            || self.session_state.scalar_functions().get(name).cloned(),
158            |func| {
159                Some(Arc::new(func.provide(FunctionContext {
160                    query_ctx: self.query_ctx.clone(),
161                    state: self.engine_state.function_state(),
162                })))
163            },
164        )
165    }
166
167    fn get_higher_order_meta(&self, name: &str) -> Option<Arc<HigherOrderUDF>> {
168        self.session_state
169            .higher_order_functions()
170            .get(name)
171            .cloned()
172    }
173
174    fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
175        self.engine_state.aggr_function(name).map_or_else(
176            || self.session_state.aggregate_functions().get(name).cloned(),
177            |func| Some(Arc::new(func)),
178        )
179    }
180
181    fn get_window_meta(&self, name: &str) -> Option<Arc<WindowUDF>> {
182        self.session_state.window_functions().get(name).cloned()
183    }
184
185    fn get_variable_type(&self, variable_names: &[String]) -> Option<DataType> {
186        if variable_names.is_empty() {
187            return None;
188        }
189
190        let provider_type = if is_system_variables(variable_names) {
191            VarType::System
192        } else {
193            VarType::UserDefined
194        };
195
196        self.session_state
197            .execution_props()
198            .var_providers
199            .as_ref()
200            .and_then(|provider| provider.get(&provider_type)?.get_type(variable_names))
201    }
202
203    fn options(&self) -> &ConfigOptions {
204        self.session_state.config_options()
205    }
206
207    fn udf_names(&self) -> Vec<String> {
208        let mut names = self.engine_state.scalar_names();
209        names.extend(self.session_state.scalar_functions().keys().cloned());
210        names
211    }
212
213    fn higher_order_function_names(&self) -> Vec<String> {
214        self.session_state
215            .higher_order_functions()
216            .keys()
217            .cloned()
218            .collect()
219    }
220
221    fn udaf_names(&self) -> Vec<String> {
222        let mut names = self.engine_state.aggr_names();
223        names.extend(self.session_state.aggregate_functions().keys().cloned());
224        names
225    }
226
227    fn udwf_names(&self) -> Vec<String> {
228        self.session_state
229            .window_functions()
230            .keys()
231            .cloned()
232            .collect()
233    }
234
235    fn get_file_type(&self, ext: &str) -> DfResult<Arc<dyn FileType>> {
236        self.file_formats
237            .get(&ext.to_lowercase())
238            .ok_or_else(|| {
239                DataFusionError::Plan(format!("There is no registered file format with ext {ext}"))
240            })
241            .map(|file_type| format_as_file_type(Arc::clone(file_type)))
242    }
243
244    fn get_table_function_source(
245        &self,
246        name: &str,
247        args: Vec<datafusion_expr::Expr>,
248    ) -> DfResult<Arc<dyn TableSource>> {
249        // Constant-fold the args before resolving the table function. DataFusion's
250        // SQL planner does not fold table-function arguments (constant folding
251        // happens later, in the analyzer), but table functions such as
252        // `generate_series`/`range` are resolved during planning and require
253        // literal bounds. Folding here lets immutable-UDF bounds like
254        // `array_upper(ARRAY[...], 1)` reach them as concrete literals.
255        // Non-constant args are returned unchanged by the simplifier.
256        let simplify_info = datafusion_expr::simplify::SimplifyContext::builder()
257            .with_config_options(Arc::clone(self.session_state.config_options()))
258            .with_query_execution_start_time(
259                self.session_state
260                    .execution_props()
261                    .query_execution_start_time,
262            )
263            .build();
264        let simplifier =
265            datafusion_optimizer::simplify_expressions::ExprSimplifier::new(simplify_info);
266        let schema = DFSchema::empty();
267        let args = args
268            .into_iter()
269            .map(|arg| {
270                simplifier
271                    .coerce(arg, &schema)
272                    .and_then(|arg| simplifier.simplify(arg))
273            })
274            .collect::<DfResult<Vec<_>>>()?;
275        let table_args = TableFunctionArgs::new(&args, &self.session_state);
276        let tbl_func = if let Some(tbl_func) = self.engine_state.table_function(name) {
277            tbl_func
278        } else {
279            self.session_state
280                .table_functions()
281                .get(name)
282                .cloned()
283                .ok_or_else(|| {
284                    DataFusionError::Plan(format!("table function '{name}' not found"))
285                })?
286        };
287        let provider = tbl_func.create_table_provider_with_args(table_args)?;
288
289        Ok(provider_as_source(provider))
290    }
291
292    fn create_cte_work_table(
293        &self,
294        name: &str,
295        schema: arrow_schema::SchemaRef,
296    ) -> DfResult<Arc<dyn TableSource>> {
297        let table = Arc::new(CteWorkTable::new(name, schema));
298        Ok(provider_as_source(table))
299    }
300
301    fn get_expr_planners(&self) -> &[Arc<dyn ExprPlanner>] {
302        &self.expr_planners
303    }
304
305    fn get_type_planner(&self) -> Option<Arc<dyn TypePlanner>> {
306        // Provide the SQL planner with Postgres oid-alias type names
307        // (`regclass`, `regproc`, `regtype`, `regnamespace`, `oid`, ...) and
308        // `pg_catalog.`-qualified builtins. DataFusion rejects these as
309        // "Unsupported SQL type" otherwise. The planner maps each to its Arrow
310        // type so reverse / column-operand casts like `prorettype::regtype::text`
311        // parse. Forward name->oid casts (`'x'::regclass`) are resolved earlier,
312        // at SQL-parse time, by the `PostgresCompatibilityParser`'s built-in
313        // `RewriteRegCastToSubquery` rule.
314        // Stateless, so a fresh instance per query is cheap.
315        Some(Arc::new(
316            datafusion_pg_catalog::pg_catalog::oid_type_planner::PgOidTypePlanner,
317        ))
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use std::sync::atomic::{AtomicBool, Ordering};
324    use std::sync::{Arc, Mutex};
325
326    use common_base::Plugins;
327    use datafusion::catalog::{TableFunction, TableFunctionArgs, TableFunctionImpl, TableProvider};
328    use datafusion::datasource::MemTable;
329    use datafusion::execution::SessionStateBuilder;
330    use datafusion::execution::context::{SessionConfig, SessionContext, SessionState};
331    use datafusion_common::ScalarValue;
332    use datafusion_expr::expr::BinaryExpr;
333    use datafusion_expr::{Expr, Operator, lit};
334    use session::context::QueryContext;
335
336    use super::*;
337    use crate::options::QueryOptions;
338
339    #[derive(Debug, Default)]
340    struct RecordingTableFunction {
341        called: AtomicBool,
342        args: Mutex<Vec<Expr>>,
343        target_partitions: Mutex<Option<usize>>,
344    }
345
346    impl TableFunctionImpl for RecordingTableFunction {
347        fn call_with_args(&self, args: TableFunctionArgs) -> DfResult<Arc<dyn TableProvider>> {
348            let session_state = args
349                .session()
350                .as_any()
351                .downcast_ref::<SessionState>()
352                .expect("table function must receive the SessionState");
353            *self.args.lock().unwrap() = args.exprs().to_vec();
354            *self.target_partitions.lock().unwrap() =
355                Some(session_state.config().target_partitions());
356            self.called.store(true, Ordering::SeqCst);
357
358            Ok(Arc::new(MemTable::try_new(
359                Arc::new(arrow_schema::Schema::empty()),
360                vec![vec![]],
361            )?))
362        }
363    }
364
365    fn query_engine_state() -> Arc<QueryEngineState> {
366        Arc::new(QueryEngineState::new(
367            catalog::memory::new_memory_catalog_manager().unwrap(),
368            None,
369            None,
370            None,
371            None,
372            None,
373            false,
374            Plugins::default(),
375            QueryOptions::default(),
376        ))
377    }
378
379    async fn context_provider(
380        engine_state: Arc<QueryEngineState>,
381        session_state: SessionState,
382    ) -> DfContextProviderAdapter {
383        DfContextProviderAdapter::try_new(engine_state, session_state, None, QueryContext::arc())
384            .await
385            .unwrap()
386    }
387
388    fn plus(left: Expr, right: Expr) -> Expr {
389        Expr::BinaryExpr(BinaryExpr {
390            left: Box::new(left),
391            op: Operator::Plus,
392            right: Box::new(right),
393        })
394    }
395
396    #[tokio::test]
397    async fn table_function_arguments_are_folded_before_engine_function_creation() {
398        let engine_state = query_engine_state();
399        let function = Arc::new(RecordingTableFunction::default());
400        engine_state.register_table_function(Arc::new(TableFunction::new(
401            "capture_engine_args".to_string(),
402            function.clone(),
403        )));
404        let provider = context_provider(engine_state.clone(), engine_state.session_state()).await;
405
406        provider
407            .get_table_function_source("capture_engine_args", vec![plus(lit(1_i64), lit(2_i64))])
408            .unwrap();
409
410        assert!(function.called.load(Ordering::SeqCst));
411        assert_eq!(
412            *function.args.lock().unwrap(),
413            vec![Expr::Literal(ScalarValue::Int64(Some(3)), None)]
414        );
415    }
416
417    #[tokio::test]
418    async fn table_function_argument_simplification_errors_are_propagated() {
419        let engine_state = query_engine_state();
420        let function = Arc::new(RecordingTableFunction::default());
421        engine_state.register_table_function(Arc::new(TableFunction::new(
422            "reject_invalid_args".to_string(),
423            function.clone(),
424        )));
425        let provider = context_provider(engine_state.clone(), engine_state.session_state()).await;
426
427        let error = match provider
428            .get_table_function_source("reject_invalid_args", vec![plus(lit(true), lit(1_i64))])
429        {
430            Ok(_) => panic!("invalid table-function argument must fail planning"),
431            Err(error) => error,
432        };
433
434        assert!(!error.to_string().is_empty());
435        assert!(!function.called.load(Ordering::SeqCst));
436    }
437
438    #[tokio::test]
439    async fn session_table_function_receives_table_function_args_session() {
440        let engine_state = query_engine_state();
441        let session_state = SessionStateBuilder::new_from_existing(engine_state.session_state())
442            .with_config(SessionConfig::new().with_target_partitions(7))
443            .build();
444        let session_context = SessionContext::new_with_state(session_state);
445        let function = Arc::new(RecordingTableFunction::default());
446        session_context.register_udtf("capture_session_args", function.clone());
447        let provider = context_provider(engine_state, session_context.state()).await;
448
449        provider
450            .get_table_function_source("capture_session_args", vec![lit(1_i64)])
451            .unwrap();
452
453        assert!(function.called.load(Ordering::SeqCst));
454        assert_eq!(*function.target_partitions.lock().unwrap(), Some(7));
455    }
456}
query/datafusion/planner.rs

query/datafusion/
planner.rs
