diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml index 76282a3d272..b72f3e72602 100644 --- a/.github/workflows/fuzz.yml +++ b/.github/workflows/fuzz.yml @@ -213,6 +213,42 @@ jobs: gh_token: ${{ secrets.GITHUB_TOKEN }} incident_io_alert_token: ${{ secrets.INCIDENT_IO_ALERT_TOKEN }} + # ============================================================================ + # Tiled Fixed-Size List Fuzzer + # ============================================================================ + tiled_fsl_fuzz: + name: "Tiled Fixed-Size List Fuzz" + uses: ./.github/workflows/run-fuzzer.yml + with: + fuzz_target: tiled_fsl + jobs: 4 + secrets: + R2_FUZZ_ACCESS_KEY_ID: ${{ secrets.R2_FUZZ_ACCESS_KEY_ID }} + R2_FUZZ_SECRET_ACCESS_KEY: ${{ secrets.R2_FUZZ_SECRET_ACCESS_KEY }} + + report-tiled-fsl-fuzz-failures: + name: "Report Tiled Fixed-Size List Fuzz Failures" + needs: tiled_fsl_fuzz + if: always() && needs.tiled_fsl_fuzz.outputs.crashes_found == 'true' + permissions: + issues: write + contents: read + id-token: write + pull-requests: read + uses: ./.github/workflows/report-fuzz-crash.yml + with: + fuzz_target: tiled_fsl + crash_file: ${{ needs.tiled_fsl_fuzz.outputs.first_crash_name }} + artifact_url: ${{ needs.tiled_fsl_fuzz.outputs.artifact_url }} + artifact_name: tiled_fsl-crash-artifacts + logs_artifact_name: tiled_fsl-logs + branch: ${{ github.ref_name }} + commit: ${{ github.sha }} + secrets: + claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }} + gh_token: ${{ secrets.GITHUB_TOKEN }} + incident_io_alert_token: ${{ secrets.INCIDENT_IO_ALERT_TOKEN }} + # ============================================================================ # Compress Roundtrip Fuzzer # ============================================================================ diff --git a/Cargo.lock b/Cargo.lock index 5be1371fac3..2e82edd4ab8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9355,6 +9355,7 @@ dependencies = [ "vortex-session", "vortex-sparse", "vortex-tensor", + "vortex-tiled-fsl", "vortex-utils", "vortex-zigzag", "vortex-zstd", @@ -9942,6 +9943,7 @@ dependencies = [ "vortex-session", "vortex-sparse", "vortex-tensor", + "vortex-tiled-fsl", "vortex-utils", "vortex-zigzag", "vortex-zstd", @@ -9998,6 +10000,7 @@ dependencies = [ "vortex-row", "vortex-runend", "vortex-session", + "vortex-tiled-fsl", "vortex-utils", ] @@ -10449,6 +10452,22 @@ dependencies = [ "vortex-cuda", ] +[[package]] +name = "vortex-tiled-fsl" +version = "0.1.0" +dependencies = [ + "codspeed-divan-compat", + "mimalloc", + "prost 0.14.4", + "rstest", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-fastlanes", + "vortex-mask", + "vortex-session", +] + [[package]] name = "vortex-tui" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 5cce58ae60d..4ff01b6b66b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,6 +62,7 @@ members = [ "encodings/bytebool", "encodings/parquet-variant", "encodings/onpair", + "encodings/tiled-fsl", # Benchmarks "benchmarks/lance-bench", "benchmarks/compress-bench", @@ -325,6 +326,7 @@ vortex-sequence = { version = "0.1.0", path = "encodings/sequence", default-feat vortex-session = { version = "0.1.0", path = "./vortex-session", default-features = false } vortex-sparse = { version = "0.1.0", path = "./encodings/sparse", default-features = false } vortex-tensor = { version = "0.1.0", path = "./vortex-tensor", default-features = false } +vortex-tiled-fsl = { version = "0.1.0", path = "./encodings/tiled-fsl", default-features = false } vortex-utils = { version = "0.1.0", path = "./vortex-utils", default-features = false } vortex-zigzag = { version = "0.1.0", path = "./encodings/zigzag", default-features = false } vortex-zstd = { version = "0.1.0", path = "./encodings/zstd", default-features = false } diff --git a/encodings/tiled-fsl/Cargo.toml b/encodings/tiled-fsl/Cargo.toml new file mode 100644 index 00000000000..918f61dd552 --- /dev/null +++ b/encodings/tiled-fsl/Cargo.toml @@ -0,0 +1,36 @@ +[package] +name = "vortex-tiled-fsl" +authors = { workspace = true } +categories = { workspace = true } +description = "Two-dimensional tiled encoding for primitive Vortex fixed-size lists" +edition = { workspace = true } +homepage = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[lints] +workspace = true + +[dependencies] +prost = { workspace = true } +vortex-array = { workspace = true } +vortex-buffer = { workspace = true } +vortex-error = { workspace = true } +vortex-mask = { workspace = true } +vortex-session = { workspace = true } + +[dev-dependencies] +divan = { workspace = true } +mimalloc = { workspace = true } +rstest = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-fastlanes = { workspace = true } + +[[bench]] +name = "tiled_fsl" +harness = false diff --git a/encodings/tiled-fsl/benches/tiled_fsl.rs b/encodings/tiled-fsl/benches/tiled_fsl.rs new file mode 100644 index 00000000000..9b4eeecec0a --- /dev/null +++ b/encodings/tiled-fsl/benches/tiled_fsl.rs @@ -0,0 +1,787 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::cast_possible_truncation)] +#![expect(clippy::unwrap_used)] + +use std::fmt; +use std::num::NonZeroU32; +use std::ops::Range; +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::FixedSizeList; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::slice::SliceReduce; +use vortex_array::assert_arrays_eq; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_fastlanes::bitpack_compress::bitpack_encode; +use vortex_session::VortexSession; +use vortex_tiled_fsl::TileGeometry; +use vortex_tiled_fsl::TiledFixedSizeList; +use vortex_tiled_fsl::TiledFixedSizeListArray; +use vortex_tiled_fsl::TiledFixedSizeListArrayExt; +use vortex_tiled_fsl::TiledFixedSizeListArraySlotsExt; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + assert_fixture_matrix(); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_tiled_fsl::initialize(&session); + vortex_fastlanes::initialize(&session); + session +}); + +#[derive(Clone, Copy)] +struct Args { + rows: usize, + dimensions: usize, + tile_rows: u32, + tile_dimensions: TileDimensions, +} + +#[derive(Clone, Copy)] +enum TileDimensions { + Full, + Fixed(u32), +} + +impl fmt::Display for Args { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let dimension_tile = match self.tile_dimensions { + TileDimensions::Full => "full".to_owned(), + TileDimensions::Fixed(dimensions) => dimensions.to_string(), + }; + write!( + f, + "rows{}_dims{}_tile{}x{}", + self.rows, self.dimensions, self.tile_rows, dimension_tile, + ) + } +} + +fn args() -> Vec { + let mut args = Vec::new(); + let geometries = [ + (32, TileDimensions::Full), + (64, TileDimensions::Full), + (32, TileDimensions::Fixed(64)), + (64, TileDimensions::Fixed(64)), + (16, TileDimensions::Fixed(4)), + ]; + for rows in [1_024, 16_384] { + for dimensions in [128, 768, 1_536] { + for (tile_rows, tile_dimensions) in geometries { + args.push(Args { + rows, + dimensions, + tile_rows, + tile_dimensions, + }); + } + } + } + for rows in [31, 33, 63, 65] { + for dimensions in [31, 33, 63, 65] { + for tile_rows in [32, 64] { + args.push(Args { + rows, + dimensions, + tile_rows, + tile_dimensions: TileDimensions::Fixed(64), + }); + } + } + } + args +} + +#[derive(Clone, Copy)] +enum SliceKind { + Small, + Half, + CrossTileBoundary, +} + +impl fmt::Display for SliceKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Small => "small", + Self::Half => "half", + Self::CrossTileBoundary => "cross_tile_boundary", + }) + } +} + +#[derive(Clone, Copy)] +struct SliceArgs { + args: Args, + kind: SliceKind, +} + +impl fmt::Display for SliceArgs { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}/{}", self.kind, self.args) + } +} + +fn slice_args() -> Vec { + args() + .into_iter() + .flat_map(|args| { + [SliceKind::Small, SliceKind::Half] + .into_iter() + .chain( + (args.rows > args.tile_rows as usize).then_some(SliceKind::CrossTileBoundary), + ) + .map(move |kind| SliceArgs { args, kind }) + }) + .collect() +} + +#[derive(Clone, Copy)] +enum TakeKind { + SortedSparse, + UnsortedDuplicated, +} + +impl fmt::Display for TakeKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::SortedSparse => "sorted_sparse", + Self::UnsortedDuplicated => "unsorted_duplicated", + }) + } +} + +#[derive(Clone, Copy)] +struct TakeArgs { + args: Args, + kind: TakeKind, +} + +impl fmt::Display for TakeArgs { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}/{}", self.kind, self.args) + } +} + +fn take_args() -> Vec { + args() + .into_iter() + .flat_map(|args| { + [TakeKind::SortedSparse, TakeKind::UnsortedDuplicated] + .into_iter() + .map(move |kind| TakeArgs { args, kind }) + }) + .collect() +} + +#[derive(Clone, Copy)] +enum PhysicalEncoding { + Raw, + Bitpacked, +} + +impl fmt::Display for PhysicalEncoding { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Raw => "raw", + Self::Bitpacked => "bitpacked_4bit", + }) + } +} + +#[derive(Clone, Copy)] +struct ScoreArgs { + args: Args, + encoding: PhysicalEncoding, +} + +impl fmt::Display for ScoreArgs { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}/{}", self.encoding, self.args) + } +} + +fn score_args() -> Vec { + args() + .into_iter() + .flat_map(|args| { + [PhysicalEncoding::Raw, PhysicalEncoding::Bitpacked] + .into_iter() + .map(move |encoding| ScoreArgs { args, encoding }) + }) + .collect() +} + +fn tile_geometry(args: Args) -> TileGeometry { + let dimensions = match args.tile_dimensions { + TileDimensions::Full => u32::try_from(args.dimensions).unwrap(), + TileDimensions::Fixed(dimensions) => dimensions, + }; + TileGeometry::new( + NonZeroU32::new(args.tile_rows).unwrap(), + NonZeroU32::new(dimensions).unwrap(), + ) +} + +fn fixture_value(row: usize, dimension: usize, _dimensions: usize) -> u8 { + // Advancing one row changes every coordinate by 5 modulo 16, regardless of row width. + ((row * 5 + dimension * 3) & 0x0f) as u8 +} + +fn representative_rows(args: Args) -> (Vec, Vec) { + ( + (0..args.dimensions) + .map(|dimension| fixture_value(0, dimension, args.dimensions)) + .collect(), + (0..args.dimensions) + .map(|dimension| fixture_value(1, dimension, args.dimensions)) + .collect(), + ) +} + +fn assert_fixture_matrix() { + let mut indistinct_widths = args() + .into_iter() + .filter_map(|args| { + let (first, second) = representative_rows(args); + (first == second).then_some(args.dimensions) + }) + .collect::>(); + indistinct_widths.sort_unstable(); + indistinct_widths.dedup(); + assert!( + indistinct_widths.is_empty(), + "adjacent fixture rows are identical for dimensions {indistinct_widths:?}", + ); +} + +fn assert_adjacent_payloads_distinct(array: &FixedSizeListArray, args: Args) { + if args.rows < 2 { + return; + } + let elements = array.elements().as_::(); + let values = elements.as_slice::(); + assert_ne!( + &values[..args.dimensions], + &values[args.dimensions..args.dimensions * 2], + "adjacent canonical rows are identical for {args}", + ); +} + +fn assert_adjacent_scores_distinct(values: &[u8], args: Args, query: &[u8]) { + if args.rows < 2 { + return; + } + let scores = + scoring::score_canonical(&values[..args.dimensions * 2], 2, args.dimensions, query); + assert_ne!( + scores[0], scores[1], + "adjacent canonical row scores are identical for {args}", + ); +} + +fn canonical_u8(args: Args) -> FixedSizeListArray { + let array = FixedSizeListArray::new( + PrimitiveArray::from_iter((0..args.rows).flat_map(|row| { + (0..args.dimensions) + .map(move |dimension| fixture_value(row, dimension, args.dimensions)) + })) + .into_array(), + u32::try_from(args.dimensions).unwrap(), + Validity::NonNullable, + args.rows, + ); + assert_adjacent_payloads_distinct(&array, args); + array +} + +fn query(args: Args) -> Vec { + (0..args.dimensions) + .map(|dimension| ((dimension * 13 + 7) & 0x0f) as u8) + .collect() +} + +fn raw_tiled(args: Args, ctx: &mut ExecutionCtx) -> VortexResult { + TiledFixedSizeList::encode(canonical_u8(args).as_view(), tile_geometry(args), ctx) +} + +fn bitpacked_tiled(args: Args, ctx: &mut ExecutionCtx) -> VortexResult { + let raw = raw_tiled(args, ctx)?; + let physical = raw.elements().clone().execute::(ctx)?; + let bitpacked = bitpack_encode(&physical, 4, None, ctx)?.into_array(); + TiledFixedSizeList::try_new( + bitpacked, + u32::try_from(args.dimensions)?, + raw.array_validity(), + args.rows, + tile_geometry(args), + ) +} + +fn canonical_f32(args: Args) -> FixedSizeListArray { + FixedSizeListArray::new( + PrimitiveArray::from_iter( + (0..args.rows * args.dimensions).map(|index| ((index * 17) % 1_009) as f32 / 1_009.0), + ) + .into_array(), + u32::try_from(args.dimensions).unwrap(), + Validity::NonNullable, + args.rows, + ) +} + +fn canonical_nullable_f32(args: Args) -> FixedSizeListArray { + let element_count = args.rows * args.dimensions; + FixedSizeListArray::new( + PrimitiveArray::new( + Buffer::from_iter( + (0..element_count).map(|index| ((index * 17) % 1_009) as f32 / 1_009.0), + ), + Validity::from_iter((0..element_count).map(|index| index % 11 != 0)), + ) + .into_array(), + u32::try_from(args.dimensions).unwrap(), + Validity::NonNullable, + args.rows, + ) +} + +fn assert_tiled_matches( + canonical: &FixedSizeListArray, + tiled: &TiledFixedSizeListArray, + ctx: &mut ExecutionCtx, +) { + assert_arrays_eq!(canonical, tiled, ctx); +} + +fn tiled_score_fixture( + args: ScoreArgs, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let tiled = match args.encoding { + PhysicalEncoding::Raw => raw_tiled(args.args, ctx), + PhysicalEncoding::Bitpacked => bitpacked_tiled(args.args, ctx), + }?; + assert_eq!(tiled.geometry(), tile_geometry(args.args)); + Ok(tiled) +} + +fn slice_range(args: SliceArgs) -> Range { + match args.kind { + SliceKind::Small => { + let start = args.args.rows / 3; + start..start + 8 + } + SliceKind::Half => { + let start = args.args.rows / 4; + start..start + args.args.rows / 2 + } + SliceKind::CrossTileBoundary => { + let boundary = args.args.tile_rows as usize; + boundary - 1..boundary + 1 + } + } +} + +fn take_indices(args: TakeArgs) -> PrimitiveArray { + let rows = u32::try_from(args.args.rows).unwrap(); + match args.kind { + TakeKind::SortedSparse => { + PrimitiveArray::from_iter([0, rows / 4, rows / 2, rows * 3 / 4, rows - 1]) + } + TakeKind::UnsortedDuplicated => { + PrimitiveArray::from_iter([rows - 1, 0, rows / 2, rows / 2, 1, rows - 1]) + } + } +} + +fn assert_scores_equal( + args: Args, + tiled: &TiledFixedSizeListArray, + tiled_values: &[u8], + query: &[u8], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let canonical = canonical_u8(args); + let canonical_values = canonical + .elements() + .clone() + .execute::(ctx)?; + assert_adjacent_scores_distinct(canonical_values.as_slice::(), args, query); + assert_eq!( + scoring::score_canonical( + canonical_values.as_slice::(), + args.rows, + args.dimensions, + query, + ), + scoring::score_tiled(tiled.as_view(), tiled_values, query), + "canonical and tiled scores differ for {args}", + ); + + if tiled.is_full_width() && args.rows > 2 { + let range = 1..args.rows - 1; + let canonical = canonical.into_array().slice(range.clone())?; + let canonical = canonical.execute::(ctx)?; + let canonical_elements = canonical.elements(); + let canonical_primitive = canonical_elements.as_::(); + let canonical_values = canonical_primitive.as_slice::(); + let tiled = tiled.clone().into_array().slice(range)?; + let tiled = tiled.as_::(); + let tiled_values = tiled.elements().clone().execute::(ctx)?; + assert_eq!( + scoring::score_canonical(canonical_values, canonical.len(), args.dimensions, query,), + scoring::score_tiled(tiled, tiled_values.as_slice::(), query), + "canonical and tiled view scores differ for {args}", + ); + } + Ok(()) +} + +mod encode { + use super::*; + + #[divan::bench(args = args())] + fn non_nullable(bencher: Bencher, args: Args) { + bench_encode(bencher, args, canonical_f32(args)); + } + + #[divan::bench(args = args())] + fn nullable_bitmap(bencher: Bencher, args: Args) { + bench_encode(bencher, args, canonical_nullable_f32(args)); + } + + fn bench_encode(bencher: Bencher, args: Args, canonical: FixedSizeListArray) { + let mut ctx = SESSION.create_execution_ctx(); + let oracle = + TiledFixedSizeList::encode(canonical.as_view(), tile_geometry(args), &mut ctx).unwrap(); + assert_tiled_matches(&canonical, &oracle, &mut ctx); + + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_values(|mut ctx| { + divan::black_box( + TiledFixedSizeList::encode(canonical.as_view(), tile_geometry(args), &mut ctx) + .unwrap(), + ) + }); + } +} + +mod execute { + use super::*; + + #[divan::bench(args = args())] + fn non_nullable(bencher: Bencher, args: Args) { + bench_execute(bencher, args, canonical_f32(args)); + } + + #[divan::bench(args = args())] + fn nullable_bitmap(bencher: Bencher, args: Args) { + bench_execute(bencher, args, canonical_nullable_f32(args)); + } + + fn bench_execute(bencher: Bencher, args: Args, canonical: FixedSizeListArray) { + let mut ctx = SESSION.create_execution_ctx(); + let tiled = + TiledFixedSizeList::encode(canonical.as_view(), tile_geometry(args), &mut ctx).unwrap(); + assert_tiled_matches(&canonical, &tiled, &mut ctx); + + bencher + .with_inputs(|| (tiled.clone().into_array(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + divan::black_box(array.execute::(&mut ctx).unwrap()) + }); + } +} + +mod slice_reduce { + use super::*; + + #[divan::bench(args = [1_024usize, 1_000_000])] + fn full_width_unaligned(bencher: Bencher, rows: usize) { + let args = Args { + rows, + dimensions: 1, + tile_rows: 64, + tile_dimensions: TileDimensions::Full, + }; + let canonical = canonical_u8(args); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = raw_tiled(args, &mut ctx).unwrap(); + let range = 1..rows - 1; + let expected = canonical.into_array().slice(range.clone()).unwrap(); + let reduced = ::slice(tiled.as_view(), range.clone()) + .unwrap() + .unwrap(); + assert_arrays_eq!(expected, reduced, &mut ctx); + + bencher + .with_inputs(|| (tiled.clone(), range.clone())) + .bench_values(|(array, range)| { + divan::black_box( + ::slice(array.as_view(), range) + .unwrap() + .unwrap(), + ) + }); + } +} + +mod slice_execute { + use super::*; + + #[divan::bench(args = [SliceKind::Small, SliceKind::Half])] + fn multi_slab_unaligned(bencher: Bencher, kind: SliceKind) { + let args = Args { + rows: 16_384, + dimensions: 128, + tile_rows: 64, + tile_dimensions: TileDimensions::Fixed(64), + }; + let slice_args = SliceArgs { args, kind }; + let range = slice_range(slice_args); + let canonical = canonical_u8(args); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = raw_tiled(args, &mut ctx).unwrap(); + let expected = canonical.into_array().slice(range.clone()).unwrap(); + let lazy_slice = tiled.into_array().slice(range).unwrap(); + let executed = lazy_slice + .clone() + .execute::(&mut ctx) + .unwrap(); + assert_arrays_eq!(expected, executed, &mut ctx); + + bencher + .with_inputs(|| (lazy_slice.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + divan::black_box(array.execute::(&mut ctx).unwrap()) + }); + } +} + +mod tile_iteration { + use super::*; + + #[divan::bench] + fn full_view(bencher: Bencher) { + bench_view(bencher, None); + } + + #[divan::bench] + fn prefix_boundary(bencher: Bencher) { + bench_view(bencher, Some(1..64)); + } + + #[divan::bench] + fn two_boundaries(bencher: Bencher) { + bench_view(bencher, Some(1..127)); + } + + fn bench_view(bencher: Bencher, range: Option>) { + let args = Args { + rows: 1_024, + dimensions: 128, + tile_rows: 64, + tile_dimensions: TileDimensions::Full, + }; + let canonical = canonical_u8(args).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = raw_tiled(args, &mut ctx).unwrap().into_array(); + let (expected, tiled) = match range { + Some(range) => ( + canonical.slice(range.clone()).unwrap(), + tiled.slice(range).unwrap(), + ), + None => (canonical, tiled), + }; + assert_arrays_eq!(expected, tiled, &mut ctx); + let tiled = tiled.as_::(); + + bencher.with_inputs(|| tiled).bench_values(|array| { + divan::black_box( + array + .tiles() + .map(|tile| tile.physical_range.len()) + .sum::(), + ) + }); + } +} + +#[divan::bench(args = slice_args())] +fn slice(bencher: Bencher, args: SliceArgs) { + let mut ctx = SESSION.create_execution_ctx(); + let tiled = raw_tiled(args.args, &mut ctx).unwrap(); + let range = slice_range(args); + bencher + .with_inputs(|| (tiled.clone(), range.clone())) + .bench_values(|(array, range)| divan::black_box(array.slice(range).unwrap())); +} + +#[divan::bench(args = take_args())] +fn take(bencher: Bencher, args: TakeArgs) { + let mut ctx = SESSION.create_execution_ctx(); + let tiled = raw_tiled(args.args, &mut ctx).unwrap(); + let indices = take_indices(args).into_array(); + bencher + .with_inputs(|| { + ( + tiled.clone().into_array(), + indices.clone(), + SESSION.create_execution_ctx(), + ) + }) + .bench_values(|(array, indices, mut ctx)| { + divan::black_box( + array + .take(indices) + .unwrap() + .execute_until::(&mut ctx) + .unwrap(), + ) + }); +} + +#[divan::bench(args = args())] +fn score_canonical(bencher: Bencher, args: Args) { + let canonical = canonical_u8(args); + let mut ctx = SESSION.create_execution_ctx(); + let values = canonical + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); + let query = query(args); + assert_adjacent_scores_distinct(values.as_slice::(), args, &query); + bencher.bench(|| { + divan::black_box(scoring::score_canonical( + values.as_slice::(), + args.rows, + args.dimensions, + &query, + )) + }); +} + +#[divan::bench(args = score_args())] +fn score_prepared(bencher: Bencher, args: ScoreArgs) { + let mut ctx = SESSION.create_execution_ctx(); + let tiled = tiled_score_fixture(args, &mut ctx).unwrap(); + let query = query(args.args); + let physical = tiled + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); + assert_scores_equal( + args.args, + &tiled, + physical.as_slice::(), + &query, + &mut ctx, + ) + .unwrap(); + + bencher.bench(|| { + divan::black_box(scoring::score_tiled( + tiled.as_view(), + physical.as_slice::(), + &query, + )) + }); +} + +#[divan::bench(args = score_args())] +fn score_end_to_end(bencher: Bencher, args: ScoreArgs) { + let mut ctx = SESSION.create_execution_ctx(); + let tiled = tiled_score_fixture(args, &mut ctx).unwrap(); + let query = query(args.args); + let physical = tiled + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); + assert_scores_equal( + args.args, + &tiled, + physical.as_slice::(), + &query, + &mut ctx, + ) + .unwrap(); + + bencher + .with_inputs(|| (tiled.elements().clone(), SESSION.create_execution_ctx())) + .bench_values(|(physical, mut ctx)| { + let physical = physical.execute::(&mut ctx).unwrap(); + divan::black_box(scoring::score_tiled( + tiled.as_view(), + physical.as_slice::(), + &query, + )) + }); +} + +mod scoring { + use vortex_array::ArrayView; + use vortex_tiled_fsl::TiledFixedSizeList; + use vortex_tiled_fsl::TiledFixedSizeListArrayExt; + + pub(super) fn score_canonical( + values: &[u8], + rows: usize, + dimensions: usize, + query: &[u8], + ) -> Vec { + let mut scores = vec![0; rows]; + for (row, score) in scores.iter_mut().enumerate() { + let row_values = &values[row * dimensions..(row + 1) * dimensions]; + *score = row_values + .iter() + .zip(query) + .map(|(&value, &weight)| u64::from(value) * u64::from(weight)) + .sum(); + } + scores + } + + pub(super) fn score_tiled( + array: ArrayView<'_, TiledFixedSizeList>, + values: &[u8], + query: &[u8], + ) -> Vec { + let mut scores = vec![0; array.len()]; + for bounds in array.tiles() { + let retained_rows = bounds.physical_range.len() / bounds.dimension_range.len(); + for (dimension_offset, dimension) in bounds.dimension_range.clone().enumerate() { + let physical_start = bounds.physical_range.start + + dimension_offset * retained_rows + + bounds.rows_within_tile.start; + let weight = u64::from(query[dimension]); + for (row_offset, row) in bounds.row_range.clone().enumerate() { + scores[row] += u64::from(values[physical_start + row_offset]) * weight; + } + } + } + scores + } +} diff --git a/encodings/tiled-fsl/goldenfiles/tiled_fsl.metadata b/encodings/tiled-fsl/goldenfiles/tiled_fsl.metadata new file mode 100644 index 00000000000..1a4b648b083 --- /dev/null +++ b/encodings/tiled-fsl/goldenfiles/tiled_fsl.metadata @@ -0,0 +1 @@ + @ÿÿÿÿ ÿÿÿÿÿÿÿÿÿ \ No newline at end of file diff --git a/encodings/tiled-fsl/src/array.rs b/encodings/tiled-fsl/src/array.rs new file mode 100644 index 00000000000..b00db1ca257 --- /dev/null +++ b/encodings/tiled-fsl/src/array.rs @@ -0,0 +1,610 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::num::NonZeroU32; +use std::sync::Arc; + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::IntoArray; +use vortex_array::TypedArrayRef; +use vortex_array::array_slots; +use vortex_array::arrays::FixedSizeList; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::require_child; +use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityVTable; +use vortex_array::vtable::child_to_validity; +use vortex_array::vtable::validity_to_child; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::TileGeometry; +use crate::geometry::TileBounds; +use crate::geometry::TileBoundsIter; +use crate::geometry::geometry_usizes; +use crate::transpose::decode_visible_elements; +use crate::transpose::encode_elements; + +/// A tiled fixed-size-list Vortex array. +pub type TiledFixedSizeListArray = Array; + +/// Wire-format metadata for [`TiledFixedSizeListArray`]. +#[derive(Clone, prost::Message)] +pub struct TiledFixedSizeListMetadata { + /// The nonzero row capacity of each physical tile. + #[prost(uint32, tag = "1")] + pub tile_rows: u32, + /// The nonzero dimension capacity of each physical tile. + #[prost(uint32, tag = "2")] + pub tile_dimensions: u32, + /// The logical row offset within the first retained physical tile. + #[prost(uint32, tag = "3")] + pub row_offset: u32, + /// The number of rows represented by the retained physical child. + #[prost(uint64, tag = "4")] + pub backing_rows: u64, +} + +#[array_slots(TiledFixedSizeList)] +/// Child slots owned by a tiled fixed-size-list array. +pub struct TiledFixedSizeListSlots { + /// The primitive physical elements in tiled order. + #[slot(0)] + pub elements: ArrayRef, + /// The optional outer-list validity bitmap. + #[slot(1)] + pub validity: Option, +} + +/// Encoding-specific state for [`TiledFixedSizeListArray`]. +#[derive(Clone, Debug)] +pub struct TiledFixedSizeListData { + geometry: TileGeometry, + row_offset: usize, + backing_rows: usize, +} + +impl TiledFixedSizeListData { + fn make_slots( + elements: &ArrayRef, + validity: &Validity, + len: usize, + ) -> vortex_array::ArraySlots { + TiledFixedSizeListSlots { + elements: elements.clone(), + validity: validity_to_child(validity, len), + } + .into_slots() + } +} + +impl Display for TiledFixedSizeListData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "tile_rows: {}, tile_dimensions: {}, row_offset: {}, backing_rows: {}", + self.geometry.rows(), + self.geometry.dimensions(), + self.row_offset, + self.backing_rows + ) + } +} + +impl ArrayHash for TiledFixedSizeListData { + fn array_hash(&self, state: &mut H, _accuracy: EqMode) { + self.geometry.hash(state); + self.row_offset.hash(state); + self.backing_rows.hash(state); + } +} + +impl ArrayEq for TiledFixedSizeListData { + fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { + self.geometry == other.geometry + && self.row_offset == other.row_offset + && self.backing_rows == other.backing_rows + } +} + +/// A two-dimensional tiled physical layout for primitive fixed-size-list values. +#[derive(Clone, Debug)] +pub struct TiledFixedSizeList; + +impl TiledFixedSizeList { + /// Encodes a canonical primitive fixed-size-list array into tiled physical order. + pub fn encode( + array: ArrayView<'_, FixedSizeList>, + geometry: TileGeometry, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let elements = array.elements().clone().execute::(ctx)?; + let tiled_elements = encode_elements( + elements.as_view(), + array.len(), + array.list_size() as usize, + geometry, + ctx, + )?; + Self::try_new( + tiled_elements.into_array(), + array.list_size(), + array.fixed_size_list_validity(), + array.len(), + geometry, + ) + } + + /// Constructs a tiled fixed-size-list array from primitive physical elements. + /// + /// The physical child must contain exactly `len * list_size` elements in tiled order. + pub fn try_new( + elements: ArrayRef, + list_size: u32, + validity: Validity, + len: usize, + geometry: TileGeometry, + ) -> VortexResult { + Self::try_new_view(elements, list_size, validity, len, geometry, 0, len) + } + + pub(crate) fn try_new_view( + elements: ArrayRef, + list_size: u32, + validity: Validity, + len: usize, + geometry: TileGeometry, + row_offset: usize, + backing_rows: usize, + ) -> VortexResult { + let dtype = DType::FixedSizeList( + Arc::new(elements.dtype().clone()), + list_size, + validity.nullability(), + ); + let data = TiledFixedSizeListData { + geometry, + row_offset, + backing_rows, + }; + let slots = TiledFixedSizeListData::make_slots(&elements, &validity, len); + Array::try_from_parts( + ArrayParts::new(TiledFixedSizeList, dtype, len, data).with_slots(slots), + ) + } +} + +impl TryFrom<&TiledFixedSizeListMetadata> for TileGeometry { + type Error = vortex_error::VortexError; + + fn try_from(metadata: &TiledFixedSizeListMetadata) -> VortexResult { + let rows = NonZeroU32::new(metadata.tile_rows) + .ok_or_else(|| vortex_err!(InvalidArgument: "tile_rows must be nonzero"))?; + let dimensions = NonZeroU32::new(metadata.tile_dimensions) + .ok_or_else(|| vortex_err!(InvalidArgument: "tile_dimensions must be nonzero"))?; + let geometry = Self::new(rows, dimensions); + geometry_usizes(geometry)?; + Ok(geometry) + } +} + +impl VTable for TiledFixedSizeList { + type TypedArrayData = TiledFixedSizeListData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.tiled_fsl"); + *ID + } + + fn validate( + &self, + data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let DType::FixedSizeList(element_dtype, list_size, nullability) = dtype else { + vortex_bail!(InvalidArgument: "tiled fixed-size-list dtype must be FixedSizeList, got {dtype}"); + }; + let (tile_rows, tile_dimensions) = geometry_usizes(data.geometry)?; + vortex_ensure!( + matches!(element_dtype.as_ref(), DType::Primitive(..)), + InvalidArgument: "tiled fixed-size-list elements must have a primitive dtype, got {element_dtype}" + ); + vortex_ensure!( + slots.len() == TiledFixedSizeListSlots::COUNT, + InvalidArgument: "tiled fixed-size-list expected {} slots, got {}", + TiledFixedSizeListSlots::COUNT, + slots.len() + ); + + let elements = slots.first().and_then(Option::as_ref).ok_or_else( + || vortex_err!(InvalidArgument: "tiled fixed-size-list elements slot is missing"), + )?; + let validity_child = slots + .get(TiledFixedSizeListSlots::VALIDITY) + .and_then(Option::as_ref); + + vortex_ensure!( + elements.dtype() == element_dtype.as_ref(), + InvalidArgument: "tiled fixed-size-list physical child dtype {} does not match logical element dtype {}", + elements.dtype(), + element_dtype + ); + + vortex_ensure!( + data.row_offset < tile_rows, + InvalidArgument: "tiled fixed-size-list row offset {} must be less than tile rows {tile_rows}", + data.row_offset + ); + let logical_end = data.row_offset.checked_add(len).ok_or_else(|| { + vortex_err!(InvalidArgument: "tiled fixed-size-list row offset {} plus length {len} overflows usize", data.row_offset) + })?; + vortex_ensure!( + logical_end <= data.backing_rows, + InvalidArgument: "tiled fixed-size-list row window {}..{logical_end} exceeds {} backing rows", + data.row_offset, + data.backing_rows + ); + if len > 0 { + let remainder = logical_end % tile_rows; + let max_backing_rows = if remainder == 0 { + logical_end + } else { + logical_end.saturating_add(tile_rows - remainder) + }; + vortex_ensure!( + data.backing_rows <= max_backing_rows, + InvalidArgument: "tiled fixed-size-list backing rows {} exceeds the retained tile extent {max_backing_rows}", + data.backing_rows + ); + } + if *list_size == 0 || (*list_size as usize) > tile_dimensions { + vortex_ensure!( + data.row_offset == 0 && data.backing_rows == len, + InvalidArgument: "tiled fixed-size-list zero-width and multi-slab arrays require row offset 0 and backing rows equal to length {len}" + ); + } + + let expected_len = data + .backing_rows + .checked_mul(*list_size as usize) + .ok_or_else(|| { + vortex_err!(InvalidArgument: "tiled fixed-size-list backing rows {} times list size {list_size} overflows usize", data.backing_rows) + })?; + vortex_ensure!( + elements.len() == expected_len, + InvalidArgument: "tiled fixed-size-list physical child length {} does not match expected {expected_len}", + elements.len() + ); + + if let Some(validity) = validity_child { + vortex_ensure!( + validity.dtype() == &Validity::DTYPE, + InvalidArgument: "tiled fixed-size-list outer validity must have dtype {}", + Validity::DTYPE + ); + vortex_ensure!( + validity.len() == len, + InvalidArgument: "tiled fixed-size-list outer validity length {} does not match {len}", + validity.len() + ); + } + + let validity = child_to_validity(validity_child, *nullability); + vortex_ensure!( + validity.nullability() == *nullability, + InvalidArgument: "tiled fixed-size-list outer validity does not match dtype nullability" + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("TiledFixedSizeListArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("TiledFixedSizeListArray buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + let row_offset = u32::try_from(array.row_offset()).map_err(|error| { + vortex_err!(InvalidArgument: "tiled fixed-size-list row offset cannot be serialized as u32: {error}") + })?; + let backing_rows = u64::try_from(array.backing_rows()).map_err(|error| { + vortex_err!(InvalidArgument: "tiled fixed-size-list backing row count cannot be serialized as u64: {error}") + })?; + Ok(Some( + TiledFixedSizeListMetadata { + tile_rows: array.geometry.rows().get(), + tile_dimensions: array.geometry.dimensions().get(), + row_offset, + backing_rows, + } + .encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + buffers.is_empty(), + InvalidArgument: "tiled fixed-size-list expects 0 buffers, got {}", + buffers.len() + ); + vortex_ensure!( + matches!(children.len(), 1 | 2), + InvalidArgument: "tiled fixed-size-list expects one elements child and an optional validity child, got {} children", + children.len() + ); + + let metadata = TiledFixedSizeListMetadata::decode(metadata)?; + let geometry = TileGeometry::try_from(&metadata)?; + let row_offset = usize::try_from(metadata.row_offset).map_err(|error| { + vortex_err!(InvalidArgument: "tiled fixed-size-list row offset cannot be represented as usize: {error}") + })?; + let backing_rows = usize::try_from(metadata.backing_rows).map_err(|error| { + vortex_err!(InvalidArgument: "tiled fixed-size-list backing row count cannot be represented as usize: {error}") + })?; + let DType::FixedSizeList(element_dtype, list_size, nullability) = dtype else { + vortex_bail!(InvalidArgument: "tiled fixed-size-list dtype must be FixedSizeList, got {dtype}"); + }; + vortex_ensure!( + matches!(element_dtype.as_ref(), DType::Primitive(..)), + InvalidArgument: "tiled fixed-size-list elements must have a primitive dtype, got {element_dtype}" + ); + vortex_ensure!( + nullability.is_nullable() || children.len() == 1, + InvalidArgument: "non-nullable tiled fixed-size-list dtype cannot have an outer validity child" + ); + let physical_len = backing_rows + .checked_mul(*list_size as usize) + .ok_or_else(|| { + vortex_err!(InvalidArgument: "tiled fixed-size-list backing rows {backing_rows} times list size {list_size} overflows usize") + })?; + let elements = children.get(0, element_dtype.as_ref(), physical_len)?; + let validity = match children.len() { + 1 => Validity::from(*nullability), + 2 => Validity::Array(children.get(1, &Validity::DTYPE, len)?), + _ => unreachable!("validated tiled fixed-size-list child count"), + }; + let array = Self::try_new_view( + elements, + *list_size, + validity, + len, + geometry, + row_offset, + backing_rows, + )?; + vortex_ensure!( + array.dtype() == dtype, + InvalidArgument: "deserialized tiled fixed-size-list dtype {} does not match supplied dtype {dtype}", + array.dtype() + ); + array.try_into_parts().map_err(|_| { + vortex_err!(InvalidArgument: "deserialized tiled fixed-size-list array is unexpectedly shared") + }) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + TiledFixedSizeListSlots::NAMES + .get(idx) + .map(ToString::to_string) + .unwrap_or_else(|| { + vortex_panic!("TiledFixedSizeListArray slot index {idx} out of bounds") + }) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!( + array, + array.elements(), + TiledFixedSizeListSlots::ELEMENTS => Primitive + ); + let elements = array + .elements() + .clone() + .try_downcast::() + .map_err(|elements| { + vortex_err!( + "tiled fixed-size-list physical child must execute to primitive, got {}", + elements.encoding_id() + ) + })?; + let decoded_elements = decode_visible_elements( + elements.as_view(), + array.len(), + array.list_size() as usize, + array.geometry(), + array.row_offset(), + array.backing_rows(), + ctx, + )?; + Ok(ExecutionResult::done( + FixedSizeListArray::new( + decoded_elements.into_array(), + array.list_size(), + array.array_validity(), + array.len(), + ) + .into_array(), + )) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + crate::rules::RULES.evaluate(array, parent, child_idx) + } +} + +impl ValidityVTable for TiledFixedSizeList { + fn validity(array: ArrayView<'_, TiledFixedSizeList>) -> VortexResult { + let validity = array + .slots() + .get(TiledFixedSizeListSlots::VALIDITY) + .and_then(Option::as_ref); + Ok(child_to_validity(validity, array.dtype().nullability())) + } +} + +/// Typed accessors for [`TiledFixedSizeListArray`]. +pub trait TiledFixedSizeListArrayExt: + TypedArrayRef + TiledFixedSizeListArraySlotsExt +{ + /// Returns the tile geometry that defines the physical layout. + fn geometry(&self) -> TileGeometry { + self.geometry + } + + /// Returns the logical row offset within the first retained physical tile. + fn row_offset(&self) -> usize { + self.row_offset + } + + /// Returns the number of rows represented by the retained physical child. + fn backing_rows(&self) -> usize { + self.backing_rows + } + + /// Returns whether every logical row is stored in a single dimension slab. + fn is_full_width(&self) -> bool { + let (_, dimensions) = validated_geometry_usizes(self.geometry()); + self.list_size() as usize <= dimensions + } + + /// Returns the number of elements in each logical fixed-size list. + fn list_size(&self) -> u32 { + match self.as_ref().dtype() { + DType::FixedSizeList(_, list_size, _) => *list_size, + _ => unreachable!("validated tiled fixed-size-list dtype"), + } + } + + /// Returns the number of row tiles needed for the logical array length. + fn row_tile_count(&self) -> usize { + let (rows, _) = validated_geometry_usizes(self.geometry()); + if self.as_ref().is_empty() { + 0 + } else { + (self.row_offset() + self.as_ref().len()).div_ceil(rows) + } + } + + /// Returns the number of dimension tiles needed for each logical list. + fn dimension_tile_count(&self) -> usize { + let (_, dimensions) = validated_geometry_usizes(self.geometry()); + (self.list_size() as usize).div_ceil(dimensions) + } + + /// Returns the bounds for one checked logical tile. + fn tile(&self, row_tile: usize, dimension_tile: usize) -> VortexResult { + crate::geometry::tile_bounds_view( + self.as_ref().len(), + self.list_size() as usize, + self.geometry(), + self.row_offset(), + self.backing_rows(), + row_tile, + dimension_tile, + ) + } + + /// Returns the logical and physical bounds of every tile in physical storage order. + fn tiles(&self) -> TileBoundsIter { + TileBoundsIter::new_view( + self.as_ref().len(), + self.list_size() as usize, + self.geometry(), + self.row_offset(), + self.backing_rows(), + self.row_tile_count(), + self.dimension_tile_count(), + ) + } + + /// Slices the physical child for inspection of one tile's elements. + /// + /// This is a cold-path convenience; scoring kernels should retain the child once and index it + /// with [`TileBounds::physical_range`]. + #[cold] + fn tile_elements(&self, bounds: &TileBounds) -> VortexResult { + self.elements().slice(bounds.physical_range.clone()) + } + + /// Returns the outer-list validity derived from the optional validity slot. + fn array_validity(&self) -> Validity { + let validity = self + .as_ref() + .slots() + .get(TiledFixedSizeListSlots::VALIDITY) + .and_then(Option::as_ref); + child_to_validity(validity, self.as_ref().dtype().nullability()) + } +} + +impl TiledFixedSizeListArrayExt for T where + T: TypedArrayRef + TiledFixedSizeListArraySlotsExt +{ +} + +fn validated_geometry_usizes(geometry: TileGeometry) -> (usize, usize) { + match geometry_usizes(geometry) { + Ok(geometry) => geometry, + Err(_) => unreachable!("validated tiled fixed-size-list geometry must fit usize"), + } +} diff --git a/encodings/tiled-fsl/src/gather.rs b/encodings/tiled-fsl/src/gather.rs new file mode 100644 index 00000000000..426f8000090 --- /dev/null +++ b/encodings/tiled-fsl/src/gather.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::TiledFixedSizeList; +use crate::TiledFixedSizeListArray; +use crate::TiledFixedSizeListArrayExt; +use crate::TiledFixedSizeListArraySlotsExt; +use crate::geometry::geometry_usizes; +use crate::geometry::physical_offset; + +/// Plans one contiguous physical child span per dimension slab for a row-tile-aligned range. +pub(crate) fn plan_physical_row_tile_spans( + array: ArrayView<'_, TiledFixedSizeList>, + range: Range, +) -> VortexResult>> { + let list_size = array.list_size() as usize; + let geometry = array.geometry(); + let (_, tile_dimensions) = geometry_usizes(geometry)?; + let dimension_slab_count = list_size.div_ceil(tile_dimensions); + let mut spans = Vec::with_capacity(dimension_slab_count); + + for dimension_start in (0..list_size).step_by(tile_dimensions) { + let dimension_width = list_size.min(dimension_start + tile_dimensions) - dimension_start; + let start = physical_offset( + array.len(), + list_size, + geometry, + range.start, + dimension_start, + )?; + let length = range.len().checked_mul(dimension_width).ok_or_else(|| { + vortex_err!( + InvalidArgument: + "row span {} times dimension width {dimension_width} overflows usize", + range.len() + ) + })?; + let end = start.checked_add(length).ok_or_else( + || vortex_err!(InvalidArgument: "physical row-tile span overflows usize"), + )?; + spans.push(start..end); + } + + Ok(spans) +} + +/// Builds one contiguous physical-index run for each dimension slab in a row-tile-aligned span. +pub(crate) fn gather_physical_row_tile_span( + array: ArrayView<'_, TiledFixedSizeList>, + range: Range, +) -> VortexResult { + let list_size = array.list_size() as usize; + let scalar_count = range.len().checked_mul(list_size).ok_or_else(|| { + vortex_err!( + InvalidArgument: + "row span {} times list size {list_size} overflows usize", + range.len() + ) + })?; + let spans = plan_physical_row_tile_spans(array, range)?; + let dimension_slab_count = spans.len(); + let mut starts = Vec::::with_capacity(dimension_slab_count); + let mut lengths = Vec::::with_capacity(dimension_slab_count); + + for span in spans { + starts.push(u64::try_from(span.start)?); + lengths.push(u64::try_from(span.len())?); + } + + Ok(PiecewiseSequenceArray::try_new( + PrimitiveArray::new(Buffer::from(starts), Validity::NonNullable).into_array(), + PrimitiveArray::new(Buffer::from(lengths), Validity::NonNullable).into_array(), + ConstantArray::new(1u64, dimension_slab_count).into_array(), + scalar_count, + )? + .into_array()) +} + +/// Gathers a row-tile-aligned logical row range while preserving the tiled encoding. +pub(crate) fn gather_tiled_slice( + array: ArrayView<'_, TiledFixedSizeList>, + range: Range, + validity: Validity, +) -> VortexResult { + let output_len = range.len(); + let geometry = array.geometry(); + let indices = gather_physical_row_tile_span(array, range)?; + let elements = array.elements().take(indices)?; + TiledFixedSizeList::try_new(elements, array.list_size(), validity, output_len, geometry) +} diff --git a/encodings/tiled-fsl/src/geometry.rs b/encodings/tiled-fsl/src/geometry.rs new file mode 100644 index 00000000000..cc5471df1e2 --- /dev/null +++ b/encodings/tiled-fsl/src/geometry.rs @@ -0,0 +1,450 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::num::NonZeroU32; +use std::ops::Range; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +/// The number of logical rows and dimensions in a physical tile. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub struct TileGeometry { + rows: NonZeroU32, + dimensions: NonZeroU32, +} + +impl TileGeometry { + /// Creates a tile geometry with nonzero row and dimension capacities. + pub const fn new(rows: NonZeroU32, dimensions: NonZeroU32) -> Self { + Self { rows, dimensions } + } + + /// Returns the number of rows in each tile. + pub const fn rows(self) -> NonZeroU32 { + self.rows + } + + /// Returns the number of dimensions in each tile. + pub const fn dimensions(self) -> NonZeroU32 { + self.dimensions + } +} + +/// The logical and physical extents of one tile. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct TileBounds { + /// The logical rows contained by this tile. + pub row_range: Range, + /// The logical dimensions contained by this tile. + pub dimension_range: Range, + /// The contiguous retained physical values for this tile. + /// + /// For a row view, this may include hidden boundary rows; [`Self::rows_within_tile`] identifies + /// the visible portion. + pub physical_range: Range, + /// The visible rows, relative to the beginning of the retained physical tile. + pub rows_within_tile: Range, + full_tile: bool, +} + +impl TileBounds { + pub(crate) fn new( + row_range: Range, + dimension_range: Range, + physical_range: Range, + rows_within_tile: Range, + full_tile: bool, + ) -> Self { + Self { + row_range, + dimension_range, + physical_range, + rows_within_tile, + full_tile, + } + } + + /// Returns whether the logical view selects every row in a full-height physical tile. + pub fn is_full_tile(&self) -> bool { + self.full_tile + } +} + +/// Iterates tiled fixed-size-list bounds in physical storage order. +/// +/// The iterator stores only scalar geometry and tile counters. Constructing or advancing it does +/// not access the array's physical child. +#[derive(Clone, Debug)] +pub struct TileBoundsIter { + len: usize, + list_size: usize, + rows: usize, + dimensions: usize, + row_offset: usize, + backing_rows: usize, + row_tile_count: usize, + dimension_tile_count: usize, + next_row_tile: usize, + next_dimension_tile: usize, +} + +#[derive(Clone, Copy, Debug)] +struct TileLayout { + len: usize, + list_size: usize, + rows: usize, + dimensions: usize, + row_offset: usize, + backing_rows: usize, +} + +impl TileBoundsIter { + pub(crate) fn new( + len: usize, + list_size: usize, + geometry: TileGeometry, + row_tile_count: usize, + dimension_tile_count: usize, + ) -> Self { + Self::new_view( + len, + list_size, + geometry, + 0, + len, + row_tile_count, + dimension_tile_count, + ) + } + + pub(crate) fn new_view( + len: usize, + list_size: usize, + geometry: TileGeometry, + row_offset: usize, + backing_rows: usize, + row_tile_count: usize, + dimension_tile_count: usize, + ) -> Self { + #[allow( + clippy::expect_used, + reason = "validated tiled fixed-size-list geometry must fit usize" + )] + let (rows, dimensions) = geometry_usizes(geometry) + .expect("validated tiled fixed-size-list geometry must fit usize"); + Self { + len, + list_size, + rows, + dimensions, + row_offset, + backing_rows, + row_tile_count, + dimension_tile_count, + next_row_tile: 0, + next_dimension_tile: 0, + } + } +} + +impl Iterator for TileBoundsIter { + type Item = TileBounds; + + fn next(&mut self) -> Option { + if self.row_tile_count == 0 || self.next_dimension_tile == self.dimension_tile_count { + return None; + } + + let bounds = tile_bounds_for_validated_array( + TileLayout { + len: self.len, + list_size: self.list_size, + rows: self.rows, + dimensions: self.dimensions, + row_offset: self.row_offset, + backing_rows: self.backing_rows, + }, + self.next_row_tile, + self.next_dimension_tile, + ); + self.next_row_tile += 1; + if self.next_row_tile == self.row_tile_count { + self.next_row_tile = 0; + self.next_dimension_tile += 1; + } + Some(bounds) + } +} + +pub(crate) fn tile_bounds( + len: usize, + list_size: usize, + geometry: TileGeometry, + row_tile: usize, + dimension_tile: usize, +) -> VortexResult { + tile_bounds_view(len, list_size, geometry, 0, len, row_tile, dimension_tile) +} + +pub(crate) fn tile_bounds_view( + len: usize, + list_size: usize, + geometry: TileGeometry, + row_offset: usize, + backing_rows: usize, + row_tile: usize, + dimension_tile: usize, +) -> VortexResult { + if len == 0 || list_size == 0 { + vortex_bail!(InvalidArgument: "cannot compute tiles for an empty logical extent"); + } + + let (rows, dimensions) = geometry_usizes(geometry)?; + let logical_end = row_offset.checked_add(len).ok_or_else( + || vortex_err!(InvalidArgument: "row offset plus length overflows logical extent"), + )?; + if row_offset >= rows || logical_end > backing_rows { + vortex_bail!(InvalidArgument: "invalid tiled fixed-size-list row window"); + } + let row_start = row_tile.checked_mul(rows).ok_or_else( + || vortex_err!(InvalidArgument: "row tile index {row_tile} overflows tile geometry"), + )?; + let dimension_start = dimension_tile.checked_mul(dimensions).ok_or_else(|| { + vortex_err!(InvalidArgument: "dimension tile index {dimension_tile} overflows tile geometry") + })?; + + if row_start >= logical_end + || row_start + .checked_add(rows) + .is_none_or(|end| end <= row_offset) + || dimension_start >= list_size + { + vortex_bail!( + InvalidArgument: + "tile ({row_tile}, {dimension_tile}) is outside logical extent ({len}, {list_size})" + ); + } + + tile_bounds_from_starts( + TileLayout { + len, + list_size, + rows, + dimensions, + row_offset, + backing_rows, + }, + row_start, + dimension_start, + ) +} + +fn tile_bounds_for_validated_array( + layout: TileLayout, + row_tile: usize, + dimension_tile: usize, +) -> TileBounds { + let TileLayout { + len, + list_size, + rows, + dimensions, + row_offset, + .. + } = layout; + let row_tile_count = if len == 0 { + 0 + } else { + (row_offset + len).div_ceil(rows) + }; + let dimension_tile_count = list_size.div_ceil(dimensions); + // Callers must provide only counters generated from this validated array's geometry. + debug_assert!(row_tile < row_tile_count); + debug_assert!(dimension_tile < dimension_tile_count); + debug_assert!(len.checked_mul(list_size).is_some()); + + let row_start = row_tile * rows; + let dimension_start = dimension_tile * dimensions; + match tile_bounds_from_starts(layout, row_start, dimension_start) { + Ok(bounds) => bounds, + Err(_) => unreachable!("validated tiled array has in-range tile bounds"), + } +} + +pub(crate) fn geometry_usizes(geometry: TileGeometry) -> VortexResult<(usize, usize)> { + let rows = usize::try_from(geometry.rows().get()).map_err(|_| { + vortex_err!( + InvalidArgument: "tile row geometry {} does not fit usize", + geometry.rows().get() + ) + })?; + let dimensions = usize::try_from(geometry.dimensions().get()).map_err(|_| { + vortex_err!( + InvalidArgument: "tile dimension geometry {} does not fit usize", + geometry.dimensions().get() + ) + })?; + Ok((rows, dimensions)) +} + +fn tile_bounds_from_starts( + layout: TileLayout, + row_start: usize, + dimension_start: usize, +) -> VortexResult { + let TileLayout { + len, + list_size, + rows, + dimensions, + row_offset, + backing_rows, + } = layout; + let retained_row_end = row_start + .checked_add(rows) + .ok_or_else(|| vortex_err!(InvalidArgument: "row tile range overflows logical extent"))? + .min(backing_rows); + let logical_end = row_offset + .checked_add(len) + .ok_or_else(|| vortex_err!(InvalidArgument: "row window overflows logical extent"))?; + let visible_row_start = row_start.max(row_offset); + let visible_row_end = retained_row_end.min(logical_end); + let dimension_end = dimension_start + .checked_add(dimensions) + .ok_or_else( + || vortex_err!(InvalidArgument: "dimension tile range overflows logical extent"), + )? + .min(list_size); + let retained_row_height = retained_row_end - row_start; + let dimension_width = dimension_end - dimension_start; + + let physical_start = dimension_start + .checked_mul(backing_rows) + .and_then(|offset| { + row_start + .checked_mul(dimension_width) + .and_then(|rows| offset.checked_add(rows)) + }) + .ok_or_else(|| vortex_err!(InvalidArgument: "tile physical offset overflows usize"))?; + let physical_len = retained_row_height + .checked_mul(dimension_width) + .ok_or_else(|| vortex_err!(InvalidArgument: "tile physical length overflows usize"))?; + let physical_end = physical_start + .checked_add(physical_len) + .ok_or_else(|| vortex_err!(InvalidArgument: "tile physical range overflows usize"))?; + + let rows_within_tile = (visible_row_start - row_start)..(visible_row_end - row_start); + let full_tile = rows_within_tile.start == 0 && rows_within_tile.end == rows; + Ok(TileBounds::new( + (visible_row_start - row_offset)..(visible_row_end - row_offset), + dimension_start..dimension_end, + physical_start..physical_end, + rows_within_tile, + full_tile, + )) +} + +pub(crate) fn physical_offset( + len: usize, + list_size: usize, + geometry: TileGeometry, + row: usize, + dimension: usize, +) -> VortexResult { + physical_offset_view(len, list_size, geometry, 0, len, row, dimension) +} + +pub(crate) fn physical_offset_view( + len: usize, + list_size: usize, + geometry: TileGeometry, + row_offset: usize, + backing_rows: usize, + row: usize, + dimension: usize, +) -> VortexResult { + if row >= len || dimension >= list_size { + vortex_bail!( + InvalidArgument: + "logical position ({row}, {dimension}) is outside extent ({len}, {list_size})" + ); + } + + let (rows, dimensions) = geometry_usizes(geometry)?; + let physical_row = row_offset + .checked_add(row) + .ok_or_else(|| vortex_err!(InvalidArgument: "physical row overflows usize"))?; + let row_tile = physical_row / rows; + let dimension_tile = dimension / dimensions; + let bounds = tile_bounds_view( + len, + list_size, + geometry, + row_offset, + backing_rows, + row_tile, + dimension_tile, + )?; + let row_within_tile = physical_row - row_tile * rows; + let dimension_within_tile = dimension - bounds.dimension_range.start; + let dimension_width = bounds.dimension_range.len(); + let row_height = bounds.physical_range.len() / dimension_width; + bounds + .physical_range + .start + .checked_add( + dimension_within_tile + .checked_mul(row_height) + .ok_or_else(|| vortex_err!(InvalidArgument: "dimension offset overflows usize"))?, + ) + .and_then(|offset| offset.checked_add(row_within_tile)) + .ok_or_else(|| vortex_err!(InvalidArgument: "physical offset overflows usize")) +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU32; + + use vortex_error::VortexResult; + + use super::TileGeometry; + use super::physical_offset; + use super::tile_bounds; + + fn geometry() -> TileGeometry { + TileGeometry::new(NonZeroU32::new(2).unwrap(), NonZeroU32::new(3).unwrap()) + } + + #[test] + fn physical_offsets_match_golden_layout() -> VortexResult<()> { + let expected = [[0, 2, 4, 9, 11], [1, 3, 5, 10, 12], [6, 7, 8, 13, 14]]; + for (row, offsets) in expected.into_iter().enumerate() { + for (dimension, expected_offset) in offsets.into_iter().enumerate() { + assert_eq!( + physical_offset(3, 5, geometry(), row, dimension)?, + expected_offset + ); + } + } + Ok(()) + } + + #[test] + fn tile_bounds_cover_unpadded_tails() -> VortexResult<()> { + let bounds = [ + tile_bounds(3, 5, geometry(), 0, 0)?, + tile_bounds(3, 5, geometry(), 1, 0)?, + tile_bounds(3, 5, geometry(), 0, 1)?, + tile_bounds(3, 5, geometry(), 1, 1)?, + ]; + assert_eq!(bounds[0].physical_range, 0..6); + assert_eq!(bounds[1].physical_range, 6..9); + assert_eq!(bounds[2].physical_range, 9..13); + assert_eq!(bounds[3].physical_range, 13..15); + assert_eq!(bounds[3].row_range, 2..3); + assert_eq!(bounds[3].dimension_range, 3..5); + Ok(()) + } +} diff --git a/encodings/tiled-fsl/src/kernel.rs b/encodings/tiled-fsl/src/kernel.rs new file mode 100644 index 00000000000..a5f356148c2 --- /dev/null +++ b/encodings/tiled-fsl/src/kernel.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Slice; +use vortex_array::arrays::slice::SliceExecuteAdaptor; +use vortex_array::arrays::slice::SliceKernel; +use vortex_array::builders::builder_with_capacity; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_session::VortexSession; + +use crate::TiledFixedSizeList; +use crate::TiledFixedSizeListArrayExt; +use crate::TiledFixedSizeListArraySlotsExt; +use crate::gather::plan_physical_row_tile_spans; +use crate::geometry::geometry_usizes; +use crate::transpose::decode_visible_elements; + +pub(crate) fn initialize(session: &VortexSession) { + session.kernels().register_execute_parent_kernel( + Slice.id(), + TiledFixedSizeList, + SliceExecuteAdaptor(TiledFixedSizeList), + ); +} + +impl SliceKernel for TiledFixedSizeList { + fn slice( + array: ArrayView<'_, Self>, + range: Range, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // Span planning below uses geometry relative to a complete, zero-offset backing array. + vortex_ensure!( + array.row_offset() == 0 && array.backing_rows() == array.len(), + InvalidArgument: "tiled fixed-size-list slice kernel requires a non-view array" + ); + let (tile_rows, _) = geometry_usizes(array.geometry())?; + let retained_start = range.start / tile_rows * tile_rows; + let retained_end = range.end.div_ceil(tile_rows) * tile_rows; + let retained_end = retained_end.min(array.len()); + let retained_range = retained_start..retained_end; + let retained_rows = retained_range.len(); + let retained_element_count = retained_rows * array.list_size() as usize; + let mut elements = builder_with_capacity(array.elements().dtype(), retained_element_count); + for span in plan_physical_row_tile_spans(array, retained_range)? { + let span_elements = array + .elements() + .slice(span)? + .execute::(ctx)?; + span_elements.append_to_builder(elements.as_mut(), ctx)?; + } + let elements = elements.finish(); + let decoded = decode_visible_elements( + elements.as_::(), + range.len(), + array.list_size() as usize, + array.geometry(), + range.start - retained_start, + retained_rows, + ctx, + )?; + + Ok(Some( + FixedSizeListArray::new( + decoded.into_array(), + array.list_size(), + array.array_validity().slice(range.clone())?, + range.len(), + ) + .into_array(), + )) + } +} diff --git a/encodings/tiled-fsl/src/lib.rs b/encodings/tiled-fsl/src/lib.rs new file mode 100644 index 00000000000..674886fa371 --- /dev/null +++ b/encodings/tiled-fsl/src/lib.rs @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Two-dimensional tiled encoding for primitive Vortex fixed-size lists. + +mod array; +mod gather; +mod geometry; +mod kernel; +mod mask; +mod operations; +mod rules; +mod slice; +mod transpose; + +pub use array::*; +pub use geometry::TileBounds; +pub use geometry::TileBoundsIter; +pub use geometry::TileGeometry; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Registers the tiled fixed-size-list array encoding in `session`. +pub fn initialize(session: &VortexSession) { + session.arrays().register(TiledFixedSizeList); + kernel::initialize(session); +} + +#[cfg(test)] +mod tests; diff --git a/encodings/tiled-fsl/src/mask.rs b/encodings/tiled-fsl/src/mask.rs new file mode 100644 index 00000000000..517e6ea36c5 --- /dev/null +++ b/encodings/tiled-fsl/src/mask.rs @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::scalar_fn::fns::mask::MaskReduce; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::TiledFixedSizeList; +use crate::TiledFixedSizeListArrayExt; +use crate::TiledFixedSizeListArraySlotsExt; + +impl MaskReduce for TiledFixedSizeList { + fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { + Ok(Some( + TiledFixedSizeList::try_new_view( + array.elements().clone(), + array.list_size(), + array.array_validity().and(Validity::Array(mask.clone()))?, + array.len(), + array.geometry(), + array.row_offset(), + array.backing_rows(), + )? + .into_array(), + )) + } +} diff --git a/encodings/tiled-fsl/src/operations.rs b/encodings/tiled-fsl/src/operations.rs new file mode 100644 index 00000000000..dcd83ca8a5c --- /dev/null +++ b/encodings/tiled-fsl/src/operations.rs @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::match_each_native_ptype; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::vtable::OperationsVTable; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::TiledFixedSizeList; +use crate::TiledFixedSizeListArrayExt; +use crate::TiledFixedSizeListArraySlotsExt; +use crate::geometry::physical_offset_view; + +impl OperationsVTable for TiledFixedSizeList { + fn scalar_at( + array: ArrayView<'_, TiledFixedSizeList>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let list_size = array.list_size() as usize; + let indices = (0..list_size) + .map(|dimension| { + let offset = physical_offset_view( + array.len(), + list_size, + array.geometry(), + array.row_offset(), + array.backing_rows(), + index, + dimension, + )?; + Ok(u64::try_from(offset)?) + }) + .collect::>>()?; + let row = array + .elements() + .take(PrimitiveArray::from_iter(indices).into_array())? + .execute::(ctx)?; + let mask = row.validity()?.execute_mask(row.len(), ctx)?; + + let tuple_values = match_each_native_ptype!(row.ptype(), |T| { + let values = row.as_slice::(); + match mask { + Mask::AllTrue(_) => values + .iter() + .copied() + .map(|value| Some(ScalarValue::Primitive(value.into()))) + .collect(), + Mask::AllFalse(_) => vec![None; values.len()], + Mask::Values(validity) => values + .iter() + .copied() + .zip(validity.bit_buffer().iter()) + .map(|(value, is_valid)| is_valid.then(|| ScalarValue::Primitive(value.into()))) + .collect(), + } + }); + + Scalar::try_new( + array.dtype().clone(), + Some(ScalarValue::Tuple(tuple_values)), + ) + } +} diff --git a/encodings/tiled-fsl/src/rules.rs b/encodings/tiled-fsl/src/rules.rs new file mode 100644 index 00000000000..8757884f5ee --- /dev/null +++ b/encodings/tiled-fsl/src/rules.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; + +use crate::TiledFixedSizeList; + +pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&MaskReduceAdaptor(TiledFixedSizeList)), + ParentRuleSet::lift(&SliceReduceAdaptor(TiledFixedSizeList)), +]); diff --git a/encodings/tiled-fsl/src/slice.rs b/encodings/tiled-fsl/src/slice.rs new file mode 100644 index 00000000000..7f9f7c59672 --- /dev/null +++ b/encodings/tiled-fsl/src/slice.rs @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::slice::SliceReduce; +use vortex_error::VortexResult; + +use crate::TiledFixedSizeList; +use crate::TiledFixedSizeListArrayExt; +use crate::TiledFixedSizeListArraySlotsExt; +use crate::gather::gather_tiled_slice; +use crate::geometry::geometry_usizes; + +impl SliceReduce for TiledFixedSizeList { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + if array.list_size() == 0 { + let validity = array.array_validity().slice(range.clone())?; + return Ok(Some( + TiledFixedSizeList::try_new_view( + array.elements().clone(), + array.list_size(), + validity, + range.len(), + array.geometry(), + 0, + range.len(), + )? + .into_array(), + )); + } + + if array.is_full_width() { + let validity = array.array_validity().slice(range.clone())?; + let (tile_rows, _) = geometry_usizes(array.geometry())?; + let list_size = array.list_size() as usize; + let absolute_start = array.row_offset() + range.start; + let absolute_end = array.row_offset() + range.end; + let retained_start = (absolute_start / tile_rows) * tile_rows; + let retained_end = + (absolute_end.div_ceil(tile_rows) * tile_rows).min(array.backing_rows()); + let physical = retained_start * list_size..retained_end * list_size; + let elements = array.elements().slice(physical)?; + return Ok(Some( + TiledFixedSizeList::try_new_view( + elements, + array.list_size(), + validity, + range.len(), + array.geometry(), + absolute_start - retained_start, + retained_end - retained_start, + )? + .into_array(), + )); + } + + let (tile_rows, _) = geometry_usizes(array.geometry())?; + if !range.start.is_multiple_of(tile_rows) + || (!range.end.is_multiple_of(tile_rows) && range.end != array.len()) + { + return Ok(None); + } + + let validity = array.array_validity().slice(range.clone())?; + Ok(Some( + gather_tiled_slice(array, range, validity)?.into_array(), + )) + } +} diff --git a/encodings/tiled-fsl/src/tests.rs b/encodings/tiled-fsl/src/tests.rs new file mode 100644 index 00000000000..8810326a4c6 --- /dev/null +++ b/encodings/tiled-fsl/src/tests.rs @@ -0,0 +1,1684 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::num::NonZeroU32; +use std::ops::Range; +use std::panic::AssertUnwindSafe; +use std::sync::Arc; +use std::sync::LazyLock; + +use prost::Message; +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Dict; +use vortex_array::arrays::FixedSizeList; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PiecewiseSequence; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Slice; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; +use vortex_array::arrays::slice::SliceKernel; +use vortex_array::arrays::slice::SliceReduce; +use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::compute::conformance::consistency::test_array_consistency; +use vortex_array::compute::conformance::filter::test_filter_conformance; +use vortex_array::compute::conformance::take::test_take_conformance; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::match_each_native_ptype; +use vortex_array::optimizer::ArrayOptimizer; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::serde::ArrayChildren; +use vortex_array::serde::SerializeOptions; +use vortex_array::serde::SerializedArray; +use vortex_array::test_harness::check_metadata; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::bitpack_compress::bitpack_encode; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use crate::TileBounds; +use crate::TileGeometry; +use crate::TiledFixedSizeList; +use crate::TiledFixedSizeListArray; +use crate::TiledFixedSizeListArrayExt; +use crate::TiledFixedSizeListArraySlotsExt; +use crate::TiledFixedSizeListMetadata; +use crate::transpose::decode_elements; +use crate::transpose::encode_elements; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + +fn geometry(rows: u32, dimensions: u32) -> TileGeometry { + TileGeometry::new( + NonZeroU32::new(rows).unwrap(), + NonZeroU32::new(dimensions).unwrap(), + ) +} + +#[test] +fn bitmap_validity_follows_value_permutation() -> VortexResult<()> { + let canonical_validity = Validity::from_iter([true, false, true, true]); + let canonical = FixedSizeListArray::new( + PrimitiveArray::new(buffer![10u16, 11, 20, 21], canonical_validity.clone()).into_array(), + 2, + Validity::NonNullable, + 2, + ); + let mut ctx = SESSION.create_execution_ctx(); + + let tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(2, 2), &mut ctx)?; + + let tiled_elements = tiled + .elements() + .clone() + .execute::(&mut ctx)?; + assert_eq!(tiled_elements.as_slice::(), &[10, 20, 11, 21]); + assert!(tiled.elements().validity()?.mask_eq( + &Validity::from_iter([true, true, false, true]), + 4, + &mut ctx, + )?); + + let decoded = tiled.into_array().execute::(&mut ctx)?; + assert!( + decoded + .elements() + .validity()? + .mask_eq(&canonical_validity, 4, &mut ctx,)? + ); + assert_fsl_equivalent(&canonical.into_array(), &decoded.into_array(), &mut ctx)?; + Ok(()) +} + +#[test] +fn constant_validity_forms_round_trip() -> VortexResult<()> { + for validity in [ + Validity::NonNullable, + Validity::AllValid, + Validity::AllInvalid, + ] { + let canonical = PrimitiveArray::new(buffer![10u16, 11, 20, 21], validity.clone()); + let mut ctx = SESSION.create_execution_ctx(); + + let tiled = encode_elements(canonical.as_view(), 2, 2, geometry(2, 2), &mut ctx)?; + assert!(same_validity_form(tiled.validity()?, &validity)); + + let decoded = decode_elements(tiled.as_view(), 2, 2, geometry(2, 2), &mut ctx)?; + assert_eq!(decoded.as_slice::(), canonical.as_slice::()); + assert!(same_validity_form( + decoded.validity()?, + &canonical.validity()? + )); + } + Ok(()) +} + +fn same_validity_form(actual: Validity, expected: &Validity) -> bool { + matches!( + (actual, expected), + (Validity::NonNullable, Validity::NonNullable) + | (Validity::AllValid, Validity::AllValid) + | (Validity::AllInvalid, Validity::AllInvalid) + ) +} + +#[test] +fn encode_elements_rejects_mismatched_extent() { + let elements = PrimitiveArray::from_iter([10u16, 11, 20]); + let mut ctx = SESSION.create_execution_ctx(); + + let error = encode_elements(elements.as_view(), 2, 2, geometry(2, 2), &mut ctx).unwrap_err(); + + assert!( + error + .to_string() + .contains("physical child length 3 does not match logical extent (2, 2)") + ); +} + +fn physical_fixture( + rows: usize, + dimensions: u32, + geometry: TileGeometry, +) -> VortexResult { + let _ = &*SESSION; + TiledFixedSizeList::try_new( + PrimitiveArray::from_iter( + (0..rows * dimensions as usize).map(|index| u16::try_from(index).unwrap_or(u16::MAX)), + ) + .into_array(), + dimensions, + Validity::NonNullable, + rows, + geometry, + ) +} + +fn offset_view_fixture() -> VortexResult { + offset_view_fixture_with_backing_rows(192) +} + +#[test] +fn mask_preserves_row_view_metadata_and_encoding() -> VortexResult<()> { + let view = offset_view_fixture()?; + let mask = Validity::from_iter((0..view.len()).map(|row| row % 5 != 2)).to_array(view.len()); + let masked = view.clone().into_array().mask(mask)?.optimize()?; + + let masked = masked.as_::(); + assert_eq!(masked.row_offset(), view.row_offset()); + assert_eq!(masked.backing_rows(), view.backing_rows()); + assert_eq!(masked.elements().len(), view.elements().len()); + + let mut ctx = SESSION.create_execution_ctx(); + assert!(masked.array_validity().mask_eq( + &Validity::from_iter((0..view.len()).map(|row| row % 5 != 2)), + view.len(), + &mut ctx, + )?); + Ok(()) +} + +fn offset_view_fixture_with_backing_rows( + backing_rows: usize, +) -> VortexResult { + let dimensions = 8u32; + let element_count = backing_rows * dimensions as usize; + let canonical = FixedSizeListArray::new( + PrimitiveArray::new( + Buffer::from_iter((0..element_count).map(|index| u16::try_from(index).unwrap())), + Validity::from_iter((0..element_count).map(|index| index % 11 != 0)), + ) + .into_array(), + dimensions, + Validity::AllValid, + backing_rows, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(64, 8), &mut ctx)?; + let view = TiledFixedSizeList::try_new_view( + tiled.elements().clone(), + dimensions, + canonical.fixed_size_list_validity().slice(10..138)?, + 128, + geometry(64, 8), + 10, + backing_rows, + )?; + Ok(view) +} + +fn fixture( + rows: usize, + dimensions: u32, + geometry: TileGeometry, +) -> VortexResult<(FixedSizeListArray, TiledFixedSizeListArray, ExecutionCtx)> { + let canonical = FixedSizeListArray::new( + PrimitiveArray::from_iter( + (0..rows * dimensions as usize) + .map(|index| u8::try_from(index % 16).unwrap_or_default()), + ) + .into_array(), + dimensions, + Validity::NonNullable, + rows, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry, &mut ctx)?; + Ok((canonical, tiled, ctx)) +} + +fn assert_fsl_equivalent( + canonical: &ArrayRef, + candidate: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + assert_eq!(candidate.dtype(), canonical.dtype()); + assert_eq!(candidate.len(), canonical.len()); + assert_arrays_eq!(canonical, candidate, ctx); + Ok(()) +} + +fn assert_array_tree_validity(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { + for descendant in array.depth_first_traversal() { + descendant.validity()?.execute_mask(descendant.len(), ctx)?; + } + Ok(()) +} + +fn mixed_validity_fixture( + rows: usize, + dimensions: u32, + tile_geometry: TileGeometry, +) -> VortexResult<(FixedSizeListArray, TiledFixedSizeListArray, ExecutionCtx)> { + let element_count = rows + .checked_mul(usize::try_from(dimensions)?) + .ok_or_else(|| vortex_err!("mixed-validity fixture extent overflows usize"))?; + let element_count_u32 = u32::try_from(element_count)?; + let canonical = FixedSizeListArray::new( + PrimitiveArray::new( + Buffer::from_iter(0..element_count_u32), + Validity::from_iter((0..element_count).map(|index| index % 11 != 0)), + ) + .into_array(), + dimensions, + Validity::from_iter((0..rows).map(|row| row % 7 != 0)), + rows, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), tile_geometry, &mut ctx)?; + Ok((canonical, tiled, ctx)) +} + +#[rstest] +#[case::full_width(128, 128, 10, 150, 60, 70, true)] +#[case::aligned_multi_slab(129, 64, 64, 192, 1, 65, true)] +#[case::unaligned_multi_slab(129, 64, 1, 130, 1, 64, false)] +fn row_view_path_conformance( + #[case] dimensions: u32, + #[case] tile_dimensions: u32, + #[case] start: usize, + #[case] stop: usize, + #[case] nested_start: usize, + #[case] nested_stop: usize, + #[case] expect_tiled: bool, +) -> VortexResult<()> { + let tile_geometry = geometry(64, tile_dimensions); + let (canonical, tiled, mut ctx) = mixed_validity_fixture(200, dimensions, tile_geometry)?; + let expected = canonical.into_array().slice(start..stop)?; + let actual = tiled.into_array().slice(start..stop)?; + + assert_eq!(actual.is::(), expect_tiled); + assert_eq!(actual.is::(), !expect_tiled); + + let tile_boundary = 64 - start % 64; + let mut probes = vec![0, actual.len() - 1]; + if tile_boundary < actual.len() { + probes.extend([tile_boundary - 1, tile_boundary]); + } + for row in probes { + assert_eq!( + expected.execute_scalar(row, &mut ctx)?, + actual.execute_scalar(row, &mut ctx)?, + ); + } + + let nested_expected = expected.clone().slice(nested_start..nested_stop)?; + let nested_actual = actual.clone().slice(nested_start..nested_stop)?; + if dimensions == 128 { + assert!(nested_actual.is::()); + assert_eq!(start + nested_start..start + nested_stop, 70..80); + } else { + assert!(nested_actual.is::()); + } + assert_fsl_equivalent(&nested_expected, &nested_actual, &mut ctx)?; + + let indices = + PrimitiveArray::from_option_iter([Some(u64::try_from(actual.len() - 1)?), None, Some(0)]) + .into_array(); + assert_arrays_eq!( + expected.clone().take(indices.clone())?, + actual.clone().take(indices)?, + &mut ctx + ); + + assert!( + actual + .validity()? + .mask_eq(&expected.validity()?, actual.len(), &mut ctx,)? + ); + let expected_fsl = expected.execute::(&mut ctx)?; + let actual_fsl = actual.clone().execute::(&mut ctx)?; + assert!(actual_fsl.elements().validity()?.mask_eq( + &expected_fsl.elements().validity()?, + actual_fsl.elements().len(), + &mut ctx, + )?); + assert_array_tree_validity(&actual, &mut ctx)?; + assert_fsl_equivalent( + &expected_fsl.into_array(), + &actual_fsl.into_array(), + &mut ctx, + ) +} + +fn encode_fixture( + values: Buffer, + element_validity: Validity, + list_size: u32, + outer_validity: Validity, + len: usize, + geometry: TileGeometry, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let canonical = FixedSizeListArray::new( + PrimitiveArray::new(values, element_validity).into_array(), + list_size, + outer_validity, + len, + ); + Ok(TiledFixedSizeList::encode(canonical.as_view(), geometry, ctx)?.into_array()) +} + +fn expected_tile_bounds( + rows: usize, + dimensions: usize, + tile_geometry: TileGeometry, +) -> VortexResult> { + let tile_rows = usize::try_from(tile_geometry.rows().get())?; + let tile_dimensions = usize::try_from(tile_geometry.dimensions().get())?; + let mut physical_cursor = 0; + let mut bounds = Vec::new(); + + for dimension_start in (0..dimensions).step_by(tile_dimensions) { + let dimension_end = dimension_start + .saturating_add(tile_dimensions) + .min(dimensions); + for row_start in (0..rows).step_by(tile_rows) { + let row_end = row_start.saturating_add(tile_rows).min(rows); + let tile_len = (row_end - row_start) * (dimension_end - dimension_start); + bounds.push(TileBounds::new( + row_start..row_end, + dimension_start..dimension_end, + physical_cursor..physical_cursor + tile_len, + 0..row_end - row_start, + row_end - row_start == tile_rows, + )); + physical_cursor += tile_len; + } + } + + Ok(bounds) +} + +fn boundary_slice_ranges(rows: usize, tile_rows: usize) -> Vec> { + let mut ranges = vec![0..0, 0..rows]; + for boundary in (tile_rows..rows).step_by(tile_rows) { + ranges.extend([ + 0..boundary - 1, + 0..boundary, + 0..boundary + 1, + boundary - 1..boundary, + boundary..boundary + 1, + boundary - 1..boundary + 1, + ]); + } + ranges.sort_by_key(|range| (range.start, range.end)); + ranges.dedup(); + ranges +} + +fn conformance_take_indices(rows: usize) -> VortexResult> { + let row_count = u32::try_from(rows)?; + let mut cases = vec![PrimitiveArray::from_iter::<[u32; 0]>([]).into_array()]; + if rows == 0 { + cases.push(PrimitiveArray::from_option_iter([None::]).into_array()); + return Ok(cases); + } + + cases.extend([ + PrimitiveArray::from_iter(0..row_count).into_array(), + PrimitiveArray::from_iter((0..row_count).rev()).into_array(), + PrimitiveArray::from_iter([0, row_count - 1, 0]).into_array(), + PrimitiveArray::from_iter([row_count - 1, 0, row_count / 2]).into_array(), + PrimitiveArray::from_option_iter([Some(row_count - 1), None, Some(0)]).into_array(), + ]); + Ok(cases) +} + +fn assert_scalar_conformance( + canonical: &FixedSizeListArray, + tiled: &TiledFixedSizeListArray, + tile_geometry: TileGeometry, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + if canonical.is_empty() { + return Ok(()); + } + + let rows = canonical.len(); + let tile_rows = usize::try_from(tile_geometry.rows().get())?; + let mut scalar_rows = vec![0, rows - 1]; + for boundary in (tile_rows..rows).step_by(tile_rows) { + scalar_rows.extend([boundary - 1, boundary]); + } + scalar_rows.sort_unstable(); + scalar_rows.dedup(); + for row in scalar_rows { + assert_eq!( + canonical.execute_scalar(row, ctx)?, + tiled.execute_scalar(row, ctx)?, + ); + } + Ok(()) +} + +fn assert_tile_conformance( + canonical: &FixedSizeListArray, + tiled: &TiledFixedSizeListArray, + tile_geometry: TileGeometry, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let rows = canonical.len(); + let dimensions = usize::try_from(canonical.list_size())?; + let tile_rows = usize::try_from(tile_geometry.rows().get())?; + let tile_dimensions = usize::try_from(tile_geometry.dimensions().get())?; + let expected_row_tile_count = rows.div_ceil(tile_rows); + let expected_dimension_tile_count = dimensions.div_ceil(tile_dimensions); + assert_eq!(tiled.row_tile_count(), expected_row_tile_count); + assert_eq!(tiled.dimension_tile_count(), expected_dimension_tile_count); + + let expected_bounds = expected_tile_bounds(rows, dimensions, tile_geometry)?; + let actual_bounds: Vec = tiled.tiles().collect(); + assert_eq!(actual_bounds, expected_bounds); + for (dimension_tile, dimension_bounds) in (0..dimensions).step_by(tile_dimensions).enumerate() { + for (row_tile, row_bounds) in (0..rows).step_by(tile_rows).enumerate() { + let expected = &expected_bounds[dimension_tile * expected_row_tile_count + row_tile]; + assert_eq!(tiled.tile(row_tile, dimension_tile)?, *expected); + + let indices = expected + .dimension_range + .clone() + .flat_map(|dimension| { + expected + .row_range + .clone() + .map(move |row| row * dimensions + dimension) + }) + .map(u64::try_from) + .collect::, _>>()?; + let expected_elements = canonical + .elements() + .clone() + .take(PrimitiveArray::from_iter(indices).into_array())?; + let actual_elements = tiled.tile_elements(expected)?; + assert_arrays_eq!(expected_elements, actual_elements, ctx); + + assert_eq!(dimension_bounds, expected.dimension_range.start); + assert_eq!(row_bounds, expected.row_range.start); + } + } + Ok(()) +} + +fn assert_slice_conformance( + canonical: &FixedSizeListArray, + tiled: &TiledFixedSizeListArray, + tile_geometry: TileGeometry, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let tile_rows = usize::try_from(tile_geometry.rows().get())?; + for range in boundary_slice_ranges(canonical.len(), tile_rows) { + let expected = canonical.clone().into_array().slice(range.clone())?; + let actual = tiled.clone().into_array().slice(range.clone())?; + if canonical.is_empty() { + assert!(actual.is::()); + assert_eq!(actual.as_::().geometry(), tile_geometry); + } else if actual.is_empty() { + assert!(actual.is::()); + } else if !tiled.is_full_width() + && (range.start % tile_rows != 0 + || (range.end % tile_rows != 0 && range.end != tiled.len())) + { + assert!(actual.is::()); + } else { + assert!(actual.is::()); + assert_eq!(actual.as_::().geometry(), tile_geometry); + } + assert_fsl_equivalent(&expected, &actual, ctx)?; + } + Ok(()) +} + +fn assert_take_oracle_conformance( + canonical: &FixedSizeListArray, + tiled: &TiledFixedSizeListArray, + _tile_geometry: TileGeometry, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + for indices in conformance_take_indices(canonical.len())? { + let expected = canonical.clone().into_array().take(indices.clone())?; + let actual = tiled.clone().into_array().take(indices)?; + let actual = if !canonical.is_empty() && !actual.is_empty() { + actual.execute_until::(ctx)? + } else if actual.is_empty() { + let actual = actual.execute_until::(ctx)?; + assert!(actual.is::()); + actual + } else { + let actual = actual.execute_until::(ctx)?; + assert!(actual.is::()); + actual + }; + assert_fsl_equivalent(&expected, &actual, ctx)?; + } + Ok(()) +} + +fn assert_tiled_conformance_case( + canonical: &FixedSizeListArray, + tile_geometry: TileGeometry, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), tile_geometry, &mut ctx)?; + assert_eq!(tiled.geometry(), tile_geometry); + + let encoded = tiled + .clone() + .into_array() + .execute::(&mut ctx)? + .into_array(); + assert_fsl_equivalent(&canonical.clone().into_array(), &encoded, &mut ctx)?; + + let reconstructed = TiledFixedSizeList::try_new( + tiled.elements().clone(), + canonical.list_size(), + canonical.fixed_size_list_validity(), + canonical.len(), + tile_geometry, + )?; + assert_fsl_equivalent( + &canonical.clone().into_array(), + &reconstructed.into_array(), + &mut ctx, + )?; + + assert_scalar_conformance(canonical, &tiled, tile_geometry, &mut ctx)?; + assert_tile_conformance(canonical, &tiled, tile_geometry, &mut ctx)?; + assert_slice_conformance(canonical, &tiled, tile_geometry, &mut ctx)?; + assert_take_oracle_conformance(canonical, &tiled, tile_geometry, &mut ctx) +} + +#[test] +fn standard_harness_conformance() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let fixtures = vec![ + encode_fixture( + buffer![0u8, 1, 2, 3, 4, 5], + Validity::NonNullable, + 2, + Validity::NonNullable, + 3, + geometry(2, 2), + &mut ctx, + )?, + encode_fixture( + buffer![0i32, 1, 2, 3, 4, 5], + Validity::NonNullable, + 2, + Validity::from_iter([true, false, true]), + 3, + geometry(2, 2), + &mut ctx, + )?, + encode_fixture( + buffer![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0], + Validity::from_iter([true, false, true, true, false, true]), + 2, + Validity::NonNullable, + 3, + geometry(2, 2), + &mut ctx, + )?, + encode_fixture( + buffer![0.0f64, 1.0, 2.0, 3.0, 4.0, 5.0], + Validity::from_iter([true, false, true, true, false, true]), + 2, + Validity::from_iter([true, false, true]), + 3, + geometry(2, 2), + &mut ctx, + )?, + encode_fixture( + buffer![0u8; 0], + Validity::NonNullable, + 5, + Validity::NonNullable, + 0, + geometry(32, 64), + &mut ctx, + )?, + encode_fixture( + buffer![0u8; 0], + Validity::NonNullable, + 0, + Validity::NonNullable, + 65, + geometry(32, 64), + &mut ctx, + )?, + encode_fixture( + buffer![7u8; 65 * 129], + Validity::NonNullable, + 129, + Validity::NonNullable, + 65, + geometry(32, 64), + &mut ctx, + )?, + encode_fixture( + buffer![7u8; 3 * 5], + Validity::NonNullable, + 5, + Validity::NonNullable, + 3, + geometry(32, 64), + &mut ctx, + )?, + ]; + + for tiled in fixtures { + test_array_consistency(&tiled, &mut ctx); + test_filter_conformance(&tiled, &mut ctx); + test_take_conformance(&tiled, &mut ctx); + } + + let (_, raw_tiled, _) = fixture(65, 128, geometry(32, 64))?; + vortex_fastlanes::initialize(ctx.session()); + let physical = raw_tiled + .elements() + .clone() + .execute::(&mut ctx)?; + let bitpacked = bitpack_encode(&physical, 4, None, &mut ctx)?.into_array(); + let bitpacked_tiled = TiledFixedSizeList::try_new( + bitpacked, + 128, + raw_tiled.array_validity(), + 65, + geometry(32, 64), + )? + .into_array(); + test_array_consistency(&bitpacked_tiled, &mut ctx); + test_filter_conformance(&bitpacked_tiled, &mut ctx); + test_take_conformance(&bitpacked_tiled, &mut ctx); + Ok(()) +} + +#[test] +fn canonical_oracle_conformance_matrix() -> VortexResult<()> { + const ROW_COUNTS: &[usize] = &[0, 1, 15, 16, 31, 32, 33, 63, 64, 65]; + const DIMENSION_COUNTS: &[u32] = &[0, 1, 3, 4, 63, 64, 65, 129]; + + for &rows in ROW_COUNTS { + for &dimensions in DIMENSION_COUNTS { + let dimension_count = usize::try_from(dimensions)?; + let values = (0..rows * dimension_count) + .map(u16::try_from) + .collect::, _>>()?; + let canonical = FixedSizeListArray::new( + PrimitiveArray::new(Buffer::from(values), Validity::NonNullable).into_array(), + dimensions, + Validity::NonNullable, + rows, + ); + let full_width = dimensions.max(1); + for tile_geometry in [ + geometry(16, 4), + geometry(32, 64), + geometry(64, 64), + geometry(64, full_width), + ] { + assert_tiled_conformance_case(&canonical, tile_geometry)?; + } + } + } + Ok(()) +} + +#[test] +fn arbitrary_take_does_not_force_tiled_preservation() { + assert!( + !SESSION + .kernels() + .has_execute_parent(Dict.id(), TiledFixedSizeList.id()) + ); +} + +fn assert_take_indices(indices: ArrayRef) -> VortexResult<()> { + let index_count = indices.len(); + let (canonical, tiled, mut ctx) = fixture(3, 5, geometry(2, 3))?; + let expected = canonical.into_array().take(indices.clone())?; + let actual = tiled + .into_array() + .take(indices)? + .execute_until::(&mut ctx)?; + assert_eq!( + actual.as_::().elements().len(), + index_count * 5 + ); + assert_fsl_equivalent(&expected, &actual, &mut ctx) +} + +#[test] +fn take_accepts_all_integer_index_ptypes() -> VortexResult<()> { + assert_take_indices(PrimitiveArray::from_iter([2u8, 0, 1]).into_array())?; + assert_take_indices(PrimitiveArray::from_iter([2u16, 0, 1]).into_array())?; + assert_take_indices(PrimitiveArray::from_iter([2u32, 0, 1]).into_array())?; + assert_take_indices(PrimitiveArray::from_iter([2u64, 0, 1]).into_array())?; + assert_take_indices(PrimitiveArray::from_iter([2i8, 0, 1]).into_array())?; + assert_take_indices(PrimitiveArray::from_iter([2i16, 0, 1]).into_array())?; + assert_take_indices(PrimitiveArray::from_iter([2i32, 0, 1]).into_array())?; + assert_take_indices(PrimitiveArray::from_iter([2i64, 0, 1]).into_array())?; + Ok(()) +} + +#[test] +fn take_fallback_preserves_order_duplicates_and_validity() -> VortexResult<()> { + let canonical = FixedSizeListArray::new( + PrimitiveArray::new( + buffer![0i32, 1, 10, 11, 20, 21], + Validity::from_iter([true, false, true, true, false, true]), + ) + .into_array(), + 2, + Validity::from_iter([true, false, true]), + 3, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(2, 2), &mut ctx)?; + let indices = PrimitiveArray::from_option_iter([Some(2u32), Some(0), None, Some(2), Some(1)]) + .into_array(); + let expected = canonical.into_array().take(indices.clone())?; + let actual = tiled + .into_array() + .take(indices)? + .execute_until::(&mut ctx)?; + assert_eq!(actual.as_::().elements().len(), 10); + assert_fsl_equivalent(&expected, &actual, &mut ctx) +} + +#[test] +fn take_all_null_indices_from_empty_source() -> VortexResult<()> { + let (canonical, tiled, mut ctx) = fixture(0, 5, geometry(2, 3))?; + let indices = PrimitiveArray::from_option_iter([None::, None]).into_array(); + let expected = canonical.into_array().take(indices.clone())?; + let actual = tiled.into_array().take(indices)?; + assert_fsl_equivalent(&expected, &actual, &mut ctx) +} + +#[test] +fn take_empty_indices_is_canonical_empty() -> VortexResult<()> { + let (_, tiled, mut ctx) = fixture(3, 5, geometry(2, 3))?; + let indices = PrimitiveArray::from_iter::<[u32; 0]>([]).into_array(); + let actual = tiled + .into_array() + .take(indices)? + .execute_until::(&mut ctx)?; + assert!(actual.is::()); + assert!(actual.is_empty()); + Ok(()) +} + +#[test] +fn take_zero_width_lists() -> VortexResult<()> { + let (canonical, tiled, mut ctx) = fixture(3, 0, geometry(2, 3))?; + let indices = PrimitiveArray::from_iter([2u32, 0]).into_array(); + let expected = canonical.into_array().take(indices.clone())?; + let actual = tiled + .into_array() + .take(indices)? + .execute_until::(&mut ctx)?; + assert_eq!(actual.as_::().elements().len(), 0); + assert_fsl_equivalent(&expected, &actual, &mut ctx) +} + +#[test] +fn nullable_take_makes_outer_dtype_nullable() -> VortexResult<()> { + let (_, tiled, mut ctx) = fixture(3, 5, geometry(2, 3))?; + let indices = PrimitiveArray::from_option_iter([Some(2u32), None]).into_array(); + let actual = tiled + .into_array() + .take(indices)? + .execute_until::(&mut ctx)?; + assert!(actual.dtype().is_nullable()); + assert_eq!(actual.as_::().elements().len(), 10); + Ok(()) +} + +#[rstest] +#[case::empty(0..0)] +#[case::first_row(0..1)] +#[case::before_boundary(0..31)] +#[case::aligned(0..32)] +#[case::past_boundary(0..33)] +#[case::unaligned_interior(1..64)] +#[case::unaligned_through_tail(31..65)] +#[case::final_row(64..65)] +fn slice_chooses_tiled_only_for_aligned_multi_slab_ranges( + #[case] range: Range, +) -> VortexResult<()> { + let (canonical, tiled, mut ctx) = fixture(65, 129, geometry(32, 64))?; + let expected = canonical.into_array().slice(range.clone())?; + let actual = tiled.into_array().slice(range.clone())?; + if actual.is_empty() { + assert!(actual.is::()); + } else if range.start.is_multiple_of(32) && (range.end.is_multiple_of(32) || range.end == 65) { + assert!(actual.is::()); + assert_eq!( + actual.as_::().geometry(), + geometry(32, 64) + ); + } else { + assert!(actual.is::()); + } + assert_arrays_eq!(expected, actual, &mut ctx); + Ok(()) +} + +#[test] +fn large_unaligned_full_width_slice_retains_two_boundary_tiles() -> VortexResult<()> { + let rows = 1_000_000; + let dimensions = 128; + let tile_rows = 64; + let range = 123_457..124_458; + let dimensions_u32 = u32::try_from(dimensions)?; + let tile_rows_u32 = u32::try_from(tile_rows)?; + let tiled = TiledFixedSizeList::try_new( + ConstantArray::new(0u8, rows * dimensions).into_array(), + dimensions_u32, + Validity::NonNullable, + rows, + geometry(tile_rows_u32, dimensions_u32), + )?; + + let sliced = tiled.into_array().slice(range.clone())?; + let sliced = sliced.as_::(); + + assert_eq!(sliced.len(), range.len()); + assert_eq!(sliced.row_offset(), 1); + assert!(sliced.backing_rows() <= range.len() + 2 * tile_rows); + assert!(sliced.backing_rows() < rows); + Ok(()) +} + +#[test] +fn large_unaligned_multi_slab_slice_allocates_no_reduce_metadata() -> VortexResult<()> { + let million_rows = TiledFixedSizeList::try_new( + ConstantArray::new(0u8, 1_000_000 * 1_536).into_array(), + 1_536, + Validity::NonNullable, + 1_000_000, + geometry(64, 64), + )?; + + assert!( + ::slice(million_rows.as_view(), 1..130)?.is_none(), + "unaligned reduction must make an O(1) decision without scalar-run metadata" + ); + Ok(()) +} + +#[test] +fn multi_slab_slice_kernel_matches_canonical() -> VortexResult<()> { + let rows = 256; + let list_size = 1_536; + let element_count = rows * list_size; + let element_count_u32 = u32::try_from(element_count)?; + let list_size_u32 = u32::try_from(list_size)?; + let canonical = FixedSizeListArray::new( + PrimitiveArray::new( + Buffer::from_iter(0..element_count_u32), + Validity::from_iter((0..element_count).map(|index| index % 11 != 0)), + ) + .into_array(), + list_size_u32, + Validity::from_iter((0..rows).map(|row| row % 7 != 0)), + rows, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(64, 64), &mut ctx)?; + let expected = canonical.into_array().slice(1..130)?; + + let unaligned = tiled.into_array().slice(1..130)?; + assert!(unaligned.is::()); + let executed = unaligned.execute::(&mut ctx)?; + + assert_eq!(executed.len(), 129); + assert_eq!(executed.elements().len(), 129 * 1_536); + assert_arrays_eq!(expected, executed, &mut ctx); + Ok(()) +} + +#[test] +fn slice_kernel_rejects_row_views() -> VortexResult<()> { + let (_, tiled, mut ctx) = fixture(200, 128, geometry(64, 128))?; + let view = tiled.into_array().slice(10..150)?; + + let error = ::slice( + view.as_::(), + 1..2, + &mut ctx, + ) + .unwrap_err(); + + assert!(error.to_string().contains("requires a non-view array")); + Ok(()) +} + +#[test] +fn aligned_multi_slab_slice_uses_one_exact_run_per_slab() -> VortexResult<()> { + let (_, tiled, mut ctx) = fixture(256, 1_536, geometry(64, 64))?; + + let retained = tiled.into_array().slice(0..192)?; + assert!(retained.is::()); + let retained_tiled = retained.as_::(); + let elements = retained_tiled.elements(); + assert!(elements.is::()); + let dict = elements.as_::(); + assert!(dict.codes().is::()); + let codes = dict.codes().as_::(); + + assert_eq!(codes.len(), 192 * 1_536); + assert_arrays_eq!( + PrimitiveArray::from_iter((0..24).map(|slab| slab * 16_384u64)).into_array(), + codes.starts(), + &mut ctx + ); + assert_arrays_eq!( + PrimitiveArray::from_iter([12_288u64; 24]).into_array(), + codes.lengths(), + &mut ctx + ); + Ok(()) +} + +#[test] +fn physical_slab_span_plan_bounds_partial_dimension_and_row_tails() -> VortexResult<()> { + let (_, tiled, _) = fixture(130, 130, geometry(64, 64))?; + + let spans = crate::gather::plan_physical_row_tile_spans(tiled.as_view(), 64..130)?; + + assert_eq!(spans, vec![4_096..8_320, 12_416..16_640, 16_768..16_900]); + Ok(()) +} + +#[test] +fn small_full_width_slice_retains_oversized_backing_tile() -> VortexResult<()> { + let (canonical, tiled, mut ctx) = fixture(3, 5, geometry(32, 64))?; + let expected = canonical.into_array().slice(1..3)?; + let actual = tiled.into_array().slice(1..3)?; + + assert!(actual.is::()); + let sliced = actual.as_::(); + assert_eq!(sliced.row_offset(), 1); + assert_eq!(sliced.backing_rows(), 3); + assert_eq!(sliced.elements().len(), 15); + assert_fsl_equivalent(&expected, &actual, &mut ctx) +} + +#[test] +fn full_width_unaligned_slice_is_offset_view() -> VortexResult<()> { + let (canonical, tiled, mut ctx) = fixture(200, 128, geometry(64, 128))?; + let expected = canonical.into_array().slice(10..150)?; + let actual = tiled.into_array().slice(10..150)?; + let sliced = actual.as_::(); + + assert_eq!(sliced.row_offset(), 10); + assert_eq!(sliced.len(), 140); + assert_eq!(sliced.backing_rows(), 192); + assert_eq!(sliced.elements().len(), 192 * 128); + let tiles = sliced.tiles().collect::>(); + assert_eq!(tiles[0].row_range, 0..54); + assert_eq!(tiles[1].row_range, 54..118); + assert_eq!(tiles[2].row_range, 118..140); + assert!(!tiles[0].is_full_tile()); + assert!(tiles[1].is_full_tile()); + assert!(!tiles[2].is_full_tile()); + assert_fsl_equivalent(&expected, &actual, &mut ctx) +} + +#[test] +fn nested_full_width_slices_rebase_and_trim() -> VortexResult<()> { + let (canonical, tiled, mut ctx) = fixture(200, 128, geometry(64, 128))?; + let expected = canonical.into_array().slice(70..80)?; + let actual = tiled.into_array().slice(10..150)?.slice(60..70)?; + let sliced = actual.as_::(); + + assert_eq!(sliced.row_offset(), 6); + assert_eq!(sliced.len(), 10); + assert_eq!(sliced.backing_rows(), 64); + assert_eq!(sliced.elements().len(), 64 * 128); + assert_fsl_equivalent(&expected, &actual, &mut ctx) +} + +#[test] +fn full_width_slice_preserves_nullable_bitpacked_child() -> VortexResult<()> { + let canonical = FixedSizeListArray::new( + PrimitiveArray::new( + Buffer::from_iter( + (0..200 * 128).map(|index| u16::try_from(index % 16).unwrap_or_default()), + ), + Validity::from_iter((0..200 * 128).map(|index| index % 11 != 0)), + ) + .into_array(), + 128, + Validity::from_iter((0..200).map(|row| row % 7 != 0)), + 200, + ); + let mut ctx = SESSION.create_execution_ctx(); + let raw_tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(64, 128), &mut ctx)?; + vortex_fastlanes::initialize(ctx.session()); + let physical = raw_tiled + .elements() + .clone() + .execute::(&mut ctx)?; + let bitpacked = bitpack_encode(&physical, 4, None, &mut ctx)?.into_array(); + let tiled = TiledFixedSizeList::try_new( + bitpacked, + 128, + raw_tiled.array_validity(), + 200, + geometry(64, 128), + )?; + + let expected = canonical.into_array().slice(10..150)?; + let actual = tiled.into_array().slice(10..150)?; + let sliced = actual.as_::(); + + assert!(sliced.elements().is::()); + assert!(sliced.elements().validity()?.mask_eq( + &raw_tiled.elements().validity()?.slice(0..192 * 128)?, + 192 * 128, + &mut ctx, + )?); + assert_fsl_equivalent(&expected, &actual, &mut ctx) +} + +#[test] +fn zero_width_nonempty_slice_preserves_tiled() -> VortexResult<()> { + let (canonical, tiled, mut ctx) = fixture(3, 0, geometry(64, 128))?; + let expected = canonical.into_array().slice(1..3)?; + let actual = tiled.into_array().slice(1..3)?; + let sliced = actual.as_::(); + + assert_eq!(sliced.row_offset(), 0); + assert_eq!(sliced.backing_rows(), 2); + assert_eq!(sliced.elements().len(), 0); + assert_fsl_equivalent(&expected, &actual, &mut ctx) +} + +#[test] +fn bitpacked_child_roundtrips_through_row_ops() -> VortexResult<()> { + let canonical = FixedSizeListArray::new( + PrimitiveArray::from_iter((0..65).flat_map(|row| { + let row_value = u8::try_from(row / 32).unwrap(); + (0..128).map(move |dimension| row_value + u8::try_from(dimension % 14).unwrap()) + })) + .into_array(), + 128, + Validity::NonNullable, + 65, + ); + let mut ctx = SESSION.create_execution_ctx(); + let raw_tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(32, 64), &mut ctx)?; + let first_row = canonical.execute_scalar(0, &mut ctx)?; + let last_row = canonical.execute_scalar(64, &mut ctx)?; + assert_ne!(first_row, last_row); + let expected_sliced = canonical.clone().into_array().slice(1..64)?; + let indices = PrimitiveArray::from_iter([64u32, 0, 32]).into_array(); + let expected_taken = canonical.clone().into_array().take(indices.clone())?; + vortex_fastlanes::initialize(ctx.session()); + let physical = raw_tiled + .elements() + .clone() + .execute::(&mut ctx)?; + let bitpacked = bitpack_encode(&physical, 4, None, &mut ctx)?.into_array(); + let tiled = TiledFixedSizeList::try_new( + bitpacked, + 128, + raw_tiled.array_validity(), + 65, + geometry(32, 64), + )?; + + assert_arrays_eq!(canonical, tiled, &mut ctx); + + let sliced = tiled.clone().into_array().slice(1..64)?; + assert!(sliced.is::()); + assert_arrays_eq!(expected_sliced, sliced, &mut ctx); + + let taken = tiled + .into_array() + .take(indices)? + .execute_until::(&mut ctx)?; + assert_arrays_eq!(expected_taken, taken, &mut ctx); + Ok(()) +} + +#[test] +fn tiles_are_range_only_and_in_physical_order() -> VortexResult<()> { + let (_, tiled, _) = fixture(3, 5, geometry(2, 3))?; + let bounds: Vec = tiled.tiles().collect(); + assert_eq!( + bounds + .iter() + .map(|bounds| bounds.physical_range.clone()) + .collect::>(), + vec![0..6, 6..9, 9..13, 13..15], + ); + assert_eq!(tiled.tile(1, 1)?, bounds[3]); + assert_eq!(tiled.tile_elements(&bounds[3])?.len(), 2); + assert!(tiled.tile(2, 0).is_err()); + assert!(tiled.tile(0, 2).is_err()); + Ok(()) +} + +#[test] +fn offset_view_tiles_expose_boundary_fragments() -> VortexResult<()> { + let view = offset_view_fixture()?; + assert_eq!(view.row_offset(), 10); + assert_eq!(view.backing_rows(), 192); + assert!(view.is_full_width()); + + let tiles = view.tiles().collect::>(); + assert_eq!(tiles[0].row_range, 0..54); + assert_eq!(tiles[0].rows_within_tile, 10..64); + assert!(!tiles[0].is_full_tile()); + assert_eq!(tiles[1].row_range, 54..118); + assert_eq!(tiles[1].rows_within_tile, 0..64); + assert!(tiles[1].is_full_tile()); + assert_eq!(tiles[2].row_range, 118..128); + assert_eq!(tiles[2].rows_within_tile, 0..10); + Ok(()) +} + +#[test] +fn offset_view_has_complete_interior_tiles() -> VortexResult<()> { + let view = offset_view_fixture_with_backing_rows(138)?; + let tiles = view.tiles().collect::>(); + + assert_eq!(tiles.len(), 3); + assert_eq!(tiles[0].physical_range, 0..512); + assert_eq!(tiles[1].physical_range, 512..1024); + assert_eq!(tiles[2].physical_range, 1024..1104); + assert_eq!(view.tile_elements(&tiles[0])?.len(), 512); + assert_eq!(view.tile_elements(&tiles[1])?.len(), 512); + assert_eq!(view.tile_elements(&tiles[2])?.len(), 80); + assert!(!tiles[2].is_full_tile()); + Ok(()) +} + +#[rstest] +#[case::window_exceeds_backing(2, 1, 1, 1, 1, 2, 2)] +#[case::multi_slab_offset_view(3, 5, 5, 2, 3, 1, 3)] +#[case::child_extent_mismatch(6, 2, 2, 5, 2, 1, 7)] +#[case::backing_extent_overflow(0, 2, 2, 0, 2, 0, usize::MAX)] +#[case::empty_nonzero_offset(64, 1, 1, 0, 1, 64, 64)] +#[case::zero_width_nonzero_offset(3, 0, 0, 2, 3, 1, 3)] +fn malformed_view_metadata_is_rejected( + #[case] element_rows: usize, + #[case] element_dimensions: usize, + #[case] list_size: u32, + #[case] len: usize, + #[case] tile_dimensions: u32, + #[case] row_offset: usize, + #[case] backing_rows: usize, +) { + let elements = PrimitiveArray::from_iter( + (0..element_rows * element_dimensions) + .map(|index| u16::try_from(index).unwrap_or(u16::MAX)), + ) + .into_array(); + + assert!( + TiledFixedSizeList::try_new_view( + elements, + list_size, + Validity::NonNullable, + len, + geometry(64, tile_dimensions), + row_offset, + backing_rows, + ) + .is_err() + ); +} + +#[test] +fn golden_physical_child_and_round_trip() -> VortexResult<()> { + let canonical = FixedSizeListArray::new( + PrimitiveArray::from_iter([0u16, 1, 2, 3, 4, 10, 11, 12, 13, 14, 20, 21, 22, 23, 24]) + .into_array(), + 5, + Validity::NonNullable, + 3, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(2, 3), &mut ctx)?; + let physical = tiled + .elements() + .clone() + .execute::(&mut ctx)?; + assert_eq!( + physical.as_slice::(), + &[0, 10, 1, 11, 2, 12, 20, 21, 22, 3, 13, 4, 14, 23, 24] + ); + assert_fsl_equivalent(&canonical.into_array(), &tiled.into_array(), &mut ctx) +} + +#[test] +fn all_native_ptypes_round_trip() -> VortexResult<()> { + for ptype in [ + PType::U8, + PType::U16, + PType::U32, + PType::U64, + PType::I8, + PType::I16, + PType::I32, + PType::I64, + PType::F16, + PType::F32, + PType::F64, + ] { + match_each_native_ptype!(ptype, |T| { + let canonical = FixedSizeListArray::new( + PrimitiveArray::new(Buffer::::zeroed(15), Validity::NonNullable).into_array(), + 5, + Validity::NonNullable, + 3, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(2, 3), &mut ctx)?; + assert_fsl_equivalent(&canonical.into_array(), &tiled.into_array(), &mut ctx) + })?; + } + Ok(()) +} + +#[test] +fn independent_outer_and_element_validity_round_trip() -> VortexResult<()> { + let cases = [ + ( + Validity::NonNullable, + Validity::NonNullable, + Validity::NonNullable, + ), + (Validity::AllValid, Validity::AllValid, Validity::AllValid), + ( + Validity::AllInvalid, + Validity::AllInvalid, + Validity::AllInvalid, + ), + ( + Validity::from_iter([true, false, true, true, false, true]), + Validity::from_iter([true, false, true]), + Validity::from_iter([true, true, false, false, true, true]), + ), + ]; + + for (element_validity, outer_validity, expected_physical_validity) in cases { + let canonical = FixedSizeListArray::new( + PrimitiveArray::new( + Buffer::copy_from([0u16, 1, 10, 11, 20, 21]), + element_validity, + ) + .into_array(), + 2, + outer_validity, + 3, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(2, 1), &mut ctx)?; + + assert!( + tiled + .elements() + .validity()? + .mask_eq(&expected_physical_validity, 6, &mut ctx,)? + ); + assert!(tiled.array_validity().mask_eq( + &canonical.fixed_size_list_validity(), + 3, + &mut ctx, + )?); + + let executed = tiled.into_array().execute::(&mut ctx)?; + assert_fsl_equivalent(&canonical.into_array(), &executed.into_array(), &mut ctx)?; + } + Ok(()) +} + +#[test] +fn mixed_validity_round_trip_with_partial_row_and_dimension_tiles() -> VortexResult<()> { + let canonical_validity = Validity::from_iter([ + true, false, false, true, true, false, true, false, true, false, true, true, false, false, + true, + ]); + let expected_physical_validity = Validity::from_iter([ + true, false, false, true, false, false, true, true, false, true, true, true, false, false, + true, + ]); + let canonical = FixedSizeListArray::new( + PrimitiveArray::new(Buffer::copy_from([0u16; 15]), canonical_validity.clone()).into_array(), + 5, + Validity::NonNullable, + 3, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled = TiledFixedSizeList::encode(canonical.as_view(), geometry(2, 3), &mut ctx)?; + + assert!( + tiled + .elements() + .validity()? + .mask_eq(&expected_physical_validity, 15, &mut ctx,)? + ); + let decoded = tiled.into_array().execute::(&mut ctx)?; + assert!( + decoded + .elements() + .validity()? + .mask_eq(&canonical_validity, 15, &mut ctx,)? + ); + assert_fsl_equivalent(&canonical.into_array(), &decoded.into_array(), &mut ctx) +} + +#[test] +fn degenerate_arrays_encode_execute_and_zero_width_scalars() -> VortexResult<()> { + let zero_rows = FixedSizeListArray::new( + PrimitiveArray::from_iter(std::iter::empty::()).into_array(), + 5, + Validity::NonNullable, + 0, + ); + let mut ctx = SESSION.create_execution_ctx(); + let tiled_zero_rows = + TiledFixedSizeList::encode(zero_rows.as_view(), geometry(32, 64), &mut ctx)?; + let decoded_zero_rows = tiled_zero_rows + .into_array() + .execute::(&mut ctx)?; + assert_fsl_equivalent( + &zero_rows.into_array(), + &decoded_zero_rows.into_array(), + &mut ctx, + )?; + + let zero_width = FixedSizeListArray::new( + PrimitiveArray::from_iter(std::iter::empty::()).into_array(), + 0, + Validity::NonNullable, + 3, + ); + let tiled_zero_width = + TiledFixedSizeList::encode(zero_width.as_view(), geometry(32, 64), &mut ctx)?; + let decoded_zero_width = tiled_zero_width + .clone() + .into_array() + .execute::(&mut ctx)?; + assert_fsl_equivalent( + &zero_width.clone().into_array(), + &decoded_zero_width.into_array(), + &mut ctx, + )?; + for row in 0..zero_width.len() { + assert_eq!( + zero_width.execute_scalar(row, &mut ctx)?, + tiled_zero_width.execute_scalar(row, &mut ctx)?, + ); + } + Ok(()) +} + +#[test] +fn try_new_derives_dtype_and_accessors() -> VortexResult<()> { + let tiled = physical_fixture(3, 5, geometry(2, 3))?; + assert_eq!(tiled.len(), 3); + assert_eq!(tiled.list_size(), 5); + assert_eq!(tiled.geometry(), geometry(2, 3)); + assert_eq!(tiled.row_tile_count(), 2); + assert_eq!(tiled.dimension_tile_count(), 2); + assert_eq!( + tiled.dtype(), + &DType::FixedSizeList( + Arc::new(DType::Primitive(PType::U16, Nullability::NonNullable)), + 5, + Nullability::NonNullable, + ) + ); + Ok(()) +} + +#[rstest] +#[case::short(3, 5, 14)] +#[case::long(3, 5, 16)] +fn rejects_wrong_child_length( + #[case] rows: usize, + #[case] dimensions: u32, + #[case] physical_len: usize, +) { + let elements = PrimitiveArray::from_iter( + (0..physical_len).map(|index| u16::try_from(index).unwrap_or(u16::MAX)), + ) + .into_array(); + assert!( + TiledFixedSizeList::try_new( + elements, + dimensions, + Validity::NonNullable, + rows, + geometry(2, 3), + ) + .is_err() + ); +} + +#[test] +fn tiled_fsl_metadata() { + check_metadata( + "tiled_fsl.metadata", + &TiledFixedSizeListMetadata { + tile_rows: 32, + tile_dimensions: 64, + row_offset: u32::MAX, + backing_rows: u64::MAX, + } + .encode_to_vec(), + ); +} + +struct TestArrayChildren(Vec); + +impl ArrayChildren for TestArrayChildren { + fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult { + let child = self + .0 + .as_slice() + .get(index) + .ok_or_else(|| vortex_err!(InvalidArgument: "missing test child {index}"))?; + if child.dtype() != dtype || child.len() != len { + return Err(vortex_err!(InvalidArgument: + "test child {index} has dtype {} and length {}, expected {dtype} and {len}", + child.dtype(), child.len() + )); + } + Ok(child.clone()) + } + + fn len(&self) -> usize { + self.0.len() + } +} + +fn deserialize_test_metadata( + metadata: TiledFixedSizeListMetadata, + len: usize, + list_size: u32, + elements: ArrayRef, +) -> VortexResult<()> { + let dtype = DType::FixedSizeList( + Arc::new(elements.dtype().clone()), + list_size, + Nullability::NonNullable, + ); + ::deserialize( + &TiledFixedSizeList, + &dtype, + len, + &metadata.encode_to_vec(), + &[] as &[BufferHandle], + &TestArrayChildren(vec![elements]), + &SESSION, + )?; + Ok(()) +} + +#[rstest] +#[case::backing_extent_overflow(2, 0, u64::MAX, 1, 2, 0, "backing")] +#[case::window_outside_backing(1, 1, 2, 2, 1, 2, "row window 1..3 exceeds 2 backing rows")] +#[case::excess_retained_tiles(1, 1, 128, 1, 1, 128, "exceeds the retained tile extent")] +fn deserialize_rejects_malformed_row_view_metadata( + #[case] tile_dimensions: u32, + #[case] row_offset: u32, + #[case] backing_rows: u64, + #[case] len: usize, + #[case] list_size: u32, + #[case] element_len: usize, + #[case] expected_error: &str, +) { + let error = deserialize_test_metadata( + TiledFixedSizeListMetadata { + tile_rows: 64, + tile_dimensions, + row_offset, + backing_rows, + }, + len, + list_size, + PrimitiveArray::from_iter((0..element_len).map(|_| 0u16)).into_array(), + ) + .unwrap_err(); + assert!(error.to_string().contains(expected_error)); +} + +#[test] +fn serialized_array_rejects_validity_child_for_nonnullable_dtype_without_panicking() +-> VortexResult<()> { + let tiled = TiledFixedSizeList::try_new( + PrimitiveArray::from_iter([10u16, 20]).into_array(), + 1, + Validity::AllInvalid, + 2, + geometry(2, 1), + )?; + let array_context = ArrayContext::empty(); + let serialized = + tiled + .into_array() + .serialize(&array_context, &SESSION, &SerializeOptions::default())?; + let mut bytes = ByteBufferMut::empty(); + for buffer in serialized { + bytes.extend_from_slice(buffer.as_ref()); + } + let serialized = SerializedArray::try_from(bytes.freeze())?; + let nonnullable_dtype = + DType::FixedSizeList(Arc::new(PType::U16.into()), 1, Nullability::NonNullable); + let read_context = ReadContext::new(array_context.to_ids()); + + let decoded = std::panic::catch_unwind(AssertUnwindSafe(|| { + serialized.decode(&nonnullable_dtype, 2, &read_context, &SESSION) + })); + let result = match decoded { + Ok(result) => result, + Err(_) => return Err(vortex_err!("malformed tiled FSL decode panicked")), + }; + let error = match result { + Ok(_) => return Err(vortex_err!("malformed tiled FSL decode succeeded")), + Err(error) => error, + }; + assert!( + error + .to_string() + .contains("cannot have an outer validity child") + ); + Ok(()) +} + +#[test] +fn rejects_non_primitive_child() { + assert!( + TiledFixedSizeList::try_new( + VarBinViewArray::from_iter_str(["x"]).into_array(), + 1, + Validity::NonNullable, + 1, + geometry(1, 1), + ) + .is_err() + ); +} + +#[test] +fn rejects_wrong_outer_validity_length() { + assert!( + TiledFixedSizeList::try_new( + PrimitiveArray::from_iter([1u8, 2, 3]).into_array(), + 1, + Validity::from_iter([true, false]), + 3, + geometry(2, 1), + ) + .is_err() + ); +} + +#[test] +fn rejects_len_times_list_size_overflow() { + assert!( + TiledFixedSizeList::try_new( + PrimitiveArray::from_iter(std::iter::empty::()).into_array(), + 2, + Validity::NonNullable, + usize::MAX, + geometry(1, 1), + ) + .is_err() + ); +} + +#[test] +fn rejects_zero_geometry_metadata() { + let metadata = TiledFixedSizeListMetadata { + tile_rows: 0, + tile_dimensions: 64, + row_offset: 0, + backing_rows: 0, + }; + assert!(TileGeometry::try_from(&metadata).is_err()); +} + +#[test] +fn max_usize_representable_geometry_constructs_and_traverses() -> VortexResult<()> { + let tiled = physical_fixture(1, 1, geometry(u32::MAX, u32::MAX))?; + let bounds: Vec = tiled.tiles().collect(); + assert_eq!(bounds.len(), 1); + assert_eq!(bounds[0].row_range, 0..1); + assert_eq!(bounds[0].dimension_range, 0..1); + assert_eq!(bounds[0].physical_range, 0..1); + Ok(()) +} diff --git a/encodings/tiled-fsl/src/transpose.rs b/encodings/tiled-fsl/src/transpose.rs new file mode 100644 index 00000000000..635077cca0a --- /dev/null +++ b/encodings/tiled-fsl/src/transpose.rs @@ -0,0 +1,193 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::match_each_native_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use crate::TileBoundsIter; +use crate::TileGeometry; +use crate::geometry::geometry_usizes; +use crate::geometry::tile_bounds; + +#[expect( + clippy::cognitive_complexity, + reason = "complexity is attributed to native-type dispatch macro expansion" +)] +pub(crate) fn encode_elements( + elements: ArrayView<'_, Primitive>, + len: usize, + list_size: usize, + geometry: TileGeometry, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let expected_len = len.checked_mul(list_size); + vortex_ensure!( + expected_len == Some(elements.len()), + InvalidArgument: + "physical child length {} does not match logical extent ({len}, {list_size})", + elements.len() + ); + + match_each_native_ptype!(elements.ptype(), |T| { + let source = elements.as_slice::(); + let mut output = BufferMut::::with_capacity(source.len()); + let (source_validity, preserved_validity) = match elements.validity()? { + validity @ Validity::Array(_) => match validity.execute_mask(elements.len(), ctx)? { + Mask::Values(values) => (Some(values), None), + mask => ( + None, + Some(Validity::from_mask( + mask, + vortex_array::dtype::Nullability::Nullable, + )), + ), + }, + validity => (None, Some(validity)), + }; + let mut output_validity = source_validity + .as_ref() + .map(|_| BitBufferMut::new_unset(elements.len())); + for dimension_tile in 0..list_size.div_ceil(usize::try_from(geometry.dimensions().get())?) { + for row_tile in 0..len.div_ceil(usize::try_from(geometry.rows().get())?) { + let bounds = tile_bounds(len, list_size, geometry, row_tile, dimension_tile)?; + for dimension in bounds.dimension_range.clone() { + for row in bounds.row_range.clone() { + let canonical = row * list_size + dimension; + let physical = output.len(); + output.push(source[canonical]); + if let (Some(source_validity), Some(output_validity)) = + (&source_validity, &mut output_validity) + { + output_validity.set_to(physical, source_validity.value(canonical)); + } + } + } + } + } + let validity = match output_validity { + Some(validity) => Validity::from(validity.freeze()), + None => preserved_validity.ok_or_else( + || vortex_err!(InvalidArgument: "validity must be preserved or transposed"), + )?, + }; + Ok(PrimitiveArray::new(output.freeze(), validity)) + }) +} + +#[cfg_attr(not(test), allow(dead_code))] +pub(crate) fn decode_elements( + elements: ArrayView<'_, Primitive>, + len: usize, + list_size: usize, + geometry: TileGeometry, + ctx: &mut ExecutionCtx, +) -> VortexResult { + decode_visible_elements(elements, len, list_size, geometry, 0, len, ctx) +} + +#[expect( + clippy::cognitive_complexity, + reason = "complexity is attributed to native-type dispatch macro expansion" +)] +pub(crate) fn decode_visible_elements( + elements: ArrayView<'_, Primitive>, + len: usize, + list_size: usize, + geometry: TileGeometry, + row_offset: usize, + backing_rows: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let expected_len = backing_rows.checked_mul(list_size); + vortex_ensure!( + expected_len == Some(elements.len()), + InvalidArgument: "physical child length {} does not match retained extent ({backing_rows}, {list_size})", + elements.len() + ); + let (tile_rows, tile_dimensions) = geometry_usizes(geometry)?; + let row_tile_count = if len == 0 { + 0 + } else { + (row_offset + len).div_ceil(tile_rows) + }; + let dimension_tile_count = list_size.div_ceil(tile_dimensions); + + match_each_native_ptype!(elements.ptype(), |T| { + let source = elements.as_slice::(); + let output_len = len.checked_mul(list_size).ok_or_else(|| { + vortex_err!(InvalidArgument: "logical length {len} times list size {list_size} overflows usize") + })?; + let mut output = BufferMut::::zeroed(output_len); + let (source_validity, preserved_validity) = match elements.validity()? { + validity @ Validity::Array(_) => match validity.execute_mask(elements.len(), ctx)? { + Mask::Values(values) => (Some(values), None), + mask => ( + None, + Some(Validity::from_mask( + mask, + vortex_array::dtype::Nullability::Nullable, + )), + ), + }, + validity => (None, Some(validity)), + }; + let mut output_validity = source_validity + .as_ref() + .map(|_| BitBufferMut::new_unset(output_len)); + let bounds = if row_offset == 0 && backing_rows == len { + TileBoundsIter::new( + len, + list_size, + geometry, + row_tile_count, + dimension_tile_count, + ) + } else { + TileBoundsIter::new_view( + len, + list_size, + geometry, + row_offset, + backing_rows, + row_tile_count, + dimension_tile_count, + ) + }; + for bounds in bounds { + let retained_rows = bounds.physical_range.len() / bounds.dimension_range.len(); + for (dimension_offset, dimension) in bounds.dimension_range.clone().enumerate() { + let mut physical = bounds.physical_range.start + + dimension_offset * retained_rows + + bounds.rows_within_tile.start; + for row in bounds.row_range.clone() { + let canonical = row * list_size + dimension; + output[canonical] = source[physical]; + if let (Some(source_validity), Some(output_validity)) = + (&source_validity, &mut output_validity) + { + output_validity.set_to(canonical, source_validity.value(physical)); + } + physical += 1; + } + } + } + let validity = match output_validity { + Some(validity) => Validity::from(validity.freeze()), + None => preserved_validity.ok_or_else( + || vortex_err!(InvalidArgument: "validity must be preserved or transposed"), + )?, + }; + Ok(PrimitiveArray::new(output.freeze(), validity)) + }) +} diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index 41675f2c03a..76fff82d9a1 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -42,6 +42,7 @@ vortex-mask = { workspace = true } vortex-row = { workspace = true } vortex-runend = { workspace = true, features = ["arbitrary"] } vortex-session = { workspace = true } +vortex-tiled-fsl = { workspace = true } vortex-utils = { workspace = true } # Native-only: libfuzzer harness and file IO (won't compile to WASM) @@ -113,3 +114,11 @@ name = "row_encode" path = "fuzz_targets/row_encode.rs" test = false required-features = ["native"] + +[[bin]] +bench = false +doc = false +name = "tiled_fsl" +path = "fuzz_targets/tiled_fsl.rs" +test = false +required-features = ["native"] diff --git a/fuzz/README.md b/fuzz/README.md index f2171456e99..a6b285891d6 100644 --- a/fuzz/README.md +++ b/fuzz/README.md @@ -15,3 +15,17 @@ input with the command `cargo fuzz run array_ops ` or `cargo f If there are any linking (on macOS) then run `cargo fuzz run --dev --sanitizer=none ...`. `--dev` runs the fuzzer in dev profile. + +## Tiled fixed-size lists + +Run the native differential property target against canonical fixed-size lists: + +```shell +cargo +nightly fuzz run tiled_fsl +``` + +Replay one saved crash input: + +```shell +cargo +nightly fuzz run tiled_fsl +``` diff --git a/fuzz/fuzz_targets/tiled_fsl.rs b/fuzz/fuzz_targets/tiled_fsl.rs new file mode 100644 index 00000000000..50971a01af1 --- /dev/null +++ b/fuzz/fuzz_targets/tiled_fsl.rs @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![no_main] + +use libfuzzer_sys::Corpus; +use libfuzzer_sys::fuzz_target; +use vortex_error::vortex_panic; +use vortex_fuzz::FuzzTiledFsl; +use vortex_fuzz::run_tiled_fsl; + +fuzz_target!(|input: FuzzTiledFsl| -> Corpus { + match run_tiled_fsl(input) { + Ok(()) => Corpus::Keep, + Err(error) => vortex_panic!("{error}"), + } +}); diff --git a/fuzz/src/lib.rs b/fuzz/src/lib.rs index 1f1b1be1778..64e6e753f53 100644 --- a/fuzz/src/lib.rs +++ b/fuzz/src/lib.rs @@ -8,6 +8,7 @@ pub mod compress; pub mod error; pub mod fsst_like; mod row; +mod tiled_fsl; // File module only available for native builds (requires vortex-file which uses tokio) #[cfg(not(target_arch = "wasm32"))] @@ -34,6 +35,10 @@ pub use gpu::FuzzCompressGpu; pub use gpu::run_compress_gpu; pub use row::FuzzRowEncode; pub use row::run_row_encode; +pub use tiled_fsl::FuzzTiledFsl; +pub use tiled_fsl::TiledFslAction; +pub use tiled_fsl::deterministic_tiled_fsl_cases; +pub use tiled_fsl::run_tiled_fsl; pub const FUZZ_ARRAY_MAX_LEN: usize = 2048; pub const FUZZ_FILE_ARRAY_MAX_LEN: usize = 16_384; diff --git a/fuzz/src/tiled_fsl.rs b/fuzz/src/tiled_fsl.rs new file mode 100644 index 00000000000..d01004ed432 --- /dev/null +++ b/fuzz/src/tiled_fsl.rs @@ -0,0 +1,778 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Differential property testing for tiled primitive fixed-size lists. + +use std::num::NonZeroU32; +use std::ops::ControlFlow; +use std::sync::Arc; +use std::sync::LazyLock; + +use arbitrary::Arbitrary; +use arbitrary::Unstructured; +use vortex_array::Array; +use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::FixedSizeList; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Slice; +use vortex_array::arrays::arbitrary::ArbitraryArray; +use vortex_array::arrays::arbitrary::ArbitraryArrayConfig; +use vortex_array::arrays::arbitrary::ArbitraryWith; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_tiled_fsl::TileGeometry; +use vortex_tiled_fsl::TiledFixedSizeList; +use vortex_tiled_fsl::TiledFixedSizeListArrayExt; +use vortex_tiled_fsl::TiledFixedSizeListArraySlotsExt; + +use crate::array::assert_array_eq; +use crate::array::assert_scalar_eq; +use crate::array::slice_canonical_array; +use crate::array::take_canonical_array; +use crate::error::Backtrace; +use crate::error::VortexFuzzError; +use crate::error::VortexFuzzResult; + +static TILED_FSL_SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_tiled_fsl::initialize(&session); + session +}); + +#[derive(Clone, Debug)] +pub enum TiledFslAction { + CheckTiles, + ScalarAt(u16), + Slice { start: u16, stop: u16 }, + Take(Vec>), + Reconstruct, + ReconstructSerde, +} + +#[derive(Debug)] +pub struct FuzzTiledFsl { + canonical: ArrayRef, + geometry: TileGeometry, + actions: Vec, +} + +impl<'a> Arbitrary<'a> for FuzzTiledFsl { + fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result { + let ptype: PType = u.arbitrary()?; + let element_nullability: Nullability = u.arbitrary()?; + let outer_nullability: Nullability = u.arbitrary()?; + let list_size = u.int_in_range(0u32..=64)?; + let row_count = u.int_in_range(0usize..=128)?; + let dtype = DType::FixedSizeList( + Arc::new(DType::Primitive(ptype, element_nullability)), + list_size, + outer_nullability, + ); + let canonical = ArbitraryArray::arbitrary_with_config( + u, + &ArbitraryArrayConfig { + dtype: Some(dtype), + len: row_count..=row_count, + }, + )? + .0; + + let geometry = TileGeometry::new( + NonZeroU32::new(u.int_in_range(1u32..=128)?) + .ok_or(arbitrary::Error::IncorrectFormat)?, + NonZeroU32::new(u.int_in_range(1u32..=128)?) + .ok_or(arbitrary::Error::IncorrectFormat)?, + ); + let action_count = u.int_in_range(1usize..=8)?; + let mut actions = Vec::with_capacity(action_count); + for _ in 0..action_count { + actions.push(match u.int_in_range(0u8..=5)? { + 0 => TiledFslAction::CheckTiles, + 1 => TiledFslAction::ScalarAt(u.arbitrary()?), + 2 => TiledFslAction::Slice { + start: u.arbitrary()?, + stop: u.arbitrary()?, + }, + 3 => { + let take_len = u.int_in_range(0usize..=64)?; + let mut seeds = Vec::with_capacity(take_len); + for _ in 0..take_len { + seeds.push(u.arbitrary()?); + } + TiledFslAction::Take(seeds) + } + 4 => TiledFslAction::Reconstruct, + 5 => TiledFslAction::ReconstructSerde, + _ => unreachable!("action tag is bounded"), + }); + } + + Ok(Self { + canonical, + geometry, + actions, + }) + } +} + +#[expect(clippy::result_large_err)] +fn fuzz(result: VortexResult) -> VortexFuzzResult { + result.map_err(|error| VortexFuzzError::VortexError(error, Backtrace::capture())) +} + +fn assert_tiled_geometry(array: &ArrayRef, expected: TileGeometry) -> VortexResult<()> { + if !array.is::() { + vortex_bail!("expected nondegenerate operation to retain tiled FSL"); + } + let actual = array.as_::().geometry(); + if actual != expected { + vortex_bail!("expected geometry {expected:?}, found {actual:?}"); + } + Ok(()) +} + +struct SerializedChildren(Vec); + +impl ArrayChildren for SerializedChildren { + fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult { + let child = self + .0 + .as_slice() + .get(index) + .ok_or_else(|| vortex_err!(InvalidArgument: "missing serialized child {index}"))?; + vortex_ensure!( + child.dtype() == dtype && child.len() == len, + InvalidArgument: + "serialized child {index} has dtype {} and len {}, expected {dtype} and {len}", + child.dtype(), + child.len() + ); + Ok(child.clone()) + } + + fn len(&self) -> usize { + self.0.len() + } +} + +fn validate_array_tree(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { + for descendant in array.depth_first_traversal() { + descendant.validity()?.execute_mask(descendant.len(), ctx)?; + } + Ok(()) +} + +fn reconstruct_serde(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let tiled = array.as_::(); + let expected_row_offset = tiled.row_offset(); + let expected_backing_rows = tiled.backing_rows(); + let metadata = ::serialize(tiled, ctx.session())? + .ok_or_else(|| vortex_err!("tiled fixed-size list did not serialize metadata"))?; + let children = SerializedChildren(array.slots().iter().flatten().cloned().collect()); + let parts = ::deserialize( + &TiledFixedSizeList, + array.dtype(), + array.len(), + &metadata, + &[] as &[BufferHandle], + &children, + ctx.session(), + )?; + let reconstructed = Array::::try_from_parts(parts)?.into_array(); + let reconstructed_tiled = reconstructed.as_::(); + vortex_ensure!( + reconstructed_tiled.row_offset() == expected_row_offset, + "serde changed row offset from {expected_row_offset} to {}", + reconstructed_tiled.row_offset() + ); + vortex_ensure!( + reconstructed_tiled.backing_rows() == expected_backing_rows, + "serde changed backing rows from {expected_backing_rows} to {}", + reconstructed_tiled.backing_rows() + ); + validate_array_tree(&reconstructed, ctx)?; + Ok(reconstructed) +} + +#[expect(clippy::result_large_err)] +fn check_tiles( + canonical: &ArrayRef, + tiled: &ArrayRef, + step: usize, + ctx: &mut ExecutionCtx, +) -> VortexFuzzResult<()> { + fuzz((|| { + vortex_ensure!( + tiled.is::(), + "tile check requires a tiled fixed-size list" + ); + Ok(()) + })())?; + let canonical_fsl = fuzz(canonical.clone().execute::(ctx))?; + let tiled_fsl = tiled.as_::(); + let rows = canonical.len(); + let list_size = canonical_fsl.list_size() as usize; + let tile_rows = tiled_fsl.geometry().rows().get() as usize; + let tile_dimensions = tiled_fsl.geometry().dimensions().get() as usize; + let row_offset = tiled_fsl.row_offset(); + let backing_rows = tiled_fsl.backing_rows(); + let logical_end = fuzz(row_offset.checked_add(rows).ok_or_else( + || vortex_err!(InvalidArgument: "row offset plus length overflows logical extent"), + ))?; + let expected_row_tile_count = if rows == 0 { + 0 + } else { + logical_end.div_ceil(tile_rows) + }; + let expected_dimension_tile_count = list_size.div_ceil(tile_dimensions); + + fuzz((|| { + vortex_ensure!( + logical_end <= backing_rows, + "logical row window {row_offset}..{logical_end} exceeds {backing_rows} backing rows" + ); + vortex_ensure!( + tiled_fsl.row_tile_count() == expected_row_tile_count, + "row tile count mismatch: expected {expected_row_tile_count}, found {}", + tiled_fsl.row_tile_count() + ); + vortex_ensure!( + tiled_fsl.dimension_tile_count() == expected_dimension_tile_count, + "dimension tile count mismatch: expected {expected_dimension_tile_count}, found {}", + tiled_fsl.dimension_tile_count() + ); + Ok(()) + })())?; + + let outer_mask = fuzz( + canonical + .validity() + .and_then(|validity| validity.execute_mask(rows, ctx)), + )?; + let mut physical_cursor = 0usize; + let mut tile_count = 0usize; + + for (tile_index, bounds) in tiled_fsl.tiles().enumerate() { + let dimension_tile = tile_index / expected_row_tile_count; + let row_tile = tile_index % expected_row_tile_count; + let retained_row_start = row_tile * tile_rows; + let retained_row_end = retained_row_start + .saturating_add(tile_rows) + .min(backing_rows); + let retained_row_count = retained_row_end - retained_row_start; + let visible_row_start = retained_row_start.max(row_offset); + let visible_row_end = retained_row_end.min(logical_end); + let expected_rows = (visible_row_start - row_offset)..(visible_row_end - row_offset); + let expected_rows_within_tile = + (visible_row_start - retained_row_start)..(visible_row_end - retained_row_start); + let dimension_start = dimension_tile * tile_dimensions; + let dimension_end = dimension_start + .saturating_add(tile_dimensions) + .min(list_size); + let physical_len = retained_row_count * (dimension_end - dimension_start); + let expected_physical_end = physical_cursor + physical_len; + let expected_dimensions = dimension_start..dimension_end; + let expected_physical = physical_cursor..expected_physical_end; + + fuzz((|| { + vortex_ensure!( + dimension_tile < expected_dimension_tile_count, + "tile iterator yielded unexpected tile {tile_index}" + ); + vortex_ensure!( + bounds.row_range == expected_rows, + "tile {tile_index} row range mismatch: expected {expected_rows:?}, found {:?}", + bounds.row_range + ); + vortex_ensure!( + bounds.dimension_range == expected_dimensions, + "tile {tile_index} dimension range mismatch: expected {expected_dimensions:?}, found {:?}", + bounds.dimension_range + ); + vortex_ensure!( + bounds.rows_within_tile == expected_rows_within_tile, + "tile {tile_index} visible rows mismatch: expected {expected_rows_within_tile:?}, found {:?}", + bounds.rows_within_tile + ); + vortex_ensure!( + bounds.physical_range == expected_physical, + "tile {tile_index} physical range mismatch: expected {expected_physical:?}, found {:?}", + bounds.physical_range + ); + vortex_ensure!( + bounds.physical_range.len() == retained_row_count * bounds.dimension_range.len(), + "tile {tile_index} retained physical cardinality mismatch" + ); + vortex_ensure!( + bounds.rows_within_tile.len() == bounds.row_range.len() + && bounds.rows_within_tile.end <= retained_row_count, + "tile {tile_index} visible rows exceed its retained physical rows" + ); + Ok(()) + })())?; + + let actual_tile = fuzz(tiled_fsl.tile_elements(&bounds))?; + let selected_expected_positions = expected_dimensions.clone().flat_map(|dimension| { + expected_rows.clone().filter_map({ + let outer_mask = &outer_mask; + move |row| { + outer_mask + .value(row) + .then_some((row * list_size + dimension) as u64) + } + }) + }); + let selected_actual_positions = + expected_dimensions + .clone() + .enumerate() + .flat_map(|(dimension_offset, _)| { + expected_rows.clone().filter_map({ + let outer_mask = &outer_mask; + let visible_start = bounds.rows_within_tile.start; + let logical_start = bounds.row_range.start; + move |row| { + outer_mask.value(row).then_some( + (dimension_offset * retained_row_count + visible_start + row + - logical_start) as u64, + ) + } + }) + }); + let expected_selection = + PrimitiveArray::from_iter(selected_expected_positions).into_array(); + let actual_selection = PrimitiveArray::from_iter(selected_actual_positions).into_array(); + let selected_expected = fuzz(canonical_fsl.elements().clone().take(expected_selection))?; + let selected_actual = fuzz(actual_tile.take(actual_selection))?; + assert_array_eq(&selected_expected, &selected_actual, step, ctx)?; + + physical_cursor = expected_physical_end; + tile_count += 1; + } + + fuzz((|| { + vortex_ensure!( + tile_count == expected_row_tile_count * expected_dimension_tile_count, + "tile iterator count mismatch: expected {}, found {tile_count}", + expected_row_tile_count * expected_dimension_tile_count + ); + vortex_ensure!( + physical_cursor == backing_rows * list_size, + "physical tile ranges cover {physical_cursor} values, expected {}", + backing_rows * list_size + ); + Ok(()) + })()) +} + +#[expect(clippy::result_large_err)] +fn execute_action( + action: TiledFslAction, + canonical: &mut ArrayRef, + tiled: &mut ArrayRef, + geometry: TileGeometry, + step: usize, + ctx: &mut ExecutionCtx, +) -> VortexFuzzResult> { + match action { + TiledFslAction::CheckTiles => { + if tiled.is::() { + check_tiles(canonical, tiled, step, ctx)?; + } + } + TiledFslAction::ScalarAt(seed) => { + if canonical.is_empty() { + return Ok(ControlFlow::Continue(())); + } + let row = usize::from(seed) % canonical.len(); + let expected = fuzz(canonical.execute_scalar(row, ctx))?; + let actual = fuzz(tiled.execute_scalar(row, ctx))?; + assert_scalar_eq(&expected, &actual, step)?; + } + TiledFslAction::Slice { start, stop } => { + let source_is_empty = canonical.is_empty(); + let source_len = canonical.len(); + let source_is_tiled = tiled.is::(); + let source_is_full_width = + source_is_tiled && tiled.as_::().is_full_width(); + let first = usize::from(start).min(usize::from(stop)); + let last = usize::from(start).max(usize::from(stop)); + let start = first % (canonical.len() + 1); + let stop = start + last % (canonical.len() - start + 1); + *canonical = fuzz(slice_canonical_array(canonical, start, stop, ctx))?; + *tiled = fuzz(tiled.clone().slice(start..stop))?; + if tiled.is_empty() { + if source_is_empty { + fuzz(assert_tiled_geometry(tiled, geometry))?; + assert_array_eq(canonical, tiled, step, ctx)?; + return Ok(ControlFlow::Continue(())); + } + fuzz((|| { + vortex_ensure!( + tiled.is::(), + "expected an empty slice of a nonempty source to be canonical FSL" + ); + Ok(()) + })())?; + assert_array_eq(canonical, tiled, step, ctx)?; + return Ok(ControlFlow::Break(())); + } + let tile_rows = geometry.rows().get() as usize; + let aligned_multi_slab = start.is_multiple_of(tile_rows) + && (stop.is_multiple_of(tile_rows) || stop == source_len); + if source_is_tiled && (source_is_full_width || aligned_multi_slab) { + fuzz(assert_tiled_geometry(tiled, geometry))?; + } else { + fuzz((|| { + vortex_ensure!( + tiled.is::(), + "expected a non-preserving slice to remain a lazy Slice, found {}", + tiled.encoding_id() + ); + Ok(()) + })())?; + } + assert_array_eq(canonical, tiled, step, ctx)?; + } + TiledFslAction::Take(seeds) => { + let source_is_empty = canonical.is_empty(); + let indices = seeds + .into_iter() + .map(|seed| { + seed.and_then(|seed| { + (!source_is_empty).then(|| usize::from(seed) % canonical.len()) + }) + }) + .collect::>(); + let index_array = if indices.contains(&None) { + PrimitiveArray::from_option_iter( + indices.iter().map(|index| index.map(|index| index as u64)), + ) + .into_array() + } else { + PrimitiveArray::from_iter(indices.iter().flatten().map(|index| *index as u64)) + .into_array() + }; + *canonical = fuzz(take_canonical_array(canonical, &indices, ctx))?; + let lazy = fuzz(tiled.clone().take(index_array))?; + *tiled = fuzz(lazy.execute::(ctx))?.into_array(); + assert_array_eq(canonical, tiled, step, ctx)?; + return Ok(ControlFlow::Break(())); + } + TiledFslAction::Reconstruct => { + if !tiled.is::() { + return Ok(ControlFlow::Continue(())); + } + let array = tiled.as_::(); + if array.row_offset() != 0 || array.backing_rows() != array.len() { + return Ok(ControlFlow::Continue(())); + } + *tiled = fuzz(TiledFixedSizeList::try_new( + array.elements().clone(), + array.list_size(), + array.array_validity(), + array.len(), + geometry, + ))? + .into_array(); + assert_array_eq(canonical, tiled, step, ctx)?; + } + TiledFslAction::ReconstructSerde => { + if !tiled.is::() { + return Ok(ControlFlow::Continue(())); + } + *tiled = fuzz(reconstruct_serde(tiled, ctx))?; + fuzz(assert_tiled_geometry(tiled, geometry))?; + assert_array_eq(canonical, tiled, step, ctx)?; + } + } + Ok(ControlFlow::Continue(())) +} + +#[expect(clippy::result_large_err)] +pub fn run_tiled_fsl(input: FuzzTiledFsl) -> VortexFuzzResult<()> { + let mut ctx = TILED_FSL_SESSION.create_execution_ctx(); + let mut canonical = fuzz( + input + .canonical + .execute::(&mut ctx) + .map(IntoArray::into_array), + )?; + let canonical_fsl = canonical.as_::(); + let mut tiled = fuzz(TiledFixedSizeList::encode( + canonical_fsl, + input.geometry, + &mut ctx, + ))? + .into_array(); + + fuzz(assert_tiled_geometry(&tiled, input.geometry))?; + assert_array_eq(&canonical, &tiled, 0, &mut ctx)?; + + if !canonical.is_empty() { + let tile_rows = input.geometry.rows().get() as usize; + let mut probe_rows = vec![0, canonical.len() - 1]; + probe_rows.extend((tile_rows..canonical.len()).step_by(tile_rows)); + probe_rows.sort_unstable(); + probe_rows.dedup(); + for row in probe_rows { + let expected = fuzz(canonical.execute_scalar(row, &mut ctx))?; + let actual = fuzz(tiled.execute_scalar(row, &mut ctx))?; + assert_scalar_eq(&expected, &actual, 0)?; + } + } + + let expected_row_tiles = canonical + .len() + .div_ceil(input.geometry.rows().get() as usize); + let expected_dimension_tiles = + (canonical_fsl.list_size() as usize).div_ceil(input.geometry.dimensions().get() as usize); + let tiled_fsl = tiled.as_::(); + fuzz((|| { + vortex_ensure!( + tiled_fsl.row_tile_count() == expected_row_tiles, + "row tile count mismatch: expected {expected_row_tiles}, found {}", + tiled_fsl.row_tile_count() + ); + vortex_ensure!( + tiled_fsl.dimension_tile_count() == expected_dimension_tiles, + "dimension tile count mismatch: expected {expected_dimension_tiles}, found {}", + tiled_fsl.dimension_tile_count() + ); + Ok(()) + })())?; + + for (step, action) in input.actions.into_iter().enumerate() { + if execute_action( + action, + &mut canonical, + &mut tiled, + input.geometry, + step, + &mut ctx, + )? + .is_break() + { + break; + } + } + Ok(()) +} + +fn geometry(rows: u32, dimensions: u32) -> TileGeometry { + let Some(rows) = NonZeroU32::new(rows) else { + unreachable!("deterministic geometry rows are nonzero"); + }; + let Some(dimensions) = NonZeroU32::new(dimensions) else { + unreachable!("deterministic geometry dimensions are nonzero"); + }; + TileGeometry::new(rows, dimensions) +} + +#[expect(clippy::result_large_err)] +pub fn deterministic_tiled_fsl_cases() -> VortexFuzzResult> { + let zero_by_zero = FixedSizeListArray::new( + PrimitiveArray::from_iter::<[i32; 0]>([]).into_array(), + 0, + Validity::NonNullable, + 0, + ) + .into_array(); + + let nullable_u16 = FixedSizeListArray::new( + PrimitiveArray::from_option_iter((0u16..15).map(|value| (value % 4 != 1).then_some(value))) + .into_array(), + 5, + Validity::from_iter([true, false, true]), + 3, + ) + .into_array(); + + let nullable_f32 = FixedSizeListArray::new( + PrimitiveArray::from_option_iter( + (0..65 * 129).map(|index| (index % 7 != 3).then_some((index as f32) * 0.25 - 1000.0)), + ) + .into_array(), + 129, + Validity::from_iter((0..65).map(|row| row % 5 != 2)), + 65, + ) + .into_array(); + + let full_width_row_view = FixedSizeListArray::new( + PrimitiveArray::new( + Buffer::from_iter((0..200 * 128).map(|index| index as u32)), + Validity::from_iter((0..200 * 128).map(|index| index % 11 != 0)), + ) + .into_array(), + 128, + Validity::from_iter((0..200).map(|row| row % 7 != 0)), + 200, + ) + .into_array(); + + let multi_slab_slices = FixedSizeListArray::new( + PrimitiveArray::new( + Buffer::from_iter((0..200 * 129).map(|index| index as u32)), + Validity::from_iter((0..200 * 129).map(|index| index % 13 != 0)), + ) + .into_array(), + 129, + Validity::from_iter((0..200).map(|row| row % 5 != 0)), + 200, + ) + .into_array(); + + Ok(vec![ + FuzzTiledFsl { + canonical: zero_by_zero, + geometry: geometry(128, 128), + actions: vec![ + TiledFslAction::CheckTiles, + TiledFslAction::Slice { start: 0, stop: 0 }, + TiledFslAction::CheckTiles, + TiledFslAction::Reconstruct, + ], + }, + FuzzTiledFsl { + canonical: nullable_u16, + geometry: geometry(2, 3), + actions: vec![ + TiledFslAction::CheckTiles, + TiledFslAction::Reconstruct, + TiledFslAction::Take(vec![Some(2), None, Some(2), Some(0)]), + ], + }, + FuzzTiledFsl { + canonical: nullable_f32, + geometry: geometry(32, 64), + actions: vec![ + TiledFslAction::CheckTiles, + TiledFslAction::Slice { + start: 31, + stop: 38, + }, + TiledFslAction::ScalarAt(1), + TiledFslAction::Take(vec![Some(2), Some(1), Some(0)]), + ], + }, + FuzzTiledFsl { + canonical: full_width_row_view, + geometry: geometry(64, 128), + actions: vec![ + TiledFslAction::CheckTiles, + TiledFslAction::Slice { + start: 10, + stop: 150, + }, + TiledFslAction::CheckTiles, + TiledFslAction::ScalarAt(53), + TiledFslAction::ScalarAt(54), + TiledFslAction::ReconstructSerde, + TiledFslAction::CheckTiles, + TiledFslAction::Slice { + start: 60, + stop: 70, + }, + TiledFslAction::CheckTiles, + TiledFslAction::ScalarAt(9), + TiledFslAction::Take(vec![Some(9), None, Some(0)]), + ], + }, + FuzzTiledFsl { + canonical: multi_slab_slices, + geometry: geometry(64, 64), + actions: vec![ + TiledFslAction::CheckTiles, + TiledFslAction::Slice { + start: 64, + stop: 192, + }, + TiledFslAction::ReconstructSerde, + TiledFslAction::ScalarAt(63), + TiledFslAction::Slice { start: 1, stop: 66 }, + TiledFslAction::ScalarAt(0), + TiledFslAction::Slice { start: 1, stop: 64 }, + TiledFslAction::Take(vec![Some(62), None, Some(0)]), + ], + }, + ]) +} + +#[cfg(test)] +mod tests { + use std::ops::ControlFlow; + + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::FixedSizeList; + use vortex_tiled_fsl::TiledFixedSizeList; + use vortex_tiled_fsl::TiledFixedSizeListArrayExt; + + use super::TILED_FSL_SESSION; + use super::TiledFslAction; + use super::deterministic_tiled_fsl_cases; + use super::execute_action; + use super::fuzz; + use super::run_tiled_fsl; + use crate::error::VortexFuzzResult; + + #[test] + #[expect(clippy::result_large_err)] + fn empty_full_range_slice_continues_with_tiled_geometry() -> VortexFuzzResult<()> { + let mut cases = deterministic_tiled_fsl_cases()?; + let input = cases.remove(0); + let geometry = input.geometry; + let mut canonical = input.canonical; + let mut ctx = TILED_FSL_SESSION.create_execution_ctx(); + let mut tiled = fuzz(TiledFixedSizeList::encode( + canonical.as_::(), + geometry, + &mut ctx, + ))? + .into_array(); + + let control = execute_action( + TiledFslAction::Slice { start: 0, stop: 0 }, + &mut canonical, + &mut tiled, + geometry, + 0, + &mut ctx, + )?; + + assert_eq!(control, ControlFlow::Continue(())); + assert!(tiled.is::()); + assert_eq!(tiled.as_::().geometry(), geometry); + Ok(()) + } + + #[test] + #[expect(clippy::result_large_err)] + fn deterministic_tiled_fsl_smoke() -> VortexFuzzResult<()> { + let cases = deterministic_tiled_fsl_cases()?; + assert_eq!(cases.len(), 5); + for input in cases { + run_tiled_fsl(input)?; + } + Ok(()) + } +} diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 5992ff06434..a2fad09e46d 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -56,6 +56,7 @@ vortex-sequence = { workspace = true } vortex-session = { workspace = true } vortex-sparse = { workspace = true } vortex-tensor = { workspace = true, optional = true } +vortex-tiled-fsl = { workspace = true, optional = true } vortex-utils = { workspace = true, features = ["dashmap"] } vortex-zigzag = { workspace = true } vortex-zstd = { workspace = true, optional = true } @@ -83,6 +84,7 @@ zstd = ["dep:vortex-zstd", "vortex-btrblocks/zstd", "vortex-btrblocks/pco"] # This feature enables unstable encodings for which we don't guarantee stability. unstable_encodings = [ "dep:vortex-tensor", + "dep:vortex-tiled-fsl", "vortex-zstd?/unstable_encodings", "vortex-btrblocks/unstable_encodings", ] diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index e57b7ff344b..c7d5361d71a 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -196,6 +196,8 @@ pub fn register_default_encodings(session: &VortexSession) { #[cfg(feature = "unstable_encodings")] vortex_tensor::initialize(session); + #[cfg(feature = "unstable_encodings")] + vortex_tiled_fsl::initialize(session); } #[cfg(test)] diff --git a/vortex-file/tests/tiled_fsl.rs b/vortex-file/tests/tiled_fsl.rs new file mode 100644 index 00000000000..79471f74793 --- /dev/null +++ b/vortex-file/tests/tiled_fsl.rs @@ -0,0 +1,172 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![cfg(feature = "unstable_encodings")] +#![expect(clippy::tests_outside_test_module)] + +mod common; + +use std::num::NonZeroU32; +use std::sync::Arc; + +use common::enable_all_registered_array_encodings; +use futures::StreamExt; +use futures::pin_mut; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::FieldNames; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_fastlanes::bitpack_compress::bitpack_encode; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_io::session::RuntimeSession; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::session::LayoutSession; +use vortex_tiled_fsl::TileGeometry; +use vortex_tiled_fsl::TiledFixedSizeList; +use vortex_tiled_fsl::TiledFixedSizeListArrayExt; +use vortex_tiled_fsl::TiledFixedSizeListArraySlotsExt; + +const ROWS: usize = 65; +const DIMENSIONS: u32 = 128; + +fn geometry(rows: u32, dimensions: u32) -> VortexResult { + let rows = NonZeroU32::new(rows) + .ok_or_else(|| vortex_err!(InvalidArgument: "tile rows must be nonzero"))?; + let dimensions = NonZeroU32::new(dimensions) + .ok_or_else(|| vortex_err!(InvalidArgument: "tile dimensions must be nonzero"))?; + Ok(TileGeometry::new(rows, dimensions)) +} + +fn row_distinguishing_input() -> VortexResult { + let values = (0..ROWS) + .flat_map(|row| { + (0..DIMENSIONS as usize).map(move |dimension| { + let value = match dimension { + // The first two dimensions retain the complete row number while remaining + // 4-bit-safe. This makes a transpose or row-order regression observable. + 0 => row & 0x0f, + 1 => (row >> 4) & 0x0f, + _ => (row * 17 + dimension * 5) & 0x0f, + }; + u8::try_from(value).map_err(|error| { + vortex_err!(InvalidArgument: "four-bit fixture value is out of range: {error}") + }) + }) + }) + .collect::>>()?; + Ok(FixedSizeListArray::new( + PrimitiveArray::new( + Buffer::from(values), + Validity::from_iter((0..ROWS * DIMENSIONS as usize).map(|index| index % 13 != 0)), + ) + .into_array(), + DIMENSIONS, + Validity::from_iter((0..ROWS).map(|row| row % 7 != 0)), + ROWS, + )) +} + +#[tokio::test] +async fn unstable_tiled_fixed_size_lists_roundtrip_through_files() -> VortexResult<()> { + let session = vortex_array::array_session() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + enable_all_registered_array_encodings(&session); + + let mut ctx = session.create_execution_ctx(); + let canonical = row_distinguishing_input()?; + let expected_geometry = geometry(64, DIMENSIONS)?; + let raw = TiledFixedSizeList::encode(canonical.as_view(), expected_geometry, &mut ctx)?; + let physical = raw.elements().clone().execute::(&mut ctx)?; + let bitpacked = bitpack_encode(&physical, 4, None, &mut ctx)?.into_array(); + let bitpacked = TiledFixedSizeList::try_new( + bitpacked, + DIMENSIONS, + raw.array_validity(), + ROWS, + expected_geometry, + )?; + let input = StructArray::new( + FieldNames::from(["raw", "bitpacked"]), + vec![ + raw.into_array().slice(10..60)?, + bitpacked.into_array().slice(10..60)?, + ], + 50, + Validity::NonNullable, + ) + .into_array(); + + let mut bytes = Vec::new(); + session + .write_options() + .with_strategy(Arc::new(FlatLayoutStrategy::default())) + .write(&mut bytes, input.clone().to_array_stream()) + .await?; + + let file = session + .open_options() + .open_buffer(ByteBuffer::from(bytes))?; + let stream = file.scan()?.into_stream()?; + pin_mut!(stream); + let chunk = stream + .next() + .await + .ok_or_else(|| vortex_err!(InvalidArgument: "written file has no chunks"))??; + assert!(stream.next().await.is_none(), "one written chunk"); + + let result = chunk.execute::(&mut ctx)?; + let raw = result.unmasked_field(0).clone(); + let bitpacked = result.unmasked_field(1).clone(); + + assert!(raw.is::()); + assert!(bitpacked.is::()); + assert_eq!( + raw.as_::().geometry(), + expected_geometry + ); + assert_eq!( + bitpacked.as_::().geometry(), + expected_geometry + ); + assert_eq!(raw.as_::().row_offset(), 10); + assert_eq!(raw.as_::().backing_rows(), 64); + assert_eq!(bitpacked.as_::().row_offset(), 10); + assert_eq!(bitpacked.as_::().backing_rows(), 64); + assert!( + bitpacked + .as_::() + .elements() + .is::() + ); + assert!( + bitpacked + .as_::() + .elements() + .dtype() + .is_nullable() + ); + assert!( + bitpacked + .as_::() + .array_validity() + .mask_eq( + &Validity::from_iter((10..60).map(|row| row % 7 != 0)), + 50, + &mut ctx, + )? + ); + assert_arrays_eq!(input, result, &mut ctx); + Ok(()) +} diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 5b26944c4ed..d36e61b5430 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -51,6 +51,7 @@ vortex-sequence = { workspace = true } vortex-session = { workspace = true } vortex-sparse = { workspace = true } vortex-tensor = { workspace = true, optional = true } +vortex-tiled-fsl = { workspace = true, optional = true } vortex-utils = { workspace = true } vortex-zigzag = { workspace = true } vortex-zstd = { workspace = true, optional = true } @@ -90,6 +91,7 @@ serde = ["vortex-array/serde", "vortex-buffer/serde", "vortex-mask/serde"] # This feature enabled unstable encodings for which we don't guarantee stability. unstable_encodings = [ "dep:vortex-tensor", + "dep:vortex-tiled-fsl", "vortex-btrblocks/unstable_encodings", "vortex-file?/unstable_encodings", "vortex-zstd?/unstable_encodings", diff --git a/vortex/src/editions/mod.rs b/vortex/src/editions/mod.rs index a4deba81a53..d902ccd85c4 100644 --- a/vortex/src/editions/mod.rs +++ b/vortex/src/editions/mod.rs @@ -37,13 +37,14 @@ pub use self::unstable::UNSTABLE_2025_05_0; pub use self::unstable::UNSTABLE_2026_02_0; pub use self::unstable::UNSTABLE_2026_04_0; pub use self::unstable::UNSTABLE_2026_06_0; +pub use self::unstable::UNSTABLE_2026_07_0; /// The `core` edition enabled for writing by the default Vortex session. pub const DEFAULT_CORE_EDITION: EditionId = CORE_2026_07_0; /// The `unstable` edition enabled for writing by the default Vortex session when the /// `unstable_encodings` feature is selected. -pub const DEFAULT_UNSTABLE_EDITION: EditionId = UNSTABLE_2026_06_0; +pub const DEFAULT_UNSTABLE_EDITION: EditionId = UNSTABLE_2026_07_0; /// The first-party Vortex edition declarations. pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ @@ -55,6 +56,7 @@ pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ &unstable::v2026_02::DECLARATION, &unstable::v2026_04::DECLARATION, &unstable::v2026_06::DECLARATION, + &unstable::v2026_07::DECLARATION, ]; /// Register the Vortex edition declarations with the session's [`EditionSession`]. diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index 6968a6c1040..c5fde460483 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -161,6 +161,51 @@ fn default_session_enables_the_write_editions() { assert!(!enabled.contains(&DEFAULT_UNSTABLE_EDITION)); } +#[cfg(feature = "unstable_encodings")] +#[test] +fn default_unstable_edition_permits_tiled_fixed_size_list() { + use vortex_array::VTable as _; + use vortex_tiled_fsl::TiledFixedSizeList; + + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + assert!( + session + .enabled_encoding_ids() + .contains(&TiledFixedSizeList.id()) + ); +} + +#[cfg(feature = "unstable_encodings")] +#[test] +fn unstable_encoding_registration_does_not_depend_on_files() { + use vortex_array::VTable as _; + use vortex_array::dtype::extension::ExtVTable as _; + use vortex_array::dtype::session::DTypeSessionExt as _; + use vortex_array::session::ArraySessionExt as _; + use vortex_tensor::fixed_shape_tensor::FixedShapeTensor; + use vortex_tensor::vector::Vector; + use vortex_tiled_fsl::TiledFixedSizeList; + + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + assert!( + session + .arrays() + .registry() + .contains_key(&TiledFixedSizeList.id()) + ); + assert!(session.dtypes().registry().contains_key(&Vector.id())); + assert!( + session + .dtypes() + .registry() + .contains_key(&FixedShapeTensor.id()) + ); +} + #[test] fn core_edition_ids_are_registered_array_encodings() { use vortex_array::session::ArraySessionExt; diff --git a/vortex/src/editions/unstable/mod.rs b/vortex/src/editions/unstable/mod.rs index 5544ba45c0f..5a3907ade17 100644 --- a/vortex/src/editions/unstable/mod.rs +++ b/vortex/src/editions/unstable/mod.rs @@ -10,8 +10,10 @@ pub mod v2025_05; pub mod v2026_02; pub mod v2026_04; pub mod v2026_06; +pub mod v2026_07; pub use v2025_05::UNSTABLE_2025_05_0; pub use v2026_02::UNSTABLE_2026_02_0; pub use v2026_04::UNSTABLE_2026_04_0; pub use v2026_06::UNSTABLE_2026_06_0; +pub use v2026_07::UNSTABLE_2026_07_0; diff --git a/vortex/src/editions/unstable/v2026_07.rs b/vortex/src/editions/unstable/v2026_07.rs new file mode 100644 index 00000000000..c26c80c31d2 --- /dev/null +++ b/vortex/src/editions/unstable/v2026_07.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionId; + +pub const UNSTABLE_2026_07_0: EditionId = EditionId::new("unstable", 2026, 7, 0); + +pub static DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: UNSTABLE_2026_07_0, + min_vortex_version: None, + }, + added: &[&"vortex.tiled_fsl"], +}; diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 5abe779ab01..c0a7569ad5a 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -281,6 +281,12 @@ pub mod encodings { pub use vortex_sparse::*; } + #[cfg(feature = "unstable_encodings")] + /// Experimental two-dimensional tiled fixed-size-list encoding. + pub mod tiled_fsl { + pub use vortex_tiled_fsl::*; + } + /// Zig-zag integer transform encoding. pub mod zigzag { pub use vortex_zigzag::*; @@ -300,6 +306,15 @@ pub trait VortexSessionDefault { fn default() -> VortexSession; } +#[cfg(all( + feature = "unstable_encodings", + any(not(feature = "files"), target_arch = "wasm32") +))] +fn register_unstable_encodings(session: &VortexSession) { + vortex_tensor::initialize(session); + vortex_tiled_fsl::initialize(session); +} + impl VortexSessionDefault for VortexSession { fn default() -> VortexSession { let session = VortexSession::empty() @@ -315,6 +330,11 @@ impl VortexSessionDefault for VortexSession { vortex_arrow::initialize(&session); editions::register_default_editions(&session); editions::enable_default_editions(&session); + #[cfg(all( + feature = "unstable_encodings", + any(not(feature = "files"), target_arch = "wasm32") + ))] + register_unstable_encodings(&session); // `MultiFileSession` holds a `moka` cache whose clock reads `std::time::Instant::now()` // when constructed. `Instant` is unsupported on `wasm32` and panics with "time not