Test: read/write ALP encoded parquet files (apache/arrow-rs#9372)

Cargo.toml

@@ -103,7 +103,7 @@ arrow-avro = { version = "59.2.0", default-features = false, features = [
 arrow-buffer = { version = "59.2.0", default-features = false }
 arrow-data = { version = "59.2.0", default-features = false }
 arrow-flight = { version = "59.2.0", features = [
-    "flight-sql-experimental",
+    "flight-sql",
 ] }
 # Both codecs are required here to make sure that code paths like
 # file-spilling have access to all compression codecs.
@@ -176,7 +176,7 @@ liblzma = { version = "0.4.6", features = ["static"] }
 log = "^0.4"
 memchr = "2.8.1"
 num-traits = { version = "0.2" }
-object_store = { version = "0.13.2", default-features = false }
+object_store = { version = "0.14.1", default-features = false }
 parking_lot = "0.12"
 parquet = { version = "59.2.0", default-features = false, features = [
     "arrow",
@@ -359,3 +359,20 @@ debug = false
 debug-assertions = false
 strip = "debuginfo"
 incremental = false
+
+## Temporary arrow-rs patch to test apache/arrow-rs#9372 (Parquet ALP encoder/decoder support)
+
+[patch.crates-io]
+arrow = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-array = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-avro = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-buffer = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-cast = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-data = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-ipc = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-schema = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-select = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-string = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-ord = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+arrow-flight = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }
+parquet = { git = "https://github.com/sdf-jkl/arrow-rs.git", branch = "alp" }

datafusion/catalog-listing/src/helpers.rs

@@ -30,7 +30,7 @@ use datafusion_expr::{BinaryExpr, Operator, lit, utils};
 
 use arrow::{
     array::AsArray,
-    datatypes::{DataType, Field},
+    datatypes::{DataType, Field, Metadata},
     record_batch::RecordBatch,
 };
 use datafusion_expr::execution_props::ExecutionProps;
@@ -423,7 +423,7 @@ pub async fn pruned_partition_list<'a>(
                 .iter()
                 .map(|(n, d)| Field::new(n, d.clone(), true))
                 .collect(),
-            Default::default(),
+            Metadata::new(),
         )?;
 
         Ok(objects

datafusion/catalog-listing/src/table.rs

@@ -20,7 +20,7 @@ use crate::helpers::{
     expr_applicable_for_cols, filter_partitioned_file, pruned_partition_list,
 };
 use crate::{ListingOptions, ListingTableConfig};
-use arrow::datatypes::{Field, Schema, SchemaBuilder, SchemaRef};
+use arrow::datatypes::{Field, Metadata, Schema, SchemaBuilder, SchemaRef};
 use async_trait::async_trait;
 use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider};
 use datafusion_common::stats::Precision;
@@ -1035,7 +1035,7 @@ impl ListingTable {
                 .iter()
                 .map(|(name, data_type)| Field::new(name, data_type.clone(), true))
                 .collect(),
-            Default::default(),
+            Metadata::new(),
         )?;
 
         file_groups

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::{BTreeSet, HashMap, HashSet};
+use std::collections::{BTreeSet, HashSet};
 use std::fmt::{Display, Formatter};
 use std::hash::Hash;
 use std::sync::{Arc, LazyLock};
@@ -31,7 +31,7 @@ use crate::{
 
 use arrow::compute::can_cast_types;
 use arrow::datatypes::{
-    DataType, Field, FieldRef, Fields, Schema, SchemaBuilder, SchemaRef,
+    DataType, Field, FieldRef, Fields, Metadata, Schema, SchemaBuilder, SchemaRef,
 };
 
 /// A reference-counted reference to a [DFSchema].
@@ -153,7 +153,7 @@ impl DFSchema {
     /// Create a `DFSchema` from an Arrow schema where all the fields have a given qualifier
     pub fn new_with_metadata(
         qualified_fields: Vec<(Option<TableReference>, Arc<Field>)>,
-        metadata: HashMap<String, String>,
+        metadata: impl Into<Metadata>,
     ) -> Result<Self> {
         let (qualifiers, fields): (Vec<Option<TableReference>>, Vec<Arc<Field>>) =
             qualified_fields.into_iter().unzip();
@@ -172,7 +172,7 @@ impl DFSchema {
     /// Create a new `DFSchema` from a list of Arrow [Field]s
     pub fn from_unqualified_fields(
         fields: Fields,
-        metadata: HashMap<String, String>,
+        metadata: impl Into<Metadata>,
     ) -> Result<Self> {
         let field_count = fields.len();
         let schema = Arc::new(Schema::new_with_metadata(fields, metadata));
@@ -864,7 +864,7 @@ impl DFSchema {
     }
 
     /// Get metadata of this schema
-    pub fn metadata(&self) -> &HashMap<String, String> {
+    pub fn metadata(&self) -> &Metadata {
         &self.inner.metadata
     }
 
@@ -1179,7 +1179,7 @@ impl ToDFSchema for Vec<Field> {
         let field_count = self.len();
         let schema = Schema {
             fields: self.into(),
-            metadata: HashMap::new(),
+            metadata: Metadata::new(),
         };
         let dfschema = DFSchema {
             inner: schema.into(),
@@ -1221,7 +1221,7 @@ pub trait ExprSchema: std::fmt::Debug {
     }
 
     /// Returns the column's optional metadata.
-    fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>> {
+    fn metadata(&self, col: &Column) -> Result<&Metadata> {
         Ok(self.field_from_column(col)?.metadata())
     }
 
@@ -1245,7 +1245,7 @@ impl<P: AsRef<DFSchema> + std::fmt::Debug> ExprSchema for P {
         self.as_ref().data_type(col)
     }
 
-    fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>> {
+    fn metadata(&self, col: &Column) -> Result<&Metadata> {
         ExprSchema::metadata(self.as_ref(), col)
     }
 
@@ -1379,6 +1379,7 @@ pub fn qualified_name(qualifier: Option<&TableReference>, name: &str) -> String
 #[cfg(test)]
 mod tests {
     use crate::assert_contains;
+    use std::collections::HashMap;
 
     use super::*;
 

datafusion/common/src/file_options/parquet_writer.rs

@@ -312,11 +312,12 @@ pub(crate) fn parse_encoding_string(
         "delta_byte_array" => Ok(parquet::basic::Encoding::DELTA_BYTE_ARRAY),
         "rle_dictionary" => Ok(parquet::basic::Encoding::RLE_DICTIONARY),
         "byte_stream_split" => Ok(parquet::basic::Encoding::BYTE_STREAM_SPLIT),
+        "alp" => Ok(parquet::basic::Encoding::ALP),
         _ => Err(DataFusionError::Configuration(format!(
             "Unknown or unsupported parquet encoding: \
         {str_setting}. Valid values are: plain, plain_dictionary, rle, \
         bit_packed, delta_binary_packed, delta_length_byte_array, \
-        delta_byte_array, rle_dictionary, and byte_stream_split."
+        delta_byte_array, rle_dictionary, byte_stream_split, and alp."
         ))),
     }
 }

datafusion/common/src/heap_size.rs

@@ -47,7 +47,7 @@ use arrow::array::{
 };
 use arrow::datatypes::{
     DataType, Field, Fields, IntervalDayTime, IntervalMonthDayNano, IntervalUnit,
-    TimeUnit, UnionFields, UnionMode, i256,
+    Metadata, TimeUnit, UnionFields, UnionMode, i256,
 };
 use chrono::{DateTime, Utc};
 use half::f16;
@@ -396,6 +396,19 @@ impl DFHeapSize for UnionFields {
     }
 }
 
+impl DFHeapSize for Metadata {
+    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
+        // `Metadata` does not expose its underlying reference-counted map, so
+        // this approximates the `BTreeMap` entries' sizes and cannot dedupe
+        // instances that share the same allocation.
+        self.iter()
+            .map(|(k, v)| {
+                size_of::<(String, String)>() + k.heap_size(ctx) + v.heap_size(ctx)
+            })
+            .sum()
+    }
+}
+
 impl DFHeapSize for Field {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
         self.name().heap_size(ctx)

datafusion/common/src/metadata.rs

@@ -17,7 +17,7 @@
 
 use std::{collections::BTreeMap, sync::Arc};
 
-use arrow::datatypes::{DataType, Field, FieldRef};
+use arrow::datatypes::{DataType, Field, FieldRef, Metadata};
 use hashbrown::HashMap;
 
 use crate::{DataFusionError, ScalarValue, error::_plan_err};
@@ -84,14 +84,8 @@ impl From<ScalarValue> for ScalarAndMetadata {
 /// Returns a planning error with suitably formatted type representations if
 /// actual and expected do not compare to equal.
 pub fn check_metadata_with_storage_equal(
-    actual: (
-        &DataType,
-        Option<&std::collections::HashMap<String, String>>,
-    ),
-    expected: (
-        &DataType,
-        Option<&std::collections::HashMap<String, String>>,
-    ),
+    actual: (&DataType, Option<&Metadata>),
+    expected: (&DataType, Option<&Metadata>),
     what: &str,
     context: &str,
 ) -> Result<(), DataFusionError> {
@@ -131,7 +125,7 @@ pub fn check_metadata_with_storage_equal(
 /// renderings.
 pub fn format_type_and_metadata(
     data_type: &DataType,
-    metadata: Option<&std::collections::HashMap<String, String>>,
+    metadata: Option<&Metadata>,
 ) -> String {
     match metadata {
         Some(metadata) if !metadata.is_empty() => {
@@ -316,6 +310,13 @@ impl FieldMetadata {
             .collect()
     }
 
+    /// Convert this `FieldMetadata` into an arrow [`Metadata`]
+    ///
+    /// This is cheap: both types share the same `Arc<BTreeMap>` representation.
+    pub fn to_metadata(&self) -> Metadata {
+        Metadata::from(Arc::clone(&self.inner))
+    }
+
     /// Updates the metadata on the Field with this metadata, if it is not empty.
     pub fn add_to_field(&self, field: Field) -> Field {
         if self.inner.is_empty() {
@@ -336,6 +337,24 @@ impl FieldMetadata {
     }
 }
 
+impl From<&FieldMetadata> for Metadata {
+    fn from(value: &FieldMetadata) -> Self {
+        value.to_metadata()
+    }
+}
+
+impl From<Metadata> for FieldMetadata {
+    fn from(value: Metadata) -> Self {
+        Self::new(value.into())
+    }
+}
+
+impl From<&Metadata> for FieldMetadata {
+    fn from(value: &Metadata) -> Self {
+        Self::from(value.clone())
+    }
+}
+
 impl From<&Field> for FieldMetadata {
     fn from(field: &Field) -> Self {
         Self::new_from_field(field)

datafusion/common/src/nested_struct.rs

@@ -1123,8 +1123,8 @@ mod tests {
                         Arc::new(non_null_field(
                             "entries",
                             struct_type(vec![
-                                non_null_field("keys", DataType::Utf8),
-                                field("values", DataType::Int32),
+                                non_null_field("key", DataType::Utf8),
+                                field("value", DataType::Int32),
                             ]),
                         )),
                         false,
@@ -1148,8 +1148,8 @@ mod tests {
                         Arc::new(non_null_field(
                             "entries",
                             struct_type(vec![
-                                non_null_field("keys", DataType::Utf8),
-                                field("values", DataType::Int32),
+                                non_null_field("key", DataType::Utf8),
+                                field("value", DataType::Int32),
                             ]),
                         )),
                         false,
@@ -1176,8 +1176,8 @@ mod tests {
         assert!(map.is_null(1));
         let map0 = map.value(0);
         let entries = map0.as_any().downcast_ref::<StructArray>().unwrap();
-        let keys = get_column_as!(entries, "keys", StringArray);
-        let vals = get_column_as!(entries, "values", Int32Array);
+        let keys = get_column_as!(entries, "key", StringArray);
+        let vals = get_column_as!(entries, "value", Int32Array);
         assert_eq!(keys.value(0), "a");
         assert_eq!(vals.value(0), 1);
     }

datafusion/common/src/param_value.rs

@@ -64,7 +64,7 @@ impl ParamValues {
                     check_metadata_with_storage_equal(
                         (
                             &lit.value.data_type(),
-                            lit.metadata.as_ref().map(|m| m.to_hashmap()).as_ref(),
+                            lit.metadata.as_ref().map(|m| m.to_metadata()).as_ref(),
                         ),
                         (param_type.data_type(), Some(param_type.metadata())),
                         "parameter",

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

@@ -136,6 +136,7 @@ mod tests {
                 },
                 range: Default::default(),
                 attributes: Attributes::default(),
+                extensions: Default::default(),
             })
         }
 

datafusion/core/tests/sql/path_partition.rs

@@ -718,6 +718,7 @@ impl ObjectStore for MirroringObjectStore {
             payload,
             meta,
             attributes: Attributes::default(),
+            extensions: Default::default(),
         })
     }
 
@@ -789,6 +790,7 @@ impl ObjectStore for MirroringObjectStore {
         Ok(ListResult {
             common_prefixes: common_prefixes.into_iter().collect(),
             objects,
+            extensions: Default::default(),
         })
     }
 

datafusion/core/tests/user_defined/user_defined_aggregates.rs

@@ -30,7 +30,7 @@ use arrow::array::{
     Array, AsArray, Int32Array, PrimitiveArray, StringArray, StructArray, UInt64Array,
     record_batch, types::UInt64Type,
 };
-use arrow::datatypes::{Fields, Schema};
+use arrow::datatypes::{Fields, Metadata, Schema};
 use arrow_schema::FieldRef;
 use datafusion::common::test_util::batches_to_string;
 use datafusion::dataframe::DataFrame;
@@ -1018,11 +1018,8 @@ async fn test_metadata_based_aggregate() -> Result<()> {
     let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
     let schema = Arc::new(Schema::new(vec![
         Field::new("no_metadata", DataType::UInt64, true),
-        Field::new("with_metadata", DataType::UInt64, true).with_metadata(
-            [("modify_values".to_string(), "double_output".to_string())]
-                .into_iter()
-                .collect(),
-        ),
+        Field::new("with_metadata", DataType::UInt64, true)
+            .with_metadata(Metadata::new().with("modify_values", "double_output")),
     ]));
 
     let batch = RecordBatch::try_new(
@@ -1093,11 +1090,8 @@ async fn test_metadata_based_aggregate_as_window() -> Result<()> {
     let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
     let schema = Arc::new(Schema::new(vec![
         Field::new("no_metadata", DataType::UInt64, true),
-        Field::new("with_metadata", DataType::UInt64, true).with_metadata(
-            [("modify_values".to_string(), "double_output".to_string())]
-                .into_iter()
-                .collect(),
-        ),
+        Field::new("with_metadata", DataType::UInt64, true)
+            .with_metadata(Metadata::new().with("modify_values", "double_output")),
     ]));
 
     let batch = RecordBatch::try_new(

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -25,7 +25,7 @@ use arrow::array::{
 };
 use arrow::array::{Int8Array, UInt64Array, as_string_array, create_array, record_batch};
 use arrow::compute::kernels::numeric::add;
-use arrow::datatypes::{DataType, Field, Schema};
+use arrow::datatypes::{DataType, Field, Metadata, Schema};
 use arrow_schema::extension::{Bool8, CanonicalExtensionType, ExtensionType};
 use arrow_schema::{ArrowError, FieldRef, SchemaRef};
 use datafusion::common::test_util::batches_to_string;
@@ -1645,11 +1645,8 @@ async fn test_metadata_based_udf() -> Result<()> {
     let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
     let schema = Arc::new(Schema::new(vec![
         Field::new("no_metadata", DataType::UInt64, true),
-        Field::new("with_metadata", DataType::UInt64, true).with_metadata(
-            [("modify_values".to_string(), "double_output".to_string())]
-                .into_iter()
-                .collect(),
-        ),
+        Field::new("with_metadata", DataType::UInt64, true)
+            .with_metadata(Metadata::new().with("modify_values", "double_output")),
     ]));
     let batch = RecordBatch::try_new(
         schema,