[Rust] Lance File Writer

rust/src/arrow.rs

@@ -26,6 +26,38 @@ use arrow_data::ArrayDataBuilder;
 use arrow_schema::{DataType, Field};
 
 use crate::error::Result;
+
+pub trait DataTypeExt {
+    /// Returns true if the data type is binary-like, such as (Large)Utf8 and (Large)Binary.
+    ///
+    /// ```
+    /// use lance::arrow::*;
+    /// use arrow_schema::DataType;
+    ///
+    /// assert!(DataType::Utf8.is_binary_like());
+    /// assert!(DataType::Binary.is_binary_like());
+    /// assert!(DataType::LargeUtf8.is_binary_like());
+    /// assert!(DataType::LargeBinary.is_binary_like());
+    /// assert!(!DataType::Int32.is_binary_like());
+    /// ```
+    fn is_binary_like(&self) -> bool;
+
+    fn is_struct(&self) -> bool;
+}
+
+impl DataTypeExt for DataType {
+    fn is_binary_like(&self) -> bool {
+        matches!(
+            self,
+            DataType::Utf8 | DataType::Binary | DataType::LargeUtf8 | DataType::LargeBinary
+        )
+    }
+
+    fn is_struct(&self) -> bool {
+        matches!(self, DataType::Struct(_))
+    }
+}
+
 pub trait ListArrayExt {
     /// Create an [`ListArray`] from values and offsets.
     ///

rust/src/datatypes.rs

@@ -1,6 +1,6 @@
 //! Lance data types, [Schema] and [Field]
 
-use std::cmp::min;
+use std::cmp::max;
 use std::collections::HashMap;
 use std::fmt;
 use std::fmt::Formatter;
@@ -9,6 +9,7 @@ use arrow_array::ArrayRef;
 use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, TimeUnit};
 use async_recursion::async_recursion;
 
+use crate::arrow::DataTypeExt;
 use crate::encodings::Encoding;
 use crate::format::pb;
 use crate::io::object_reader::ObjectReader;
@@ -201,19 +202,6 @@ impl TryFrom<&LogicalType> for DataType {
     }
 }
 
-fn is_numeric(data_type: &DataType) -> bool {
-    use DataType::*;
-    matches!(
-        data_type,
-        UInt8 | UInt16 | UInt32 | UInt64 | Int8 | Int16 | Int32 | Int64 | Float32 | Float64
-    )
-}
-
-fn is_binary(data_type: &DataType) -> bool {
-    use DataType::*;
-    matches!(data_type, Binary | Utf8 | LargeBinary | LargeUtf8)
-}
-
 #[derive(Debug, Clone, PartialEq)]
 pub struct Dictionary {
     offset: usize,
@@ -242,7 +230,7 @@ pub struct Field {
     parent_id: i32,
     logical_type: LogicalType,
     extension_name: String,
-    encoding: Option<Encoding>,
+    pub(crate) encoding: Option<Encoding>,
     nullable: bool,
 
     pub children: Vec<Field>,
@@ -313,9 +301,13 @@ impl Field {
 
     // Get the max field id of itself and all children.
     fn max_id(&self) -> i32 {
-        min(
+        max(
             self.id,
-            self.children.iter().map(|c| c.id).min().unwrap_or(i32::MAX),
+            self.children
+                .iter()
+                .map(|c| c.max_id())
+                .max()
+                .unwrap_or(i32::MIN),
         )
     }
 
@@ -420,8 +412,8 @@ impl TryFrom<&ArrowField> for Field {
             name: field.name().clone(),
             logical_type: LogicalType::try_from(field.data_type())?,
             encoding: match field.data_type() {
-                dt if is_numeric(dt) => Some(Encoding::Plain),
-                dt if is_binary(dt) => Some(Encoding::VarBinary),
+                dt if dt.is_numeric() => Some(Encoding::Plain),
+                dt if dt.is_binary_like() => Some(Encoding::VarBinary),
                 DataType::Dictionary(_, _) => Some(Encoding::Dictionary),
                 _ => None,
             },

rust/src/encodings/binary.rs

@@ -25,11 +25,11 @@ use crate::io::object_writer::ObjectWriter;
 
 /// Encoder for Var-binary encoding.
 pub struct BinaryEncoder<'a> {
-    writer: &'a mut ObjectWriter<'a>,
+    writer: &'a mut ObjectWriter,
 }
 
 impl<'a> BinaryEncoder<'a> {
-    pub fn new(writer: &'a mut ObjectWriter<'a>) -> Self {
+    pub fn new(writer: &'a mut ObjectWriter) -> Self {
         Self { writer }
     }
 
@@ -190,16 +190,12 @@ mod tests {
     async fn test_round_trips<O: OffsetSizeTrait>(arr: &GenericStringArray<O>) {
         let store = ObjectStore::new(":memory:").unwrap();
         let path = Path::from("/foo");
-        let (_, mut writer) = store.inner.put_multipart(&path).await.unwrap();
-
-        let mut object_writer = ObjectWriter::new(writer.as_mut());
+        let mut object_writer = ObjectWriter::new(&store, &path).await.unwrap();
         // Write some gabage to reset "tell()".
         object_writer.write_all(b"1234").await.unwrap();
         let mut encoder = BinaryEncoder::new(&mut object_writer);
-
         let pos = encoder.encode(&arr).await.unwrap();
-
-        writer.shutdown().await.unwrap();
+        object_writer.shutdown().await.unwrap();
 
         let mut reader = store.open(&path).await.unwrap();
         let decoder = BinaryDecoder::<GenericStringType<O>>::new(&mut reader, pos, arr.len());

rust/src/encodings/dictionary.rs

@@ -126,7 +126,6 @@ mod tests {
     };
     use arrow_array::Array;
     use object_store::path::Path;
-    use tokio::io::AsyncWriteExt;
 
     async fn test_dict_decoder_for_type<T: ArrowDictionaryKeyType>() {
         let values = vec!["a", "b", "b", "a", "c"];
@@ -137,11 +136,10 @@ mod tests {
 
         let pos;
         {
-            let (_, mut writer) = store.inner.put_multipart(&path).await.unwrap();
-            let mut object_writer = ObjectWriter::new(writer.as_mut());
+            let mut object_writer = ObjectWriter::new(&store, &path).await.unwrap();
             let mut encoder = PlainEncoder::new(&mut object_writer, arr.keys().data_type());
             pos = encoder.encode(arr.keys()).await.unwrap();
-            writer.shutdown().await.unwrap();
+            object_writer.shutdown().await.unwrap();
         }
 
         let reader = store.open(&path).await.unwrap();

rust/src/encodings/plain.rs

@@ -49,12 +49,12 @@ use crate::io::object_writer::ObjectWriter;
 /// Encoder for plain encoding.
 ///
 pub struct PlainEncoder<'a> {
-    writer: &'a mut ObjectWriter<'a>,
+    writer: &'a mut ObjectWriter,
     data_type: &'a DataType,
 }
 
 impl<'a> PlainEncoder<'a> {
-    pub fn new(writer: &'a mut ObjectWriter<'a>, data_type: &'a DataType) -> PlainEncoder<'a> {
+    pub fn new(writer: &'a mut ObjectWriter, data_type: &'a DataType) -> PlainEncoder<'a> {
         PlainEncoder { writer, data_type }
     }
 
@@ -229,7 +229,6 @@ mod tests {
     use rand::prelude::*;
     use std::borrow::Borrow;
     use std::sync::Arc;
-    use tokio::io::AsyncWriteExt;
 
     use super::*;
     use crate::arrow::*;
@@ -267,15 +266,11 @@ mod tests {
     async fn test_round_trip(expected: ArrayRef, data_type: DataType) {
         let store = ObjectStore::new(":memory:").unwrap();
         let path = Path::from("/foo");
-        let (_, mut writer) = store.inner.put_multipart(&path).await.unwrap();
+        let mut object_writer = ObjectWriter::new(&store, &path).await.unwrap();
+        let mut encoder = PlainEncoder::new(&mut object_writer, &data_type);
 
-        {
-            let mut object_writer = ObjectWriter::new(writer.as_mut());
-            let mut encoder = PlainEncoder::new(&mut object_writer, &data_type);
-
-            assert_eq!(encoder.encode(expected.as_ref()).await.unwrap(), 0);
-        }
-        writer.shutdown().await.unwrap();
+        assert_eq!(encoder.encode(expected.as_ref()).await.unwrap(), 0);
+        object_writer.shutdown().await.unwrap();
 
         let mut reader = store.open(&path).await.unwrap();
         assert!(reader.size().await.unwrap() > 0);

rust/src/format.rs

@@ -1,3 +1,5 @@
+//! On-disk format
+
 mod fragment;
 mod manifest;
 mod metadata;
@@ -15,6 +17,11 @@ pub mod pb {
     include!(concat!(env!("OUT_DIR"), "/lance.format.pb.rs"));
 }
 
+pub const MAJOR_VERSION: i16 = 0;
+pub const MINOR_VERSION: i16 = 1;
+pub const MAGIC: &[u8; 4] = b"LANC";
+pub const INDEX_MAGIC: &[u8; 8] = b"LANC_IDX";
+
 /// Annotation on a struct that can be converted a Protobuf message.
 pub trait ProtoStruct {
     type Proto: Message;

rust/src/format/fragment.rs

@@ -79,3 +79,12 @@ impl From<&pb::DataFragment> for Fragment {
         }
     }
 }
+
+impl From<&Fragment> for pb::DataFragment {
+    fn from(f: &Fragment) -> Self {
+        Self {
+            id: f.id,
+            files: f.files.iter().map(pb::DataFile::from).collect(),
+        }
+    }
+}

rust/src/format/manifest.rs

@@ -18,6 +18,7 @@
 use super::Fragment;
 use crate::datatypes::Schema;
 use crate::format::{pb, ProtoStruct};
+use std::collections::HashMap;
 
 /// Manifest of a dataset
 ///
@@ -36,6 +37,16 @@ pub struct Manifest {
     pub fragments: Vec<Fragment>,
 }
 
+impl Manifest {
+    pub fn new(schema: &Schema) -> Self {
+        Self {
+            schema: schema.clone(),
+            version: 1,
+            fragments: vec![],
+        }
+    }
+}
+
 impl ProtoStruct for Manifest {
     type Proto = pb::Manifest;
 }
@@ -49,3 +60,15 @@ impl From<pb::Manifest> for Manifest {
         }
     }
 }
+
+impl From<&Manifest> for pb::Manifest {
+    fn from(m: &Manifest) -> Self {
+        Self {
+            fields: (&m.schema).into(),
+            version: m.version,
+            fragments: m.fragments.iter().map(pb::DataFragment::from).collect(),
+            metadata: HashMap::default(),
+            version_aux_data: 0,
+        }
+    }
+}