Skip to main content

iceberg/arrow/reader/
pipeline.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! The main `ArrowReader` pipeline: reading a stream of `FileScanTask`s,
19//! opening Parquet files and resolving schemas, then wiring projection,
20//! predicates, row-group / row selection, and delete handling into a stream
21//! of transformed Arrow `RecordBatch`es.
22
23use std::borrow::Cow;
24use std::collections::HashMap;
25use std::sync::Arc;
26use std::sync::atomic::AtomicU64;
27
28use arrow_schema::{DataType, Field};
29use futures::{StreamExt, TryStreamExt};
30use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions};
31use parquet::arrow::{
32    PARQUET_FIELD_ID_META_KEY, ParquetRecordBatchStreamBuilder, ProjectionMask, RowNumber,
33};
34use parquet::encryption::decrypt::FileDecryptionProperties;
35
36use super::row_lineage::synthesize_row_id_column;
37use super::{
38    ArrowFileReader, ArrowReader, ParquetReadOptions, add_fallback_field_ids_to_arrow_schema,
39    apply_name_mapping_to_arrow_schema, find_leaf_by_field_id,
40};
41use crate::arrow::build_partition_constant;
42use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader;
43use crate::arrow::int96::coerce_int96_timestamps;
44use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder;
45use crate::arrow::scan_metrics::{CountingFileRead, ScanMetrics, ScanResult};
46use crate::encryption::StandardKeyMetadata;
47use crate::error::{Result, invalid_data};
48use crate::expr::BoundPredicate;
49use crate::expr::visitors::bloom_filter_evaluator::{
50    BloomFilterEvaluator, ColumnBloomFilter, collect_bloom_filter_field_ids,
51};
52use crate::io::{FileIO, FileMetadata, FileRead};
53use crate::metadata_columns::{
54    RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER, RESERVED_COL_NAME_POS,
55    RESERVED_COL_NAME_ROW_ID, RESERVED_FIELD_ID_FILE,
56    RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER, RESERVED_FIELD_ID_PARTITION,
57    RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID, RESERVED_FIELD_ID_SPEC_ID, is_metadata_field,
58};
59use crate::runtime::Runtime;
60use crate::scan::{ArrowRecordBatchStream, FileScanTask, FileScanTaskStream};
61use crate::spec::{Datum, PartitionSpec, Struct};
62use crate::{Error, ErrorKind};
63
64impl ArrowReader {
65    /// Take a stream of FileScanTasks and reads all the files.
66    /// Returns a [`ScanResult`] containing the record batch stream and scan metrics.
67    pub fn read(self, tasks: FileScanTaskStream) -> Result<ScanResult> {
68        let concurrency_limit_data_files = self.concurrency_limit_data_files;
69        let scan_metrics = ScanMetrics::new();
70
71        let task_reader = FileScanTaskReader {
72            batch_size: self.batch_size,
73            file_io: self.file_io,
74            delete_file_loader: self
75                .delete_file_loader
76                .with_scan_metrics(scan_metrics.clone()),
77            row_group_filtering_enabled: self.row_group_filtering_enabled,
78            row_selection_enabled: self.row_selection_enabled,
79            bloom_filter_enabled: self.bloom_filter_enabled,
80            parquet_read_options: self.parquet_read_options,
81            scan_metrics: scan_metrics.clone(),
82            runtime: self.runtime,
83        };
84
85        // Fast-path for single concurrency to avoid overhead of try_flatten_unordered
86        let stream: ArrowRecordBatchStream = if concurrency_limit_data_files == 1 {
87            Box::pin(
88                tasks
89                    .and_then(move |task| task_reader.clone().process(task))
90                    .map_err(|err| {
91                        Error::new(ErrorKind::Unexpected, "file scan task generate failed")
92                            .with_source(err)
93                    })
94                    .try_flatten(),
95            )
96        } else {
97            Box::pin(
98                tasks
99                    .map_ok(move |task| task_reader.clone().process(task))
100                    .map_err(|err| {
101                        Error::new(ErrorKind::Unexpected, "file scan task generate failed")
102                            .with_source(err)
103                    })
104                    .try_buffer_unordered(concurrency_limit_data_files)
105                    .try_flatten_unordered(concurrency_limit_data_files),
106            )
107        };
108
109        Ok(ScanResult::new(stream, scan_metrics))
110    }
111}
112
113// Metadata columns synthesized without reading any data column, so a projection of only
114// these (including the empty projection, e.g. `COUNT(*)`) can be pruned to zero data
115// columns. Narrower than `is_metadata_field`, which also matches `_deleted` -- excluded here
116// because it has no synthesis handler.
117const PRUNABLE_METADATA_FIELDS: &[i32] = &[
118    RESERVED_FIELD_ID_FILE,
119    RESERVED_FIELD_ID_SPEC_ID,
120    RESERVED_FIELD_ID_PARTITION,
121    RESERVED_FIELD_ID_POS,
122    RESERVED_FIELD_ID_ROW_ID,
123    RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
124];
125
126/// Per-scan state for processing [`FileScanTask`]s. Created once per
127/// [`ArrowReader::read`] call and cloned per task.
128#[derive(Clone)]
129struct FileScanTaskReader {
130    batch_size: Option<usize>,
131    file_io: FileIO,
132    delete_file_loader: CachingDeleteFileLoader,
133    row_group_filtering_enabled: bool,
134    row_selection_enabled: bool,
135    bloom_filter_enabled: bool,
136    parquet_read_options: ParquetReadOptions,
137    scan_metrics: ScanMetrics,
138    runtime: Runtime,
139}
140
141// Borrow an existing selection, or enumerate all groups for a whole-file scan.
142fn candidate_row_groups(
143    selected_row_groups: Option<&[usize]>,
144    num_row_groups: usize,
145) -> Cow<'_, [usize]> {
146    match selected_row_groups {
147        Some(indices) => Cow::Borrowed(indices),
148        None => Cow::Owned((0..num_row_groups).collect()),
149    }
150}
151
152impl FileScanTaskReader {
153    async fn process(self, task: FileScanTask) -> Result<ArrowRecordBatchStream> {
154        let should_load_page_index = (self.row_selection_enabled && task.predicate().is_some())
155            || !task.deletes().is_empty();
156        let mut parquet_read_options = self.parquet_read_options;
157        parquet_read_options.preload_page_index = should_load_page_index;
158
159        let delete_filter_rx = self
160            .delete_file_loader
161            .load_deletes(task.deletes(), task.schema_ref());
162
163        // `try_flatten_unordered` stops polling the buffered opens once it holds
164        // `concurrency_limit_data_files` streams, so spawn the open (#3276).
165        // Footer and page-index decode also run on the IO runtime.
166        // Aborts on drop.
167        let path = task.data_file_path().to_string();
168        let file_io = self.file_io.clone();
169        let file_size_in_bytes = task.file_size_in_bytes();
170        let bytes_read = Arc::clone(self.scan_metrics.bytes_read_counter());
171        let key_metadata = task.key_metadata().map(|km| km.to_vec());
172        let opening = self
173            .runtime
174            .io()
175            .spawn(async move {
176                ArrowReader::open_parquet_file(
177                    &path,
178                    &file_io,
179                    file_size_in_bytes,
180                    parquet_read_options,
181                    &bytes_read,
182                    key_metadata.as_deref(),
183                )
184                .await
185            })
186            .abort_on_drop();
187        let (parquet_file_reader, arrow_metadata) = opening.await??;
188
189        // Check if Parquet file has embedded field IDs
190        // Corresponds to Java's ParquetSchemaUtil.hasIds()
191        // Reference: parquet/src/main/java/org/apache/iceberg/parquet/ParquetSchemaUtil.java:118
192        let missing_field_ids = arrow_metadata
193            .schema()
194            .fields()
195            .iter()
196            .next()
197            .is_some_and(|f| f.metadata().get(PARQUET_FIELD_ID_META_KEY).is_none());
198
199        // Position-based fallback applies only when the file has no embedded field IDs
200        // AND no name mapping is available. With a name mapping, field IDs are assigned
201        // to the Arrow schema below, and projection/predicate planning must use them
202        // (see #2403).
203        let use_position_fallback = missing_field_ids && task.name_mapping().is_none();
204
205        let project_pos = task.project_field_ids().contains(&RESERVED_FIELD_ID_POS);
206        let project_row_id = task.project_field_ids().contains(&RESERVED_FIELD_ID_ROW_ID);
207
208        // The RowNumber virtual column materializes `_pos`. It is also the per-row
209        // positional fallback for `_row_id` (`first_row_id + pos`), so add it whenever
210        // `_row_id` is synthesized. A null `first_row_id` nulls the whole `_row_id`
211        // column, so nothing is synthesized and the column is not needed.
212        let need_row_number = project_pos || (project_row_id && task.first_row_id().is_some());
213
214        let field_ids = task.project_field_ids();
215        // Vacuously true for an empty projection (e.g. `COUNT(*)`), so no `is_empty()` guard.
216        let no_data_columns_projected = field_ids
217            .iter()
218            .all(|id| PRUNABLE_METADATA_FIELDS.contains(id));
219
220        let install_row_number = need_row_number || no_data_columns_projected;
221
222        let arrow_metadata = Self::configure_arrow_reader_metadata(
223            arrow_metadata,
224            &task,
225            missing_field_ids,
226            install_row_number,
227        )?;
228
229        // Build the stream reader, reusing the already-opened file reader
230        let mut record_batch_stream_builder =
231            ParquetRecordBatchStreamBuilder::new_with_metadata(parquet_file_reader, arrow_metadata);
232
233        // Whether the file physically carries the `_last_updated_sequence_number` column
234        // (some engines, e.g. Iceberg Java on rewrite, write it per-row), resolved by its
235        // embedded field id against the Parquet schema.
236        let project_last_updated_seq = task
237            .project_field_ids()
238            .contains(&RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER);
239
240        // Parquet leaf index of the physically-stored column, resolved by its embedded
241        // reserved field id. `find_leaf_by_field_id` tolerates id-less leaves (e.g. a
242        // Variant column's internal metadata/value leaves, which the spec requires to have
243        // no id), so an unprojected variant alongside a metadata column with correct ID does
244        // not hide it.
245        let phys_last_updated_seq_leaf = if project_last_updated_seq {
246            find_leaf_by_field_id(
247                record_batch_stream_builder.parquet_schema(),
248                RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
249            )
250        } else {
251            None
252        };
253
254        // Present by name but not by the embedded id (only meaningful when no by-id column
255        // was found). An unthreadable shape we reject rather than coalesce incorrectly.
256        let last_updated_seq_present_by_name_only = project_last_updated_seq
257            && phys_last_updated_seq_leaf.is_none()
258            && record_batch_stream_builder
259                .schema()
260                .fields()
261                .iter()
262                .any(|f| f.name() == RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER);
263
264        // Read the physical column only when first_row_id is set (with a data sequence
265        // number to fall back to). A null first_row_id drops the leaf and nulls the whole
266        // column below, discarding any per-row values the file carries -- matching Java
267        // (`ValueReaders.lastUpdated` nulls when the base row id is null).
268        let coalesce_last_updated_seq_leaf = phys_last_updated_seq_leaf
269            .filter(|_| task.first_row_id().is_some() && task.data_sequence_number().is_some());
270
271        let phys_row_id_leaf = if project_row_id {
272            find_leaf_by_field_id(
273                record_batch_stream_builder.parquet_schema(),
274                RESERVED_FIELD_ID_ROW_ID,
275            )
276        } else {
277            None
278        };
279
280        // A column named `_row_id` that carries no embedded field id, in a file that DOES
281        // use embedded ids -- an unthreadable physically-stored `_row_id`, rejected below.
282        // Gated on `!use_position_fallback`: under positional fallback every column lacks an
283        // embedded id and synthetic ids are assigned by position, so a user column that
284        // happens to be named `_row_id` is real data, not a reserved metadata column.
285        let row_id_present_by_name_only = project_row_id
286            && !use_position_fallback
287            && phys_row_id_leaf.is_none()
288            && record_batch_stream_builder
289                .schema()
290                .fields()
291                .iter()
292                .any(|f| f.name() == RESERVED_COL_NAME_ROW_ID);
293
294        // Read the physical column only when first_row_id is set. A null first_row_id
295        // nulls the whole column below (matching Java `ValueReaders.rowIds`).
296        let coalesce_row_id_leaf = phys_row_id_leaf.filter(|_| task.first_row_id().is_some());
297
298        // Filter out metadata fields for Parquet projection (they don't exist in files)
299        let project_field_ids_without_metadata: Vec<i32> = task
300            .project_field_ids()
301            .iter()
302            .filter(|&&id| !is_metadata_field(id))
303            .copied()
304            .collect();
305
306        // Create projection mask based on field IDs
307        // - If file has embedded IDs: field-ID-based projection
308        // - If name mapping applied: field-ID-based projection using the IDs the name
309        //   mapping assigned to the Arrow schema
310        // - Otherwise: position-based fallback projection
311        let mut projection_mask = ArrowReader::get_arrow_projection_mask(
312            &project_field_ids_without_metadata,
313            task.schema(),
314            record_batch_stream_builder.parquet_schema(),
315            record_batch_stream_builder.schema(),
316            use_position_fallback, // Whether to use position-based (true) or field-ID-based (false) projection
317        )?;
318
319        // A projection with no data columns leaves `project_field_ids_without_metadata`
320        // empty, which `get_arrow_projection_mask` maps to "read all columns". Downgrade that
321        // to "read no data columns": `install_row_number` put the RowNumber virtual column on
322        // every such projection as a row-count source independent of the data columns, so the
323        // count survives with zero data columns read.
324        //
325        // This runs BEFORE the union so any physical metadata leaf is added onto a `none`
326        // base, pruning the read to just that leaf (`union` with an `all` base stays `all`).
327        if project_field_ids_without_metadata.is_empty() && install_row_number {
328            projection_mask =
329                ProjectionMask::none(record_batch_stream_builder.parquet_schema().num_columns());
330        }
331
332        // Union in the physical leaves of any metadata columns we will coalesce. Their
333        // reserved field ids are not in the task schema, so they can't be requested through
334        // `get_arrow_projection_mask` (which resolves ids against the task schema); add
335        // their Parquet leaves directly.
336        for leaf in [coalesce_last_updated_seq_leaf, coalesce_row_id_leaf]
337            .into_iter()
338            .flatten()
339        {
340            let phys_mask =
341                ProjectionMask::leaves(record_batch_stream_builder.parquet_schema(), vec![leaf]);
342            projection_mask.union(&phys_mask);
343        }
344
345        record_batch_stream_builder =
346            record_batch_stream_builder.with_projection(projection_mask.clone());
347
348        // RecordBatchTransformer performs any transformations required on the RecordBatches
349        // that come back from the file, such as type promotion, default column insertion,
350        // column re-ordering, partition constants, and virtual field addition (like _file)
351        let mut record_batch_transformer_builder =
352            RecordBatchTransformerBuilder::new(task.schema_ref(), task.project_field_ids());
353
354        // Add the _file metadata column if it's in the projected fields
355        if task.project_field_ids().contains(&RESERVED_FIELD_ID_FILE) {
356            let file_datum = Datum::string(task.data_file_path().to_string());
357            record_batch_transformer_builder =
358                record_batch_transformer_builder.with_constant(RESERVED_FIELD_ID_FILE, file_datum);
359        }
360
361        if task
362            .project_field_ids()
363            .contains(&RESERVED_FIELD_ID_SPEC_ID)
364        {
365            let partition_spec = task
366                .partition_spec()
367                .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Partition spec is missing"))?;
368
369            let spec_id_datum = Datum::int(partition_spec.spec_id());
370            record_batch_transformer_builder = record_batch_transformer_builder
371                .with_constant(RESERVED_FIELD_ID_SPEC_ID, spec_id_datum);
372        }
373
374        if project_last_updated_seq {
375            // Materialize the column, gated on the data file's `first_row_id`. Java gates
376            // it this way (`ValueReaders.lastUpdated` returns nulls when the base row id is
377            // null); the spec itself only says the column is assigned the manifest entry's
378            // sequence number on read.
379            record_batch_transformer_builder =
380                match (task.first_row_id(), task.data_sequence_number()) {
381                    (Some(_), Some(seq)) => {
382                        let datum = Datum::long(seq);
383                        if coalesce_last_updated_seq_leaf.is_some() {
384                            // The file physically carries the column: read the per-row value,
385                            // falling back to the data sequence number only where null.
386                            record_batch_transformer_builder.with_coalesced_last_updated_seq_column(
387                                RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
388                                datum,
389                            )
390                        } else if last_updated_seq_present_by_name_only {
391                            // Present by name but without the embedded field id (name mapping /
392                            // positional fallback). The transformer keys the source column by
393                            // field id, so we can't thread it; no real writer produces this, so
394                            // reject loudly rather than silently overwrite with the constant.
395                            // Arm-local by design: only this arm reads the physical column, so
396                            // only here can a name-only column defeat us. The `(None, _)` arm
397                            // nulls the column without reading it, so it needs no such guard.
398                            return Err(Error::new(
399                                ErrorKind::FeatureUnsupported,
400                                "Reading a physically-stored _last_updated_sequence_number column \
401                             without an embedded field id is not supported",
402                            ));
403                        } else {
404                            // Column absent: derive it from the data sequence number.
405                            record_batch_transformer_builder.with_constant(
406                                RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
407                                datum,
408                            )
409                        }
410                    }
411                    // Null first_row_id (v1/v2, or a pre-upgrade v3 snapshot): the column is null.
412                    (None, _) => record_batch_transformer_builder.with_null_metadata_column(
413                        RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
414                    )?,
415                    // first_row_id present but no data sequence number: after manifest
416                    // inheritance a committed entry always has one, so this is a malformed
417                    // manifest rather than a legitimate null.
418                    (Some(_), None) => {
419                        return Err(invalid_data!(
420                            "Data file {} has a first_row_id but no data sequence number",
421                            task.data_file_path()
422                        ));
423                    }
424                };
425        }
426
427        if project_row_id {
428            // A name-only physical `_row_id` can't be threaded (synthesis keys the leaf by
429            // its reserved field id). Reject only when `first_row_id` is set; otherwise the
430            // leaf is never read and `_row_id` is nulled out downstream (matching Java
431            // `ValueReaders.rowIds`), so a pre-v3 file with a user column named `_row_id`
432            // reads back as null rather than erroring.
433            if task.first_row_id().is_some() && row_id_present_by_name_only {
434                return Err(Error::new(
435                    ErrorKind::FeatureUnsupported,
436                    "Reading a physically-stored _row_id column without an embedded field id \
437                     is not supported",
438                ));
439            }
440
441            // `_row_id` is synthesized downstream over the record-batch stream (see
442            // `row_lineage::synthesize_row_id_column`); the transformer only passes the
443            // resulting column through, like `_pos`.
444            record_batch_transformer_builder =
445                record_batch_transformer_builder.with_virtual_field(RESERVED_FIELD_ID_ROW_ID);
446        }
447
448        if let (Some(partition_spec), Some(partition_data)) =
449            (task.partition_spec(), task.partition())
450        {
451            record_batch_transformer_builder = record_batch_transformer_builder
452                .with_partition(Arc::clone(partition_spec), partition_data.clone())?;
453        }
454
455        if project_pos {
456            record_batch_transformer_builder =
457                record_batch_transformer_builder.with_virtual_field(RESERVED_FIELD_ID_POS);
458        }
459
460        // Add the _partition metadata struct column if it's in the projected fields.
461        // Computed lazily here at read time from the unified partition type + task's spec + data.
462        if task
463            .project_field_ids()
464            .contains(&RESERVED_FIELD_ID_PARTITION)
465            && let Some(unified_type) = task.unified_partition_type()
466        {
467            let (spec, partition_data) = match (task.partition_spec(), task.partition()) {
468                (Some(spec), Some(data)) => (Arc::clone(spec), data.clone()),
469                // A missing spec/data is only acceptable when there are no partition
470                // fields to fill (unpartitioned table). If the unified type has fields
471                // but we lack a spec or data, the task is inconsistent and we cannot
472                // build the _partition column.
473                _ if unified_type.fields().is_empty() => {
474                    (Arc::new(PartitionSpec::unpartition_spec()), Struct::empty())
475                }
476                _ => {
477                    return Err(Error::new(
478                        ErrorKind::Unexpected,
479                        "cannot build _partition column: unified partition type has fields \
480                         but the scan task is missing its partition spec or data",
481                    ));
482                }
483            };
484            let constant = build_partition_constant(unified_type, &spec, &partition_data)?;
485            record_batch_transformer_builder =
486                record_batch_transformer_builder.with_partition_constant(constant);
487        }
488
489        let mut record_batch_transformer = record_batch_transformer_builder.build();
490
491        if let Some(batch_size) = self.batch_size {
492            record_batch_stream_builder = record_batch_stream_builder.with_batch_size(batch_size);
493        }
494
495        let delete_filter = delete_filter_rx.await.unwrap()?;
496        let delete_predicate = delete_filter.build_equality_delete_predicate(&task).await?;
497
498        // In addition to the optional predicate supplied in the `FileScanTask`,
499        // we also have an optional predicate resulting from equality delete files.
500        // If both are present, we logical-AND them together to form a single filter
501        // predicate that we can pass to the `RecordBatchStreamBuilder`.
502        let final_predicate = match (task.predicate(), delete_predicate) {
503            (None, None) => None,
504            (Some(predicate), None) => Some(predicate.clone()),
505            (None, Some(ref predicate)) => Some(predicate.clone()),
506            (Some(filter_predicate), Some(delete_predicate)) => {
507                Some(filter_predicate.clone().and(delete_predicate))
508            }
509        };
510
511        // There are three possible sources for potential lists of selected RowGroup indices,
512        // and two for `RowSelection`s.
513        // Selected RowGroup index lists can come from three sources:
514        //   * When task.start and task.length specify a byte range (file splitting);
515        //   * When there are equality delete files that are applicable;
516        //   * When there is a scan predicate and row_group_filtering_enabled = true.
517        // `RowSelection`s can be created in either or both of the following cases:
518        //   * When there are positional delete files that are applicable;
519        //   * When there is a scan predicate and row_selection_enabled = true
520        // Note that row group filtering from predicates only happens when
521        // there is a scan predicate AND row_group_filtering_enabled = true,
522        // but we perform row selection filtering if there are applicable
523        // equality delete files OR (there is a scan predicate AND row_selection_enabled),
524        // since the only implemented method of applying positional deletes is
525        // by using a `RowSelection`.
526        let mut selected_row_group_indices = None;
527        let mut row_selection = None;
528
529        // Filter row groups based on byte range from task.start and task.length.
530        // If both start and length are 0, read the entire file (backwards compatibility).
531        if task.start() != 0 || task.length() != 0 {
532            let byte_range_filtered_row_groups = ArrowReader::filter_row_groups_by_byte_range(
533                record_batch_stream_builder.metadata(),
534                task.start(),
535                task.length(),
536            )?;
537            selected_row_group_indices = Some(byte_range_filtered_row_groups);
538        }
539
540        if let Some(predicate) = final_predicate {
541            // Normalize NOT out of the predicate once for every consumer below. The
542            // conservative row-group and page-index evaluators cannot soundly negate
543            // a NOT node (negating a "might match" bound can prune matching rows, and
544            // the page-index evaluator rejects NOT outright). A FileScanTask built or
545            // deserialized outside TableScan carries an un-normalized predicate, so
546            // rewrite it here. Matches TableScan::with_filter.
547            let predicate = predicate.rewrite_not();
548
549            let (iceberg_field_ids, field_id_map) = ArrowReader::build_field_id_set_and_map(
550                record_batch_stream_builder.parquet_schema(),
551                record_batch_stream_builder.schema(),
552                &predicate,
553                use_position_fallback,
554            )?;
555
556            let row_filter = ArrowReader::get_row_filter(
557                &predicate,
558                record_batch_stream_builder.parquet_schema(),
559                &iceberg_field_ids,
560                &field_id_map,
561            )?;
562            record_batch_stream_builder = record_batch_stream_builder.with_row_filter(row_filter);
563
564            if self.row_group_filtering_enabled {
565                let candidates = candidate_row_groups(
566                    selected_row_group_indices.as_deref(),
567                    record_batch_stream_builder.metadata().num_row_groups(),
568                );
569
570                selected_row_group_indices = Some(ArrowReader::get_selected_row_group_indices(
571                    &predicate,
572                    record_batch_stream_builder.metadata(),
573                    &field_id_map,
574                    task.schema(),
575                    candidates.as_ref(),
576                )?);
577            }
578
579            if self.bloom_filter_enabled {
580                let candidates = candidate_row_groups(
581                    selected_row_group_indices.as_deref(),
582                    record_batch_stream_builder.metadata().num_row_groups(),
583                );
584
585                // Spawned for the same reason as the open (#3276).
586                let bloom_predicate = predicate.clone();
587                let bloom_candidates = candidates.as_ref().to_vec();
588                let bloom_field_ids = field_id_map.clone();
589                let mut builder = record_batch_stream_builder;
590                let (builder, bloom_filtered) = self
591                    .runtime
592                    .io()
593                    .spawn(async move {
594                        let filtered = Self::filter_row_groups_by_bloom_filter(
595                            &bloom_predicate,
596                            &mut builder,
597                            &bloom_candidates,
598                            &bloom_field_ids,
599                        )
600                        .await;
601                        (builder, filtered)
602                    })
603                    .abort_on_drop()
604                    .await?;
605                record_batch_stream_builder = builder;
606                let bloom_filtered = bloom_filtered?;
607
608                if bloom_filtered.len() < candidates.len() {
609                    selected_row_group_indices = Some(bloom_filtered);
610                }
611            }
612
613            if self.row_selection_enabled {
614                row_selection = ArrowReader::get_row_selection_for_filter_predicate(
615                    &predicate,
616                    record_batch_stream_builder.metadata(),
617                    &selected_row_group_indices,
618                    &field_id_map,
619                    task.schema(),
620                )?;
621            }
622        }
623
624        let positional_delete_indexes = delete_filter.get_delete_vector(&task);
625
626        if let Some(positional_delete_indexes) = positional_delete_indexes {
627            let delete_row_selection = {
628                let positional_delete_indexes = positional_delete_indexes.lock().unwrap();
629
630                ArrowReader::build_deletes_row_selection(
631                    record_batch_stream_builder.metadata().row_groups(),
632                    &selected_row_group_indices,
633                    &positional_delete_indexes,
634                )
635            }?;
636
637            // merge the row selection from the delete files with the row selection
638            // from the filter predicate, if there is one from the filter predicate
639            row_selection = match row_selection {
640                None => Some(delete_row_selection),
641                Some(filter_row_selection) => {
642                    Some(filter_row_selection.intersection(&delete_row_selection))
643                }
644            };
645        }
646
647        if let Some(row_selection) = row_selection {
648            record_batch_stream_builder =
649                record_batch_stream_builder.with_row_selection(row_selection);
650        }
651
652        if let Some(selected_row_group_indices) = selected_row_group_indices {
653            record_batch_stream_builder =
654                record_batch_stream_builder.with_row_groups(selected_row_group_indices);
655        }
656
657        // Build the batch stream and send all the RecordBatches that it generates
658        // to the requester. When `_row_id` is projected, synthesize it over the raw parquet
659        // batches (using the reader-produced `_pos` position) before the transformer, which
660        // then passes it through as a virtual field.
661        let first_row_id = task.first_row_id();
662        let record_batch_stream = record_batch_stream_builder.build()?.map(move |batch| {
663            let mut batch = batch.map_err(|err| -> Error { err.into() })?;
664            if project_row_id {
665                batch = synthesize_row_id_column(batch, first_row_id)?;
666            }
667            // Process the record batch (type promotion, column reordering, virtual fields, etc.)
668            record_batch_transformer.process_record_batch(batch)
669        });
670
671        Ok(Box::pin(record_batch_stream) as ArrowRecordBatchStream)
672    }
673
674    /// Applies all task-specific schema and virtual-column options, rebuilding the
675    /// Arrow reader metadata at most once.
676    fn configure_arrow_reader_metadata(
677        arrow_metadata: ArrowReaderMetadata,
678        task: &FileScanTask,
679        missing_field_ids: bool,
680        install_row_number: bool,
681    ) -> Result<ArrowReaderMetadata> {
682        // Schema resolution follows the Iceberg Column Projection rule:
683        // "Columns in Iceberg data files are selected by field id."
684        // https://iceberg.apache.org/spec/#column-projection
685        //
686        // This mirrors Java's ReadConf strategy: use embedded IDs with pruneColumns();
687        // otherwise, use applyNameMapping() followed by pruneColumns() when configured, or
688        // addFallbackIds() followed by pruneColumnsFallback() for position-based fallback.
689        // The fast path returns early without materializing an owned schema.
690        let arrow_schema = if missing_field_ids {
691            let schema = if let Some(name_mapping) = task.name_mapping() {
692                apply_name_mapping_to_arrow_schema(
693                    Arc::clone(arrow_metadata.schema()),
694                    name_mapping,
695                )?
696            } else {
697                add_fallback_field_ids_to_arrow_schema(arrow_metadata.schema())
698            };
699            // Coerce INT96 timestamp columns before building the stream reader to avoid
700            // i64 overflow in arrow-rs. Apply this after assigning any missing field IDs
701            // so the final schema contains both changes.
702            coerce_int96_timestamps(&schema, task.schema()).unwrap_or(schema)
703        } else if let Some(coerced) =
704            coerce_int96_timestamps(arrow_metadata.schema(), task.schema())
705        {
706            coerced
707        } else if install_row_number {
708            Arc::clone(arrow_metadata.schema())
709        } else {
710            return Ok(arrow_metadata);
711        };
712
713        let mut options = ArrowReaderOptions::new().with_schema(Arc::clone(&arrow_schema));
714        if install_row_number {
715            let row_number_field = Arc::new(
716                Field::new(RESERVED_COL_NAME_POS, DataType::Int64, false)
717                    .with_metadata(HashMap::from([(
718                        PARQUET_FIELD_ID_META_KEY.to_string(),
719                        RESERVED_FIELD_ID_POS.to_string(),
720                    )]))
721                    .with_extension_type(RowNumber),
722            );
723            options = options.with_virtual_columns(vec![row_number_field])?;
724        }
725
726        ArrowReaderMetadata::try_new(Arc::clone(arrow_metadata.metadata()), options).map_err(|e| {
727            Error::new(
728                ErrorKind::Unexpected,
729                format!(
730                    "Failed to create ArrowReaderMetadata with the configured reader options \
731                     (missing_field_ids: {missing_field_ids}, \
732                      install_row_number: {install_row_number}, schema: {arrow_schema})"
733                ),
734            )
735            .with_source(e)
736        })
737    }
738
739    /// Reads bloom filters for relevant columns and evaluates the predicate
740    /// against them to filter out row groups that definitely don't match.
741    async fn filter_row_groups_by_bloom_filter(
742        predicate: &BoundPredicate,
743        builder: &mut ParquetRecordBatchStreamBuilder<ArrowFileReader>,
744        candidate_row_groups: &[usize],
745        field_id_map: &HashMap<i32, usize>,
746    ) -> Result<Vec<usize>> {
747        // Only collect field IDs from eq/in predicates — the only types
748        // bloom filters can help with. Skip columns not in the parquet schema.
749        let bloom_filter_field_ids: Vec<i32> = collect_bloom_filter_field_ids(predicate)?
750            .into_iter()
751            .filter(|id| field_id_map.contains_key(id))
752            .collect();
753
754        if bloom_filter_field_ids.is_empty() {
755            return Ok(candidate_row_groups.to_vec());
756        }
757
758        let mut result = Vec::with_capacity(candidate_row_groups.len());
759
760        for &rg_idx in candidate_row_groups {
761            let mut bloom_filters: HashMap<i32, ColumnBloomFilter> = HashMap::new();
762
763            for &field_id in &bloom_filter_field_ids {
764                let col_idx = field_id_map[&field_id];
765                let col_meta = builder.metadata().row_group(rg_idx).column(col_idx);
766
767                // Only attempt to load if this column chunk actually has a bloom filter
768                if col_meta.bloom_filter_offset().is_none() {
769                    continue;
770                }
771
772                let physical_type = col_meta.column_type();
773                let type_length = col_meta.column_descr().type_length();
774
775                match builder
776                    .get_row_group_column_bloom_filter(rg_idx, col_idx)
777                    .await
778                {
779                    Ok(Some(sbbf)) => {
780                        bloom_filters.insert(
781                            field_id,
782                            ColumnBloomFilter::new(sbbf, physical_type, type_length),
783                        );
784                    }
785                    Ok(None) => {}
786                    Err(e) => {
787                        // Left absent from the map, so the evaluator treats the column
788                        // as might-match and the row group survives.
789                        tracing::debug!(
790                            "Bloom filter for field {field_id} in row group {rg_idx} could not be read: {e}"
791                        );
792                    }
793                }
794            }
795
796            match BloomFilterEvaluator::eval(predicate, &bloom_filters) {
797                Ok(true) => result.push(rg_idx),
798                Ok(false) => { /* Row group pruned by bloom filter */ }
799                Err(e) => {
800                    tracing::debug!(
801                        "Bloom filter evaluation failed for row group {rg_idx}, including it: {e}"
802                    );
803                    result.push(rg_idx);
804                }
805            }
806        }
807
808        Ok(result)
809    }
810}
811
812impl ArrowReader {
813    /// Opens a Parquet file and loads its metadata, wrapping the reader with
814    /// [`CountingFileRead`] so all I/O is accumulated into `bytes_read`.
815    pub(crate) async fn open_parquet_file(
816        data_file_path: &str,
817        file_io: &FileIO,
818        file_size_in_bytes: u64,
819        parquet_read_options: ParquetReadOptions,
820        bytes_read: &Arc<AtomicU64>,
821        key_metadata: Option<&[u8]>,
822    ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> {
823        let parquet_file = file_io.new_input(data_file_path)?;
824        let counting_reader =
825            CountingFileRead::new(parquet_file.reader().await?, Arc::clone(bytes_read));
826        Self::build_parquet_reader(
827            Box::new(counting_reader),
828            file_size_in_bytes,
829            parquet_read_options,
830            key_metadata,
831        )
832        .await
833    }
834
835    async fn build_parquet_reader(
836        parquet_reader: Box<dyn FileRead>,
837        file_size_in_bytes: u64,
838        parquet_read_options: ParquetReadOptions,
839        key_metadata: Option<&[u8]>,
840    ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> {
841        let mut reader = ArrowFileReader::new(
842            FileMetadata {
843                size: file_size_in_bytes,
844            },
845            parquet_reader,
846        )
847        .with_parquet_read_options(parquet_read_options);
848
849        let arrow_reader_options = Self::build_arrow_reader_options(key_metadata)?;
850
851        let arrow_metadata = ArrowReaderMetadata::load_async(&mut reader, arrow_reader_options)
852            .await
853            .map_err(|e| {
854                Error::new(ErrorKind::Unexpected, "Failed to load Parquet metadata").with_source(e)
855            })?;
856
857        Ok((reader, arrow_metadata))
858    }
859
860    /// Builds `ArrowReaderOptions`, adding `FileDecryptionProperties` when
861    /// key metadata is present for Parquet Modular Encryption.
862    fn build_arrow_reader_options(key_metadata: Option<&[u8]>) -> Result<ArrowReaderOptions> {
863        match key_metadata {
864            Some(km) => {
865                let standard_key_metadata = StandardKeyMetadata::decode(km)?;
866                let mut builder = FileDecryptionProperties::builder(
867                    standard_key_metadata.encryption_key().as_bytes().to_vec(),
868                );
869                if let Some(aad) = standard_key_metadata.aad_prefix() {
870                    builder = builder.with_aad_prefix(aad.to_vec());
871                }
872                let decryption_properties = builder.build().map_err(|e| {
873                    Error::new(
874                        ErrorKind::Unexpected,
875                        "Failed to build Parquet file decryption properties",
876                    )
877                    .with_source(e)
878                })?;
879                Ok(
880                    ArrowReaderOptions::new()
881                        .with_file_decryption_properties(decryption_properties),
882                )
883            }
884            None => Ok(ArrowReaderOptions::default()),
885        }
886    }
887}
888
889#[cfg(test)]
890mod tests {
891    use std::collections::HashMap;
892    use std::fs::File;
893    use std::sync::Arc;
894
895    use arrow_array::cast::AsArray;
896    use arrow_array::{Array, ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray};
897    use arrow_cast::cast;
898    use arrow_schema::{DataType, Field, Schema as ArrowSchema};
899    use futures::TryStreamExt;
900    use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
901    use parquet::basic::Compression;
902    use parquet::file::properties::WriterProperties;
903    use tempfile::TempDir;
904
905    use crate::Runtime;
906    use crate::arrow::ArrowReaderBuilder;
907    use crate::arrow::test_utils::write_encrypted_parquet;
908    use crate::io::FileIO;
909    use crate::metadata_columns::{
910        RESERVED_COL_NAME_POS, RESERVED_COL_NAME_ROW_ID, RESERVED_FIELD_ID_FILE,
911        RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID,
912    };
913    use crate::scan::{FileScanTask, FileScanTaskDeleteFile, FileScanTaskStream};
914    use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Schema, SchemaRef, Type};
915
916    // INT96 encoding: [nanos_low_u32, nanos_high_u32, julian_day_u32]
917    // Julian day 2_440_588 = Unix epoch (1970-01-01)
918    const UNIX_EPOCH_JULIAN: i64 = 2_440_588;
919    const MICROS_PER_DAY: i64 = 86_400_000_000;
920    // Noon on 3333-01-01 (Julian day 2_953_529) — outside the i64 nanosecond range (~1677-2262).
921    const INT96_TEST_NANOS_WITHIN_DAY: u64 = 43_200_000_000_000;
922    const INT96_TEST_JULIAN_DAY: u32 = 2_953_529;
923
924    fn make_int96_test_value() -> (parquet::data_type::Int96, i64) {
925        let mut val = parquet::data_type::Int96::new();
926        val.set_data(
927            (INT96_TEST_NANOS_WITHIN_DAY & 0xFFFFFFFF) as u32,
928            (INT96_TEST_NANOS_WITHIN_DAY >> 32) as u32,
929            INT96_TEST_JULIAN_DAY,
930        );
931        let expected_micros = (INT96_TEST_JULIAN_DAY as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY
932            + (INT96_TEST_NANOS_WITHIN_DAY / 1_000) as i64;
933        (val, expected_micros)
934    }
935
936    async fn read_int96_batches(
937        file_path: &str,
938        schema: SchemaRef,
939        project_field_ids: Vec<i32>,
940    ) -> Vec<RecordBatch> {
941        let file_io = FileIO::new_with_fs();
942        let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
943
944        let file_size = std::fs::metadata(file_path).unwrap().len();
945        let task = FileScanTask::builder()
946            .with_file_size_in_bytes(file_size)
947            .with_start(0)
948            .with_length(file_size)
949            .with_data_file_path(file_path.to_string())
950            .with_data_file_format(DataFileFormat::Parquet)
951            .with_schema(schema)
952            .with_project_field_ids(project_field_ids)
953            .with_case_sensitive(false)
954            .build()
955            .unwrap();
956
957        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
958        reader
959            .read(tasks)
960            .unwrap()
961            .stream()
962            .try_collect()
963            .await
964            .unwrap()
965    }
966
967    // ArrowWriter cannot write INT96, so we use SerializedFileWriter directly.
968    fn write_int96_parquet_file(
969        table_location: &str,
970        filename: &str,
971        with_field_ids: bool,
972    ) -> (String, Vec<i64>) {
973        use parquet::basic::{Repetition, Type as PhysicalType};
974        use parquet::data_type::{Int32Type, Int96, Int96Type};
975        use parquet::file::writer::SerializedFileWriter;
976        use parquet::schema::types::Type as SchemaType;
977
978        let file_path = format!("{table_location}/{filename}");
979
980        let mut ts_builder = SchemaType::primitive_type_builder("ts", PhysicalType::INT96)
981            .with_repetition(Repetition::OPTIONAL);
982        let mut id_builder = SchemaType::primitive_type_builder("id", PhysicalType::INT32)
983            .with_repetition(Repetition::REQUIRED);
984
985        if with_field_ids {
986            ts_builder = ts_builder.with_id(Some(1));
987            id_builder = id_builder.with_id(Some(2));
988        }
989
990        let schema = SchemaType::group_type_builder("schema")
991            .with_fields(vec![
992                Arc::new(ts_builder.build().unwrap()),
993                Arc::new(id_builder.build().unwrap()),
994            ])
995            .build()
996            .unwrap();
997
998        // Dates outside the i64 nanosecond range (~1677-2262) overflow without coercion.
999        const NOON_NANOS: u64 = INT96_TEST_NANOS_WITHIN_DAY;
1000        const JULIAN_3333: u32 = INT96_TEST_JULIAN_DAY;
1001        const JULIAN_2100: u32 = 2_488_070;
1002
1003        let test_data: Vec<(u32, u32, u32, i64)> = vec![
1004            // 3333-01-01 00:00:00
1005            (
1006                0,
1007                0,
1008                JULIAN_3333,
1009                (JULIAN_3333 as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY,
1010            ),
1011            // 3333-01-01 12:00:00
1012            (
1013                (NOON_NANOS & 0xFFFFFFFF) as u32,
1014                (NOON_NANOS >> 32) as u32,
1015                JULIAN_3333,
1016                (JULIAN_3333 as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY
1017                    + (NOON_NANOS / 1_000) as i64,
1018            ),
1019            // 2100-01-01 00:00:00
1020            (
1021                0,
1022                0,
1023                JULIAN_2100,
1024                (JULIAN_2100 as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY,
1025            ),
1026        ];
1027
1028        let int96_values: Vec<Int96> = test_data
1029            .iter()
1030            .map(|(lo, hi, day, _)| {
1031                let mut v = Int96::new();
1032                v.set_data(*lo, *hi, *day);
1033                v
1034            })
1035            .collect();
1036
1037        let id_values: Vec<i32> = (0..test_data.len() as i32).collect();
1038        let expected_micros: Vec<i64> = test_data.iter().map(|(_, _, _, m)| *m).collect();
1039
1040        let file = File::create(&file_path).unwrap();
1041        let mut writer =
1042            SerializedFileWriter::new(file, Arc::new(schema), Default::default()).unwrap();
1043
1044        let mut row_group = writer.next_row_group().unwrap();
1045        {
1046            // def=1: ts is OPTIONAL and present. No repetition levels (top-level columns).
1047            let mut col = row_group.next_column().unwrap().unwrap();
1048            col.typed::<Int96Type>()
1049                .write_batch(&int96_values, Some(&vec![1; test_data.len()]), None)
1050                .unwrap();
1051            col.close().unwrap();
1052        }
1053        {
1054            let mut col = row_group.next_column().unwrap().unwrap();
1055            col.typed::<Int32Type>()
1056                .write_batch(&id_values, None, None)
1057                .unwrap();
1058            col.close().unwrap();
1059        }
1060        row_group.close().unwrap();
1061        writer.close().unwrap();
1062
1063        (file_path, expected_micros)
1064    }
1065
1066    async fn assert_int96_read_matches(
1067        file_path: &str,
1068        schema: SchemaRef,
1069        project_field_ids: Vec<i32>,
1070        expected_micros: &[i64],
1071    ) {
1072        use arrow_array::TimestampMicrosecondArray;
1073
1074        let batches = read_int96_batches(file_path, schema, project_field_ids).await;
1075
1076        assert_eq!(batches.len(), 1);
1077        let ts_array = batches[0]
1078            .column(0)
1079            .as_any()
1080            .downcast_ref::<TimestampMicrosecondArray>()
1081            .expect("Expected TimestampMicrosecondArray");
1082
1083        for (i, expected) in expected_micros.iter().enumerate() {
1084            assert_eq!(
1085                ts_array.value(i),
1086                *expected,
1087                "Row {i}: got {}, expected {expected}",
1088                ts_array.value(i)
1089            );
1090        }
1091    }
1092
1093    /// Writes a single-column Parquet file encrypted with `encryption_key`, then reads it
1094    /// back through `ArrowReader` and asserts the round-tripped values. The key length
1095    /// selects the AES-GCM variant in arrow-rs (16 -> AES-128, 32 -> AES-256).
1096    async fn assert_encrypted_parquet_roundtrip(encryption_key: &[u8]) {
1097        let aad_prefix = b"aad_prefix";
1098
1099        let schema = Arc::new(
1100            Schema::builder()
1101                .with_schema_id(1)
1102                .with_fields(vec![
1103                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1104                ])
1105                .build()
1106                .unwrap(),
1107        );
1108
1109        let arrow_schema = Arc::new(ArrowSchema::new(vec![
1110            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
1111                PARQUET_FIELD_ID_META_KEY.to_string(),
1112                "1".to_string(),
1113            )])),
1114        ]));
1115
1116        let tmp_dir = TempDir::new().unwrap();
1117        let table_location = tmp_dir.path().to_str().unwrap().to_string();
1118        let file_io = FileIO::new_with_fs();
1119
1120        let id_data = Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef;
1121        let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap();
1122
1123        let file_path = format!("{table_location}/encrypted.parquet");
1124        write_encrypted_parquet(&file_path, &batch, encryption_key, Some(aad_prefix));
1125
1126        let key_metadata = crate::encryption::StandardKeyMetadata::try_new(encryption_key)
1127            .unwrap()
1128            .with_aad_prefix(aad_prefix)
1129            .encode()
1130            .unwrap();
1131
1132        let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
1133
1134        let task = FileScanTask::builder()
1135            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1136            .with_start(0)
1137            .with_length(0)
1138            .with_data_file_path(file_path)
1139            .with_data_file_format(DataFileFormat::Parquet)
1140            .with_schema(schema)
1141            .with_project_field_ids(vec![1])
1142            .with_case_sensitive(false)
1143            .with_key_metadata(Some(key_metadata))
1144            .build()
1145            .unwrap();
1146
1147        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1148        let batches: Vec<RecordBatch> = reader
1149            .read(tasks)
1150            .unwrap()
1151            .stream()
1152            .try_collect()
1153            .await
1154            .unwrap();
1155
1156        assert_eq!(batches.len(), 1);
1157        let ids = batches[0]
1158            .column(0)
1159            .as_any()
1160            .downcast_ref::<Int32Array>()
1161            .unwrap();
1162        assert_eq!(ids.values(), &[10, 20, 30]);
1163    }
1164
1165    #[tokio::test]
1166    async fn test_read_encrypted_parquet_aes_128() {
1167        assert_encrypted_parquet_roundtrip(b"0123456789abcdef").await;
1168    }
1169
1170    #[tokio::test]
1171    async fn test_read_encrypted_parquet_aes_256() {
1172        assert_encrypted_parquet_roundtrip(b"0123456789abcdef0123456789abcdef").await;
1173    }
1174
1175    #[tokio::test]
1176    async fn test_read_encrypted_parquet_without_key_metadata_fails() {
1177        let encryption_key = b"0123456789abcdef";
1178
1179        let schema = Arc::new(
1180            Schema::builder()
1181                .with_schema_id(1)
1182                .with_fields(vec![
1183                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1184                ])
1185                .build()
1186                .unwrap(),
1187        );
1188
1189        let arrow_schema = Arc::new(ArrowSchema::new(vec![
1190            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
1191                PARQUET_FIELD_ID_META_KEY.to_string(),
1192                "1".to_string(),
1193            )])),
1194        ]));
1195
1196        let tmp_dir = TempDir::new().unwrap();
1197        let table_location = tmp_dir.path().to_str().unwrap().to_string();
1198        let file_io = FileIO::new_with_fs();
1199
1200        let id_data = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
1201        let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap();
1202
1203        let file_path = format!("{table_location}/encrypted_no_key.parquet");
1204        write_encrypted_parquet(&file_path, &batch, encryption_key, None);
1205
1206        let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
1207
1208        let task = FileScanTask::builder()
1209            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1210            .with_start(0)
1211            .with_length(0)
1212            .with_data_file_path(file_path)
1213            .with_data_file_format(DataFileFormat::Parquet)
1214            .with_schema(schema)
1215            .with_project_field_ids(vec![1])
1216            .with_case_sensitive(false)
1217            .build()
1218            .unwrap();
1219
1220        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1221        let result: Result<Vec<RecordBatch>, _> =
1222            reader.read(tasks).unwrap().stream().try_collect().await;
1223
1224        let err = result.unwrap_err();
1225        assert_eq!(err.kind(), crate::ErrorKind::Unexpected);
1226        let err_str = format!("{err}");
1227        assert!(
1228            err_str.contains("encrypted footer"),
1229            "Expected error about encrypted footer, got: {err_str}"
1230        );
1231        assert!(
1232            err_str.contains("decryption properties were not provided"),
1233            "Expected error about missing decryption properties, got: {err_str}"
1234        );
1235    }
1236
1237    /// Writes a plain (unencrypted) single-column Int32 "id" parquet file with the
1238    /// given extra Arrow fields/columns appended, returning the file path.
1239    fn write_plain_parquet(
1240        dir: &str,
1241        name: &str,
1242        extra_fields: Vec<Field>,
1243        extra_columns: Vec<ArrayRef>,
1244    ) -> String {
1245        let props = WriterProperties::builder()
1246            .set_compression(Compression::SNAPPY)
1247            .build();
1248        write_parquet_with(dir, name, extra_fields, extra_columns, props)
1249    }
1250
1251    fn write_parquet_with(
1252        dir: &str,
1253        name: &str,
1254        extra_fields: Vec<Field>,
1255        extra_columns: Vec<ArrayRef>,
1256        props: WriterProperties,
1257    ) -> String {
1258        let mut fields =
1259            vec![
1260                Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
1261                    PARQUET_FIELD_ID_META_KEY.to_string(),
1262                    "1".to_string(),
1263                )])),
1264            ];
1265        fields.extend(extra_fields);
1266        let arrow_schema = Arc::new(ArrowSchema::new(fields));
1267
1268        let mut columns: Vec<ArrayRef> = vec![Arc::new(Int32Array::from(vec![1, 2, 3]))];
1269        columns.extend(extra_columns);
1270        let batch = RecordBatch::try_new(arrow_schema.clone(), columns).unwrap();
1271
1272        let file_path = format!("{dir}/{name}");
1273        let file = File::create(&file_path).unwrap();
1274        let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
1275        writer.write(&batch).unwrap();
1276        writer.close().unwrap();
1277        file_path
1278    }
1279
1280    fn last_updated_seq_task(
1281        file_path: String,
1282        first_row_id: Option<i64>,
1283        data_sequence_number: Option<i64>,
1284    ) -> FileScanTask {
1285        use crate::metadata_columns::RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER;
1286
1287        let schema = Arc::new(
1288            Schema::builder()
1289                .with_schema_id(1)
1290                .with_fields(vec![
1291                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1292                ])
1293                .build()
1294                .unwrap(),
1295        );
1296
1297        FileScanTask::builder()
1298            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1299            .with_start(0)
1300            .with_length(0)
1301            .with_data_file_path(file_path)
1302            .with_data_file_format(DataFileFormat::Parquet)
1303            .with_schema(schema)
1304            .with_project_field_ids(vec![1, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER])
1305            .with_first_row_id(first_row_id)
1306            .with_data_sequence_number(data_sequence_number)
1307            .with_case_sensitive(false)
1308            .build()
1309            .unwrap()
1310    }
1311
1312    /// Asserts the logical per-row values of the `_last_updated_sequence_number`
1313    /// column across all batches, independent of the physical (run-end) encoding.
1314    fn assert_last_updated_seq_column(batches: &[RecordBatch], expected: &[Option<i64>]) {
1315        use arrow_array::cast::AsArray;
1316        use arrow_cast::cast;
1317        use arrow_schema::DataType;
1318
1319        use crate::metadata_columns::RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER;
1320
1321        let mut actual = Vec::new();
1322        for batch in batches {
1323            let col = batch
1324                .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
1325                .expect("_last_updated_sequence_number column should be present");
1326            let logical = cast(col, &DataType::Int64).unwrap();
1327            let values = logical.as_primitive::<arrow_array::types::Int64Type>();
1328            for i in 0..values.len() {
1329                actual.push((!values.is_null(i)).then(|| values.value(i)));
1330            }
1331        }
1332        assert_eq!(actual, expected);
1333    }
1334
1335    #[tokio::test]
1336    async fn test_last_updated_sequence_number_null_when_no_first_row_id() {
1337        let tmp_dir = TempDir::new().unwrap();
1338        let dir = tmp_dir.path().to_str().unwrap();
1339        let file_path = write_plain_parquet(dir, "no_first_row_id.parquet", vec![], vec![]);
1340
1341        // A file with a null first_row_id (v1/v2, or a pre-upgrade v3 snapshot) produces
1342        // a null _last_updated_sequence_number column, even though it has a data
1343        // sequence number; the spec gates both lineage columns on first_row_id.
1344        let task = last_updated_seq_task(file_path, None, Some(9));
1345
1346        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1347        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1348        let batches: Vec<RecordBatch> = reader
1349            .read(tasks)
1350            .unwrap()
1351            .stream()
1352            .try_collect()
1353            .await
1354            .unwrap();
1355
1356        assert_last_updated_seq_column(&batches, &[None, None, None]);
1357    }
1358
1359    #[tokio::test]
1360    async fn test_last_updated_sequence_number_error_when_no_data_seq() {
1361        let tmp_dir = TempDir::new().unwrap();
1362        let dir = tmp_dir.path().to_str().unwrap();
1363        let file_path = write_plain_parquet(dir, "no_data_seq.parquet", vec![], vec![]);
1364
1365        // first_row_id present but data_sequence_number absent: after manifest
1366        // inheritance a committed entry always has one, so this is a malformed
1367        // manifest and must error rather than fabricate or null the column.
1368        let task = last_updated_seq_task(file_path, Some(42), None);
1369
1370        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1371        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1372        let result: Result<Vec<RecordBatch>, _> =
1373            reader.read(tasks).unwrap().stream().try_collect().await;
1374
1375        let err = result.unwrap_err();
1376        assert_eq!(err.kind(), crate::ErrorKind::DataInvalid);
1377        assert!(
1378            format!("{err}").contains("no data sequence number"),
1379            "unexpected error: {err}"
1380        );
1381    }
1382
1383    #[tokio::test]
1384    async fn test_last_updated_sequence_number_derived_from_data_seq() {
1385        let tmp_dir = TempDir::new().unwrap();
1386        let dir = tmp_dir.path().to_str().unwrap();
1387        let file_path = write_plain_parquet(dir, "with_first_row_id.parquet", vec![], vec![]);
1388
1389        // Non-null first_row_id + data sequence number -> the derived value (the data
1390        // sequence number) for every row. This is the only value-producing arm.
1391        let task = last_updated_seq_task(file_path, Some(42), Some(7));
1392
1393        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1394        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1395        let batches: Vec<RecordBatch> = reader
1396            .read(tasks)
1397            .unwrap()
1398            .stream()
1399            .try_collect()
1400            .await
1401            .unwrap();
1402
1403        assert_last_updated_seq_column(&batches, &[Some(7), Some(7), Some(7)]);
1404    }
1405
1406    #[tokio::test]
1407    async fn test_last_updated_sequence_number_mixed_files_share_schema() {
1408        use arrow_select::concat::concat_batches;
1409
1410        let tmp_dir = TempDir::new().unwrap();
1411        let dir = tmp_dir.path().to_str().unwrap();
1412
1413        // Three files in one scan exercising all three column paths, which must all
1414        // produce the SAME Arrow type (run-end-encoded) or concatenation fails:
1415        //   - constant: first_row_id set, no physical column -> derived constant
1416        //   - null gate: no first_row_id -> null column
1417        //   - coalesce: first_row_id set, physical column present -> per-row + fallback
1418        let constant = last_updated_seq_task(
1419            write_plain_parquet(dir, "constant.parquet", vec![], vec![]),
1420            Some(42),
1421            Some(7),
1422        );
1423        let nulled = last_updated_seq_task(
1424            write_plain_parquet(dir, "nulled.parquet", vec![], vec![]),
1425            None,
1426            Some(7),
1427        );
1428        let coalesced = last_updated_seq_task(
1429            write_plain_parquet(
1430                dir,
1431                "coalesced.parquet",
1432                vec![physical_last_updated_seq_field()],
1433                vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
1434            ),
1435            Some(50),
1436            Some(7),
1437        );
1438
1439        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1440        let tasks = Box::pin(futures::stream::iter(vec![
1441            Ok(constant),
1442            Ok(nulled),
1443            Ok(coalesced),
1444        ])) as FileScanTaskStream;
1445        let batches: Vec<RecordBatch> = reader
1446            .read(tasks)
1447            .unwrap()
1448            .stream()
1449            .try_collect()
1450            .await
1451            .unwrap();
1452
1453        assert_eq!(batches.len(), 3);
1454        // Identical schema across all three paths -> concat succeeds.
1455        let schema = batches[0].schema();
1456        concat_batches(&schema, &batches)
1457            .expect("constant, null and coalesce files must share one column type");
1458    }
1459
1460    /// A parquet field carrying the embedded `_last_updated_sequence_number` field id.
1461    fn physical_last_updated_seq_field() -> Field {
1462        use crate::metadata_columns::{
1463            RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
1464            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1465        };
1466        Field::new(
1467            RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
1468            DataType::Int64,
1469            true,
1470        )
1471        .with_metadata(HashMap::from([(
1472            PARQUET_FIELD_ID_META_KEY.to_string(),
1473            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER.to_string(),
1474        )]))
1475    }
1476
1477    #[tokio::test]
1478    async fn test_last_updated_sequence_number_physical_column_coalesced() {
1479        let tmp_dir = TempDir::new().unwrap();
1480        let dir = tmp_dir.path().to_str().unwrap();
1481        // A file that physically carries the column, as Iceberg Java writes when
1482        // carrying rows forward across a rewrite: some rows have a stored value, some
1483        // are null (added/modified rows, inherited on read).
1484        let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1485        let file_path = write_plain_parquet(
1486            dir,
1487            "with_seq.parquet",
1488            vec![physical_last_updated_seq_field()],
1489            vec![seq_col],
1490        );
1491
1492        let task = last_updated_seq_task(file_path, Some(100), Some(9));
1493
1494        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1495        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1496        let batches: Vec<RecordBatch> = reader
1497            .read(tasks)
1498            .unwrap()
1499            .stream()
1500            .try_collect()
1501            .await
1502            .unwrap();
1503
1504        // Per-row value where non-null; the data sequence number (9) where null.
1505        assert_last_updated_seq_column(&batches, &[Some(5), Some(9), Some(8)]);
1506    }
1507
1508    #[tokio::test]
1509    async fn test_last_updated_sequence_number_coalesced_with_pos_column() {
1510        use crate::metadata_columns::{
1511            RESERVED_COL_NAME_POS, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1512            RESERVED_FIELD_ID_POS,
1513        };
1514
1515        let tmp_dir = TempDir::new().unwrap();
1516        let dir = tmp_dir.path().to_str().unwrap();
1517        let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1518        let file_path = write_plain_parquet(
1519            dir,
1520            "with_seq_and_pos.parquet",
1521            vec![physical_last_updated_seq_field()],
1522            vec![seq_col],
1523        );
1524
1525        // Co-project `_pos` (a virtual column appended to the Arrow output schema) with the
1526        // physical coalesce column. This guards that the physical column's index is
1527        // resolved in the Parquet schema, not the Arrow schema (whose indices shift once
1528        // virtual columns are appended).
1529        let schema = Arc::new(
1530            Schema::builder()
1531                .with_schema_id(1)
1532                .with_fields(vec![
1533                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1534                ])
1535                .build()
1536                .unwrap(),
1537        );
1538        let task = FileScanTask::builder()
1539            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1540            .with_start(0)
1541            .with_length(0)
1542            .with_data_file_path(file_path)
1543            .with_data_file_format(DataFileFormat::Parquet)
1544            .with_schema(schema)
1545            .with_project_field_ids(vec![
1546                1,
1547                RESERVED_FIELD_ID_POS,
1548                RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1549            ])
1550            .with_first_row_id(Some(100))
1551            .with_data_sequence_number(Some(9))
1552            .with_case_sensitive(false)
1553            .build()
1554            .unwrap();
1555
1556        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1557        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1558        let batches: Vec<RecordBatch> = reader
1559            .read(tasks)
1560            .unwrap()
1561            .stream()
1562            .try_collect()
1563            .await
1564            .unwrap();
1565
1566        // The seq column still coalesces correctly...
1567        assert_last_updated_seq_column(&batches, &[Some(5), Some(9), Some(8)]);
1568        // ...and `_pos` is the row position, unaffected by the physical-column union.
1569        let pos_col = batches[0]
1570            .column_by_name(RESERVED_COL_NAME_POS)
1571            .expect("_pos column should be present")
1572            .as_primitive::<arrow_array::types::Int64Type>();
1573        assert_eq!(pos_col.values(), &[0, 1, 2]);
1574    }
1575
1576    #[tokio::test]
1577    async fn test_last_updated_sequence_number_physical_column_nulled_without_first_row_id() {
1578        let tmp_dir = TempDir::new().unwrap();
1579        let dir = tmp_dir.path().to_str().unwrap();
1580        let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1581        let file_path = write_plain_parquet(
1582            dir,
1583            "with_seq_no_first_row_id.parquet",
1584            vec![physical_last_updated_seq_field()],
1585            vec![seq_col],
1586        );
1587
1588        // Null first_row_id: the whole column is null even though the file physically
1589        // carries per-row values -- the gate wins, and the physical column is not read.
1590        let task = last_updated_seq_task(file_path, None, Some(9));
1591
1592        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1593        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1594        let batches: Vec<RecordBatch> = reader
1595            .read(tasks)
1596            .unwrap()
1597            .stream()
1598            .try_collect()
1599            .await
1600            .unwrap();
1601
1602        assert_last_updated_seq_column(&batches, &[None, None, None]);
1603    }
1604
1605    #[tokio::test]
1606    async fn test_last_updated_sequence_number_present_by_name_without_id_unsupported() {
1607        let tmp_dir = TempDir::new().unwrap();
1608        let dir = tmp_dir.path().to_str().unwrap();
1609        // Column present by name but WITHOUT the embedded field id (e.g. name mapping /
1610        // positional fallback). The transformer keys the source column by field id, so
1611        // this shape can't be threaded and is rejected loudly.
1612        let seq_field = Field::new(
1613            crate::metadata_columns::RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
1614            DataType::Int64,
1615            true,
1616        );
1617        let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1618        let file_path =
1619            write_plain_parquet(dir, "with_seq_by_name.parquet", vec![seq_field], vec![
1620                seq_col,
1621            ]);
1622
1623        let task = last_updated_seq_task(file_path, Some(100), Some(9));
1624
1625        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1626        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1627        let result: Result<Vec<RecordBatch>, _> =
1628            reader.read(tasks).unwrap().stream().try_collect().await;
1629
1630        let err = result.unwrap_err();
1631        assert_eq!(err.kind(), crate::ErrorKind::FeatureUnsupported);
1632        assert!(
1633            format!("{err}").contains("without an embedded field id"),
1634            "unexpected error: {err}"
1635        );
1636    }
1637
1638    #[tokio::test]
1639    async fn test_last_updated_sequence_number_physical_column_first_row_id_without_data_seq() {
1640        let tmp_dir = TempDir::new().unwrap();
1641        let dir = tmp_dir.path().to_str().unwrap();
1642        let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1643        let file_path = write_plain_parquet(
1644            dir,
1645            "with_seq_no_data_seq.parquet",
1646            vec![physical_last_updated_seq_field()],
1647            vec![seq_col],
1648        );
1649
1650        // first_row_id set but no data sequence number: after manifest inheritance a
1651        // committed entry always has one, so this is a malformed manifest, rejected loudly
1652        // rather than nulled.
1653        let task = last_updated_seq_task(file_path, Some(100), None);
1654
1655        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1656        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1657        let result: Result<Vec<RecordBatch>, _> =
1658            reader.read(tasks).unwrap().stream().try_collect().await;
1659
1660        let err = result.unwrap_err();
1661        assert_eq!(err.kind(), crate::ErrorKind::DataInvalid);
1662        assert!(
1663            format!("{err}").contains("no data sequence number"),
1664            "unexpected error: {err}"
1665        );
1666    }
1667
1668    /// A scan task projecting `id` + `_row_id`, with the given `first_row_id`.
1669    fn row_id_task(file_path: String, first_row_id: Option<i64>) -> FileScanTask {
1670        row_id_task_with_options(file_path, first_row_id, 0, 0, vec![])
1671    }
1672
1673    fn row_id_task_with_options(
1674        file_path: String,
1675        first_row_id: Option<i64>,
1676        start: u64,
1677        length: u64,
1678        deletes: Vec<FileScanTaskDeleteFile>,
1679    ) -> FileScanTask {
1680        let schema = Arc::new(
1681            Schema::builder()
1682                .with_schema_id(1)
1683                .with_fields(vec![
1684                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1685                ])
1686                .build()
1687                .unwrap(),
1688        );
1689
1690        FileScanTask::builder()
1691            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1692            .with_start(start)
1693            .with_length(length)
1694            .with_data_file_path(file_path)
1695            .with_data_file_format(DataFileFormat::Parquet)
1696            .with_schema(schema)
1697            .with_project_field_ids(vec![1, RESERVED_FIELD_ID_ROW_ID])
1698            .with_first_row_id(first_row_id)
1699            .with_deletes(deletes)
1700            .with_case_sensitive(false)
1701            .build()
1702            .unwrap()
1703    }
1704
1705    /// Asserts the logical per-row values of the `_row_id` column across all batches,
1706    /// independent of the physical (run-end) encoding.
1707    fn assert_row_id_column(batches: &[RecordBatch], expected: &[Option<i64>]) {
1708        use arrow_array::cast::AsArray;
1709        use arrow_cast::cast;
1710        use arrow_schema::DataType;
1711
1712        let mut actual = Vec::new();
1713        for batch in batches {
1714            let col = batch
1715                .column_by_name(RESERVED_COL_NAME_ROW_ID)
1716                .expect("_row_id column should be present");
1717            let logical = cast(col, &DataType::Int64).unwrap();
1718            let values = logical.as_primitive::<arrow_array::types::Int64Type>();
1719            for i in 0..values.len() {
1720                actual.push((!values.is_null(i)).then(|| values.value(i)));
1721            }
1722        }
1723        assert_eq!(actual, expected);
1724    }
1725
1726    /// A parquet field carrying the embedded `_row_id` field id.
1727    fn physical_row_id_field() -> Field {
1728        Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true).with_metadata(HashMap::from([
1729            (
1730                PARQUET_FIELD_ID_META_KEY.to_string(),
1731                RESERVED_FIELD_ID_ROW_ID.to_string(),
1732            ),
1733        ]))
1734    }
1735
1736    #[tokio::test]
1737    async fn test_row_id_synthesized_from_first_row_id_and_pos() {
1738        let tmp_dir = TempDir::new().unwrap();
1739        let dir = tmp_dir.path().to_str().unwrap();
1740        let file_path = write_plain_parquet(dir, "row_id_synth.parquet", vec![], vec![]);
1741
1742        // No physical column: every row is first_row_id + pos.
1743        let task = row_id_task(file_path, Some(100));
1744
1745        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1746        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1747        let batches: Vec<RecordBatch> = reader
1748            .read(tasks)
1749            .unwrap()
1750            .stream()
1751            .try_collect()
1752            .await
1753            .unwrap();
1754
1755        assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]);
1756    }
1757
1758    #[tokio::test]
1759    async fn test_row_id_physical_column_coalesced() {
1760        let tmp_dir = TempDir::new().unwrap();
1761        let dir = tmp_dir.path().to_str().unwrap();
1762        // A file that physically carries `_row_id`, as written when carrying rows forward
1763        // across a rewrite: some rows have a stored value, some are null.
1764        let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1765        let file_path = write_plain_parquet(
1766            dir,
1767            "row_id_phys.parquet",
1768            vec![physical_row_id_field()],
1769            vec![id_col],
1770        );
1771
1772        let task = row_id_task(file_path, Some(100));
1773
1774        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1775        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1776        let batches: Vec<RecordBatch> = reader
1777            .read(tasks)
1778            .unwrap()
1779            .stream()
1780            .try_collect()
1781            .await
1782            .unwrap();
1783
1784        // Per-row value where non-null; first_row_id + pos (101) where null.
1785        assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
1786    }
1787
1788    #[tokio::test]
1789    async fn test_row_id_only_synthesis_reads_no_data_columns() {
1790        // The common v3 case: a new-row file with `first_row_id` set and NO physically
1791        // stored `_row_id`, projecting only `_row_id`. `_row_id` synthesis installs the
1792        // RowNumber virtual column (via `need_row_number`), so the row count comes from it
1793        // -- the scan must read no data columns, not fall back to reading everything.
1794        let tmp_dir = TempDir::new().unwrap();
1795        let dir = tmp_dir.path().to_str().unwrap();
1796
1797        let meta_only = metadata_projection_task_with_first_row_id(
1798            write_parquet_with_wide_column(dir, "row_id_only.parquet", vec![], vec![]),
1799            id_and_wide_schema(),
1800            vec![RESERVED_FIELD_ID_ROW_ID],
1801            Some(100),
1802        );
1803        let (batches, meta_only_bytes) = scan_task(meta_only).await;
1804
1805        assert_eq!(batches[0].num_columns(), 1);
1806        assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]);
1807
1808        // A scan that also projects the wide data column must read materially more.
1809        let with_data = metadata_projection_task_with_first_row_id(
1810            write_parquet_with_wide_column(dir, "row_id_only_ref.parquet", vec![], vec![]),
1811            id_and_wide_schema(),
1812            vec![2, RESERVED_FIELD_ID_ROW_ID],
1813            Some(100),
1814        );
1815        let (_, with_data_bytes) = scan_task(with_data).await;
1816
1817        assert!(
1818            meta_only_bytes < with_data_bytes,
1819            "_row_id-only synthesis should read fewer bytes than a scan of the wide column: \
1820             {meta_only_bytes} vs {with_data_bytes}"
1821        );
1822    }
1823
1824    #[tokio::test]
1825    async fn test_row_id_only_null_first_row_id_reads_no_data_columns() {
1826        // A null `first_row_id` (v1/v2, or a pre-upgrade v3 snapshot) nulls the whole
1827        // `_row_id` column, so nothing is synthesized -- but the column still has a length.
1828        // The RowNumber counter installed for the metadata-only projection supplies it, so
1829        // the scan reads no data columns instead of reading everything just for the count.
1830        let tmp_dir = TempDir::new().unwrap();
1831        let dir = tmp_dir.path().to_str().unwrap();
1832
1833        let meta_only = metadata_projection_task_with_first_row_id(
1834            write_parquet_with_wide_column(dir, "row_id_null.parquet", vec![], vec![]),
1835            id_and_wide_schema(),
1836            vec![RESERVED_FIELD_ID_ROW_ID],
1837            None,
1838        );
1839        let (batches, meta_only_bytes) = scan_task(meta_only).await;
1840
1841        assert_eq!(batches[0].num_columns(), 1);
1842        assert_row_id_column(&batches, &[None, None, None]);
1843
1844        let with_data = metadata_projection_task_with_first_row_id(
1845            write_parquet_with_wide_column(dir, "row_id_null_ref.parquet", vec![], vec![]),
1846            id_and_wide_schema(),
1847            vec![2, RESERVED_FIELD_ID_ROW_ID],
1848            None,
1849        );
1850        let (_, with_data_bytes) = scan_task(with_data).await;
1851
1852        assert!(
1853            meta_only_bytes < with_data_bytes,
1854            "_row_id-only scan with a null first_row_id should read fewer bytes than a scan \
1855             of the wide column: {meta_only_bytes} vs {with_data_bytes}"
1856        );
1857    }
1858
1859    #[tokio::test]
1860    async fn test_last_updated_seq_only_reads_no_data_columns() {
1861        use crate::metadata_columns::RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER;
1862
1863        // `SELECT _last_updated_sequence_number` derives a per-row constant (the data
1864        // sequence number) with no physical column to read. The metadata-only RowNumber
1865        // counter supplies the row count, so the scan prunes the data columns rather than
1866        // reading them all just to size the constant.
1867        let tmp_dir = TempDir::new().unwrap();
1868        let dir = tmp_dir.path().to_str().unwrap();
1869
1870        let lusn_only_task = |file_path: String, project_field_ids: Vec<i32>| {
1871            FileScanTask::builder()
1872                .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1873                .with_start(0)
1874                .with_length(0)
1875                .with_data_file_path(file_path)
1876                .with_data_file_format(DataFileFormat::Parquet)
1877                .with_schema(id_and_wide_schema())
1878                .with_project_field_ids(project_field_ids)
1879                .with_first_row_id(Some(42))
1880                .with_data_sequence_number(Some(7))
1881                .with_case_sensitive(false)
1882                .build()
1883                .unwrap()
1884        };
1885
1886        let meta_only = lusn_only_task(
1887            write_parquet_with_wide_column(dir, "lusn_only.parquet", vec![], vec![]),
1888            vec![RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER],
1889        );
1890        let (batches, meta_only_bytes) = scan_task(meta_only).await;
1891
1892        assert_eq!(batches[0].num_columns(), 1);
1893        assert_last_updated_seq_column(&batches, &[Some(7), Some(7), Some(7)]);
1894
1895        let with_data = lusn_only_task(
1896            write_parquet_with_wide_column(dir, "lusn_only_ref.parquet", vec![], vec![]),
1897            vec![2, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER],
1898        );
1899        let (_, with_data_bytes) = scan_task(with_data).await;
1900
1901        assert!(
1902            meta_only_bytes < with_data_bytes,
1903            "_last_updated_sequence_number-only scan should read fewer bytes than a scan of \
1904             the wide column: {meta_only_bytes} vs {with_data_bytes}"
1905        );
1906    }
1907
1908    #[tokio::test]
1909    async fn test_partition_only_reads_no_data_columns() {
1910        use arrow_array::StructArray;
1911
1912        use crate::metadata_columns::{RESERVED_COL_NAME_PARTITION, RESERVED_FIELD_ID_PARTITION};
1913        use crate::spec::{Literal, PartitionSpec, Struct, Transform};
1914
1915        // `SELECT _partition` materializes a struct constant from the task's partition
1916        // metadata, with no physical column to read. It is the only struct-constant metadata
1917        // column; the RowNumber counter sizes it, so the scan prunes the data columns.
1918        let tmp_dir = TempDir::new().unwrap();
1919        let dir = tmp_dir.path().to_str().unwrap();
1920        let schema = id_and_wide_schema();
1921        let spec = Arc::new(
1922            PartitionSpec::builder(schema.clone())
1923                .with_spec_id(7)
1924                .add_partition_field("id", "id", Transform::Identity)
1925                .unwrap()
1926                .build()
1927                .unwrap(),
1928        );
1929        let unified_type = Arc::new(spec.partition_type(&schema).unwrap());
1930        let partition_data = Struct::from_iter(vec![Some(Literal::int(42))]);
1931
1932        let partition_task = |file_path: String, project_field_ids: Vec<i32>| {
1933            FileScanTask::builder()
1934                .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1935                .with_start(0)
1936                .with_length(0)
1937                .with_data_file_path(file_path)
1938                .with_data_file_format(DataFileFormat::Parquet)
1939                .with_schema(schema.clone())
1940                .with_project_field_ids(project_field_ids)
1941                .with_partition_spec(Some(spec.clone()))
1942                .with_partition(Some(partition_data.clone()))
1943                .with_unified_partition_type(Some(unified_type.clone()))
1944                .with_case_sensitive(false)
1945                .build()
1946                .unwrap()
1947        };
1948
1949        let meta_only = partition_task(
1950            write_parquet_with_wide_column(dir, "partition_only.parquet", vec![], vec![]),
1951            vec![RESERVED_FIELD_ID_PARTITION],
1952        );
1953        let (batches, meta_only_bytes) = scan_task(meta_only).await;
1954
1955        assert_eq!(batches[0].num_columns(), 1);
1956        let partition_col = batches[0]
1957            .column_by_name(RESERVED_COL_NAME_PARTITION)
1958            .expect("_partition column should be present")
1959            .as_any()
1960            .downcast_ref::<StructArray>()
1961            .unwrap();
1962        assert_eq!(partition_col.len(), 3);
1963        let inner = partition_col
1964            .column(0)
1965            .as_any()
1966            .downcast_ref::<Int32Array>()
1967            .unwrap();
1968        assert_eq!(inner.values(), &[42, 42, 42]);
1969
1970        let with_data = partition_task(
1971            write_parquet_with_wide_column(dir, "partition_only_ref.parquet", vec![], vec![]),
1972            vec![2, RESERVED_FIELD_ID_PARTITION],
1973        );
1974        let (_, with_data_bytes) = scan_task(with_data).await;
1975
1976        assert!(
1977            meta_only_bytes < with_data_bytes,
1978            "_partition-only scan should read fewer bytes than a scan of the wide column: \
1979             {meta_only_bytes} vs {with_data_bytes}"
1980        );
1981    }
1982
1983    #[tokio::test]
1984    async fn test_row_id_resolves_alongside_id_less_leaf() {
1985        // A file with an id-less leaf (mimicking a Variant column's internal metadata/value
1986        // leaves, which the spec requires to have no field id) plus a physical `_row_id`
1987        // that carries its embedded id. The reserved id must still resolve -- an
1988        // all-or-nothing field map would bail on the id-less leaf and wrongly reject the file.
1989        let tmp_dir = TempDir::new().unwrap();
1990        let dir = tmp_dir.path().to_str().unwrap();
1991        let idless_field = Field::new("variant_internal", DataType::Utf8, true);
1992        let idless_col = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
1993        let row_id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1994        let file_path = write_plain_parquet(
1995            dir,
1996            "row_id_with_idless_leaf.parquet",
1997            vec![idless_field, physical_row_id_field()],
1998            vec![idless_col, row_id_col],
1999        );
2000
2001        let task = row_id_task(file_path, Some(100));
2002
2003        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2004        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2005        let batches: Vec<RecordBatch> = reader
2006            .read(tasks)
2007            .unwrap()
2008            .stream()
2009            .try_collect()
2010            .await
2011            .unwrap();
2012
2013        // Physical value where non-null; first_row_id + pos (101) where null.
2014        assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
2015    }
2016
2017    #[tokio::test]
2018    async fn test_row_id_and_last_updated_seq_co_projected() {
2019        use crate::metadata_columns::RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER;
2020
2021        // Both lineage columns projected together over a file carrying both physical
2022        // leaves. Each must materialize independently -- neither leaf's mask clobbers the
2023        // other, and the two synthesized columns keep their own values.
2024        let tmp_dir = TempDir::new().unwrap();
2025        let dir = tmp_dir.path().to_str().unwrap();
2026        let row_id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2027        let seq_col = Arc::new(Int64Array::from(vec![Some(50), None, Some(70)])) as ArrayRef;
2028        let file_path = write_plain_parquet(
2029            dir,
2030            "row_id_and_seq.parquet",
2031            vec![physical_row_id_field(), physical_last_updated_seq_field()],
2032            vec![row_id_col, seq_col],
2033        );
2034
2035        let schema = Arc::new(
2036            Schema::builder()
2037                .with_schema_id(1)
2038                .with_fields(vec![
2039                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2040                ])
2041                .build()
2042                .unwrap(),
2043        );
2044        let task = FileScanTask::builder()
2045            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2046            .with_start(0)
2047            .with_length(0)
2048            .with_data_file_path(file_path)
2049            .with_data_file_format(DataFileFormat::Parquet)
2050            .with_schema(schema)
2051            .with_project_field_ids(vec![
2052                1,
2053                RESERVED_FIELD_ID_ROW_ID,
2054                RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
2055            ])
2056            .with_first_row_id(Some(100))
2057            .with_data_sequence_number(Some(9))
2058            .with_case_sensitive(false)
2059            .build()
2060            .unwrap();
2061
2062        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2063        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2064        let batches: Vec<RecordBatch> = reader
2065            .read(tasks)
2066            .unwrap()
2067            .stream()
2068            .try_collect()
2069            .await
2070            .unwrap();
2071
2072        // _row_id: physical value where non-null, else first_row_id + pos (101).
2073        assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
2074        // _last_updated_sequence_number: physical value where non-null, else data seq (9).
2075        assert_last_updated_seq_column(&batches, &[Some(50), Some(9), Some(70)]);
2076    }
2077
2078    #[tokio::test]
2079    async fn test_row_id_null_when_no_first_row_id() {
2080        let tmp_dir = TempDir::new().unwrap();
2081        let dir = tmp_dir.path().to_str().unwrap();
2082        // Physically carries `_row_id`, but the file has a null first_row_id.
2083        let id_col = Arc::new(Int64Array::from(vec![Some(5), Some(6), Some(7)])) as ArrayRef;
2084        let file_path = write_plain_parquet(
2085            dir,
2086            "row_id_no_first.parquet",
2087            vec![physical_row_id_field()],
2088            vec![id_col],
2089        );
2090
2091        // Null first_row_id: the whole column is null; the physical values are not read.
2092        let task = row_id_task(file_path, None);
2093
2094        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2095        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2096        let batches: Vec<RecordBatch> = reader
2097            .read(tasks)
2098            .unwrap()
2099            .stream()
2100            .try_collect()
2101            .await
2102            .unwrap();
2103
2104        assert_row_id_column(&batches, &[None, None, None]);
2105    }
2106
2107    #[tokio::test]
2108    async fn test_row_id_with_pos_column() {
2109        use crate::metadata_columns::RESERVED_COL_NAME_POS;
2110
2111        let tmp_dir = TempDir::new().unwrap();
2112        let dir = tmp_dir.path().to_str().unwrap();
2113        let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2114        let file_path = write_plain_parquet(
2115            dir,
2116            "row_id_and_pos.parquet",
2117            vec![physical_row_id_field()],
2118            vec![id_col],
2119        );
2120
2121        // Co-project `_pos` and `_row_id`. `_row_id` synthesis consumes the position, and
2122        // `_pos` is also emitted -- the RowNumber column must be added once and the two
2123        // must not interfere.
2124        let schema = Arc::new(
2125            Schema::builder()
2126                .with_schema_id(1)
2127                .with_fields(vec![
2128                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2129                ])
2130                .build()
2131                .unwrap(),
2132        );
2133        let task = FileScanTask::builder()
2134            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2135            .with_start(0)
2136            .with_length(0)
2137            .with_data_file_path(file_path)
2138            .with_data_file_format(DataFileFormat::Parquet)
2139            .with_schema(schema)
2140            .with_project_field_ids(vec![1, RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID])
2141            .with_first_row_id(Some(100))
2142            .with_case_sensitive(false)
2143            .build()
2144            .unwrap();
2145
2146        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2147        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2148        let batches: Vec<RecordBatch> = reader
2149            .read(tasks)
2150            .unwrap()
2151            .stream()
2152            .try_collect()
2153            .await
2154            .unwrap();
2155
2156        // `_row_id` coalesces correctly...
2157        assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
2158        // ...and `_pos` is the row position, not double-counted.
2159        let pos_col = batches[0]
2160            .column_by_name(RESERVED_COL_NAME_POS)
2161            .expect("_pos column should be present")
2162            .as_primitive::<arrow_array::types::Int64Type>();
2163        assert_eq!(pos_col.values(), &[0, 1, 2]);
2164    }
2165
2166    #[tokio::test]
2167    async fn test_row_id_mixed_files_share_schema() {
2168        use arrow_select::concat::concat_batches;
2169
2170        let tmp_dir = TempDir::new().unwrap();
2171        let dir = tmp_dir.path().to_str().unwrap();
2172
2173        // Three files in one scan exercising all three column paths, which must all
2174        // produce the SAME Arrow type (plain Int64) or concatenation fails:
2175        //   - synthesis: first_row_id set, no physical column -> first_row_id + pos
2176        //   - null gate: no first_row_id -> null column
2177        //   - coalesce: first_row_id set, physical column present -> per-row + fallback
2178        let synth = row_id_task(
2179            write_plain_parquet(dir, "row_id_synth2.parquet", vec![], vec![]),
2180            Some(42),
2181        );
2182        let nulled = row_id_task(
2183            write_plain_parquet(dir, "row_id_null2.parquet", vec![], vec![]),
2184            None,
2185        );
2186        let coalesced = row_id_task(
2187            write_plain_parquet(
2188                dir,
2189                "row_id_coalesced2.parquet",
2190                vec![physical_row_id_field()],
2191                vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
2192            ),
2193            Some(50),
2194        );
2195
2196        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2197        let tasks = Box::pin(futures::stream::iter(vec![
2198            Ok(synth),
2199            Ok(nulled),
2200            Ok(coalesced),
2201        ])) as FileScanTaskStream;
2202        let batches: Vec<RecordBatch> = reader
2203            .read(tasks)
2204            .unwrap()
2205            .stream()
2206            .try_collect()
2207            .await
2208            .unwrap();
2209
2210        assert_eq!(batches.len(), 3);
2211        let schema = batches[0].schema();
2212        concat_batches(&schema, &batches)
2213            .expect("synthesis, null and coalesce files must share one column type");
2214    }
2215
2216    #[tokio::test]
2217    async fn test_row_id_present_by_name_without_id_unsupported() {
2218        let tmp_dir = TempDir::new().unwrap();
2219        let dir = tmp_dir.path().to_str().unwrap();
2220        // Column present by name but WITHOUT the embedded field id. The transformer keys
2221        // the source column by field id, so this shape can't be threaded and is rejected.
2222        let id_field = Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true);
2223        let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2224        let file_path =
2225            write_plain_parquet(dir, "row_id_by_name.parquet", vec![id_field], vec![id_col]);
2226
2227        let task = row_id_task(file_path, Some(100));
2228
2229        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2230        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2231        let result: Result<Vec<RecordBatch>, _> =
2232            reader.read(tasks).unwrap().stream().try_collect().await;
2233
2234        let err = result.unwrap_err();
2235        assert_eq!(err.kind(), crate::ErrorKind::FeatureUnsupported);
2236        assert!(
2237            format!("{err}").contains("without an embedded field id"),
2238            "unexpected error: {err}"
2239        );
2240    }
2241
2242    #[tokio::test]
2243    async fn test_row_id_present_by_name_without_id_nulls_when_no_lineage() {
2244        let tmp_dir = TempDir::new().unwrap();
2245        let dir = tmp_dir.path().to_str().unwrap();
2246        // Same name-only shape as the reject test above, but the file carries no row lineage
2247        // (`first_row_id = None`), as a migrated pre-v3 file with a user column named
2248        // `_row_id` would. The physical leaf is never read, so `_row_id` is nulled out
2249        // rather than rejected (matching Java `ValueReaders.rowIds`).
2250        let id_field = Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true);
2251        let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2252        let file_path = write_plain_parquet(
2253            dir,
2254            "row_id_by_name_no_lineage.parquet",
2255            vec![id_field],
2256            vec![id_col],
2257        );
2258
2259        let task = row_id_task(file_path, None);
2260
2261        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2262        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2263        let batches: Vec<RecordBatch> = reader
2264            .read(tasks)
2265            .unwrap()
2266            .stream()
2267            .try_collect()
2268            .await
2269            .unwrap();
2270
2271        assert_row_id_column(&batches, &[None, None, None]);
2272    }
2273
2274    /// Builds a `row_id_task` (see above) that additionally carries a bound predicate,
2275    /// so a `RowSelection` is applied when the reader has row selection enabled.
2276    fn row_id_task_with_predicate(
2277        file_path: String,
2278        first_row_id: Option<i64>,
2279        extra_project_field_ids: Vec<i32>,
2280        predicate: crate::expr::Predicate,
2281    ) -> FileScanTask {
2282        use crate::expr::Bind;
2283
2284        let schema = Arc::new(
2285            Schema::builder()
2286                .with_schema_id(1)
2287                .with_fields(vec![
2288                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2289                ])
2290                .build()
2291                .unwrap(),
2292        );
2293        let bound = predicate.bind(Arc::clone(&schema), false).unwrap();
2294
2295        let mut project_field_ids = vec![1];
2296        project_field_ids.extend(extra_project_field_ids);
2297        project_field_ids.push(RESERVED_FIELD_ID_ROW_ID);
2298
2299        FileScanTask::builder()
2300            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2301            .with_start(0)
2302            .with_length(0)
2303            .with_data_file_path(file_path)
2304            .with_data_file_format(DataFileFormat::Parquet)
2305            .with_schema(schema)
2306            .with_project_field_ids(project_field_ids)
2307            .with_predicate(Some(bound))
2308            .with_first_row_id(first_row_id)
2309            .with_case_sensitive(false)
2310            .build()
2311            .unwrap()
2312    }
2313
2314    #[tokio::test]
2315    async fn test_row_id_stable_under_row_selection() {
2316        use crate::expr::Reference;
2317        use crate::spec::Datum;
2318
2319        let tmp_dir = TempDir::new().unwrap();
2320        let dir = tmp_dir.path().to_str().unwrap();
2321        // id = [1, 2, 3]; drop the middle physical row via a predicate + row selection.
2322        let file_path = write_plain_parquet(dir, "row_id_selection.parquet", vec![], vec![]);
2323
2324        let task = row_id_task_with_predicate(
2325            file_path,
2326            Some(100),
2327            vec![],
2328            Reference::new("id").not_equal_to(Datum::int(2)),
2329        );
2330
2331        // Row selection must be enabled for the predicate to produce a RowSelection.
2332        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current())
2333            .with_row_selection_enabled(true)
2334            .build();
2335        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2336        let batches: Vec<RecordBatch> = reader
2337            .read(tasks)
2338            .unwrap()
2339            .stream()
2340            .try_collect()
2341            .await
2342            .unwrap();
2343
2344        // The survivors are physical rows 0 and 2, so their _row_id is first_row_id + the
2345        // PHYSICAL position: [100, 102]. A dense output index would wrongly give [100, 101].
2346        assert_row_id_column(&batches, &[Some(100), Some(102)]);
2347    }
2348
2349    #[tokio::test]
2350    async fn test_row_id_coalesce_stable_under_row_selection() {
2351        use crate::expr::Reference;
2352        use crate::spec::Datum;
2353
2354        let tmp_dir = TempDir::new().unwrap();
2355        let dir = tmp_dir.path().to_str().unwrap();
2356        // Physical _row_id = [Some(5), None, Some(8)] over id = [1, 2, 3]. Dropping the
2357        // middle row must keep the physical column and the RowNumber fallback row-aligned.
2358        let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2359        let file_path = write_plain_parquet(
2360            dir,
2361            "row_id_coalesce_selection.parquet",
2362            vec![physical_row_id_field()],
2363            vec![id_col],
2364        );
2365
2366        let task = row_id_task_with_predicate(
2367            file_path,
2368            Some(100),
2369            vec![],
2370            Reference::new("id").not_equal_to(Datum::int(2)),
2371        );
2372
2373        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current())
2374            .with_row_selection_enabled(true)
2375            .build();
2376        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2377        let batches: Vec<RecordBatch> = reader
2378            .read(tasks)
2379            .unwrap()
2380            .stream()
2381            .try_collect()
2382            .await
2383            .unwrap();
2384
2385        // Rows 0 and 2 survive: their stored values (5, 8) pass through. The dropped
2386        // row's null (which would have fallen back to 100 + 1) is gone -- proving the
2387        // physical column and the positional fallback are filtered by the same selection.
2388        assert_row_id_column(&batches, &[Some(5), Some(8)]);
2389    }
2390
2391    #[tokio::test]
2392    async fn test_row_id_global_across_row_groups() {
2393        let tmp_dir = TempDir::new().unwrap();
2394        let dir = tmp_dir.path().to_str().unwrap();
2395
2396        // 5 rows written with max_row_group_size = 2 -> 3 row groups. `_pos` must be the
2397        // GLOBAL file position, so `_row_id` continues across row-group boundaries rather
2398        // than restarting per group (which would silently duplicate ids).
2399        let file_path = format!("{dir}/row_id_multi_rg.parquet");
2400        let field = Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2401            PARQUET_FIELD_ID_META_KEY.to_string(),
2402            "1".to_string(),
2403        )]));
2404        let arrow_schema = Arc::new(ArrowSchema::new(vec![field]));
2405        let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(Int32Array::from(
2406            vec![1, 2, 3, 4, 5],
2407        ))])
2408        .unwrap();
2409        let props = WriterProperties::builder()
2410            .set_compression(Compression::SNAPPY)
2411            .set_max_row_group_row_count(Some(2))
2412            .build();
2413        let file = File::create(&file_path).unwrap();
2414        let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2415        writer.write(&batch).unwrap();
2416        writer.close().unwrap();
2417
2418        let task = row_id_task(file_path, Some(0));
2419        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2420        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2421        let batches: Vec<RecordBatch> = reader
2422            .read(tasks)
2423            .unwrap()
2424            .stream()
2425            .try_collect()
2426            .await
2427            .unwrap();
2428
2429        assert_row_id_column(&batches, &[Some(0), Some(1), Some(2), Some(3), Some(4)]);
2430    }
2431
2432    #[tokio::test]
2433    async fn test_row_id_global_when_first_row_group_pruned() {
2434        use parquet::file::reader::{FileReader, SerializedFileReader};
2435
2436        let tmp_dir = TempDir::new().unwrap();
2437        let dir = tmp_dir.path().to_str().unwrap();
2438
2439        // 6 rows with max_row_group_size = 2 -> 3 row groups. A byte-range split that prunes
2440        // row group 0 reaches the reader via `with_row_groups()` -- a different path than the
2441        // `RowSelection` cases above. The survivors must keep their GLOBAL positions (starting
2442        // at 2), so a per-group RowNumber restart would surface as duplicate ids here.
2443        let file_path = format!("{dir}/row_id_prune_rg0.parquet");
2444        let field = Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2445            PARQUET_FIELD_ID_META_KEY.to_string(),
2446            "1".to_string(),
2447        )]));
2448        let arrow_schema = Arc::new(ArrowSchema::new(vec![field]));
2449        let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(Int32Array::from(
2450            vec![1, 2, 3, 4, 5, 6],
2451        ))])
2452        .unwrap();
2453        let props = WriterProperties::builder()
2454            .set_compression(Compression::SNAPPY)
2455            .set_max_row_group_row_count(Some(2))
2456            .build();
2457        let file = File::create(&file_path).unwrap();
2458        let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2459        writer.write(&batch).unwrap();
2460        writer.close().unwrap();
2461
2462        // A byte range starting just past row group 0 prunes it (its midpoint falls below
2463        // `start`) while keeping groups 1 and 2 (physical rows 2..6).
2464        let metadata = SerializedFileReader::new(File::open(&file_path).unwrap())
2465            .unwrap()
2466            .metadata()
2467            .clone();
2468        assert_eq!(metadata.num_row_groups(), 3);
2469        let start = 4 + metadata.row_group(0).compressed_size() as u64;
2470        let file_size = std::fs::metadata(&file_path).unwrap().len();
2471
2472        let task = row_id_task_with_options(file_path, Some(100), start, file_size - start, vec![]);
2473
2474        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2475        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2476        let batches: Vec<RecordBatch> = reader
2477            .read(tasks)
2478            .unwrap()
2479            .stream()
2480            .try_collect()
2481            .await
2482            .unwrap();
2483
2484        // Groups 1 and 2 survive at physical positions 2..6 -> _row_id 102..106, not a
2485        // per-group restart at 100.
2486        assert_row_id_column(&batches, &[Some(102), Some(103), Some(104), Some(105)]);
2487    }
2488
2489    /// Writes a positional delete file (`file_path` + `pos` reserved columns) marking the
2490    /// given `positions` of `data_file_path` as deleted.
2491    fn write_positional_delete(
2492        dir: &str,
2493        name: &str,
2494        data_file_path: &str,
2495        positions: &[i64],
2496    ) -> String {
2497        use arrow_array::StringArray;
2498
2499        let file_path_field =
2500            Field::new("file_path", DataType::Utf8, false).with_metadata(HashMap::from([(
2501                PARQUET_FIELD_ID_META_KEY.to_string(),
2502                "2147483546".to_string(),
2503            )]));
2504        let pos_field = Field::new("pos", DataType::Int64, false).with_metadata(HashMap::from([(
2505            PARQUET_FIELD_ID_META_KEY.to_string(),
2506            "2147483545".to_string(),
2507        )]));
2508        let arrow_schema = Arc::new(ArrowSchema::new(vec![file_path_field, pos_field]));
2509        let batch = RecordBatch::try_new(arrow_schema.clone(), vec![
2510            Arc::new(StringArray::from(vec![data_file_path; positions.len()])),
2511            Arc::new(Int64Array::from(positions.to_vec())),
2512        ])
2513        .unwrap();
2514
2515        let path = format!("{dir}/{name}");
2516        let file = File::create(&path).unwrap();
2517        let props = WriterProperties::builder()
2518            .set_compression(Compression::SNAPPY)
2519            .build();
2520        let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2521        writer.write(&batch).unwrap();
2522        writer.close().unwrap();
2523        path
2524    }
2525
2526    #[tokio::test]
2527    async fn test_row_id_survives_positional_delete() {
2528        use crate::spec::DataContentType;
2529
2530        let tmp_dir = TempDir::new().unwrap();
2531        let dir = tmp_dir.path().to_str().unwrap();
2532        // id = [1, 2, 3]; a positional delete drops the middle physical row (pos = 1).
2533        let data_path = write_plain_parquet(dir, "row_id_posdel_data.parquet", vec![], vec![]);
2534        let del_path = write_positional_delete(dir, "row_id_posdel.parquet", &data_path, &[1]);
2535
2536        let delete = FileScanTaskDeleteFile::builder()
2537            .with_file_path(del_path.clone())
2538            .with_file_size_in_bytes(std::fs::metadata(&del_path).unwrap().len())
2539            .with_file_type(DataContentType::PositionDeletes)
2540            .with_file_format(DataFileFormat::Parquet)
2541            .with_partition_spec_id(0)
2542            .build()
2543            .unwrap();
2544        let task = row_id_task_with_options(data_path, Some(100), 0, 0, vec![delete]);
2545
2546        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2547        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2548        let batches: Vec<RecordBatch> = reader
2549            .read(tasks)
2550            .unwrap()
2551            .stream()
2552            .try_collect()
2553            .await
2554            .unwrap();
2555
2556        // Survivors keep their ABSOLUTE positions -- [100, 102], not renumbered [100, 101].
2557        // Positional-delete selection reaches the reader via a different path than predicate
2558        // selection, so this covers it explicitly.
2559        assert_row_id_column(&batches, &[Some(100), Some(102)]);
2560    }
2561
2562    #[tokio::test]
2563    async fn test_row_id_name_collision_under_positional_fallback() {
2564        let tmp_dir = TempDir::new().unwrap();
2565        let dir = tmp_dir.path().to_str().unwrap();
2566
2567        // A file with NO embedded field ids (positional fallback) whose columns include one
2568        // literally named `_row_id` (user data). Projecting `_row_id` must NOT be rejected as
2569        // an unthreadable physical metadata column -- under fallback the reserved column is
2570        // synthesized and the same-named user column is just data.
2571        let file_path = format!("{dir}/fallback_row_id_name.parquet");
2572        let arrow_schema = Arc::new(ArrowSchema::new(vec![
2573            Field::new("id", DataType::Int32, false),
2574            Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true),
2575        ]));
2576        let batch = RecordBatch::try_new(arrow_schema.clone(), vec![
2577            Arc::new(Int32Array::from(vec![1, 2, 3])),
2578            Arc::new(Int64Array::from(vec![7i64, 8, 9])),
2579        ])
2580        .unwrap();
2581        let props = WriterProperties::builder()
2582            .set_compression(Compression::SNAPPY)
2583            .build();
2584        let file = File::create(&file_path).unwrap();
2585        let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2586        writer.write(&batch).unwrap();
2587        writer.close().unwrap();
2588
2589        let task = row_id_task(file_path, Some(100));
2590        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2591        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2592        let batches: Vec<RecordBatch> = reader
2593            .read(tasks)
2594            .unwrap()
2595            .stream()
2596            .try_collect()
2597            .await
2598            .unwrap();
2599
2600        // Not rejected; `_row_id` is synthesized as first_row_id + pos.
2601        assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]);
2602    }
2603
2604    #[tokio::test]
2605    async fn test_read_encrypted_parquet_with_wrong_key_fails() {
2606        let encryption_key = b"0123456789abcdef";
2607        let wrong_key = b"fedcba9876543210";
2608
2609        let schema = Arc::new(
2610            Schema::builder()
2611                .with_schema_id(1)
2612                .with_fields(vec![
2613                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2614                ])
2615                .build()
2616                .unwrap(),
2617        );
2618
2619        let arrow_schema = Arc::new(ArrowSchema::new(vec![
2620            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2621                PARQUET_FIELD_ID_META_KEY.to_string(),
2622                "1".to_string(),
2623            )])),
2624        ]));
2625
2626        let tmp_dir = TempDir::new().unwrap();
2627        let table_location = tmp_dir.path().to_str().unwrap().to_string();
2628        let file_io = FileIO::new_with_fs();
2629
2630        let id_data = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
2631        let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap();
2632
2633        let file_path = format!("{table_location}/encrypted_wrong_key.parquet");
2634        write_encrypted_parquet(&file_path, &batch, encryption_key, None);
2635
2636        let wrong_key_metadata = crate::encryption::StandardKeyMetadata::try_new(wrong_key)
2637            .unwrap()
2638            .encode()
2639            .unwrap();
2640
2641        let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
2642
2643        let task = FileScanTask::builder()
2644            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2645            .with_start(0)
2646            .with_length(0)
2647            .with_data_file_path(file_path)
2648            .with_data_file_format(DataFileFormat::Parquet)
2649            .with_schema(schema)
2650            .with_project_field_ids(vec![1])
2651            .with_case_sensitive(false)
2652            .with_key_metadata(Some(wrong_key_metadata))
2653            .build()
2654            .unwrap();
2655
2656        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2657        let result: Result<Vec<RecordBatch>, _> =
2658            reader.read(tasks).unwrap().stream().try_collect().await;
2659
2660        let err = result.unwrap_err();
2661        assert_eq!(err.kind(), crate::ErrorKind::Unexpected);
2662        let err_str = format!("{err}");
2663        assert!(
2664            err_str.contains("unable to decrypt parquet footer"),
2665            "Expected error about decryption failure, got: {err_str}"
2666        );
2667    }
2668
2669    /// Test that concurrency=1 reads all files correctly and in deterministic order.
2670    /// This verifies the fast-path optimization for single concurrency.
2671    #[tokio::test]
2672    async fn test_read_with_concurrency_one() {
2673        use arrow_array::Int32Array;
2674
2675        let schema = Arc::new(
2676            Schema::builder()
2677                .with_schema_id(1)
2678                .with_fields(vec![
2679                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2680                    NestedField::required(2, "file_num", Type::Primitive(PrimitiveType::Int))
2681                        .into(),
2682                ])
2683                .build()
2684                .unwrap(),
2685        );
2686
2687        let arrow_schema = Arc::new(ArrowSchema::new(vec![
2688            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2689                PARQUET_FIELD_ID_META_KEY.to_string(),
2690                "1".to_string(),
2691            )])),
2692            Field::new("file_num", DataType::Int32, false).with_metadata(HashMap::from([(
2693                PARQUET_FIELD_ID_META_KEY.to_string(),
2694                "2".to_string(),
2695            )])),
2696        ]));
2697
2698        let tmp_dir = TempDir::new().unwrap();
2699        let table_location = tmp_dir.path().to_str().unwrap().to_string();
2700        let file_io = FileIO::new_with_fs();
2701
2702        // Create 3 parquet files with different data
2703        let props = WriterProperties::builder()
2704            .set_compression(Compression::SNAPPY)
2705            .build();
2706
2707        for file_num in 0..3 {
2708            let id_data = Arc::new(Int32Array::from_iter_values(
2709                file_num * 10..(file_num + 1) * 10,
2710            )) as ArrayRef;
2711            let file_num_data = Arc::new(Int32Array::from(vec![file_num; 10])) as ArrayRef;
2712
2713            let to_write =
2714                RecordBatch::try_new(arrow_schema.clone(), vec![id_data, file_num_data]).unwrap();
2715
2716            let file = File::create(format!("{table_location}/file_{file_num}.parquet")).unwrap();
2717            let mut writer =
2718                ArrowWriter::try_new(file, to_write.schema(), Some(props.clone())).unwrap();
2719            writer.write(&to_write).expect("Writing batch");
2720            writer.close().unwrap();
2721        }
2722
2723        // Read with concurrency=1 (fast-path)
2724        let reader = ArrowReaderBuilder::new(file_io, Runtime::current())
2725            .with_data_file_concurrency_limit(1)
2726            .build();
2727
2728        // Create tasks in a specific order: file_0, file_1, file_2
2729        let tasks = vec![
2730            Ok(FileScanTask::builder()
2731                .with_file_size_in_bytes(
2732                    std::fs::metadata(format!("{table_location}/file_0.parquet"))
2733                        .unwrap()
2734                        .len(),
2735                )
2736                .with_start(0)
2737                .with_length(0)
2738                .with_data_file_path(format!("{table_location}/file_0.parquet"))
2739                .with_data_file_format(DataFileFormat::Parquet)
2740                .with_schema(schema.clone())
2741                .with_project_field_ids(vec![1, 2])
2742                .with_case_sensitive(false)
2743                .build()
2744                .unwrap()),
2745            Ok(FileScanTask::builder()
2746                .with_file_size_in_bytes(
2747                    std::fs::metadata(format!("{table_location}/file_1.parquet"))
2748                        .unwrap()
2749                        .len(),
2750                )
2751                .with_start(0)
2752                .with_length(0)
2753                .with_data_file_path(format!("{table_location}/file_1.parquet"))
2754                .with_data_file_format(DataFileFormat::Parquet)
2755                .with_schema(schema.clone())
2756                .with_project_field_ids(vec![1, 2])
2757                .with_case_sensitive(false)
2758                .build()
2759                .unwrap()),
2760            Ok(FileScanTask::builder()
2761                .with_file_size_in_bytes(
2762                    std::fs::metadata(format!("{table_location}/file_2.parquet"))
2763                        .unwrap()
2764                        .len(),
2765                )
2766                .with_start(0)
2767                .with_length(0)
2768                .with_data_file_path(format!("{table_location}/file_2.parquet"))
2769                .with_data_file_format(DataFileFormat::Parquet)
2770                .with_schema(schema.clone())
2771                .with_project_field_ids(vec![1, 2])
2772                .with_case_sensitive(false)
2773                .build()
2774                .unwrap()),
2775        ];
2776
2777        let tasks_stream = Box::pin(futures::stream::iter(tasks)) as FileScanTaskStream;
2778
2779        let result = reader
2780            .read(tasks_stream)
2781            .unwrap()
2782            .stream()
2783            .try_collect::<Vec<RecordBatch>>()
2784            .await
2785            .unwrap();
2786
2787        // Verify we got all 30 rows (10 from each file)
2788        let total_rows: usize = result.iter().map(|b| b.num_rows()).sum();
2789        assert_eq!(total_rows, 30, "Should have 30 total rows");
2790
2791        // Collect all ids and file_nums to verify data
2792        let mut all_ids = Vec::new();
2793        let mut all_file_nums = Vec::new();
2794
2795        for batch in &result {
2796            let id_col = batch
2797                .column(0)
2798                .as_primitive::<arrow_array::types::Int32Type>();
2799            let file_num_col = batch
2800                .column(1)
2801                .as_primitive::<arrow_array::types::Int32Type>();
2802
2803            for i in 0..batch.num_rows() {
2804                all_ids.push(id_col.value(i));
2805                all_file_nums.push(file_num_col.value(i));
2806            }
2807        }
2808
2809        assert_eq!(all_ids.len(), 30);
2810        assert_eq!(all_file_nums.len(), 30);
2811
2812        // With concurrency=1 and sequential processing, files should be processed in order
2813        // file_0: ids 0-9, file_num=0
2814        // file_1: ids 10-19, file_num=1
2815        // file_2: ids 20-29, file_num=2
2816        for i in 0..10 {
2817            assert_eq!(all_file_nums[i], 0, "First 10 rows should be from file_0");
2818            assert_eq!(all_ids[i], i as i32, "IDs should be 0-9");
2819        }
2820        for i in 10..20 {
2821            assert_eq!(all_file_nums[i], 1, "Next 10 rows should be from file_1");
2822            assert_eq!(all_ids[i], i as i32, "IDs should be 10-19");
2823        }
2824        for i in 20..30 {
2825            assert_eq!(all_file_nums[i], 2, "Last 10 rows should be from file_2");
2826            assert_eq!(all_ids[i], i as i32, "IDs should be 20-29");
2827        }
2828    }
2829
2830    #[tokio::test]
2831    async fn test_read_int96_timestamps_with_field_ids() {
2832        let schema = Arc::new(
2833            Schema::builder()
2834                .with_schema_id(1)
2835                .with_fields(vec![
2836                    NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2837                        .into(),
2838                    NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2839                ])
2840                .build()
2841                .unwrap(),
2842        );
2843
2844        let tmp_dir = TempDir::new().unwrap();
2845        let table_location = tmp_dir.path().to_str().unwrap().to_string();
2846        let (file_path, expected_micros) =
2847            write_int96_parquet_file(&table_location, "with_ids.parquet", true);
2848
2849        assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await;
2850    }
2851
2852    #[tokio::test]
2853    async fn test_read_int96_timestamps_without_field_ids() {
2854        let schema = Arc::new(
2855            Schema::builder()
2856                .with_schema_id(1)
2857                .with_fields(vec![
2858                    NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2859                        .into(),
2860                    NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2861                ])
2862                .build()
2863                .unwrap(),
2864        );
2865
2866        let tmp_dir = TempDir::new().unwrap();
2867        let table_location = tmp_dir.path().to_str().unwrap().to_string();
2868        let (file_path, expected_micros) =
2869            write_int96_parquet_file(&table_location, "no_ids.parquet", false);
2870
2871        assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await;
2872    }
2873
2874    #[tokio::test]
2875    async fn test_read_int96_timestamps_with_fallback_ids_and_pos() {
2876        use arrow_array::TimestampMicrosecondArray;
2877
2878        // Regression test for the combined path this refactor introduced: a field-id-less
2879        // file (positional fallback IDs) with an INT96 column and a `_pos` projection.
2880        // All three transforms -- field-ID assignment, INT96 coercion, and the row-number
2881        // virtual column -- apply in the single ArrowReaderMetadata rebuild.
2882        let schema = Arc::new(
2883            Schema::builder()
2884                .with_schema_id(1)
2885                .with_fields(vec![
2886                    NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2887                        .into(),
2888                    NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2889                ])
2890                .build()
2891                .unwrap(),
2892        );
2893
2894        let tmp_dir = TempDir::new().unwrap();
2895        let table_location = tmp_dir.path().to_str().unwrap().to_string();
2896        let (file_path, expected_micros) =
2897            write_int96_parquet_file(&table_location, "no_ids_with_pos.parquet", false);
2898
2899        let batches =
2900            read_int96_batches(&file_path, schema, vec![1, 2, RESERVED_FIELD_ID_POS]).await;
2901
2902        assert_eq!(batches.len(), 1);
2903        // The INT96 timestamps are coerced to micros...
2904        let ts_col = batches[0]
2905            .column_by_name("ts")
2906            .expect("ts column should be present")
2907            .as_any()
2908            .downcast_ref::<TimestampMicrosecondArray>()
2909            .expect("Expected TimestampMicrosecondArray");
2910        for (i, expected) in expected_micros.iter().enumerate() {
2911            assert_eq!(ts_col.value(i), *expected, "Row {i}");
2912        }
2913        // ...and `_pos`, materialized by the RowNumber virtual column, counts rows 0,1,2.
2914        let pos_col = batches[0]
2915            .column_by_name(RESERVED_COL_NAME_POS)
2916            .expect("_pos column should be present")
2917            .as_primitive::<arrow_array::types::Int64Type>();
2918        assert_eq!(pos_col.values(), &[0, 1, 2]);
2919    }
2920
2921    #[tokio::test]
2922    async fn test_read_int96_timestamps_with_field_ids_and_pos() {
2923        use arrow_array::TimestampMicrosecondArray;
2924
2925        let schema = Arc::new(
2926            Schema::builder()
2927                .with_schema_id(1)
2928                .with_fields(vec![
2929                    NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2930                        .into(),
2931                    NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2932                ])
2933                .build()
2934                .unwrap(),
2935        );
2936
2937        let tmp_dir = TempDir::new().unwrap();
2938        let table_location = tmp_dir.path().to_str().unwrap().to_string();
2939        let (file_path, expected_micros) =
2940            write_int96_parquet_file(&table_location, "with_ids_with_pos.parquet", true);
2941
2942        let batches =
2943            read_int96_batches(&file_path, schema, vec![1, 2, RESERVED_FIELD_ID_POS]).await;
2944
2945        assert_eq!(batches.len(), 1);
2946        let ts_col = batches[0]
2947            .column_by_name("ts")
2948            .expect("ts column should be present")
2949            .as_any()
2950            .downcast_ref::<TimestampMicrosecondArray>()
2951            .expect("Expected TimestampMicrosecondArray");
2952
2953        for (i, expected) in expected_micros.iter().enumerate() {
2954            assert_eq!(ts_col.value(i), *expected, "Row {i}");
2955        }
2956
2957        let pos_col = batches[0]
2958            .column_by_name(RESERVED_COL_NAME_POS)
2959            .expect("_pos column should be present")
2960            .as_primitive::<arrow_array::types::Int64Type>();
2961
2962        assert_eq!(pos_col.values(), &[0, 1, 2]);
2963    }
2964
2965    #[tokio::test]
2966    async fn test_read_int96_timestamps_in_struct() {
2967        use arrow_array::{StructArray, TimestampMicrosecondArray};
2968        use parquet::basic::{Repetition, Type as PhysicalType};
2969        use parquet::data_type::Int96Type;
2970        use parquet::file::writer::SerializedFileWriter;
2971        use parquet::schema::types::Type as SchemaType;
2972
2973        let tmp_dir = TempDir::new().unwrap();
2974        let table_location = tmp_dir.path().to_str().unwrap().to_string();
2975        let file_path = format!("{table_location}/struct_int96.parquet");
2976
2977        let ts_type = SchemaType::primitive_type_builder("ts", PhysicalType::INT96)
2978            .with_repetition(Repetition::OPTIONAL)
2979            .with_id(Some(2))
2980            .build()
2981            .unwrap();
2982
2983        let struct_type = SchemaType::group_type_builder("data")
2984            .with_repetition(Repetition::REQUIRED)
2985            .with_id(Some(1))
2986            .with_fields(vec![Arc::new(ts_type)])
2987            .build()
2988            .unwrap();
2989
2990        let parquet_schema = SchemaType::group_type_builder("schema")
2991            .with_fields(vec![Arc::new(struct_type)])
2992            .build()
2993            .unwrap();
2994
2995        let (int96_val, expected_micros) = make_int96_test_value();
2996
2997        let file = File::create(&file_path).unwrap();
2998        let mut writer =
2999            SerializedFileWriter::new(file, Arc::new(parquet_schema), Default::default()).unwrap();
3000
3001        // def=1: struct is REQUIRED so no level, ts is OPTIONAL and present (1).
3002        // No repetition levels needed (no repeated groups).
3003        let mut row_group = writer.next_row_group().unwrap();
3004        {
3005            let mut col = row_group.next_column().unwrap().unwrap();
3006            col.typed::<Int96Type>()
3007                .write_batch(&[int96_val], Some(&[1]), None)
3008                .unwrap();
3009            col.close().unwrap();
3010        }
3011        row_group.close().unwrap();
3012        writer.close().unwrap();
3013
3014        let iceberg_schema = Arc::new(
3015            Schema::builder()
3016                .with_schema_id(1)
3017                .with_fields(vec![
3018                    NestedField::required(
3019                        1,
3020                        "data",
3021                        Type::Struct(crate::spec::StructType::new(vec![
3022                            NestedField::optional(
3023                                2,
3024                                "ts",
3025                                Type::Primitive(PrimitiveType::Timestamp),
3026                            )
3027                            .into(),
3028                        ])),
3029                    )
3030                    .into(),
3031                ])
3032                .build()
3033                .unwrap(),
3034        );
3035
3036        let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await;
3037
3038        assert_eq!(batches.len(), 1);
3039        let struct_array = batches[0]
3040            .column(0)
3041            .as_any()
3042            .downcast_ref::<StructArray>()
3043            .expect("Expected StructArray");
3044        let ts_array = struct_array
3045            .column(0)
3046            .as_any()
3047            .downcast_ref::<TimestampMicrosecondArray>()
3048            .expect("Expected TimestampMicrosecondArray inside struct");
3049
3050        assert_eq!(
3051            ts_array.value(0),
3052            expected_micros,
3053            "INT96 in struct: got {}, expected {expected_micros}",
3054            ts_array.value(0)
3055        );
3056    }
3057
3058    #[tokio::test]
3059    async fn test_read_int96_timestamps_in_list() {
3060        use arrow_array::{ListArray, TimestampMicrosecondArray};
3061        use parquet::basic::{Repetition, Type as PhysicalType};
3062        use parquet::data_type::Int96Type;
3063        use parquet::file::writer::SerializedFileWriter;
3064        use parquet::schema::types::Type as SchemaType;
3065
3066        let tmp_dir = TempDir::new().unwrap();
3067        let table_location = tmp_dir.path().to_str().unwrap().to_string();
3068        let file_path = format!("{table_location}/list_int96.parquet");
3069
3070        // 3-level LIST encoding:
3071        //   optional group timestamps (LIST) {
3072        //     repeated group list {
3073        //       optional int96 element;
3074        //     }
3075        //   }
3076        let element_type = SchemaType::primitive_type_builder("element", PhysicalType::INT96)
3077            .with_repetition(Repetition::OPTIONAL)
3078            .with_id(Some(2))
3079            .build()
3080            .unwrap();
3081
3082        let list_group = SchemaType::group_type_builder("list")
3083            .with_repetition(Repetition::REPEATED)
3084            .with_fields(vec![Arc::new(element_type)])
3085            .build()
3086            .unwrap();
3087
3088        let list_type = SchemaType::group_type_builder("timestamps")
3089            .with_repetition(Repetition::OPTIONAL)
3090            .with_id(Some(1))
3091            .with_logical_type(Some(parquet::basic::LogicalType::List))
3092            .with_fields(vec![Arc::new(list_group)])
3093            .build()
3094            .unwrap();
3095
3096        let parquet_schema = SchemaType::group_type_builder("schema")
3097            .with_fields(vec![Arc::new(list_type)])
3098            .build()
3099            .unwrap();
3100
3101        let (int96_val, expected_micros) = make_int96_test_value();
3102
3103        let file = File::create(&file_path).unwrap();
3104        let mut writer =
3105            SerializedFileWriter::new(file, Arc::new(parquet_schema), Default::default()).unwrap();
3106
3107        // Write a single row with a list containing one INT96 element.
3108        // def=3: list present (1) + repeated group (2) + element present (3)
3109        // rep=0: start of a new list
3110        let mut row_group = writer.next_row_group().unwrap();
3111        {
3112            let mut col = row_group.next_column().unwrap().unwrap();
3113            col.typed::<Int96Type>()
3114                .write_batch(&[int96_val], Some(&[3]), Some(&[0]))
3115                .unwrap();
3116            col.close().unwrap();
3117        }
3118        row_group.close().unwrap();
3119        writer.close().unwrap();
3120
3121        let iceberg_schema = Arc::new(
3122            Schema::builder()
3123                .with_schema_id(1)
3124                .with_fields(vec![
3125                    NestedField::optional(
3126                        1,
3127                        "timestamps",
3128                        Type::List(crate::spec::ListType {
3129                            element_field: NestedField::optional(
3130                                2,
3131                                "element",
3132                                Type::Primitive(PrimitiveType::Timestamp),
3133                            )
3134                            .into(),
3135                        }),
3136                    )
3137                    .into(),
3138                ])
3139                .build()
3140                .unwrap(),
3141        );
3142
3143        let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await;
3144
3145        assert_eq!(batches.len(), 1);
3146        let list_array = batches[0]
3147            .column(0)
3148            .as_any()
3149            .downcast_ref::<ListArray>()
3150            .expect("Expected ListArray");
3151        let ts_array = list_array
3152            .values()
3153            .as_any()
3154            .downcast_ref::<TimestampMicrosecondArray>()
3155            .expect("Expected TimestampMicrosecondArray inside list");
3156
3157        assert_eq!(
3158            ts_array.value(0),
3159            expected_micros,
3160            "INT96 in list: got {}, expected {expected_micros}",
3161            ts_array.value(0)
3162        );
3163    }
3164
3165    #[tokio::test]
3166    async fn test_read_int96_timestamps_in_map() {
3167        use arrow_array::{MapArray, TimestampMicrosecondArray};
3168        use parquet::basic::{Repetition, Type as PhysicalType};
3169        use parquet::data_type::{ByteArrayType, Int96Type};
3170        use parquet::file::writer::SerializedFileWriter;
3171        use parquet::schema::types::Type as SchemaType;
3172
3173        let tmp_dir = TempDir::new().unwrap();
3174        let table_location = tmp_dir.path().to_str().unwrap().to_string();
3175        let file_path = format!("{table_location}/map_int96.parquet");
3176
3177        // MAP encoding:
3178        //   optional group ts_map (MAP) {
3179        //     repeated group key_value {
3180        //       required binary key (UTF8);
3181        //       optional int96 value;
3182        //     }
3183        //   }
3184        let key_type = SchemaType::primitive_type_builder("key", PhysicalType::BYTE_ARRAY)
3185            .with_repetition(Repetition::REQUIRED)
3186            .with_logical_type(Some(parquet::basic::LogicalType::String))
3187            .with_id(Some(2))
3188            .build()
3189            .unwrap();
3190
3191        let value_type = SchemaType::primitive_type_builder("value", PhysicalType::INT96)
3192            .with_repetition(Repetition::OPTIONAL)
3193            .with_id(Some(3))
3194            .build()
3195            .unwrap();
3196
3197        let key_value_group = SchemaType::group_type_builder("key_value")
3198            .with_repetition(Repetition::REPEATED)
3199            .with_fields(vec![Arc::new(key_type), Arc::new(value_type)])
3200            .build()
3201            .unwrap();
3202
3203        let map_type = SchemaType::group_type_builder("ts_map")
3204            .with_repetition(Repetition::OPTIONAL)
3205            .with_id(Some(1))
3206            .with_logical_type(Some(parquet::basic::LogicalType::Map))
3207            .with_fields(vec![Arc::new(key_value_group)])
3208            .build()
3209            .unwrap();
3210
3211        let parquet_schema = SchemaType::group_type_builder("schema")
3212            .with_fields(vec![Arc::new(map_type)])
3213            .build()
3214            .unwrap();
3215
3216        let (int96_val, expected_micros) = make_int96_test_value();
3217
3218        let file = File::create(&file_path).unwrap();
3219        let mut writer =
3220            SerializedFileWriter::new(file, Arc::new(parquet_schema), Default::default()).unwrap();
3221
3222        // Write a single row with a map containing one key-value pair.
3223        // rep=0 for both columns: start of a new map.
3224        // key def=2: map present (1) + key_value entry present (2), key is REQUIRED.
3225        // value def=3: map present (1) + key_value entry present (2) + value present (3).
3226        let mut row_group = writer.next_row_group().unwrap();
3227        {
3228            let mut col = row_group.next_column().unwrap().unwrap();
3229            col.typed::<ByteArrayType>()
3230                .write_batch(
3231                    &[parquet::data_type::ByteArray::from("event_time")],
3232                    Some(&[2]),
3233                    Some(&[0]),
3234                )
3235                .unwrap();
3236            col.close().unwrap();
3237        }
3238        {
3239            let mut col = row_group.next_column().unwrap().unwrap();
3240            col.typed::<Int96Type>()
3241                .write_batch(&[int96_val], Some(&[3]), Some(&[0]))
3242                .unwrap();
3243            col.close().unwrap();
3244        }
3245        row_group.close().unwrap();
3246        writer.close().unwrap();
3247
3248        let iceberg_schema = Arc::new(
3249            Schema::builder()
3250                .with_schema_id(1)
3251                .with_fields(vec![
3252                    NestedField::optional(
3253                        1,
3254                        "ts_map",
3255                        Type::Map(crate::spec::MapType {
3256                            key_field: NestedField::required(
3257                                2,
3258                                "key",
3259                                Type::Primitive(PrimitiveType::String),
3260                            )
3261                            .into(),
3262                            value_field: NestedField::optional(
3263                                3,
3264                                "value",
3265                                Type::Primitive(PrimitiveType::Timestamp),
3266                            )
3267                            .into(),
3268                        }),
3269                    )
3270                    .into(),
3271                ])
3272                .build()
3273                .unwrap(),
3274        );
3275
3276        let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await;
3277
3278        assert_eq!(batches.len(), 1);
3279        let map_array = batches[0]
3280            .column(0)
3281            .as_any()
3282            .downcast_ref::<MapArray>()
3283            .expect("Expected MapArray");
3284        let ts_array = map_array
3285            .values()
3286            .as_any()
3287            .downcast_ref::<TimestampMicrosecondArray>()
3288            .expect("Expected TimestampMicrosecondArray as map values");
3289
3290        assert_eq!(
3291            ts_array.value(0),
3292            expected_micros,
3293            "INT96 in map: got {}, expected {expected_micros}",
3294            ts_array.value(0)
3295        );
3296    }
3297
3298    /// Writes `id` (Int32) plus a wide string column (field id 2) whose bytes dominate
3299    /// the file, so that reading it is visible in `bytes_read`.
3300    ///
3301    /// `extra_fields`/`extra_columns` (e.g. a physical metadata leaf) are appended after
3302    /// the `id` and wide columns, mirroring `write_plain_parquet`'s shape.
3303    fn write_parquet_with_wide_column(
3304        dir: &str,
3305        name: &str,
3306        extra_fields: Vec<Field>,
3307        extra_columns: Vec<ArrayRef>,
3308    ) -> String {
3309        let wide_field =
3310            Field::new("wide", DataType::Utf8, false).with_metadata(HashMap::from([(
3311                PARQUET_FIELD_ID_META_KEY.to_string(),
3312                "2".to_string(),
3313            )]));
3314        // Varied bytes so the column chunk does not compress away under SNAPPY, keeping
3315        // the `bytes_read` difference between projecting it and not unambiguous.
3316        let wide_values: Vec<String> = (0..3)
3317            .map(|i| {
3318                (0..2048)
3319                    .map(|j| ((i * 2048 + j) % 251) as u8 as char)
3320                    .collect()
3321            })
3322            .collect();
3323
3324        let mut fields = vec![wide_field];
3325        fields.extend(extra_fields);
3326        let mut columns: Vec<ArrayRef> = vec![Arc::new(StringArray::from(wide_values))];
3327        columns.extend(extra_columns);
3328        write_plain_parquet(dir, name, fields, columns)
3329    }
3330
3331    /// Schema with `id` (field 1, Int) and `wide` (field 2, String), matching
3332    /// `write_parquet_with_wide_column`.
3333    fn id_and_wide_schema() -> SchemaRef {
3334        Arc::new(
3335            Schema::builder()
3336                .with_schema_id(1)
3337                .with_fields(vec![
3338                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3339                    NestedField::required(2, "wide", Type::Primitive(PrimitiveType::String)).into(),
3340                ])
3341                .build()
3342                .unwrap(),
3343        )
3344    }
3345
3346    /// Builds a scan task over `file_path` projecting `project_field_ids`.
3347    fn metadata_projection_task(
3348        file_path: String,
3349        schema: SchemaRef,
3350        project_field_ids: Vec<i32>,
3351    ) -> FileScanTask {
3352        metadata_projection_task_with_first_row_id(file_path, schema, project_field_ids, None)
3353    }
3354
3355    fn metadata_projection_task_with_first_row_id(
3356        file_path: String,
3357        schema: SchemaRef,
3358        project_field_ids: Vec<i32>,
3359        first_row_id: Option<i64>,
3360    ) -> FileScanTask {
3361        FileScanTask::builder()
3362            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3363            .with_start(0)
3364            .with_length(0)
3365            .with_data_file_path(file_path)
3366            .with_data_file_format(DataFileFormat::Parquet)
3367            .with_schema(schema)
3368            .with_project_field_ids(project_field_ids)
3369            .with_first_row_id(first_row_id)
3370            .with_case_sensitive(false)
3371            .build()
3372            .unwrap()
3373    }
3374
3375    /// Runs a single-task scan and returns the batches plus the bytes read from storage.
3376    async fn scan_task(task: FileScanTask) -> (Vec<RecordBatch>, u64) {
3377        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
3378        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
3379        let scan = reader.read(tasks).unwrap();
3380        let metrics = scan.metrics().clone();
3381        let batches = scan.stream().try_collect().await.unwrap();
3382        (batches, metrics.bytes_read())
3383    }
3384
3385    #[tokio::test]
3386    async fn test_pos_only_projection_reads_no_data_columns() {
3387        let tmp_dir = TempDir::new().unwrap();
3388        let dir = tmp_dir.path().to_str().unwrap();
3389
3390        let pos_only = metadata_projection_task(
3391            write_parquet_with_wide_column(dir, "pos_only.parquet", vec![], vec![]),
3392            id_and_wide_schema(),
3393            vec![RESERVED_FIELD_ID_POS],
3394        );
3395        let (batches, pos_only_bytes) = scan_task(pos_only).await;
3396
3397        // Only `_pos` is materialized -- no data columns.
3398        assert_eq!(batches[0].num_columns(), 1);
3399        let pos_col = batches[0]
3400            .column_by_name(RESERVED_COL_NAME_POS)
3401            .expect("_pos column should be present")
3402            .as_primitive::<arrow_array::types::Int64Type>();
3403        assert_eq!(pos_col.values(), &[0, 1, 2]);
3404
3405        // A scan of the same-shaped file that also projects the wide data column must read
3406        // materially more, proving the wide column chunk was not fetched above.
3407        let with_data = metadata_projection_task(
3408            write_parquet_with_wide_column(dir, "pos_only_ref.parquet", vec![], vec![]),
3409            id_and_wide_schema(),
3410            vec![2, RESERVED_FIELD_ID_POS],
3411        );
3412        let (_, with_data_bytes) = scan_task(with_data).await;
3413
3414        assert!(
3415            pos_only_bytes < with_data_bytes,
3416            "_pos-only scan should read fewer bytes than a scan of the wide column: \
3417             {pos_only_bytes} vs {with_data_bytes}"
3418        );
3419    }
3420
3421    #[tokio::test]
3422    async fn test_empty_projection_reads_no_data_columns() {
3423        let tmp_dir = TempDir::new().unwrap();
3424        let dir = tmp_dir.path().to_str().unwrap();
3425
3426        let empty = metadata_projection_task(
3427            write_parquet_with_wide_column(dir, "empty.parquet", vec![], vec![]),
3428            id_and_wide_schema(),
3429            vec![],
3430        );
3431        let (batches, empty_bytes) = scan_task(empty).await;
3432
3433        // No data columns come out, but the fixture's 3 rows are still counted.
3434        assert!(batches.iter().all(|b| b.num_columns() == 0));
3435        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
3436
3437        let with_data = metadata_projection_task(
3438            write_parquet_with_wide_column(dir, "empty_ref.parquet", vec![], vec![]),
3439            id_and_wide_schema(),
3440            vec![2],
3441        );
3442        let (_, with_data_bytes) = scan_task(with_data).await;
3443
3444        assert!(
3445            empty_bytes < with_data_bytes,
3446            "empty projection should read fewer bytes than the wide column: \
3447             {empty_bytes} vs {with_data_bytes}"
3448        );
3449    }
3450
3451    #[tokio::test]
3452    async fn test_empty_projection_counts_rows_after_positional_delete() {
3453        use crate::spec::DataContentType;
3454
3455        let tmp_dir = TempDir::new().unwrap();
3456        let dir = tmp_dir.path().to_str().unwrap();
3457        let data_path = write_plain_parquet(dir, "empty_posdel_data.parquet", vec![], vec![]);
3458        let del_path = write_positional_delete(dir, "empty_posdel.parquet", &data_path, &[1]);
3459        let delete = FileScanTaskDeleteFile::builder()
3460            .with_file_path(del_path.clone())
3461            .with_file_size_in_bytes(std::fs::metadata(&del_path).unwrap().len())
3462            .with_file_type(DataContentType::PositionDeletes)
3463            .with_file_format(DataFileFormat::Parquet)
3464            .with_partition_spec_id(0)
3465            .build()
3466            .unwrap();
3467        let task = FileScanTask::builder()
3468            .with_file_size_in_bytes(std::fs::metadata(&data_path).unwrap().len())
3469            .with_start(0)
3470            .with_length(0)
3471            .with_data_file_path(data_path)
3472            .with_data_file_format(DataFileFormat::Parquet)
3473            .with_schema(id_and_wide_schema())
3474            .with_project_field_ids(vec![])
3475            .with_deletes(vec![delete])
3476            .with_case_sensitive(false)
3477            .build()
3478            .unwrap();
3479
3480        let (batches, _) = scan_task(task).await;
3481
3482        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
3483    }
3484
3485    #[tokio::test]
3486    async fn test_empty_projection_counts_rows_after_predicate() {
3487        use crate::expr::{Bind, Reference};
3488        use crate::spec::Datum;
3489
3490        let tmp_dir = TempDir::new().unwrap();
3491        let dir = tmp_dir.path().to_str().unwrap();
3492        let file_path = write_plain_parquet(dir, "empty_predicate.parquet", vec![], vec![]);
3493        let schema = id_and_wide_schema();
3494        // `id` is filtered on but not projected, so the row filter reads it on its own mask.
3495        let bound = Reference::new("id")
3496            .not_equal_to(Datum::int(2))
3497            .bind(Arc::clone(&schema), false)
3498            .unwrap();
3499        let task = FileScanTask::builder()
3500            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3501            .with_start(0)
3502            .with_length(0)
3503            .with_data_file_path(file_path)
3504            .with_data_file_format(DataFileFormat::Parquet)
3505            .with_schema(schema)
3506            .with_project_field_ids(vec![])
3507            .with_predicate(Some(bound))
3508            .with_case_sensitive(false)
3509            .build()
3510            .unwrap();
3511
3512        let (batches, _) = scan_task(task).await;
3513
3514        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
3515    }
3516
3517    #[tokio::test]
3518    async fn test_pos_only_projection_keeps_absolute_pos_under_predicate() {
3519        use crate::expr::{Bind, Reference};
3520        use crate::spec::Datum;
3521
3522        let tmp_dir = TempDir::new().unwrap();
3523        let dir = tmp_dir.path().to_str().unwrap();
3524        // id = [1, 2, 3]; drop the middle physical row via a predicate + row selection.
3525        let file_path = write_plain_parquet(dir, "pos_only_predicate.parquet", vec![], vec![]);
3526
3527        let schema = Arc::new(
3528            Schema::builder()
3529                .with_schema_id(1)
3530                .with_fields(vec![
3531                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3532                ])
3533                .build()
3534                .unwrap(),
3535        );
3536        let bound = Reference::new("id")
3537            .not_equal_to(Datum::int(2))
3538            .bind(Arc::clone(&schema), false)
3539            .unwrap();
3540        let task = FileScanTask::builder()
3541            .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3542            .with_start(0)
3543            .with_length(0)
3544            .with_data_file_path(file_path)
3545            .with_data_file_format(DataFileFormat::Parquet)
3546            .with_schema(schema)
3547            .with_project_field_ids(vec![RESERVED_FIELD_ID_POS])
3548            .with_predicate(Some(bound))
3549            .with_case_sensitive(false)
3550            .build()
3551            .unwrap();
3552
3553        // Row selection must be enabled for the predicate to filter rows. The row filter
3554        // reads `id` for its own evaluation even though `id` is not projected; the surviving
3555        // rows must keep their ABSOLUTE positions (0 and 2), not renumbered (0 and 1).
3556        let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current())
3557            .with_row_selection_enabled(true)
3558            .build();
3559        let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
3560        let batches: Vec<RecordBatch> = reader
3561            .read(tasks)
3562            .unwrap()
3563            .stream()
3564            .try_collect()
3565            .await
3566            .unwrap();
3567
3568        let pos: Vec<i64> = batches
3569            .iter()
3570            .flat_map(|b| {
3571                b.column_by_name(RESERVED_COL_NAME_POS)
3572                    .expect("_pos column should be present")
3573                    .as_primitive::<arrow_array::types::Int64Type>()
3574                    .values()
3575                    .to_vec()
3576            })
3577            .collect();
3578        assert_eq!(pos, vec![0, 2]);
3579    }
3580
3581    #[tokio::test]
3582    async fn test_pos_and_file_projection() {
3583        use crate::metadata_columns::RESERVED_COL_NAME_FILE;
3584
3585        let tmp_dir = TempDir::new().unwrap();
3586        let dir = tmp_dir.path().to_str().unwrap();
3587        // The motivating row-lineage shape: a synthesized position column (mask -> none)
3588        // alongside a materialized per-file constant.
3589        let file_path = write_parquet_with_wide_column(dir, "pos_and_file.parquet", vec![], vec![]);
3590        let task = metadata_projection_task(file_path.clone(), id_and_wide_schema(), vec![
3591            RESERVED_FIELD_ID_POS,
3592            RESERVED_FIELD_ID_FILE,
3593        ]);
3594        let (batches, _) = scan_task(task).await;
3595
3596        // Both metadata columns materialize; no data column is read.
3597        assert_eq!(batches[0].num_columns(), 2);
3598        let pos_col = batches[0]
3599            .column_by_name(RESERVED_COL_NAME_POS)
3600            .expect("_pos column should be present")
3601            .as_primitive::<arrow_array::types::Int64Type>();
3602        assert_eq!(pos_col.values(), &[0, 1, 2]);
3603        let file_col = batches[0]
3604            .column_by_name(RESERVED_COL_NAME_FILE)
3605            .expect("_file column should be present");
3606        let file_col = cast(file_col, &DataType::Utf8).unwrap();
3607        let file_col = file_col.as_any().downcast_ref::<StringArray>().unwrap();
3608        assert_eq!(file_col.value(0), file_path);
3609    }
3610
3611    #[tokio::test]
3612    async fn test_pos_and_physical_seq_projection_reads_only_the_leaf() {
3613        use crate::metadata_columns::{
3614            RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
3615            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3616        };
3617
3618        // A v3 rewrite that carried rows forward stores `_last_updated_sequence_number`
3619        // per-row. Projecting only `_pos` + the sequence column must read just that one
3620        // physical leaf, not every data column.
3621        let tmp_dir = TempDir::new().unwrap();
3622        let dir = tmp_dir.path().to_str().unwrap();
3623
3624        // File: id (1), wide data column (2), physical _last_updated_sequence_number.
3625        let write = |name: &str| {
3626            write_parquet_with_wide_column(
3627                dir,
3628                name,
3629                vec![physical_last_updated_seq_field()],
3630                vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
3631            )
3632        };
3633
3634        let seq_task = |path: String, ids: Vec<i32>| {
3635            FileScanTask::builder()
3636                .with_file_size_in_bytes(std::fs::metadata(&path).unwrap().len())
3637                .with_start(0)
3638                .with_length(0)
3639                .with_data_file_path(path)
3640                .with_data_file_format(DataFileFormat::Parquet)
3641                .with_schema(id_and_wide_schema())
3642                .with_project_field_ids(ids)
3643                .with_first_row_id(Some(100))
3644                .with_data_sequence_number(Some(9))
3645                .with_case_sensitive(false)
3646                .build()
3647                .unwrap()
3648        };
3649
3650        let meta_only = seq_task(write("pos_seq.parquet"), vec![
3651            RESERVED_FIELD_ID_POS,
3652            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3653        ]);
3654        let (batches, meta_only_bytes) = scan_task(meta_only).await;
3655
3656        // `_pos` and the coalesced sequence column materialize; the wide column does not.
3657        let pos_col = batches[0]
3658            .column_by_name(RESERVED_COL_NAME_POS)
3659            .expect("_pos column should be present")
3660            .as_primitive::<arrow_array::types::Int64Type>();
3661        assert_eq!(pos_col.values(), &[0, 1, 2]);
3662        let seq_col = batches[0]
3663            .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
3664            .expect("_last_updated_sequence_number column should be present");
3665        let seq_col = cast(seq_col, &DataType::Int64).unwrap();
3666        let seq_col = seq_col.as_any().downcast_ref::<Int64Array>().unwrap();
3667        // Per-row stored value where non-null, else the data sequence number (9).
3668        assert_eq!(seq_col.value(0), 5);
3669        assert_eq!(seq_col.value(1), 9);
3670        assert_eq!(seq_col.value(2), 8);
3671        assert!(batches[0].column_by_name("wide").is_none());
3672
3673        // A scan that also projects the wide data column must read materially more,
3674        // proving the metadata-only scan pruned to just the sequence leaf.
3675        let with_data = seq_task(write("pos_seq_ref.parquet"), vec![
3676            2,
3677            RESERVED_FIELD_ID_POS,
3678            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3679        ]);
3680        let (_, with_data_bytes) = scan_task(with_data).await;
3681
3682        assert!(
3683            meta_only_bytes < with_data_bytes,
3684            "_pos + physical sequence scan should read fewer bytes than one that also \
3685             reads the wide column: {meta_only_bytes} vs {with_data_bytes}"
3686        );
3687    }
3688
3689    #[tokio::test]
3690    async fn test_seq_only_projection_reads_only_the_leaf() {
3691        use crate::metadata_columns::{
3692            RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
3693            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3694        };
3695
3696        // `_last_updated_sequence_number` alone (no `_pos`, no data column). This is a
3697        // metadata-only projection, so RowNumber drives the `none()` downgrade and the
3698        // physical leaf is unioned back in, pruning the read to just that leaf.
3699        let tmp_dir = TempDir::new().unwrap();
3700        let dir = tmp_dir.path().to_str().unwrap();
3701        let write = |name: &str| {
3702            write_parquet_with_wide_column(
3703                dir,
3704                name,
3705                vec![physical_last_updated_seq_field()],
3706                vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
3707            )
3708        };
3709        let seq_task = |path: String, ids: Vec<i32>| {
3710            FileScanTask::builder()
3711                .with_file_size_in_bytes(std::fs::metadata(&path).unwrap().len())
3712                .with_start(0)
3713                .with_length(0)
3714                .with_data_file_path(path)
3715                .with_data_file_format(DataFileFormat::Parquet)
3716                .with_schema(id_and_wide_schema())
3717                .with_project_field_ids(ids)
3718                .with_first_row_id(Some(100))
3719                .with_data_sequence_number(Some(9))
3720                .with_case_sensitive(false)
3721                .build()
3722                .unwrap()
3723        };
3724
3725        let meta_only = seq_task(write("seq_only.parquet"), vec![
3726            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3727        ]);
3728        let (batches, meta_only_bytes) = scan_task(meta_only).await;
3729
3730        let seq_col = batches[0]
3731            .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
3732            .expect("_last_updated_sequence_number column should be present");
3733        let seq_col = cast(seq_col, &DataType::Int64).unwrap();
3734        let seq_col = seq_col.as_any().downcast_ref::<Int64Array>().unwrap();
3735        assert_eq!(seq_col.value(0), 5);
3736        assert_eq!(seq_col.value(1), 9);
3737        assert_eq!(seq_col.value(2), 8);
3738        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3739        assert_eq!(total_rows, 3);
3740        assert!(batches[0].column_by_name("wide").is_none());
3741
3742        let with_data = seq_task(write("seq_only_ref.parquet"), vec![
3743            2,
3744            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3745        ]);
3746        let (_, with_data_bytes) = scan_task(with_data).await;
3747
3748        assert!(
3749            meta_only_bytes < with_data_bytes,
3750            "seq-only scan should read fewer bytes than one that also reads the wide \
3751             column: {meta_only_bytes} vs {with_data_bytes}"
3752        );
3753    }
3754
3755    #[tokio::test]
3756    async fn test_seq_only_null_first_row_id_reads_no_data_columns() {
3757        use crate::metadata_columns::{
3758            RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
3759            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3760        };
3761
3762        // Seq-only projection with a null first_row_id: the column is nulled and the
3763        // physical leaf is NOT read (the gated `coalesce_last_updated_seq_leaf` is None).
3764        // This is a metadata-only projection, so RowNumber supplies the row count and the
3765        // data columns are pruned; the row count is still 3 and the values all null.
3766        let tmp_dir = TempDir::new().unwrap();
3767        let dir = tmp_dir.path().to_str().unwrap();
3768        let seq_task = |name: &str, ids: Vec<i32>| {
3769            let file_path = write_parquet_with_wide_column(
3770                dir,
3771                name,
3772                vec![physical_last_updated_seq_field()],
3773                vec![Arc::new(Int64Array::from(vec![Some(5), Some(6), Some(7)])) as ArrayRef],
3774            );
3775            FileScanTask::builder()
3776                .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3777                .with_start(0)
3778                .with_length(0)
3779                .with_data_file_path(file_path)
3780                .with_data_file_format(DataFileFormat::Parquet)
3781                .with_schema(id_and_wide_schema())
3782                .with_project_field_ids(ids)
3783                .with_first_row_id(None)
3784                .with_data_sequence_number(Some(9))
3785                .with_case_sensitive(false)
3786                .build()
3787                .unwrap()
3788        };
3789
3790        let meta_only = seq_task("seq_only_null_first.parquet", vec![
3791            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3792        ]);
3793        let (batches, meta_only_bytes) = scan_task(meta_only).await;
3794
3795        assert_eq!(batches[0].num_columns(), 1);
3796        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3797        assert_eq!(total_rows, 3);
3798        let seq_col = batches[0]
3799            .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
3800            .expect("_last_updated_sequence_number column should be present");
3801        let seq_col = cast(seq_col, &DataType::Int64).unwrap();
3802        let seq_col = seq_col.as_any().downcast_ref::<Int64Array>().unwrap();
3803        assert!((0..3).all(|i| seq_col.is_null(i)));
3804
3805        let with_data = seq_task("seq_only_null_first_ref.parquet", vec![
3806            2,
3807            RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3808        ]);
3809        let (_, with_data_bytes) = scan_task(with_data).await;
3810
3811        assert!(
3812            meta_only_bytes < with_data_bytes,
3813            "seq-only scan with a null first_row_id should read fewer bytes than a scan of \
3814             the wide column: {meta_only_bytes} vs {with_data_bytes}"
3815        );
3816    }
3817
3818    #[tokio::test]
3819    async fn test_file_only_reads_no_data_columns() {
3820        use crate::metadata_columns::RESERVED_COL_NAME_FILE;
3821
3822        // `SELECT _file` materializes a per-file string constant with no physical column to
3823        // read. The metadata-only RowNumber counter sizes it, so the scan prunes the data
3824        // columns rather than reading them all just to recover the row count.
3825        let tmp_dir = TempDir::new().unwrap();
3826        let dir = tmp_dir.path().to_str().unwrap();
3827
3828        let file_path = write_parquet_with_wide_column(dir, "file_only.parquet", vec![], vec![]);
3829        let meta_only = metadata_projection_task(file_path.clone(), id_and_wide_schema(), vec![
3830            RESERVED_FIELD_ID_FILE,
3831        ]);
3832        let (batches, meta_only_bytes) = scan_task(meta_only).await;
3833
3834        assert_eq!(batches[0].num_columns(), 1);
3835        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3836        assert_eq!(total_rows, 3);
3837        let file_col = batches[0]
3838            .column_by_name(RESERVED_COL_NAME_FILE)
3839            .expect("_file column should be present");
3840        let file_col = cast(file_col, &DataType::Utf8).unwrap();
3841        let file_col = file_col.as_any().downcast_ref::<StringArray>().unwrap();
3842        assert_eq!(file_col.value(0), file_path);
3843
3844        let with_data = metadata_projection_task(
3845            write_parquet_with_wide_column(dir, "file_only_ref.parquet", vec![], vec![]),
3846            id_and_wide_schema(),
3847            vec![2, RESERVED_FIELD_ID_FILE],
3848        );
3849        let (_, with_data_bytes) = scan_task(with_data).await;
3850
3851        assert!(
3852            meta_only_bytes < with_data_bytes,
3853            "_file-only scan should read fewer bytes than a scan of the wide column: \
3854             {meta_only_bytes} vs {with_data_bytes}"
3855        );
3856    }
3857
3858    #[tokio::test]
3859    async fn test_spec_id_only_reads_no_data_columns() {
3860        use crate::metadata_columns::{RESERVED_COL_NAME_SPEC_ID, RESERVED_FIELD_ID_SPEC_ID};
3861        use crate::spec::{Literal, PartitionSpec, Struct, Transform};
3862
3863        // `SELECT _spec_id` materializes an int constant from the task's partition spec id,
3864        // with no physical column to read. The RowNumber counter sizes it, so the scan
3865        // prunes the data columns rather than reading them all to recover the row count.
3866        let tmp_dir = TempDir::new().unwrap();
3867        let dir = tmp_dir.path().to_str().unwrap();
3868        let schema = id_and_wide_schema();
3869        let spec = Arc::new(
3870            PartitionSpec::builder(schema.clone())
3871                .with_spec_id(7)
3872                .add_partition_field("id", "id", Transform::Identity)
3873                .unwrap()
3874                .build()
3875                .unwrap(),
3876        );
3877
3878        let spec_id_task = |file_path: String, project_field_ids: Vec<i32>| {
3879            FileScanTask::builder()
3880                .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3881                .with_start(0)
3882                .with_length(0)
3883                .with_data_file_path(file_path)
3884                .with_data_file_format(DataFileFormat::Parquet)
3885                .with_schema(schema.clone())
3886                .with_project_field_ids(project_field_ids)
3887                .with_partition(Some(Struct::from_iter([Some(Literal::int(42))])))
3888                .with_partition_spec(Some(spec.clone()))
3889                .with_case_sensitive(false)
3890                .build()
3891                .unwrap()
3892        };
3893
3894        let meta_only = spec_id_task(
3895            write_parquet_with_wide_column(dir, "spec_id_only.parquet", vec![], vec![]),
3896            vec![RESERVED_FIELD_ID_SPEC_ID],
3897        );
3898        let (batches, meta_only_bytes) = scan_task(meta_only).await;
3899
3900        assert_eq!(batches[0].num_columns(), 1);
3901        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3902        assert_eq!(total_rows, 3);
3903        let spec_id_col = batches[0]
3904            .column_by_name(RESERVED_COL_NAME_SPEC_ID)
3905            .expect("_spec_id column should be present");
3906        let spec_id_col = cast(spec_id_col, &DataType::Int32).unwrap();
3907        let spec_id_col = spec_id_col.as_any().downcast_ref::<Int32Array>().unwrap();
3908        assert_eq!(spec_id_col.values(), &[7, 7, 7]);
3909
3910        let with_data = spec_id_task(
3911            write_parquet_with_wide_column(dir, "spec_id_only_ref.parquet", vec![], vec![]),
3912            vec![2, RESERVED_FIELD_ID_SPEC_ID],
3913        );
3914        let (_, with_data_bytes) = scan_task(with_data).await;
3915
3916        assert!(
3917            meta_only_bytes < with_data_bytes,
3918            "_spec_id-only scan should read fewer bytes than a scan of the wide column: \
3919             {meta_only_bytes} vs {with_data_bytes}"
3920        );
3921    }
3922
3923    #[tokio::test]
3924    async fn test_empty_projection_preserves_row_count() {
3925        let tmp_dir = TempDir::new().unwrap();
3926        let dir = tmp_dir.path().to_str().unwrap();
3927        let file_path = write_plain_parquet(dir, "empty_projection.parquet", vec![], vec![]);
3928        let schema = Arc::new(
3929            Schema::builder()
3930                .with_schema_id(1)
3931                .with_fields(vec![
3932                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3933                ])
3934                .build()
3935                .unwrap(),
3936        );
3937        let task = metadata_projection_task(file_path, schema, vec![]);
3938        let (batches, _) = scan_task(task).await;
3939
3940        // A bare COUNT(*)-style empty projection must still report the row count.
3941        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3942        assert_eq!(total_rows, 3);
3943    }
3944
3945    /// The reads a data file needs must finish even when nothing polls the scan,
3946    /// and a dropped scan must stop reading (#3276).
3947    mod open_progress {
3948        use std::ops::Range;
3949        use std::sync::Mutex;
3950
3951        use async_trait::async_trait;
3952        use bytes::Bytes;
3953        use futures::StreamExt;
3954        use futures::stream::BoxStream;
3955        use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
3956        use serde::{Deserialize, Serialize};
3957
3958        use super::*;
3959        use crate::Result;
3960        use crate::arrow::scan_metrics::ScanMetrics;
3961        use crate::expr::{Bind, Predicate, Reference};
3962        use crate::io::{
3963            FileIOBuilder, FileMetadata, FileRead, FileWrite, InputFile, MemoryStorage, OutputFile,
3964            Storage, StorageConfig, StorageFactory,
3965        };
3966        use crate::scan::ArrowRecordBatchStream;
3967        use crate::spec::Datum;
3968
3969        const CONCURRENCY: usize = 2;
3970        /// Smaller than a footer, so one open takes more than one request.
3971        const METADATA_SIZE_HINT: usize = 64;
3972        /// Bytes one open reads: the hint, then the footer and its indexes.
3973        /// The footer stores the writer's `created_by` version, so a `parquet`
3974        /// upgrade can change this.
3975        const OPEN_BYTES: u64 = 547;
3976        /// Scheduler turns, far more than a request needs. A bound only because
3977        /// `#[tokio::test]` runs the test and the spawned reads on one thread.
3978        const YIELDS: usize = 100;
3979
3980        #[derive(Debug, Default)]
3981        struct Reads {
3982            last_start: Option<u64>,
3983            started: usize,
3984            finished: usize,
3985        }
3986
3987        /// What the storage was asked for, and how much of it came back.
3988        /// `ScanMetrics` counts a read when the request goes out, so it cannot
3989        /// tell a read that hangs from one that answers.
3990        #[derive(Debug, Clone, Default)]
3991        struct ReadLog(Arc<Mutex<Reads>>);
3992
3993        impl ReadLog {
3994            /// Where the newest read began, which says which read it is.
3995            fn last_start(&self) -> Option<u64> {
3996                self.0.lock().unwrap().last_start
3997            }
3998
3999            fn in_flight(&self) -> usize {
4000                let reads = self.0.lock().unwrap();
4001                reads.started - reads.finished
4002            }
4003        }
4004
4005        /// In-memory files whose reads take a scheduler turn to answer.
4006        #[derive(Debug, Clone, Default, Serialize, Deserialize)]
4007        struct DeferredStorage {
4008            #[serde(skip)]
4009            files: MemoryStorage,
4010            #[serde(skip)]
4011            log: ReadLog,
4012        }
4013
4014        struct DeferredRead {
4015            contents: Bytes,
4016            log: ReadLog,
4017        }
4018
4019        #[async_trait]
4020        impl FileRead for DeferredRead {
4021            async fn read(&self, range: Range<u64>) -> Result<Bytes> {
4022                {
4023                    let mut reads = self.log.0.lock().unwrap();
4024                    reads.last_start = Some(range.start);
4025                    reads.started += 1;
4026                }
4027                // The yield stands in for a request in flight: it finishes only
4028                // if something still drives the read.
4029                tokio::task::yield_now().await;
4030                self.log.0.lock().unwrap().finished += 1;
4031                Ok(self
4032                    .contents
4033                    .slice(range.start as usize..range.end as usize))
4034            }
4035        }
4036
4037        #[async_trait]
4038        #[typetag::serde]
4039        impl Storage for DeferredStorage {
4040            async fn reader(&self, path: &str) -> Result<Box<dyn FileRead>> {
4041                Ok(Box::new(DeferredRead {
4042                    contents: self.files.read(path).await?,
4043                    log: self.log.clone(),
4044                }))
4045            }
4046
4047            fn new_input(&self, path: &str) -> Result<InputFile> {
4048                Ok(InputFile::new(Arc::new(self.clone()), path.to_string()))
4049            }
4050
4051            async fn exists(&self, path: &str) -> Result<bool> {
4052                self.files.exists(path).await
4053            }
4054
4055            async fn metadata(&self, path: &str) -> Result<FileMetadata> {
4056                self.files.metadata(path).await
4057            }
4058
4059            async fn read(&self, path: &str) -> Result<Bytes> {
4060                self.files.read(path).await
4061            }
4062
4063            async fn write(&self, path: &str, bs: Bytes) -> Result<()> {
4064                self.files.write(path, bs).await
4065            }
4066
4067            async fn writer(&self, path: &str) -> Result<Box<dyn FileWrite>> {
4068                self.files.writer(path).await
4069            }
4070
4071            async fn delete(&self, path: &str) -> Result<()> {
4072                self.files.delete(path).await
4073            }
4074
4075            async fn delete_prefix(&self, path: &str) -> Result<()> {
4076                self.files.delete_prefix(path).await
4077            }
4078
4079            async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> {
4080                self.files.delete_stream(paths).await
4081            }
4082
4083            fn new_output(&self, path: &str) -> Result<OutputFile> {
4084                Ok(OutputFile::new(Arc::new(self.clone()), path.to_string()))
4085            }
4086        }
4087
4088        #[derive(Debug, Default, Serialize, Deserialize)]
4089        struct DeferredStorageFactory {
4090            #[serde(skip)]
4091            storage: DeferredStorage,
4092        }
4093
4094        #[typetag::serde]
4095        impl StorageFactory for DeferredStorageFactory {
4096            fn build(&self, _config: &StorageConfig) -> Result<Arc<dyn Storage>> {
4097                Ok(Arc::new(self.storage.clone()))
4098            }
4099        }
4100
4101        /// A scan of `files` copies of `contents`, the metrics that count the bytes
4102        /// it has read, and the log of what the storage saw.
4103        async fn deferred_scan_of(
4104            contents: Bytes,
4105            files: usize,
4106            predicate: Option<Predicate>,
4107        ) -> (ScanMetrics, ReadLog, ArrowRecordBatchStream) {
4108            let schema = Arc::new(
4109                Schema::builder()
4110                    .with_schema_id(1)
4111                    .with_fields(vec![
4112                        NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
4113                    ])
4114                    .build()
4115                    .unwrap(),
4116            );
4117            let bound = predicate.map(|p| p.bind(Arc::clone(&schema), false).unwrap());
4118
4119            let storage = DeferredStorage::default();
4120            let log = storage.log.clone();
4121            let mut tasks = Vec::with_capacity(files);
4122            for file_num in 0..files {
4123                let path = format!("memory:/data/file_{file_num}.parquet");
4124                storage.write(&path, contents.clone()).await.unwrap();
4125                tasks.push(
4126                    FileScanTask::builder()
4127                        .with_data_file_path(path)
4128                        .with_file_size_in_bytes(contents.len() as u64)
4129                        .with_start(0)
4130                        .with_length(0)
4131                        .with_data_file_format(DataFileFormat::Parquet)
4132                        .with_schema(Arc::clone(&schema))
4133                        .with_project_field_ids(vec![1])
4134                        .with_predicate(bound.clone())
4135                        .with_case_sensitive(false)
4136                        .build()
4137                        .unwrap(),
4138                );
4139            }
4140
4141            let file_io = FileIOBuilder::new(Arc::new(DeferredStorageFactory { storage })).build();
4142
4143            let reader = ArrowReaderBuilder::new(file_io, Runtime::current())
4144                .with_data_file_concurrency_limit(CONCURRENCY)
4145                .with_metadata_size_hint(METADATA_SIZE_HINT)
4146                .with_bloom_filter_enabled(bound.is_some())
4147                .build();
4148            let tasks =
4149                Box::pin(futures::stream::iter(tasks.into_iter().map(Ok))) as FileScanTaskStream;
4150            let result = reader.read(tasks).unwrap();
4151            (result.metrics().clone(), log, result.stream())
4152        }
4153
4154        /// `CONCURRENCY` identical files and no predicate, so the open is the only read.
4155        async fn deferred_scan() -> (ScanMetrics, ArrowRecordBatchStream) {
4156            let tmp_dir = TempDir::new().unwrap();
4157            let on_disk = write_plain_parquet(
4158                tmp_dir.path().to_str().unwrap(),
4159                "data.parquet",
4160                vec![],
4161                vec![],
4162            );
4163            let contents = Bytes::from(std::fs::read(on_disk).unwrap());
4164            let (metrics, _, scan) = deferred_scan_of(contents, CONCURRENCY, None).await;
4165            (metrics, scan)
4166        }
4167
4168        /// Gives the scheduler turns until `done` is true, or `YIELDS` turns pass.
4169        async fn spin(done: impl Fn() -> bool) {
4170            for _ in 0..YIELDS {
4171                if done() {
4172                    return;
4173                }
4174                tokio::task::yield_now().await;
4175            }
4176        }
4177
4178        /// An open must finish its requests even when the scan is never polled again.
4179        #[tokio::test]
4180        async fn test_open_proceeds_without_the_scan_being_polled() {
4181            let (metrics, scan) = deferred_scan().await;
4182            let mut scan = Box::pin(scan);
4183            let every_open = OPEN_BYTES * CONCURRENCY as u64;
4184
4185            // One poll starts an open per concurrency slot. Nothing polls the
4186            // scan after this, so only a spawned open can read any further.
4187            let _ = futures::poll!(scan.next());
4188            spin(|| metrics.bytes_read() >= every_open).await;
4189            assert_eq!(
4190                metrics.bytes_read(),
4191                every_open,
4192                "the opens stopped where the scan stopped polling"
4193            );
4194        }
4195
4196        /// A dropped scan must stop opening files.
4197        #[tokio::test]
4198        async fn test_dropped_scan_stops_opening_files() {
4199            let (metrics, scan) = deferred_scan().await;
4200            let mut scan = Box::pin(scan);
4201            let every_open = OPEN_BYTES * CONCURRENCY as u64;
4202
4203            // Drop while the opens are part way through, or there is nothing to abort.
4204            let _ = futures::poll!(scan.next());
4205            spin(|| metrics.bytes_read() > 0).await;
4206            let in_flight = metrics.bytes_read();
4207            assert!(in_flight > 0, "no open reached the storage");
4208            assert!(in_flight < every_open, "the opens finished before the drop");
4209
4210            drop(scan);
4211            spin(|| false).await;
4212            assert_eq!(
4213                metrics.bytes_read(),
4214                in_flight,
4215                "a dropped scan read {} more bytes",
4216                metrics.bytes_read() - in_flight
4217            );
4218        }
4219
4220        /// The bloom filter fetch is the next read after the open, and runs in the
4221        /// same buffered future, so it must finish on its own too.
4222        #[tokio::test]
4223        async fn test_bloom_filter_fetch_proceeds_without_the_scan_being_polled() {
4224            let tmp_dir = TempDir::new().unwrap();
4225            let props = WriterProperties::builder()
4226                .set_compression(Compression::SNAPPY)
4227                .set_bloom_filter_enabled(true)
4228                .build();
4229            let on_disk = write_parquet_with(
4230                tmp_dir.path().to_str().unwrap(),
4231                "bloom.parquet",
4232                vec![],
4233                vec![],
4234                props,
4235            );
4236            let bloom_offset =
4237                ParquetRecordBatchReaderBuilder::try_new(File::open(&on_disk).unwrap())
4238                    .unwrap()
4239                    .metadata()
4240                    .row_group(0)
4241                    .column(0)
4242                    .bloom_filter_offset()
4243                    .expect("the writer put a bloom filter on `id`") as u64;
4244
4245            let contents = Bytes::from(std::fs::read(&on_disk).unwrap());
4246            let predicate = Reference::new("id").equal_to(Datum::int(2));
4247            let (_, log, scan) = deferred_scan_of(contents, 1, Some(predicate)).await;
4248            let mut scan = Box::pin(scan);
4249
4250            // Poll only until the open is done and the bloom filter fetch is the
4251            // read in flight. Any later poll would drive that fetch itself.
4252            let at_the_bloom_filter = || log.last_start() == Some(bloom_offset);
4253            for _ in 0..YIELDS {
4254                if at_the_bloom_filter() {
4255                    break;
4256                }
4257                let _ = futures::poll!(scan.next());
4258                tokio::task::yield_now().await;
4259            }
4260            assert!(at_the_bloom_filter(), "the bloom filter was never read");
4261
4262            // Nothing polls the scan from here, so only a spawned fetch can finish.
4263            spin(|| log.in_flight() == 0).await;
4264            assert_eq!(
4265                log.in_flight(),
4266                0,
4267                "a read was left in flight where the scan stopped polling"
4268            );
4269        }
4270    }
4271}