Skip to content

Commit d58f579

Browse files
alambclaude
andcommitted
Fix compilation for arrow-rs main (Metadata type, rand 0.10, object_store 0.14)
- Adapt to the new arrow_schema::Metadata type (was HashMap<String, String>): DFSchema::metadata()/ExprSchema::metadata() now return &Metadata, DFSchema constructors take impl Into<Metadata>, SchemaFieldMetadata is now an alias for Metadata, and check_metadata_with_storage_equal/format_type_and_metadata accept Option<&Metadata> - Add FieldMetadata <-> Metadata conversions and a DFHeapSize impl for Metadata - Replace deprecated fb_to_schema with try_fb_to_schema - Rework two benches that mixed DataFusion's rand 0.9 with arrow's rand 0.10 - Add new object_store 0.14 GetResult/ListResult extensions field - Update map field name expectations for new arrow spec defaults (key/value) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 441ec42 commit d58f579

40 files changed

Lines changed: 192 additions & 181 deletions

File tree

datafusion/catalog-listing/src/helpers.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ use datafusion_expr::{BinaryExpr, Operator, lit, utils};
3030

3131
use arrow::{
3232
array::AsArray,
33-
datatypes::{DataType, Field},
33+
datatypes::{DataType, Field, Metadata},
3434
record_batch::RecordBatch,
3535
};
3636
use datafusion_expr::execution_props::ExecutionProps;
@@ -423,7 +423,7 @@ pub async fn pruned_partition_list<'a>(
423423
.iter()
424424
.map(|(n, d)| Field::new(n, d.clone(), true))
425425
.collect(),
426-
Default::default(),
426+
Metadata::new(),
427427
)?;
428428

429429
Ok(objects

datafusion/catalog-listing/src/table.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use crate::helpers::{
2020
expr_applicable_for_cols, filter_partitioned_file, pruned_partition_list,
2121
};
2222
use crate::{ListingOptions, ListingTableConfig};
23-
use arrow::datatypes::{Field, Schema, SchemaBuilder, SchemaRef};
23+
use arrow::datatypes::{Field, Metadata, Schema, SchemaBuilder, SchemaRef};
2424
use async_trait::async_trait;
2525
use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider};
2626
use datafusion_common::stats::Precision;
@@ -1035,7 +1035,7 @@ impl ListingTable {
10351035
.iter()
10361036
.map(|(name, data_type)| Field::new(name, data_type.clone(), true))
10371037
.collect(),
1038-
Default::default(),
1038+
Metadata::new(),
10391039
)?;
10401040

10411041
file_groups

datafusion/common/src/dfschema.rs

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
//! DFSchema is an extended schema struct that DataFusion uses to provide support for
1919
//! fields with optional relation names.
2020
21-
use std::collections::{BTreeSet, HashMap, HashSet};
21+
use std::collections::{BTreeSet, HashSet};
2222
use std::fmt::{Display, Formatter};
2323
use std::hash::Hash;
2424
use std::sync::{Arc, LazyLock};
@@ -31,7 +31,7 @@ use crate::{
3131

3232
use arrow::compute::can_cast_types;
3333
use arrow::datatypes::{
34-
DataType, Field, FieldRef, Fields, Schema, SchemaBuilder, SchemaRef,
34+
DataType, Field, FieldRef, Fields, Metadata, Schema, SchemaBuilder, SchemaRef,
3535
};
3636

3737
/// A reference-counted reference to a [DFSchema].
@@ -153,7 +153,7 @@ impl DFSchema {
153153
/// Create a `DFSchema` from an Arrow schema where all the fields have a given qualifier
154154
pub fn new_with_metadata(
155155
qualified_fields: Vec<(Option<TableReference>, Arc<Field>)>,
156-
metadata: HashMap<String, String>,
156+
metadata: impl Into<Metadata>,
157157
) -> Result<Self> {
158158
let (qualifiers, fields): (Vec<Option<TableReference>>, Vec<Arc<Field>>) =
159159
qualified_fields.into_iter().unzip();
@@ -172,7 +172,7 @@ impl DFSchema {
172172
/// Create a new `DFSchema` from a list of Arrow [Field]s
173173
pub fn from_unqualified_fields(
174174
fields: Fields,
175-
metadata: HashMap<String, String>,
175+
metadata: impl Into<Metadata>,
176176
) -> Result<Self> {
177177
let field_count = fields.len();
178178
let schema = Arc::new(Schema::new_with_metadata(fields, metadata));
@@ -864,7 +864,7 @@ impl DFSchema {
864864
}
865865

866866
/// Get metadata of this schema
867-
pub fn metadata(&self) -> &HashMap<String, String> {
867+
pub fn metadata(&self) -> &Metadata {
868868
&self.inner.metadata
869869
}
870870

@@ -1179,7 +1179,7 @@ impl ToDFSchema for Vec<Field> {
11791179
let field_count = self.len();
11801180
let schema = Schema {
11811181
fields: self.into(),
1182-
metadata: HashMap::new(),
1182+
metadata: Metadata::new(),
11831183
};
11841184
let dfschema = DFSchema {
11851185
inner: schema.into(),
@@ -1221,7 +1221,7 @@ pub trait ExprSchema: std::fmt::Debug {
12211221
}
12221222

12231223
/// Returns the column's optional metadata.
1224-
fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>> {
1224+
fn metadata(&self, col: &Column) -> Result<&Metadata> {
12251225
Ok(self.field_from_column(col)?.metadata())
12261226
}
12271227

@@ -1245,7 +1245,7 @@ impl<P: AsRef<DFSchema> + std::fmt::Debug> ExprSchema for P {
12451245
self.as_ref().data_type(col)
12461246
}
12471247

1248-
fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>> {
1248+
fn metadata(&self, col: &Column) -> Result<&Metadata> {
12491249
ExprSchema::metadata(self.as_ref(), col)
12501250
}
12511251

@@ -1379,6 +1379,7 @@ pub fn qualified_name(qualifier: Option<&TableReference>, name: &str) -> String
13791379
#[cfg(test)]
13801380
mod tests {
13811381
use crate::assert_contains;
1382+
use std::collections::HashMap;
13821383

13831384
use super::*;
13841385

datafusion/common/src/heap_size.rs

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ use arrow::array::{
4747
};
4848
use arrow::datatypes::{
4949
DataType, Field, Fields, IntervalDayTime, IntervalMonthDayNano, IntervalUnit,
50-
TimeUnit, UnionFields, UnionMode, i256,
50+
Metadata, TimeUnit, UnionFields, UnionMode, i256,
5151
};
5252
use chrono::{DateTime, Utc};
5353
use half::f16;
@@ -396,6 +396,19 @@ impl DFHeapSize for UnionFields {
396396
}
397397
}
398398

399+
impl DFHeapSize for Metadata {
400+
fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
401+
// `Metadata` does not expose its underlying reference-counted map, so
402+
// this approximates the `BTreeMap` entries' sizes and cannot dedupe
403+
// instances that share the same allocation.
404+
self.iter()
405+
.map(|(k, v)| {
406+
size_of::<(String, String)>() + k.heap_size(ctx) + v.heap_size(ctx)
407+
})
408+
.sum()
409+
}
410+
}
411+
399412
impl DFHeapSize for Field {
400413
fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
401414
self.name().heap_size(ctx)

datafusion/common/src/metadata.rs

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
use std::{collections::BTreeMap, sync::Arc};
1919

20-
use arrow::datatypes::{DataType, Field, FieldRef};
20+
use arrow::datatypes::{DataType, Field, FieldRef, Metadata};
2121
use hashbrown::HashMap;
2222

2323
use crate::{DataFusionError, ScalarValue, error::_plan_err};
@@ -84,14 +84,8 @@ impl From<ScalarValue> for ScalarAndMetadata {
8484
/// Returns a planning error with suitably formatted type representations if
8585
/// actual and expected do not compare to equal.
8686
pub fn check_metadata_with_storage_equal(
87-
actual: (
88-
&DataType,
89-
Option<&std::collections::HashMap<String, String>>,
90-
),
91-
expected: (
92-
&DataType,
93-
Option<&std::collections::HashMap<String, String>>,
94-
),
87+
actual: (&DataType, Option<&Metadata>),
88+
expected: (&DataType, Option<&Metadata>),
9589
what: &str,
9690
context: &str,
9791
) -> Result<(), DataFusionError> {
@@ -131,7 +125,7 @@ pub fn check_metadata_with_storage_equal(
131125
/// renderings.
132126
pub fn format_type_and_metadata(
133127
data_type: &DataType,
134-
metadata: Option<&std::collections::HashMap<String, String>>,
128+
metadata: Option<&Metadata>,
135129
) -> String {
136130
match metadata {
137131
Some(metadata) if !metadata.is_empty() => {
@@ -316,6 +310,13 @@ impl FieldMetadata {
316310
.collect()
317311
}
318312

313+
/// Convert this `FieldMetadata` into an arrow [`Metadata`]
314+
///
315+
/// This is cheap: both types share the same `Arc<BTreeMap>` representation.
316+
pub fn to_metadata(&self) -> Metadata {
317+
Metadata::from(Arc::clone(&self.inner))
318+
}
319+
319320
/// Updates the metadata on the Field with this metadata, if it is not empty.
320321
pub fn add_to_field(&self, field: Field) -> Field {
321322
if self.inner.is_empty() {
@@ -336,6 +337,24 @@ impl FieldMetadata {
336337
}
337338
}
338339

340+
impl From<&FieldMetadata> for Metadata {
341+
fn from(value: &FieldMetadata) -> Self {
342+
value.to_metadata()
343+
}
344+
}
345+
346+
impl From<Metadata> for FieldMetadata {
347+
fn from(value: Metadata) -> Self {
348+
Self::new(value.into())
349+
}
350+
}
351+
352+
impl From<&Metadata> for FieldMetadata {
353+
fn from(value: &Metadata) -> Self {
354+
Self::from(value.clone())
355+
}
356+
}
357+
339358
impl From<&Field> for FieldMetadata {
340359
fn from(field: &Field) -> Self {
341360
Self::new_from_field(field)

datafusion/common/src/nested_struct.rs

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1123,8 +1123,8 @@ mod tests {
11231123
Arc::new(non_null_field(
11241124
"entries",
11251125
struct_type(vec![
1126-
non_null_field("keys", DataType::Utf8),
1127-
field("values", DataType::Int32),
1126+
non_null_field("key", DataType::Utf8),
1127+
field("value", DataType::Int32),
11281128
]),
11291129
)),
11301130
false,
@@ -1148,8 +1148,8 @@ mod tests {
11481148
Arc::new(non_null_field(
11491149
"entries",
11501150
struct_type(vec![
1151-
non_null_field("keys", DataType::Utf8),
1152-
field("values", DataType::Int32),
1151+
non_null_field("key", DataType::Utf8),
1152+
field("value", DataType::Int32),
11531153
]),
11541154
)),
11551155
false,
@@ -1176,8 +1176,8 @@ mod tests {
11761176
assert!(map.is_null(1));
11771177
let map0 = map.value(0);
11781178
let entries = map0.as_any().downcast_ref::<StructArray>().unwrap();
1179-
let keys = get_column_as!(entries, "keys", StringArray);
1180-
let vals = get_column_as!(entries, "values", Int32Array);
1179+
let keys = get_column_as!(entries, "key", StringArray);
1180+
let vals = get_column_as!(entries, "value", Int32Array);
11811181
assert_eq!(keys.value(0), "a");
11821182
assert_eq!(vals.value(0), 1);
11831183
}

datafusion/common/src/param_value.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ impl ParamValues {
6464
check_metadata_with_storage_equal(
6565
(
6666
&lit.value.data_type(),
67-
lit.metadata.as_ref().map(|m| m.to_hashmap()).as_ref(),
67+
lit.metadata.as_ref().map(|m| m.to_metadata()).as_ref(),
6868
),
6969
(param_type.data_type(), Some(param_type.metadata())),
7070
"parameter",

datafusion/core/src/datasource/file_format/csv.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@ mod tests {
136136
},
137137
range: Default::default(),
138138
attributes: Attributes::default(),
139+
extensions: Default::default(),
139140
})
140141
}
141142

datafusion/core/tests/sql/path_partition.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -718,6 +718,7 @@ impl ObjectStore for MirroringObjectStore {
718718
payload,
719719
meta,
720720
attributes: Attributes::default(),
721+
extensions: Default::default(),
721722
})
722723
}
723724

@@ -789,6 +790,7 @@ impl ObjectStore for MirroringObjectStore {
789790
Ok(ListResult {
790791
common_prefixes: common_prefixes.into_iter().collect(),
791792
objects,
793+
extensions: Default::default(),
792794
})
793795
}
794796

datafusion/core/tests/user_defined/user_defined_aggregates.rs

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ use arrow::array::{
3030
Array, AsArray, Int32Array, PrimitiveArray, StringArray, StructArray, UInt64Array,
3131
record_batch, types::UInt64Type,
3232
};
33-
use arrow::datatypes::{Fields, Schema};
33+
use arrow::datatypes::{Fields, Metadata, Schema};
3434
use arrow_schema::FieldRef;
3535
use datafusion::common::test_util::batches_to_string;
3636
use datafusion::dataframe::DataFrame;
@@ -1018,11 +1018,8 @@ async fn test_metadata_based_aggregate() -> Result<()> {
10181018
let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
10191019
let schema = Arc::new(Schema::new(vec![
10201020
Field::new("no_metadata", DataType::UInt64, true),
1021-
Field::new("with_metadata", DataType::UInt64, true).with_metadata(
1022-
[("modify_values".to_string(), "double_output".to_string())]
1023-
.into_iter()
1024-
.collect(),
1025-
),
1021+
Field::new("with_metadata", DataType::UInt64, true)
1022+
.with_metadata(Metadata::new().with("modify_values", "double_output")),
10261023
]));
10271024

10281025
let batch = RecordBatch::try_new(
@@ -1093,11 +1090,8 @@ async fn test_metadata_based_aggregate_as_window() -> Result<()> {
10931090
let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
10941091
let schema = Arc::new(Schema::new(vec![
10951092
Field::new("no_metadata", DataType::UInt64, true),
1096-
Field::new("with_metadata", DataType::UInt64, true).with_metadata(
1097-
[("modify_values".to_string(), "double_output".to_string())]
1098-
.into_iter()
1099-
.collect(),
1100-
),
1093+
Field::new("with_metadata", DataType::UInt64, true)
1094+
.with_metadata(Metadata::new().with("modify_values", "double_output")),
11011095
]));
11021096

11031097
let batch = RecordBatch::try_new(

0 commit comments

Comments
 (0)