lance_io/
utils.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{cmp::min, num::NonZero, sync::atomic::AtomicU64};
5
6use arrow_array::{
7    ArrayRef,
8    types::{BinaryType, LargeBinaryType, LargeUtf8Type, Utf8Type},
9};
10use arrow_schema::DataType;
11use byteorder::{ByteOrder, LittleEndian};
12use bytes::Bytes;
13use deepsize::DeepSizeOf;
14use lance_arrow::*;
15use prost::Message;
16use serde::{Deserialize, Serialize};
17
18use crate::{ReadBatchParams, traits::Reader};
19use crate::{
20    encodings::{AsyncIndex, Decoder, binary::BinaryDecoder, plain::PlainDecoder},
21    traits::ProtoStruct,
22};
23use lance_core::{Error, Result};
24
25pub mod tracking_store;
26
27/// Read a binary array from a [Reader].
28///
29pub async fn read_binary_array(
30    reader: &dyn Reader,
31    data_type: &DataType,
32    nullable: bool,
33    position: usize,
34    length: usize,
35    params: impl Into<ReadBatchParams>,
36) -> Result<ArrayRef> {
37    use arrow_schema::DataType::*;
38    let decoder: Box<dyn Decoder<Output = Result<ArrayRef>> + Send> = match data_type {
39        Utf8 => Box::new(BinaryDecoder::<Utf8Type>::new(
40            reader, position, length, nullable,
41        )),
42        Binary => Box::new(BinaryDecoder::<BinaryType>::new(
43            reader, position, length, nullable,
44        )),
45        LargeUtf8 => Box::new(BinaryDecoder::<LargeUtf8Type>::new(
46            reader, position, length, nullable,
47        )),
48        LargeBinary => Box::new(BinaryDecoder::<LargeBinaryType>::new(
49            reader, position, length, nullable,
50        )),
51        _ => {
52            return Err(Error::invalid_input(format!(
53                "Unsupported binary type: {}",
54                data_type
55            )));
56        }
57    };
58    let fut = decoder.as_ref().get(params.into());
59    fut.await
60}
61
62/// Read a fixed stride array from disk.
63///
64pub async fn read_fixed_stride_array(
65    reader: &dyn Reader,
66    data_type: &DataType,
67    position: usize,
68    length: usize,
69    params: impl Into<ReadBatchParams>,
70) -> Result<ArrayRef> {
71    if !data_type.is_fixed_stride() {
72        return Err(Error::schema(format!(
73            "{data_type} is not a fixed stride type"
74        )));
75    }
76    // TODO: support more than plain encoding here.
77    let decoder = PlainDecoder::new(reader, data_type, position, length)?;
78    decoder.get(params.into()).await
79}
80
81/// Read a protobuf message at file position 'pos'.
82///
83/// We write protobuf by first writing the length of the message as a u32,
84/// followed by the message itself.
85pub async fn read_message<M: Message + Default>(reader: &dyn Reader, pos: usize) -> Result<M> {
86    let file_size = reader.size().await?;
87    if pos > file_size {
88        return Err(Error::io("file size is too small".to_string()));
89    }
90
91    let range = pos..min(pos + reader.block_size(), file_size);
92    let buf = reader.get_range(range.clone()).await?;
93    let msg_len = LittleEndian::read_u32(&buf) as usize;
94
95    if msg_len + 4 > buf.len() {
96        let remaining_range = range.end..min(4 + pos + msg_len, file_size);
97        let remaining_bytes = reader.get_range(remaining_range).await?;
98        let buf = [buf, remaining_bytes].concat();
99        assert!(buf.len() >= msg_len + 4);
100        Ok(M::decode(&buf[4..4 + msg_len])?)
101    } else {
102        Ok(M::decode(&buf[4..4 + msg_len])?)
103    }
104}
105
106/// Read a Protobuf-backed struct at file position: `pos`.
107// TODO: pub(crate)
108pub async fn read_struct<
109    M: Message + Default + 'static,
110    T: ProtoStruct<Proto = M> + TryFrom<M, Error = Error>,
111>(
112    reader: &dyn Reader,
113    pos: usize,
114) -> Result<T> {
115    let msg = read_message::<M>(reader, pos).await?;
116    T::try_from(msg)
117}
118
119pub async fn read_last_block(reader: &dyn Reader) -> object_store::Result<Bytes> {
120    let file_size = reader.size().await?;
121    let block_size = reader.block_size();
122    let begin = file_size.saturating_sub(block_size);
123    reader.get_range(begin..file_size).await
124}
125
126pub fn read_metadata_offset(bytes: &Bytes) -> Result<usize> {
127    let len = bytes.len();
128    if len < 16 {
129        return Err(Error::io(format!(
130            "does not have sufficient data, len: {}, bytes: {:?}",
131            len, bytes
132        )));
133    }
134    let offset_bytes = bytes.slice(len - 16..len - 8);
135    Ok(LittleEndian::read_u64(offset_bytes.as_ref()) as usize)
136}
137
138/// Read the version from the footer bytes
139pub fn read_version(bytes: &Bytes) -> Result<(u16, u16)> {
140    let len = bytes.len();
141    if len < 8 {
142        return Err(Error::io(format!(
143            "does not have sufficient data, len: {}, bytes: {:?}",
144            len, bytes
145        )));
146    }
147
148    let major_version = LittleEndian::read_u16(bytes.slice(len - 8..len - 6).as_ref());
149    let minor_version = LittleEndian::read_u16(bytes.slice(len - 6..len - 4).as_ref());
150    Ok((major_version, minor_version))
151}
152
153/// Read protobuf from a buffer.
154pub fn read_message_from_buf<M: Message + Default>(buf: &Bytes) -> Result<M> {
155    let msg_len = LittleEndian::read_u32(buf) as usize;
156    Ok(M::decode(&buf[4..4 + msg_len])?)
157}
158
159/// Read a Protobuf-backed struct from a buffer.
160pub fn read_struct_from_buf<
161    M: Message + Default,
162    T: ProtoStruct<Proto = M> + TryFrom<M, Error = Error>,
163>(
164    buf: &Bytes,
165) -> Result<T> {
166    let msg: M = read_message_from_buf(buf)?;
167    T::try_from(msg)
168}
169
170/// A cached file size.
171///
172/// This wraps an atomic u64 to allow setting the cached file size without
173/// needed a mutable reference.
174///
175/// Zero is interpreted as unknown.
176#[derive(Debug, DeepSizeOf)]
177pub struct CachedFileSize(AtomicU64);
178
179impl<'de> Deserialize<'de> for CachedFileSize {
180    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
181    where
182        D: serde::Deserializer<'de>,
183    {
184        let size = Option::<u64>::deserialize(deserializer)?.unwrap_or(0);
185        Ok(Self::new(size))
186    }
187}
188
189impl Serialize for CachedFileSize {
190    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
191    where
192        S: serde::Serializer,
193    {
194        let size = self.0.load(std::sync::atomic::Ordering::Relaxed);
195        if size == 0 {
196            serializer.serialize_none()
197        } else {
198            serializer.serialize_u64(size)
199        }
200    }
201}
202
203impl From<Option<NonZero<u64>>> for CachedFileSize {
204    fn from(size: Option<NonZero<u64>>) -> Self {
205        match size {
206            Some(size) => Self(AtomicU64::new(size.into())),
207            None => Self(AtomicU64::new(0)),
208        }
209    }
210}
211
212impl Default for CachedFileSize {
213    fn default() -> Self {
214        Self(AtomicU64::new(0))
215    }
216}
217
218impl Clone for CachedFileSize {
219    fn clone(&self) -> Self {
220        Self(AtomicU64::new(
221            self.0.load(std::sync::atomic::Ordering::Relaxed),
222        ))
223    }
224}
225
226impl PartialEq for CachedFileSize {
227    fn eq(&self, other: &Self) -> bool {
228        self.0.load(std::sync::atomic::Ordering::Relaxed)
229            == other.0.load(std::sync::atomic::Ordering::Relaxed)
230    }
231}
232
233impl Eq for CachedFileSize {}
234
235impl CachedFileSize {
236    pub fn new(size: u64) -> Self {
237        Self(AtomicU64::new(size))
238    }
239
240    pub fn unknown() -> Self {
241        Self(AtomicU64::new(0))
242    }
243
244    pub fn get(&self) -> Option<NonZero<u64>> {
245        NonZero::new(self.0.load(std::sync::atomic::Ordering::Relaxed))
246    }
247
248    pub fn set(&self, size: NonZero<u64>) {
249        self.0
250            .store(size.into(), std::sync::atomic::Ordering::Relaxed);
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use bytes::Bytes;
257    use object_store::path::Path;
258
259    use crate::{
260        Error, Result,
261        object_reader::CloudObjectReader,
262        object_store::{DEFAULT_DOWNLOAD_RETRY_COUNT, ObjectStore},
263        object_writer::ObjectWriter,
264        traits::{ProtoStruct, WriteExt, Writer},
265        utils::read_struct,
266    };
267
268    // Bytes is a prost::Message, since we don't have any .proto files in this crate we
269    // can use it to simulate a real message object.
270    #[derive(Debug, PartialEq)]
271    struct BytesWrapper(Bytes);
272
273    impl ProtoStruct for BytesWrapper {
274        type Proto = Bytes;
275    }
276
277    impl From<&BytesWrapper> for Bytes {
278        fn from(value: &BytesWrapper) -> Self {
279            value.0.clone()
280        }
281    }
282
283    impl TryFrom<Bytes> for BytesWrapper {
284        type Error = Error;
285        fn try_from(value: Bytes) -> Result<Self> {
286            Ok(Self(value))
287        }
288    }
289
290    #[tokio::test]
291    async fn test_write_proto_structs() {
292        let store = ObjectStore::memory();
293        let path = Path::from("/foo");
294
295        let mut object_writer = ObjectWriter::new(&store, &path).await.unwrap();
296        assert_eq!(object_writer.tell().await.unwrap(), 0);
297
298        let some_message = BytesWrapper(Bytes::from(vec![10, 20, 30]));
299
300        let pos = object_writer.write_struct(&some_message).await.unwrap();
301        assert_eq!(pos, 0);
302        object_writer.shutdown().await.unwrap();
303
304        let object_reader =
305            CloudObjectReader::new(store.inner, path, 1024, None, DEFAULT_DOWNLOAD_RETRY_COUNT)
306                .unwrap();
307        let actual: BytesWrapper = read_struct(&object_reader, pos).await.unwrap();
308        assert_eq!(some_message, actual);
309    }
310}
lance_io/utils.rs

lance_io/
utils.rs
