Skip to main content

lance_table/format/
row_ids.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::ops::Deref;
5use std::sync::{Arc, OnceLock};
6
7use lance_core::deepsize::{Context, DeepSizeOf};
8use lance_core::{Error, Result};
9use serde::{Deserialize, Deserializer, Serialize, Serializer};
10
11use super::pb;
12
13/// A reference to a part of a file.
14#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
15pub struct ExternalFile {
16    pub path: String,
17    pub offset: u64,
18    pub size: u64,
19}
20
21/// A fragment's row id sequence, encoded inline in the manifest.
22///
23/// Carries a memoized [`digest`](Self::digest) of the encoded bytes. The digest
24/// identifies *which* sequence these bytes are, which is what the row id
25/// sequence cache keys on: a fragment id alone does not identify a sequence,
26/// because fragment ids are reused across dataset generations.
27///
28/// The digest is memoized because computing it is proportional to the encoded
29/// size. A run-encoded sequence is a handful of bytes per run, but a heavily
30/// fragmented one is array-encoded at 8 bytes per row, and the cache is
31/// consulted on every scan, count, prefilter and index load.
32///
33/// Bytes and memo share one immutable allocation, so cloning shares both. That
34/// is what makes the memo worth having: callers clone `Fragment` before loading
35/// its sequence (see `count_from_mask`), and a memo held per clone would be
36/// filled and dropped by each scan, rehashing the whole sequence every time.
37/// Sharing also keeps a cloned fragment from duplicating the encoded bytes.
38///
39/// The digest lives *with* the bytes rather than beside them so the two cannot
40/// drift: several write paths replace a fragment's `row_id_meta` after the
41/// fragment is built, and a digest that outlived its bytes would silently
42/// resolve to another generation's sequence.
43#[derive(Clone)]
44pub struct InlineRowIds {
45    inner: Arc<InlineRowIdsInner>,
46}
47
48struct InlineRowIdsInner {
49    data: Vec<u8>,
50    digest: OnceLock<[u8; 32]>,
51}
52
53impl InlineRowIds {
54    /// Digest of the encoded bytes, computed on first use and shared by clones.
55    pub fn digest(&self) -> &[u8; 32] {
56        self.inner
57            .digest
58            .get_or_init(|| blake3::hash(&self.inner.data).into())
59    }
60}
61
62impl From<Vec<u8>> for InlineRowIds {
63    fn from(data: Vec<u8>) -> Self {
64        Self {
65            inner: Arc::new(InlineRowIdsInner {
66                data,
67                digest: OnceLock::new(),
68            }),
69        }
70    }
71}
72
73impl Deref for InlineRowIds {
74    type Target = [u8];
75
76    fn deref(&self) -> &Self::Target {
77        &self.inner.data
78    }
79}
80
81// Debug, equality and serialization all present the bytes alone: the memo is a
82// derived value and must not show up in output, comparisons or the manifest.
83impl std::fmt::Debug for InlineRowIds {
84    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85        self.inner.data.fmt(f)
86    }
87}
88
89impl PartialEq for InlineRowIds {
90    fn eq(&self, other: &Self) -> bool {
91        Arc::ptr_eq(&self.inner, &other.inner) || self.inner.data == other.inner.data
92    }
93}
94
95impl Eq for InlineRowIds {}
96
97impl Serialize for InlineRowIds {
98    fn serialize<S: Serializer>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> {
99        self.inner.data.serialize(serializer)
100    }
101}
102
103impl<'de> Deserialize<'de> for InlineRowIds {
104    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
105        Vec::<u8>::deserialize(deserializer).map(Self::from)
106    }
107}
108
109impl DeepSizeOf for InlineRowIds {
110    fn deep_size_of_children(&self, context: &mut Context) -> usize {
111        // Delegate to the `Arc` so clones sharing one allocation are counted once.
112        self.inner.deep_size_of_children(context)
113    }
114}
115
116impl DeepSizeOf for InlineRowIdsInner {
117    fn deep_size_of_children(&self, context: &mut Context) -> usize {
118        self.data.deep_size_of_children(context)
119    }
120}
121
122/// Metadata about location of the row id sequence.
123#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
124pub enum RowIdMeta {
125    Inline(InlineRowIds),
126    External(ExternalFile),
127}
128
129impl TryFrom<pb::data_fragment::RowIdSequence> for RowIdMeta {
130    type Error = Error;
131
132    fn try_from(value: pb::data_fragment::RowIdSequence) -> Result<Self> {
133        match value {
134            pb::data_fragment::RowIdSequence::InlineRowIds(data) => Ok(Self::Inline(data.into())),
135            pb::data_fragment::RowIdSequence::ExternalRowIds(file) => {
136                Ok(Self::External(ExternalFile {
137                    path: file.path.clone(),
138                    offset: file.offset,
139                    size: file.size,
140                }))
141            }
142        }
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149    use crate::format::Fragment;
150
151    #[test]
152    fn inline_row_ids_digest_identifies_contents() {
153        let first = InlineRowIds::from(vec![1, 2, 3]);
154        let same = InlineRowIds::from(vec![1, 2, 3]);
155        let other = InlineRowIds::from(vec![1, 2, 4]);
156
157        // The digest is what the row id sequence cache keys on, so it must
158        // follow the bytes exactly: same bytes, same key; any change, new key.
159        assert_eq!(first.digest(), same.digest());
160        assert_ne!(first.digest(), other.digest());
161        assert_eq!(first, same);
162        assert_ne!(first, other);
163
164        // Memoized on first use, so repeat use must return the same digest.
165        let memoized = *first.digest();
166        assert_eq!(first.digest(), &memoized);
167    }
168
169    #[test]
170    fn inline_row_ids_clones_share_bytes_and_memo() {
171        let first = InlineRowIds::from(vec![1, 2, 3]);
172        let cloned = first.clone();
173
174        // Callers clone `Fragment` before loading its row id sequence, so a memo
175        // held per clone would be filled and dropped by each scan and rehash the
176        // whole sequence every time. Sharing one allocation is what makes the
177        // memo pay off, and it keeps clones from duplicating the bytes.
178        assert!(std::ptr::eq(first.as_ptr(), cloned.as_ptr()));
179        assert!(std::ptr::eq(first.digest(), cloned.digest()));
180
181        // Guard against the assertions above passing vacuously: only clones
182        // share storage, equal bytes built separately do not.
183        let separate = InlineRowIds::from(vec![1, 2, 3]);
184        assert_eq!(first, separate);
185        assert!(!std::ptr::eq(first.as_ptr(), separate.as_ptr()));
186
187        // Computing through one clone must be visible from the other.
188        let fresh = InlineRowIds::from(vec![4, 5, 6]);
189        let fresh_clone = fresh.clone();
190        let via_clone = *fresh_clone.digest();
191        assert_eq!(fresh.digest(), &via_clone);
192    }
193
194    #[test]
195    fn inline_row_ids_serializes_as_bare_bytes() {
196        // The manifest is a stable format: the memo must not reach the wire.
197        let meta = RowIdMeta::Inline(InlineRowIds::from(vec![7, 8, 9]));
198        let json = serde_json::to_string(&meta).unwrap();
199        assert_eq!(json, r#"{"Inline":[7,8,9]}"#);
200        assert_eq!(serde_json::from_str::<RowIdMeta>(&json).unwrap(), meta);
201
202        // ...and round-trips through protobuf unchanged.
203        let fragment = Fragment {
204            row_id_meta: Some(meta.clone()),
205            ..Fragment::new(0)
206        };
207        let restored = Fragment::try_from(pb::DataFragment::from(&fragment)).unwrap();
208        assert_eq!(restored.row_id_meta, Some(meta));
209    }
210}