feat: add support for empty structs to the 2.0 format (#3499)

Author: westonpace

python/python/tests/test_file.py

@@ -359,6 +359,17 @@ def round_trip(arr):
         assert round_tripped.type == dict_arr.type
 
 
+def test_empty_structs(tmp_path):
+    schema = pa.schema([pa.field("empties", pa.struct([]))])
+    table = pa.table({"empties": [{}] * 3}, schema=schema)
+    path = tmp_path / "foo.lance"
+    with LanceFileWriter(str(path)) as writer:
+        writer.write_batch(table)
+    reader = LanceFileReader(str(path))
+    round_tripped = reader.read_all().to_table()
+    assert round_tripped == table
+
+
 def test_write_read_global_buffer(tmp_path):
     table = pa.table({"a": [1, 2, 3]})
     path = tmp_path / "foo.lance"

rust/lance-datagen/src/generator.rs

@@ -1183,6 +1183,12 @@ impl ArrayGenerator for RandomStructGenerator {
         length: RowCount,
         rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
     ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
+        if self.child_gens.is_empty() {
+            // Have to create empty struct arrays specially to ensure they have the correct
+            // row count
+            let struct_arr = StructArray::new_empty_fields(length.0 as usize, None);
+            return Ok(Arc::new(struct_arr));
+        }
         let child_arrays = self
             .child_gens
             .iter_mut()

rust/lance-encoding/src/decoder.rs

@@ -959,6 +959,11 @@ impl CoreFieldDecoderStrategy {
                 } else {
                     // use default struct encoding
                     Self::check_simple_struct(column_info, &field.name).unwrap();
+                    let num_rows = column_info
+                        .page_infos
+                        .iter()
+                        .map(|page| page.num_rows)
+                        .sum();
                     let mut child_schedulers = Vec::with_capacity(field.children.len());
                     for field in &field.children {
                         column_infos.next_top_level();
@@ -971,6 +976,7 @@ impl CoreFieldDecoderStrategy {
                     Ok(Box::new(SimpleStructScheduler::new(
                         child_schedulers,
                         fields,
+                        num_rows,
                     )))
                 }
             }

rust/lance-encoding/src/encodings/logical/list.rs

@@ -365,7 +365,7 @@ async fn indirect_schedule_task(
     // Create a new root scheduler, which has one column, which is our items data
     let root_fields = Fields::from(vec![Field::new("item", items_type, true)]);
     let indirect_root_scheduler =
-        SimpleStructScheduler::new(vec![items_scheduler], root_fields.clone());
+        SimpleStructScheduler::new(vec![items_scheduler], root_fields.clone(), num_items);
     let mut indirect_scheduler = DecodeBatchScheduler::from_scheduler(
         Arc::new(indirect_root_scheduler),
         root_fields.clone(),

rust/lance-encoding/src/encodings/logical/struct.rs

@@ -8,7 +8,7 @@ use std::{
 };
 
 use arrow_array::{cast::AsArray, Array, ArrayRef, StructArray};
-use arrow_schema::{DataType, Fields};
+use arrow_schema::{DataType, Field, Fields};
 use futures::{
     future::BoxFuture,
     stream::{FuturesOrdered, FuturesUnordered},
@@ -64,6 +64,89 @@ impl Ord for SchedulingJobWithStatus<'_> {
     }
 }
 
+#[derive(Debug)]
+struct EmptyStructDecodeTask {
+    num_rows: u64,
+}
+
+impl DecodeArrayTask for EmptyStructDecodeTask {
+    fn decode(self: Box<Self>) -> Result<ArrayRef> {
+        Ok(Arc::new(StructArray::new_empty_fields(
+            self.num_rows as usize,
+            None,
+        )))
+    }
+}
+
+#[derive(Debug)]
+struct EmptyStructDecoder {
+    num_rows: u64,
+    rows_drained: u64,
+    data_type: DataType,
+}
+
+impl EmptyStructDecoder {
+    fn new(num_rows: u64) -> Self {
+        Self {
+            num_rows,
+            rows_drained: 0,
+            data_type: DataType::Struct(Fields::from(Vec::<Field>::default())),
+        }
+    }
+}
+
+impl LogicalPageDecoder for EmptyStructDecoder {
+    fn wait_for_loaded(&mut self, _loaded_need: u64) -> BoxFuture<Result<()>> {
+        Box::pin(std::future::ready(Ok(())))
+    }
+    fn rows_loaded(&self) -> u64 {
+        self.num_rows
+    }
+    fn rows_unloaded(&self) -> u64 {
+        0
+    }
+    fn num_rows(&self) -> u64 {
+        self.num_rows
+    }
+    fn rows_drained(&self) -> u64 {
+        self.rows_drained
+    }
+    fn drain(&mut self, num_rows: u64) -> Result<NextDecodeTask> {
+        self.rows_drained += num_rows;
+        Ok(NextDecodeTask {
+            num_rows,
+            task: Box::new(EmptyStructDecodeTask { num_rows }),
+        })
+    }
+    fn data_type(&self) -> &DataType {
+        &self.data_type
+    }
+}
+
+#[derive(Debug)]
+struct EmptyStructSchedulerJob {
+    num_rows: u64,
+}
+
+impl SchedulingJob for EmptyStructSchedulerJob {
+    fn schedule_next(
+        &mut self,
+        context: &mut SchedulerContext,
+        _priority: &dyn PriorityRange,
+    ) -> Result<ScheduledScanLine> {
+        let empty_decoder = Box::new(EmptyStructDecoder::new(self.num_rows));
+        let struct_decoder = context.locate_decoder(empty_decoder);
+        Ok(ScheduledScanLine {
+            decoders: vec![MessageType::DecoderReady(struct_decoder)],
+            rows_scheduled: self.num_rows,
+        })
+    }
+
+    fn num_rows(&self) -> u64 {
+        self.num_rows
+    }
+}
+
 /// Scheduling job for struct data
 ///
 /// The order in which we schedule the children is important.  We want to schedule the child
@@ -175,9 +258,15 @@ pub struct SimpleStructScheduler {
 }
 
 impl SimpleStructScheduler {
-    pub fn new(children: Vec<Arc<dyn FieldScheduler>>, child_fields: Fields) -> Self {
-        debug_assert!(!children.is_empty());
-        let num_rows = children[0].num_rows();
+    pub fn new(
+        children: Vec<Arc<dyn FieldScheduler>>,
+        child_fields: Fields,
+        num_rows: u64,
+    ) -> Self {
+        let num_rows = children
+            .first()
+            .map(|child| child.num_rows())
+            .unwrap_or(num_rows);
         debug_assert!(children.iter().all(|child| child.num_rows() == num_rows));
         Self {
             children,
@@ -193,6 +282,11 @@ impl FieldScheduler for SimpleStructScheduler {
         ranges: &[Range<u64>],
         filter: &FilterExpression,
     ) -> Result<Box<dyn SchedulingJob + 'a>> {
+        if self.children.is_empty() {
+            return Ok(Box::new(EmptyStructSchedulerJob {
+                num_rows: ranges.iter().map(|r| r.end - r.start).sum(),
+            }));
+        }
         let child_schedulers = self
             .children
             .iter()
@@ -1120,6 +1214,15 @@ mod tests {
         check_round_trip_encoding_random(field, LanceFileVersion::V2_0).await;
     }
 
+    #[test_log::test(tokio::test)]
+    async fn test_empty_struct() {
+        // It's technically legal for a struct to have 0 children, need to
+        // make sure we support that
+        let data_type = DataType::Struct(Fields::from(Vec::<Field>::default()));
+        let field = Field::new("row", data_type, false);
+        check_round_trip_encoding_random(field, LanceFileVersion::V2_0).await;
+    }
+
     #[test_log::test(tokio::test)]
     async fn test_complicated_struct() {
         let data_type = DataType::Struct(Fields::from(vec![

rust/lance-encoding/src/testing.rs

@@ -4,7 +4,7 @@
 use std::{cmp::Ordering, collections::HashMap, ops::Range, sync::Arc};
 
 use arrow::array::make_comparator;
-use arrow_array::{Array, UInt64Array};
+use arrow_array::{Array, StructArray, UInt64Array};
 use arrow_schema::{DataType, Field, FieldRef, Schema, SortOptions};
 use arrow_select::concat::concat;
 use bytes::{Bytes, BytesMut};
@@ -651,9 +651,23 @@ async fn check_round_trip_encoding_inner(
         }
         let num_rows = indices.len() as u64;
         let indices_arr = UInt64Array::from(indices.clone());
-        let expected = concat_data
-            .as_ref()
-            .map(|concat_data| arrow_select::take::take(&concat_data, &indices_arr, None).unwrap());
+
+        // There is a bug in arrow_select::take::take that causes it to return empty arrays
+        // if the data type is an empty struct.  This is a workaround for that.
+        let is_empty_struct = if let DataType::Struct(fields) = field.data_type() {
+            fields.is_empty()
+        } else {
+            false
+        };
+
+        let expected = if is_empty_struct {
+            Some(Arc::new(StructArray::new_empty_fields(indices_arr.len(), None)) as Arc<dyn Array>)
+        } else {
+            concat_data.as_ref().map(|concat_data| {
+                arrow_select::take::take(&concat_data, &indices_arr, None).unwrap()
+            })
+        };
+
         let scheduler = scheduler.clone();
         let indices = indices.clone();
         test_decode(

rust/lance-file/src/v2/reader.rs

@@ -1525,6 +1525,11 @@ pub mod tests {
             .copied()
             .collect::<BTreeMap<_, _>>();
 
+        let empty_projection = ReaderProjection {
+            column_indices: Vec::default(),
+            schema: Arc::new(Schema::default()),
+        };
+
         for columns in [
             vec!["score"],
             vec!["location"],
@@ -1606,12 +1611,17 @@ pub mod tests {
                 })),
             )
             .await;
-        }
 
-        let empty_projection = ReaderProjection {
-            column_indices: Vec::default(),
-            schema: Arc::new(Schema::default()),
-        };
+            assert!(file_reader
+                .read_stream_projected(
+                    lance_io::ReadBatchParams::RangeFull,
+                    1024,
+                    16,
+                    empty_projection.clone(),
+                    FilterExpression::no_filter(),
+                )
+                .is_err());
+        }
 
         assert!(FileReader::try_open(
             file_scheduler.clone(),