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
137
138
139
140
141
142
143
144
145
146

//! Locating stable row ids within one immutable data file.

use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};

use futures::TryStreamExt;

use crate::{
    data_file::{
        ParquetFile, RowIdSource, ScopedRows, auxiliary_cache, carries_embedded_row_ids,
        row_set::{FileRowSet, PositionedRowSet, RowOrder},
        scoped_read_entry_batches,
    },
    error::Result,
};

/// One file's row-id membership, and what it cost to obtain.
#[derive(Clone)]
pub(crate) struct FileSummary {
    rows: Arc<PositionedRowSet>,
    /// Whether this call read the file's row-id column and cached the
    /// result.
    pub(crate) built: bool,
}

impl FileSummary {
    /// Estimated resident bytes of the membership set.
    pub(crate) fn estimated_bytes(&self) -> u64 {
        self.rows.estimated_bytes()
    }

    /// Verified dense membership, irrespective of how the file records row IDs.
    pub(crate) fn dense_range(&self) -> Option<std::ops::Range<u64>> {
        match &self.rows.rows {
            FileRowSet::Range { start, end } => Some(*start..*end),
            _ => None,
        }
    }

    /// Which of `requested` this file holds, in request order.
    #[cfg(test)]
    pub(crate) fn matching(&self, requested: &[u64]) -> Vec<u64> {
        self.rows.rows.matching(requested)
    }

    pub(crate) fn contains(&self, row_id: u64) -> bool {
        self.rows.rows.contains(row_id)
    }

    /// The least and greatest row id this file holds, `None` when empty.
    pub(crate) fn bounds(&self) -> Option<(u64, u64)> {
        self.rows.rows.bounds()
    }

    /// File positions of `requested` rows within this file; `None` for rows
    /// this file does not hold.
    #[cfg(test)]
    pub(crate) fn positions_of(&self, requested: &[u64]) -> Vec<Option<u64>> {
        self.rows.positions_of(requested)
    }

    /// Visits every physical version of `row_id`; false means it is absent.
    pub(crate) fn visit_positions(&self, row_id: u64, visit: impl FnMut(u64)) -> bool {
        self.rows.visit_positions(row_id, visit)
    }
}

/// This file's row-id membership, from the cache when it is resident. A
/// file whose ids are the recorded dense range answers from its footer and
/// is remembered as that range; one carrying the reserved row-id column
/// reads only that column and is cached.
pub(crate) async fn file_summary(
    file: ParquetFile,
    table_id: u64,
    data_file_id: u64,
    row_id_start: Option<u64>,
    record_count: u64,
) -> Result<FileSummary> {
    let path = file.path.clone();
    let file_size = file.file_size;
    let store = file.store.clone();
    let key = auxiliary_cache::FileSummaryKey {
        table_id,
        data_file_id,
        path: &path,
        file_size,
    };

    if let Some(rows) = auxiliary_cache::shared().summary(&store, &key).await {
        return Ok(FileSummary { rows, built: false });
    }

    if let Some(start) = row_id_start
        && !carries_embedded_row_ids(file.clone()).await?
    {
        let rows = Arc::new(PositionedRowSet {
            rows: FileRowSet::range(start, record_count)?,
            order: RowOrder::Ascending,
        });
        auxiliary_cache::shared().insert_summary(&store, &key, &rows);

        Ok(FileSummary { rows, built: false })
    } else {
        let built = Arc::new(AtomicBool::new(false));
        let read = {
            let built = Arc::clone(&built);
            move || async move {
                built.store(true, Ordering::Relaxed);
                read_row_ids(file, row_id_start).await
            }
        };

        let rows = auxiliary_cache::shared()
            .fetch_summary(&store, &key, read)
            .await?;

        Ok(FileSummary {
            rows,
            built: built.load(Ordering::Relaxed),
        })
    }
}

async fn read_row_ids(
    file: ParquetFile,
    row_id_start: Option<u64>,
) -> Result<Arc<PositionedRowSet>> {
    let row_ids_in_file_order = scoped_read_entry_batches(
        file,
        &[],
        ScopedRows::All,
        RowIdSource::Resolve { row_id_start },
    )
    .await?
    .try_fold(Vec::new(), |mut row_ids, batch| async move {
        row_ids.extend(batch.iter().map(|entry| entry.row_id));
        Ok(row_ids)
    })
    .await?;

    Ok(Arc::new(PositionedRowSet::from_file_order(
        row_ids_in_file_order,
    )?))
}