array tree layouts

Signed-off-by: Onur Satici <onur@spiraldb.com>

Author: onursatici

vortex-file/public-api.lock

@@ -41,6 +41,7 @@ use vortex_fastlanes::FoR;
 use vortex_fastlanes::RLE;
 use vortex_fsst::FSST;
 use vortex_layout::LayoutStrategy;
+use vortex_layout::layouts::array_tree::writer as array_tree_writer;
 use vortex_layout::layouts::buffered::BufferedStrategy;
 use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy;
 use vortex_layout::layouts::collect::CollectStrategy;
@@ -149,6 +150,7 @@ pub struct WriteStrategyBuilder {
     field_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
     allow_encodings: Option<HashSet<ArrayId>>,
     flat_strategy: Option<Arc<dyn LayoutStrategy>>,
+    array_tree: bool,
 }
 
 impl Default for WriteStrategyBuilder {
@@ -161,6 +163,7 @@ impl Default for WriteStrategyBuilder {
             field_writers: HashMap::new(),
             allow_encodings: Some(ALLOWED_ENCODINGS.clone()),
             flat_strategy: None,
+            array_tree: false,
         }
     }
 }
@@ -193,11 +196,32 @@ impl WriteStrategyBuilder {
     ///
     /// By default, this uses [`FlatLayoutStrategy`]. This can be used to substitute a custom
     /// layout strategy, e.g. one that inlines constant array buffers for GPU reads.
+    ///
+    /// Passing a custom flat strategy implicitly disables the array-tree outlining feature
+    /// (see [`Self::with_array_tree`]), since the custom strategy owns the leaf format.
     pub fn with_flat_strategy(mut self, flat: Arc<dyn LayoutStrategy>) -> Self {
         self.flat_strategy = Some(flat);
         self
     }
 
+    /// Enable array-tree outlining: each chunk's encoding tree (without per-chunk statistics)
+    /// is collected into a single auxiliary segment per column rather than being inlined
+    /// alongside the chunk's data.
+    ///
+    /// Disabled by default. When enabled, the written file uses two encodings that older
+    /// readers will not understand:
+    /// [`vortex_layout::layouts::array_tree::ArrayTreeFlatLayout`] at the data leaves, and a
+    /// wrapping [`vortex_layout::layouts::array_tree::ArrayTreeLayout`] that owns the
+    /// consolidated auxiliary segment. Files written by this builder with the feature on
+    /// require a reader that recognizes both encodings.
+    ///
+    /// Has no effect if a custom flat strategy is provided via
+    /// [`Self::with_flat_strategy`] — the user-supplied leaf format wins.
+    pub fn with_array_tree(mut self, array_tree: bool) -> Self {
+        self.array_tree = array_tree;
+        self
+    }
+
     /// Override the default [`BtrBlocksCompressorBuilder`] used for compression.
     ///
     /// The builder is finalized during [`build`](Self::build), producing two compressors: one for
@@ -218,23 +242,17 @@ impl WriteStrategyBuilder {
     /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides
     /// applied.
     pub fn build(self) -> Arc<dyn LayoutStrategy> {
-        let flat: Arc<dyn LayoutStrategy> = if let Some(flat) = self.flat_strategy {
-            flat
-        } else if let Some(allow_encodings) = self.allow_encodings {
-            Arc::new(FlatLayoutStrategy::default().with_allow_encodings(allow_encodings))
+        let flat: Arc<dyn LayoutStrategy> = if let Some(flat) = &self.flat_strategy {
+            Arc::clone(flat)
+        } else if let Some(allow_encodings) = &self.allow_encodings {
+            Arc::new(FlatLayoutStrategy::default().with_allow_encodings(allow_encodings.clone()))
         } else {
             Arc::new(FlatLayoutStrategy::default())
         };
 
-        // 7. for each chunk create a flat layout
-        let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat));
-        // 6. buffer chunks so they end up with closer segment ids physically
-        let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB
-
-        // 5. compress each chunk.
-        // Exclude IntDictScheme from the data compressor because DictStrategy (step 3) already
-        // dictionary-encodes columns. Allowing IntDictScheme here would redundantly
-        // dictionary-encode the integer codes produced by that earlier step.
+        // Data compressor: excludes IntDictScheme because DictStrategy (step 3 below) already
+        // dictionary-encodes columns; allowing it here would redundantly dictionary-encode the
+        // integer codes produced by that earlier step.
         let data_compressor: Arc<dyn CompressorPlugin> = match &self.compressor {
             CompressorConfig::BtrBlocks(builder) => Arc::new(
                 builder
@@ -244,6 +262,37 @@ impl WriteStrategyBuilder {
             ),
             CompressorConfig::Opaque(compressor) => Arc::clone(compressor),
         };
+        // Stats compressor: used for zone-map tables, dict values, and (when enabled) the
+        // consolidated array-trees segment.
+        let stats_compressor: Arc<dyn CompressorPlugin> = match &self.compressor {
+            CompressorConfig::BtrBlocks(builder) => Arc::new(builder.clone().build()),
+            CompressorConfig::Opaque(compressor) => Arc::clone(compressor),
+        };
+        let compress_then_flat =
+            CompressingStrategy::new(Arc::clone(&flat), Arc::clone(&stats_compressor));
+        let compress_then_flat_arc: Arc<dyn LayoutStrategy> = Arc::new(compress_then_flat.clone());
+
+        let array_tree_enabled = self.array_tree && self.flat_strategy.is_none();
+        let (data_leaf, array_tree_collector): (Arc<dyn LayoutStrategy>, _) = if !array_tree_enabled
+        {
+            (Arc::clone(&flat), None)
+        } else {
+            let data_flat = if let Some(allow_encodings) = &self.allow_encodings {
+                FlatLayoutStrategy::default().with_allow_encodings(allow_encodings.clone())
+            } else {
+                FlatLayoutStrategy::default()
+            };
+            let (collector, leaf) =
+                array_tree_writer::writer(data_flat, Arc::clone(&compress_then_flat_arc));
+            (Arc::new(leaf), Some(collector))
+        };
+
+        // 7. for each chunk create a flat layout
+        let chunked = ChunkedLayoutStrategy::new(data_leaf);
+        // 6. buffer chunks so they end up with closer segment ids physically
+        let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB
+
+        // 5. compress each chunk.
         let compressing = CompressingStrategy::new(buffered, data_compressor);
 
         // 4. prior to compression, coalesce up to a minimum size
@@ -263,13 +312,6 @@ impl WriteStrategyBuilder {
             },
         );
 
-        // 2.1. | 3.1. compress stats tables and dict values.
-        let stats_compressor: Arc<dyn CompressorPlugin> = match self.compressor {
-            CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()),
-            CompressorConfig::Opaque(compressor) => compressor,
-        };
-        let compress_then_flat = CompressingStrategy::new(flat, stats_compressor);
-
         // 3. apply dict encoding or fallback
         let dict = DictStrategy::new(
             coalescing.clone(),
@@ -278,9 +320,16 @@ impl WriteStrategyBuilder {
             Default::default(),
         );
 
+        // 2.5. wrap dict in the array-tree collector if outlining is enabled.
+        let data_pipeline: Arc<dyn LayoutStrategy> = if let Some(collector) = array_tree_collector {
+            Arc::new(collector.wrap(dict))
+        } else {
+            Arc::new(dict)
+        };
+
         // 2. calculate stats for each row group
         let stats = ZonedStrategy::new(
-            dict,
+            data_pipeline,
             compress_then_flat.clone(),
             ZonedLayoutOptions {
                 block_size: self.row_block_size,

vortex-file/src/strategy.rs

@@ -1953,3 +1953,233 @@ async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn test_segment_ordering_array_trees_consolidated_and_after_data() -> VortexResult<()> {
+    // Multi-column struct large enough to produce chunked data, so each column will have
+    // many ArrayTreeFlat leaves. The collector should consolidate their compact trees into a
+    // single segment per ArrayTreeLayout.
+    let n = 100_000;
+    let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect();
+    let strings = VarBinArray::from(values).into_array();
+    let numbers = PrimitiveArray::from_iter(0..n as i32).into_array();
+    let floats = PrimitiveArray::from_iter((0..n).map(|i| i as f64 * 0.1)).into_array();
+
+    let st = StructArray::from_fields(&[
+        ("strings", strings),
+        ("numbers", numbers),
+        ("floats", floats),
+    ])
+    .unwrap();
+
+    let mut buf = ByteBufferMut::empty();
+    let strategy = crate::WriteStrategyBuilder::default()
+        .with_array_tree(true)
+        .build();
+    let summary = SESSION
+        .write_options()
+        .with_strategy(strategy)
+        .write(&mut buf, st.into_array().to_array_stream())
+        .await?;
+
+    let footer = summary.footer();
+    let segment_specs = footer.segment_map();
+    let root = footer.layout();
+
+    // For each ArrayTreeLayout in the tree, assert:
+    //   1. **Consolidation:** the auxiliary `array_trees` child writes exactly one segment.
+    //   2. **Per-column ordering:** every data segment under child 0 appears before the
+    //      array_trees segment under child 1.
+    fn check_array_tree_layouts(
+        layout: &dyn Layout,
+        segment_specs: &[SegmentSpec],
+        found_any: &mut bool,
+    ) {
+        if layout.encoding_id().as_ref() == "vortex.array_tree" {
+            *found_any = true;
+
+            let data_child = layout.child(0).unwrap();
+            let array_trees_child = layout.child(1).unwrap();
+
+            let data_offsets = collect_segment_offsets(data_child.as_ref(), segment_specs);
+            let array_trees_offsets =
+                collect_segment_offsets(array_trees_child.as_ref(), segment_specs);
+
+            assert_eq!(
+                array_trees_offsets.len(),
+                1,
+                "array_tree: auxiliary child must consolidate to exactly 1 segment, got {} segments at offsets {:?}",
+                array_trees_offsets.len(),
+                array_trees_offsets,
+            );
+
+            assert!(
+                !data_offsets.is_empty(),
+                "array_tree: data child must have at least one segment"
+            );
+
+            assert_offsets_ordered(
+                &data_offsets,
+                &array_trees_offsets,
+                "array_tree: all data segments should come before the array_trees segment",
+            );
+        }
+
+        for child in layout.children().unwrap() {
+            check_array_tree_layouts(child.as_ref(), segment_specs, found_any);
+        }
+    }
+
+    let mut found_any = false;
+    check_array_tree_layouts(root.as_ref(), segment_specs, &mut found_any);
+    assert!(
+        found_any,
+        "test setup expected the default write strategy to produce at least one ArrayTreeLayout"
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn test_segment_ordering_array_trees_before_zones() -> VortexResult<()> {
+    // The write strategy wraps every column in `ZonedStrategy { data: ArrayTree, zones }`.
+    // Assert per-Zoned-layout that the array_trees segment (inside the data child) appears
+    // before every zone-map segment in the same column.
+    let n = 100_000;
+    let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect();
+    let strings = VarBinArray::from(values).into_array();
+    let numbers = PrimitiveArray::from_iter(0..n as i32).into_array();
+    let floats = PrimitiveArray::from_iter((0..n).map(|i| i as f64 * 0.1)).into_array();
+
+    let st = StructArray::from_fields(&[
+        ("strings", strings),
+        ("numbers", numbers),
+        ("floats", floats),
+    ])
+    .unwrap();
+
+    let mut buf = ByteBufferMut::empty();
+    let strategy = crate::WriteStrategyBuilder::default()
+        .with_array_tree(true)
+        .build();
+    let summary = SESSION
+        .write_options()
+        .with_strategy(strategy)
+        .write(&mut buf, st.into_array().to_array_stream())
+        .await?;
+
+    let footer = summary.footer();
+    let segment_specs = footer.segment_map();
+    let root = footer.layout();
+
+    fn check_zoned_with_array_tree(
+        layout: &dyn Layout,
+        segment_specs: &[SegmentSpec],
+        found_any: &mut bool,
+    ) {
+        if layout.encoding_id().as_ref() == "vortex.stats" {
+            let data_child = layout.child(0).unwrap();
+            let zones_child = layout.child(1).unwrap();
+
+            if data_child.encoding_id().as_ref() == "vortex.array_tree" {
+                *found_any = true;
+                let array_trees_offsets =
+                    collect_segment_offsets(data_child.child(1).unwrap().as_ref(), segment_specs);
+                let zones_offsets = collect_segment_offsets(zones_child.as_ref(), segment_specs);
+
+                assert_offsets_ordered(
+                    &array_trees_offsets,
+                    &zones_offsets,
+                    "zoned wrapping array_tree: the array_trees segment should come before zone-map segments",
+                );
+            }
+        }
+
+        for child in layout.children().unwrap() {
+            check_zoned_with_array_tree(child.as_ref(), segment_specs, found_any);
+        }
+    }
+
+    let mut found_any = false;
+    check_zoned_with_array_tree(root.as_ref(), segment_specs, &mut found_any);
+    assert!(
+        found_any,
+        "test setup expected the default write strategy to produce at least one Zoned wrapping an ArrayTree"
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn test_roundtrip_array_tree_layout() -> VortexResult<()> {
+    // End-to-end coverage: write with `with_array_tree(true)`, then read back through the
+    // source-publishing reader-ctx path and assert the data matches. Exercises:
+    //   - ArrayTreeCollectorStrategy collecting compact trees from leaf transient state
+    //   - ArrayTreeLayout::new_reader publishing the ArrayTreesSource into the reader ctx
+    //   - ArrayTreeFlatReader pulling the source and resolving its tree by segment_id
+    //   - ColumnarSerializedArray::from_segment_and_tree decoding the data segment
+    let mut ctx = SESSION.create_execution_ctx();
+
+    let n = 10_000;
+    let strings_in: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect();
+    let strings = VarBinArray::from(strings_in.clone()).into_array();
+    let numbers_in: Vec<i32> = (0..n as i32).collect();
+    let numbers = PrimitiveArray::from_iter(numbers_in.iter().copied()).into_array();
+
+    let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)])?.into_array();
+    let dtype = st.dtype().clone();
+
+    let mut buf = ByteBufferMut::empty();
+    let strategy = crate::WriteStrategyBuilder::default()
+        .with_array_tree(true)
+        .build();
+    SESSION
+        .write_options()
+        .with_strategy(strategy)
+        .write(&mut buf, st.to_array_stream())
+        .await?;
+
+    // Sanity-check we actually wrote ArrayTreeLayout nodes — otherwise the test would
+    // silently pass on the default code path.
+    let file = SESSION.open_options().open_buffer(buf)?;
+    fn has_array_tree(layout: &dyn Layout) -> bool {
+        if layout.encoding_id().as_ref() == "vortex.array_tree" {
+            return true;
+        }
+        layout
+            .children()
+            .map(|cs| cs.iter().any(|c| has_array_tree(c.as_ref())))
+            .unwrap_or(false)
+    }
+    assert!(
+        has_array_tree(file.footer().layout().as_ref()),
+        "test expected ArrayTreeLayout in the written file"
+    );
+
+    let result = file.scan()?.into_array_stream()?.read_all().await?;
+    assert_eq!(result.len(), n);
+    assert_eq!(result.dtype(), &dtype);
+
+    let struct_array = result.execute::<StructArray>(&mut ctx)?;
+
+    let read_numbers = struct_array.unmasked_field_by_name("numbers").cloned()?;
+    let expected_numbers = PrimitiveArray::from_iter(numbers_in.iter().copied()).into_array();
+    assert_arrays_eq!(read_numbers, expected_numbers);
+
+    let read_strings = struct_array
+        .unmasked_field_by_name("strings")
+        .cloned()?
+        .execute::<VarBinViewArray>(&mut ctx)?
+        .with_iterator(|iter| {
+            iter.map(|s| s.map(|st| unsafe { String::from_utf8_unchecked(st.to_vec()) }))
+                .collect::<Vec<_>>()
+        });
+    let expected_strings: Vec<Option<String>> =
+        strings_in.iter().map(|s| Some((*s).to_string())).collect();
+    assert_eq!(read_strings, expected_strings);
+
+    Ok(())
+}