[query] refactor: add DatabendQueryContext::build_table_from_source_plan()

Author: drmingdrmer

common/meta/types/src/table_info.rs

@@ -55,7 +55,13 @@ impl Display for TableIdent {
 pub struct TableInfo {
     pub ident: TableIdent,
 
+    /// For a table it is `db_name.table_name`.
+    /// For a table function, it is `table_name(args)`
     pub desc: String,
+
+    /// `name` is meant to be used with table-function.
+    /// Table-function is identified by `name`.
+    /// A table in the contrast, can only be identified by table-id.
     pub name: String,
 
     /// The essential information about a table definition.

query/src/interpreters/plan_scheduler.rs

@@ -48,7 +48,6 @@ use common_tracing::tracing;
 use crate::api::BroadcastAction;
 use crate::api::FlightAction;
 use crate::api::ShuffleAction;
-use crate::catalogs::Catalog;
 use crate::catalogs::TablePtr;
 use crate::catalogs::ToReadDataSourcePlan;
 use crate::sessions::DatabendQueryContext;
@@ -784,14 +783,7 @@ impl PlanScheduler {
     }
 
     fn visit_data_source(&mut self, plan: &ReadDataSourcePlan, _: &mut Tasks) -> Result<()> {
-        let table = if plan.tbl_args.is_none() {
-            let catalog = self.query_context.get_catalog();
-            catalog.build_table(&plan.table_info)?
-        } else {
-            self.query_context
-                .get_table_function(&plan.table_info.name, plan.tbl_args.clone())?
-                .as_table()
-        };
+        let table = self.query_context.build_table_from_source_plan(plan)?;
 
         match table.is_local() {
             true => self.visit_local_data_source(plan),

query/src/optimizers/optimizer_scatters.rs

@@ -32,7 +32,6 @@ use common_planners::SortPlan;
 use common_planners::StageKind;
 use common_planners::StagePlan;
 
-use crate::catalogs::Catalog;
 use crate::optimizers::Optimizer;
 use crate::sessions::DatabendQueryContext;
 use crate::sessions::DatabendQueryContextRef;
@@ -281,17 +280,9 @@ impl PlanRewriter for ScattersOptimizerImpl {
     }
 
     fn rewrite_read_data_source(&mut self, plan: &ReadDataSourcePlan) -> Result<PlanNode> {
-        let context = self.ctx.clone();
-        let select_table = if plan.tbl_args.is_none() {
-            let catalog = context.get_catalog();
-            catalog.build_table(&plan.table_info)?
-        } else {
-            context
-                .get_table_function(&plan.table_info.name, plan.tbl_args.clone())?
-                .as_table()
-        };
+        let t = self.ctx.build_table_from_source_plan(plan)?;
 
-        match select_table.is_local() {
+        match t.is_local() {
             false => self.running_mode = RunningMode::Cluster,
             true => self.running_mode = RunningMode::Standalone,
         }

query/src/pipelines/transforms/transform_source.rs

@@ -23,7 +23,6 @@ use common_streams::ProgressStream;
 use common_streams::SendableDataBlockStream;
 use common_tracing::tracing;
 
-use crate::catalogs::Catalog;
 use crate::pipelines::processors::EmptyProcessor;
 use crate::pipelines::processors::Processor;
 use crate::sessions::DatabendQueryContextRef;
@@ -42,16 +41,7 @@ impl SourceTransform {
     }
 
     async fn read_table(&self) -> Result<SendableDataBlockStream> {
-        let table = if self.source_plan.tbl_args.is_none() {
-            let catalog = self.ctx.get_catalog();
-            catalog.build_table(&self.source_plan.table_info)?
-        } else {
-            let func_meta = self.ctx.get_table_function(
-                &self.source_plan.table_info.name,
-                self.source_plan.tbl_args.clone(),
-            )?;
-            func_meta.as_table()
-        };
+        let table = self.ctx.build_table_from_source_plan(&self.source_plan)?;
 
         // TODO(xp): get_single_node_table_io_context() or
         //           get_cluster_table_io_context()?

query/src/sessions/context.rs

@@ -33,18 +33,17 @@ use common_meta_types::NodeInfo;
 use common_planners::Part;
 use common_planners::Partitions;
 use common_planners::PlanNode;
+use common_planners::ReadDataSourcePlan;
 use common_planners::Statistics;
 use common_streams::AbortStream;
 use common_streams::SendableDataBlockStream;
 
 use crate::catalogs::impls::DatabaseCatalog;
 use crate::catalogs::Catalog;
 use crate::catalogs::Table;
-use crate::catalogs::TableFunction;
 use crate::clusters::ClusterRef;
 use crate::configs::Config;
 use crate::datasources::common::ContextDalBuilder;
-use crate::datasources::table_func_engine::TableArgs;
 use crate::sessions::context_shared::DatabendQueryContextShared;
 use crate::sessions::SessionManagerRef;
 use crate::sessions::Settings;
@@ -157,13 +156,25 @@ impl DatabendQueryContext {
         self.shared.get_table(database, table)
     }
 
-    pub fn get_table_function(
+    /// Build a table instance the plan wants to operate on.
+    ///
+    /// A plan just contains raw information about a table or table function.
+    /// This method builds a `dyn Table`, which provides table specific io methods the plan needs.
+    pub fn build_table_from_source_plan(
         &self,
-        function_name: &str,
-        tbl_args: TableArgs,
-    ) -> Result<Arc<dyn TableFunction>> {
-        self.get_catalog()
-            .get_table_function(function_name, tbl_args)
+        plan: &ReadDataSourcePlan,
+    ) -> Result<Arc<dyn Table>> {
+        let catalog = self.get_catalog();
+
+        let t = if plan.tbl_args.is_none() {
+            catalog.build_table(&plan.table_info)?
+        } else {
+            catalog
+                .get_table_function(&plan.table_info.name, plan.tbl_args.clone())?
+                .as_table()
+        };
+
+        Ok(t)
     }
 
     pub fn get_id(&self) -> String {

query/src/sql/plan_parser.rs

@@ -65,6 +65,7 @@ use sqlparser::ast::Statement;
 use sqlparser::ast::TableFactor;
 use sqlparser::ast::UnaryOperator;
 
+use crate::catalogs::Catalog;
 use crate::catalogs::ToReadDataSourcePlan;
 use crate::functions::ContextFunction;
 use crate::sessions::DatabendQueryContextRef;
@@ -788,7 +789,10 @@ impl PlanParser {
                         },
                     )?;
 
-                    let table_func = self.ctx.get_table_function(&table_name, Some(table_args))?;
+                    let table_func = self
+                        .ctx
+                        .get_catalog()
+                        .get_table_function(&table_name, Some(table_args))?;
                     meta_id = table_func.get_id();
                     meta_version = table_func.get_table_info().ident.version;
                     table = table_func.as_table();