diff --git a/datafusion/common/src/functional_dependencies.rs b/datafusion/common/src/functional_dependencies.rs index 6a08c4fd35890..324374f557b48 100644 --- a/datafusion/common/src/functional_dependencies.rs +++ b/datafusion/common/src/functional_dependencies.rs @@ -18,11 +18,14 @@ //! FunctionalDependencies keeps track of functional dependencies //! inside DFSchema. -use crate::{DFSchema, DFSchemaRef, DataFusionError, JoinType, Result}; -use sqlparser::ast::TableConstraint; use std::collections::HashSet; use std::fmt::{Display, Formatter}; use std::ops::Deref; +use std::vec::IntoIter; + +use crate::{DFSchema, DFSchemaRef, DataFusionError, JoinType, Result}; + +use sqlparser::ast::TableConstraint; /// This object defines a constraint on a table. #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -43,13 +46,15 @@ pub struct Constraints { impl Constraints { /// Create empty constraints pub fn empty() -> Self { - Constraints::new(vec![]) + Constraints::new_unverified(vec![]) } - // This method is private. - // Outside callers can either create empty constraint using `Constraints::empty` API. - // or create constraint from table constraints using `Constraints::new_from_table_constraints` API. - fn new(constraints: Vec) -> Self { + /// Create a new `Constraints` object from the given `constraints`. + /// Users should use the `empty` or `new_from_table_constraints` functions + /// for constructing `Constraints`. This constructor is for internal + /// purposes only and does not check whether the argument is valid. The user + /// is responsible for supplying a valid vector of `Constraint` objects. + pub fn new_unverified(constraints: Vec) -> Self { Self { inner: constraints } } @@ -104,7 +109,7 @@ impl Constraints { )), }) .collect::>>()?; - Ok(Constraints::new(constraints)) + Ok(Constraints::new_unverified(constraints)) } /// Check whether constraints is empty @@ -113,6 +118,15 @@ impl Constraints { } } +impl IntoIterator for Constraints { + type Item = Constraint; + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.inner.into_iter() + } +} + impl Display for Constraints { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { let pk: Vec = self.inner.iter().map(|c| format!("{:?}", c)).collect(); @@ -534,7 +548,7 @@ mod tests { #[test] fn constraints_iter() { - let constraints = Constraints::new(vec![ + let constraints = Constraints::new_unverified(vec![ Constraint::PrimaryKey(vec![10]), Constraint::Unique(vec![20]), ]); diff --git a/datafusion/common/src/lib.rs b/datafusion/common/src/lib.rs index 71782f67046d5..a939cf73dc9c0 100644 --- a/datafusion/common/src/lib.rs +++ b/datafusion/common/src/lib.rs @@ -15,45 +15,47 @@ // specific language governing permissions and limitations // under the License. +mod column; +mod dfschema; +mod error; +mod functional_dependencies; +mod join_type; +#[cfg(feature = "pyarrow")] +mod pyarrow; +mod schema_reference; +mod table_reference; +mod unnest; + pub mod alias; pub mod cast; -mod column; pub mod config; -mod dfschema; pub mod display; -mod error; pub mod file_options; pub mod format; -mod functional_dependencies; pub mod hash_utils; -mod join_type; pub mod parsers; -#[cfg(feature = "pyarrow")] -mod pyarrow; pub mod scalar; -mod schema_reference; pub mod stats; -mod table_reference; pub mod test_util; pub mod tree_node; -mod unnest; pub mod utils; +/// Reexport arrow crate +pub use arrow; pub use column::Column; pub use dfschema::{DFField, DFSchema, DFSchemaRef, ExprSchema, SchemaExt, ToDFSchema}; pub use error::{ field_not_found, unqualified_field_not_found, DataFusionError, Result, SchemaError, SharedResult, }; - pub use file_options::file_type::{ FileType, GetExt, DEFAULT_ARROW_EXTENSION, DEFAULT_AVRO_EXTENSION, DEFAULT_CSV_EXTENSION, DEFAULT_JSON_EXTENSION, DEFAULT_PARQUET_EXTENSION, }; pub use file_options::FileTypeWriterOptions; pub use functional_dependencies::{ - aggregate_functional_dependencies, get_target_functional_dependencies, Constraints, - Dependency, FunctionalDependence, FunctionalDependencies, + aggregate_functional_dependencies, get_target_functional_dependencies, Constraint, + Constraints, Dependency, FunctionalDependence, FunctionalDependencies, }; pub use join_type::{JoinConstraint, JoinType}; pub use scalar::{ScalarType, ScalarValue}; @@ -63,9 +65,6 @@ pub use table_reference::{OwnedTableReference, ResolvedTableReference, TableRefe pub use unnest::UnnestOptions; pub use utils::project_schema; -/// Reexport arrow crate -pub use arrow; - /// Downcast an Arrow Array to a concrete type, return an `DataFusionError::Internal` if the cast is /// not possible. In normal usage of DataFusion the downcast should always succeed. /// diff --git a/datafusion/core/src/catalog/listing_schema.rs b/datafusion/core/src/catalog/listing_schema.rs index e7b4d8dec03c4..7e527642be164 100644 --- a/datafusion/core/src/catalog/listing_schema.rs +++ b/datafusion/core/src/catalog/listing_schema.rs @@ -16,21 +16,25 @@ // under the License. //! listing_schema contains a SchemaProvider that scans ObjectStores for tables automatically + +use std::any::Any; +use std::collections::{HashMap, HashSet}; +use std::path::Path; +use std::sync::{Arc, Mutex}; + use crate::catalog::schema::SchemaProvider; use crate::datasource::provider::TableProviderFactory; use crate::datasource::TableProvider; use crate::execution::context::SessionState; -use async_trait::async_trait; + use datafusion_common::parsers::CompressionTypeVariant; -use datafusion_common::{DFSchema, DataFusionError, OwnedTableReference}; +use datafusion_common::{Constraints, DFSchema, DataFusionError, OwnedTableReference}; use datafusion_expr::CreateExternalTable; + +use async_trait::async_trait; use futures::TryStreamExt; use itertools::Itertools; use object_store::ObjectStore; -use std::any::Any; -use std::collections::{HashMap, HashSet}; -use std::path::Path; -use std::sync::{Arc, Mutex}; /// A [`SchemaProvider`] that scans an [`ObjectStore`] to automatically discover tables /// @@ -149,6 +153,7 @@ impl ListingSchemaProvider { order_exprs: vec![], unbounded: false, options: Default::default(), + constraints: Constraints::empty(), }, ) .await?; diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 5b1710d344ee2..228168d38a655 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -23,19 +23,19 @@ use std::{any::Any, sync::Arc}; use super::helpers::{expr_applicable_for_cols, pruned_partition_list, split_files}; use super::PartitionedFile; -use crate::datasource::file_format::file_compression_type::{ - FileCompressionType, FileTypeExt, -}; -use crate::datasource::physical_plan::{ - is_plan_streaming, FileScanConfig, FileSinkConfig, -}; use crate::datasource::{ file_format::{ - arrow::ArrowFormat, avro::AvroFormat, csv::CsvFormat, json::JsonFormat, - parquet::ParquetFormat, FileFormat, + arrow::ArrowFormat, + avro::AvroFormat, + csv::CsvFormat, + file_compression_type::{FileCompressionType, FileTypeExt}, + json::JsonFormat, + parquet::ParquetFormat, + FileFormat, }, get_statistics_with_limit, listing::ListingTableUrl, + physical_plan::{is_plan_streaming, FileScanConfig, FileSinkConfig}, TableProvider, TableType, }; use crate::logical_expr::TableProviderFilterPushDown; @@ -46,12 +46,13 @@ use crate::{ logical_expr::Expr, physical_plan::{empty::EmptyExec, ExecutionPlan, Statistics}, }; + use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, SchemaBuilder, SchemaRef}; use arrow_schema::Schema; use datafusion_common::{ - internal_err, plan_err, project_schema, FileType, FileTypeWriterOptions, SchemaExt, - ToDFSchema, + internal_err, plan_err, project_schema, Constraints, FileType, FileTypeWriterOptions, + SchemaExt, ToDFSchema, }; use datafusion_execution::cache::cache_manager::FileStatisticsCache; use datafusion_execution::cache::cache_unit::DefaultFileStatisticsCache; @@ -590,6 +591,7 @@ pub struct ListingTable { definition: Option, collected_statistics: FileStatisticsCache, infinite_source: bool, + constraints: Constraints, } impl ListingTable { @@ -627,11 +629,18 @@ impl ListingTable { definition: None, collected_statistics: Arc::new(DefaultFileStatisticsCache::default()), infinite_source, + constraints: Constraints::empty(), }; Ok(table) } + /// Assign constraints + pub fn with_constraints(mut self, constraints: Constraints) -> Self { + self.constraints = constraints; + self + } + /// Set the [`FileStatisticsCache`] used to cache parquet file statistics. /// /// Setting a statistics cache on the `SessionContext` can avoid refetching statistics @@ -703,6 +712,10 @@ impl TableProvider for ListingTable { Arc::clone(&self.table_schema) } + fn constraints(&self) -> Option<&Constraints> { + Some(&self.constraints) + } + fn table_type(&self) -> TableType { TableType::Base } diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs index 9c438a47943f4..ebfb589f179e1 100644 --- a/datafusion/core/src/datasource/listing_table_factory.rs +++ b/datafusion/core/src/datasource/listing_table_factory.rs @@ -21,28 +21,26 @@ use std::path::Path; use std::str::FromStr; use std::sync::Arc; -use arrow::datatypes::{DataType, SchemaRef}; -use async_trait::async_trait; -use datafusion_common::file_options::{FileTypeWriterOptions, StatementOptions}; -use datafusion_common::DataFusionError; -use datafusion_expr::CreateExternalTable; +use super::listing::ListingTableInsertMode; -use crate::datasource::file_format::arrow::ArrowFormat; -use crate::datasource::file_format::avro::AvroFormat; -use crate::datasource::file_format::csv::CsvFormat; -use crate::datasource::file_format::file_compression_type::FileCompressionType; -use crate::datasource::file_format::json::JsonFormat; -use crate::datasource::file_format::parquet::ParquetFormat; -use crate::datasource::file_format::FileFormat; +use crate::datasource::file_format::{ + arrow::ArrowFormat, avro::AvroFormat, csv::CsvFormat, + file_compression_type::FileCompressionType, json::JsonFormat, parquet::ParquetFormat, + FileFormat, +}; use crate::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; use crate::datasource::provider::TableProviderFactory; use crate::datasource::TableProvider; use crate::execution::context::SessionState; -use datafusion_common::FileType; -use super::listing::ListingTableInsertMode; +use arrow::datatypes::{DataType, SchemaRef}; +use datafusion_common::file_options::{FileTypeWriterOptions, StatementOptions}; +use datafusion_common::{DataFusionError, FileType}; +use datafusion_expr::CreateExternalTable; + +use async_trait::async_trait; /// A `TableProviderFactory` capable of creating new `ListingTable`s pub struct ListingTableFactory {} @@ -232,7 +230,9 @@ impl TableProviderFactory for ListingTableFactory { .with_schema(resolved_schema); let provider = ListingTable::try_new(config)? .with_cache(state.runtime_env().cache_manager.get_file_statistic_cache()); - let table = provider.with_definition(cmd.definition.clone()); + let table = provider + .with_definition(cmd.definition.clone()) + .with_constraints(cmd.constraints.clone()); Ok(Arc::new(table)) } } @@ -248,13 +248,13 @@ fn get_extension(path: &str) -> String { #[cfg(test)] mod tests { - use super::*; - use std::collections::HashMap; + use super::*; use crate::execution::context::SessionContext; + use datafusion_common::parsers::CompressionTypeVariant; - use datafusion_common::{DFSchema, OwnedTableReference}; + use datafusion_common::{Constraints, DFSchema, OwnedTableReference}; #[tokio::test] async fn test_create_using_non_std_file_ext() { @@ -282,6 +282,7 @@ mod tests { order_exprs: vec![], unbounded: false, options: HashMap::new(), + constraints: Constraints::empty(), }; let table_provider = factory.create(&state, &cmd).await.unwrap(); let listing_table = table_provider diff --git a/datafusion/core/src/datasource/memory.rs b/datafusion/core/src/datasource/memory.rs index 337a8cabc269c..d499ed679f933 100644 --- a/datafusion/core/src/datasource/memory.rs +++ b/datafusion/core/src/datasource/memory.rs @@ -54,7 +54,7 @@ pub type PartitionData = Arc>>; pub struct MemTable { schema: SchemaRef, pub(crate) batches: Vec, - constraints: Option, + constraints: Constraints, } impl MemTable { @@ -77,15 +77,13 @@ impl MemTable { .into_iter() .map(|e| Arc::new(RwLock::new(e))) .collect::>(), - constraints: None, + constraints: Constraints::empty(), }) } /// Assign constraints pub fn with_constraints(mut self, constraints: Constraints) -> Self { - if !constraints.is_empty() { - self.constraints = Some(constraints); - } + self.constraints = constraints; self } @@ -164,7 +162,7 @@ impl TableProvider for MemTable { } fn constraints(&self) -> Option<&Constraints> { - self.constraints.as_ref() + Some(&self.constraints) } fn table_type(&self) -> TableType { diff --git a/datafusion/core/src/datasource/provider.rs b/datafusion/core/src/datasource/provider.rs index 5ebcc45b572b7..111e3a9057f5a 100644 --- a/datafusion/core/src/datasource/provider.rs +++ b/datafusion/core/src/datasource/provider.rs @@ -42,6 +42,11 @@ pub trait TableProvider: Sync + Send { fn schema(&self) -> SchemaRef; /// Get a reference to the constraints of the table. + /// Returns: + /// - `None` for tables that do not support constraints. + /// - `Some(&Constraints)` for tables supporting constraints. + /// Therefore, a `Some(&Constraints::empty())` return value indicates that + /// this table supports constraints, but there are no constraints. fn constraints(&self) -> Option<&Constraints> { None } diff --git a/datafusion/expr/src/logical_plan/ddl.rs b/datafusion/expr/src/logical_plan/ddl.rs index dc247da3642c0..2c90a3aca7543 100644 --- a/datafusion/expr/src/logical_plan/ddl.rs +++ b/datafusion/expr/src/logical_plan/ddl.rs @@ -112,9 +112,10 @@ impl DdlStatement { match self.0 { DdlStatement::CreateExternalTable(CreateExternalTable { ref name, + constraints, .. }) => { - write!(f, "CreateExternalTable: {name:?}") + write!(f, "CreateExternalTable: {name:?}{constraints}") } DdlStatement::CreateMemoryTable(CreateMemoryTable { name, @@ -191,6 +192,8 @@ pub struct CreateExternalTable { pub unbounded: bool, /// Table(provider) specific options pub options: HashMap, + /// The list of constraints in the schema, such as primary key, unique, etc. + pub constraints: Constraints, } // Hashing refers to a subset of fields considered in PartialEq. diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 0ebcf2537ddab..bda0f7828726d 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -180,6 +180,25 @@ message EmptyRelationNode { bool produce_one_row = 1; } +message PrimaryKeyConstraint{ + repeated uint64 indices = 1; +} + +message UniqueConstraint{ + repeated uint64 indices = 1; +} + +message Constraint{ + oneof constraint_mode{ + PrimaryKeyConstraint primary_key = 1; + UniqueConstraint unique = 2; + } +} + +message Constraints{ + repeated Constraint constraints = 1; +} + message CreateExternalTableNode { reserved 1; // was string name OwnedTableReference name = 12; @@ -195,6 +214,7 @@ message CreateExternalTableNode { repeated LogicalExprNodeCollection order_exprs = 13; bool unbounded = 14; map options = 11; + Constraints constraints = 15; } message PrepareNode { diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index d1e9e886e7d58..ced0c8bd7c7a0 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -3425,6 +3425,207 @@ impl<'de> serde::Deserialize<'de> for ColumnStats { deserializer.deserialize_struct("datafusion.ColumnStats", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for Constraint { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.constraint_mode.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.Constraint", len)?; + if let Some(v) = self.constraint_mode.as_ref() { + match v { + constraint::ConstraintMode::PrimaryKey(v) => { + struct_ser.serialize_field("primaryKey", v)?; + } + constraint::ConstraintMode::Unique(v) => { + struct_ser.serialize_field("unique", v)?; + } + } + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for Constraint { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "primary_key", + "primaryKey", + "unique", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + PrimaryKey, + Unique, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "primaryKey" | "primary_key" => Ok(GeneratedField::PrimaryKey), + "unique" => Ok(GeneratedField::Unique), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = Constraint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.Constraint") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut constraint_mode__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::PrimaryKey => { + if constraint_mode__.is_some() { + return Err(serde::de::Error::duplicate_field("primaryKey")); + } + constraint_mode__ = map_.next_value::<::std::option::Option<_>>()?.map(constraint::ConstraintMode::PrimaryKey) +; + } + GeneratedField::Unique => { + if constraint_mode__.is_some() { + return Err(serde::de::Error::duplicate_field("unique")); + } + constraint_mode__ = map_.next_value::<::std::option::Option<_>>()?.map(constraint::ConstraintMode::Unique) +; + } + } + } + Ok(Constraint { + constraint_mode: constraint_mode__, + }) + } + } + deserializer.deserialize_struct("datafusion.Constraint", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for Constraints { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.constraints.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.Constraints", len)?; + if !self.constraints.is_empty() { + struct_ser.serialize_field("constraints", &self.constraints)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for Constraints { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "constraints", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Constraints, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "constraints" => Ok(GeneratedField::Constraints), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = Constraints; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.Constraints") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut constraints__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Constraints => { + if constraints__.is_some() { + return Err(serde::de::Error::duplicate_field("constraints")); + } + constraints__ = Some(map_.next_value()?); + } + } + } + Ok(Constraints { + constraints: constraints__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.Constraints", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for CreateCatalogNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -3726,6 +3927,9 @@ impl serde::Serialize for CreateExternalTableNode { if !self.options.is_empty() { len += 1; } + if self.constraints.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.CreateExternalTableNode", len)?; if let Some(v) = self.name.as_ref() { struct_ser.serialize_field("name", v)?; @@ -3766,6 +3970,9 @@ impl serde::Serialize for CreateExternalTableNode { if !self.options.is_empty() { struct_ser.serialize_field("options", &self.options)?; } + if let Some(v) = self.constraints.as_ref() { + struct_ser.serialize_field("constraints", v)?; + } struct_ser.end() } } @@ -3795,6 +4002,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { "orderExprs", "unbounded", "options", + "constraints", ]; #[allow(clippy::enum_variant_names)] @@ -3812,6 +4020,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { OrderExprs, Unbounded, Options, + Constraints, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -3846,6 +4055,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { "orderExprs" | "order_exprs" => Ok(GeneratedField::OrderExprs), "unbounded" => Ok(GeneratedField::Unbounded), "options" => Ok(GeneratedField::Options), + "constraints" => Ok(GeneratedField::Constraints), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -3878,6 +4088,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { let mut order_exprs__ = None; let mut unbounded__ = None; let mut options__ = None; + let mut constraints__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Name => { @@ -3960,6 +4171,12 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { map_.next_value::>()? ); } + GeneratedField::Constraints => { + if constraints__.is_some() { + return Err(serde::de::Error::duplicate_field("constraints")); + } + constraints__ = map_.next_value()?; + } } } Ok(CreateExternalTableNode { @@ -3976,6 +4193,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { order_exprs: order_exprs__.unwrap_or_default(), unbounded: unbounded__.unwrap_or_default(), options: options__.unwrap_or_default(), + constraints: constraints__, }) } } @@ -18228,6 +18446,100 @@ impl<'de> serde::Deserialize<'de> for PrepareNode { deserializer.deserialize_struct("datafusion.PrepareNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PrimaryKeyConstraint { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.indices.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PrimaryKeyConstraint", len)?; + if !self.indices.is_empty() { + struct_ser.serialize_field("indices", &self.indices.iter().map(ToString::to_string).collect::>())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PrimaryKeyConstraint { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "indices", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Indices, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "indices" => Ok(GeneratedField::Indices), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PrimaryKeyConstraint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PrimaryKeyConstraint") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut indices__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Indices => { + if indices__.is_some() { + return Err(serde::de::Error::duplicate_field("indices")); + } + indices__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) + ; + } + } + } + Ok(PrimaryKeyConstraint { + indices: indices__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PrimaryKeyConstraint", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for ProjectionColumns { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -23166,6 +23478,100 @@ impl<'de> serde::Deserialize<'de> for UnionNode { deserializer.deserialize_struct("datafusion.UnionNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for UniqueConstraint { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.indices.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.UniqueConstraint", len)?; + if !self.indices.is_empty() { + struct_ser.serialize_field("indices", &self.indices.iter().map(ToString::to_string).collect::>())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for UniqueConstraint { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "indices", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Indices, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "indices" => Ok(GeneratedField::Indices), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = UniqueConstraint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.UniqueConstraint") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut indices__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Indices => { + if indices__.is_some() { + return Err(serde::de::Error::duplicate_field("indices")); + } + indices__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) + ; + } + } + } + Ok(UniqueConstraint { + indices: indices__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.UniqueConstraint", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for ValuesNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index 3382fa17fe58a..ca20cd35cb557 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -291,6 +291,41 @@ pub struct EmptyRelationNode { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PrimaryKeyConstraint { + #[prost(uint64, repeated, tag = "1")] + pub indices: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct UniqueConstraint { + #[prost(uint64, repeated, tag = "1")] + pub indices: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct Constraint { + #[prost(oneof = "constraint::ConstraintMode", tags = "1, 2")] + pub constraint_mode: ::core::option::Option, +} +/// Nested message and enum types in `Constraint`. +pub mod constraint { + #[allow(clippy::derive_partial_eq_without_eq)] + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum ConstraintMode { + #[prost(message, tag = "1")] + PrimaryKey(super::PrimaryKeyConstraint), + #[prost(message, tag = "2")] + Unique(super::UniqueConstraint), + } +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct Constraints { + #[prost(message, repeated, tag = "1")] + pub constraints: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct CreateExternalTableNode { #[prost(message, optional, tag = "12")] pub name: ::core::option::Option, @@ -321,6 +356,8 @@ pub struct CreateExternalTableNode { ::prost::alloc::string::String, ::prost::alloc::string::String, >, + #[prost(message, optional, tag = "15")] + pub constraints: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index f8746ef4fd6cb..aa56c7b19d0c2 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -31,8 +31,8 @@ use arrow::datatypes::{ }; use datafusion::execution::registry::FunctionRegistry; use datafusion_common::{ - internal_err, Column, DFField, DFSchema, DFSchemaRef, DataFusionError, - OwnedTableReference, Result, ScalarValue, + internal_err, Column, Constraint, Constraints, DFField, DFSchema, DFSchemaRef, + DataFusionError, OwnedTableReference, Result, ScalarValue, }; use datafusion_expr::{ abs, acos, acosh, array, array_append, array_concat, array_dims, array_element, @@ -880,6 +880,33 @@ impl From for JoinConstraint { } } +impl From for Constraints { + fn from(constraints: protobuf::Constraints) -> Self { + Constraints::new_unverified( + constraints + .constraints + .into_iter() + .map(|item| item.into()) + .collect(), + ) + } +} + +impl From for Constraint { + fn from(value: protobuf::Constraint) -> Self { + match value.constraint_mode.unwrap() { + protobuf::constraint::ConstraintMode::PrimaryKey(elem) => { + Constraint::PrimaryKey( + elem.indices.into_iter().map(|item| item as usize).collect(), + ) + } + protobuf::constraint::ConstraintMode::Unique(elem) => Constraint::Unique( + elem.indices.into_iter().map(|item| item as usize).collect(), + ), + } + } +} + pub fn parse_i32_to_time_unit(value: &i32) -> Result { protobuf::TimeUnit::try_from(*value) .map(|t| t.into()) diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 099352cb589e9..76d44186b5b5e 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -15,6 +15,10 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::Debug; +use std::str::FromStr; +use std::sync::Arc; + use crate::common::{byte_to_string, proto_error, str_to_byte}; use crate::protobuf::logical_plan_node::LogicalPlanType::CustomScan; use crate::protobuf::{CustomTableScanNode, LogicalExprNodeCollection}; @@ -25,6 +29,7 @@ use crate::{ logical_plan_node::LogicalPlanType, LogicalExtensionNode, LogicalPlanNode, }, }; + use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion::{ datasource::{ @@ -38,27 +43,22 @@ use datafusion::{ datasource::{provider_as_source, source_as_provider}, prelude::SessionContext, }; -use datafusion_common::not_impl_err; use datafusion_common::{ - context, internal_err, parsers::CompressionTypeVariant, DataFusionError, - OwnedTableReference, Result, + context, internal_err, not_impl_err, parsers::CompressionTypeVariant, + DataFusionError, OwnedTableReference, Result, }; -use datafusion_expr::logical_plan::DdlStatement; -use datafusion_expr::DropView; use datafusion_expr::{ logical_plan::{ builder::project, Aggregate, CreateCatalog, CreateCatalogSchema, - CreateExternalTable, CreateView, CrossJoin, Distinct, EmptyRelation, Extension, - Join, JoinConstraint, Limit, Prepare, Projection, Repartition, Sort, - SubqueryAlias, TableScan, Values, Window, + CreateExternalTable, CreateView, CrossJoin, DdlStatement, Distinct, + EmptyRelation, Extension, Join, JoinConstraint, Limit, Prepare, Projection, + Repartition, Sort, SubqueryAlias, TableScan, Values, Window, }, - Expr, LogicalPlan, LogicalPlanBuilder, + DropView, Expr, LogicalPlan, LogicalPlanBuilder, }; + use prost::bytes::BufMut; use prost::Message; -use std::fmt::Debug; -use std::str::FromStr; -use std::sync::Arc; pub mod from_proto; pub mod to_proto; @@ -493,6 +493,11 @@ impl AsLogicalPlan for LogicalPlanNode { )) })?; + let constraints = (create_extern_table.constraints.clone()).ok_or_else(|| { + DataFusionError::Internal(String::from( + "Protobuf deserialization error, CreateExternalTableNode was missing required table constraints.", + )) + })?; let definition = if !create_extern_table.definition.is_empty() { Some(create_extern_table.definition.clone()) } else { @@ -532,6 +537,7 @@ impl AsLogicalPlan for LogicalPlanNode { definition, unbounded: create_extern_table.unbounded, options: create_extern_table.options.clone(), + constraints: constraints.into(), }))) } LogicalPlanType::CreateView(create_view) => { @@ -1205,6 +1211,7 @@ impl AsLogicalPlan for LogicalPlanNode { order_exprs, unbounded, options, + constraints, }, )) => { let mut converted_order_exprs: Vec = vec![]; @@ -1235,6 +1242,7 @@ impl AsLogicalPlan for LogicalPlanNode { file_compression_type: file_compression_type.to_string(), unbounded: *unbounded, options: options.clone(), + constraints: Some(constraints.clone().into()), }, )), }) diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 8a8550d05d133..c10855bf25140 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -30,12 +30,14 @@ use crate::protobuf::{ AnalyzedLogicalPlanType, CubeNode, EmptyMessage, GroupingSetNode, LogicalExprList, OptimizedLogicalPlanType, OptimizedPhysicalPlanType, PlaceholderNode, RollupNode, }; + use arrow::datatypes::{ DataType, Field, IntervalMonthDayNanoType, IntervalUnit, Schema, SchemaRef, TimeUnit, UnionMode, }; use datafusion_common::{ - Column, DFField, DFSchema, DFSchemaRef, OwnedTableReference, ScalarValue, + Column, Constraint, Constraints, DFField, DFSchema, DFSchemaRef, OwnedTableReference, + ScalarValue, }; use datafusion_expr::expr::{ self, Alias, Between, BinaryExpr, Cast, GetFieldAccess, GetIndexedField, GroupingSet, @@ -1623,6 +1625,35 @@ impl From for protobuf::JoinConstraint { } } +impl From for protobuf::Constraints { + fn from(value: Constraints) -> Self { + let constraints = value.into_iter().map(|item| item.into()).collect(); + protobuf::Constraints { constraints } + } +} + +impl From for protobuf::Constraint { + fn from(value: Constraint) -> Self { + let res = match value { + Constraint::PrimaryKey(indices) => { + let indices = indices.into_iter().map(|item| item as u64).collect(); + protobuf::constraint::ConstraintMode::PrimaryKey( + protobuf::PrimaryKeyConstraint { indices }, + ) + } + Constraint::Unique(indices) => { + let indices = indices.into_iter().map(|item| item as u64).collect(); + protobuf::constraint::ConstraintMode::PrimaryKey( + protobuf::PrimaryKeyConstraint { indices }, + ) + } + }; + protobuf::Constraint { + constraint_mode: Some(res), + } + } +} + /// Creates a scalar protobuf value from an optional value (T), and /// encoding None as the appropriate datatype fn create_proto_scalar protobuf::scalar_value::Value>( diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 5b53946379c86..11ee8c0876bc8 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -185,6 +185,94 @@ async fn roundtrip_custom_tables() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_custom_memory_tables() -> Result<()> { + let ctx = SessionContext::new(); + // Make sure during round-trip, constraint information is preserved + let query = "CREATE TABLE sales_global_with_pk (zip_code INT, + country VARCHAR(3), + sn INT, + ts TIMESTAMP, + currency VARCHAR(3), + amount FLOAT, + primary key(sn) + ) as VALUES + (0, 'GRC', 0, '2022-01-01 06:00:00'::timestamp, 'EUR', 30.0), + (1, 'FRA', 1, '2022-01-01 08:00:00'::timestamp, 'EUR', 50.0), + (1, 'TUR', 2, '2022-01-01 11:30:00'::timestamp, 'TRY', 75.0), + (1, 'FRA', 3, '2022-01-02 12:00:00'::timestamp, 'EUR', 200.0), + (1, 'TUR', 4, '2022-01-03 10:00:00'::timestamp, 'TRY', 100.0)"; + + let plan = ctx.sql(query).await?.into_optimized_plan()?; + + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx)?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + + Ok(()) +} + +#[tokio::test] +async fn roundtrip_custom_listing_tables() -> Result<()> { + let ctx = SessionContext::new(); + + // Make sure during round-trip, constraint information is preserved + let query = "CREATE EXTERNAL TABLE multiple_ordered_table_with_pk ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER, + d INTEGER, + primary key(c) + ) + STORED AS CSV + WITH HEADER ROW + WITH ORDER (a ASC, b ASC) + WITH ORDER (c ASC) + LOCATION '../core/tests/data/window_2.csv';"; + + let plan = ctx.sql(query).await?.into_optimized_plan()?; + + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx)?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + + Ok(()) +} + +#[tokio::test] +async fn roundtrip_logical_plan_aggregation_with_pk() -> Result<()> { + let ctx = SessionContext::new(); + + ctx.sql( + "CREATE EXTERNAL TABLE multiple_ordered_table_with_pk ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER, + d INTEGER, + primary key(c) + ) + STORED AS CSV + WITH HEADER ROW + WITH ORDER (a ASC, b ASC) + WITH ORDER (c ASC) + LOCATION '../core/tests/data/window_2.csv';", + ) + .await?; + + let query = "SELECT c, b, SUM(d) + FROM multiple_ordered_table_with_pk + GROUP BY c"; + let plan = ctx.sql(query).await?.into_optimized_plan()?; + + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx)?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + + Ok(()) +} + #[tokio::test] async fn roundtrip_logical_plan_aggregation() -> Result<()> { let ctx = SessionContext::new(); diff --git a/datafusion/sql/src/parser.rs b/datafusion/sql/src/parser.rs index f52fdec7f7044..9c104ff18a9b3 100644 --- a/datafusion/sql/src/parser.rs +++ b/datafusion/sql/src/parser.rs @@ -197,6 +197,8 @@ pub struct CreateExternalTable { pub unbounded: bool, /// Table(provider) specific options pub options: HashMap, + /// A table-level constraint + pub constraints: Vec, } impl fmt::Display for CreateExternalTable { @@ -629,7 +631,7 @@ impl<'a> DFParser<'a> { self.parser .parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); let table_name = self.parser.parse_object_name()?; - let (columns, _) = self.parse_columns()?; + let (columns, constraints) = self.parse_columns()?; #[derive(Default)] struct Builder { @@ -748,6 +750,7 @@ impl<'a> DFParser<'a> { .unwrap_or(CompressionTypeVariant::UNCOMPRESSED), unbounded, options: builder.options.unwrap_or(HashMap::new()), + constraints, }; Ok(Statement::CreateExternalTable(create)) } @@ -899,6 +902,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -917,6 +921,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -936,6 +941,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -955,6 +961,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -974,6 +981,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -996,6 +1004,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; } @@ -1023,6 +1032,7 @@ mod tests { )?, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; } @@ -1042,6 +1052,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -1060,6 +1071,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -1078,6 +1090,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -1097,6 +1110,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -1121,6 +1135,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::from([("k1".into(), "v1".into())]), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -1143,6 +1158,7 @@ mod tests { ("k1".into(), "v1".into()), ("k2".into(), "v2".into()), ]), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -1188,6 +1204,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; } @@ -1228,6 +1245,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -1264,6 +1282,7 @@ mod tests { file_compression_type: UNCOMPRESSED, unbounded: false, options: HashMap::new(), + constraints: vec![], }); expect_parse_ok(sql, expected)?; @@ -1312,6 +1331,7 @@ mod tests { ("ROW_GROUP_SIZE".into(), "1024".into()), ("TRUNCATE".into(), "NO".into()), ]), + constraints: vec![], }); expect_parse_ok(sql, expected)?; diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index ab19fa716c9b8..67165d8b931eb 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -15,6 +15,9 @@ // specific language governing permissions and limitations // under the License. +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::Arc; + use crate::parser::{ CopyToSource, CopyToStatement, CreateExternalTable, DFParser, DescribeTableStmt, ExplainStatement, LexOrdering, Statement as DFStatement, @@ -28,8 +31,8 @@ use arrow_schema::DataType; use datafusion_common::file_options::StatementOptions; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::{ - not_impl_err, unqualified_field_not_found, Column, Constraints, DFField, DFSchema, - DFSchemaRef, DataFusionError, ExprSchema, OwnedTableReference, Result, + not_impl_err, plan_err, unqualified_field_not_found, Column, Constraints, DFField, + DFSchema, DFSchemaRef, DataFusionError, ExprSchema, OwnedTableReference, Result, SchemaReference, TableReference, ToDFSchema, }; use datafusion_expr::dml::{CopyOptions, CopyTo}; @@ -49,16 +52,12 @@ use datafusion_expr::{ }; use sqlparser::ast; use sqlparser::ast::{ - Assignment, Expr as SQLExpr, Expr, Ident, ObjectName, ObjectType, Query, SchemaName, - SetExpr, ShowCreateObject, ShowStatementFilter, Statement, TableFactor, - TableWithJoins, TransactionMode, UnaryOperator, Value, + Assignment, ColumnDef, Expr as SQLExpr, Expr, Ident, ObjectName, ObjectType, Query, + SchemaName, SetExpr, ShowCreateObject, ShowStatementFilter, Statement, + TableConstraint, TableFactor, TableWithJoins, TransactionMode, UnaryOperator, Value, }; use sqlparser::parser::ParserError::ParserError; -use datafusion_common::plan_err; -use std::collections::{BTreeMap, HashMap, HashSet}; -use std::sync::Arc; - fn ident_to_string(ident: &Ident) -> String { normalize_ident(ident.to_owned()) } @@ -84,6 +83,54 @@ fn get_schema_name(schema_name: &SchemaName) -> String { } } +/// Construct `TableConstraint`(s) for the given columns by iterating over +/// `columns` and extracting individual inline constraint definitions. +fn calc_inline_constraints_from_columns(columns: &[ColumnDef]) -> Vec { + let mut constraints = vec![]; + for column in columns { + for ast::ColumnOptionDef { name, option } in &column.options { + match option { + ast::ColumnOption::Unique { is_primary } => { + constraints.push(ast::TableConstraint::Unique { + name: name.clone(), + columns: vec![column.name.clone()], + is_primary: *is_primary, + }) + } + ast::ColumnOption::ForeignKey { + foreign_table, + referred_columns, + on_delete, + on_update, + } => constraints.push(ast::TableConstraint::ForeignKey { + name: name.clone(), + columns: vec![], + foreign_table: foreign_table.clone(), + referred_columns: referred_columns.to_vec(), + on_delete: *on_delete, + on_update: *on_update, + }), + ast::ColumnOption::Check(expr) => { + constraints.push(ast::TableConstraint::Check { + name: name.clone(), + expr: Box::new(expr.clone()), + }) + } + // Other options are not constraint related. + ast::ColumnOption::Default(_) + | ast::ColumnOption::Null + | ast::ColumnOption::NotNull + | ast::ColumnOption::DialectSpecific(_) + | ast::ColumnOption::CharacterSet(_) + | ast::ColumnOption::Generated { .. } + | ast::ColumnOption::Comment(_) + | ast::ColumnOption::OnUpdate(_) => {} + } + } + } + constraints +} + impl<'a, S: ContextProvider> SqlToRel<'a, S> { /// Generate a logical plan from an DataFusion SQL statement pub fn statement_to_plan(&self, statement: DFStatement) -> Result { @@ -154,18 +201,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { or_replace, .. } if table_properties.is_empty() && with_options.is_empty() => { - let mut constraints = constraints; - for column in &columns { - for option in &column.options { - if let ast::ColumnOption::Unique { is_primary } = option.option { - constraints.push(ast::TableConstraint::Unique { - name: None, - columns: vec![column.name.clone()], - is_primary, - }) - } - } - } + // Merge inline constraints and existing constraints + let mut all_constraints = constraints; + let inline_constraints = calc_inline_constraints_from_columns(&columns); + all_constraints.extend(inline_constraints); match query { Some(query) => { let plan = self.query_to_plan(*query, planner_context)?; @@ -201,7 +240,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { }; let constraints = Constraints::new_from_table_constraints( - &constraints, + &all_constraints, plan.schema(), )?; @@ -224,7 +263,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { }; let plan = LogicalPlan::EmptyRelation(plan); let constraints = Constraints::new_from_table_constraints( - &constraints, + &all_constraints, plan.schema(), )?; Ok(LogicalPlan::Ddl(DdlStatement::CreateMemoryTable( @@ -681,8 +720,14 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { order_exprs, unbounded, options, + constraints, } = statement; + // Merge inline constraints and existing constraints + let mut all_constraints = constraints; + let inline_constraints = calc_inline_constraints_from_columns(&columns); + all_constraints.extend(inline_constraints); + if (file_type == "PARQUET" || file_type == "AVRO" || file_type == "ARROW") && file_compression_type != CompressionTypeVariant::UNCOMPRESSED { @@ -699,7 +744,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // External tables do not support schemas at the moment, so the name is just a table name let name = OwnedTableReference::bare(name); - + let constraints = + Constraints::new_from_table_constraints(&all_constraints, &df_schema)?; Ok(LogicalPlan::Ddl(DdlStatement::CreateExternalTable( PlanCreateExternalTable { schema: df_schema, @@ -715,6 +761,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { order_exprs: ordered_exprs, unbounded, options, + constraints, }, ))) } diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index f4ffea06b7a13..b1de5a12bcd0e 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -1806,6 +1806,14 @@ fn create_external_table_csv() { quick_test(sql, expected); } +#[test] +fn create_external_table_with_pk() { + let sql = "CREATE EXTERNAL TABLE t(c1 int, primary key(c1)) STORED AS CSV LOCATION 'foo.csv'"; + let expected = + "CreateExternalTable: Bare { table: \"t\" } constraints=[PrimaryKey([0])]"; + quick_test(sql, expected); +} + #[test] fn create_schema_with_quoted_name() { let sql = "CREATE SCHEMA \"quoted_schema_name\""; diff --git a/datafusion/sqllogictest/test_files/groupby.slt b/datafusion/sqllogictest/test_files/groupby.slt index 5bb0f31ed5423..b7070a8d7efbe 100644 --- a/datafusion/sqllogictest/test_files/groupby.slt +++ b/datafusion/sqllogictest/test_files/groupby.slt @@ -3093,6 +3093,55 @@ CREATE TABLE sales_global_with_pk_alternate (zip_code INT, (1, 'FRA', 3, '2022-01-02 12:00:00'::timestamp, 'EUR', 200.0), (1, 'TUR', 4, '2022-01-03 10:00:00'::timestamp, 'TRY', 100.0) +# we do not currently support foreign key constraints. +statement error DataFusion error: Error during planning: Foreign key constraints are not currently supported +CREATE TABLE sales_global_with_foreign_key (zip_code INT, + country VARCHAR(3), + sn INT references sales_global_with_pk_alternate(sn), + ts TIMESTAMP, + currency VARCHAR(3), + amount FLOAT +) as VALUES + (0, 'GRC', 0, '2022-01-01 06:00:00'::timestamp, 'EUR', 30.0), + (1, 'FRA', 1, '2022-01-01 08:00:00'::timestamp, 'EUR', 50.0), + (1, 'TUR', 2, '2022-01-01 11:30:00'::timestamp, 'TRY', 75.0), + (1, 'FRA', 3, '2022-01-02 12:00:00'::timestamp, 'EUR', 200.0), + (1, 'TUR', 4, '2022-01-03 10:00:00'::timestamp, 'TRY', 100.0) + +# we do not currently support foreign key +statement error DataFusion error: Error during planning: Foreign key constraints are not currently supported +CREATE TABLE sales_global_with_foreign_key (zip_code INT, + country VARCHAR(3), + sn INT REFERENCES sales_global_with_pk_alternate(sn), + ts TIMESTAMP, + currency VARCHAR(3), + amount FLOAT +) as VALUES + (0, 'GRC', 0, '2022-01-01 06:00:00'::timestamp, 'EUR', 30.0), + (1, 'FRA', 1, '2022-01-01 08:00:00'::timestamp, 'EUR', 50.0), + (1, 'TUR', 2, '2022-01-01 11:30:00'::timestamp, 'TRY', 75.0), + (1, 'FRA', 3, '2022-01-02 12:00:00'::timestamp, 'EUR', 200.0), + (1, 'TUR', 4, '2022-01-03 10:00:00'::timestamp, 'TRY', 100.0) + +# we do not currently support foreign key +# foreign key can be defined with a different syntax. +# we should get the same error. +statement error DataFusion error: Error during planning: Foreign key constraints are not currently supported +CREATE TABLE sales_global_with_foreign_key (zip_code INT, + country VARCHAR(3), + sn INT, + ts TIMESTAMP, + currency VARCHAR(3), + amount FLOAT, + FOREIGN KEY (sn) + REFERENCES sales_global_with_pk_alternate(sn) +) as VALUES + (0, 'GRC', 0, '2022-01-01 06:00:00'::timestamp, 'EUR', 30.0), + (1, 'FRA', 1, '2022-01-01 08:00:00'::timestamp, 'EUR', 50.0), + (1, 'TUR', 2, '2022-01-01 11:30:00'::timestamp, 'TRY', 75.0), + (1, 'FRA', 3, '2022-01-02 12:00:00'::timestamp, 'EUR', 200.0), + (1, 'TUR', 4, '2022-01-03 10:00:00'::timestamp, 'TRY', 100.0) + # create a table for testing, where primary key is composite statement ok CREATE TABLE sales_global_with_composite_pk (zip_code INT, @@ -3563,3 +3612,81 @@ logical_plan Sort: multiple_ordered_table.c ASC NULLS LAST --TableScan: multiple_ordered_table projection=[c] physical_plan CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[c], output_ordering=[c@0 ASC NULLS LAST], has_header=true + +# Create an external table with primary key +# column c +statement ok +CREATE EXTERNAL TABLE multiple_ordered_table_with_pk ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER, + d INTEGER, + primary key(c) +) +STORED AS CSV +WITH HEADER ROW +WITH ORDER (a ASC, b ASC) +WITH ORDER (c ASC) +LOCATION '../core/tests/data/window_2.csv'; + +# We can use column b during selection +# even if it is not among group by expressions +# because column c is primary key. +query TT +EXPLAIN SELECT c, b, SUM(d) +FROM multiple_ordered_table_with_pk +GROUP BY c; +---- +logical_plan +Aggregate: groupBy=[[multiple_ordered_table_with_pk.c, multiple_ordered_table_with_pk.b]], aggr=[[SUM(CAST(multiple_ordered_table_with_pk.d AS Int64))]] +--TableScan: multiple_ordered_table_with_pk projection=[b, c, d] +physical_plan +AggregateExec: mode=FinalPartitioned, gby=[c@0 as c, b@1 as b], aggr=[SUM(multiple_ordered_table_with_pk.d)], ordering_mode=PartiallyOrdered +--SortExec: expr=[c@0 ASC NULLS LAST] +----CoalesceBatchesExec: target_batch_size=8192 +------RepartitionExec: partitioning=Hash([c@0, b@1], 8), input_partitions=8 +--------AggregateExec: mode=Partial, gby=[c@1 as c, b@0 as b], aggr=[SUM(multiple_ordered_table_with_pk.d)], ordering_mode=PartiallyOrdered +----------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1 +------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[b, c, d], output_ordering=[c@1 ASC NULLS LAST], has_header=true + +# drop table multiple_ordered_table_with_pk +statement ok +drop table multiple_ordered_table_with_pk; + +# Create an external table with primary key +# column c, in this case use alternative syntax +# for defining primary key +statement ok +CREATE EXTERNAL TABLE multiple_ordered_table_with_pk ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER primary key, + d INTEGER +) +STORED AS CSV +WITH HEADER ROW +WITH ORDER (a ASC, b ASC) +WITH ORDER (c ASC) +LOCATION '../core/tests/data/window_2.csv'; + +# We can use column b during selection +# even if it is not among group by expressions +# because column c is primary key. +query TT +EXPLAIN SELECT c, b, SUM(d) +FROM multiple_ordered_table_with_pk +GROUP BY c; +---- +logical_plan +Aggregate: groupBy=[[multiple_ordered_table_with_pk.c, multiple_ordered_table_with_pk.b]], aggr=[[SUM(CAST(multiple_ordered_table_with_pk.d AS Int64))]] +--TableScan: multiple_ordered_table_with_pk projection=[b, c, d] +physical_plan +AggregateExec: mode=FinalPartitioned, gby=[c@0 as c, b@1 as b], aggr=[SUM(multiple_ordered_table_with_pk.d)], ordering_mode=PartiallyOrdered +--SortExec: expr=[c@0 ASC NULLS LAST] +----CoalesceBatchesExec: target_batch_size=8192 +------RepartitionExec: partitioning=Hash([c@0, b@1], 8), input_partitions=8 +--------AggregateExec: mode=Partial, gby=[c@1 as c, b@0 as b], aggr=[SUM(multiple_ordered_table_with_pk.d)], ordering_mode=PartiallyOrdered +----------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1 +------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[b, c, d], output_ordering=[c@1 ASC NULLS LAST], has_header=true