Add fields_index for searching column index

Author: Evgeny Maruschenko

datafusion/common/src/dfschema.rs

@@ -18,7 +18,7 @@
 //! DFSchema is an extended schema struct that DataFusion uses to provide support for
 //! fields with optional relation names.
 
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::convert::TryFrom;
 use std::fmt::{Display, Formatter};
 use std::hash::Hash;
@@ -42,6 +42,8 @@ pub type DFSchemaRef = Arc<DFSchema>;
 pub struct DFSchema {
     /// Fields
     fields: Vec<DFField>,
+    // For searching column id by field name
+    fields_index: BTreeMap<String, Vec<usize>>,
     /// Additional metadata in form of key value pairs
     metadata: HashMap<String, String>,
     /// Stores functional dependencies in the schema.
@@ -53,6 +55,7 @@ impl DFSchema {
     pub fn empty() -> Self {
         Self {
             fields: vec![],
+            fields_index: BTreeMap::new(),
             metadata: HashMap::new(),
             functional_dependencies: FunctionalDependencies::empty(),
         }
@@ -104,9 +107,24 @@ impl DFSchema {
                 ));
             }
         }
+
+        let mut fields_index: BTreeMap<String, Vec<usize>> = BTreeMap::new();
+        for (idx, field) in fields.iter().enumerate() {
+            let key = match field.qualifier {
+                Some(_) => field.name().to_owned(),
+                None => Column::from_qualified_name(field.name()).name,
+            };
+            field.qualified_name();
+            fields_index
+                .entry(key)
+                .and_modify(|vec| vec.push(idx))
+                .or_insert(vec![idx]);
+        }
+
         Ok(Self {
             fields,
             metadata,
+            fields_index,
             functional_dependencies: FunctionalDependencies::empty(),
         })
     }
@@ -161,6 +179,16 @@ impl DFSchema {
             };
             if !duplicated_field {
                 self.fields.push(field.clone());
+
+                let idx = self.fields.len() - 1;
+                let key = match field.qualifier {
+                    Some(_) => field.name().to_owned(),
+                    None => Column::from_qualified_name(field.name()).name,
+                };
+                self.fields_index
+                    .entry(key)
+                    .and_modify(|vec| vec.push(idx))
+                    .or_insert(vec![idx]);
             }
         }
         self.metadata.extend(other_schema.metadata.clone())
@@ -208,32 +236,50 @@ impl DFSchema {
         qualifier: Option<&TableReference>,
         name: &str,
     ) -> Result<Option<usize>> {
-        let mut matches = self
-            .fields
+        let key = match qualifier {
+            Some(_) => name.to_owned(),
+            None => Column::from_qualified_name(name).name,
+        };
+        let Some(matched_fields_idx) = self.fields_index.get(&key) else {
+            return Ok(None);
+        };
+
+        let Some(qualifier) = qualifier else {
+            // field to lookup is unqualified, no need to compare qualifier
+            let mut matches = matched_fields_idx
+                .iter()
+                .map(|idx| (*idx, self.field(*idx)))
+                .filter(|(_, field)| field.name() == name)
+                .map(|(idx, _)| idx);
+            return Ok(matches.next());
+        };
+
+        // field to lookup is qualified.
+        let mut matches = matched_fields_idx
             .iter()
-            .enumerate()
-            .filter(|(_, field)| match (qualifier, &field.qualifier) {
+            .map(|idx| (*idx, self.field(*idx)))
+            .filter(|(_, field)| match &field.qualifier {
                 // 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.name() == name
+                Some(field_qualifier) => {
+                    qualifier.resolved_eq(field_qualifier) && field.name() == name
                 }
                 // field to lookup is qualified but current field is unqualified.
-                (Some(qq), None) => {
+                None => {
                     // the original field may now be aliased with a name that matches the
                     // original qualified name
                     let column = Column::from_qualified_name(field.name());
                     match column {
                         Column {
-                            relation: Some(r),
+                            relation: Some(column_qualifier),
                             name: column_name,
-                        } => &r == qq && column_name == name,
+                        } => {
+                            column_qualifier.resolved_eq(qualifier) && column_name == name
+                        }
                         _ => false,
                     }
                 }
-                // field to lookup is unqualified, no need to compare qualifier
-                (None, Some(_)) | (None, None) => field.name() == name,
             })
             .map(|(idx, _)| idx);
         Ok(matches.next())

datafusion/core/src/physical_planner.rs

@@ -2326,7 +2326,7 @@ mod tests {
                 dict_id: 0, \
                 dict_is_ordered: false, \
                 metadata: {} } }\
-        ], metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }, \
+        ], fields_index: {\"a\": [0]}, metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }, \
         ExecutionPlan schema: Schema { fields: [\
             Field { \
                 name: \"b\", \

datafusion/optimizer/src/optimizer.rs

@@ -504,12 +504,12 @@ mod tests {
         assert_eq!(
             "Optimizer rule 'get table_scan rule' failed\ncaused by\nget table_scan rule\ncaused by\n\
              Internal error: Failed due to generate a different schema, \
-             original schema: DFSchema { fields: [], metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }, \
+             original schema: DFSchema { fields: [], fields_index: {}, metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }, \
              new schema: DFSchema { fields: [\
              DFField { qualifier: Some(Bare { table: \"test\" }), field: Field { name: \"a\", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }, \
              DFField { qualifier: Some(Bare { table: \"test\" }), field: Field { name: \"b\", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }, \
              DFField { qualifier: Some(Bare { table: \"test\" }), field: Field { name: \"c\", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }], \
-             metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }.\
+             fields_index: {\"a\": [0], \"b\": [1], \"c\": [2]}, metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }.\
              \nThis was likely caused by a bug in DataFusion's code \
              and we would welcome that you file an bug report in our issue tracker",
             err.strip_backtrace()