Draft Column refactor

Author: Evgeny Maruschenko

datafusion/common/src/column.rs

@@ -27,11 +27,14 @@ use std::sync::Arc;
 
 /// A named reference to a qualified field in a schema.
 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
-pub struct Column {
-    /// relation/table reference.
-    pub relation: Option<OwnedTableReference>,
-    /// field/column name.
-    pub name: String,
+pub enum Column {
+    Unqualified {
+        name: String,
+    },
+    Qualified {
+        relation: OwnedTableReference,
+        name: String,
+    },
 }
 
 impl Column {
@@ -45,26 +48,23 @@ impl Column {
         relation: Option<impl Into<OwnedTableReference>>,
         name: impl Into<String>,
     ) -> Self {
-        Self {
-            relation: relation.map(|r| r.into()),
-            name: name.into(),
+        match relation {
+            Some(relation) => Self::Qualified {
+                relation: relation.into(),
+                name: name.into(),
+            },
+            None => Self::Unqualified { name: name.into() },
         }
     }
 
     /// Convenience method for when there is no qualifier
     pub fn new_unqualified(name: impl Into<String>) -> Self {
-        Self {
-            relation: None,
-            name: name.into(),
-        }
+        Self::Unqualified { name: name.into() }
     }
 
     /// Create Column from unqualified name.
     pub fn from_name(name: impl Into<String>) -> Self {
-        Self {
-            relation: None,
-            name: name.into(),
-        }
+        Self::Unqualified { name: name.into() }
     }
 
     fn from_idents(idents: &mut Vec<String>) -> Option<Self> {
@@ -95,52 +95,81 @@ impl Column {
             // identifiers will be treated as an unqualified column name
             _ => return None,
         };
-        Some(Self { relation, name })
+        match relation {
+            Some(relation) => Some(Self::Qualified { relation, name }),
+            None => Some(Self::Unqualified { name }),
+        }
     }
 
     /// Deserialize a fully qualified name string into a column
     ///
     /// Treats the name as a SQL identifier. For example
     /// `foo.BAR` would be parsed to a reference to relation `foo`, column name `bar` (lower case)
     /// where `"foo.BAR"` would be parsed to a reference to column named `foo.BAR`
-    pub fn from_qualified_name(flat_name: impl Into<String>) -> Self {
-        let flat_name: &str = &flat_name.into();
-        Self::from_idents(&mut parse_identifiers_normalized(flat_name, false))
-            .unwrap_or_else(|| Self {
-                relation: None,
-                name: flat_name.to_owned(),
-            })
+    pub fn from_qualified_name(qualified_name: impl Into<String>) -> Self {
+        let qualified_name: &str = &qualified_name.into();
+        Self::from_idents(&mut parse_identifiers_normalized(qualified_name, false))
+            .unwrap_or_else(|| Self::new_unqualified(qualified_name.to_owned()))
     }
 
     /// Deserialize a fully qualified name string into a column preserving column text case
-    pub fn from_qualified_name_ignore_case(flat_name: impl Into<String>) -> Self {
-        let flat_name: &str = &flat_name.into();
-        Self::from_idents(&mut parse_identifiers_normalized(flat_name, true))
-            .unwrap_or_else(|| Self {
-                relation: None,
-                name: flat_name.to_owned(),
-            })
+    pub fn from_qualified_name_ignore_case(qualified_name: impl Into<String>) -> Self {
+        let qualified_name: &str = &qualified_name.into();
+        Self::from_idents(&mut parse_identifiers_normalized(qualified_name, true))
+            .unwrap_or_else(|| Self::new_unqualified(qualified_name.to_owned()))
+    }
+
+    pub fn relation(&self) -> Option<&OwnedTableReference> {
+        match self {
+            Self::Unqualified { .. } => None,
+            Self::Qualified { relation, .. } => Some(relation),
+        }
+    }
+
+    pub fn unqualified_column(&self) -> Self {
+        Self::Unqualified {
+            name: self.unqualified_name(),
+        }
+    }
+
+    pub fn with_new_qualifier(
+        &self,
+        qualifier: Option<impl Into<OwnedTableReference>>,
+    ) -> Self {
+        match self {
+            Self::Unqualified { name } | Self::Qualified { name, .. } => {
+                Self::new(qualifier, name)
+            }
+        }
+    }
+
+    /// Unqualified name string
+    pub fn unqualified_name(&self) -> String {
+        match self {
+            Self::Unqualified { name } | Self::Qualified { name, .. } => name.clone(),
+        }
+    }
+
+    /// Serialize column into a qualified name string
+    pub fn qualified_name(&self) -> String {
+        match self {
+            Self::Unqualified { name } => name.clone(),
+            Self::Qualified { relation, name } => format!("{}.{}", relation, name),
+        }
     }
 
     /// Serialize column into a flat name string
     pub fn flat_name(&self) -> String {
-        match &self.relation {
-            Some(r) => format!("{}.{}", r, self.name),
-            None => self.name.clone(),
-        }
+        self.qualified_name()
     }
 
     /// Serialize column into a quoted flat name string
     pub fn quoted_flat_name(&self) -> String {
-        match &self.relation {
-            Some(r) => {
-                format!(
-                    "{}.{}",
-                    r.to_quoted_string(),
-                    quote_identifier(self.name.as_str())
-                )
+        match self {
+            Self::Unqualified { name } => quote_identifier(name).to_string(),
+            Self::Qualified { relation, name } => {
+                format!("{}.{}", relation.to_quoted_string(), quote_identifier(name))
             }
-            None => quote_identifier(&self.name).to_string(),
         }
     }
 
@@ -172,12 +201,12 @@ impl Column {
         schemas: &[&Arc<DFSchema>],
         using_columns: &[HashSet<Column>],
     ) -> Result<Self> {
-        if self.relation.is_some() {
+        if let Self::Qualified { .. } = self {
             return Ok(self);
         }
 
         for schema in schemas {
-            let fields = schema.fields_with_unqualified_name(&self.name);
+            let fields = schema.fields_with_unqualified_name(&self.unqualified_name());
             match fields.len() {
                 0 => continue,
                 1 => {
@@ -212,7 +241,7 @@ impl Column {
         }
 
         Err(DataFusionError::SchemaError(SchemaError::FieldNotFound {
-            field: Box::new(Column::new(self.relation.clone(), self.name)),
+            field: Box::new(self.clone()),
             valid_fields: schemas
                 .iter()
                 .flat_map(|s| s.fields().iter().map(|f| f.qualified_column()))
@@ -261,14 +290,14 @@ impl Column {
         schemas: &[&[&DFSchema]],
         using_columns: &[HashSet<Column>],
     ) -> Result<Self> {
-        if self.relation.is_some() {
+        if let Self::Qualified { .. } = self {
             return Ok(self);
         }
 
         for schema_level in schemas {
             let fields = schema_level
                 .iter()
-                .flat_map(|s| s.fields_with_unqualified_name(&self.name))
+                .flat_map(|schema| schema.fields_with_unqualified_name(&self.flat_name()))
                 .collect::<Vec<_>>();
             match fields.len() {
                 0 => continue,
@@ -301,7 +330,7 @@ impl Column {
                     // If not due to USING columns then due to ambiguous column name
                     return Err(DataFusionError::SchemaError(
                         SchemaError::AmbiguousReference {
-                            field: Column::new_unqualified(self.name),
+                            field: Column::new_unqualified(self.flat_name()),
                         },
                     ));
                 }

datafusion/common/src/dfschema.rs

@@ -99,10 +99,7 @@ impl DFSchema {
             if unqualified_names.contains(name) {
                 return Err(DataFusionError::SchemaError(
                     SchemaError::AmbiguousReference {
-                        field: Column {
-                            relation: Some((*qualifier).clone()),
-                            name: name.to_string(),
-                        },
+                        field: Column::new(Some((*qualifier).clone()), name.to_string()),
                     },
                 ));
             }
@@ -112,7 +109,7 @@ impl DFSchema {
         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,
+                None => Column::from_qualified_name(field.name()).unqualified_name(),
             };
             field.qualified_name();
             fields_index
@@ -183,7 +180,7 @@ impl DFSchema {
                 let idx = self.fields.len() - 1;
                 let key = match field.qualifier {
                     Some(_) => field.name().to_owned(),
-                    None => Column::from_qualified_name(field.name()).name,
+                    None => Column::from_qualified_name(field.name()).unqualified_name(),
                 };
                 self.fields_index
                     .entry(key)
@@ -238,7 +235,7 @@ impl DFSchema {
     ) -> Result<Option<usize>> {
         let key = match qualifier {
             Some(_) => name.to_owned(),
-            None => Column::from_qualified_name(name).name,
+            None => Column::from_qualified_name(name).unqualified_name(),
         };
         let Some(matched_fields_idx) = self.fields_index.get(&key) else {
             return Ok(None);
@@ -271,8 +268,8 @@ impl DFSchema {
                     // original qualified name
                     let column = Column::from_qualified_name(field.name());
                     match column {
-                        Column {
-                            relation: Some(column_qualifier),
+                        Column::Qualified {
+                            relation: column_qualifier,
                             name: column_name,
                         } => {
                             column_qualifier.resolved_eq(qualifier) && column_name == name
@@ -286,14 +283,20 @@ impl DFSchema {
     }
 
     /// Find the index of the column with the given qualifier and name
-    pub fn index_of_column(&self, col: &Column) -> Result<usize> {
-        self.index_of_column_by_name(col.relation.as_ref(), &col.name)?
-            .ok_or_else(|| field_not_found(col.relation.clone(), &col.name, self))
+    pub fn index_of_column(&self, column: &Column) -> Result<usize> {
+        self.index_of_column_by_name(column.relation(), &column.unqualified_name())?
+            .ok_or_else(|| {
+                field_not_found(
+                    column.relation().map(|q| (*q).clone()),
+                    &column.unqualified_name(),
+                    self,
+                )
+            })
     }
 
     /// Check if the column is in the current schema
-    pub fn is_column_from_schema(&self, col: &Column) -> Result<bool> {
-        self.index_of_column_by_name(col.relation.as_ref(), &col.name)
+    pub fn is_column_from_schema(&self, column: &Column) -> Result<bool> {
+        self.index_of_column_by_name(column.relation(), &column.unqualified_name())
             .map(|idx| idx.is_some())
     }
 
@@ -349,8 +352,7 @@ impl DFSchema {
                 } else {
                     Err(DataFusionError::SchemaError(
                         SchemaError::AmbiguousReference {
-                            field: Column {
-                                relation: None,
+                            field: Column::Unqualified {
                                 name: name.to_string(),
                             },
                         },
@@ -375,9 +377,12 @@ impl DFSchema {
 
     /// Find the field with the given qualified column
     pub fn field_from_column(&self, column: &Column) -> Result<&DFField> {
-        match &column.relation {
-            Some(r) => self.field_with_qualified_name(r, &column.name),
-            None => self.field_with_unqualified_name(&column.name),
+        match column {
+            Column::Qualified {
+                relation: qualifier,
+                name,
+            } => self.field_with_qualified_name(qualifier, name),
+            Column::Unqualified { name } => self.field_with_unqualified_name(name),
         }
     }
 
@@ -400,9 +405,12 @@ impl DFSchema {
 
     /// Find if the field exists with the given qualified column
     pub fn has_column(&self, column: &Column) -> bool {
-        match &column.relation {
-            Some(r) => self.has_column_with_qualified_name(r, &column.name),
-            None => self.has_column_with_unqualified_name(&column.name),
+        match column {
+            Column::Qualified {
+                relation: qualifier,
+                name,
+            } => self.has_column_with_qualified_name(qualifier, name),
+            Column::Unqualified { name } => self.has_column_with_unqualified_name(name),
         }
     }
 
@@ -757,18 +765,12 @@ impl DFField {
 
     /// Builds a qualified column based on self
     pub fn qualified_column(&self) -> Column {
-        Column {
-            relation: self.qualifier.clone(),
-            name: self.field.name().to_string(),
-        }
+        Column::new(self.qualifier().cloned(), self.field.name())
     }
 
     /// Builds an unqualified column based on self
     pub fn unqualified_column(&self) -> Column {
-        Column {
-            relation: None,
-            name: self.field.name().to_string(),
-        }
+        Column::new_unqualified(self.field.name())
     }
 
     /// Get the optional qualifier

datafusion/common/src/error.rs

@@ -188,12 +188,12 @@ impl Display for SchemaError {
                 )
             }
             Self::AmbiguousReference { field } => {
-                if field.relation.is_some() {
+                if field.relation().is_some() {
                     write!(
                         f,
                         "Schema contains qualified field name {} and unqualified field name {} which would be ambiguous",
                         field.quoted_flat_name(),
-                        quote_identifier(&field.name)
+                        quote_identifier(&field.unqualified_name())
                     )
                 } else {
                     write!(

datafusion/core/src/datasource/listing/helpers.rs

@@ -55,7 +55,8 @@ pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
     let mut is_applicable = true;
     expr.apply(&mut |expr| {
         Ok(match expr {
-            Expr::Column(Column { ref name, .. }) => {
+            Expr::Column(Column::Unqualified { ref name })
+            | Expr::Column(Column::Qualified { ref name, .. }) => {
                 is_applicable &= col_names.contains(name);
                 if is_applicable {
                     VisitRecursion::Skip

datafusion/core/src/datasource/listing/table.rs

@@ -684,7 +684,7 @@ impl ListingTable {
                 .map(|expr| {
                     if let Expr::Sort(Sort { expr, asc, nulls_first }) = expr {
                         if let Expr::Column(col) = expr.as_ref() {
-                            let expr = physical_plan::expressions::col(&col.name, self.table_schema.as_ref())?;
+                            let expr = physical_plan::expressions::col(&col.unqualified_name(), self.table_schema.as_ref())?;
                             Ok(PhysicalSortExpr {
                                 expr,
                                 options: SortOptions {