Skip to content

Commit 1b55509

Browse files
committed
list: top-level streaming decomposition behind VORTEX_EXPERIMENTAL_LIST_LAYOUT
Rework ListLayoutStrategy into a streaming transpose driven by the TableStrategy dispatcher (global u64 offsets, memory-bounded), add bounded element reads and is_null pruning routing in ListReader, and gate the whole thing off by default behind the VORTEX_EXPERIMENTAL_LIST_LAYOUT env var / with_list_layout(). Signed-off-by: Matt Katz <mhkatz97@gmail.com>
1 parent 86d3068 commit 1b55509

6 files changed

Lines changed: 630 additions & 118 deletions

File tree

vortex-file/src/strategy.rs

Lines changed: 26 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,10 @@ use vortex_layout::layouts::compressed::CompressingStrategy;
5050
use vortex_layout::layouts::compressed::CompressorPlugin;
5151
use vortex_layout::layouts::dict::writer::DictStrategy;
5252
use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
53-
use vortex_layout::layouts::list::writer::ListLayoutStrategy;
5453
use vortex_layout::layouts::repartition::RepartitionStrategy;
5554
use vortex_layout::layouts::repartition::RepartitionWriterOptions;
5655
use vortex_layout::layouts::table::TableStrategy;
56+
use vortex_layout::layouts::table::use_experimental_list_layout;
5757
use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions;
5858
use vortex_layout::layouts::zoned::writer::ZonedStrategy;
5959
#[cfg(feature = "unstable_encodings")]
@@ -158,6 +158,10 @@ pub struct WriteStrategyBuilder {
158158
allow_encodings: Option<HashSet<ArrayId>>,
159159
flat_strategy: Option<Arc<dyn LayoutStrategy>>,
160160
probe_compressor: Option<Arc<dyn CompressorPlugin>>,
161+
/// Force list-column decomposition on, overriding the
162+
/// [`use_experimental_list_layout`](vortex_layout::layouts::table::use_experimental_list_layout)
163+
/// env-var default of off.
164+
list_layout: bool,
161165
}
162166

163167
impl Default for WriteStrategyBuilder {
@@ -171,6 +175,7 @@ impl Default for WriteStrategyBuilder {
171175
allow_encodings: Some(ALLOWED_ENCODINGS.clone()),
172176
flat_strategy: None,
173177
probe_compressor: None,
178+
list_layout: false,
174179
}
175180
}
176181
}
@@ -185,6 +190,14 @@ impl WriteStrategyBuilder {
185190
self
186191
}
187192

193+
/// Enable list-column decomposition, overriding the env-var default of off. When enabled, list
194+
/// columns are written as a `ListLayout` (independently compressed/chunked
195+
/// elements/offsets/validity) rather than as flat arrays.
196+
pub fn with_list_layout(mut self) -> Self {
197+
self.list_layout = true;
198+
self
199+
}
200+
188201
/// Override the write layout for a specific field somewhere in the nested schema tree.
189202
///
190203
/// The field path is matched after the root struct is split into columns. This is useful when a
@@ -251,20 +264,10 @@ impl WriteStrategyBuilder {
251264
Arc::new(FlatLayoutStrategy::default())
252265
};
253266

254-
// 7. for each chunk create a layout. List-typed chunks route through
255-
// `ListLayoutStrategy` (separately-addressable elements/offsets/validity sub-layouts;
256-
// non-list chunks fall through its built-in fallback to `flat`).
257-
let leaf: Arc<dyn LayoutStrategy> = Arc::new(
258-
// Thread the configured `flat` (which carries `allow_encodings` / any custom flat
259-
// override) through every child.
260-
ListLayoutStrategy::default()
261-
.with_elements(Arc::clone(&flat))
262-
.with_offsets(Arc::clone(&flat))
263-
.with_validity(Arc::clone(&flat))
264-
.with_fallback(Arc::clone(&flat)),
265-
);
266-
267-
let chunked = ChunkedLayoutStrategy::new(leaf);
267+
// 7. for each chunk create a flat layout. List columns are decomposed above this point by
268+
// the `TableStrategy` dispatcher (into independently-compressed elements/offsets/validity
269+
// sub-columns), so the per-chunk leaf only ever sees flat, non-list chunks.
270+
let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat));
268271
// 6. buffer chunks so they end up with closer segment ids physically
269272
let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB
270273

@@ -348,8 +351,14 @@ impl WriteStrategyBuilder {
348351
let validity_strategy = CollectStrategy::new(compress_then_flat);
349352

350353
// Take any field overrides from the builder and apply them to the final strategy.
351-
let table_strategy = TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition))
352-
.with_field_writers(self.field_writers);
354+
let mut table_strategy =
355+
TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition))
356+
.with_field_writers(self.field_writers);
357+
// List decomposition is experimental: enabled by an explicit builder opt-in or the
358+
// `VORTEX_EXPERIMENTAL_LIST_LAYOUT` env var; off otherwise.
359+
if self.list_layout || use_experimental_list_layout() {
360+
table_strategy = table_strategy.with_list_layout();
361+
}
353362

354363
Arc::new(table_strategy)
355364
}

vortex-file/src/tests.rs

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ use vortex_array::ArrayRef;
1515
use vortex_array::IntoArray;
1616
use vortex_array::VortexSessionExecute;
1717
use vortex_array::array_session;
18+
use vortex_array::arrays::BoolArray;
1819
use vortex_array::arrays::ChunkedArray;
1920
use vortex_array::arrays::ConstantArray;
2021
use vortex_array::arrays::DecimalArray;
@@ -1671,6 +1672,75 @@ async fn test_writer_with_complex_types() -> VortexResult<()> {
16711672
Ok(())
16721673
}
16731674

1675+
/// Write `array` with list decomposition forced on (through the full compress/zone pipeline) and
1676+
/// read the whole thing back.
1677+
async fn write_read_roundtrip(array: ArrayRef) -> VortexResult<ArrayRef> {
1678+
let strategy = crate::strategy::WriteStrategyBuilder::default()
1679+
.with_list_layout()
1680+
.build();
1681+
let mut buf = ByteBufferMut::empty();
1682+
SESSION
1683+
.write_options()
1684+
.with_strategy(strategy)
1685+
.write(&mut buf, array.to_array_stream())
1686+
.await?;
1687+
SESSION
1688+
.open_options()
1689+
.open_buffer(buf)?
1690+
.scan()?
1691+
.into_array_stream()?
1692+
.read_all()
1693+
.await
1694+
}
1695+
1696+
/// A `list<list<i32>>` column round-trips through the `TableStrategy` dispatcher, exercising list
1697+
/// decomposition recursing into itself (the outer list's `elements` are themselves lists).
1698+
#[tokio::test]
1699+
#[cfg_attr(miri, ignore)]
1700+
async fn nested_list_of_list_roundtrip() -> VortexResult<()> {
1701+
let inner = ListArray::try_new(
1702+
buffer![1i32, 2, 3, 4, 5, 6].into_array(),
1703+
buffer![0u32, 2, 5, 5, 6].into_array(),
1704+
Validity::NonNullable,
1705+
)?
1706+
.into_array();
1707+
let outer = ListArray::try_new(
1708+
inner,
1709+
buffer![0u32, 2, 4].into_array(),
1710+
Validity::NonNullable,
1711+
)?
1712+
.into_array();
1713+
let st = StructArray::from_fields(&[("nested", outer)])?.into_array();
1714+
1715+
let result = write_read_roundtrip(st.clone()).await?;
1716+
assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx());
1717+
Ok(())
1718+
}
1719+
1720+
/// A `struct<{ items: list<struct<{a,b}>>? }>` column round-trips, exercising list decomposition
1721+
/// recursing into struct decomposition (list `elements` are structs) plus a nullable list validity
1722+
/// child.
1723+
#[tokio::test]
1724+
#[cfg_attr(miri, ignore)]
1725+
async fn nested_struct_list_struct_roundtrip() -> VortexResult<()> {
1726+
let inner_struct = StructArray::from_fields(&[
1727+
("a", buffer![1i32, 2, 3, 4, 5].into_array()),
1728+
("b", buffer![10i32, 20, 30, 40, 50].into_array()),
1729+
])?
1730+
.into_array();
1731+
let items = ListArray::try_new(
1732+
inner_struct,
1733+
buffer![0u32, 2, 5, 5].into_array(),
1734+
Validity::Array(BoolArray::from_iter([true, false, true]).into_array()),
1735+
)?
1736+
.into_array();
1737+
let st = StructArray::from_fields(&[("items", items)])?.into_array();
1738+
1739+
let result = write_read_roundtrip(st.clone()).await?;
1740+
assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx());
1741+
Ok(())
1742+
}
1743+
16741744
#[tokio::test]
16751745
async fn test_writer_with_statistics() -> VortexResult<()> {
16761746
let array = StructArray::from_fields(&[("numbers", buffer![1u32, 2, 3, 4, 5].into_array())])?

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4-
mod reader;
4+
pub(crate) mod reader;
55
pub mod writer;
66

77
use std::sync::Arc;

0 commit comments

Comments
 (0)