moraine 0.6.2

A SlateDB-backed catalog for DuckLake
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

//! What a read fetches: the projection over a file's schema covering the
//! indexed and row-id columns, and how each row's id resolves.

use object_store::path::Path;
use parquet::arrow::ProjectionMask;

use crate::{
    data_file::{IndexProjection, ROW_ID_FIELD_ID, RowIdSource},
    error::{Error, Result},
};

/// The root position of the embedded row-id column, if the file has one.
pub(super) fn embedded_row_id_position(
    schema: &parquet::schema::types::SchemaDescriptor,
) -> Option<usize> {
    schema.root_schema().get_fields().iter().position(|field| {
        let info = field.get_basic_info();
        info.has_id() && u64::try_from(info.id()).ok() == Some(ROW_ID_FIELD_ID)
    })
}

/// Resolves `source` against the file's schema into the projection
/// inputs: the row-id column's position (when it is to be read) and the
/// dense start.
pub(super) fn resolve_row_id_source(
    schema: &parquet::schema::types::SchemaDescriptor,
    source: RowIdSource,
    path: &Path,
) -> Result<(Option<usize>, u64)> {
    match (source, embedded_row_id_position(schema)) {
        (RowIdSource::Ordinal, Some(_)) => Err(Error::Constraint(format!(
            "scoped read: {path} carries an embedded row-id column and cannot be read by ordinal"
        ))),
        (RowIdSource::Ordinal, None) => Ok((None, 0)),
        (RowIdSource::Resolve { .. }, Some(position)) => Ok((Some(position), 0)),
        (
            RowIdSource::Resolve {
                row_id_start: Some(start),
            },
            None,
        ) => Ok((None, start)),
        (RowIdSource::Resolve { row_id_start: None }, None) => Err(Error::Corruption(format!(
            "scoped read: {path} is recorded as carrying per-row ids but has no embedded \
             row-id column"
        ))),
    }
}

/// The projection over `schema` covering only the indexed and row-id
/// columns, with each requested position mapped to its index in the
/// projected output batch.
pub(super) fn projection(
    schema: &parquet::schema::types::SchemaDescriptor,
    indexed_positions: &[usize],
    row_id_position: Option<usize>,
) -> Result<(ProjectionMask, Vec<usize>, Option<usize>)> {
    let mut projected: Vec<usize> = indexed_positions.to_vec();
    if let Some(position) = row_id_position {
        projected.push(position);
    }
    projected.sort_unstable();
    projected.dedup();
    let mask = ProjectionMask::roots(schema, projected.iter().copied());
    // Output-batch column index for an original file position.
    let output_index = |position: usize| {
        projected
            .iter()
            .position(|&candidate| candidate == position)
            .ok_or_else(|| Error::Corruption("scoped read: projected column vanished".to_owned()))
    };
    let indexed_output = indexed_positions
        .iter()
        .map(|&position| output_index(position))
        .collect::<Result<Vec<_>>>()?;
    let row_id_output = row_id_position.map(output_index).transpose()?;

    Ok((mask, indexed_output, row_id_output))
}

/// Unique Arrow positions needed by every index plan, sorted.
pub(super) fn index_positions(projections: &[IndexProjection]) -> Vec<usize> {
    let mut positions: Vec<_> = projections
        .iter()
        .flat_map(|projection| projection.positions.iter().copied())
        .collect();
    positions.sort_unstable();
    positions.dedup();
    positions
}

/// Remaps plans from source-schema positions to columns in a projected
/// `RecordBatch`, refusing a plan whose directions or null orders do not
/// pair one-to-one with its positions.
pub(super) fn remap_index_projections(
    projections: Vec<IndexProjection>,
    source_positions: &[usize],
    output_positions: &[usize],
) -> Result<Vec<IndexProjection>> {
    projections
        .into_iter()
        .map(|mut projection| {
            let columns = projection.positions.len();
            if projection.directions.len() != columns || projection.nulls.len() != columns {
                return Err(Error::Corruption(format!(
                    "scoped read: index {} names {columns} columns but {} directions and {} \
                     null orders",
                    projection.index_id,
                    projection.directions.len(),
                    projection.nulls.len()
                )));
            }
            for position in &mut projection.positions {
                let source = source_positions.binary_search(position).map_err(|_| {
                    Error::Corruption(
                        "scoped read: index column vanished from projection".to_owned(),
                    )
                })?;
                *position = output_positions.get(source).copied().ok_or_else(|| {
                    Error::Corruption(
                        "scoped read: projected index column has no output".to_owned(),
                    )
                })?;
            }

            Ok(projection)
        })
        .collect()
}