Add docs and fixes

Author: trueleo

datafusion-federation/examples/df-csv-advanced.rs

@@ -78,7 +78,7 @@ async fn main() {
     let sqlite_federation_provider = Arc::new(SQLFederationProvider::new(sqlite_executor));
     // Create the schema provider
     let sqlite_schema_provider = Arc::new(
-        SQLSchemaProvider::new_with_table_names(sqlite_federation_provider, sqlite_known_tables)
+        SQLSchemaProvider::new_with_tables(sqlite_federation_provider, sqlite_known_tables)
             .await
             .expect("Create new schema provider with tables"),
     );
@@ -107,12 +107,9 @@ async fn main() {
     let postgres_federation_provider = Arc::new(SQLFederationProvider::new(postgres_executor));
     // Create the schema provider
     let postgres_schema_provider = Arc::new(
-        SQLSchemaProvider::new_with_table_names(
-            postgres_federation_provider,
-            postgres_known_tables,
-        )
-        .await
-        .expect("Create new schema provider with tables"),
+        SQLSchemaProvider::new_with_tables(postgres_federation_provider, postgres_known_tables)
+            .await
+            .expect("Create new schema provider with tables"),
     );
 
     /////////////////////

datafusion-federation/examples/df-csv.rs

@@ -30,7 +30,7 @@ async fn main() -> Result<()> {
 
     // Get the schema
     let schema_provider =
-        Arc::new(SQLSchemaProvider::new_with_table_names(provider, known_tables).await?);
+        Arc::new(SQLSchemaProvider::new_with_tables(provider, known_tables).await?);
 
     // Main context
     let state = datafusion_federation::default_session_state();

datafusion-federation/src/sql/analyzer.rs

@@ -20,6 +20,7 @@ use super::SQLTableSource;
 
 type Result<T> = std::result::Result<T, datafusion::error::DataFusionError>;
 
+/// Rewrite LogicalPlan's table scans and expressions to use the federated table name.
 #[derive(Debug)]
 pub struct RewriteTableScanAnalyzer;
 
@@ -31,7 +32,7 @@ impl RewriteTableScanAnalyzer {
 }
 
 /// Rewrite table scans to use the original federated table name.
-pub fn rewrite_table_scans(
+fn rewrite_table_scans(
     plan: &LogicalPlan,
     known_rewrites: &mut HashMap<TableReference, TableReference>,
 ) -> Result<LogicalPlan> {
@@ -176,7 +177,7 @@ pub fn rewrite_column_name_in_expr(
     }
 }
 
-pub fn rewrite_table_scans_in_expr(
+fn rewrite_table_scans_in_expr(
     expr: Expr,
     known_rewrites: &mut HashMap<TableReference, TableReference>,
 ) -> Result<Expr> {

datafusion-federation/src/sql/ast_analyzer.rs

@@ -1,7 +1,7 @@
 use std::ops::ControlFlow;
 
 use datafusion::sql::{
-    sqlparser::ast::{TableFactor, TableFunctionArgs, VisitorMut},
+    sqlparser::ast::{FunctionArg, TableFactor, TableFunctionArgs, VisitorMut},
     TableReference,
 };
 
@@ -15,21 +15,60 @@ pub fn replace_table_args_analyzer(mut visitor: TableArgReplace) -> AstAnalyzer
     Box::new(x)
 }
 
+/// Used to construct a AstAnalyzer that can replace table arguments.
+///
+/// ```rust
+/// use datafusion::sql::sqlparser::ast::{FunctionArg, Expr, Value};
+/// use datafusion_federation::sql::ast_analyzer::TableArgReplace;
+///
+/// let mut analyzer = TableArgReplace::default().with(
+///     TableReference::parse_str("table1"),
+///     vec![FunctionArg::Unnamed(
+///         Expr::Value(
+///             Value::Number("1".to_string(), false),
+///         )
+///         .into(),
+///     )],
+/// );
+/// let analyzer = analyzer.into_analyzer();
+/// ```
 #[derive(Debug, Clone, PartialEq, Eq, Default)]
 pub struct TableArgReplace {
     pub tables: Vec<(TableReference, TableFunctionArgs)>,
 }
 
 impl TableArgReplace {
-    pub fn new(tables: Vec<(TableReference, TableFunctionArgs)>) -> Self {
-        Self { tables }
+    /// Constructs a new `TableArgReplace` instance.
+    pub fn new(tables: Vec<(TableReference, Vec<FunctionArg>)>) -> Self {
+        Self {
+            tables: tables
+                .into_iter()
+                .map(|(table, args)| {
+                    (
+                        table,
+                        TableFunctionArgs {
+                            args,
+                            settings: None,
+                        },
+                    )
+                })
+                .collect(),
+        }
     }
 
-    pub fn with(mut self, table: TableReference, args: TableFunctionArgs) -> Self {
-        self.tables.push((table, args));
+    /// Adds a new table argument replacement.
+    pub fn with(mut self, table: TableReference, args: Vec<FunctionArg>) -> Self {
+        self.tables.push((
+            table,
+            TableFunctionArgs {
+                args,
+                settings: None,
+            },
+        ));
         self
     }
 
+    /// Converts the `TableArgReplace` instance into an `AstAnalyzer`.
     pub fn into_analyzer(self) -> AstAnalyzer {
         replace_table_args_analyzer(self)
     }

datafusion-federation/src/sql/mod.rs

@@ -1,4 +1,4 @@
-pub mod analyzer;
+mod analyzer;
 pub mod ast_analyzer;
 mod executor;
 mod schema;
@@ -368,7 +368,6 @@ mod tests {
     #[derive(Debug, Clone)]
     struct TestExecutor {
         compute_context: String,
-        rewrite_table: bool,
     }
 
     #[async_trait]
@@ -389,11 +388,6 @@ mod tests {
             unimplemented!()
         }
 
-        fn logical_optimizer(&self) -> Option<LogicalOptimizer> {
-            self.rewrite_table
-                .then_some(Box::new(analyzer::RewriteTableScanAnalyzer::rewrite))
-        }
-
         async fn table_names(&self) -> Result<Vec<String>> {
             unimplemented!()
         }
@@ -448,12 +442,10 @@ mod tests {
     async fn basic_sql_federation_test() -> Result<(), DataFusionError> {
         let test_executor_a = TestExecutor {
             compute_context: "a".into(),
-            rewrite_table: false,
         };
 
         let test_executor_b = TestExecutor {
             compute_context: "b".into(),
-            rewrite_table: true,
         };
 
         let table_a1_ref = "table_a1".to_string();

datafusion-federation/src/sql/schema.rs

@@ -2,19 +2,25 @@ use std::{any::Any, str::FromStr, sync::Arc};
 
 use async_trait::async_trait;
 use datafusion::{
-    catalog::SchemaProvider, datasource::TableProvider, error::Result, sql::TableReference,
+    catalog::SchemaProvider,
+    datasource::TableProvider,
+    error::{DataFusionError, Result},
+    sql::TableReference,
 };
 use futures::future::join_all;
 
 use super::{table::SQLTable, RemoteTableRef, SQLTableSource};
 use crate::{sql::SQLFederationProvider, FederatedTableProviderAdaptor};
 
+/// An in-memory schema provider for SQL tables.
 #[derive(Debug)]
 pub struct SQLSchemaProvider {
     tables: Vec<Arc<SQLTableSource>>,
 }
 
 impl SQLSchemaProvider {
+    /// Creates a new SQLSchemaProvider from a [`SQLFederationProvider`].
+    /// Initializes the schema provider by fetching table names and schema from the federation provider's executor,
     pub async fn new(provider: Arc<SQLFederationProvider>) -> Result<Self> {
         let executor = Arc::clone(&provider.executor);
         let tables = executor
@@ -41,14 +47,19 @@ impl SQLSchemaProvider {
         Ok(Self { tables })
     }
 
-    pub async fn new_with_table_names(
+    /// Creates a new SQLSchemaProvider from a SQLFederationProvider and a list of table references.
+    /// Fetches the schema for each table using the executor's implementation.
+    pub async fn new_with_tables<T>(
         provider: Arc<SQLFederationProvider>,
-        tables: Vec<String>,
-    ) -> Result<Self> {
+        tables: Vec<T>,
+    ) -> Result<Self>
+    where
+        T: TryInto<RemoteTableRef, Error = DataFusionError>,
+    {
         let tables = tables
             .into_iter()
-            .map(|t| RemoteTableRef::from_str(&t))
-            .collect::<Result<Vec<_>>>()?;
+            .map(|t| t.try_into())
+            .collect::<Result<Vec<RemoteTableRef>>>()?;
         let futures: Vec<_> = tables
             .into_iter()
             .map(|t| SQLTableSource::new(Arc::clone(&provider), t))
@@ -58,7 +69,8 @@ impl SQLSchemaProvider {
         Ok(Self { tables })
     }
 
-    pub fn new_with_tables(
+    /// Creates a new SQLSchemaProvider from a SQLFederationProvider and a list of custom table instances.
+    pub fn new_with_custom_tables(
         provider: Arc<SQLFederationProvider>,
         tables: Vec<Arc<dyn SQLTable>>,
     ) -> Self {

datafusion-federation/src/sql/table.rs

@@ -5,7 +5,6 @@ use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::error::Result;
 use datafusion::logical_expr::TableSource;
 use datafusion::logical_expr::TableType;
-use datafusion::sql::sqlparser::ast::TableFunctionArgs;
 use datafusion::sql::TableReference;
 use std::any::Any;
 use std::sync::Arc;
@@ -15,16 +14,22 @@ use super::executor::LogicalOptimizer;
 use super::AstAnalyzer;
 use super::RemoteTableRef;
 
+/// Trait to represent a SQL remote table inside [`SQLTableSource`].
+/// A remote table provides information such as schema, table reference, and
+/// provides hooks for rewriting the logical plan and AST before execution.
+/// This crate provides [`RemoteTable`] as a default ready-to-use type.
 pub trait SQLTable: std::fmt::Debug + Send + Sync {
     /// Returns a reference as a trait object.
     fn as_any(&self) -> &dyn Any;
-    /// Table reference for the table, which can be used to identify the table in the SQL query.
-    /// This method should return a unique identifier for the table.
-    /// Used for registering the table to sql schema provider.
+    /// Provides the [`TableReference`](`datafusion::sql::TableReference`) used to identify the table in SQL queries.
+    /// This TableReference is used for registering the table with the [`SQLSchemaProvider`](`super::SQLSchemaProvider`).
+    /// If the table provider is registered in the Datafusion context under a different name,
+    /// the logical plan will be rewritten to use this table reference during execution.
+    /// Therefore, any AST analyzer should match against this table reference.
     fn table_reference(&self) -> TableReference;
     /// Schema of the remote table
     fn schema(&self) -> SchemaRef;
-    /// Returns an logical optimizer specific to this table, will be used to modify the logical plan before execution
+    /// Returns a logical optimizer specific to this table, will be used to modify the logical plan before execution
     fn logical_optimizer(&self) -> Option<LogicalOptimizer> {
         None
     }
@@ -34,20 +39,33 @@ pub trait SQLTable: std::fmt::Debug + Send + Sync {
     }
 }
 
-#[derive(Debug)]
+/// Represents a remote table with a reference and schema.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct RemoteTable {
     table_ref: RemoteTableRef,
     schema: SchemaRef,
 }
 
 impl RemoteTable {
+    /// Creates a new `RemoteTable` instance.
+    ///
+    /// Examples:
+    /// ```rust
+    /// RemoteTable::new("myschema.table", schema);
+    /// RemoteTable::new(r#"myschema."Table""#, schema);
+    /// RemoteTable::new("myschema.view('obj')", schema);
+    /// RemoteTable::new("myschema.view(name => 'obj')", schema);
+    /// RemoteTable::new("myschema.view(name = 'obj')", schema);
+    /// ```
     pub fn new(table_ref: impl Into<RemoteTableRef>, schema: SchemaRef) -> Self {
         Self {
             table_ref: table_ref.into(),
             schema,
         }
     }
 
+    /// Return table reference of this remote table.
+    /// Only returns the object name, ignoring functional params if any
     pub fn table_reference(&self) -> TableReference {
         TableReference::from(&self.table_ref)
     }
@@ -70,17 +88,12 @@ impl SQLTable for RemoteTable {
         None
     }
 
+    /// Returns ast analyzer that modifies table that contains functional args after table ident
     fn ast_analyzer(&self) -> Option<AstAnalyzer> {
         if let RemoteTableRef::TableReferenceWithArgs(name, args) = &self.table_ref {
             Some(
                 ast_analyzer::TableArgReplace::default()
-                    .with(
-                        name.clone(),
-                        TableFunctionArgs {
-                            args: args.iter().cloned().collect(),
-                            settings: None,
-                        },
-                    )
+                    .with(name.clone(), args.iter().cloned().collect())
                     .into_analyzer(),
             )
         } else {
@@ -107,6 +120,7 @@ impl SQLTableSource {
         Ok(Self::new_with_schema(provider, remote_table_ref, schema))
     }
 
+    /// Create a SQLTableSource with a table reference and schema
     pub fn new_with_schema(
         provider: Arc<SQLFederationProvider>,
         table_ref: impl Into<RemoteTableRef>,
@@ -118,10 +132,12 @@ impl SQLTableSource {
         }
     }
 
+    /// Create new with a custom SQLtable instance.
     pub fn new_with_table(provider: Arc<SQLFederationProvider>, table: Arc<dyn SQLTable>) -> Self {
         Self { provider, table }
     }
 
+    /// Return associated table reference of stored remote table
     pub fn table_reference(&self) -> TableReference {
         self.table.table_reference()
     }

datafusion-federation/src/sql/table_reference.rs

@@ -8,13 +8,24 @@ use datafusion::{
     },
 };
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+/// A multipart identifier to a remote table, view or parameterized view.
+///
+/// RemoteTableRef can be created by parsing from a string represeting a table obbject with optional
+/// ```rust
+/// RemoteTableRef::from_str("myschema.table")
+/// RemoteTableRef::from_str(r#"myschema."Table""#)
+/// RemoteTableRef::from_str("myschema.view('obj')")
+/// RemoteTableRef::from_str("myschema.view(name => 'obj')")
+/// RemoteTableRef::from_str("myschema.view(name = 'obj')")
+/// ```
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum RemoteTableRef {
     TableReference(TableReference),
     TableReferenceWithArgs(TableReference, Arc<[FunctionArg]>),
 }
 
 impl RemoteTableRef {
+    /// Get quoted_string representation for the table it is referencing, this is same as calling to_quoted_string on the inner table reference.
     pub fn to_quoted_string(&self) -> String {
         match self {
             RemoteTableRef::TableReference(table_ref) => table_ref.to_quoted_string(),
@@ -60,6 +71,20 @@ impl TryFrom<&str> for RemoteTableRef {
     }
 }
 
+impl TryFrom<String> for RemoteTableRef {
+    type Error = datafusion::error::DataFusionError;
+    fn try_from(s: String) -> Result<Self, Self::Error> {
+        s.parse()
+    }
+}
+
+impl TryFrom<&String> for RemoteTableRef {
+    type Error = datafusion::error::DataFusionError;
+    fn try_from(s: &String) -> Result<Self, Self::Error> {
+        s.parse()
+    }
+}
+
 impl FromStr for RemoteTableRef {
     type Err = datafusion::error::DataFusionError;