moraine 0.6.5

//! Reading DuckLake's data files.
//!
//! A scoped read derives index entries from a registered Parquet file by
//! fetching only the indexed columns, the row positions, and — when the file
//! carries one — the row-id column. A bounded, merge-free projection, not
//! the scan path the no-Parquet-read rule guards.
//!
//! This is the only layer that touches Parquet. [`crate::store`] reads the
//! catalog's own keyspace out of SlateDB, and [`crate::catalog`] interprets
//! what either returns.

mod auxiliary_cache;
mod columns;
mod data_store;
mod delete_file;
mod entries;
mod inline_batch;
mod inline_rows;
mod located_rows;
mod metrics;
mod reader;
mod row_location;
mod row_set;
mod schema;
mod selection;
mod values;

#[cfg(test)]
mod auxiliary_cache_tests;
#[cfg(test)]
mod row_location_tests;
#[cfg(test)]
mod row_set_tests;
#[cfg(test)]
mod tests;

use std::{sync::Arc, time::Instant};

use bytes::Bytes;
use futures::{
    StreamExt, TryStreamExt,
    stream::{self, BoxStream},
};
use object_store::path::Path;
use parquet::{
    arrow::{
        arrow_reader::ArrowReaderOptions,
        async_reader::{AsyncFileReader, ParquetRecordBatchStreamBuilder},
    },
    file::metadata::PageIndexPolicy,
};

pub use crate::data_file::data_store::DataStore;
#[cfg(test)]
pub(crate) use crate::data_file::inline_batch::{
    inline_batch_decode_count, inline_schema_decode_count,
};
pub(crate) use crate::data_file::{
    auxiliary_cache::{
        install as install_auxiliary, occupancy as auxiliary_occupancy, row_summary_occupancy,
    },
    delete_file::{delete_file_positions, delete_file_positions_at},
    inline_batch::{decode_inline_schema, inline_batch_entries, inline_batch_index_entries},
    inline_rows::InlineRows,
    located_rows::{encode_located_batch, inline_rows_batch, scoped_read_row_batches},
    metrics::{DataStoreCounters, ScopedReadMetrics, ScopedReadTally, run_bounded_index_encoding},
    row_location::{FileSummary, file_summary},
    schema::ReadColumn,
    selection::{RowPositions, ScopedRows},
};
use crate::{
    data_file::{
        columns::{
            embedded_row_id_position, index_positions, projection, remap_index_projections,
            resolve_row_id_source,
        },
        entries::{record_batch_entries, record_batch_index_entries},
        metrics::index_encoding_concurrency,
        reader::ObjectStoreReader,
        selection::scoped_selection,
    },
    error::{Error, Result},
    store::index_encoding::{Direction, IndexKeyValue, NullOrder},
};

/// DuckDB's reserved Parquet field id for `_ducklake_internal_row_id`.
pub(crate) const ROW_ID_FIELD_ID: u64 = 2_147_483_540;

/// One entry derived from a registered file: the row id and the canonical
/// values of the indexed columns, in the index's column order.
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct ScopedReadEntry {
    /// Physical position of this row in its Parquet file.
    pub(crate) ordinal: u64,
    /// The row this entry points at.
    pub(crate) row_id: u64,
    /// The indexed column values; a `None` is SQL NULL (stored multi-shaped).
    pub(crate) values: Vec<Option<IndexKeyValue>>,
}

/// One index's projection over an Arrow batch. Positions name Arrow columns
/// in index-column order; ordering vectors run in parallel to them.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct IndexProjection {
    /// Persistent id used in the physical entry prefix.
    pub(crate) index_id: u64,
    /// Whether non-NULL values use the unique entry shape.
    pub(crate) unique: bool,
    /// Arrow column positions in index-column order.
    pub(crate) positions: Vec<usize>,
    /// Per-column sort directions.
    pub(crate) directions: Vec<Direction>,
    /// Per-column NULL placement.
    pub(crate) nulls: Vec<NullOrder>,
}

/// A final physical index key derived directly from an Arrow row.
pub(crate) struct ScopedIndexEntry {
    /// Position of the owning index in the supplied projection plans.
    pub(crate) index: usize,
    /// Physical position of the row in its Parquet file.
    pub(crate) ordinal: u64,
    /// Stable row id carried by the entry.
    pub(crate) row_id: u64,
    /// Fully encoded SlateDB entry key.
    pub(crate) key: Bytes,
    /// Whether the key uses the unique physical shape.
    pub(crate) unique: bool,
}

/// How a scoped read resolves each row's id.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RowIdSource {
    /// The embedded row-id column when the file carries one (even if a dense
    /// start is recorded), else `start + ordinal`, else refusal.
    Resolve {
        /// The catalog row's dense start, if it records one.
        row_id_start: Option<u64>,
    },
    /// Row ids are file ordinals from 0; a file carrying the embedded
    /// column is refused.
    Ordinal,
}

/// A `map_err` builder turning any displayable reader error into
/// `Error::Corruption("{what}: {err}")`.
fn corrupt<E: std::fmt::Display>(what: &'static str) -> impl Fn(E) -> Error {
    move |err| Error::Corruption(format!("{what}: {err}"))
}

/// A `usize` row count or offset as a `u64`.
fn usize_as_u64(value: usize) -> u64 {
    // `usize` is at most 64 bits on every supported target.
    #[allow(clippy::expect_used)]
    u64::try_from(value).expect("usize fits in u64 on supported targets")
}

/// Reads entries from one recorded Parquet object.
pub(crate) async fn scoped_read_recorded_entries(
    file: ParquetFile,
    indexed_positions: &[usize],
    rows: ScopedRows<'_>,
    row_id_source: RowIdSource,
) -> Result<Vec<ScopedReadEntry>> {
    scoped_read_entry_batches(file, indexed_positions, rows, row_id_source)
        .await?
        .try_fold(Vec::new(), |mut entries, batch| async move {
            entries.extend(batch);
            Ok(entries)
        })
        .await
}

/// One immutable Parquet object's location and recorded sizes.
#[derive(Clone)]
pub(crate) struct ParquetFile {
    store: DataStore,
    path: Path,
    file_size: u64,
    footer_size: u64,
    metrics: Arc<ScopedReadMetrics>,
    columns: Option<Arc<Vec<ReadColumn>>>,
    single_touch: bool,
}

impl ParquetFile {
    /// Describes one object using DuckLake's recorded sizes.
    pub(crate) fn new(store: DataStore, path: Path, file_size: u64, footer_size: u64) -> Self {
        Self {
            store,
            path,
            file_size,
            footer_size,
            metrics: Arc::new(ScopedReadMetrics::default()),
            columns: None,
            single_touch: false,
        }
    }

    /// Marks this file's data ranges as read once: they are fetched but not
    /// retained, while its footer stays cached as for any other read.
    pub(crate) fn single_touch(mut self) -> Self {
        self.single_touch = true;
        self
    }

    /// Resolves logical column positions against this file's physical schema.
    pub(crate) fn with_columns(mut self, columns: Vec<ReadColumn>) -> Self {
        self.columns = Some(Arc::new(columns));
        self
    }

    /// Records this file's reads in the supplied commit-wide tally.
    pub(crate) fn with_metrics(mut self, metrics: Arc<ScopedReadMetrics>) -> Self {
        self.metrics = metrics;
        self
    }

    /// The tally this file's reads are recorded in.
    #[cfg(test)]
    pub(crate) fn metrics(&self) -> &Arc<ScopedReadMetrics> {
        &self.metrics
    }
}

/// Row counts of `file`'s row groups, in file order.
pub(crate) async fn row_group_row_counts(file: &ParquetFile) -> Result<Vec<u64>> {
    let reader = ObjectStoreReader::new(file, PageIndexPolicy::Skip);
    let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Skip);
    let builder = ParquetRecordBatchStreamBuilder::new_with_options(reader, options)
        .await
        .map_err(corrupt("row-group probe"))?;

    selection::row_group_row_counts(builder.metadata(), &file.path)
}

/// Whether `file` carries the reserved embedded row-id column.
pub(crate) async fn carries_embedded_row_ids(file: ParquetFile) -> Result<bool> {
    let reader = ObjectStoreReader::new(&file, PageIndexPolicy::Skip);
    let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Skip);
    let builder = ParquetRecordBatchStreamBuilder::new_with_options(reader, options)
        .await
        .map_err(corrupt("row-id probe"))?;

    Ok(embedded_row_id_position(builder.parquet_schema()).is_some())
}

/// Loads `file`'s footer into the process cache, with its page index when
/// `page_index`, so a reader opened on the file afterwards finds it
/// resident.
pub(crate) async fn load_metadata(file: &ParquetFile, page_index: bool) -> Result<()> {
    let policy = if page_index {
        PageIndexPolicy::Optional
    } else {
        PageIndexPolicy::Skip
    };

    ObjectStoreReader::new(file, policy)
        .get_metadata(None)
        .await
        .map(drop)
        .map_err(corrupt("footer read"))
}

/// Rows decoded at once by a streamed scoped read.
pub(crate) const BUILD_READ_BATCH_ROWS: usize = 8_192;

/// Streams a file's fused index entries in bounded Arrow batches. Shared
/// columns are projected once and borrowed directly into their final keys.
pub(crate) async fn scoped_read_index_entry_batches(
    file: ParquetFile,
    projections: Vec<IndexProjection>,
    rows: ScopedRows<'_>,
    row_id_source: RowIdSource,
    excluded_ordinals: Option<&RowPositions>,
) -> Result<BoxStream<'static, Result<Vec<ScopedIndexEntry>>>> {
    if rows.is_empty() || projections.is_empty() {
        return Ok(stream::empty().boxed());
    }

    file.metrics.parquet_file();
    let source_positions = index_positions(&projections);
    let excluded_ordinals = excluded_ordinals.cloned();

    let reader = ObjectStoreReader::new(&file, rows.page_index_policy());
    let options = ArrowReaderOptions::new().with_page_index_policy(rows.page_index_policy());
    let builder = ParquetRecordBatchStreamBuilder::new_with_options(reader, options)
        .await
        .map_err(corrupt("scoped read"))?;
    let (row_id_position, row_id_start) =
        resolve_row_id_source(builder.parquet_schema(), row_id_source, &file.path)?;
    let (mask, indexed_output, row_id_output, normalization) = read_projection(
        &builder,
        file.columns.as_deref(),
        &source_positions,
        row_id_position,
    )?;
    let projections = Arc::new(remap_index_projections(
        projections,
        &source_positions,
        &indexed_output,
    )?);
    let (scope, ordinals) = scoped_selection(rows, builder.metadata(), &file.path)?;
    let ordinals = Arc::new(ordinals);
    let builder = scope.narrow(
        builder
            .with_projection(mask)
            .with_batch_size(BUILD_READ_BATCH_ROWS),
    );
    let arrow_reader = builder.build().map_err(corrupt("scoped read"))?;
    let mut emitted = 0usize;

    let metrics = Arc::clone(&file.metrics);
    Ok(arrow_reader
        .map(move |batch| {
            let batch = batch.map_err(corrupt("scoped read"))?;
            let batch = normalize_batch(batch, normalization.as_ref())?;
            file.metrics.arrow_batch();
            let batch_start = emitted;
            emitted = emitted.saturating_add(batch.num_rows());
            Result::Ok((batch, batch_start))
        })
        .map(move |prepared| {
            let projections = Arc::clone(&projections);
            let ordinals = Arc::clone(&ordinals);
            let excluded_ordinals = excluded_ordinals.clone();
            let metrics = Arc::clone(&metrics);
            async move {
                let (batch, batch_start) = prepared?;
                run_bounded_index_encoding(move || {
                    let started = Instant::now();
                    let result = record_batch_index_entries(
                        &batch,
                        &projections,
                        row_id_output,
                        row_id_start,
                        ordinals.borrowed(),
                        batch_start,
                        excluded_ordinals.as_ref(),
                        None,
                    );
                    metrics.encoded(started.elapsed());
                    result
                })
                .await
            }
        })
        .buffered(index_encoding_concurrency())
        .boxed())
}

/// Streams a file's projected index entries in bounded Arrow batches.
pub(crate) async fn scoped_read_entry_batches(
    file: ParquetFile,
    indexed_positions: &[usize],
    rows: ScopedRows<'_>,
    row_id_source: RowIdSource,
) -> Result<BoxStream<'static, Result<Vec<ScopedReadEntry>>>> {
    if rows.is_empty() {
        return Ok(stream::empty().boxed());
    }

    let reader = ObjectStoreReader::new(&file, rows.page_index_policy());
    let options = ArrowReaderOptions::new().with_page_index_policy(rows.page_index_policy());
    let builder = ParquetRecordBatchStreamBuilder::new_with_options(reader, options)
        .await
        .map_err(corrupt("scoped read"))?;
    let (scope, ordinals) = scoped_selection(rows, builder.metadata(), &file.path)?;

    let (row_id_position, row_id_start) =
        resolve_row_id_source(builder.parquet_schema(), row_id_source, &file.path)?;
    let (mask, indexed_output, row_id_output, normalization) = read_projection(
        &builder,
        file.columns.as_deref(),
        indexed_positions,
        row_id_position,
    )?;
    let builder = scope.narrow(
        builder
            .with_projection(mask)
            .with_batch_size(BUILD_READ_BATCH_ROWS),
    );
    let arrow_reader = builder.build().map_err(corrupt("scoped read"))?;

    let mut emitted = 0usize;
    let selection_stream = arrow_reader.map(move |batch| {
        let batch = batch.map_err(corrupt("scoped read"))?;
        let batch = normalize_batch(batch, normalization.as_ref())?;
        let batch_start = emitted;
        emitted = emitted.saturating_add(batch.num_rows());

        record_batch_entries(
            &batch,
            &indexed_output,
            row_id_output,
            row_id_start,
            ordinals.borrowed(),
            batch_start,
        )
    });

    Ok(selection_stream.boxed())
}

fn normalize_batch(
    batch: arrow::array::RecordBatch,
    projection: Option<&schema::BatchProjection>,
) -> Result<arrow::array::RecordBatch> {
    match projection {
        Some(projection) => projection.apply(&batch),
        None => Ok(batch),
    }
}

#[allow(clippy::type_complexity)]
fn read_projection(
    builder: &ParquetRecordBatchStreamBuilder<ObjectStoreReader>,
    columns: Option<&Vec<ReadColumn>>,
    requested: &[usize],
    row_id: Option<usize>,
) -> Result<(
    parquet::arrow::ProjectionMask,
    Vec<usize>,
    Option<usize>,
    Option<schema::BatchProjection>,
)> {
    let Some(columns) = columns else {
        let (mask, output, row_id) = projection(builder.parquet_schema(), requested, row_id)?;
        return Ok((mask, output, row_id, None));
    };
    let (mut normalization, positions) =
        schema::BatchProjection::resolve(builder.schema(), columns, requested, row_id)?;
    let (mask, _, _) = projection(builder.parquet_schema(), &positions, None)?;
    normalization.remap(&positions)?;
    Ok((
        mask,
        (0..requested.len()).collect(),
        row_id.map(|_| requested.len()),
        Some(normalization),
    ))
}