fix: remove a few more dropped metadata points

influxdata · Sep 28, 2024 · c153b38 · c153b38
1 parent 6447c7e
commit c153b38
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 6 deletions.
diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs
@@ -674,13 +674,14 @@ impl ParquetSink {
                 .iter()
                 .map(|(s, _)| s)
                 .collect();
-            Arc::new(Schema::new(
+            Arc::new(Schema::new_with_metadata(
                 schema
                     .fields()
                     .iter()
                     .filter(|f| !partition_names.contains(&f.name()))
                     .map(|f| (**f).clone())
                     .collect::<Vec<_>>(),
+                schema.metadata.clone(),
             ))
         } else {
             self.config.output_schema().clone()

diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs
@@ -431,7 +431,10 @@ pub fn expand_qualified_wildcard(
         return plan_err!("Invalid qualifier {qualifier}");
     }
 
-    let qualified_schema = Arc::new(Schema::new(fields_with_qualified));
+    let qualified_schema = Arc::new(Schema::new_with_metadata(
+        fields_with_qualified,
+        schema.metadata().clone(),
+    ));
     let qualified_dfschema =
         DFSchema::try_from_qualified_schema(qualifier.clone(), &qualified_schema)?
             .with_functional_dependencies(projected_func_dependencies)?;

diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -822,12 +822,18 @@ fn create_schema(
         }
     }
 
-    Ok(Schema::new(fields))
+    Ok(Schema::new_with_metadata(
+        fields,
+        input_schema.metadata.clone(),
+    ))
 }
 
 fn group_schema(schema: &Schema, group_count: usize) -> SchemaRef {
     let group_fields = schema.fields()[0..group_count].to_vec();
-    Arc::new(Schema::new(group_fields))
+    Arc::new(Schema::new_with_metadata(
+        group_fields,
+        schema.metadata.clone(),
+    ))
 }
 
 /// Determines the lexical ordering requirement for an aggregate expression.

diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs
@@ -79,14 +79,20 @@ pub fn schema_add_window_field(
         .collect_vec();
     // Skip extending schema for UDAF
     if let WindowFunctionDefinition::AggregateUDF(_) = window_fn {
-        Ok(Arc::new(Schema::new(window_fields)))
+        Ok(Arc::new(Schema::new_with_metadata(
+            window_fields,
+            schema.metadata.clone(),
+        )))
     } else {
         window_fields.extend_from_slice(&[Field::new(
             fn_name,
             window_expr_return_type,
             false,
         )]);
-        Ok(Arc::new(Schema::new(window_fields)))
+        Ok(Arc::new(Schema::new_with_metadata(
+            window_fields,
+            schema.metadata.clone(),
+        )))
     }
 }