buoyant_kernel/engine/default/
parquet.rs

1//! Default Parquet handler implementation
2
3use std::collections::HashMap;
4use std::ops::Range;
5use std::sync::Arc;
6
7use futures::stream::{self, BoxStream};
8use futures::{StreamExt, TryStreamExt};
9use uuid::Uuid;
10
11use super::file_stream::{FileOpenFuture, FileOpener, FileStream};
12use super::stats::collect_stats;
13use super::UrlExt;
14use crate::arrow::array::builder::{MapBuilder, MapFieldNames, StringBuilder};
15use crate::arrow::array::{Array, Int64Array, RecordBatch, StringArray, StructArray};
16use crate::arrow::datatypes::{DataType, Field, Schema};
17use crate::engine::arrow_conversion::{TryFromArrow as _, TryIntoArrow as _};
18use crate::engine::arrow_data::ArrowEngineData;
19use crate::engine::arrow_utils::{
20    fixup_parquet_read, generate_mask, get_requested_indices, ordering_needs_row_indexes,
21    RowIndexBuilder,
22};
23use crate::engine::default::executor::TaskExecutor;
24use crate::engine::parquet_row_group_skipping::ParquetRowGroupSkipping;
25use crate::engine::{reader_options, writer_options};
26use crate::expressions::ColumnName;
27use crate::metrics::emit_parquet_read_completed;
28use crate::object_store::path::Path;
29use crate::object_store::{DynObjectStore, ObjectStoreExt as _};
30use crate::parquet::arrow::arrow_reader::{ArrowReaderMetadata, ParquetRecordBatchReaderBuilder};
31use crate::parquet::arrow::arrow_writer::ArrowWriter;
32use crate::parquet::arrow::async_reader::{ParquetObjectReader, ParquetRecordBatchStreamBuilder};
33use crate::parquet::arrow::async_writer::{AsyncArrowWriter, ParquetObjectWriter};
34use crate::schema::{SchemaRef, StructType};
35use crate::transaction::WriteContext;
36use crate::{
37    DeltaResult, EngineData, Error, FileDataReadResultIterator, FileMeta, ParquetFooter,
38    ParquetHandler, PredicateRef,
39};
40
41#[derive(Debug)]
42pub struct DefaultParquetHandler<E: TaskExecutor> {
43    store: Arc<DynObjectStore>,
44    task_executor: Arc<E>,
45    readahead: usize,
46}
47
48/// Metadata of a data file (typically a parquet file).
49#[derive(Debug)]
50pub struct DataFileMetadata {
51    file_meta: FileMeta,
52    /// Collected statistics for this file (includes numRecords, tightBounds, etc.).
53    stats: StructArray,
54}
55
56impl DataFileMetadata {
57    pub fn new(file_meta: FileMeta, stats: StructArray) -> Self {
58        Self { file_meta, stats }
59    }
60
61    /// Returns the absolute URL of the written file.
62    pub fn location(&self) -> &url::Url {
63        &self.file_meta.location
64    }
65
66    /// Converts this `DataFileMetadata` into an [`EngineData`] record batch matching the schema
67    /// returned by [`Transaction::add_files_schema`].
68    ///
69    /// The `partition_values` map uses physical column names as keys and protocol-serialized
70    /// strings as values. `None` represents a null partition value. The serialization layer
71    /// converts nulls and empty strings to `None` before reaching this method, so `Some("")`
72    /// is not expected in normal usage.
73    ///
74    /// The `log_path` is the path string written to the Delta log's `add.path` field.
75    ///
76    /// [`Transaction::add_files_schema`]: crate::transaction::Transaction::add_files_schema
77    pub(crate) fn as_record_batch(
78        &self,
79        partition_values: &HashMap<String, Option<String>>,
80        log_path: &str,
81    ) -> DeltaResult<Box<dyn EngineData>> {
82        let path = Arc::new(StringArray::from(vec![log_path]));
83        let key_builder = StringBuilder::new();
84        let val_builder = StringBuilder::new();
85        let names = MapFieldNames {
86            entry: "key_value".to_string(),
87            key: "key".to_string(),
88            value: "value".to_string(),
89        };
90        let mut builder = MapBuilder::new(Some(names), key_builder, val_builder);
91        for (k, v) in partition_values {
92            builder.keys().append_value(k);
93            match v.as_deref() {
94                // The serialization layer already converts empty strings to None, so
95                // Some("") should not occur. The empty check is purely defensive.
96                Some(val) if !val.is_empty() => builder.values().append_value(val),
97                _ => builder.values().append_null(),
98            }
99        }
100        builder.append(true)?;
101        let partitions = Arc::new(builder.finish());
102        // this means max size we can write is i64::MAX (~8EB)
103        let size: i64 = self
104            .file_meta
105            .size
106            .try_into()
107            .map_err(|_| Error::generic("Failed to convert parquet metadata 'size' to i64"))?;
108        let size = Arc::new(Int64Array::from(vec![size]));
109        let modification_time = Arc::new(Int64Array::from(vec![self.file_meta.last_modified]));
110
111        let stats_array = Arc::new(self.stats.clone());
112
113        // Build schema dynamically based on stats (stats schema varies based on collected
114        // statistics)
115        let key_value_struct = DataType::Struct(
116            vec![
117                Field::new("key", DataType::Utf8, false),
118                Field::new("value", DataType::Utf8, true),
119            ]
120            .into(),
121        );
122        let schema = Schema::new(vec![
123            Field::new("path", DataType::Utf8, false),
124            Field::new(
125                "partitionValues",
126                DataType::Map(
127                    Arc::new(Field::new("key_value", key_value_struct, false)),
128                    false,
129                ),
130                false,
131            ),
132            Field::new("size", DataType::Int64, false),
133            Field::new("modificationTime", DataType::Int64, false),
134            Field::new("stats", stats_array.data_type().clone(), true),
135        ]);
136
137        Ok(Box::new(ArrowEngineData::new(RecordBatch::try_new(
138            Arc::new(schema),
139            vec![path, partitions, size, modification_time, stats_array],
140        )?)))
141    }
142}
143
144impl<E: TaskExecutor> DefaultParquetHandler<E> {
145    pub fn new(store: Arc<DynObjectStore>, task_executor: Arc<E>) -> Self {
146        Self {
147            store,
148            task_executor,
149            readahead: 10,
150        }
151    }
152
153    /// Max number of batches to read ahead while executing [Self::read_parquet_files()].
154    ///
155    /// Defaults to 10.
156    pub fn with_readahead(mut self, readahead: usize) -> Self {
157        self.readahead = readahead;
158        self
159    }
160
161    // Write `data` to `{path}/<uuid>.parquet` as parquet using ArrowWriter and return the parquet
162    // metadata (where `<uuid>` is a generated UUIDv4).
163    //
164    // Note: after encoding the data as parquet, this issues a PUT followed by a HEAD to storage in
165    // order to obtain metadata about the object just written.
166    async fn write_parquet(
167        &self,
168        path: &url::Url,
169        data: Box<dyn EngineData>,
170        stats_columns: &[ColumnName],
171    ) -> DeltaResult<DataFileMetadata> {
172        let batch: Box<_> = ArrowEngineData::try_from_engine_data(data)?;
173        let record_batch = batch.record_batch();
174
175        // Collect statistics before writing (includes numRecords)
176        let stats = collect_stats(record_batch, stats_columns)?;
177
178        let mut buffer = vec![];
179        let mut writer = ArrowWriter::try_new_with_options(
180            &mut buffer,
181            record_batch.schema(),
182            writer_options(),
183        )?;
184        writer.write(record_batch)?;
185        writer.close()?; // writer must be closed to write footer
186
187        let size: u64 = buffer
188            .len()
189            .try_into()
190            .map_err(|_| Error::generic("unable to convert usize to u64"))?;
191        let name: String = format!("{}.parquet", Uuid::new_v4());
192        // fail if path does not end with a trailing slash
193        if !path.path().ends_with('/') {
194            return Err(Error::generic(format!(
195                "Path must end with a trailing slash: {path}"
196            )));
197        }
198        let path = path.join(&name)?;
199
200        self.store
201            .put(&Path::from_url_path(path.path())?, buffer.into())
202            .await?;
203
204        let metadata = self.store.head(&Path::from_url_path(path.path())?).await?;
205        let modification_time = metadata.last_modified.timestamp_millis();
206        if size != metadata.size {
207            return Err(Error::generic(format!(
208                "Size mismatch after writing parquet file: expected {}, got {}",
209                size, metadata.size
210            )));
211        }
212
213        let file_meta = FileMeta::new(path, modification_time, size);
214        Ok(DataFileMetadata::new(file_meta, stats))
215    }
216
217    /// Write `data` to a new parquet file under the [`WriteContext::write_dir`] and return
218    /// Add action metadata ready for [`Transaction::add_files`].
219    ///
220    /// Note that the schema does not contain the dataChange column. In order to set `data_change`
221    /// flag, use [`crate::transaction::Transaction::with_data_change`].
222    ///
223    /// [`WriteContext::write_dir`]: crate::transaction::WriteContext::write_dir
224    /// [`Transaction::add_files`]: crate::transaction::Transaction::add_files
225    pub async fn write_parquet_file(
226        &self,
227        data: Box<dyn EngineData>,
228        write_context: &WriteContext,
229    ) -> DeltaResult<Box<dyn EngineData>> {
230        let file_metadata = self
231            .write_parquet(
232                &write_context.write_dir(),
233                data,
234                write_context.stats_columns(),
235            )
236            .await?;
237        super::build_add_file_metadata(file_metadata, write_context)
238    }
239}
240
241/// Internal async implementation of read_parquet_files
242async fn read_parquet_files_impl(
243    store: Arc<DynObjectStore>,
244    files: Vec<FileMeta>,
245    physical_schema: SchemaRef,
246    predicate: Option<PredicateRef>,
247) -> DeltaResult<BoxStream<'static, DeltaResult<Box<dyn EngineData>>>> {
248    if files.is_empty() {
249        return Ok(Box::pin(stream::empty()));
250    }
251
252    let arrow_schema = Arc::new(physical_schema.as_ref().try_into_arrow()?);
253
254    // get the first FileMeta to decide how to fetch the file.
255    // NB: This means that every file in `FileMeta` _must_ have the same scheme or things will break
256    // s3://    -> aws   (ParquetOpener)
257    // nothing  -> local (ParquetOpener)
258    // https:// -> assume presigned URL (and fetch without object_store)
259    //   -> reqwest to get data
260    //   -> parse to parquet
261    // SAFETY: we did is_empty check above, this is ok.
262    if files[0].location.is_presigned() {
263        let file_opener = Box::new(PresignedUrlOpener::new(
264            1024,
265            physical_schema.clone(),
266            predicate,
267        ));
268        let stream = FileStream::new(files, arrow_schema, file_opener)?.map_ok(
269            |record_batch| -> Box<dyn EngineData> { Box::new(ArrowEngineData::new(record_batch)) },
270        );
271        return Ok(Box::pin(stream));
272    }
273
274    // an iterator of futures that open each file
275    let file_futures = files.into_iter().map(move |file| {
276        let store = store.clone();
277        let schema = physical_schema.clone();
278        let predicate = predicate.clone();
279        async move {
280            open_parquet_file(
281                store,
282                schema,
283                predicate,
284                None,
285                super::DEFAULT_BATCH_SIZE,
286                file,
287            )
288            .await
289        }
290    });
291    // create a stream from that iterator which buffers up to `buffer_size` futures at a time
292    let result_stream = stream::iter(file_futures)
293        .buffered(super::DEFAULT_BUFFER_SIZE)
294        .try_flatten()
295        .map_ok(|record_batch| -> Box<dyn EngineData> {
296            Box::new(ArrowEngineData::new(record_batch))
297        });
298
299    Ok(Box::pin(result_stream))
300}
301
302impl<E: TaskExecutor> ParquetHandler for DefaultParquetHandler<E> {
303    fn read_parquet_files(
304        &self,
305        files: &[FileMeta],
306        physical_schema: SchemaRef,
307        predicate: Option<PredicateRef>,
308    ) -> DeltaResult<FileDataReadResultIterator> {
309        let num_files = files.len() as u64;
310        let bytes_read = files.iter().map(|f| f.size).sum();
311        let future = read_parquet_files_impl(
312            self.store.clone(),
313            files.to_vec(),
314            physical_schema,
315            predicate,
316        );
317        let inner = super::stream_future_to_iter(self.task_executor.clone(), future)?;
318        Ok(Box::new(super::ReadMetricsIterator::new(
319            inner,
320            num_files,
321            bytes_read,
322            emit_parquet_read_completed,
323        )))
324    }
325
326    /// Writes engine data to a Parquet file at the specified location.
327    ///
328    /// This implementation uses asynchronous file I/O with object_store to write the Parquet file.
329    /// If a file already exists at the given location, it will be overwritten.
330    ///
331    /// # Parameters
332    ///
333    /// - `location` - The full URL path where the Parquet file should be written (e.g., `s3://bucket/path/file.parquet`,
334    ///   `file:///path/to/file.parquet`).
335    /// - `data` - An iterator of engine data to be written to the Parquet file.
336    ///
337    /// # Returns
338    ///
339    /// A [`DeltaResult`] indicating success or failure.
340    fn write_parquet_file(
341        &self,
342        location: url::Url,
343        mut data: Box<dyn Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send>,
344    ) -> DeltaResult<()> {
345        use tokio::sync::Mutex;
346
347        let store = self.store.clone();
348        let path = Path::from_url_path(location.path())?;
349
350        // Get first batch to initialize writer with schema
351        let first_batch = data.next().ok_or_else(|| {
352            Error::generic("Cannot write parquet file with empty data iterator")
353        })??;
354        let first_arrow = ArrowEngineData::try_from_engine_data(first_batch)?;
355        let first_record_batch: RecordBatch = (*first_arrow).into();
356
357        let object_writer = ParquetObjectWriter::new(store, path);
358        let schema = first_record_batch.schema();
359        let writer = Arc::new(Mutex::new(AsyncArrowWriter::try_new_with_options(
360            object_writer,
361            schema,
362            writer_options(),
363        )?));
364
365        let w = writer.clone();
366        self.task_executor.block_on(async move {
367            // Write the first batch
368            let mut writer = w.lock().await;
369            writer.write(&first_record_batch).await
370        })?;
371
372        // In order to avoid a deadlock when running inside a single-threaded runtime, this code
373        // iterator over the blocking stream of `data` separately from the `write` awaits` .
374        // See: https://github.com/delta-io/delta-kernel-rs/issues/2399>
375        for result in data {
376            let engine_data = result?;
377            let arrow_data = ArrowEngineData::try_from_engine_data(engine_data)?;
378            let batch: RecordBatch = (*arrow_data).into();
379            let w = writer.clone();
380            self.task_executor.block_on(async move {
381                let mut writer = w.lock().await;
382                writer.write(&batch).await
383            })?;
384        }
385
386        let w = writer.clone();
387        self.task_executor.block_on(async move {
388            let mut writer = w.lock().await;
389            writer.finish().await
390        })?;
391
392        Ok(())
393    }
394
395    fn read_parquet_footer(&self, file: &FileMeta) -> DeltaResult<ParquetFooter> {
396        let store = self.store.clone();
397        let location = file.location.clone();
398        let file_size = file.size;
399
400        self.task_executor.block_on(async move {
401            let metadata = if location.is_presigned() {
402                let client = reqwest::Client::new();
403                let response =
404                    client.get(location.as_str()).send().await.map_err(|e| {
405                        Error::generic(format!("Failed to fetch presigned URL: {e}"))
406                    })?;
407                let bytes = response
408                    .bytes()
409                    .await
410                    .map_err(|e| Error::generic(format!("Failed to read response bytes: {e}")))?;
411                ArrowReaderMetadata::load(&bytes, reader_options())?
412            } else {
413                let path = Path::from_url_path(location.path())?;
414                let mut reader = ParquetObjectReader::new(store, path).with_file_size(file_size);
415                ArrowReaderMetadata::load_async(&mut reader, reader_options()).await?
416            };
417
418            let schema = StructType::try_from_arrow(metadata.schema().as_ref())
419                .map(Arc::new)
420                .map_err(Error::Arrow)?;
421            Ok(ParquetFooter { schema })
422        })
423    }
424}
425
426/// Opens a Parquet file and returns a stream of record batches
427async fn open_parquet_file(
428    store: Arc<DynObjectStore>,
429    table_schema: SchemaRef,
430    predicate: Option<PredicateRef>,
431    limit: Option<usize>,
432    batch_size: usize,
433    file_meta: FileMeta,
434) -> DeltaResult<BoxStream<'static, DeltaResult<RecordBatch>>> {
435    let file_location = file_meta.location.to_string();
436    let path = Path::from_url_path(file_meta.location.path())?;
437
438    let mut reader = {
439        use crate::object_store::ObjectStoreScheme;
440        // HACK: unfortunately, `ParquetObjectReader` under the hood does a suffix range
441        // request which isn't supported by Azure. For now we just detect if the URL is
442        // pointing to azure and if so, do a HEAD request so we can pass in file size to the
443        // reader which will cause the reader to avoid a suffix range request.
444        // see also: https://github.com/delta-io/delta-kernel-rs/issues/968
445
446        // Since the `Remove` action's size value is optional as specified in the delta protocol
447        // https://github.com/delta-io/delta/blob/master/PROTOCOL.md#add-file-and-remove-file,
448        // the extracted size will be zero in this case. Thus, this function
449        // need to handle the case of zero file_meta.size.
450        if file_meta.size != 0 {
451            ParquetObjectReader::new(store, path).with_file_size(file_meta.size)
452        } else if let Ok((ObjectStoreScheme::MicrosoftAzure, _)) =
453            ObjectStoreScheme::parse(&file_meta.location)
454        {
455            // also note doing HEAD then actual GET isn't atomic, and leaves us vulnerable
456            // to file changing between the two calls.
457            let meta = store.head(&path).await?;
458            ParquetObjectReader::new(store, path).with_file_size(meta.size)
459        } else {
460            ParquetObjectReader::new(store, path)
461        }
462    };
463
464    let reader_options = reader_options();
465    let metadata = ArrowReaderMetadata::load_async(&mut reader, reader_options.clone()).await?;
466    let parquet_schema = metadata.schema();
467    let (indices, requested_ordering) = get_requested_indices(&table_schema, parquet_schema)?;
468    let mut builder =
469        ParquetRecordBatchStreamBuilder::new_with_options(reader, reader_options).await?;
470    if let Some(mask) = generate_mask(
471        &table_schema,
472        parquet_schema,
473        builder.parquet_schema(),
474        &indices,
475    ) {
476        builder = builder.with_projection(mask)
477    }
478
479    // Only create RowIndexBuilder if row indexes are actually needed
480    let mut row_indexes = ordering_needs_row_indexes(&requested_ordering)
481        .then(|| RowIndexBuilder::new(builder.metadata().row_groups()));
482
483    // Filter row groups and row indexes if a predicate is provided
484    if let Some(ref predicate) = predicate {
485        builder = builder.with_row_group_filter(predicate, row_indexes.as_mut());
486    }
487    if let Some(limit) = limit {
488        builder = builder.with_limit(limit)
489    }
490
491    let mut row_indexes = row_indexes.map(|rb| rb.build()).transpose()?;
492    let stream = builder.with_batch_size(batch_size).build()?;
493
494    let stream = stream.map(move |rbr| {
495        fixup_parquet_read(
496            rbr?,
497            &requested_ordering,
498            row_indexes.as_mut(),
499            Some(&file_location),
500            Some(&table_schema),
501        )
502        .map(Into::into)
503    });
504    Ok(stream.boxed())
505}
506
507/// Implements [`FileOpener`] for a opening a parquet file from a presigned URL
508struct PresignedUrlOpener {
509    batch_size: usize,
510    predicate: Option<PredicateRef>,
511    limit: Option<usize>,
512    table_schema: SchemaRef,
513    client: reqwest::Client,
514}
515
516impl PresignedUrlOpener {
517    pub(crate) fn new(
518        batch_size: usize,
519        schema: SchemaRef,
520        predicate: Option<PredicateRef>,
521    ) -> Self {
522        Self {
523            batch_size,
524            table_schema: schema,
525            predicate,
526            limit: None,
527            client: reqwest::Client::new(),
528        }
529    }
530}
531
532impl FileOpener for PresignedUrlOpener {
533    fn open(&self, file_meta: FileMeta, _range: Option<Range<i64>>) -> DeltaResult<FileOpenFuture> {
534        let batch_size = self.batch_size;
535        let table_schema = self.table_schema.clone();
536        let predicate = self.predicate.clone();
537        let limit = self.limit;
538        let client = self.client.clone(); // uses Arc internally according to reqwest docs
539        let file_location = file_meta.location.to_string();
540
541        Ok(Box::pin(async move {
542            // fetch the file from the interweb
543            let reader = client.get(&file_location).send().await?.bytes().await?;
544            let reader_options = reader_options();
545            let metadata = ArrowReaderMetadata::load(&reader, reader_options.clone())?;
546            let parquet_schema = metadata.schema();
547            let (indices, requested_ordering) =
548                get_requested_indices(&table_schema, parquet_schema)?;
549
550            let mut builder =
551                ParquetRecordBatchReaderBuilder::try_new_with_options(reader, reader_options)?;
552            if let Some(mask) = generate_mask(
553                &table_schema,
554                parquet_schema,
555                builder.parquet_schema(),
556                &indices,
557            ) {
558                builder = builder.with_projection(mask)
559            }
560
561            // Only create RowIndexBuilder if row indexes are actually needed
562            let mut row_indexes = ordering_needs_row_indexes(&requested_ordering)
563                .then(|| RowIndexBuilder::new(builder.metadata().row_groups()));
564
565            // Filter row groups and row indexes if a predicate is provided
566            if let Some(ref predicate) = predicate {
567                builder = builder.with_row_group_filter(predicate, row_indexes.as_mut());
568            }
569            if let Some(limit) = limit {
570                builder = builder.with_limit(limit)
571            }
572
573            let reader = builder.with_batch_size(batch_size).build()?;
574
575            let mut row_indexes = row_indexes.map(|rb| rb.build()).transpose()?;
576            let stream = futures::stream::iter(reader);
577            let stream = stream.map(move |rbr| {
578                fixup_parquet_read(
579                    rbr?,
580                    &requested_ordering,
581                    row_indexes.as_mut(),
582                    Some(&file_location),
583                    Some(&table_schema),
584                )
585                .map(Into::into)
586            });
587            Ok(stream.boxed())
588        }))
589    }
590}
591
592#[cfg(test)]
593mod tests {
594    use std::collections::HashMap;
595    use std::path::PathBuf;
596    use std::slice;
597
598    use itertools::Itertools;
599    use url::Url;
600
601    use super::*;
602
603    #[cfg(feature = "nanosecond-timestamps")]
604    use crate::arrow::array::TimestampNanosecondArray;
605    use crate::actions::{NUM_RECORDS, TIGHT_BOUNDS};
606    use crate::arrow::array::{
607        Array, BinaryArray, BooleanArray, Date32Array, Decimal128Array, Float32Array, Float64Array,
608        Int16Array, Int32Array, Int64Array, Int8Array, RecordBatch, StringArray,
609        TimestampMicrosecondArray,
610    };
611    use crate::arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
612    use crate::engine::arrow_conversion::TryIntoKernel as _;
613    use crate::engine::arrow_data::ArrowEngineData;
614    use crate::engine::default::executor::tokio::TokioBackgroundExecutor;
615    use crate::engine::default::DEFAULT_BATCH_SIZE;
616    use crate::object_store::local::LocalFileSystem;
617    use crate::object_store::memory::InMemory;
618    use crate::parquet::arrow::{ARROW_SCHEMA_META_KEY, PARQUET_FIELD_ID_META_KEY};
619    use crate::schema::{ColumnMetadataKey, MetadataValue, StructField, StructType};
620    use crate::utils::current_time_ms;
621    use crate::utils::test_utils::assert_result_error_with_message;
622    use crate::EngineData;
623
624    fn into_record_batch(
625        engine_data: DeltaResult<Box<dyn EngineData>>,
626    ) -> DeltaResult<RecordBatch> {
627        engine_data
628            .and_then(ArrowEngineData::try_from_engine_data)
629            .map(Into::into)
630    }
631
632    async fn read_all_rows_helper(file_meta: FileMeta) -> DeltaResult<Vec<RecordBatch>> {
633        let store = Arc::new(LocalFileSystem::new());
634        let path = Path::from_url_path(file_meta.location.path()).unwrap();
635        let reader = ParquetObjectReader::new(store.clone(), path);
636        let physical_schema = ParquetRecordBatchStreamBuilder::new(reader)
637            .await
638            .unwrap()
639            .schema()
640            .clone();
641        let stream = open_parquet_file(
642            store,
643            Arc::new(physical_schema.try_into_kernel().unwrap()),
644            None,
645            None,
646            DEFAULT_BATCH_SIZE,
647            file_meta,
648        )
649        .await
650        .unwrap();
651
652        let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
653        Ok(batches)
654    }
655
656    #[tokio::test]
657    async fn test_open_parquet_file_with_size() {
658        let path = std::fs::canonicalize(PathBuf::from(
659            "./tests/data/table-with-dv-small/part-00000-fae5310a-a37d-4e51-827b-c3d5516560ca-c000.snappy.parquet"
660        )).unwrap();
661        let file_size = std::fs::metadata(&path).unwrap().len();
662        let url = Url::from_file_path(path).unwrap();
663        let file_meta = FileMeta {
664            location: url,
665            last_modified: 0,
666            size: file_size,
667        };
668        let data = read_all_rows_helper(file_meta).await.unwrap();
669
670        assert_eq!(data.len(), 1);
671        assert_eq!(data[0].num_rows(), 10);
672    }
673
674    #[tokio::test]
675    async fn test_open_parquet_file_without_size() {
676        let path = std::fs::canonicalize(PathBuf::from(
677            "./tests/data/table-with-dv-small/part-00000-fae5310a-a37d-4e51-827b-c3d5516560ca-c000.snappy.parquet"
678        )).unwrap();
679        let url = Url::from_file_path(path).unwrap();
680        let file_meta = FileMeta {
681            location: url,
682            last_modified: 0,
683            size: 0,
684        };
685        let data = read_all_rows_helper(file_meta).await.unwrap();
686
687        assert_eq!(data.len(), 1);
688        assert_eq!(data[0].num_rows(), 10);
689    }
690
691    #[tokio::test]
692    async fn test_read_parquet_files() {
693        let store = Arc::new(LocalFileSystem::new());
694
695        let path = std::fs::canonicalize(PathBuf::from(
696            "./tests/data/table-with-dv-small/part-00000-fae5310a-a37d-4e51-827b-c3d5516560ca-c000.snappy.parquet"
697        )).unwrap();
698        let url = url::Url::from_file_path(path).unwrap();
699        let location = Path::from_url_path(url.path()).unwrap();
700        let meta = store.head(&location).await.unwrap();
701
702        let reader = ParquetObjectReader::new(store.clone(), location);
703        let physical_schema = ParquetRecordBatchStreamBuilder::new(reader)
704            .await
705            .unwrap()
706            .schema()
707            .clone();
708
709        let files = &[FileMeta {
710            location: url.clone(),
711            last_modified: meta.last_modified.timestamp(),
712            size: meta.size,
713        }];
714
715        let handler = DefaultParquetHandler::new(store, Arc::new(TokioBackgroundExecutor::new()));
716        let data: Vec<RecordBatch> = handler
717            .read_parquet_files(
718                files,
719                Arc::new(physical_schema.try_into_kernel().unwrap()),
720                None,
721            )
722            .unwrap()
723            .map(into_record_batch)
724            .try_collect()
725            .unwrap();
726
727        assert_eq!(data.len(), 1);
728        assert_eq!(data[0].num_rows(), 10);
729    }
730
731    #[rstest::rstest]
732    fn test_as_record_batch(
733        #[values(None, Some("a".to_string()))] partition_value: Option<String>,
734    ) {
735        let location = Url::parse("file:///test_url").unwrap();
736        let size = 1_000_000;
737        let last_modified = 10000000000;
738        let num_records = 10;
739        let file_metadata = FileMeta::new(location.clone(), last_modified, size);
740        let stats = StructArray::try_new(
741            vec![
742                Field::new(NUM_RECORDS, ArrowDataType::Int64, true),
743                Field::new(TIGHT_BOUNDS, ArrowDataType::Boolean, true),
744            ]
745            .into(),
746            vec![
747                Arc::new(Int64Array::from(vec![num_records as i64])),
748                Arc::new(BooleanArray::from(vec![true])),
749            ],
750            None,
751        )
752        .unwrap();
753        let data_file_metadata = DataFileMetadata::new(file_metadata, stats.clone());
754        let partition_values = HashMap::from([("partition1".to_string(), partition_value.clone())]);
755        let actual = data_file_metadata
756            .as_record_batch(&partition_values, "test_url")
757            .unwrap();
758        let actual = ArrowEngineData::try_from_engine_data(actual).unwrap();
759
760        let mut partition_values_builder = MapBuilder::new(
761            Some(MapFieldNames {
762                entry: "key_value".to_string(),
763                key: "key".to_string(),
764                value: "value".to_string(),
765            }),
766            StringBuilder::new(),
767            StringBuilder::new(),
768        );
769
770        partition_values_builder.keys().append_value("partition1");
771        match &partition_value {
772            None => partition_values_builder.values().append_null(),
773            Some(v) => partition_values_builder.values().append_value(v),
774        }
775        partition_values_builder.append(true).unwrap();
776        let partition_values = partition_values_builder.finish();
777
778        // Build expected schema dynamically based on stats
779        let stats_field = Field::new("stats", stats.data_type().clone(), true);
780        let schema = Arc::new(crate::arrow::datatypes::Schema::new(vec![
781            Field::new("path", ArrowDataType::Utf8, false),
782            Field::new(
783                "partitionValues",
784                ArrowDataType::Map(
785                    Arc::new(Field::new(
786                        "key_value",
787                        ArrowDataType::Struct(
788                            vec![
789                                Field::new("key", ArrowDataType::Utf8, false),
790                                Field::new("value", ArrowDataType::Utf8, true),
791                            ]
792                            .into(),
793                        ),
794                        false,
795                    )),
796                    false,
797                ),
798                false,
799            ),
800            Field::new("size", ArrowDataType::Int64, false),
801            Field::new("modificationTime", ArrowDataType::Int64, false),
802            stats_field,
803        ]));
804
805        let expected = RecordBatch::try_new(
806            schema,
807            vec![
808                Arc::new(StringArray::from(vec!["test_url"])),
809                Arc::new(partition_values),
810                Arc::new(Int64Array::from(vec![size as i64])),
811                Arc::new(Int64Array::from(vec![last_modified])),
812                Arc::new(stats),
813            ],
814        )
815        .unwrap();
816
817        assert_eq!(actual.record_batch(), &expected);
818    }
819
820    #[tokio::test]
821    async fn test_write_parquet() {
822        let store = Arc::new(InMemory::new());
823        let parquet_handler =
824            DefaultParquetHandler::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
825
826        let data = Box::new(ArrowEngineData::new(
827            RecordBatch::try_from_iter(vec![(
828                "a",
829                Arc::new(Int64Array::from(vec![1, 2, 3])) as Arc<dyn Array>,
830            )])
831            .unwrap(),
832        ));
833
834        let write_metadata = parquet_handler
835            .write_parquet(&Url::parse("memory:///data/").unwrap(), data, &[])
836            .await
837            .unwrap();
838
839        let DataFileMetadata {
840            file_meta:
841                ref parquet_file @ FileMeta {
842                    ref location,
843                    last_modified,
844                    size,
845                },
846            ref stats,
847        } = write_metadata;
848        let expected_location = Url::parse("memory:///data/").unwrap();
849
850        // head the object to get metadata
851        let meta = store
852            .head(&Path::from_url_path(location.path()).unwrap())
853            .await
854            .unwrap();
855        let expected_size = meta.size;
856
857        // check that last_modified is within 10s of now
858        let now: i64 = current_time_ms().unwrap();
859
860        let filename = location.path().split('/').next_back().unwrap();
861        assert_eq!(&expected_location.join(filename).unwrap(), location);
862        assert_eq!(expected_size, size);
863        assert!(now - last_modified < 10_000);
864
865        // Check numRecords from stats
866        let num_records = stats
867            .column_by_name(NUM_RECORDS)
868            .unwrap()
869            .as_any()
870            .downcast_ref::<Int64Array>()
871            .unwrap()
872            .value(0);
873        assert_eq!(num_records, 3);
874
875        // check we can read back
876        let path = Path::from_url_path(location.path()).unwrap();
877        let reader = ParquetObjectReader::new(store.clone(), path);
878        let physical_schema = ParquetRecordBatchStreamBuilder::new(reader)
879            .await
880            .unwrap()
881            .schema()
882            .clone();
883
884        let data: Vec<RecordBatch> = parquet_handler
885            .read_parquet_files(
886                slice::from_ref(parquet_file),
887                Arc::new(physical_schema.try_into_kernel().unwrap()),
888                None,
889            )
890            .unwrap()
891            .map(into_record_batch)
892            .try_collect()
893            .unwrap();
894
895        assert_eq!(data.len(), 1);
896        assert_eq!(data[0].num_rows(), 3);
897    }
898
899    #[tokio::test]
900    async fn test_disallow_non_trailing_slash() {
901        let store = Arc::new(InMemory::new());
902        let parquet_handler =
903            DefaultParquetHandler::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
904
905        let data = Box::new(ArrowEngineData::new(
906            RecordBatch::try_from_iter(vec![(
907                "a",
908                Arc::new(Int64Array::from(vec![1, 2, 3])) as Arc<dyn Array>,
909            )])
910            .unwrap(),
911        ));
912
913        assert_result_error_with_message(
914            parquet_handler
915                .write_parquet(&Url::parse("memory:///data").unwrap(), data, &[])
916                .await,
917            "Generic delta kernel error: Path must end with a trailing slash: memory:///data",
918        );
919    }
920
921    #[tokio::test]
922    async fn test_parquet_handler_trait_write() {
923        let store = Arc::new(InMemory::new());
924        let parquet_handler: Arc<dyn ParquetHandler> = Arc::new(DefaultParquetHandler::new(
925            store.clone(),
926            Arc::new(TokioBackgroundExecutor::new()),
927        ));
928
929        let engine_data: Box<dyn EngineData> = Box::new(ArrowEngineData::new(
930            RecordBatch::try_from_iter(vec![
931                (
932                    "x",
933                    Arc::new(Int64Array::from(vec![10, 20, 30])) as Arc<dyn Array>,
934                ),
935                (
936                    "y",
937                    Arc::new(Int64Array::from(vec![100, 200, 300])) as Arc<dyn Array>,
938                ),
939            ])
940            .unwrap(),
941        ));
942
943        // Create iterator with single batch
944        let data_iter: Box<dyn Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send> =
945            Box::new(std::iter::once(Ok(engine_data)));
946
947        // Test writing through the trait method
948        let file_url = Url::parse("memory:///test/data.parquet").unwrap();
949        parquet_handler
950            .write_parquet_file(file_url.clone(), data_iter)
951            .unwrap();
952
953        // Verify we can read the file back
954        let path = Path::from_url_path(file_url.path()).unwrap();
955        let metadata = store.head(&path).await.unwrap();
956        let reader = ParquetObjectReader::new(store.clone(), path);
957        let physical_schema = ParquetRecordBatchStreamBuilder::new(reader)
958            .await
959            .unwrap()
960            .schema()
961            .clone();
962
963        let file_meta = FileMeta {
964            location: file_url,
965            last_modified: 0,
966            size: metadata.size,
967        };
968
969        let data: Vec<RecordBatch> = parquet_handler
970            .read_parquet_files(
971                slice::from_ref(&file_meta),
972                Arc::new(physical_schema.try_into_kernel().unwrap()),
973                None,
974            )
975            .unwrap()
976            .map(into_record_batch)
977            .try_collect()
978            .unwrap();
979
980        assert_eq!(data.len(), 1);
981        assert_eq!(data[0].num_rows(), 3);
982        assert_eq!(data[0].num_columns(), 2);
983    }
984
985    #[tokio::test]
986    async fn test_parquet_handler_trait_write_and_read_roundtrip() {
987        let store = Arc::new(InMemory::new());
988        let parquet_handler: Arc<dyn ParquetHandler> = Arc::new(DefaultParquetHandler::new(
989            store.clone(),
990            Arc::new(TokioBackgroundExecutor::new()),
991        ));
992
993        // Create test data with all Delta-supported primitive types
994        let engine_data: Box<dyn EngineData> = Box::new(ArrowEngineData::new(
995            RecordBatch::try_from_iter(vec![
996                // Byte (i8)
997                (
998                    "byte_col",
999                    Arc::new(Int8Array::from(vec![1i8, 2, 3, 4, 5])) as Arc<dyn Array>,
1000                ),
1001                // Short (i16)
1002                (
1003                    "short_col",
1004                    Arc::new(Int16Array::from(vec![100i16, 200, 300, 400, 500])) as Arc<dyn Array>,
1005                ),
1006                // Integer (i32)
1007                (
1008                    "int_col",
1009                    Arc::new(Int32Array::from(vec![1000i32, 2000, 3000, 4000, 5000]))
1010                        as Arc<dyn Array>,
1011                ),
1012                // Long (i64)
1013                (
1014                    "long_col",
1015                    Arc::new(Int64Array::from(vec![10000i64, 20000, 30000, 40000, 50000]))
1016                        as Arc<dyn Array>,
1017                ),
1018                // Float (f32)
1019                (
1020                    "float_col",
1021                    Arc::new(Float32Array::from(vec![1.1f32, 2.2, 3.3, 4.4, 5.5]))
1022                        as Arc<dyn Array>,
1023                ),
1024                // Double (f64)
1025                (
1026                    "double_col",
1027                    Arc::new(Float64Array::from(vec![1.11f64, 2.22, 3.33, 4.44, 5.55]))
1028                        as Arc<dyn Array>,
1029                ),
1030                // Boolean
1031                (
1032                    "bool_col",
1033                    Arc::new(BooleanArray::from(vec![true, false, true, false, true]))
1034                        as Arc<dyn Array>,
1035                ),
1036                // String
1037                (
1038                    "string_col",
1039                    Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])) as Arc<dyn Array>,
1040                ),
1041                // Binary
1042                (
1043                    "binary_col",
1044                    Arc::new(BinaryArray::from_vec(vec![
1045                        b"bin1", b"bin2", b"bin3", b"bin4", b"bin5",
1046                    ])) as Arc<dyn Array>,
1047                ),
1048                // Date
1049                (
1050                    "date_col",
1051                    Arc::new(Date32Array::from(vec![18262, 18263, 18264, 18265, 18266]))
1052                        as Arc<dyn Array>, // Days since epoch (2020-01-01 onwards)
1053                ),
1054                #[cfg(feature = "nanosecond-timestamps")]
1055                // Timestamps in nanoseconds (with UTC timezone)
1056                (
1057                    "timestamp_nanos_col",
1058                    Arc::new(
1059                        TimestampNanosecondArray::from(vec![
1060                            1609459200000000000i64, // 2021-01-01 00:00:00 UTC
1061                            1609545600000000000i64,
1062                            1609632000000000000i64,
1063                            1609718400000000000i64,
1064                            1609804800000000123i64,
1065                        ])
1066                        .with_timezone("UTC"),
1067                    ) as Arc<dyn Array>,
1068                ),
1069                // Timestamp (with UTC timezone)
1070                (
1071                    "timestamp_col",
1072                    Arc::new(
1073                        TimestampMicrosecondArray::from(vec![
1074                            1609459200000000i64, // 2021-01-01 00:00:00 UTC
1075                            1609545600000000i64,
1076                            1609632000000000i64,
1077                            1609718400000000i64,
1078                            1609804800000000i64,
1079                        ])
1080                        .with_timezone("UTC"),
1081                    ) as Arc<dyn Array>,
1082                ),
1083                // TimestampNtz (without timezone)
1084                (
1085                    "timestamp_ntz_col",
1086                    Arc::new(TimestampMicrosecondArray::from(vec![
1087                        1609459200000000i64, // 2021-01-01 00:00:00
1088                        1609545600000000i64,
1089                        1609632000000000i64,
1090                        1609718400000000i64,
1091                        1609804800000000i64,
1092                    ])) as Arc<dyn Array>,
1093                ),
1094                // Decimal (precision 10, scale 2)
1095                (
1096                    "decimal_col",
1097                    Arc::new(
1098                        Decimal128Array::from(vec![12345i128, 23456, 34567, 45678, 56789])
1099                            .with_precision_and_scale(10, 2)
1100                            .unwrap(),
1101                    ) as Arc<dyn Array>,
1102                ),
1103            ])
1104            .unwrap(),
1105        ));
1106
1107        // Create iterator with single batch
1108        let data_iter: Box<dyn Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send> =
1109            Box::new(std::iter::once(Ok(engine_data)));
1110
1111        // Write the data
1112        let file_url = Url::parse("memory:///roundtrip/test.parquet").unwrap();
1113        parquet_handler
1114            .write_parquet_file(file_url.clone(), data_iter)
1115            .unwrap();
1116
1117        // Read it back
1118        let path = Path::from_url_path(file_url.path()).unwrap();
1119        let metadata = store.head(&path).await.unwrap();
1120        let reader = ParquetObjectReader::new(store.clone(), path);
1121        let physical_schema = ParquetRecordBatchStreamBuilder::new(reader)
1122            .await
1123            .unwrap()
1124            .schema()
1125            .clone();
1126
1127        let file_meta = FileMeta {
1128            location: file_url.clone(),
1129            last_modified: 0,
1130            size: metadata.size,
1131        };
1132
1133        let data: Vec<RecordBatch> = parquet_handler
1134            .read_parquet_files(
1135                slice::from_ref(&file_meta),
1136                Arc::new(physical_schema.try_into_kernel().unwrap()),
1137                None,
1138            )
1139            .unwrap()
1140            .map(into_record_batch)
1141            .try_collect()
1142            .unwrap();
1143
1144        // Verify the data
1145        assert_eq!(data.len(), 1);
1146        assert_eq!(data[0].num_rows(), 5);
1147        assert_eq!(
1148            data[0].num_columns(),
1149            13 + cfg!(feature = "nanosecond-timestamps") as usize
1150        );
1151
1152        let mut col_idx = 0;
1153
1154        // Verify byte column
1155        let byte_col = data[0]
1156            .column(col_idx)
1157            .as_any()
1158            .downcast_ref::<Int8Array>()
1159            .unwrap();
1160        assert_eq!(byte_col.values(), &[1i8, 2, 3, 4, 5]);
1161        col_idx += 1;
1162
1163        // Verify short column
1164        let short_col = data[0]
1165            .column(col_idx)
1166            .as_any()
1167            .downcast_ref::<Int16Array>()
1168            .unwrap();
1169        assert_eq!(short_col.values(), &[100i16, 200, 300, 400, 500]);
1170        col_idx += 1;
1171
1172        // Verify int column
1173        let int_col = data[0]
1174            .column(col_idx)
1175            .as_any()
1176            .downcast_ref::<Int32Array>()
1177            .unwrap();
1178        assert_eq!(int_col.values(), &[1000i32, 2000, 3000, 4000, 5000]);
1179        col_idx += 1;
1180
1181        // Verify long column
1182        let long_col = data[0]
1183            .column(col_idx)
1184            .as_any()
1185            .downcast_ref::<Int64Array>()
1186            .unwrap();
1187        assert_eq!(long_col.values(), &[10000i64, 20000, 30000, 40000, 50000]);
1188        col_idx += 1;
1189
1190        // Verify float column
1191        let float_col = data[0]
1192            .column(col_idx)
1193            .as_any()
1194            .downcast_ref::<Float32Array>()
1195            .unwrap();
1196        assert_eq!(float_col.values(), &[1.1f32, 2.2, 3.3, 4.4, 5.5]);
1197        col_idx += 1;
1198
1199        // Verify double column
1200        let double_col = data[0]
1201            .column(col_idx)
1202            .as_any()
1203            .downcast_ref::<Float64Array>()
1204            .unwrap();
1205        assert_eq!(double_col.values(), &[1.11f64, 2.22, 3.33, 4.44, 5.55]);
1206        col_idx += 1;
1207
1208        // Verify bool column
1209        let bool_col = data[0]
1210            .column(col_idx)
1211            .as_any()
1212            .downcast_ref::<BooleanArray>()
1213            .unwrap();
1214        assert!(bool_col.value(0));
1215        assert!(!bool_col.value(1));
1216        col_idx += 1;
1217
1218        // Verify string column
1219        let string_col = data[0]
1220            .column(col_idx)
1221            .as_any()
1222            .downcast_ref::<StringArray>()
1223            .unwrap();
1224        assert_eq!(string_col.value(0), "a");
1225        assert_eq!(string_col.value(4), "e");
1226        col_idx += 1;
1227
1228        // Verify binary column
1229        let binary_col = data[0]
1230            .column(col_idx)
1231            .as_any()
1232            .downcast_ref::<BinaryArray>()
1233            .unwrap();
1234        assert_eq!(binary_col.value(0), b"bin1");
1235        assert_eq!(binary_col.value(4), b"bin5");
1236        col_idx += 1;
1237
1238        // Verify date column
1239        let date_col = data[0]
1240            .column(col_idx)
1241            .as_any()
1242            .downcast_ref::<Date32Array>()
1243            .unwrap();
1244        assert_eq!(date_col.values(), &[18262, 18263, 18264, 18265, 18266]);
1245        col_idx += 1;
1246
1247        #[cfg(feature = "nanosecond-timestamps")]
1248        {
1249            // Verify timestamp column (with UTC timezone)
1250            let timestamp_col = data[0]
1251                .column(col_idx)
1252                .as_any()
1253                .downcast_ref::<TimestampNanosecondArray>()
1254                .unwrap();
1255            assert_eq!(timestamp_col.value(0), 1609459200000000000i64);
1256            assert_eq!(timestamp_col.value(4), 1609804800000000123i64);
1257            assert!(timestamp_col
1258                .timezone()
1259                .is_some_and(|tz| tz.eq_ignore_ascii_case("utc")));
1260            col_idx += 1;
1261        }
1262
1263        // Verify timestamp column (with UTC timezone)
1264        let timestamp_col = data[0]
1265            .column(col_idx)
1266            .as_any()
1267            .downcast_ref::<TimestampMicrosecondArray>()
1268            .unwrap();
1269        assert_eq!(timestamp_col.value(0), 1609459200000000i64);
1270        assert_eq!(timestamp_col.value(4), 1609804800000000i64);
1271        assert!(timestamp_col
1272            .timezone()
1273            .is_some_and(|tz| tz.eq_ignore_ascii_case("utc")));
1274        col_idx += 1;
1275
1276        // Verify timestamp_ntz column (without timezone)
1277        let timestamp_ntz_col = data[0]
1278            .column(col_idx)
1279            .as_any()
1280            .downcast_ref::<TimestampMicrosecondArray>()
1281            .unwrap();
1282        assert_eq!(timestamp_ntz_col.value(0), 1609459200000000i64);
1283        assert_eq!(timestamp_ntz_col.value(4), 1609804800000000i64);
1284        assert!(timestamp_ntz_col.timezone().is_none());
1285        col_idx += 1;
1286
1287        // Verify decimal column
1288        let decimal_col = data[0]
1289            .column(col_idx)
1290            .as_any()
1291            .downcast_ref::<Decimal128Array>()
1292            .unwrap();
1293        assert_eq!(decimal_col.value(0), 12345i128);
1294        assert_eq!(decimal_col.value(4), 56789i128);
1295        assert_eq!(decimal_col.precision(), 10);
1296        assert_eq!(decimal_col.scale(), 2);
1297    }
1298
1299    /// Test that field IDs are accessible via ColumnMetadataKey::ParquetFieldId as documented.
1300    ///
1301    /// Per trait definitions in lib.rs, field IDs should be accessible via
1302    /// StructField::get_config_value with ColumnMetadataKey::ParquetFieldId.
1303    #[test]
1304    fn test_parquet_footer_read_with_field_id() {
1305        // Write parquet file with field ID
1306        let field = Field::new("value", ArrowDataType::Int64, false).with_metadata(HashMap::from(
1307            [(PARQUET_FIELD_ID_META_KEY.to_string(), "42".to_string())],
1308        ));
1309        let arrow_schema = Arc::new(ArrowSchema::new(vec![field]));
1310
1311        let temp_dir = tempfile::tempdir().unwrap();
1312        let file_path = temp_dir.path().join("field_id_test.parquet");
1313        let batch = RecordBatch::try_new(
1314            arrow_schema.clone(),
1315            vec![Arc::new(Int64Array::from(vec![1, 2, 3]))],
1316        )
1317        .unwrap();
1318
1319        let file = std::fs::File::create(&file_path).unwrap();
1320        let mut writer = ArrowWriter::try_new(file, arrow_schema, None).unwrap();
1321        writer.write(&batch).unwrap();
1322        writer.close().unwrap();
1323
1324        // Read footer and verify field ID accessibility
1325        let store = Arc::new(LocalFileSystem::new());
1326        let handler = DefaultParquetHandler::new(store, Arc::new(TokioBackgroundExecutor::new()));
1327        let file_size = std::fs::metadata(&file_path).unwrap().len();
1328        let file_meta = FileMeta {
1329            location: Url::from_file_path(&file_path).unwrap(),
1330            last_modified: 0,
1331            size: file_size,
1332        };
1333
1334        let footer = handler.read_parquet_footer(&file_meta).unwrap();
1335        let field = footer
1336            .schema
1337            .fields()
1338            .find(|f| f.name() == "value")
1339            .unwrap();
1340
1341        // Field ID is transformed to kernel key when reading. arrow->kernel parses the
1342        // `PARQUET:field_id` string back into kernel's canonical `MetadataValue::Number(i64)`.
1343        assert_eq!(
1344            field
1345                .metadata()
1346                .get(ColumnMetadataKey::ParquetFieldId.as_ref()),
1347            Some(&MetadataValue::Number(42))
1348        );
1349
1350        // Field ID should be accessible via documented API
1351        let field_id = field.get_config_value(&ColumnMetadataKey::ParquetFieldId)
1352            .expect("Field ID should be accessible via ColumnMetadataKey::ParquetFieldId per lib.rs:836-837");
1353
1354        match field_id {
1355            crate::schema::MetadataValue::String(id) => assert_eq!(id, "42"),
1356            crate::schema::MetadataValue::Number(id) => assert_eq!(*id, 42),
1357            other => panic!("Expected String or Number, got {other:?}"),
1358        }
1359    }
1360
1361    /// Test that columns are matched by field ID when column names differ.
1362    ///
1363    /// Per lib.rs:676-680, field IDs (via [`ColumnMetadataKey::ParquetFieldId`]) should take
1364    /// precedence over field names for column matching.
1365    ///
1366    /// [`ColumnMetadataKey::ParquetFieldId`]: crate::schema::ColumnMetadataKey::ParquetFieldId
1367    #[test]
1368    fn test_read_parquet_with_field_id_matching() {
1369        // Write parquet with field IDs using PARQUET_FIELD_ID_META_KEY (Parquet's native key)
1370        // The kernel will transform these to parquet.field.id when reading
1371        let fields = vec![
1372            Field::new("id", ArrowDataType::Int64, false).with_metadata(HashMap::from([(
1373                PARQUET_FIELD_ID_META_KEY.to_string(),
1374                "1".to_string(),
1375            )])),
1376            Field::new("name", ArrowDataType::Utf8, false).with_metadata(HashMap::from([(
1377                PARQUET_FIELD_ID_META_KEY.to_string(),
1378                "2".to_string(),
1379            )])),
1380        ];
1381        let arrow_schema = Arc::new(ArrowSchema::new(fields));
1382
1383        let temp_dir = tempfile::tempdir().unwrap();
1384        let file_path = temp_dir.path().join("field_id_matching.parquet");
1385        let batch = RecordBatch::try_new(
1386            arrow_schema.clone(),
1387            vec![
1388                Arc::new(Int64Array::from(vec![1, 2, 3])),
1389                Arc::new(StringArray::from(vec!["alice", "bob", "charlie"])),
1390            ],
1391        )
1392        .unwrap();
1393
1394        let file = std::fs::File::create(&file_path).unwrap();
1395        let mut writer = ArrowWriter::try_new(file, arrow_schema, None).unwrap();
1396        writer.write(&batch).unwrap();
1397        writer.close().unwrap();
1398
1399        // Create kernel schema with DIFFERENT names but SAME field IDs
1400        let kernel_schema = Arc::new(
1401            StructType::try_new(vec![
1402                StructField::new("user_id", crate::schema::DataType::LONG, false).with_metadata([
1403                    (
1404                        ColumnMetadataKey::ParquetFieldId.as_ref(),
1405                        MetadataValue::Number(1),
1406                    ),
1407                ]),
1408                StructField::new("user_name", crate::schema::DataType::STRING, false)
1409                    .with_metadata([(
1410                        ColumnMetadataKey::ParquetFieldId.as_ref(),
1411                        MetadataValue::Number(2),
1412                    )]),
1413            ])
1414            .unwrap(),
1415        );
1416
1417        // Read using kernel schema with different column names
1418        let store = Arc::new(LocalFileSystem::new());
1419        let handler = DefaultParquetHandler::new(store, Arc::new(TokioBackgroundExecutor::new()));
1420        let file_meta = FileMeta {
1421            location: Url::from_file_path(&file_path).unwrap(),
1422            last_modified: 0,
1423            size: std::fs::metadata(&file_path).unwrap().len(),
1424        };
1425
1426        // Should successfully match by field ID despite different names
1427        let data: Vec<RecordBatch> = handler
1428            .read_parquet_files(slice::from_ref(&file_meta), kernel_schema, None)
1429            .unwrap()
1430            .map(into_record_batch)
1431            .try_collect()
1432            .unwrap();
1433
1434        // Verify data was correctly matched by field ID
1435        assert_eq!(data.len(), 1);
1436        let batch = &data[0];
1437
1438        // Verify columns were renamed to match the kernel schema (the names from the parquet
1439        // file's schema are discarded; the matching agreed on field IDs only).
1440        let schema = batch.schema();
1441        assert_eq!(schema.field(0).name(), "user_id");
1442        assert_eq!(schema.field(1).name(), "user_name");
1443
1444        let id_col = batch
1445            .column(0)
1446            .as_any()
1447            .downcast_ref::<Int64Array>()
1448            .unwrap();
1449        assert_eq!(id_col.values(), &[1, 2, 3], "Should match by field ID 1");
1450
1451        let name_col = batch
1452            .column(1)
1453            .as_any()
1454            .downcast_ref::<StringArray>()
1455            .unwrap();
1456        assert_eq!(name_col.value(0), "alice", "Should match by field ID 2");
1457        assert_eq!(name_col.value(1), "bob");
1458        assert_eq!(name_col.value(2), "charlie");
1459    }
1460
1461    // Verifies that write_parquet (the internal stats-collecting path) does not embed the Arrow
1462    // IPC schema in the Parquet file metadata.
1463    #[tokio::test]
1464    async fn write_parquet_omits_arrow_schema_metadata() {
1465        let store = Arc::new(InMemory::new());
1466        let parquet_handler =
1467            DefaultParquetHandler::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
1468
1469        let data = Box::new(ArrowEngineData::new(
1470            RecordBatch::try_from_iter(vec![(
1471                "a",
1472                Arc::new(Int64Array::from(vec![1, 2, 3])) as Arc<dyn Array>,
1473            )])
1474            .unwrap(),
1475        ));
1476        let metadata = parquet_handler
1477            .write_parquet(&Url::parse("memory:///data/").unwrap(), data, &[])
1478            .await
1479            .unwrap();
1480
1481        let path = Path::from_url_path(metadata.file_meta.location.path()).unwrap();
1482        let reader = ParquetObjectReader::new(store, path);
1483        let builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap();
1484        let kv = builder.metadata().file_metadata().key_value_metadata();
1485        let has = kv
1486            .map(|kv| kv.iter().any(|e| e.key == ARROW_SCHEMA_META_KEY))
1487            .unwrap_or(false);
1488        assert!(
1489            !has,
1490            "Parquet file should not contain embedded Arrow schema metadata"
1491        );
1492    }
1493
1494    #[tokio::test]
1495    async fn write_parquet_file_creates_parent_directories() {
1496        // GIVEN a file path whose parent directories do not exist
1497        let temp_dir = tempfile::tempdir().unwrap();
1498        let nested_path = temp_dir.path().join("a/b/c/output.parquet");
1499        assert!(!nested_path.parent().unwrap().exists());
1500
1501        let store = Arc::new(LocalFileSystem::new());
1502        let parquet_handler: Arc<dyn ParquetHandler> = Arc::new(DefaultParquetHandler::new(
1503            store.clone(),
1504            Arc::new(TokioBackgroundExecutor::new()),
1505        ));
1506
1507        let engine_data: Box<dyn EngineData> = Box::new(ArrowEngineData::new(
1508            RecordBatch::try_from_iter(vec![(
1509                "x",
1510                Arc::new(Int64Array::from(vec![1, 2, 3])) as Arc<dyn Array>,
1511            )])
1512            .unwrap(),
1513        ));
1514        let data_iter: Box<dyn Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send> =
1515            Box::new(std::iter::once(Ok(engine_data)));
1516
1517        // WHEN we write a parquet file to that path
1518        let file_url = Url::from_file_path(&nested_path).unwrap();
1519        parquet_handler
1520            .write_parquet_file(file_url.clone(), data_iter)
1521            .unwrap();
1522
1523        // THEN the file is created and contains the expected data
1524        assert!(nested_path.exists());
1525
1526        let path = Path::from_url_path(file_url.path()).unwrap();
1527        let reader = ParquetObjectReader::new(store.clone(), path);
1528        let batches: Vec<RecordBatch> = ParquetRecordBatchStreamBuilder::new(reader)
1529            .await
1530            .unwrap()
1531            .build()
1532            .unwrap()
1533            .try_collect()
1534            .await
1535            .unwrap();
1536        assert_eq!(batches.len(), 1);
1537        let col = batches[0]
1538            .column(0)
1539            .as_any()
1540            .downcast_ref::<Int64Array>()
1541            .unwrap();
1542        assert_eq!(col.values(), &[1, 2, 3]);
1543    }
1544
1545    // === ParquetHandler contract tests ===
1546    //
1547    // These call the shared contract helpers in `engine::tests` against `DefaultParquetHandler`
1548    // (the matching `SyncParquetHandler` invocations live in `engine/sync/parquet.rs`).
1549
1550    fn default_parquet_handler() -> DefaultParquetHandler<TokioBackgroundExecutor> {
1551        DefaultParquetHandler::new(
1552            Arc::new(LocalFileSystem::new()),
1553            Arc::new(TokioBackgroundExecutor::new()),
1554        )
1555    }
1556
1557    #[test]
1558    fn parquet_handler_reads_footer() {
1559        crate::engine::tests::test_parquet_handler_reads_footer(&default_parquet_handler());
1560    }
1561
1562    #[test]
1563    fn parquet_handler_footer_errors_on_missing_file() {
1564        crate::engine::tests::test_parquet_handler_footer_errors_on_missing_file(
1565            &default_parquet_handler(),
1566        );
1567    }
1568
1569    #[test]
1570    fn parquet_handler_footer_preserves_field_ids() {
1571        crate::engine::tests::test_parquet_handler_footer_preserves_field_ids(
1572            &default_parquet_handler(),
1573        );
1574    }
1575
1576    #[test]
1577    fn parquet_handler_write_always_overwrites() {
1578        crate::engine::tests::test_parquet_handler_write_always_overwrites(
1579            &default_parquet_handler(),
1580        );
1581    }
1582
1583    #[test]
1584    fn parquet_handler_write_omits_arrow_schema() {
1585        crate::engine::tests::test_parquet_handler_write_omits_arrow_schema(
1586            &default_parquet_handler(),
1587        );
1588    }
1589
1590    #[test]
1591    fn parquet_handler_reads_file_with_arrow_schema() {
1592        crate::engine::tests::test_parquet_handler_reads_file_with_arrow_schema(
1593            &default_parquet_handler(),
1594        );
1595    }
1596}
buoyant_kernel/engine/default/parquet.rs

buoyant_kernel/engine/default/
parquet.rs
