diff --git a/benchmarks/src/tpcds/run.rs b/benchmarks/src/tpcds/run.rs index 3eaaf172c0f16..abea3102c96c0 100644 --- a/benchmarks/src/tpcds/run.rs +++ b/benchmarks/src/tpcds/run.rs @@ -35,7 +35,9 @@ use datafusion::physical_plan::{collect, displayable}; use datafusion::prelude::*; use datafusion_common::instant::Instant; use datafusion_common::utils::get_available_parallelism; -use datafusion_common::{Constraint, Constraints, DEFAULT_PARQUET_EXTENSION, plan_err}; +use datafusion_common::{ + Constraint, Constraints, DEFAULT_PARQUET_EXTENSION, TableReference, plan_err, +}; use clap::Args; use log::info; @@ -102,8 +104,213 @@ static TPCDS_PRIMARY_KEYS: &[(&str, &[&str])] = &[ ("web_site", &["web_site_sk"]), ]; -/// Get the constraints for a TPC-DS table. Only primary keys are returned; -/// TPC-DS also defines foreign keys, but those are currently unsupported. +/// The foreign keys TPC-DS defines, as (table, columns, referenced table, +/// referenced columns). Only the single-column surrogate key references are +/// listed; the composite fact-to-fact references are not foreign keys. +static TPCDS_FOREIGN_KEYS: &[(&str, &[&str], &str, &[&str])] = &[ + ( + "catalog_returns", + &["cr_returned_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ("catalog_returns", &["cr_item_sk"], "item", &["i_item_sk"]), + ( + "catalog_returns", + &["cr_reason_sk"], + "reason", + &["r_reason_sk"], + ), + ( + "catalog_returns", + &["cr_ship_mode_sk"], + "ship_mode", + &["sm_ship_mode_sk"], + ), + ( + "catalog_returns", + &["cr_warehouse_sk"], + "warehouse", + &["w_warehouse_sk"], + ), + ( + "catalog_sales", + &["cs_sold_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ( + "catalog_sales", + &["cs_ship_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ("catalog_sales", &["cs_item_sk"], "item", &["i_item_sk"]), + ( + "catalog_sales", + &["cs_promo_sk"], + "promotion", + &["p_promo_sk"], + ), + ( + "catalog_sales", + &["cs_ship_mode_sk"], + "ship_mode", + &["sm_ship_mode_sk"], + ), + ( + "catalog_sales", + &["cs_warehouse_sk"], + "warehouse", + &["w_warehouse_sk"], + ), + ( + "catalog_sales", + &["cs_call_center_sk"], + "call_center", + &["cc_call_center_sk"], + ), + ( + "catalog_sales", + &["cs_catalog_page_sk"], + "catalog_page", + &["cp_catalog_page_sk"], + ), + ( + "customer", + &["c_current_addr_sk"], + "customer_address", + &["ca_address_sk"], + ), + ( + "customer", + &["c_current_cdemo_sk"], + "customer_demographics", + &["cd_demo_sk"], + ), + ( + "customer", + &["c_current_hdemo_sk"], + "household_demographics", + &["hd_demo_sk"], + ), + ( + "customer", + &["c_first_sales_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ( + "customer", + &["c_first_shipto_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ( + "household_demographics", + &["hd_income_band_sk"], + "income_band", + &["ib_income_band_sk"], + ), + ("inventory", &["inv_date_sk"], "date_dim", &["d_date_sk"]), + ("inventory", &["inv_item_sk"], "item", &["i_item_sk"]), + ( + "inventory", + &["inv_warehouse_sk"], + "warehouse", + &["w_warehouse_sk"], + ), + ("promotion", &["p_item_sk"], "item", &["i_item_sk"]), + ( + "store_returns", + &["sr_returned_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ("store_returns", &["sr_item_sk"], "item", &["i_item_sk"]), + ( + "store_returns", + &["sr_reason_sk"], + "reason", + &["r_reason_sk"], + ), + ("store_returns", &["sr_store_sk"], "store", &["s_store_sk"]), + ( + "store_sales", + &["ss_sold_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ("store_sales", &["ss_item_sk"], "item", &["i_item_sk"]), + ( + "store_sales", + &["ss_promo_sk"], + "promotion", + &["p_promo_sk"], + ), + ("store_sales", &["ss_store_sk"], "store", &["s_store_sk"]), + ( + "web_returns", + &["wr_returned_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ("web_returns", &["wr_item_sk"], "item", &["i_item_sk"]), + ("web_returns", &["wr_reason_sk"], "reason", &["r_reason_sk"]), + ( + "web_sales", + &["ws_sold_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ( + "web_sales", + &["ws_ship_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ("web_sales", &["ws_item_sk"], "item", &["i_item_sk"]), + ("web_sales", &["ws_promo_sk"], "promotion", &["p_promo_sk"]), + ( + "web_sales", + &["ws_ship_mode_sk"], + "ship_mode", + &["sm_ship_mode_sk"], + ), + ( + "web_sales", + &["ws_warehouse_sk"], + "warehouse", + &["w_warehouse_sk"], + ), + ( + "web_sales", + &["ws_web_page_sk"], + "web_page", + &["wp_web_page_sk"], + ), + ( + "web_sales", + &["ws_web_site_sk"], + "web_site", + &["web_site_sk"], + ), + ( + "web_site", + &["web_open_date_sk"], + "date_dim", + &["d_date_sk"], + ), + ( + "web_site", + &["web_close_date_sk"], + "date_dim", + &["d_date_sk"], + ), +]; + +/// Get the constraints for a TPC-DS table: its primary key and its foreign +/// keys. fn table_constraints(table: &str, schema: &Schema) -> Constraints { let columns = TPCDS_PRIMARY_KEYS .iter() @@ -111,7 +318,35 @@ fn table_constraints(table: &str, schema: &Schema) -> Constraints { .map(|(_, columns)| *columns) .unwrap_or_else(|| unimplemented!("unknown TPC-DS table: {table}")); - Constraints::new_unverified(vec![primary_key(schema, columns)]) + let mut constraints = vec![primary_key(schema, columns)]; + constraints.extend( + TPCDS_FOREIGN_KEYS + .iter() + .filter(|(name, ..)| *name == table) + .filter_map(|(_, columns, referenced_table, referenced_columns)| { + foreign_key(schema, columns, referenced_table, referenced_columns) + }), + ); + Constraints::new_unverified(constraints) +} + +/// Builds a foreign key constraint, or `None` when the referencing columns are +/// not in the schema (some generators omit columns). +fn foreign_key( + schema: &Schema, + column_names: &[&str], + referenced_table: &str, + referenced_columns: &[&str], +) -> Option { + let columns = column_names + .iter() + .map(|column_name| schema.index_of(column_name).ok()) + .collect::>>()?; + Some(Constraint::ForeignKey { + columns, + referenced_table: TableReference::bare(referenced_table.to_string()), + referenced_columns: referenced_columns.iter().map(|c| c.to_string()).collect(), + }) } fn primary_key(schema: &Schema, column_names: &[&str]) -> Constraint { diff --git a/datafusion/common/src/functional_dependencies.rs b/datafusion/common/src/functional_dependencies.rs index e8275aac2da4c..9d8c16d291bda 100644 --- a/datafusion/common/src/functional_dependencies.rs +++ b/datafusion/common/src/functional_dependencies.rs @@ -23,7 +23,7 @@ use std::ops::Deref; use std::vec::IntoIter; use crate::utils::{merge_and_order_indices, set_difference}; -use crate::{DFSchema, HashSet, JoinType}; +use crate::{DFSchema, HashSet, JoinType, TableReference}; /// This object defines a constraint on a table. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] @@ -33,6 +33,17 @@ pub enum Constraint { PrimaryKey(Vec), /// Columns with the given indices form a composite unique key: Unique(Vec), + /// Columns with the given indices reference `referenced_columns` of + /// `referenced_table`: every non-NULL tuple of values here also occurs + /// there. Like the other constraints, this is taken on trust. + ForeignKey { + /// Indices of the referencing columns in this table's schema. + columns: Vec, + /// The table the columns reference. + referenced_table: TableReference, + /// Names of the referenced columns, which form a key of that table. + referenced_columns: Vec, + }, } /// This object encapsulates a list of functional constraints: @@ -81,6 +92,22 @@ impl Constraints { (new_indices.len() == indices.len()) .then_some(Constraint::Unique(new_indices)) } + Constraint::ForeignKey { + columns, + referenced_table, + referenced_columns, + } => { + let new_indices = + update_elements_with_matching_indices(columns, proj_indices); + // Only keep the constraint if all columns are preserved: + (new_indices.len() == columns.len()).then_some( + Constraint::ForeignKey { + columns: new_indices, + referenced_table: referenced_table.clone(), + referenced_columns: referenced_columns.clone(), + }, + ) + } } }) .collect::>(); @@ -206,7 +233,7 @@ impl FunctionalDependencies { // Construct dependency objects based on each individual constraint: let dependencies = constraints .iter() - .map(|constraint| { + .filter_map(|constraint| { // All the field indices are associated with the whole table // since we are dealing with table level constraints: let dependency = match constraint { @@ -220,10 +247,13 @@ impl FunctionalDependencies { (0..n_field).collect::>(), true, ), + // A foreign key says where the values also occur, not + // that they determine anything in this table. + Constraint::ForeignKey { .. } => return None, }; // As primary keys are guaranteed to be unique, set the // functional dependency mode to `Dependency::Single`: - dependency.with_mode(Dependency::Single) + Some(dependency.with_mode(Dependency::Single)) }) .collect::>(); Self::new(dependencies) diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 73f8b8afa1904..03f97c22bf426 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -575,7 +575,11 @@ impl FileScanConfigBuilder { fn add_key_distinct_counts(constraints: &Constraints, statistics: &mut Statistics) { let num_rows = statistics.num_rows; for constraint in constraints.iter() { - let (Constraint::PrimaryKey(indices) | Constraint::Unique(indices)) = constraint; + // A foreign key says nothing about how many distinct values are here. + let (Constraint::PrimaryKey(indices) | Constraint::Unique(indices)) = constraint + else { + continue; + }; let [index] = indices[..] else { continue; }; diff --git a/datafusion/optimizer/src/eliminate_join.rs b/datafusion/optimizer/src/eliminate_join.rs index 56aa8887065be..f8e8389313d45 100644 --- a/datafusion/optimizer/src/eliminate_join.rs +++ b/datafusion/optimizer/src/eliminate_join.rs @@ -71,12 +71,14 @@ //! duplicate-sensitivity (projection, aggregate, sort, ...) adjust it first. use crate::utils::for_each_referenced_index; use crate::{OptimizerConfig, OptimizerRule}; -use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{ - DFSchema, Dependency, HashSet, NullEquality, Result, ScalarValue, + Constraint, DFSchema, Dependency, HashSet, NullEquality, Result, ScalarValue, + TableReference, }; +use datafusion_expr::utils::conjunction; use datafusion_expr::{ - Expr, JoinType, + Expr, JoinType, LogicalPlanBuilder, TableSource, logical_plan::{ Aggregate, Distinct, DistinctOn, EmptyRelation, Filter, Join, Limit, LogicalPlan, Partitioning, Projection, Repartition, Sort, SubqueryAlias, @@ -431,6 +433,18 @@ fn rewrite_join( )?; return Ok(Transformed::yes(right.data)); } + JoinRewrite::ReplaceWithLeftWhereNotNull(predicate) => { + let left = rewrite_subtree( + Arc::unwrap_or_clone(join.left), + visible_left, + duplicate_insensitive, + )?; + return Ok(Transformed::yes( + LogicalPlanBuilder::from(left.data) + .filter(predicate)? + .build()?, + )); + } JoinRewrite::Join(join_type) => join_type, }; @@ -522,6 +536,10 @@ enum JoinRewrite { ReplaceWithLeft, /// The join has no observable effect; replace it with its right input. ReplaceWithRight, + /// A foreign key guarantees every left row matches, so the join only tests + /// that the referencing columns are non-NULL. Replace it with its left + /// input under this filter. + ReplaceWithLeftWhereNotNull(Expr), } /// Chooses a cheaper form for a join: removes an outer join whose non-preserved @@ -570,6 +588,11 @@ fn rewritten_join_type( } if can_remove_right { + // A foreign key from the left to the right's table means every left row + // already has a match, so the join only filters out NULL keys. + if let Some(predicate) = foreign_key_makes_join_redundant(join) { + return JoinRewrite::ReplaceWithLeftWhereNotNull(predicate); + } return JoinRewrite::Join(JoinType::LeftSemi); } if can_remove_left { @@ -615,6 +638,151 @@ fn split_join_output_columns( } } +/// When the left side declares a foreign key to the right side's table on +/// exactly the join keys, every left row already matches, so an inner join only +/// filters out rows whose foreign key is NULL. Returns that filter. +/// +/// The right side has to be an unfiltered scan of the referenced table: a +/// foreign key promises the value occurs somewhere in that table, not that it +/// survives a predicate. The NULL filter is emitted even for columns declared +/// NOT NULL, because an outer join below the left side could have padded them. +fn foreign_key_makes_join_redundant(join: &Join) -> Option { + if join.filter.is_some() || join.on.is_empty() { + return None; + } + let referenced_table = unfiltered_scan(&join.right)?; + + // The pairs this join equates, as (referencing column, referenced column). + let mut referencing = Vec::with_capacity(join.on.len()); + let mut joined_pairs = Vec::with_capacity(join.on.len()); + for (left, right) in &join.on { + let left = left.try_as_col()?; + let right = right.try_as_col()?; + referencing.push(left.clone()); + joined_pairs.push((left.name.clone(), right.name.clone())); + } + + // They must all come from the one table that declares the foreign key. + let relation = referencing.first()?.relation.as_ref()?; + if referencing + .iter() + .any(|column| column.relation.as_ref() != Some(relation)) + { + return None; + } + let source = scan_source(&join.left, relation)?; + let schema = source.schema(); + joined_pairs.sort(); + + let declared = source.constraints()?.iter().any(|constraint| { + let Constraint::ForeignKey { + columns, + referenced_table: table, + referenced_columns, + } = constraint + else { + return false; + }; + if table.table() != referenced_table.table() + || columns.len() != joined_pairs.len() + || columns.len() != referenced_columns.len() + { + return false; + } + let mut declared_pairs = columns + .iter() + .zip(referenced_columns) + .map(|(&index, referenced)| { + schema + .fields() + .get(index) + .map(|field| (field.name().clone(), referenced.clone())) + }) + .collect::>>() + .unwrap_or_default(); + declared_pairs.sort(); + declared_pairs == joined_pairs + }); + if !declared { + return None; + } + + conjunction( + referencing + .into_iter() + .map(|column| Expr::Column(column).is_not_null()), + ) +} + +/// The table a plan scans, when it is a bare scan of a single table: only +/// projections and aliases may sit above it, so every row of the table reaches +/// the join. +fn unfiltered_scan(plan: &LogicalPlan) -> Option { + match plan { + LogicalPlan::TableScan(scan) + if scan.filters.is_empty() && scan.fetch.is_none() => + { + Some(scan.table_name.clone()) + } + LogicalPlan::Projection(Projection { input, .. }) + | LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => { + unfiltered_scan(input) + } + _ => None, + } +} + +/// Finds the source of the relation the join keys are qualified with, anywhere +/// below `plan`. The qualifier is either a table name or an alias introduced by +/// a `SubqueryAlias`, so both are matched. +fn scan_source( + plan: &LogicalPlan, + relation: &TableReference, +) -> Option> { + let mut found = None; + plan.apply(|node| { + match node { + LogicalPlan::TableScan(scan) + if scan.table_name.table() == relation.table() => + { + found = Some(Arc::clone(&scan.source)); + return Ok(TreeNodeRecursion::Stop); + } + LogicalPlan::SubqueryAlias(alias) + if alias.alias.table() == relation.table() => + { + // The alias renames whatever it wraps, so look inside it for + // the scan the columns actually come from. + found = only_scan_source(&alias.input); + return Ok(TreeNodeRecursion::Stop); + } + _ => {} + } + Ok(TreeNodeRecursion::Continue) + }) + .ok()?; + found +} + +/// The source of the single table a plan scans, or `None` when it scans none or +/// several. +fn only_scan_source(plan: &LogicalPlan) -> Option> { + let mut found = None; + let mut several = false; + plan.apply(|node| { + if let LogicalPlan::TableScan(scan) = node { + if found.is_some() { + several = true; + return Ok(TreeNodeRecursion::Stop); + } + found = Some(Arc::clone(&scan.source)); + } + Ok(TreeNodeRecursion::Continue) + }) + .ok()?; + (!several).then_some(found).flatten() +} + fn side_unique_on_join<'a>( schema: &DFSchema, join_exprs: impl Iterator, diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 08c05efe0ccc0..794761fac38e9 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -779,6 +779,8 @@ impl EquivalenceProperties { normal_exprs: &[PhysicalSortExpr], ) -> bool { self.constraints.iter().any(|constraint| match constraint { + // A foreign key says nothing about uniqueness or ordering here. + Constraint::ForeignKey { .. } => false, Constraint::PrimaryKey(indices) | Constraint::Unique(indices) => { let check_null = matches!(constraint, Constraint::Unique(_)); let normalized_size = normal_exprs.len(); @@ -823,6 +825,8 @@ impl EquivalenceProperties { /// unique constraints, also verifies nullable columns. fn satisfied_by_constraints(&self, normal_reqs: &[PhysicalSortRequirement]) -> bool { self.constraints.iter().any(|constraint| match constraint { + // A foreign key says nothing about uniqueness or ordering here. + Constraint::ForeignKey { .. } => false, Constraint::PrimaryKey(indices) | Constraint::Unique(indices) => { let check_null = matches!(constraint, Constraint::Unique(_)); let normalized_size = normal_reqs.len(); diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 27d1101036d9b..a989a283ab599 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -66,10 +66,17 @@ message UniqueConstraint{ repeated uint64 indices = 1; } +message ForeignKeyConstraint{ + repeated uint64 indices = 1; + string referenced_table = 2; + repeated string referenced_columns = 3; +} + message Constraint{ oneof constraint_mode{ PrimaryKeyConstraint primary_key = 1; UniqueConstraint unique = 2; + ForeignKeyConstraint foreign_key = 3; } } diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index cf3a3cb75a0f0..a491148e0818b 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -761,6 +761,13 @@ impl From for Constraint { protobuf::constraint::ConstraintMode::Unique(elem) => Constraint::Unique( elem.indices.into_iter().map(|item| item as usize).collect(), ), + protobuf::constraint::ConstraintMode::ForeignKey(elem) => { + Constraint::ForeignKey { + columns: elem.indices.into_iter().map(|item| item as usize).collect(), + referenced_table: TableReference::from(elem.referenced_table), + referenced_columns: elem.referenced_columns, + } + } } } } @@ -897,6 +904,13 @@ impl From<&protobuf::Constraint> for Constraint { elem.indices.iter().map(|&item| item as usize).collect(), ) } + Some(protobuf::constraint::ConstraintMode::ForeignKey(elem)) => { + Constraint::ForeignKey { + columns: elem.indices.iter().map(|&item| item as usize).collect(), + referenced_table: TableReference::from(elem.referenced_table.clone()), + referenced_columns: elem.referenced_columns.clone(), + } + } None => panic!("constraint_mode not set"), } } diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index c222cd1cb8687..30777df809492 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -1392,6 +1392,9 @@ impl serde::Serialize for Constraint { constraint::ConstraintMode::Unique(v) => { struct_ser.serialize_field("unique", v)?; } + constraint::ConstraintMode::ForeignKey(v) => { + struct_ser.serialize_field("foreignKey", v)?; + } } } struct_ser.end() @@ -1407,12 +1410,15 @@ impl<'de> serde::Deserialize<'de> for Constraint { "primary_key", "primaryKey", "unique", + "foreign_key", + "foreignKey", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { PrimaryKey, Unique, + ForeignKey, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -1436,6 +1442,7 @@ impl<'de> serde::Deserialize<'de> for Constraint { match value { "primaryKey" | "primary_key" => Ok(GeneratedField::PrimaryKey), "unique" => Ok(GeneratedField::Unique), + "foreignKey" | "foreign_key" => Ok(GeneratedField::ForeignKey), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -1470,6 +1477,13 @@ impl<'de> serde::Deserialize<'de> for Constraint { return Err(serde::de::Error::duplicate_field("unique")); } constraint_mode__ = map_.next_value::<::std::option::Option<_>>()?.map(constraint::ConstraintMode::Unique) +; + } + GeneratedField::ForeignKey => { + if constraint_mode__.is_some() { + return Err(serde::de::Error::duplicate_field("foreignKey")); + } + constraint_mode__ = map_.next_value::<::std::option::Option<_>>()?.map(constraint::ConstraintMode::ForeignKey) ; } } @@ -4441,6 +4455,136 @@ impl<'de> serde::Deserialize<'de> for FixedSizeList { deserializer.deserialize_struct("datafusion_common.FixedSizeList", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for ForeignKeyConstraint { + #[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; + } + if !self.referenced_table.is_empty() { + len += 1; + } + if !self.referenced_columns.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion_common.ForeignKeyConstraint", len)?; + if !self.indices.is_empty() { + struct_ser.serialize_field("indices", &self.indices.iter().map(ToString::to_string).collect::>())?; + } + if !self.referenced_table.is_empty() { + struct_ser.serialize_field("referencedTable", &self.referenced_table)?; + } + if !self.referenced_columns.is_empty() { + struct_ser.serialize_field("referencedColumns", &self.referenced_columns)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ForeignKeyConstraint { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "indices", + "referenced_table", + "referencedTable", + "referenced_columns", + "referencedColumns", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Indices, + ReferencedTable, + ReferencedColumns, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> 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), + "referencedTable" | "referenced_table" => Ok(GeneratedField::ReferencedTable), + "referencedColumns" | "referenced_columns" => Ok(GeneratedField::ReferencedColumns), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ForeignKeyConstraint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion_common.ForeignKeyConstraint") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut indices__ = None; + let mut referenced_table__ = None; + let mut referenced_columns__ = 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()) + ; + } + GeneratedField::ReferencedTable => { + if referenced_table__.is_some() { + return Err(serde::de::Error::duplicate_field("referencedTable")); + } + referenced_table__ = Some(map_.next_value()?); + } + GeneratedField::ReferencedColumns => { + if referenced_columns__.is_some() { + return Err(serde::de::Error::duplicate_field("referencedColumns")); + } + referenced_columns__ = Some(map_.next_value()?); + } + } + } + Ok(ForeignKeyConstraint { + indices: indices__.unwrap_or_default(), + referenced_table: referenced_table__.unwrap_or_default(), + referenced_columns: referenced_columns__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion_common.ForeignKeyConstraint", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for IntervalDayTimeValue { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index bdbe38538e1d7..267d4c5527618 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -58,8 +58,17 @@ pub struct UniqueConstraint { pub indices: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ForeignKeyConstraint { + #[prost(uint64, repeated, tag = "1")] + pub indices: ::prost::alloc::vec::Vec, + #[prost(string, tag = "2")] + pub referenced_table: ::prost::alloc::string::String, + #[prost(string, repeated, tag = "3")] + pub referenced_columns: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct Constraint { - #[prost(oneof = "constraint::ConstraintMode", tags = "1, 2")] + #[prost(oneof = "constraint::ConstraintMode", tags = "1, 2, 3")] pub constraint_mode: ::core::option::Option, } /// Nested message and enum types in `Constraint`. @@ -70,6 +79,8 @@ pub mod constraint { PrimaryKey(super::PrimaryKeyConstraint), #[prost(message, tag = "2")] Unique(super::UniqueConstraint), + #[prost(message, tag = "3")] + ForeignKey(super::ForeignKeyConstraint), } } #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index 360981746585b..cbf99a925be83 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -750,6 +750,17 @@ impl From for protobuf::Constraint { protobuf::PrimaryKeyConstraint { indices }, ) } + Constraint::ForeignKey { + columns, + referenced_table, + referenced_columns, + } => protobuf::constraint::ConstraintMode::ForeignKey( + protobuf::ForeignKeyConstraint { + indices: columns.into_iter().map(|item| item as u64).collect(), + referenced_table: referenced_table.to_string(), + referenced_columns, + }, + ), }; protobuf::Constraint { constraint_mode: Some(res), diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index bdbe38538e1d7..267d4c5527618 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -58,8 +58,17 @@ pub struct UniqueConstraint { pub indices: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ForeignKeyConstraint { + #[prost(uint64, repeated, tag = "1")] + pub indices: ::prost::alloc::vec::Vec, + #[prost(string, tag = "2")] + pub referenced_table: ::prost::alloc::string::String, + #[prost(string, repeated, tag = "3")] + pub referenced_columns: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct Constraint { - #[prost(oneof = "constraint::ConstraintMode", tags = "1, 2")] + #[prost(oneof = "constraint::ConstraintMode", tags = "1, 2, 3")] pub constraint_mode: ::core::option::Option, } /// Nested message and enum types in `Constraint`. @@ -70,6 +79,8 @@ pub mod constraint { PrimaryKey(super::PrimaryKeyConstraint), #[prost(message, tag = "2")] Unique(super::UniqueConstraint), + #[prost(message, tag = "3")] + ForeignKey(super::ForeignKeyConstraint), } } #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 33791227b3f81..ed4dc87c12c85 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -181,7 +181,9 @@ fn calc_inline_constraints_from_columns(columns: &[ColumnDef]) -> Vec SqlToRel<'_, S> { )?; Ok(Constraint::PrimaryKey(indices)) } - TableConstraint::ForeignKey { .. } => { - _plan_err!("Foreign key constraints are not currently supported") + TableConstraint::ForeignKey(ForeignKeyConstraint { + columns, + foreign_table, + referred_columns, + .. + }) => { + let field_names = df_schema.field_names(); + let indices = columns + .iter() + .map(|ident| { + let column = + self.ident_normalizer.normalize(ident.clone()); + field_names + .iter() + .position(|item| *item == column) + .ok_or_else(|| { + plan_datafusion_err!( + "Column for foreign key not found in schema: {column}" + ) + }) + }) + .collect::>>()?; + if referred_columns.len() != indices.len() { + return _plan_err!( + "Foreign key must reference as many columns as it has" + ); + } + Ok(Constraint::ForeignKey { + columns: indices, + referenced_table: self.object_name_to_table_reference( + foreign_table.clone(), + )?, + referenced_columns: referred_columns + .iter() + .map(|column| column.value.clone()) + .collect(), + }) } TableConstraint::Check { .. } => { _plan_err!("Check constraints are not currently supported") diff --git a/datafusion/sqllogictest/test_files/functional_dependencies.slt b/datafusion/sqllogictest/test_files/functional_dependencies.slt index c49004190dc60..478b8659be6fc 100644 --- a/datafusion/sqllogictest/test_files/functional_dependencies.slt +++ b/datafusion/sqllogictest/test_files/functional_dependencies.slt @@ -296,6 +296,65 @@ drop table t_null; statement ok drop table t_probe; +########## +## Foreign keys +########## + +statement ok +CREATE TABLE dim (k INT, name TEXT, PRIMARY KEY (k)) AS VALUES (1, 'one'), (2, 'two'); + +statement ok +CREATE TABLE fact (id INT, k INT, FOREIGN KEY (k) REFERENCES dim(k)) AS VALUES +(10, 1), +(11, 2), +(12, NULL); + +# Every fact row matches, so the join only removes NULL keys. Without the +# foreign key this is a semi join instead. +query TT +EXPLAIN SELECT f.id FROM fact f JOIN dim d ON f.k = d.k; +---- +logical_plan +01)SubqueryAlias: f +02)--Projection: fact.id +03)----Filter: fact.k IS NOT NULL +04)------TableScan: fact projection=[id, k] + +query I rowsort +SELECT f.id FROM fact f JOIN dim d ON f.k = d.k; +---- +10 +11 + +# A predicate on the referenced side breaks the guarantee: the key exists in the +# table, but not necessarily in the filtered subset. +query TT +EXPLAIN SELECT f.id FROM fact f JOIN dim d ON f.k = d.k WHERE d.name = 'one'; +---- +logical_plan +01)Projection: f.id +02)--LeftSemi Join: f.k = d.k +03)----SubqueryAlias: f +04)------TableScan: fact projection=[id, k] +05)----SubqueryAlias: d +06)------Projection: dim.k +07)--------Filter: dim.name = Utf8View("one") +08)----------TableScan: dim projection=[k, name] + +# Using a column of the referenced table keeps the join: it supplies data, not +# just existence. +query IT rowsort +SELECT f.id, d.name FROM fact f JOIN dim d ON f.k = d.k; +---- +10 one +11 two + +statement ok +drop table fact; + +statement ok +drop table dim; + ########## ## Cleanup ########## diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 38d1b7821451d..dbe9aff23ee10 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -3317,8 +3317,8 @@ 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 +# a foreign key declared inline on the column +statement ok CREATE TABLE sales_global_with_foreign_key (zip_code INT, country VARCHAR(3), sn INT references sales_global_with_pk_alternate(sn), @@ -3332,9 +3332,9 @@ CREATE TABLE sales_global_with_foreign_key (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 -statement error DataFusion error: Error during planning: Foreign key constraints are not currently supported -CREATE TABLE sales_global_with_foreign_key (zip_code INT, +# the same, spelled with uppercase REFERENCES +statement ok +CREATE TABLE sales_global_with_foreign_key_2 (zip_code INT, country VARCHAR(3), sn INT REFERENCES sales_global_with_pk_alternate(sn), ts TIMESTAMP, @@ -3347,11 +3347,9 @@ CREATE TABLE sales_global_with_foreign_key (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 -# 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, +# a foreign key declared as a table constraint +statement ok +CREATE TABLE sales_global_with_foreign_key_3 (zip_code INT, country VARCHAR(3), sn INT, ts TIMESTAMP,