From d043a7e000dbd1d33506a9dd5b683a9535e5e545 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sun, 30 Aug 2026 11:21:03 +0200 Subject: [PATCH] perf: Eliminate a join a foreign key makes redundant When a fact table joins a dimension on a foreign key, uses none of the dimension's columns, and does not filter it, the join only checks that a row that is guaranteed to exist does exist. `eliminate_join` already turns that into a semi join; with the foreign key declared it can drop the join entirely. DataFusion had no way to declare one. `Constraint` gains a `ForeignKey` variant, the SQL planner accepts `FOREIGN KEY (..) REFERENCES t(..)` and the inline `REFERENCES` form instead of rejecting them, and both round-trip through proto. The join is replaced by the left input under a `col IS NOT NULL` filter rather than dropped outright. A foreign key column may be nullable, and a semi join drops rows whose key is NULL, so the filter preserves that. It also covers an outer join further down having padded the column with NULLs. The rewrite only applies when the referenced side is an unfiltered scan of the referenced table: a foreign key promises the value exists in the table, not that it survives a predicate. `dfbench` declares the TPC-DS foreign keys, which its comment previously noted as unsupported. On TPC-DS SF1 that removes 6 joins from q64, taking it from 347 ms to 239 ms (1.45x), and q18 from 92 ms to 76 ms (1.22x), for -1.2% on the suite. All 99 queries return unchanged results. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SgqwvctZdvR1ZCz2hbkEJC --- benchmarks/src/tpcds/run.rs | 243 +++++++++++++++++- .../common/src/functional_dependencies.rs | 36 ++- .../datasource/src/file_scan_config/mod.rs | 6 +- datafusion/optimizer/src/eliminate_join.rs | 174 ++++++++++++- .../src/equivalence/properties/mod.rs | 4 + .../proto/datafusion_common.proto | 7 + datafusion/proto-common/src/from_proto/mod.rs | 14 + .../proto-common/src/generated/pbjson.rs | 144 +++++++++++ .../proto-common/src/generated/prost.rs | 13 +- datafusion/proto-common/src/to_proto/mod.rs | 11 + .../src/generated/datafusion_proto_common.rs | 13 +- datafusion/sql/src/statement.rs | 43 +++- .../test_files/functional_dependencies.slt | 59 +++++ .../sqllogictest/test_files/group_by.slt | 18 +- 14 files changed, 759 insertions(+), 26 deletions(-) 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,