moraine 0.7.0

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
129
130
131
132
133
134
135
136

//! 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();
    if projected
        .iter()
        .any(|&position| position >= schema.root_schema().get_fields().len())
    {
        return Err(Error::Corruption(
            "scoped read: projected column is out of bounds".to_owned(),
        ));
    }
    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()
}