You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2022/03/10 14:35:58 UTC
[arrow-datafusion] branch master updated: add metadata to DFSchema, close #1806. (#1914)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/master by this push:
new 43b7d49 add metadata to DFSchema, close #1806. (#1914)
43b7d49 is described below
commit 43b7d49d021f176a18de3373388181ddae1c42fe
Author: Jiacai Liu <de...@liujiacai.net>
AuthorDate: Thu Mar 10 22:35:50 2022 +0800
add metadata to DFSchema, close #1806. (#1914)
* add metadata to DFSchema, close #1806.
* add docs
* fix CI
* fix test
* fix clippy
* fix ballista clippy
---
datafusion-common/src/dfschema.rs | 87 +++++++++++++++++-----
datafusion-proto/src/from_proto.rs | 1 +
datafusion/src/logical_plan/builder.rs | 36 ++++++---
datafusion/src/logical_plan/expr_rewriter.rs | 35 +++++----
.../src/optimizer/common_subexpr_eliminate.rs | 2 +-
datafusion/src/optimizer/projection_push_down.rs | 58 ++++++++++-----
datafusion/src/optimizer/simplify_expressions.rs | 16 ++--
.../src/optimizer/single_distinct_to_groupby.rs | 9 ++-
datafusion/src/physical_plan/planner.rs | 10 ++-
datafusion/src/sql/planner.rs | 4 +-
10 files changed, 185 insertions(+), 73 deletions(-)
diff --git a/datafusion-common/src/dfschema.rs b/datafusion-common/src/dfschema.rs
index f6a2b07..6a3dcb0 100644
--- a/datafusion-common/src/dfschema.rs
+++ b/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::HashSet;
+use std::collections::{HashMap, HashSet};
use std::convert::TryFrom;
use std::sync::Arc;
@@ -36,16 +36,30 @@ pub type DFSchemaRef = Arc<DFSchema>;
pub struct DFSchema {
/// Fields
fields: Vec<DFField>,
+ /// Additional metadata in form of key value pairs
+ metadata: HashMap<String, String>,
}
impl DFSchema {
/// Creates an empty `DFSchema`
pub fn empty() -> Self {
- Self { fields: vec![] }
+ Self {
+ fields: vec![],
+ metadata: HashMap::new(),
+ }
}
+ #[deprecated(since = "7.0.0", note = "please use `new_with_metadata` instead")]
/// Create a new `DFSchema`
pub fn new(fields: Vec<DFField>) -> Result<Self> {
+ Self::new_with_metadata(fields, HashMap::new())
+ }
+
+ /// Create a new `DFSchema`
+ pub fn new_with_metadata(
+ fields: Vec<DFField>,
+ metadata: HashMap<String, String>,
+ ) -> Result<Self> {
let mut qualified_names = HashSet::new();
let mut unqualified_names = HashSet::new();
@@ -86,25 +100,28 @@ impl DFSchema {
)));
}
}
- Ok(Self { fields })
+ Ok(Self { fields, metadata })
}
/// Create a `DFSchema` from an Arrow schema
pub fn try_from_qualified_schema(qualifier: &str, schema: &Schema) -> Result<Self> {
- Self::new(
+ Self::new_with_metadata(
schema
.fields()
.iter()
.map(|f| DFField::from_qualified(qualifier, f.clone()))
.collect(),
+ schema.metadata().clone(),
)
}
/// Combine two schemas
pub fn join(&self, schema: &DFSchema) -> Result<Self> {
let mut fields = self.fields.clone();
+ let mut metadata = self.metadata.clone();
fields.extend_from_slice(schema.fields().as_slice());
- Self::new(fields)
+ metadata.extend(schema.metadata.clone());
+ Self::new_with_metadata(fields, metadata)
}
/// Merge a schema into self
@@ -120,6 +137,7 @@ impl DFSchema {
self.fields.push(field.clone());
}
}
+ self.metadata.extend(other_schema.metadata.clone())
}
/// Get a list of fields
@@ -263,6 +281,7 @@ impl DFSchema {
.into_iter()
.map(|f| f.strip_qualifier())
.collect(),
+ ..self
}
}
@@ -281,6 +300,7 @@ impl DFSchema {
)
})
.collect(),
+ ..self
}
}
@@ -295,12 +315,17 @@ impl DFSchema {
.collect::<Vec<_>>()
.join(", ")
}
+
+ /// Get metadata of this schema
+ pub fn metadata(&self) -> &HashMap<String, String> {
+ &self.metadata
+ }
}
impl From<DFSchema> for Schema {
/// Convert DFSchema into a Schema
fn from(df_schema: DFSchema) -> Self {
- Schema::new(
+ Schema::new_with_metadata(
df_schema
.fields
.into_iter()
@@ -316,6 +341,7 @@ impl From<DFSchema> for Schema {
}
})
.collect(),
+ df_schema.metadata,
)
}
}
@@ -323,7 +349,10 @@ impl From<DFSchema> for Schema {
impl From<&DFSchema> for Schema {
/// Convert DFSchema reference into a Schema
fn from(df_schema: &DFSchema) -> Self {
- Schema::new(df_schema.fields.iter().map(|f| f.field.clone()).collect())
+ Schema::new_with_metadata(
+ df_schema.fields.iter().map(|f| f.field.clone()).collect(),
+ df_schema.metadata.clone(),
+ )
}
}
@@ -331,12 +360,13 @@ impl From<&DFSchema> for Schema {
impl TryFrom<Schema> for DFSchema {
type Error = DataFusionError;
fn try_from(schema: Schema) -> std::result::Result<Self, Self::Error> {
- Self::new(
+ Self::new_with_metadata(
schema
.fields()
.iter()
.map(|f| DFField::from(f.clone()))
.collect(),
+ schema.metadata().clone(),
)
}
}
@@ -384,7 +414,7 @@ impl ToDFSchema for SchemaRef {
impl ToDFSchema for Vec<DFField> {
fn to_dfschema(self) -> Result<DFSchema> {
- DFSchema::new(self)
+ DFSchema::new_with_metadata(self, HashMap::new())
}
}
@@ -392,12 +422,13 @@ impl Display for DFSchema {
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
write!(
f,
- "{}",
+ "fields:[{}], metadata:{:?}",
self.fields
.iter()
.map(|field| field.qualified_name())
.collect::<Vec<String>>()
- .join(", ")
+ .join(", "),
+ self.metadata
)
}
}
@@ -556,14 +587,14 @@ mod tests {
#[test]
fn from_unqualified_schema() -> Result<()> {
let schema = DFSchema::try_from(test_schema_1())?;
- assert_eq!("c0, c1", schema.to_string());
+ assert_eq!("fields:[c0, c1], metadata:{}", schema.to_string());
Ok(())
}
#[test]
fn from_qualified_schema() -> Result<()> {
let schema = DFSchema::try_from_qualified_schema("t1", &test_schema_1())?;
- assert_eq!("t1.c0, t1.c1", schema.to_string());
+ assert_eq!("fields:[t1.c0, t1.c1], metadata:{}", schema.to_string());
Ok(())
}
@@ -582,7 +613,10 @@ mod tests {
let left = DFSchema::try_from_qualified_schema("t1", &test_schema_1())?;
let right = DFSchema::try_from_qualified_schema("t2", &test_schema_1())?;
let join = left.join(&right)?;
- assert_eq!("t1.c0, t1.c1, t2.c0, t2.c1", join.to_string());
+ assert_eq!(
+ "fields:[t1.c0, t1.c1, t2.c0, t2.c1], metadata:{}",
+ join.to_string()
+ );
// test valid access
assert!(join.field_with_qualified_name("t1", "c0").is_ok());
assert!(join.field_with_qualified_name("t2", "c0").is_ok());
@@ -626,7 +660,10 @@ mod tests {
let left = DFSchema::try_from_qualified_schema("t1", &test_schema_1())?;
let right = DFSchema::try_from(test_schema_2())?;
let join = left.join(&right)?;
- assert_eq!("t1.c0, t1.c1, c100, c101", join.to_string());
+ assert_eq!(
+ "fields:[t1.c0, t1.c1, c100, c101], metadata:{}",
+ join.to_string()
+ );
// test valid access
assert!(join.field_with_qualified_name("t1", "c0").is_ok());
assert!(join.field_with_unqualified_name("c0").is_ok());
@@ -678,11 +715,18 @@ mod tests {
#[test]
fn into() {
// Demonstrate how to convert back and forth between Schema, SchemaRef, DFSchema, and DFSchemaRef
- let arrow_schema = Schema::new(vec![Field::new("c0", DataType::Int64, true)]);
+ let metadata = test_metadata();
+ let arrow_schema = Schema::new_with_metadata(
+ vec![Field::new("c0", DataType::Int64, true)],
+ metadata.clone(),
+ );
let arrow_schema_ref = Arc::new(arrow_schema.clone());
- let df_schema =
- DFSchema::new(vec![DFField::new(None, "c0", DataType::Int64, true)]).unwrap();
+ let df_schema = DFSchema::new_with_metadata(
+ vec![DFField::new(None, "c0", DataType::Int64, true)],
+ metadata,
+ )
+ .unwrap();
let df_schema_ref = Arc::new(df_schema.clone());
{
@@ -719,4 +763,11 @@ mod tests {
Field::new("c101", DataType::Boolean, true),
])
}
+
+ fn test_metadata() -> HashMap<String, String> {
+ vec![("k1", "v1"), ("k2", "v2")]
+ .into_iter()
+ .map(|(k, v)| (k.to_string(), v.to_string()))
+ .collect::<HashMap<_, _>>()
+ }
}
diff --git a/datafusion-proto/src/from_proto.rs b/datafusion-proto/src/from_proto.rs
index d589ef3..2bf65d4 100644
--- a/datafusion-proto/src/from_proto.rs
+++ b/datafusion-proto/src/from_proto.rs
@@ -155,6 +155,7 @@ impl TryFrom<&protobuf::DfSchema> for DFSchema {
.iter()
.map(|c| c.try_into())
.collect::<Result<Vec<DFField>, _>>()?;
+ #[allow(deprecated)]
Ok(DFSchema::new(fields)?)
}
}
diff --git a/datafusion/src/logical_plan/builder.rs b/datafusion/src/logical_plan/builder.rs
index 086d331..d0bfb5c 100644
--- a/datafusion/src/logical_plan/builder.rs
+++ b/datafusion/src/logical_plan/builder.rs
@@ -188,7 +188,8 @@ impl LogicalPlanBuilder {
for (i, j) in nulls {
values[i][j] = Expr::Literal(ScalarValue::try_from(fields[j].data_type())?);
}
- let schema = DFSchemaRef::new(DFSchema::new(fields)?);
+ let schema =
+ DFSchemaRef::new(DFSchema::new_with_metadata(fields, HashMap::new())?);
Ok(Self::from(LogicalPlan::Values(Values { schema, values })))
}
@@ -392,12 +393,13 @@ impl LogicalPlanBuilder {
let projected_schema = projection
.as_ref()
.map(|p| {
- DFSchema::new(
+ DFSchema::new_with_metadata(
p.iter()
.map(|i| {
DFField::from_qualified(&table_name, schema.field(*i).clone())
})
.collect(),
+ schema.metadata().clone(),
)
})
.unwrap_or_else(|| {
@@ -500,7 +502,10 @@ impl LogicalPlanBuilder {
expr.extend(missing_exprs);
- let new_schema = DFSchema::new(exprlist_to_fields(&expr, input_schema)?)?;
+ let new_schema = DFSchema::new_with_metadata(
+ exprlist_to_fields(&expr, input_schema)?,
+ input_schema.metadata().clone(),
+ )?;
Ok(LogicalPlan::Projection(Projection {
expr,
@@ -569,7 +574,10 @@ impl LogicalPlanBuilder {
.iter()
.map(|f| Expr::Column(f.qualified_column()))
.collect();
- let new_schema = DFSchema::new(exprlist_to_fields(&new_expr, schema)?)?;
+ let new_schema = DFSchema::new_with_metadata(
+ exprlist_to_fields(&new_expr, schema)?,
+ schema.metadata().clone(),
+ )?;
Ok(Self::from(LogicalPlan::Projection(Projection {
expr: new_expr,
@@ -787,7 +795,10 @@ impl LogicalPlanBuilder {
Ok(Self::from(LogicalPlan::Window(Window {
input: Arc::new(self.plan.clone()),
window_expr,
- schema: Arc::new(DFSchema::new(window_fields)?),
+ schema: Arc::new(DFSchema::new_with_metadata(
+ window_fields,
+ self.plan.schema().metadata().clone(),
+ )?),
})))
}
@@ -803,8 +814,10 @@ impl LogicalPlanBuilder {
let aggr_expr = normalize_cols(aggr_expr, &self.plan)?;
let all_expr = group_expr.iter().chain(aggr_expr.iter());
validate_unique_names("Aggregations", all_expr.clone(), self.plan.schema())?;
- let aggr_schema =
- DFSchema::new(exprlist_to_fields(all_expr, self.plan.schema())?)?;
+ let aggr_schema = DFSchema::new_with_metadata(
+ exprlist_to_fields(all_expr, self.plan.schema())?,
+ self.plan.schema().metadata().clone(),
+ )?;
Ok(Self::from(LogicalPlan::Aggregate(Aggregate {
input: Arc::new(self.plan.clone()),
group_expr,
@@ -927,7 +940,9 @@ pub fn build_join_schema(
}
};
- DFSchema::new(fields)
+ let mut metadata = left.metadata().clone();
+ metadata.extend(right.metadata().clone());
+ DFSchema::new_with_metadata(fields, metadata)
}
/// Errors if one or more expressions have equal names.
@@ -1019,7 +1034,10 @@ pub fn project_with_alias(
}
}
validate_unique_names("Projections", projected_expr.iter(), input_schema)?;
- let input_schema = DFSchema::new(exprlist_to_fields(&projected_expr, input_schema)?)?;
+ let input_schema = DFSchema::new_with_metadata(
+ exprlist_to_fields(&projected_expr, input_schema)?,
+ plan.schema().metadata().clone(),
+ )?;
let schema = match alias {
Some(ref alias) => input_schema.replace_qualifier(alias.as_str()),
None => input_schema,
diff --git a/datafusion/src/logical_plan/expr_rewriter.rs b/datafusion/src/logical_plan/expr_rewriter.rs
index cee7c66..2457939 100644
--- a/datafusion/src/logical_plan/expr_rewriter.rs
+++ b/datafusion/src/logical_plan/expr_rewriter.rs
@@ -503,17 +503,20 @@ mod test {
let expr = col("a") + col("b") + col("c");
// Schemas with some matching and some non matching cols
- let schema_a =
- DFSchema::new(vec![make_field("tableA", "a"), make_field("tableA", "aa")])
- .unwrap();
- let schema_c =
- DFSchema::new(vec![make_field("tableC", "cc"), make_field("tableC", "c")])
- .unwrap();
- let schema_b = DFSchema::new(vec![make_field("tableB", "b")]).unwrap();
+ let schema_a = make_schema_with_empty_metadata(vec![
+ make_field("tableA", "a"),
+ make_field("tableA", "aa"),
+ ]);
+ let schema_c = make_schema_with_empty_metadata(vec![
+ make_field("tableC", "cc"),
+ make_field("tableC", "c"),
+ ]);
+ let schema_b = make_schema_with_empty_metadata(vec![make_field("tableB", "b")]);
// non matching
- let schema_f =
- DFSchema::new(vec![make_field("tableC", "f"), make_field("tableC", "ff")])
- .unwrap();
+ let schema_f = make_schema_with_empty_metadata(vec![
+ make_field("tableC", "f"),
+ make_field("tableC", "ff"),
+ ]);
let schemas = vec![schema_c, schema_f, schema_b, schema_a]
.into_iter()
.map(Arc::new)
@@ -531,9 +534,9 @@ mod test {
fn normalize_cols_priority() {
let expr = col("a") + col("b");
// Schemas with multiple matches for column a, first takes priority
- let schema_a = DFSchema::new(vec![make_field("tableA", "a")]).unwrap();
- let schema_b = DFSchema::new(vec![make_field("tableB", "b")]).unwrap();
- let schema_a2 = DFSchema::new(vec![make_field("tableA2", "a")]).unwrap();
+ let schema_a = make_schema_with_empty_metadata(vec![make_field("tableA", "a")]);
+ let schema_b = make_schema_with_empty_metadata(vec![make_field("tableB", "b")]);
+ let schema_a2 = make_schema_with_empty_metadata(vec![make_field("tableA2", "a")]);
let schemas = vec![schema_a2, schema_b, schema_a]
.into_iter()
.map(Arc::new)
@@ -548,7 +551,7 @@ mod test {
fn normalize_cols_non_exist() {
// test normalizing columns when the name doesn't exist
let expr = col("a") + col("b");
- let schema_a = DFSchema::new(vec![make_field("tableA", "a")]).unwrap();
+ let schema_a = make_schema_with_empty_metadata(vec![make_field("tableA", "a")]);
let schemas = vec![schema_a].into_iter().map(Arc::new).collect::<Vec<_>>();
let schemas = schemas.iter().collect::<Vec<_>>();
@@ -568,6 +571,10 @@ mod test {
assert_eq!(unnormalized_expr, col("a") + col("b"));
}
+ fn make_schema_with_empty_metadata(fields: Vec<DFField>) -> DFSchema {
+ DFSchema::new_with_metadata(fields, HashMap::new()).unwrap()
+ }
+
fn make_field(relation: &str, column: &str) -> DFField {
DFField::new(Some(relation), column, DataType::Int8, false)
}
diff --git a/datafusion/src/optimizer/common_subexpr_eliminate.rs b/datafusion/src/optimizer/common_subexpr_eliminate.rs
index 895bcf4..48b6abe 100644
--- a/datafusion/src/optimizer/common_subexpr_eliminate.rs
+++ b/datafusion/src/optimizer/common_subexpr_eliminate.rs
@@ -277,7 +277,7 @@ fn build_project_plan(
project_exprs.push(Expr::Column(field.qualified_column()));
});
- let mut schema = DFSchema::new(fields)?;
+ let mut schema = DFSchema::new_with_metadata(fields, HashMap::new())?;
schema.merge(input.schema());
Ok(LogicalPlan::Projection(Projection {
diff --git a/datafusion/src/optimizer/projection_push_down.rs b/datafusion/src/optimizer/projection_push_down.rs
index 1a64a44..7debb7a 100644
--- a/datafusion/src/optimizer/projection_push_down.rs
+++ b/datafusion/src/optimizer/projection_push_down.rs
@@ -185,10 +185,13 @@ fn optimize_plan(
// no need for an expression at all
Ok(new_input)
} else {
+ let metadata = new_input.schema().metadata().clone();
Ok(LogicalPlan::Projection(Projection {
expr: new_expr,
input: Arc::new(new_input),
- schema: DFSchemaRef::new(DFSchema::new(new_fields)?),
+ schema: DFSchemaRef::new(DFSchema::new_with_metadata(
+ new_fields, metadata,
+ )?),
alias: alias.clone(),
}))
}
@@ -307,13 +310,14 @@ fn optimize_plan(
}
})?;
- let new_schema = DFSchema::new(
+ let new_schema = DFSchema::new_with_metadata(
schema
.fields()
.iter()
.filter(|x| new_required_columns.contains(&x.qualified_column()))
.cloned()
.collect(),
+ schema.metadata().clone(),
)?;
Ok(LogicalPlan::Aggregate(Aggregate {
@@ -413,13 +417,14 @@ fn optimize_plan(
)
})
.collect::<Result<Vec<_>>>()?;
- let new_schema = DFSchema::new(
+ let new_schema = DFSchema::new_with_metadata(
schema
.fields()
.iter()
.filter(|f| union_required_fields.contains(f.field()))
.cloned()
.collect(),
+ schema.metadata().clone(),
)?;
Ok(LogicalPlan::Union(Union {
inputs: new_inputs,
@@ -467,6 +472,8 @@ fn optimize_plan(
#[cfg(test)]
mod tests {
+ use std::collections::HashMap;
+
use super::*;
use crate::logical_plan::{
col, exprlist_to_fields, lit, max, min, Expr, JoinType, LogicalPlanBuilder,
@@ -605,11 +612,14 @@ mod tests {
let optimized_join = optimized_plan.inputs()[0];
assert_eq!(
**optimized_join.schema(),
- DFSchema::new(vec![
- DFField::new(Some("test"), "a", DataType::UInt32, false),
- DFField::new(Some("test"), "b", DataType::UInt32, false),
- DFField::new(Some("test2"), "c1", DataType::UInt32, false),
- ])?,
+ DFSchema::new_with_metadata(
+ vec![
+ DFField::new(Some("test"), "a", DataType::UInt32, false),
+ DFField::new(Some("test"), "b", DataType::UInt32, false),
+ DFField::new(Some("test2"), "c1", DataType::UInt32, false),
+ ],
+ HashMap::new()
+ )?,
);
Ok(())
@@ -646,11 +656,14 @@ mod tests {
let optimized_join = optimized_plan.inputs()[0];
assert_eq!(
**optimized_join.schema(),
- DFSchema::new(vec![
- DFField::new(Some("test"), "a", DataType::UInt32, false),
- DFField::new(Some("test"), "b", DataType::UInt32, false),
- DFField::new(Some("test2"), "c1", DataType::UInt32, false),
- ])?,
+ DFSchema::new_with_metadata(
+ vec![
+ DFField::new(Some("test"), "a", DataType::UInt32, false),
+ DFField::new(Some("test"), "b", DataType::UInt32, false),
+ DFField::new(Some("test2"), "c1", DataType::UInt32, false),
+ ],
+ HashMap::new()
+ )?,
);
Ok(())
@@ -685,11 +698,14 @@ mod tests {
let optimized_join = optimized_plan.inputs()[0];
assert_eq!(
**optimized_join.schema(),
- DFSchema::new(vec![
- DFField::new(Some("test"), "a", DataType::UInt32, false),
- DFField::new(Some("test"), "b", DataType::UInt32, false),
- DFField::new(Some("test2"), "a", DataType::UInt32, false),
- ])?,
+ DFSchema::new_with_metadata(
+ vec![
+ DFField::new(Some("test"), "a", DataType::UInt32, false),
+ DFField::new(Some("test"), "b", DataType::UInt32, false),
+ DFField::new(Some("test2"), "a", DataType::UInt32, false),
+ ],
+ HashMap::new()
+ )?,
);
Ok(())
@@ -746,7 +762,11 @@ mod tests {
// relation is `None`). PlanBuilder resolves the expressions
let expr = vec![col("a"), col("b")];
let projected_fields = exprlist_to_fields(&expr, input_schema).unwrap();
- let projected_schema = DFSchema::new(projected_fields).unwrap();
+ let projected_schema = DFSchema::new_with_metadata(
+ projected_fields,
+ input_schema.metadata().clone(),
+ )
+ .unwrap();
let plan = LogicalPlan::Projection(Projection {
expr,
input: Arc::new(table_scan),
diff --git a/datafusion/src/optimizer/simplify_expressions.rs b/datafusion/src/optimizer/simplify_expressions.rs
index e1eacb6..521d976 100644
--- a/datafusion/src/optimizer/simplify_expressions.rs
+++ b/datafusion/src/optimizer/simplify_expressions.rs
@@ -729,6 +729,7 @@ impl<'a, S: SimplifyInfo> ExprRewriter for Simplifier<'a, S> {
#[cfg(test)]
mod tests {
+ use std::collections::HashMap;
use std::sync::Arc;
use arrow::array::{ArrayRef, Int32Array};
@@ -1207,12 +1208,15 @@ mod tests {
fn expr_test_schema() -> DFSchemaRef {
Arc::new(
- DFSchema::new(vec![
- DFField::new(None, "c1", DataType::Utf8, true),
- DFField::new(None, "c2", DataType::Boolean, true),
- DFField::new(None, "c1_non_null", DataType::Utf8, false),
- DFField::new(None, "c2_non_null", DataType::Boolean, false),
- ])
+ DFSchema::new_with_metadata(
+ vec![
+ DFField::new(None, "c1", DataType::Utf8, true),
+ DFField::new(None, "c2", DataType::Boolean, true),
+ DFField::new(None, "c1_non_null", DataType::Utf8, false),
+ DFField::new(None, "c2_non_null", DataType::Boolean, false),
+ ],
+ HashMap::new(),
+ )
.unwrap(),
)
}
diff --git a/datafusion/src/optimizer/single_distinct_to_groupby.rs b/datafusion/src/optimizer/single_distinct_to_groupby.rs
index 2e0bd5f..4f01438 100644
--- a/datafusion/src/optimizer/single_distinct_to_groupby.rs
+++ b/datafusion/src/optimizer/single_distinct_to_groupby.rs
@@ -91,7 +91,11 @@ fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
.map(|expr| expr.to_field(input.schema()).unwrap())
.collect::<Vec<_>>();
- let grouped_schema = DFSchema::new(all_field).unwrap();
+ let grouped_schema = DFSchema::new_with_metadata(
+ all_field,
+ input.schema().metadata().clone(),
+ )
+ .unwrap();
let grouped_agg = LogicalPlan::Aggregate(Aggregate {
input: input.clone(),
group_expr: all_group_args,
@@ -100,12 +104,13 @@ fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
});
let grouped_agg = optimize_children(&grouped_agg);
let final_agg_schema = Arc::new(
- DFSchema::new(
+ DFSchema::new_with_metadata(
group_expr
.iter()
.chain(new_aggr_expr.iter())
.map(|expr| expr.to_field(&grouped_schema).unwrap())
.collect::<Vec<_>>(),
+ input.schema().metadata().clone(),
)
.unwrap(),
);
diff --git a/datafusion/src/physical_plan/planner.rs b/datafusion/src/physical_plan/planner.rs
index 88ff9fa..4055b14 100644
--- a/datafusion/src/physical_plan/planner.rs
+++ b/datafusion/src/physical_plan/planner.rs
@@ -1444,6 +1444,7 @@ mod tests {
use datafusion_expr::sum;
use datafusion_expr::{col, lit};
use fmt::Debug;
+ use std::collections::HashMap;
use std::convert::TryFrom;
use std::{any::Any, fmt};
@@ -1628,7 +1629,7 @@ mod tests {
dict_id: 0, \
dict_is_ordered: false, \
metadata: None } }\
- ] }, \
+ ], metadata: {} }, \
ExecutionPlan schema: Schema { fields: [\
Field { \
name: \"b\", \
@@ -1816,8 +1817,11 @@ mod tests {
fn default() -> Self {
Self {
schema: DFSchemaRef::new(
- DFSchema::new(vec![DFField::new(None, "a", DataType::Int32, false)])
- .unwrap(),
+ DFSchema::new_with_metadata(
+ vec![DFField::new(None, "a", DataType::Int32, false)],
+ HashMap::new(),
+ )
+ .unwrap(),
),
}
}
diff --git a/datafusion/src/sql/planner.rs b/datafusion/src/sql/planner.rs
index 0ae05f9..b528e56 100644
--- a/datafusion/src/sql/planner.rs
+++ b/datafusion/src/sql/planner.rs
@@ -712,10 +712,12 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
Some(predicate_expr) => {
// build join schema
let mut fields = vec![];
+ let mut metadata = std::collections::HashMap::new();
for plan in &plans {
fields.extend_from_slice(plan.schema().fields());
+ metadata.extend(plan.schema().metadata().clone());
}
- let join_schema = DFSchema::new(fields)?;
+ let join_schema = DFSchema::new_with_metadata(fields, metadata)?;
let filter_expr = self.sql_to_rex(predicate_expr, &join_schema)?;