[WIP] Experiment with DataFusion against Arrow with Extension DataType support

Cargo.toml

@@ -87,19 +87,19 @@ ahash = { version = "0.8", default-features = false, features = [
     "runtime-rng",
 ] }
 apache-avro = { version = "0.17", default-features = false }
-arrow = { version = "54.3.1", features = [
+arrow = { git = "https://github.com/paleolimbot/arrow-rs.git", branch = "type-extension-maybe", features = [
     "prettyprint",
     "chrono-tz",
 ] }
-arrow-buffer = { version = "54.3.0", default-features = false }
-arrow-flight = { version = "54.3.1", features = [
+arrow-buffer = { git = "https://github.com/paleolimbot/arrow-rs.git", branch = "type-extension-maybe", default-features = false }
+arrow-flight = { git = "https://github.com/paleolimbot/arrow-rs.git", branch = "type-extension-maybe", features = [
     "flight-sql-experimental",
 ] }
-arrow-ipc = { version = "54.3.0", default-features = false, features = [
+arrow-ipc = { git = "https://github.com/paleolimbot/arrow-rs.git", branch = "type-extension-maybe", default-features = false, features = [
     "lz4",
 ] }
-arrow-ord = { version = "54.3.0", default-features = false }
-arrow-schema = { version = "54.3.0", default-features = false }
+arrow-ord = { git = "https://github.com/paleolimbot/arrow-rs.git", branch = "type-extension-maybe", default-features = false }
+arrow-schema = { git = "https://github.com/paleolimbot/arrow-rs.git", branch = "type-extension-maybe", default-features = false }
 async-trait = "0.1.88"
 bigdecimal = "0.4.8"
 bytes = "1.10"
@@ -147,9 +147,9 @@ hashbrown = { version = "0.14.5", features = ["raw"] }
 indexmap = "2.8.0"
 itertools = "0.14"
 log = "^0.4"
-object_store = { version = "0.11.0", default-features = false }
+object_store = { version = "0.12.0", default-features = false, features = ["fs"] }
 parking_lot = "0.12"
-parquet = { version = "54.3.1", default-features = false, features = [
+parquet = { git = "https://github.com/paleolimbot/arrow-rs.git", branch = "type-extension-maybe", default-features = false, features = [
     "arrow",
     "async",
     "object_store",
@@ -159,7 +159,7 @@ pbjson-types = "0.7"
 # Should match arrow-flight's version of prost.
 insta = { version = "1.41.1", features = ["glob", "filters"] }
 prost = "0.13.1"
-rand = "0.8.5"
+rand = { version = "0.9", features = ["std_rng"] }
 recursive = "0.1.1"
 regex = "1.8"
 rstest = "0.24.0"

benchmarks/src/cancellation.rs

@@ -38,7 +38,7 @@ use futures::TryStreamExt;
 use object_store::ObjectStore;
 use parquet::arrow::async_writer::ParquetObjectWriter;
 use parquet::arrow::AsyncArrowWriter;
-use rand::distributions::Alphanumeric;
+use rand::distr::Alphanumeric;
 use rand::rngs::ThreadRng;
 use rand::Rng;
 use structopt::StructOpt;

datafusion/common/Cargo.toml

@@ -63,7 +63,7 @@ log = { workspace = true }
 object_store = { workspace = true, optional = true }
 parquet = { workspace = true, optional = true, default-features = true }
 paste = "1.0.15"
-pyo3 = { version = "0.23.5", optional = true }
+pyo3 = { version = "0.24.1", optional = true }
 recursive = { workspace = true, optional = true }
 sqlparser = { workspace = true }
 tokio = { workspace = true }

datafusion/common/src/config.rs

@@ -428,7 +428,7 @@ config_namespace! {
         /// bytes of the parquet file optimistically. If not specified, two reads are required:
         /// One read to fetch the 8-byte parquet footer and
         /// another to fetch the metadata length encoded in the footer
-        pub metadata_size_hint: Option<usize>, default = None
+        pub metadata_size_hint: Option<u64>, default = None
 
         /// (reading) If true, filter expressions are be applied during the parquet decoding operation to
         /// reduce the number of rows decoded. This optimization is sometimes called "late materialization".

datafusion/common/src/scalar/mod.rs

@@ -2103,7 +2103,8 @@ impl ScalarValue {
             | DataType::Time64(TimeUnit::Millisecond)
             | DataType::RunEndEncoded(_, _)
             | DataType::ListView(_)
-            | DataType::LargeListView(_) => {
+            | DataType::LargeListView(_)
+            | DataType::Extension(_) => {
                 return _not_impl_err!(
                     "Unsupported creation of {:?} array from ScalarValue {:?}",
                     data_type,

datafusion/common/src/types/native.rs

@@ -411,6 +411,7 @@ impl From<DataType> for NativeType {
             DataType::Map(field, _) => Map(Arc::new(field.as_ref().into())),
             DataType::Dictionary(_, data_type) => data_type.as_ref().clone().into(),
             DataType::RunEndEncoded(_, field) => field.data_type().clone().into(),
+            DataType::Extension(extension) => extension.storage_type().clone().into(),
         }
     }
 }

datafusion/core/Cargo.toml

@@ -156,7 +156,7 @@ env_logger = { workspace = true }
 insta = { workspace = true }
 paste = "^1.0"
 rand = { workspace = true, features = ["small_rng"] }
-rand_distr = "0.4.3"
+rand_distr = "0.5.1"
 regex = { workspace = true }
 rstest = { workspace = true }
 serde_json = { workspace = true }

datafusion/core/benches/parquet_query_sql.rs

@@ -29,9 +29,9 @@ use datafusion_common::instant::Instant;
 use futures::stream::StreamExt;
 use parquet::arrow::ArrowWriter;
 use parquet::file::properties::{WriterProperties, WriterVersion};
-use rand::distributions::uniform::SampleUniform;
-use rand::distributions::Alphanumeric;
-use rand::prelude::*;
+use rand::distr::uniform::SampleUniform;
+use rand::distr::Alphanumeric;
+use rand::{prelude::*, thread_rng};
 use std::fs::File;
 use std::io::Read;
 use std::ops::Range;

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

@@ -442,7 +442,7 @@ mod tests {
         let object_meta = ObjectMeta {
             location,
             last_modified: DateTime::default(),
-            size: usize::MAX,
+            size: u64::MAX,
             e_tag: None,
             version: None,
         };
@@ -485,7 +485,7 @@ mod tests {
         let object_meta = ObjectMeta {
             location,
             last_modified: DateTime::default(),
-            size: usize::MAX,
+            size: u64::MAX,
             e_tag: None,
             version: None,
         };

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

@@ -103,7 +103,7 @@ mod tests {
 
         async fn get(&self, location: &Path) -> object_store::Result<GetResult> {
             let bytes = self.bytes_to_repeat.clone();
-            let range = 0..bytes.len() * self.max_iterations;
+            let range = 0..bytes.len() as u64 * self.max_iterations as u64;
             let arc = self.iterations_detected.clone();
             let stream = futures::stream::repeat_with(move || {
                 let arc_inner = arc.clone();
@@ -138,7 +138,7 @@ mod tests {
         async fn get_ranges(
             &self,
             _location: &Path,
-            _ranges: &[Range<usize>],
+            _ranges: &[Range<u64>],
         ) -> object_store::Result<Vec<Bytes>> {
             unimplemented!()
         }
@@ -153,8 +153,8 @@ mod tests {
 
         fn list(
             &self,
-            _prefix: Option<&Path>,
-        ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+            _: Option<&Path>,
+        ) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
             unimplemented!()
         }
 
@@ -371,7 +371,7 @@ mod tests {
         let object_meta = ObjectMeta {
             location: Path::parse("/")?,
             last_modified: DateTime::default(),
-            size: usize::MAX,
+            size: u64::MAX,
             e_tag: None,
             version: None,
         };
@@ -429,7 +429,7 @@ mod tests {
         let object_meta = ObjectMeta {
             location: Path::parse("/")?,
             last_modified: DateTime::default(),
-            size: usize::MAX,
+            size: u64::MAX,
             e_tag: None,
             version: None,
         };

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

@@ -330,8 +330,8 @@ mod tests {
 
         fn list(
             &self,
-            _prefix: Option<&Path>,
-        ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+            _: Option<&Path>,
+        ) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
             Box::pin(futures::stream::once(async {
                 Err(object_store::Error::NotImplemented)
             }))

datafusion/core/src/datasource/mod.rs

@@ -106,7 +106,7 @@ mod tests {
         let meta = ObjectMeta {
             location,
             last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(),
-            size: metadata.len() as usize,
+            size: metadata.len(),
             e_tag: None,
             version: None,
         };

datafusion/core/src/datasource/physical_plan/arrow_file.rs

@@ -305,7 +305,7 @@ impl FileOpener for ArrowOpener {
                     )?;
                     // read footer according to footer_len
                     let get_option = GetOptions {
-                        range: Some(GetRange::Suffix(10 + footer_len)),
+                        range: Some(GetRange::Suffix(10 + footer_len as u64)),
                         ..Default::default()
                     };
                     let get_result = object_store
@@ -332,9 +332,9 @@ impl FileOpener for ArrowOpener {
                         .iter()
                         .flatten()
                         .map(|block| {
-                            let block_len = block.bodyLength() as usize
-                                + block.metaDataLength() as usize;
-                            let block_offset = block.offset() as usize;
+                            let block_len =
+                                block.bodyLength() as u64 + block.metaDataLength() as u64;
+                            let block_offset = block.offset() as u64;
                             block_offset..block_offset + block_len
                         })
                         .collect_vec();
@@ -364,9 +364,9 @@ impl FileOpener for ArrowOpener {
                     let recordbatch_ranges = recordbatches
                         .iter()
                         .map(|block| {
-                            let block_len = block.bodyLength() as usize
-                                + block.metaDataLength() as usize;
-                            let block_offset = block.offset() as usize;
+                            let block_len =
+                                block.bodyLength() as u64 + block.metaDataLength() as u64;
+                            let block_offset = block.offset() as u64;
                             block_offset..block_offset + block_len
                         })
                         .collect_vec();

datafusion/core/src/datasource/physical_plan/parquet.rs

@@ -1806,7 +1806,7 @@ mod tests {
     #[derive(Debug, Clone)]
     struct TrackingParquetFileReaderFactory {
         inner: Arc<dyn ParquetFileReaderFactory>,
-        metadata_size_hint_calls: Arc<Mutex<Vec<Option<usize>>>>,
+        metadata_size_hint_calls: Arc<Mutex<Vec<Option<u64>>>>,
     }
 
     impl TrackingParquetFileReaderFactory {
@@ -1823,7 +1823,7 @@ mod tests {
             &self,
             partition_index: usize,
             file_meta: FileMeta,
-            metadata_size_hint: Option<usize>,
+            metadata_size_hint: Option<u64>,
             metrics: &ExecutionPlanMetricsSet,
         ) -> Result<Box<dyn parquet::arrow::async_reader::AsyncFileReader + Send>>
         {
@@ -1856,8 +1856,8 @@ mod tests {
         let schema = batch.schema();
         let name_1 = "test1.parquet";
         let name_2 = "test2.parquet";
-        let total_size_1 = write_batch(name_1, store.clone(), batch.clone()).await;
-        let total_size_2 = write_batch(name_2, store.clone(), batch.clone()).await;
+        let total_size_1 = write_batch(name_1, store.clone(), batch.clone()).await as u64;
+        let total_size_2 = write_batch(name_2, store.clone(), batch.clone()).await as u64;
 
         let reader_factory =
             Arc::new(TrackingParquetFileReaderFactory::new(store.clone()));

datafusion/core/src/test/object_store.rs

@@ -66,7 +66,7 @@ pub fn local_unpartitioned_file(path: impl AsRef<std::path::Path>) -> ObjectMeta
     ObjectMeta {
         location,
         last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(),
-        size: metadata.len() as usize,
+        size: metadata.len(),
         e_tag: None,
         version: None,
     }
@@ -166,7 +166,7 @@ impl ObjectStore for BlockingObjectStore {
     fn list(
         &self,
         prefix: Option<&Path>,
-    ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+    ) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
         self.inner.list(prefix)
     }
 

datafusion/core/src/test_util/parquet.rs

@@ -102,7 +102,7 @@ impl TestParquetFile {
 
         println!("Generated test dataset with {num_rows} rows");
 
-        let size = std::fs::metadata(&path)?.len() as usize;
+        let size = std::fs::metadata(&path)?.len();
 
         let mut canonical_path = path.canonicalize()?;
 

datafusion/datasource-avro/src/avro_to_arrow/schema.rs

@@ -239,6 +239,7 @@ fn default_field_name(dt: &DataType) -> &str {
         }
         DataType::Decimal128(_, _) => "decimal",
         DataType::Decimal256(_, _) => "decimal",
+        DataType::Extension(_) => unimplemented!("Extension support not implemented"),
     }
 }
 

datafusion/datasource-parquet/src/file_format.rs

@@ -193,13 +193,13 @@ impl ParquetFormat {
     /// another read to fetch the metadata length encoded in the footer.
     ///
     /// - If `None`, defaults to value on `config_options`
-    pub fn with_metadata_size_hint(mut self, size_hint: Option<usize>) -> Self {
+    pub fn with_metadata_size_hint(mut self, size_hint: Option<u64>) -> Self {
         self.options.global.metadata_size_hint = size_hint;
         self
     }
 
     /// Return the metadata size hint if set
-    pub fn metadata_size_hint(&self) -> Option<usize> {
+    pub fn metadata_size_hint(&self) -> Option<u64> {
         self.options.global.metadata_size_hint
     }
 
@@ -290,7 +290,7 @@ fn clear_metadata(
 async fn fetch_schema_with_location(
     store: &dyn ObjectStore,
     file: &ObjectMeta,
-    metadata_size_hint: Option<usize>,
+    metadata_size_hint: Option<u64>,
 ) -> Result<(Path, Schema)> {
     let loc_path = file.location.clone();
     let schema = fetch_schema(store, file, metadata_size_hint).await?;
@@ -735,15 +735,13 @@ impl<'a> ObjectStoreFetch<'a> {
 }
 
 impl MetadataFetch for ObjectStoreFetch<'_> {
-    fn fetch(
-        &mut self,
-        range: Range<usize>,
-    ) -> BoxFuture<'_, Result<Bytes, ParquetError>> {
-        async {
+    fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, ParquetError>> {
+        async move {
+            let range_usize: Range<u64> = range.start..range.end;
             self.store
-                .get_range(&self.meta.location, range)
+                .get_range(&self.meta.location, range_usize)
                 .await
-                .map_err(ParquetError::from)
+                .map_err(|e| ParquetError::External(Box::new(e)))
         }
         .boxed()
     }
@@ -758,13 +756,13 @@ impl MetadataFetch for ObjectStoreFetch<'_> {
 pub async fn fetch_parquet_metadata(
     store: &dyn ObjectStore,
     meta: &ObjectMeta,
-    size_hint: Option<usize>,
+    size_hint: Option<u64>,
 ) -> Result<ParquetMetaData> {
     let file_size = meta.size;
     let fetch = ObjectStoreFetch::new(store, meta);
 
     ParquetMetaDataReader::new()
-        .with_prefetch_hint(size_hint)
+        .with_prefetch_hint(size_hint.map(|n| n.try_into().unwrap()))
         .load_and_finish(fetch, file_size)
         .await
         .map_err(DataFusionError::from)
@@ -774,7 +772,7 @@ pub async fn fetch_parquet_metadata(
 async fn fetch_schema(
     store: &dyn ObjectStore,
     file: &ObjectMeta,
-    metadata_size_hint: Option<usize>,
+    metadata_size_hint: Option<u64>,
 ) -> Result<Schema> {
     let metadata = fetch_parquet_metadata(store, file, metadata_size_hint).await?;
     let file_metadata = metadata.file_metadata();
@@ -792,7 +790,7 @@ pub async fn fetch_statistics(
     store: &dyn ObjectStore,
     table_schema: SchemaRef,
     file: &ObjectMeta,
-    metadata_size_hint: Option<usize>,
+    metadata_size_hint: Option<u64>,
 ) -> Result<Statistics> {
     let metadata = fetch_parquet_metadata(store, file, metadata_size_hint).await?;
     statistics_from_parquet_meta_calc(&metadata, table_schema)