Skip to content

perf: Eliminate a join a foreign key makes redundant #24793

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Dandandan wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

perf: Eliminate a join a foreign key makes redundant #24793

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 239 additions & 4 deletions benchmarks/src/tpcds/run.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -102,16 +104,249 @@ 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()
.find(|(name, _)| *name == table)
.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<Constraint> {
let columns = column_names
.iter()
.map(|column_name| schema.index_of(column_name).ok())
.collect::<Option<Vec<_>>>()?;
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 {
Expand Down
36 changes: 33 additions & 3 deletions datafusion/common/src/functional_dependencies.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -33,6 +33,17 @@ pub enum Constraint {
PrimaryKey(Vec<usize>),
/// Columns with the given indices form a composite unique key:
Unique(Vec<usize>),
/// 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<usize>,
/// The table the columns reference.
referenced_table: TableReference,
/// Names of the referenced columns, which form a key of that table.
referenced_columns: Vec<String>,
},
}

/// This object encapsulates a list of functional constraints:
Expand Down Expand Up @@ -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::<Vec<_>>();
Expand Down Expand Up @@ -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 {
Expand All @@ -220,10 +247,13 @@ impl FunctionalDependencies {
(0..n_field).collect::<Vec<_>>(),
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::<Vec<_>>();
Self::new(dependencies)
Expand Down
6 changes: 5 additions & 1 deletion datafusion/datasource/src/file_scan_config/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
};
Expand Down
Loading