From 4d45b6bad0d26ba5fa833a6c4b9b859df00e840f Mon Sep 17 00:00:00 2001 From: jonahgao Date: Mon, 20 Nov 2023 14:58:22 +0800 Subject: [PATCH 1/9] parse column default values --- datafusion/core/src/execution/context/mod.rs | 1 + datafusion/expr/src/logical_plan/ddl.rs | 2 ++ datafusion/expr/src/logical_plan/plan.rs | 2 ++ datafusion/expr/src/table_source.rs | 5 ++++ datafusion/sql/src/planner.rs | 27 ++++++++++++++++++++ datafusion/sql/src/query.rs | 1 + datafusion/sql/src/statement.rs | 15 +++++++++-- 7 files changed, 51 insertions(+), 2 deletions(-) diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index b8e111d361b15..dd3342a861ddb 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -529,6 +529,7 @@ impl SessionContext { if_not_exists, or_replace, constraints, + column_defaults, } = cmd; let input = Arc::try_unwrap(input).unwrap_or_else(|e| e.as_ref().clone()); diff --git a/datafusion/expr/src/logical_plan/ddl.rs b/datafusion/expr/src/logical_plan/ddl.rs index 2c90a3aca7543..97551a941abf7 100644 --- a/datafusion/expr/src/logical_plan/ddl.rs +++ b/datafusion/expr/src/logical_plan/ddl.rs @@ -228,6 +228,8 @@ pub struct CreateMemoryTable { pub if_not_exists: bool, /// Option to replace table content if table already exists pub or_replace: bool, + /// Default values for columns + pub column_defaults: Vec<(String, Expr)>, } /// Creates a view. diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index a024824c7a5a8..69ba42d34a705 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -811,6 +811,7 @@ impl LogicalPlan { name, if_not_exists, or_replace, + column_defaults, .. })) => Ok(LogicalPlan::Ddl(DdlStatement::CreateMemoryTable( CreateMemoryTable { @@ -819,6 +820,7 @@ impl LogicalPlan { name: name.clone(), if_not_exists: *if_not_exists, or_replace: *or_replace, + column_defaults: column_defaults.clone(), }, ))), LogicalPlan::Ddl(DdlStatement::CreateView(CreateView { diff --git a/datafusion/expr/src/table_source.rs b/datafusion/expr/src/table_source.rs index 94f26d9158cd1..565f48c1c5a9e 100644 --- a/datafusion/expr/src/table_source.rs +++ b/datafusion/expr/src/table_source.rs @@ -103,4 +103,9 @@ pub trait TableSource: Sync + Send { fn get_logical_plan(&self) -> Option<&LogicalPlan> { None } + + /// Get the default value for a column, if available. + fn get_column_default(&self, _column: &str) -> Option<&Expr> { + None + } } diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index ca5e260aee050..38d54d2aacbb4 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -230,6 +230,33 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { Ok(Schema::new(fields)) } + pub(super) fn build_column_defaults( + &self, + columns: &Vec, + planner_context: &mut PlannerContext, + ) -> Result> { + let mut column_defaults = vec![]; + // Default expressions are restricted, they must be variable-free + let empty_schema = DFSchema::empty(); + for column in columns { + if let Some(default_sql_expr) = + column.options.iter().find_map(|o| match &o.option { + ColumnOption::Default(expr) => Some(expr), + _ => None, + }) + { + let default_expr = self.sql_to_expr( + default_sql_expr.clone(), + &empty_schema, + planner_context, + )?; + column_defaults + .push((self.normalizer.normalize(column.name.clone()), default_expr)); + } + } + Ok(column_defaults) + } + /// Apply the given TableAlias to the input plan pub(crate) fn apply_table_alias( &self, diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index 832e2da9c6ec9..643f41d844858 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -90,6 +90,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { input: Arc::new(plan), if_not_exists: false, or_replace: false, + column_defaults: vec![], })) } _ => plan, diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 49755729d2d50..a4eebd74804b9 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -204,6 +204,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let mut all_constraints = constraints; let inline_constraints = calc_inline_constraints_from_columns(&columns); all_constraints.extend(inline_constraints); + // Build default column values + let column_defaults = + self.build_column_defaults(&columns, planner_context)?; match query { Some(query) => { let plan = self.query_to_plan(*query, planner_context)?; @@ -250,6 +253,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { input: Arc::new(plan), if_not_exists, or_replace, + column_defaults, }, ))) } @@ -272,6 +276,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { input: Arc::new(plan), if_not_exists, or_replace, + column_defaults, }, ))) } @@ -1170,8 +1175,14 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { datafusion_expr::Expr::Column(source_field.qualified_column()) .cast_to(target_field.data_type(), source.schema())? } - // Fill the default value for the column, currently only supports NULL. - None => datafusion_expr::Expr::Literal(ScalarValue::Null) + // The value is not specified. Fill in the default value for the column. + None => table_source + .get_column_default(&target_field.name()) + .cloned() + .unwrap_or_else(|| { + // If there is no default for the column, then the default is NULL + datafusion_expr::Expr::Literal(ScalarValue::Null) + }) .cast_to(target_field.data_type(), &DFSchema::empty())?, }; Ok(expr.alias(target_field.name())) From e37a2935449d4d2dc5a91a74aa4f4a34082bf776 Mon Sep 17 00:00:00 2001 From: jonahgao Date: Mon, 20 Nov 2023 16:15:37 +0800 Subject: [PATCH 2/9] fix clippy --- datafusion/core/src/execution/context/mod.rs | 2 +- datafusion/sql/src/statement.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index dd3342a861ddb..fe485dce46415 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -529,7 +529,7 @@ impl SessionContext { if_not_exists, or_replace, constraints, - column_defaults, + column_defaults: _, } = cmd; let input = Arc::try_unwrap(input).unwrap_or_else(|e| e.as_ref().clone()); diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index a4eebd74804b9..04f8ce518edfb 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -1177,7 +1177,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } // The value is not specified. Fill in the default value for the column. None => table_source - .get_column_default(&target_field.name()) + .get_column_default(target_field.name()) .cloned() .unwrap_or_else(|| { // If there is no default for the column, then the default is NULL From 1007b4f70c7c68f0734a09be37b947ca2444c068 Mon Sep 17 00:00:00 2001 From: jonahgao Date: Mon, 20 Nov 2023 17:57:49 +0800 Subject: [PATCH 3/9] Impl for memroy table --- .../core/src/datasource/default_table_source.rs | 4 ++++ datafusion/core/src/datasource/memory.rs | 16 ++++++++++++++++ datafusion/core/src/datasource/provider.rs | 5 +++++ datafusion/core/src/execution/context/mod.rs | 15 +++++++++++---- 4 files changed, 36 insertions(+), 4 deletions(-) diff --git a/datafusion/core/src/datasource/default_table_source.rs b/datafusion/core/src/datasource/default_table_source.rs index 00a9c123ceeec..fadf01c74c5d4 100644 --- a/datafusion/core/src/datasource/default_table_source.rs +++ b/datafusion/core/src/datasource/default_table_source.rs @@ -73,6 +73,10 @@ impl TableSource for DefaultTableSource { fn get_logical_plan(&self) -> Option<&datafusion_expr::LogicalPlan> { self.table_provider.get_logical_plan() } + + fn get_column_default(&self, column: &str) -> Option<&Expr> { + self.table_provider.get_column_default(column) + } } /// Wrap TableProvider in TableSource diff --git a/datafusion/core/src/datasource/memory.rs b/datafusion/core/src/datasource/memory.rs index 6bcaa97a408fc..a841518d9c8fa 100644 --- a/datafusion/core/src/datasource/memory.rs +++ b/datafusion/core/src/datasource/memory.rs @@ -19,6 +19,7 @@ use datafusion_physical_plan::metrics::MetricsSet; use futures::StreamExt; +use hashbrown::HashMap; use log::debug; use std::any::Any; use std::fmt::{self, Debug}; @@ -56,6 +57,7 @@ pub struct MemTable { schema: SchemaRef, pub(crate) batches: Vec, constraints: Constraints, + column_defaults: HashMap, } impl MemTable { @@ -79,6 +81,7 @@ impl MemTable { .map(|e| Arc::new(RwLock::new(e))) .collect::>(), constraints: Constraints::empty(), + column_defaults: HashMap::new(), }) } @@ -88,6 +91,15 @@ impl MemTable { self } + /// Assign column defaults + pub fn with_column_defaults( + mut self, + column_defaults: HashMap, + ) -> Self { + self.column_defaults = column_defaults; + self + } + /// Create a mem table by reading from another data source pub async fn load( t: Arc, @@ -228,6 +240,10 @@ impl TableProvider for MemTable { None, ))) } + + fn get_column_default(&self, column: &str) -> Option<&Expr> { + self.column_defaults.get(column) + } } /// Implements for writing to a [`MemTable`] diff --git a/datafusion/core/src/datasource/provider.rs b/datafusion/core/src/datasource/provider.rs index 4fe433044e6ca..275523405a094 100644 --- a/datafusion/core/src/datasource/provider.rs +++ b/datafusion/core/src/datasource/provider.rs @@ -66,6 +66,11 @@ pub trait TableProvider: Sync + Send { None } + /// Get the default value for a column, if available. + fn get_column_default(&self, _column: &str) -> Option<&Expr> { + None + } + /// Create an [`ExecutionPlan`] for scanning the table with optionally /// specified `projection`, `filter` and `limit`, described below. /// diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index fe485dce46415..b25d120dbdc88 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -529,7 +529,7 @@ impl SessionContext { if_not_exists, or_replace, constraints, - column_defaults: _, + column_defaults, } = cmd; let input = Arc::try_unwrap(input).unwrap_or_else(|e| e.as_ref().clone()); @@ -543,7 +543,12 @@ impl SessionContext { let physical = DataFrame::new(self.state(), input); let batches: Vec<_> = physical.collect_partitioned().await?; - let table = Arc::new(MemTable::try_new(schema, batches)?); + let table = Arc::new( + // pass constraints and column defaults to the mem table. + MemTable::try_new(schema, batches)? + .with_constraints(constraints) + .with_column_defaults(column_defaults.into_iter().collect()), + ); self.register_table(&name, table)?; self.return_empty_dataframe() @@ -558,8 +563,10 @@ impl SessionContext { let batches: Vec<_> = physical.collect_partitioned().await?; let table = Arc::new( - // pass constraints to the mem table. - MemTable::try_new(schema, batches)?.with_constraints(constraints), + // pass constraints and column defaults to the mem table. + MemTable::try_new(schema, batches)? + .with_constraints(constraints) + .with_column_defaults(column_defaults.into_iter().collect()), ); self.register_table(&name, table)?; From b837d468cd3e09318525761ffea1dc2295feaa7d Mon Sep 17 00:00:00 2001 From: jonahgao Date: Mon, 20 Nov 2023 23:50:36 +0800 Subject: [PATCH 4/9] Add tests --- datafusion/sql/src/planner.rs | 23 ++++++++----- datafusion/sql/src/statement.rs | 2 +- datafusion/sqllogictest/test_files/insert.slt | 34 +++++++++++++++++++ 3 files changed, 50 insertions(+), 9 deletions(-) diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 38d54d2aacbb4..d86f24853e63a 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -21,8 +21,7 @@ use std::sync::Arc; use std::vec; use arrow_schema::*; -use datafusion_common::field_not_found; -use datafusion_common::internal_err; +use datafusion_common::{field_not_found, internal_err, SchemaError}; use datafusion_expr::WindowUDF; use sqlparser::ast::TimezoneInfo; use sqlparser::ast::{ArrayElemTypeDef, ExactNumberInfo}; @@ -236,8 +235,18 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { planner_context: &mut PlannerContext, ) -> Result> { let mut column_defaults = vec![]; - // Default expressions are restricted, they must be variable-free + // Default expressions are restricted, column references are not allowed let empty_schema = DFSchema::empty(); + let human_readable_error = |e: DataFusionError| match e { + DataFusionError::SchemaError(SchemaError::FieldNotFound { .. }) => { + DataFusionError::Plan(format!( + "Column reference is not allowed in the DEFAULT expression : {}", + e + )) + } + _ => e, + }; + for column in columns { if let Some(default_sql_expr) = column.options.iter().find_map(|o| match &o.option { @@ -245,11 +254,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { _ => None, }) { - let default_expr = self.sql_to_expr( - default_sql_expr.clone(), - &empty_schema, - planner_context, - )?; + let default_expr = self + .sql_to_expr(default_sql_expr.clone(), &empty_schema, planner_context) + .map_err(human_readable_error)?; column_defaults .push((self.normalizer.normalize(column.name.clone()), default_expr)); } diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 04f8ce518edfb..aa2f0583cb99f 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -204,7 +204,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let mut all_constraints = constraints; let inline_constraints = calc_inline_constraints_from_columns(&columns); all_constraints.extend(inline_constraints); - // Build default column values + // Build column default values let column_defaults = self.build_column_defaults(&columns, planner_context)?; match query { diff --git a/datafusion/sqllogictest/test_files/insert.slt b/datafusion/sqllogictest/test_files/insert.slt index a100b5ac6b854..47e29272ef493 100644 --- a/datafusion/sqllogictest/test_files/insert.slt +++ b/datafusion/sqllogictest/test_files/insert.slt @@ -314,3 +314,37 @@ select * from table_without_values; statement ok drop table table_without_values; + +statement ok +create table test_column_defaults( + a int, + b int not null default null, + c int default 100*2+300, + d text default lower('DEFAULT_TEXT'), + e timestamp default now() +) + +query IIITP +insert into test_column_defaults values(1, 10, 100, 'ABC', now()) +---- +1 + +statement error DataFusion error: Execution error: Invalid batch column at '1' has null but schema specifies non-nullable +insert into test_column_defaults(a) values(2) + +query IIITP +insert into test_column_defaults(b) values(20) +---- +1 + +query IIIT rowsort +select a,b,c,d from test_column_defaults +---- +1 10 100 ABC +NULL 20 500 default_text + +statement ok +drop table test_column_defaults + +statement error DataFusion error: Error during planning: Column reference is not allowed in the DEFAULT expression : Schema error: No field named a. +create table test_column_defaults(a int, b int default a+1) \ No newline at end of file From 63a8df6234e4ce399e8057164296a66ec9e6947c Mon Sep 17 00:00:00 2001 From: jonahgao Date: Mon, 20 Nov 2023 23:53:16 +0800 Subject: [PATCH 5/9] Add test --- datafusion/sqllogictest/test_files/insert.slt | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/datafusion/sqllogictest/test_files/insert.slt b/datafusion/sqllogictest/test_files/insert.slt index 47e29272ef493..e0195165753be 100644 --- a/datafusion/sqllogictest/test_files/insert.slt +++ b/datafusion/sqllogictest/test_files/insert.slt @@ -346,5 +346,30 @@ NULL 20 500 default_text statement ok drop table test_column_defaults + +# test create table as +statement ok +create table test_column_defaults( + a int, + b int not null default null, + c int default 100*2+300, + d text default lower('DEFAULT_TEXT'), + e timestamp default now() +) as values(1, 10, 100, 'ABC', now()) + +query IIITP +insert into test_column_defaults(b) values(20) +---- +1 + +query IIIT rowsort +select a,b,c,d from test_column_defaults +---- +1 10 100 ABC +NULL 20 500 default_text + +statement ok +drop table test_column_defaults + statement error DataFusion error: Error during planning: Column reference is not allowed in the DEFAULT expression : Schema error: No field named a. create table test_column_defaults(a int, b int default a+1) \ No newline at end of file From 670d8ccbb89621c8c631685365435fc01882e3b3 Mon Sep 17 00:00:00 2001 From: jonahgao Date: Tue, 21 Nov 2023 10:31:12 +0800 Subject: [PATCH 6/9] Use plan_datafusion_err --- datafusion/sql/src/planner.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index d86f24853e63a..70b5342b7f9e6 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -21,7 +21,9 @@ use std::sync::Arc; use std::vec; use arrow_schema::*; -use datafusion_common::{field_not_found, internal_err, SchemaError}; +use datafusion_common::{ + field_not_found, internal_err, plan_datafusion_err, SchemaError, +}; use datafusion_expr::WindowUDF; use sqlparser::ast::TimezoneInfo; use sqlparser::ast::{ArrayElemTypeDef, ExactNumberInfo}; @@ -239,10 +241,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let empty_schema = DFSchema::empty(); let human_readable_error = |e: DataFusionError| match e { DataFusionError::SchemaError(SchemaError::FieldNotFound { .. }) => { - DataFusionError::Plan(format!( + plan_datafusion_err!( "Column reference is not allowed in the DEFAULT expression : {}", e - )) + ) } _ => e, }; From f8617327aba501570e443eebc0df19a2192b984c Mon Sep 17 00:00:00 2001 From: jonahgao Date: Tue, 21 Nov 2023 14:43:18 +0800 Subject: [PATCH 7/9] Add comment --- datafusion/sql/src/planner.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 70b5342b7f9e6..c01d703c65917 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -231,6 +231,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { Ok(Schema::new(fields)) } + /// Returns pairs of (column_name, default_expr) pub(super) fn build_column_defaults( &self, columns: &Vec, From 81d15624bb1cfde6b14f1c868a4c5f4196c567d7 Mon Sep 17 00:00:00 2001 From: Jonah Gao Date: Tue, 21 Nov 2023 14:43:43 +0800 Subject: [PATCH 8/9] Update datafusion/sql/src/planner.rs Co-authored-by: comphead --- datafusion/sql/src/planner.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index c01d703c65917..63af54f8d2f0c 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -240,7 +240,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let mut column_defaults = vec![]; // Default expressions are restricted, column references are not allowed let empty_schema = DFSchema::empty(); - let human_readable_error = |e: DataFusionError| match e { + let error_desc = |e: DataFusionError| match e { DataFusionError::SchemaError(SchemaError::FieldNotFound { .. }) => { plan_datafusion_err!( "Column reference is not allowed in the DEFAULT expression : {}", From 0b30bcb235bdb1b1dcd933b12cb9163f8f546c6d Mon Sep 17 00:00:00 2001 From: jonahgao Date: Tue, 21 Nov 2023 14:49:31 +0800 Subject: [PATCH 9/9] Fix ci --- datafusion/sql/src/planner.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 63af54f8d2f0c..622e5aca799ab 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -231,7 +231,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { Ok(Schema::new(fields)) } - /// Returns pairs of (column_name, default_expr) + /// Returns a vector of (column_name, default_expr) pairs pub(super) fn build_column_defaults( &self, columns: &Vec, @@ -259,7 +259,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { { let default_expr = self .sql_to_expr(default_sql_expr.clone(), &empty_schema, planner_context) - .map_err(human_readable_error)?; + .map_err(error_desc)?; column_defaults .push((self.normalizer.normalize(column.name.clone()), default_expr)); }