From f31cc019631688719ad0c579f2651f36c9f104d3 Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Thu, 13 Aug 2026 10:30:51 +0800 Subject: [PATCH 1/3] perf: resolve schema fields by name index instead of scanning Field lookup by name was linear in the schema width, and the alias branch of `Expr::to_field` paid for it twice. Together that made deriving a projection's schema quadratic in the number of columns. Two independent changes: - `DFSchema` gains a lazily built map from field name to the ascending indices carrying it. `index_of_column_by_name` and `qualified_fields_with_unqualified_name` consult it instead of walking every field, and the latter no longer allocates a `Vec` per lookup. Every arm of the lookup rules already required the field name to match, so restricting the walk to same-named candidates and applying the qualifier rules in index order returns exactly what the scan returned, including which duplicate wins. - `Expr::to_field`'s `Expr::Alias` arm resolved the aliased expression twice: `Expr::metadata` is itself `to_field(..).1.metadata()`, so calling it alongside `to_field` walked the inner expression, and thus the schema, a second time for no extra information. The index is derived state and takes no part in equality, and `Debug` is now written by hand so it keeps printing exactly the three real fields. Plan snapshots compare that string and a `HashMap`'s iteration order is not deterministic. `Clone` starts a fresh cache rather than copying one. Timings for `to_field` over W `col AS col` aliases against a W-column schema, which is the shape wide view-matcher style projections produce (per iteration, 3000 iterations): | W | before | after | |-----|----------|----------| | 18 | 26.74 us | 12.03 us | | 40 | 75.40 us | 18.64 us | | 100 | 394.8 us | 47.21 us | | 300 | 3124 us | 141.4 us | Per-expression cost goes from 0.98 us at W=18 to 5.32 us at W=300 before, and holds at about 0.47 us after, i.e. the quadratic term is gone. At narrow widths the alias change is what pays; the index takes over as the schema widens. `name_index_matches_linear_scan` pins the equivalence by keeping the previous scan as a reference implementation and comparing both lookups across qualifier and name combinations, over a schema with the same name under two relations, qualified and unqualified fields, a non-ASCII name and absent names. `name_index_is_derived_state` covers clone, `strip_qualifiers` and `replace_qualifier`. --- datafusion/common/src/dfschema.rs | 241 ++++++++++++++++++++++++++--- datafusion/expr/src/expr_schema.rs | 13 +- 2 files changed, 230 insertions(+), 24 deletions(-) diff --git a/datafusion/common/src/dfschema.rs b/datafusion/common/src/dfschema.rs index 262f1dcf619d9..c0ff0b73b2185 100644 --- a/datafusion/common/src/dfschema.rs +++ b/datafusion/common/src/dfschema.rs @@ -21,7 +21,7 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::fmt::{Display, Formatter}; use std::hash::Hash; -use std::sync::{Arc, LazyLock}; +use std::sync::{Arc, LazyLock, OnceLock}; use crate::error::{_plan_err, _schema_err, DataFusionError, Result}; use crate::{ @@ -108,7 +108,6 @@ pub type DFSchemaRef = Arc; /// let schema: &Schema = df_schema.as_arrow(); /// assert_eq!(schema.fields().len(), 1); /// ``` -#[derive(Debug, Clone, PartialEq, Eq)] pub struct DFSchema { /// Inner Arrow schema reference. inner: SchemaRef, @@ -117,8 +116,47 @@ pub struct DFSchema { field_qualifiers: Vec>, /// Stores functional dependencies in the schema. functional_dependencies: FunctionalDependencies, + /// Lazily built accelerator for name lookups: maps a field name to the + /// ascending list of indices carrying it. Purely derived from `inner`, so + /// it takes no part in equality or `Debug`. + name_index: OnceLock>>, } +impl std::fmt::Debug for DFSchema { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + // Mirrors the previously derived output: `name_index` is a derived + // lookup accelerator and must not appear (its iteration order is not + // deterministic, and plan snapshots compare this string). + f.debug_struct("DFSchema") + .field("inner", &self.inner) + .field("field_qualifiers", &self.field_qualifiers) + .field("functional_dependencies", &self.functional_dependencies) + .finish() + } +} + +impl Clone for DFSchema { + fn clone(&self) -> Self { + // Derived state; rebuild on demand rather than copying on every clone. + Self { + inner: Arc::clone(&self.inner), + field_qualifiers: self.field_qualifiers.clone(), + functional_dependencies: self.functional_dependencies.clone(), + name_index: OnceLock::new(), + } + } +} + +impl PartialEq for DFSchema { + fn eq(&self, other: &Self) -> bool { + self.inner == other.inner + && self.field_qualifiers == other.field_qualifiers + && self.functional_dependencies == other.functional_dependencies + } +} + +impl Eq for DFSchema {} + impl DFSchema { /// Creates an empty `DFSchema` pub fn empty() -> Self { @@ -126,6 +164,7 @@ impl DFSchema { inner: Arc::new(Schema::new([])), field_qualifiers: vec![], functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), } } @@ -164,6 +203,7 @@ impl DFSchema { inner: schema, field_qualifiers: qualifiers, functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), }; dfschema.check_names()?; Ok(dfschema) @@ -180,6 +220,7 @@ impl DFSchema { inner: schema, field_qualifiers: vec![None; field_count], functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), }; dfschema.check_names()?; Ok(dfschema) @@ -198,6 +239,7 @@ impl DFSchema { inner: schema.clone().into(), field_qualifiers: vec![Some(qualifier); schema.fields.len()], functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), }; schema.check_names()?; Ok(schema) @@ -212,6 +254,7 @@ impl DFSchema { inner: Arc::clone(schema), field_qualifiers: qualifiers, functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), }; dfschema.check_names()?; Ok(dfschema) @@ -233,6 +276,7 @@ impl DFSchema { inner: Arc::clone(&self.inner), field_qualifiers: qualifiers, functional_dependencies: self.functional_dependencies.clone(), + name_index: OnceLock::new(), }) } @@ -301,6 +345,7 @@ impl DFSchema { inner: Arc::new(new_schema_with_metadata), field_qualifiers: new_qualifiers, functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), }; new_self.check_names()?; Ok(new_self) @@ -378,26 +423,35 @@ impl DFSchema { (self.field_qualifiers[i].as_ref(), self.field(i)) } + /// Lazily built map from field name to the ascending indices carrying it. + fn name_index(&self) -> &HashMap> { + self.name_index.get_or_init(|| { + let mut map: HashMap> = + HashMap::with_capacity(self.inner.fields().len()); + for (idx, field) in self.inner.fields().iter().enumerate() { + map.entry(field.name().to_owned()).or_default().push(idx); + } + map + }) + } + pub fn index_of_column_by_name( &self, qualifier: Option<&TableReference>, name: &str, ) -> Option { - let mut matches = self - .iter() - .enumerate() - .filter(|(_, (q, f))| match (qualifier, q) { - // field to lookup is qualified. - // current field is qualified and not shared between relations, compare both - // qualifier and name. - (Some(q), Some(field_q)) => q.resolved_eq(field_q) && f.name() == name, - // field to lookup is qualified but current field is unqualified. + // Every arm below requires the field name to match, so only indices + // carrying `name` can match. Look those up instead of scanning the + // whole schema, then apply the qualifier rules in index order so the + // first match is still the one returned. + let candidates = self.name_index().get(name)?; + candidates.iter().copied().find(|&idx| { + match (qualifier, self.field_qualifiers[idx].as_ref()) { + (Some(q), Some(field_q)) => q.resolved_eq(field_q), (Some(_), None) => false, - // field to lookup is unqualified, no need to compare qualifier - (None, Some(_)) | (None, None) => f.name() == name, - }) - .map(|(idx, _)| idx); - matches.next() + (None, Some(_)) | (None, None) => true, + } + }) } /// Find the index of the column with the given qualifier and name, @@ -486,9 +540,17 @@ impl DFSchema { &self, name: &str, ) -> Vec<(Option<&TableReference>, &FieldRef)> { - self.iter() - .filter(|(_, field)| field.name() == name) - .collect() + // Fields are looked up by name far more often than schemas are built, + // so go through the name index rather than scanning every field. + self.name_index() + .get(name) + .map(|indices| { + indices + .iter() + .map(|&idx| self.qualified_field(idx)) + .collect() + }) + .unwrap_or_default() } /// Find all fields that match the given name and convert to column @@ -843,6 +905,7 @@ impl DFSchema { field_qualifiers: vec![None; self.inner.fields.len()], inner: self.inner, functional_dependencies: self.functional_dependencies, + name_index: OnceLock::new(), } } @@ -853,6 +916,7 @@ impl DFSchema { field_qualifiers: vec![Some(qualifier); self.inner.fields.len()], inner: self.inner, functional_dependencies: self.functional_dependencies, + name_index: OnceLock::new(), } } @@ -1124,6 +1188,7 @@ impl TryFrom for DFSchema { inner: schema, field_qualifiers: vec![None; field_count], functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), }; // Without checking names, because schema here may have duplicate field names. // For example, Partial AggregateMode will generate duplicate field names from @@ -1185,6 +1250,7 @@ impl ToDFSchema for Vec { inner: schema.into(), field_qualifiers: vec![None; field_count], functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), }; Ok(dfschema) } @@ -1382,6 +1448,141 @@ mod tests { use super::*; + /// Reference implementation of `index_of_column_by_name` as it was before + /// the name index was introduced: a linear scan over every field, taking + /// the first match in index order. + fn reference_index_of_column_by_name( + schema: &DFSchema, + qualifier: Option<&TableReference>, + name: &str, + ) -> Option { + schema + .iter() + .enumerate() + .filter(|(_, (q, f))| match (qualifier, q) { + (Some(q), Some(field_q)) => q.resolved_eq(field_q) && f.name() == name, + (Some(_), None) => false, + (None, Some(_)) | (None, None) => f.name() == name, + }) + .map(|(idx, _)| idx) + .next() + } + + /// Reference implementation of `qualified_fields_with_unqualified_name` + /// before the name index: a linear scan keeping schema order. + fn reference_qualified_fields_with_unqualified_name<'a>( + schema: &'a DFSchema, + name: &str, + ) -> Vec<(Option<&'a TableReference>, &'a FieldRef)> { + schema + .iter() + .filter(|(_, field)| field.name() == name) + .collect() + } + + /// A schema exercising the cases the lookup rules distinguish: the same + /// name under different qualifiers, the same name both qualified and + /// unqualified, a name unique to one relation, and a non-ASCII name. + fn ambiguous_schema() -> DFSchema { + let f = |name: &str| Arc::new(Field::new(name, DataType::Int32, true)); + let t1 = TableReference::bare("t1"); + let t2 = TableReference::partial("s", "t2"); + DFSchema::new_with_metadata( + vec![ + // Same name under two different relations: the qualifier is + // what disambiguates, and an unqualified lookup must still + // return the first one in schema order. + (Some(t1.clone()), f("a")), + (Some(t1.clone()), f("b")), + (Some(t2.clone()), f("a")), + (Some(t2.clone()), f("d")), + (None, f("c")), + (None, f("\u{1f600}")), + ], + HashMap::new(), + ) + .unwrap() + } + + /// The name index must be a pure lookup accelerator: for every qualifier + /// and name it has to return exactly what the previous linear scan did, + /// including which duplicate wins and which lookups miss. + #[test] + fn name_index_matches_linear_scan() { + let schema = ambiguous_schema(); + + let qualifiers: Vec> = vec![ + None, + Some(TableReference::bare("t1")), + Some(TableReference::partial("s", "t2")), + // Same table, spelled with and without its schema. + Some(TableReference::bare("t2")), + Some(TableReference::full("c", "s", "t2")), + // A relation that is not in the schema at all. + Some(TableReference::bare("nope")), + ]; + let names = ["a", "b", "c", "d", "\u{1f600}", "missing", "A"]; + + for q in &qualifiers { + for name in names { + assert_eq!( + schema.index_of_column_by_name(q.as_ref(), name), + reference_index_of_column_by_name(&schema, q.as_ref(), name), + "index_of_column_by_name disagrees for qualifier {q:?} name {name}" + ); + } + } + + for name in names { + assert_eq!( + schema.qualified_fields_with_unqualified_name(name), + reference_qualified_fields_with_unqualified_name(&schema, name), + "qualified_fields_with_unqualified_name disagrees for {name}" + ); + } + } + + /// The index is derived state, so it must survive the operations that + /// rebuild a schema's qualifiers and must not leak into equality. + #[test] + fn name_index_is_derived_state() { + let schema = ambiguous_schema(); + + // Populate the cache, then check a clone agrees with a fresh schema. + let _ = schema.index_of_column_by_name(None, "a"); + let cloned = schema.clone(); + assert_eq!(schema, cloned); + assert_eq!( + cloned.index_of_column_by_name(None, "a"), + reference_index_of_column_by_name(&cloned, None, "a") + ); + + // Stripping qualifiers changes which field a qualified lookup finds; + // the rebuilt schema must not reuse the old schema's answers. + let stripped = schema.clone().strip_qualifiers(); + for name in ["a", "b", "c"] { + assert_eq!( + stripped.index_of_column_by_name(Some(&TableReference::bare("t1")), name), + reference_index_of_column_by_name( + &stripped, + Some(&TableReference::bare("t1")), + name + ), + "stripped schema disagrees for {name}" + ); + } + + let replaced = schema.replace_qualifier("t9"); + assert_eq!( + replaced.index_of_column_by_name(Some(&TableReference::bare("t9")), "a"), + reference_index_of_column_by_name( + &replaced, + Some(&TableReference::bare("t9")), + "a" + ) + ); + } + /// `qualified_name` doesn't use `TableReference::Display` for performance /// reasons, but check that the output is consistent. #[test] @@ -1671,6 +1872,7 @@ mod tests { inner: Arc::clone(&arrow_schema_ref), field_qualifiers: vec![None; arrow_schema_ref.fields.len()], functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), }; let df_schema_ref = Arc::new(df_schema.clone()); @@ -1717,6 +1919,7 @@ mod tests { inner: Arc::clone(&schema), field_qualifiers: vec![None; schema.fields.len()], functional_dependencies: FunctionalDependencies::empty(), + name_index: OnceLock::new(), }; assert_eq!(df_schema.inner.metadata(), schema.metadata()) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 8927fcf4d0bbe..29006472303ae 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -491,15 +491,18 @@ impl ExprSchemable for Expr { metadata, .. }) => { - let mut combined_metadata = expr.metadata(schema)?; + // Resolve the aliased expression once. `Expr::metadata` is + // itself `to_field(..).1.metadata()`, so calling both walks the + // inner expression -- and, for columns, linearly scans the + // input schema -- a second time for no extra information. + let inner_field = expr.to_field(schema).map(|(_, f)| f)?; + + let mut combined_metadata = FieldMetadata::from(inner_field.metadata()); if let Some(metadata) = metadata { combined_metadata.extend(metadata.clone()); } - Ok(expr - .to_field(schema) - .map(|(_, f)| f)? - .with_field_metadata(&combined_metadata)) + Ok(inner_field.with_field_metadata(&combined_metadata)) } Expr::Negative(expr) => expr.to_field(schema).map(|(_, f)| f), Expr::Column(c) => schema.field_from_column(c).map(Arc::clone), From 8df0e2dd43e73414b876650862d8647c9d86bc7e Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Thu, 13 Aug 2026 11:07:07 +0800 Subject: [PATCH 2/3] fix: reset the name index when a schema is merged in place `merge` is the one method that mutates a `DFSchema` in place, replacing `inner` and extending `field_qualifiers`. A name index built before the merge kept describing the old field set, so every field merged in was invisible to later lookups: `index_of_column_by_name` returned `None` for them and `qualified_fields_with_unqualified_name` left them out. Drop the index at the end of `merge` and let the next lookup rebuild it. `name_index_survives_merge` covers it, comparing against the linear-scan reference after a merge that both adds a new name and reuses an existing one under a different qualifier. It fails without this change. --- datafusion/common/src/dfschema.rs | 45 +++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/datafusion/common/src/dfschema.rs b/datafusion/common/src/dfschema.rs index c0ff0b73b2185..14757a97f74b0 100644 --- a/datafusion/common/src/dfschema.rs +++ b/datafusion/common/src/dfschema.rs @@ -402,6 +402,9 @@ impl DFSchema { let finished_with_metadata = finished.with_metadata(metadata); self.inner = finished_with_metadata.into(); self.field_qualifiers.extend(qualifiers); + // The fields just changed, so the name index no longer describes this + // schema. Drop it and let the next lookup rebuild it. + self.name_index = OnceLock::new(); } /// Get a list of fields for this schema @@ -1542,6 +1545,48 @@ mod tests { } } + /// `merge` mutates a schema in place, so a cache populated beforehand + /// must not survive it: the merged-in fields have to be visible to every + /// later lookup. + #[test] + fn name_index_survives_merge() { + let f = |name: &str| Arc::new(Field::new(name, DataType::Int32, true)); + let t1 = TableReference::bare("t1"); + let t2 = TableReference::bare("t2"); + + let mut schema = + DFSchema::new_with_metadata(vec![(Some(t1.clone()), f("a"))], HashMap::new()) + .unwrap(); + let other = DFSchema::new_with_metadata( + vec![(Some(t2.clone()), f("b")), (Some(t2.clone()), f("a"))], + HashMap::new(), + ) + .unwrap(); + + // Populate the cache before mutating, which is what makes a stale + // cache observable. + assert_eq!(schema.index_of_column_by_name(Some(&t1), "a"), Some(0)); + + schema.merge(&other); + + for (qualifier, name) in [ + (Some(&t1), "a"), + (Some(&t2), "b"), + (Some(&t2), "a"), + (None, "b"), + ] { + assert_eq!( + schema.index_of_column_by_name(qualifier, name), + reference_index_of_column_by_name(&schema, qualifier, name), + "stale lookup after merge for qualifier {qualifier:?} name {name}" + ); + } + assert_eq!( + schema.qualified_fields_with_unqualified_name("b"), + reference_qualified_fields_with_unqualified_name(&schema, "b") + ); + } + /// The index is derived state, so it must survive the operations that /// rebuild a schema's qualifiers and must not leak into equality. #[test] From aabab115d21513f800fba0a3dc37fe03a6135b3d Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Thu, 13 Aug 2026 13:57:52 +0800 Subject: [PATCH 3/3] perf: build the name index only for repeatedly probed schemas Building the index allocates a `String` and a `Vec` per distinct field name, which costs more than several linear scans. Planning creates a very large number of schemas that are looked up once or twice, so paying that up front is a net loss for them: `sql_planner` showed roughly 25% slower logical planning across narrow-schema cases, and even `logical_select_one_from_700` regressed, since `SELECT c1 FROM t700` does a single lookup but was paying for a 700-entry map. Count by-name probes instead and only build the index once a schema has been probed more than `NAME_INDEX_PROBE_THRESHOLD` times, keeping the original scan until then. One-shot lookups no longer pay anything, while the schemas that dominate the wide cases, probed once per column, still switch over and amortise the build. `name_index_and_scan_paths_agree` pins that both paths answer identically: it probes fresh schemas once each so they stay on the scan, probes one schema past the threshold so it switches, asserts the index was actually built, and compares both against the linear-scan reference. --- datafusion/common/src/dfschema.rs | 162 +++++++++++++++++++++++++----- 1 file changed, 135 insertions(+), 27 deletions(-) diff --git a/datafusion/common/src/dfschema.rs b/datafusion/common/src/dfschema.rs index 14757a97f74b0..d77436522f231 100644 --- a/datafusion/common/src/dfschema.rs +++ b/datafusion/common/src/dfschema.rs @@ -21,6 +21,7 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::fmt::{Display, Formatter}; use std::hash::Hash; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock, OnceLock}; use crate::error::{_plan_err, _schema_err, DataFusionError, Result}; @@ -108,6 +109,12 @@ pub type DFSchemaRef = Arc; /// let schema: &Schema = df_schema.as_arrow(); /// assert_eq!(schema.fields().len(), 1); /// ``` +/// Number of by-name probes a schema takes on the linear path before it is +/// worth building [`DFSchema::name_index`]. Building allocates per distinct +/// field name, so a schema looked up only a couple of times is cheaper to +/// scan. +const NAME_INDEX_PROBE_THRESHOLD: usize = 8; + pub struct DFSchema { /// Inner Arrow schema reference. inner: SchemaRef, @@ -116,10 +123,13 @@ pub struct DFSchema { field_qualifiers: Vec>, /// Stores functional dependencies in the schema. functional_dependencies: FunctionalDependencies, - /// Lazily built accelerator for name lookups: maps a field name to the - /// ascending list of indices carrying it. Purely derived from `inner`, so - /// it takes no part in equality or `Debug`. + /// Accelerator for name lookups: maps a field name to the ascending list + /// of indices carrying it. Built only once a schema has been probed often + /// enough to pay for it (see `NAME_INDEX_PROBE_THRESHOLD`). Purely derived + /// from `inner`, so it takes no part in equality or `Debug`. name_index: OnceLock>>, + /// How many times this schema has been probed by name while unindexed. + name_probes: AtomicUsize, } impl std::fmt::Debug for DFSchema { @@ -143,6 +153,7 @@ impl Clone for DFSchema { field_qualifiers: self.field_qualifiers.clone(), functional_dependencies: self.functional_dependencies.clone(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), } } } @@ -165,6 +176,7 @@ impl DFSchema { field_qualifiers: vec![], functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), } } @@ -204,6 +216,7 @@ impl DFSchema { field_qualifiers: qualifiers, functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }; dfschema.check_names()?; Ok(dfschema) @@ -221,6 +234,7 @@ impl DFSchema { field_qualifiers: vec![None; field_count], functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }; dfschema.check_names()?; Ok(dfschema) @@ -240,6 +254,7 @@ impl DFSchema { field_qualifiers: vec![Some(qualifier); schema.fields.len()], functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }; schema.check_names()?; Ok(schema) @@ -255,6 +270,7 @@ impl DFSchema { field_qualifiers: qualifiers, functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }; dfschema.check_names()?; Ok(dfschema) @@ -277,6 +293,7 @@ impl DFSchema { field_qualifiers: qualifiers, functional_dependencies: self.functional_dependencies.clone(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }) } @@ -346,6 +363,7 @@ impl DFSchema { field_qualifiers: new_qualifiers, functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }; new_self.check_names()?; Ok(new_self) @@ -405,6 +423,7 @@ impl DFSchema { // The fields just changed, so the name index no longer describes this // schema. Drop it and let the next lookup rebuild it. self.name_index = OnceLock::new(); + self.name_probes = AtomicUsize::new(0); } /// Get a list of fields for this schema @@ -426,16 +445,30 @@ impl DFSchema { (self.field_qualifiers[i].as_ref(), self.field(i)) } - /// Lazily built map from field name to the ascending indices carrying it. - fn name_index(&self) -> &HashMap> { - self.name_index.get_or_init(|| { + /// Map from field name to the ascending indices carrying it, or `None` + /// while this schema has not been probed enough to justify building it. + /// + /// Building the map allocates per distinct field name, which costs more + /// than several linear scans. Most schemas produced during planning are + /// looked up once or twice, so paying that up front is a net loss; the + /// schemas that dominate are the ones probed once per column. Counting + /// probes lets both cases win: one-shot lookups keep scanning, repeated + /// lookups switch over and amortise the build. + fn name_index(&self) -> Option<&HashMap>> { + if let Some(index) = self.name_index.get() { + return Some(index); + } + if self.name_probes.fetch_add(1, Ordering::Relaxed) < NAME_INDEX_PROBE_THRESHOLD { + return None; + } + Some(self.name_index.get_or_init(|| { let mut map: HashMap> = HashMap::with_capacity(self.inner.fields().len()); for (idx, field) in self.inner.fields().iter().enumerate() { map.entry(field.name().to_owned()).or_default().push(idx); } map - }) + })) } pub fn index_of_column_by_name( @@ -443,18 +476,35 @@ impl DFSchema { qualifier: Option<&TableReference>, name: &str, ) -> Option { - // Every arm below requires the field name to match, so only indices - // carrying `name` can match. Look those up instead of scanning the - // whole schema, then apply the qualifier rules in index order so the - // first match is still the one returned. - let candidates = self.name_index().get(name)?; - candidates.iter().copied().find(|&idx| { - match (qualifier, self.field_qualifiers[idx].as_ref()) { + let matches_qualifier = + |idx: usize| match (qualifier, self.field_qualifiers[idx].as_ref()) { + // field to lookup is qualified. + // current field is qualified and not shared between relations, compare both + // qualifier and name. (Some(q), Some(field_q)) => q.resolved_eq(field_q), + // field to lookup is qualified but current field is unqualified. (Some(_), None) => false, + // field to lookup is unqualified, no need to compare qualifier (None, Some(_)) | (None, None) => true, - } - }) + }; + + // Every arm above also requires the field name to match, so once the + // index exists only the indices carrying `name` can match. Either way + // the first match in index order is the one returned. + match self.name_index() { + Some(index) => index + .get(name)? + .iter() + .copied() + .find(|&idx| matches_qualifier(idx)), + None => self + .inner + .fields() + .iter() + .enumerate() + .find(|(idx, f)| f.name() == name && matches_qualifier(*idx)) + .map(|(idx, _)| idx), + } } /// Find the index of the column with the given qualifier and name, @@ -543,17 +593,21 @@ impl DFSchema { &self, name: &str, ) -> Vec<(Option<&TableReference>, &FieldRef)> { - // Fields are looked up by name far more often than schemas are built, - // so go through the name index rather than scanning every field. - self.name_index() - .get(name) - .map(|indices| { - indices - .iter() - .map(|&idx| self.qualified_field(idx)) - .collect() - }) - .unwrap_or_default() + match self.name_index() { + Some(index) => index + .get(name) + .map(|indices| { + indices + .iter() + .map(|&idx| self.qualified_field(idx)) + .collect() + }) + .unwrap_or_default(), + None => self + .iter() + .filter(|(_, field)| field.name() == name) + .collect(), + } } /// Find all fields that match the given name and convert to column @@ -909,6 +963,7 @@ impl DFSchema { inner: self.inner, functional_dependencies: self.functional_dependencies, name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), } } @@ -920,6 +975,7 @@ impl DFSchema { inner: self.inner, functional_dependencies: self.functional_dependencies, name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), } } @@ -1192,6 +1248,7 @@ impl TryFrom for DFSchema { field_qualifiers: vec![None; field_count], functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }; // Without checking names, because schema here may have duplicate field names. // For example, Partial AggregateMode will generate duplicate field names from @@ -1254,6 +1311,7 @@ impl ToDFSchema for Vec { field_qualifiers: vec![None; field_count], functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }; Ok(dfschema) } @@ -1545,6 +1603,54 @@ mod tests { } } + /// The scanning path and the indexed path must agree. A schema only + /// switches over after `NAME_INDEX_PROBE_THRESHOLD` probes, so both are + /// live in production and a divergence would be data dependent. + #[test] + fn name_index_and_scan_paths_agree() { + let t1 = TableReference::bare("t1"); + let t2 = TableReference::partial("s", "t2"); + let probes: Vec<(Option<&TableReference>, &str)> = vec![ + (None, "a"), + (Some(&t1), "a"), + (Some(&t2), "a"), + (Some(&t1), "missing"), + (None, "c"), + (Some(&t2), "d"), + (None, "\u{1f600}"), + ]; + + // A fresh schema per probe never reaches the threshold, so each of + // these is answered by the scan. + let scanned: Vec> = probes + .iter() + .map(|(q, name)| ambiguous_schema().index_of_column_by_name(*q, name)) + .collect(); + + // One schema probed past the threshold answers from the index. + let indexed_schema = ambiguous_schema(); + for _ in 0..=NAME_INDEX_PROBE_THRESHOLD { + let _ = indexed_schema.index_of_column_by_name(None, "a"); + } + assert!( + indexed_schema.name_index.get().is_some(), + "schema should have switched to the index by now" + ); + let indexed: Vec> = probes + .iter() + .map(|(q, name)| indexed_schema.index_of_column_by_name(*q, name)) + .collect(); + + assert_eq!(scanned, indexed); + for ((q, name), got) in probes.iter().zip(&scanned) { + assert_eq!( + *got, + reference_index_of_column_by_name(&ambiguous_schema(), *q, name), + "both paths disagree with the reference for {q:?} {name}" + ); + } + } + /// `merge` mutates a schema in place, so a cache populated beforehand /// must not survive it: the merged-in fields have to be visible to every /// later lookup. @@ -1918,6 +2024,7 @@ mod tests { field_qualifiers: vec![None; arrow_schema_ref.fields.len()], functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }; let df_schema_ref = Arc::new(df_schema.clone()); @@ -1965,6 +2072,7 @@ mod tests { field_qualifiers: vec![None; schema.fields.len()], functional_dependencies: FunctionalDependencies::empty(), name_index: OnceLock::new(), + name_probes: AtomicUsize::new(0), }; assert_eq!(df_schema.inner.metadata(), schema.metadata())