From 7c35ca56ffff07fb2c8dadb80f70ec9baa41354c Mon Sep 17 00:00:00 2001 From: AnuragRaut08 Date: Tue, 15 Sep 2026 21:39:44 +0530 Subject: [PATCH] fix: preserve dotted column relation qualifiers --- .../proto/datafusion_common.proto | 1 + datafusion/proto-common/src/from_proto/mod.rs | 125 +++++++++++++++++- .../proto-common/src/generated/pbjson.rs | 19 ++- .../proto-common/src/generated/prost.rs | 2 + datafusion/proto-common/src/to_proto/mod.rs | 2 + .../src/generated/datafusion_proto_common.rs | 2 + 6 files changed, 149 insertions(+), 2 deletions(-) diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 712212f6b6ae5..503f5fddcab5f 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -22,6 +22,7 @@ package datafusion_common; message ColumnRelation { string relation = 1; + repeated string parts = 2; } message Column { diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index 92506ad92bad0..89856487736e0 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -145,7 +145,14 @@ where impl From for TableReference { fn from(rel: protobuf::ColumnRelation) -> Self { - Self::parse_str_normalized(rel.relation.as_str(), true) + match rel.parts.as_slice() { + [table] => Self::bare(table.as_str()), + [schema, table] => Self::partial(schema.as_str(), table.as_str()), + [catalog, schema, table] => { + Self::full(catalog.as_str(), schema.as_str(), table.as_str()) + } + _ => Self::parse_str_normalized(rel.relation.as_str(), true), + } } } @@ -1386,6 +1393,7 @@ pub(crate) fn csv_writer_options_from_proto( #[cfg(test)] mod tests { + use datafusion_common::TableReference; use datafusion_common::config::{ MaxRowGroupBytes, ParquetCdcOptions, ParquetOptions, TableParquetOptions, }; @@ -1405,6 +1413,121 @@ mod tests { ); } + #[test] + fn column_relation_round_trip_preserves_dotted_bare_table() { + let column = datafusion_common::Column::new( + Some(TableReference::bare("has.dot")), + "column", + ); + + let proto: crate::protobuf_common::Column = (&column).into(); + let relation = proto.relation.expect("relation should be present"); + + assert_eq!(relation.relation, "has.dot"); + assert_eq!(relation.parts, vec!["has.dot".to_string()]); + + let recovered = TableReference::from(relation); + assert_eq!(recovered, TableReference::bare("has.dot")); + } + + #[test] + fn column_relation_round_trip_preserves_dotted_partial_reference() { + let column = datafusion_common::Column::new( + Some(TableReference::partial("my.schema", "table")), + "column", + ); + + let proto: crate::protobuf_common::Column = (&column).into(); + let relation = proto.relation.expect("relation should be present"); + + assert_eq!(relation.relation, "my.schema.table"); + assert_eq!( + relation.parts, + vec!["my.schema".to_string(), "table".to_string()] + ); + + let recovered = TableReference::from(relation); + assert_eq!(recovered, TableReference::partial("my.schema", "table")); + } + + #[test] + fn column_relation_round_trip_preserves_dotted_full_reference() { + let column = datafusion_common::Column::new( + Some(TableReference::full("catalog", "my.schema", "table")), + "column", + ); + + let proto: crate::protobuf_common::Column = (&column).into(); + let relation = proto.relation.expect("relation should be present"); + + assert_eq!(relation.relation, "catalog.my.schema.table"); + assert_eq!( + relation.parts, + vec![ + "catalog".to_string(), + "my.schema".to_string(), + "table".to_string() + ] + ); + + let recovered = TableReference::from(relation); + assert_eq!( + recovered, + TableReference::full("catalog", "my.schema", "table") + ); + } + + #[test] + fn df_schema_round_trip_preserves_dotted_qualifiers() { + use std::sync::Arc; + + use datafusion_common::DFSchema; + use datafusion_common::arrow::datatypes::{DataType, Field}; + + let schema = DFSchema::new_with_metadata( + vec![ + ( + Some(TableReference::partial("my.schema", "table")), + Arc::new(Field::new("partial_col", DataType::Utf8, true)), + ), + ( + Some(TableReference::full("catalog", "my.schema", "table")), + Arc::new(Field::new("full_col", DataType::Utf8, true)), + ), + ], + std::collections::HashMap::::new(), + ) + .unwrap(); + + let proto: crate::protobuf_common::DfSchema = (&schema).try_into().unwrap(); + let recovered = DFSchema::try_from(&proto).unwrap(); + + let qualifiers = recovered + .iter() + .map(|(qualifier, _)| qualifier.cloned()) + .collect::>(); + + assert_eq!( + qualifiers, + vec![ + Some(TableReference::partial("my.schema", "table")), + Some(TableReference::full("catalog", "my.schema", "table")), + ] + ); + } + + #[test] + fn column_relation_decodes_legacy_relation() { + let proto = crate::protobuf_common::ColumnRelation { + relation: "schema.table".to_string(), + parts: vec![], + }; + + let recovered = TableReference::from(proto); + + assert_eq!(recovered, TableReference::partial("schema", "table")); + } + #[test] fn table_parquet_options_defaults_missing_global() { let recovered = TableParquetOptions::try_from( diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index 2e94368bfd01e..295caec0de7c6 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -1030,10 +1030,16 @@ impl serde::Serialize for ColumnRelation { if !self.relation.is_empty() { len += 1; } + if !self.parts.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion_common.ColumnRelation", len)?; if !self.relation.is_empty() { struct_ser.serialize_field("relation", &self.relation)?; } + if !self.parts.is_empty() { + struct_ser.serialize_field("parts", &self.parts)?; + } struct_ser.end() } } @@ -1045,11 +1051,13 @@ impl<'de> serde::Deserialize<'de> for ColumnRelation { { const FIELDS: &[&str] = &[ "relation", + "parts", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Relation, + Parts, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -1072,6 +1080,7 @@ impl<'de> serde::Deserialize<'de> for ColumnRelation { { match value { "relation" => Ok(GeneratedField::Relation), + "parts" => Ok(GeneratedField::Parts), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -1092,6 +1101,7 @@ impl<'de> serde::Deserialize<'de> for ColumnRelation { V: serde::de::MapAccess<'de>, { let mut relation__ = None; + let mut parts__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Relation => { @@ -1100,10 +1110,17 @@ impl<'de> serde::Deserialize<'de> for ColumnRelation { } relation__ = Some(map_.next_value()?); } + GeneratedField::Parts => { + if parts__.is_some() { + return Err(serde::de::Error::duplicate_field("parts")); + } + parts__ = Some(map_.next_value()?); + } } } Ok(ColumnRelation { relation: relation__.unwrap_or_default(), + parts: parts__.unwrap_or_default(), }) } } @@ -3997,7 +4014,7 @@ impl serde::Serialize for ExplainAnalyzeCategoriesNode { struct_ser.serialize_field("all", &self.all)?; } if !self.only.is_empty() { - let v = self.only.iter().copied().map(|v| { + let v = self.only.iter().cloned().map(|v| { MetricCategory::try_from(v) .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", v))) }).collect::, _>>()?; diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index c0e79aec6d873..8a2b2969229b1 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -3,6 +3,8 @@ pub struct ColumnRelation { #[prost(string, tag = "1")] pub relation: ::prost::alloc::string::String, + #[prost(string, repeated, tag = "2")] + pub parts: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct Column { diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index 7f95e03f41db4..aad9dfa308a73 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -247,6 +247,7 @@ impl From for protobuf::Column { Self { relation: c.relation.map(|relation| protobuf::ColumnRelation { relation: relation.to_string(), + parts: relation.to_vec(), }), name: c.name, } @@ -292,6 +293,7 @@ impl TryFrom<&DFSchema> for protobuf::DfSchema { field: Some(field.as_ref().try_into()?), qualifier: qualifier.map(|r| protobuf::ColumnRelation { relation: r.to_string(), + parts: r.to_vec(), }), }) }) diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index c0e79aec6d873..8a2b2969229b1 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -3,6 +3,8 @@ pub struct ColumnRelation { #[prost(string, tag = "1")] pub relation: ::prost::alloc::string::String, + #[prost(string, repeated, tag = "2")] + pub parts: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct Column {