move to explicit but hardcoded id

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

Author: onursatici

vortex-layout/src/layouts/array_tree/flat.rs

@@ -21,6 +21,7 @@ use crate::LayoutReaderRef;
 use crate::LayoutRef;
 use crate::VTable;
 use crate::children::LayoutChildren;
+use crate::layouts::array_tree::ARRAY_TREES_SOURCE_ID;
 use crate::layouts::array_tree::ArrayTreesSource;
 use crate::layouts::array_tree::reader::ArrayTreeFlatReader;
 use crate::layouts::flat::FlatLayout;
@@ -131,14 +132,16 @@ impl VTable for ArrayTreeFlat {
         session: &VortexSession,
         ctx: &LayoutReaderContext,
     ) -> VortexResult<LayoutReaderRef> {
-        let source = ctx.get::<ArrayTreesSource>().ok_or_else(|| {
-            vortex_error::vortex_err!(
-                "ArrayTreeFlatLayout requires an ancestor ArrayTreeLayout to publish an \
+        let source = ctx
+            .get::<ArrayTreesSource>(*ARRAY_TREES_SOURCE_ID)
+            .ok_or_else(|| {
+                vortex_error::vortex_err!(
+                    "ArrayTreeFlatLayout requires an ancestor ArrayTreeLayout to publish an \
                  ArrayTreesSource into the reader context; call \
                  ArrayTreeLayout::derive_reader_ctx on each ArrayTreeLayout ancestor before \
                  constructing a reader for this layout"
-            )
-        })?;
+                )
+            })?;
         Ok(Arc::new(ArrayTreeFlatReader::new(
             layout.clone(),
             name,

vortex-layout/src/layouts/array_tree/mod.rs

@@ -14,6 +14,7 @@ mod reader;
 pub mod writer;
 
 use std::sync::Arc;
+use std::sync::LazyLock;
 use std::sync::OnceLock;
 
 use futures::FutureExt;
@@ -38,15 +39,17 @@ use vortex_array::dtype::StructFields;
 use vortex_array::expr::root;
 use vortex_array::serde::BUFFER_COLUMNS_DTYPE;
 use vortex_array::serde::ColumnarArrayTree;
-use vortex_array::serde::NODES_COLUMNS_DTYPE;
+use vortex_array::serde::DEFAULT_STATS;
 use vortex_array::serde::StatsColumns;
+use vortex_array::serde::nodes_columns_dtype;
 use vortex_error::SharedVortexResult;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
+use vortex_session::registry::Id;
 use vortex_session::registry::ReadContext;
 use vortex_utils::aliases::hash_map::HashMap;
 
@@ -69,6 +72,17 @@ use crate::vtable;
 
 vtable!(ArrayTree);
 
+/// Well-known [`LayoutReaderContext`] key under which [`ArrayTreeLayout::derive_reader_ctx`]
+/// publishes its [`ArrayTreesSource`].
+///
+/// Both the publisher (parent [`ArrayTreeLayout`]) and the consumer
+/// ([`ArrayTreeFlatLayout`]'s `new_reader`) hardcode this id, so no metadata persistence is
+/// needed to bind them. Two stacked `ArrayTreeLayouts` both publish under this id; the
+/// inner one overrides the outer in the descendant's view — exactly the "nearest ancestor
+/// wins" semantic each `ArrayTreeFlat` leaf wants.
+pub static ARRAY_TREES_SOURCE_ID: LazyLock<Id> =
+    LazyLock::new(|| Id::new_static("vortex.array_tree.source"));
+
 /// Encoding marker for [`ArrayTreeLayout`].
 #[derive(Debug)]
 pub struct ArrayTreeLayoutEncoding;
@@ -109,10 +123,15 @@ impl ArrayTreeLayout {
     /// }
     /// ```
     ///
-    /// The List<> element types are exactly the canonical per-tree schemas defined by
-    /// [`vortex_array::serde::NODES_COLUMNS_DTYPE`] / [`vortex_array::serde::BUFFER_COLUMNS_DTYPE`],
-    /// so slicing one row's nodes (resp. buffers) yields a [`StructArray`] matching the inner
+    /// The List<> element types are produced by
+    /// [`vortex_array::serde::nodes_columns_dtype`] / [`vortex_array::serde::BUFFER_COLUMNS_DTYPE`].
+    /// Slicing one row's nodes (resp. buffers) yields a [`StructArray`] matching the inner
     /// shape of a [`ColumnarArrayTree`].
+    ///
+    /// The `nodes` schema is parameterized by the stat menu the layout tracks; today this
+    /// defaults to [`vortex_array::serde::DEFAULT_STATS`] (the historical 11-stat set), but
+    /// future writers can pick a different menu and the schema field names carry that menu
+    /// across the wire.
     pub fn array_trees_dtype() -> DType {
         let nn = Nullability::NonNullable;
         DType::Struct(
@@ -125,7 +144,7 @@ impl ArrayTreeLayout {
                 .into(),
                 vec![
                     DType::Primitive(PType::U32, nn),
-                    DType::List(Arc::new(NODES_COLUMNS_DTYPE.clone()), nn),
+                    DType::List(Arc::new(nodes_columns_dtype(DEFAULT_STATS)), nn),
                     DType::List(Arc::new(BUFFER_COLUMNS_DTYPE.clone()), nn),
                 ],
             ),
@@ -134,8 +153,9 @@ impl ArrayTreeLayout {
     }
 
     /// Derive a [`LayoutReaderContext`] that publishes an [`ArrayTreesSource`] backed by this
-    /// layout's auxiliary `array_trees` child. Descendant [`ArrayTreeFlatLayout`] readers
-    /// pull the source via `ctx.get::<ArrayTreesSource>()` to resolve their compact trees.
+    /// layout's auxiliary `array_trees` child under [`ARRAY_TREES_SOURCE_ID`]. Descendant
+    /// [`ArrayTreeFlatLayout`] readers pull the source by the same id to resolve their
+    /// compact trees.
     ///
     /// Used by:
     /// - The normal [`crate::VTable::new_reader`] dispatch on `ArrayTreeLayout` (production path).
@@ -158,7 +178,7 @@ impl ArrayTreeLayout {
             ctx,
         )?;
         let source = Arc::new(ArrayTreesSource::new(trees_reader, session.clone()));
-        Ok(ctx.with(source))
+        Ok(ctx.with(*ARRAY_TREES_SOURCE_ID, source))
     }
 }
 
@@ -259,8 +279,8 @@ impl VTable for ArrayTree {
 /// the cached map.
 ///
 /// Published by [`ArrayTreeLayout::derive_reader_ctx`] into the [`LayoutReaderContext`]
-/// passed to descendants; pulled by [`ArrayTreeFlatLayout`]'s reader via
-/// `ctx.get::<ArrayTreesSource>()`.
+/// passed to descendants under [`ARRAY_TREES_SOURCE_ID`]; pulled by
+/// [`ArrayTreeFlatLayout`]'s reader by the same id.
 pub struct ArrayTreesSource {
     reader: LayoutReaderRef,
     /// Session used to create execution contexts when canonicalizing the consolidated array

vortex-layout/src/layouts/array_tree/reader.rs

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::collections::BTreeSet;
 use std::ops::BitAnd;
 use std::ops::Range;
 use std::sync::Arc;
@@ -24,6 +23,7 @@ use crate::layouts::SharedArrayFuture;
 use crate::layouts::array_tree::ArrayTreesSource;
 use crate::layouts::array_tree::flat::ArrayTreeFlatLayout;
 use crate::reader::ArrayFuture;
+use crate::reader::RowSplits;
 use crate::reader::SplitRange;
 use crate::segments::SegmentSource;
 
@@ -62,7 +62,7 @@ impl LayoutReader for ArrayTreeReader {
         &self,
         field_mask: &[FieldMask],
         split_range: &SplitRange,
-        splits: &mut BTreeSet<u64>,
+        splits: &mut RowSplits,
     ) -> VortexResult<()> {
         self.data_reader
             .register_splits(field_mask, split_range, splits)
@@ -178,10 +178,10 @@ impl LayoutReader for ArrayTreeFlatReader {
         &self,
         _field_mask: &[FieldMask],
         split_range: &SplitRange,
-        splits: &mut BTreeSet<u64>,
+        splits: &mut RowSplits,
     ) -> VortexResult<()> {
         split_range.check_bounds(self.layout.inner().row_count())?;
-        splits.insert(split_range.root_row_range().end);
+        splits.push(split_range.root_row_range().end);
         Ok(())
     }