From 0395567ca86aee75bfe3ac150efd0beed37e2b4f Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 13:51:13 -0400 Subject: [PATCH 1/8] Test public PCO selection contracts without depending on optimized kernels Signed-off-by: Daniel King --- encodings/pco/src/lib.rs | 3 + encodings/pco/src/selection_contract_tests.rs | 148 ++++++++++++++++++ 2 files changed, 151 insertions(+) create mode 100644 encodings/pco/src/selection_contract_tests.rs diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index df6581ec1d1..c153d9adef0 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -62,3 +62,6 @@ pub struct PcoMetadata { #[cfg(test)] mod tests; + +#[cfg(test)] +mod selection_contract_tests; diff --git a/encodings/pco/src/selection_contract_tests.rs b/encodings/pco/src/selection_contract_tests.rs new file mode 100644 index 00000000000..e3c8ef3c2ae --- /dev/null +++ b/encodings/pco/src/selection_contract_tests.rs @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public selection contracts, independent of optional PCO metadata or execution strategy. + +use rstest::rstest; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::Pco; + +#[rstest] +fn selected_values_and_nulls(#[values(false, true)] nullable: bool) -> VortexResult<()> { + let session = VortexSession::empty(); + session.arrays().register(Pco); + let mut ctx = session.create_execution_ctx(); + let input = PrimitiveArray::new( + (0..4096i64).map(|i| i * 19).collect::>(), + if nullable { + Validity::from_iter((0..4096).map(|i| i % 7 != 0)) + } else { + Validity::NonNullable + }, + ); + let compressed = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?.into_array(); + let source = compressed.slice(17..4001)?.slice(13..3900)?; + let expected = input.slice(30..3917)?; + let indices = PrimitiveArray::from_option_iter([ + Some(3000u64), + Some(0), + None, + Some(127), + Some(1), + Some(3000), + Some(128), + None, + ]) + .into_array(); + assert_arrays_eq!( + source.take(indices.clone())?, + expected.take(indices)?, + &mut ctx + ); + let mut probe = source.repeated_probe(); + for i in [3000, 0, 127, 1, 3000, 128, 0] { + assert_eq!( + probe.execute_scalar(i, &mut ctx)?, + expected.execute_scalar(i, &mut ctx)? + ); + } + assert!(probe.execute_scalar(source.len(), &mut ctx).is_err()); + Ok(()) +} + +#[rstest] +#[case(vec![0u64, 40, 0, 70], vec![15u64, 20, 3, 0], vec![1u64, 1, 0, 1])] +#[case(vec![100u64, 0, 20], vec![20u64, 20, 20], vec![1u64, 1, 2])] +fn compact_ranges( + #[case] starts: Vec, + #[case] lengths: Vec, + #[case] strides: Vec, +) -> VortexResult<()> { + let session = VortexSession::empty(); + let mut ctx = session.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..256i64).map(|i| (i % 5 != 0).then_some(i * i))); + let compressed = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let expected_indices = starts + .iter() + .zip(&lengths) + .zip(&strides) + .flat_map(|((&s, &n), &k)| (0..n).map(move |i| s + i * k)) + .collect::>(); + let indices = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(lengths).into_array(), + PrimitiveArray::from_iter(strides).into_array(), + expected_indices.len(), + )? + .into_array(); + assert_arrays_eq!( + compressed.take(indices)?, + input.take(PrimitiveArray::from_iter(expected_indices).into_array())?, + &mut ctx + ); + Ok(()) +} + +#[rstest] +fn null_runs_and_empty(#[values(false, true)] all_null: bool) -> VortexResult<()> { + let session = VortexSession::empty(); + let mut ctx = session.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter( + (0..65536i32).map(|i| (!all_null && !(100..65000).contains(&i)).then_some(i)), + ); + let compressed = Pco::from_primitive(input.as_view(), 3, 128, &mut ctx)?.into_array(); + let indices = + PrimitiveArray::from_option_iter([Some(0u64), None, Some(65001), Some(30000), Some(0)]) + .into_array(); + assert_arrays_eq!( + compressed.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + assert_arrays_eq!( + compressed.slice(99..65002)?, + input.slice(99..65002)?, + &mut ctx + ); + assert_arrays_eq!( + compressed.take(PrimitiveArray::from_iter([] as [u64; 0]).into_array())?, + input.take(PrimitiveArray::from_iter([] as [u64; 0]).into_array())?, + &mut ctx + ); + Ok(()) +} + +#[test] +fn invalid_indices_under_nulls_are_ignored() -> VortexResult<()> { + let session = VortexSession::empty(); + let mut ctx = session.create_execution_ctx(); + let input = PrimitiveArray::from_iter([11i32, 22, 33]); + let compressed = Pco::from_primitive(input.as_view(), 3, 2, &mut ctx)?.into_array(); + let indices = PrimitiveArray::new( + vec![0i64, i64::MIN, i64::MAX, 2], + Validity::from_iter([true, false, false, true]), + ) + .into_array(); + assert_arrays_eq!( + compressed.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + assert!( + compressed + .take(PrimitiveArray::from_iter([u64::MAX]).into_array())? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) +} From e6afd4102d75c848573568a4ca4f9143127ef082 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 14:04:01 -0400 Subject: [PATCH 2/8] Keep defined out-of-bounds rejection tests with the later Take guard change Signed-off-by: Daniel King --- encodings/pco/src/selection_contract_tests.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/encodings/pco/src/selection_contract_tests.rs b/encodings/pco/src/selection_contract_tests.rs index e3c8ef3c2ae..f68612b877e 100644 --- a/encodings/pco/src/selection_contract_tests.rs +++ b/encodings/pco/src/selection_contract_tests.rs @@ -138,11 +138,5 @@ fn invalid_indices_under_nulls_are_ignored() -> VortexResult<()> { input.take(indices)?, &mut ctx ); - assert!( - compressed - .take(PrimitiveArray::from_iter([u64::MAX]).into_array())? - .execute::(&mut ctx) - .is_err() - ); Ok(()) } From 03cec85eaab0c5716395ee6feb81ad8663f30aa5 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 14:16:37 -0400 Subject: [PATCH 3/8] Preserve the public compressed-validity selection oracle Signed-off-by: Daniel King --- vortex-file/tests/pco_validity_selection.rs | 58 +++++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 vortex-file/tests/pco_validity_selection.rs diff --git a/vortex-file/tests/pco_validity_selection.rs b/vortex-file/tests/pco_validity_selection.rs new file mode 100644 index 00000000000..2b82a07cb57 --- /dev/null +++ b/vortex-file/tests/pco_validity_selection.rs @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public selection oracle for a compressed logical validity child. + +#![expect(clippy::tests_outside_test_module)] + +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_pco::Pco; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +#[test] +fn compressed_validity_selection_oracle() -> VortexResult<()> { + let session = VortexSession::empty(); + vortex_file::register_default_encodings(&session); + let mut ctx = session.create_execution_ctx(); + let validity = Validity::Array( + RunEnd::try_new( + PrimitiveArray::from_iter([32768u64, 65536, 98304, 131072]).into_array(), + BoolArray::from_iter([false, true, false, true]).into_array(), + &mut ctx, + )? + .into_array(), + ); + let input = PrimitiveArray::new((0..131072i64).collect::>(), validity); + let expected = + PrimitiveArray::from_option_iter((0..131072i64).map(|i| (i / 32768 % 2 == 1).then_some(i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 127, &mut ctx)?.into_array(); + let indices = + PrimitiveArray::from_iter([0u32, 32767, 32768, 65535, 65536, 98304, 131071, 32768]) + .into_array(); + assert_arrays_eq!( + encoded.take(indices.clone())?, + expected.take(indices)?, + &mut ctx + ); + assert_arrays_eq!( + encoded.slice(30001..130001)?, + expected.slice(30001..130001)?, + &mut ctx + ); + let mut probe = encoded.repeated_probe(); + for index in [0, 32767, 32768, 65535, 65536, 98304, 131071, 32768] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + expected.execute_scalar(index, &mut ctx)? + ); + } + Ok(()) +} + From 1ba431e15e5c654d94278ead843d7172d8bf5a38 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 14:42:05 -0400 Subject: [PATCH 4/8] Import referenced modules at file scope and document page coordinates Signed-off-by: Daniel King --- vortex-file/tests/pco_validity_selection.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-file/tests/pco_validity_selection.rs b/vortex-file/tests/pco_validity_selection.rs index 2b82a07cb57..e3c8c8c5817 100644 --- a/vortex-file/tests/pco_validity_selection.rs +++ b/vortex-file/tests/pco_validity_selection.rs @@ -12,6 +12,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; use vortex_error::VortexResult; +use vortex_file::register_default_encodings; use vortex_pco::Pco; use vortex_runend::RunEnd; use vortex_session::VortexSession; @@ -19,7 +20,7 @@ use vortex_session::VortexSession; #[test] fn compressed_validity_selection_oracle() -> VortexResult<()> { let session = VortexSession::empty(); - vortex_file::register_default_encodings(&session); + register_default_encodings(&session); let mut ctx = session.create_execution_ctx(); let validity = Validity::Array( RunEnd::try_new( @@ -55,4 +56,3 @@ fn compressed_validity_selection_oracle() -> VortexResult<()> { } Ok(()) } - From 27800c47f2bba84fee9fb2d83e7aadb5d4c7f25e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 15:29:49 -0400 Subject: [PATCH 5/8] Use independent selection expectations and cover chunk boundaries Signed-off-by: Daniel King --- encodings/pco/src/selection_contract_tests.rs | 64 +++++++++++++++---- vortex-file/tests/pco_validity_selection.rs | 13 +++- 2 files changed, 61 insertions(+), 16 deletions(-) diff --git a/encodings/pco/src/selection_contract_tests.rs b/encodings/pco/src/selection_contract_tests.rs index f68612b877e..07d235cf53f 100644 --- a/encodings/pco/src/selection_contract_tests.rs +++ b/encodings/pco/src/selection_contract_tests.rs @@ -31,7 +31,6 @@ fn selected_values_and_nulls(#[values(false, true)] nullable: bool) -> VortexRes ); let compressed = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?.into_array(); let source = compressed.slice(17..4001)?.slice(13..3900)?; - let expected = input.slice(30..3917)?; let indices = PrimitiveArray::from_option_iter([ Some(3000u64), Some(0), @@ -41,18 +40,34 @@ fn selected_values_and_nulls(#[values(false, true)] nullable: bool) -> VortexRes Some(3000), Some(128), None, + Some(5), ]) .into_array(); - assert_arrays_eq!( - source.take(indices.clone())?, - expected.take(indices)?, - &mut ctx + let expected = PrimitiveArray::from_option_iter([ + Some(57570i64), + Some(570), + None, + Some(2983), + Some(589), + Some(57570), + Some(3002), + None, + if nullable { None } else { Some(665) }, + ]); + assert_arrays_eq!(source.take(indices)?, expected, &mut ctx); + let expected_probes = PrimitiveArray::new( + vec![57570i64, 570, 2983, 589, 57570, 3002, 570, 665], + if nullable { + Validity::from_iter([true, true, true, true, true, true, true, false]) + } else { + Validity::NonNullable + }, ); let mut probe = source.repeated_probe(); - for i in [3000, 0, 127, 1, 3000, 128, 0] { + for (position, index) in [3000, 0, 127, 1, 3000, 128, 0, 5].into_iter().enumerate() { assert_eq!( - probe.execute_scalar(i, &mut ctx)?, - expected.execute_scalar(i, &mut ctx)? + probe.execute_scalar(index, &mut ctx)?, + expected_probes.execute_scalar(position, &mut ctx)?, ); } assert!(probe.execute_scalar(source.len(), &mut ctx).is_err()); @@ -87,7 +102,9 @@ fn compact_ranges( .into_array(); assert_arrays_eq!( compressed.take(indices)?, - input.take(PrimitiveArray::from_iter(expected_indices).into_array())?, + PrimitiveArray::from_option_iter(expected_indices.into_iter().map(|i| { + (i % 5 != 0).then(|| i64::try_from(i * i).unwrap()) + })), &mut ctx ); Ok(()) @@ -105,8 +122,12 @@ fn null_runs_and_empty(#[values(false, true)] all_null: bool) -> VortexResult<() PrimitiveArray::from_option_iter([Some(0u64), None, Some(65001), Some(30000), Some(0)]) .into_array(); assert_arrays_eq!( - compressed.take(indices.clone())?, - input.take(indices)?, + compressed.take(indices)?, + PrimitiveArray::from_option_iter(if all_null { + [None, None, None, None, None] + } else { + [Some(0i32), None, Some(65001), None, Some(0)] + }), &mut ctx ); assert_arrays_eq!( @@ -116,7 +137,7 @@ fn null_runs_and_empty(#[values(false, true)] all_null: bool) -> VortexResult<() ); assert_arrays_eq!( compressed.take(PrimitiveArray::from_iter([] as [u64; 0]).into_array())?, - input.take(PrimitiveArray::from_iter([] as [u64; 0]).into_array())?, + PrimitiveArray::from_option_iter([] as [Option; 0]), &mut ctx ); Ok(()) @@ -134,8 +155,23 @@ fn invalid_indices_under_nulls_are_ignored() -> VortexResult<()> { ) .into_array(); assert_arrays_eq!( - compressed.take(indices.clone())?, - input.take(indices)?, + compressed.take(indices)?, + PrimitiveArray::from_option_iter([Some(11i32), None, None, Some(33)]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn nonnullable_selection_crosses_chunks() -> VortexResult<()> { + let session = VortexSession::empty(); + let mut ctx = session.create_execution_ctx(); + let input = PrimitiveArray::from_iter((0..300000i64).map(|i| i * 3 + 1)); + let compressed = Pco::from_primitive(input.as_view(), 3, 127, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_iter([299999u64, 262144, 0, 262143, 262144]).into_array(); + assert_arrays_eq!( + compressed.take(indices)?, + PrimitiveArray::from_iter([899998i64, 786433, 1, 786430, 786433]), &mut ctx ); Ok(()) diff --git a/vortex-file/tests/pco_validity_selection.rs b/vortex-file/tests/pco_validity_selection.rs index e3c8c8c5817..4bbe881288a 100644 --- a/vortex-file/tests/pco_validity_selection.rs +++ b/vortex-file/tests/pco_validity_selection.rs @@ -38,8 +38,17 @@ fn compressed_validity_selection_oracle() -> VortexResult<()> { PrimitiveArray::from_iter([0u32, 32767, 32768, 65535, 65536, 98304, 131071, 32768]) .into_array(); assert_arrays_eq!( - encoded.take(indices.clone())?, - expected.take(indices)?, + encoded.take(indices)?, + PrimitiveArray::from_option_iter([ + None, + None, + Some(32768i64), + Some(65535), + None, + Some(98304), + Some(131071), + Some(32768), + ]), &mut ctx ); assert_arrays_eq!( From 39dadb2f3a89b53f01a59ca9745d170eb34abf97 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 16:03:59 -0400 Subject: [PATCH 6/8] Apply pinned upstream Rust formatting Signed-off-by: Daniel King --- encodings/pco/src/selection_contract_tests.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/encodings/pco/src/selection_contract_tests.rs b/encodings/pco/src/selection_contract_tests.rs index 07d235cf53f..7d97732c0cc 100644 --- a/encodings/pco/src/selection_contract_tests.rs +++ b/encodings/pco/src/selection_contract_tests.rs @@ -102,9 +102,11 @@ fn compact_ranges( .into_array(); assert_arrays_eq!( compressed.take(indices)?, - PrimitiveArray::from_option_iter(expected_indices.into_iter().map(|i| { - (i % 5 != 0).then(|| i64::try_from(i * i).unwrap()) - })), + PrimitiveArray::from_option_iter( + expected_indices + .into_iter() + .map(|i| { (i % 5 != 0).then(|| i64::try_from(i * i).unwrap()) }) + ), &mut ctx ); Ok(()) From 08c48e3681c5c111bf7b67b742d0e6d368223d8e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 16:57:50 -0400 Subject: [PATCH 7/8] Move public PCO selection semantics ahead of optimization features Signed-off-by: Daniel King --- encodings/pco/src/lib.rs | 6 + encodings/pco/src/probe/tests.rs | 20 ++ encodings/pco/src/range_semantics_tests.rs | 211 +++++++++++++++ encodings/pco/src/selection_contract_tests.rs | 1 - .../pco/src/selection_semantics_tests.rs | 240 ++++++++++++++++++ vortex-file/tests/pco_file_compatibility.rs | 90 +++++++ 6 files changed, 567 insertions(+), 1 deletion(-) create mode 100644 encodings/pco/src/range_semantics_tests.rs create mode 100644 encodings/pco/src/selection_semantics_tests.rs create mode 100644 vortex-file/tests/pco_file_compatibility.rs diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index c153d9adef0..982f6dffc34 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -65,3 +65,9 @@ mod tests; #[cfg(test)] mod selection_contract_tests; + +#[cfg(test)] +mod selection_semantics_tests; + +#[cfg(test)] +mod range_semantics_tests; diff --git a/encodings/pco/src/probe/tests.rs b/encodings/pco/src/probe/tests.rs index c5386f1aa00..bf06069e2af 100644 --- a/encodings/pco/src/probe/tests.rs +++ b/encodings/pco/src/probe/tests.rs @@ -110,3 +110,23 @@ fn probe_outlives_source_and_moves() -> VortexResult<()> { assert!(moved.0.execute_scalar(4096, &mut ctx).is_err()); Ok(()) } + +#[test] +fn retained_all_null_probes_preserve_sliced_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::new(vec![0i64; 100000], Validity::AllInvalid); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let sliced = encoded.slice(17..145)?; + let mut probe = sliced.repeated_probe(); + for index in [0, 127, 64, 0, 127] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + vortex_array::scalar::Scalar::null(input.dtype().clone()) + ); + } + assert_eq!( + probe.execute_scalar(0, &mut ctx)?, + vortex_array::scalar::Scalar::null(input.dtype().clone()) + ); + Ok(()) +} diff --git a/encodings/pco/src/range_semantics_tests.rs b/encodings/pco/src/range_semantics_tests.rs new file mode 100644 index 00000000000..841fa9e76ba --- /dev/null +++ b/encodings/pco/src/range_semantics_tests.rs @@ -0,0 +1,211 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public slice and compact-range behavior, independent of compressed metadata. + +use std::iter::repeat_n; +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::Pco; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + session.arrays().register(Pco); + session +}); + +#[rstest] +#[case(vec![0u64, 40, 0, 70], vec![15u64, 20, 3, 0], vec![1u64, 1, 0, 1])] +#[case(vec![100u64, 0, 20], vec![20u64, 20, 20], vec![1u64, 1, 2])] +#[case(vec![99u64, 0], vec![0u64, 0], vec![u64::MAX, 0])] +fn piecewise_ranges( + #[case] starts: Vec, + #[case] lengths: Vec, + #[case] strides: Vec, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..256i64).map(|i| (i % 5 != 0).then_some(i * i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let mut expected_indices = Vec::new(); + for ((&start, &length), &stride) in starts.iter().zip(&lengths).zip(&strides) { + expected_indices.extend((0..length).map(|i| start + i * stride)); + } + let indices = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(lengths).into_array(), + PrimitiveArray::from_iter(strides).into_array(), + expected_indices.len(), + )? + .into_array(); + assert_arrays_eq!( + encoded.take(indices)?, + PrimitiveArray::from_option_iter( + expected_indices + .into_iter() + .map(|row| { (row % 5 != 0).then_some((row * row) as i64) }) + ), + &mut ctx + ); + Ok(()) +} + +#[test] +fn sliced_piecewise_coordinates_are_visible_rows() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..256i64).map(|i| (i % 5 != 0).then_some(i * i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let indices = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([0u64, 40, 3]).into_array(), + PrimitiveArray::from_iter([15u64, 20, 4]).into_array(), + PrimitiveArray::from_iter([1u64, 1, 0]).into_array(), + 39, + )? + .into_array(); + let expected = + PrimitiveArray::from_option_iter((0..15).chain(40..60).chain(repeat_n(3, 4)).map(|i| { + let row = i + 17i64; + (row % 5 != 0).then_some(row * row) + })); + assert_arrays_eq!(encoded.slice(17..200)?.take(indices)?, expected, &mut ctx); + Ok(()) +} + +#[test] +fn large_public_compressor_selections() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let len = 262_401usize; + let input = PrimitiveArray::from_iter( + (0..len).map(|i| u32::try_from(i).vortex_expect("fixture fits u32")), + ); + let encoded = Pco::from_primitive(input.as_view(), 0, 4096, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_iter([ + 0u32, + 900, + 260, + 3000, + u32::try_from(len - 1).vortex_expect("fixture fits u32"), + 4095, + 900, + ]) + .into_array(); + assert_arrays_eq!( + encoded.take(indices)?, + PrimitiveArray::from_iter([0u32, 900, 260, 3000, 262400, 4095, 900]), + &mut ctx + ); + assert_arrays_eq!( + encoded.slice(len - 1000..len - 1)?, + PrimitiveArray::from_iter((len - 1000..len - 1).map(|row| row as u32)), + &mut ctx + ); + Ok(()) +} + +#[test] +fn mixed_fragments_narrow_slices_and_compact_repeats() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter( + (0..32768i64).map(|row| (row % 5 != 0).then_some(row * 11)), + ); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let mixed = PrimitiveArray::from_iter( + (0..4096u64).map(|i| if i < 2 { i } else { (i * 104729) % 32768 }), + ) + .into_array(); + assert_arrays_eq!( + encoded.take(mixed)?, + PrimitiveArray::from_option_iter((0..4096u64).map(|i| { + let row = if i < 2 { i } else { (i * 104729) % 32768 }; + (row % 5 != 0).then_some(row as i64 * 11) + })), + &mut ctx + ); + let repeated = + PrimitiveArray::from_iter((0..100000u64).map(|i| if i % 2 == 0 { 0 } else { 127 })) + .into_array(); + assert_arrays_eq!( + encoded.slice(30000..30128)?.take(repeated)?, + PrimitiveArray::from_option_iter((0..100000).map(|i| { + if i % 2 == 0 { + None + } else { + Some(30127i64 * 11) + } + })), + &mut ctx + ); + let compact = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([1u64, 5]).into_array(), + PrimitiveArray::from_iter([500000u64, 500000]).into_array(), + PrimitiveArray::from_iter([0u64, 0]).into_array(), + 1000000, + )? + .into_array(); + let expected = PrimitiveArray::from_option_iter( + repeat_n(Some(11i64), 500000).chain(repeat_n(None, 500000)), + ); + assert_arrays_eq!(encoded.take(compact)?, expected, &mut ctx); + let nulls = PrimitiveArray::new(vec![0i64; 100000], Validity::AllInvalid); + let encoded_nulls = Pco::from_primitive(nulls.as_view(), 3, 1024, &mut ctx)?.into_array(); + let two = PrimitiveArray::from_iter([0u64, 99999]).into_array(); + assert_arrays_eq!( + encoded_nulls.take(two)?, + PrimitiveArray::from_option_iter([None::, None]), + &mut ctx + ); + Ok(()) +} + +#[rstest] +fn sliced_values_preserve_float_bits_and_validity( + #[values(false, true)] dense: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let bits = [ + 0u64, + 1u64 << 63, + 0x7ff8000000000042, + 0x7ff0000000000000, + 0xfff0000000000000, + 0x3ff4000000000000, + ]; + let values = (0..2048) + .map(|row| f64::from_bits(bits[row % bits.len()])) + .collect::>(); + let validity = if dense { + Validity::NonNullable + } else { + Validity::from_iter((0..2048).map(|row| row % 7 != 0)) + }; + let input = PrimitiveArray::new(values.clone(), validity); + let encoded = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?.into_array(); + for range in [1..2047, 89..129, 1000..1001] { + let actual = encoded + .slice(range.clone())? + .execute::(&mut ctx)?; + assert_eq!(actual.dtype(), input.dtype()); + let mask = actual.validity()?.execute_mask(actual.len(), &mut ctx)?; + for (offset, &value) in actual.as_slice::().iter().enumerate() { + let row = range.start + offset; + assert_eq!(mask.value(offset), dense || row % 7 != 0); + if mask.value(offset) { + assert_eq!(value.to_bits(), values[row].to_bits()); + } + } + } + Ok(()) +} diff --git a/encodings/pco/src/selection_contract_tests.rs b/encodings/pco/src/selection_contract_tests.rs index 7d97732c0cc..bba142d0672 100644 --- a/encodings/pco/src/selection_contract_tests.rs +++ b/encodings/pco/src/selection_contract_tests.rs @@ -70,7 +70,6 @@ fn selected_values_and_nulls(#[values(false, true)] nullable: bool) -> VortexRes expected_probes.execute_scalar(position, &mut ctx)?, ); } - assert!(probe.execute_scalar(source.len(), &mut ctx).is_err()); Ok(()) } diff --git a/encodings/pco/src/selection_semantics_tests.rs b/encodings/pco/src/selection_semantics_tests.rs new file mode 100644 index 00000000000..159c52559d3 --- /dev/null +++ b/encodings/pco/src/selection_semantics_tests.rs @@ -0,0 +1,240 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public take behavior, independent of compressed metadata and execution strategy. + +use std::iter::empty; +use std::sync::LazyLock; + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::Pco; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + session.arrays().register(Pco); + session +}); + +#[test] +fn reordered_nullable_take() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..4096i32).map(|i| (i % 7 != 0).then_some(i * 19))); + let encoded = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?.into_array(); + let source = encoded.slice(17..4001)?.slice(13..3900)?; + let expected_probes = PrimitiveArray::from_option_iter([ + Some(57570i32), + Some(570), + Some(2983), + Some(589), + Some(57570), + Some(3002), + Some(19570), + Some(570), + ]); + let mut probe = source.repeated_probe(); + for (position, index) in [3000, 0, 127, 1, 3000, 128, 1000, 0] + .into_iter() + .enumerate() + { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + expected_probes.execute_scalar(position, &mut ctx)? + ); + } + let indices = PrimitiveArray::from_option_iter([ + Some(3000u32), + Some(0), + None, + Some(127), + Some(1), + Some(3000), + Some(128), + None, + Some(1000), + ]) + .into_array(); + assert_arrays_eq!( + source.take(indices)?, + PrimitiveArray::from_option_iter([ + Some(57570i32), + Some(570), + None, + Some(2983), + Some(589), + Some(57570), + Some(3002), + None, + Some(19570), + ]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn all_null_and_long_null_runs() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + for input in [ + PrimitiveArray::new(vec![0i32; 1024], Validity::AllInvalid), + PrimitiveArray::from_option_iter( + (0..100_000).map(|i| [1, 90_000, 99_990].contains(&i).then_some(i)), + ), + ] { + let encoded = Pco::from_primitive(input.as_view(), 3, 1, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_iter([0u32, 1, 100, 1000, 2, 1]).into_array(); + assert_arrays_eq!( + encoded.take(indices)?, + PrimitiveArray::from_option_iter([0i32, 1, 100, 1000, 2, 1].map(|row| { + if input.len() == 1024 { + None + } else { + [1, 90_000, 99_990].contains(&row).then_some(row) + } + })), + &mut ctx + ); + assert_arrays_eq!( + encoded.slice(1..1000)?, + PrimitiveArray::from_option_iter((1..1000).map(|row| { + if input.len() == 1024 { + None + } else { + (row == 1).then_some(row) + } + })), + &mut ctx + ); + } + Ok(()) +} + +#[test] +fn registry_only_session_preserves_take_semantics() -> VortexResult<()> { + let session = VortexSession::empty(); + session.arrays().register(Pco); + let mut ctx = session.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter([Some(5i32), None, Some(9), Some(2)]); + let encoded = Pco::from_primitive(input.as_view(), 3, 1, &mut ctx)?.into_array(); + let indices: ArrayRef = PrimitiveArray::from_iter([3u32, 1, 0, 3]).into_array(); + assert_arrays_eq!( + encoded.take(indices)?, + PrimitiveArray::from_option_iter([Some(2i32), None, Some(5), Some(2)]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn repeated_scattered_requests_preserve_order_and_nulls() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..32768i64) + .map(|row| (row % 7 != 0).then_some(row * 13)) + .collect::>(); + let input = PrimitiveArray::from_option_iter(values.iter().copied()); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let requests = (0..1024usize) + .map(|i| (i % 11 != 0).then_some((i * 97) % 16)) + .collect::>(); + let indices = PrimitiveArray::from_option_iter( + requests + .iter() + .map(|row| row.map(|row| u64::try_from(row).vortex_expect("small index"))), + ) + .into_array(); + let expected = PrimitiveArray::from_option_iter( + requests.iter().map(|row| row.and_then(|row| values[row])), + ); + assert_arrays_eq!(encoded.take(indices)?, expected, &mut ctx); + Ok(()) +} + +#[test] +fn empty_take_dtype_and_mostly_null_request_storage() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..32768i64); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + for source in [encoded.clone(), encoded.slice(0..0)?] { + for indices in [ + PrimitiveArray::from_iter(empty::()).into_array(), + PrimitiveArray::from_option_iter(empty::>()).into_array(), + ] { + let result = source + .take(indices.clone())? + .execute::(&mut ctx)?; + let expected = if indices.dtype().is_nullable() { + PrimitiveArray::from_option_iter(empty::>()) + } else { + PrimitiveArray::from_iter(empty::()) + }; + assert_eq!(result.dtype(), expected.dtype()); + assert_arrays_eq!(result, expected, &mut ctx); + } + } + let indices = + PrimitiveArray::from_option_iter((0..200000).map(|i| (i == 12345).then_some(1u64))) + .into_array(); + let result = encoded.take(indices)?.execute::(&mut ctx)?; + let expected = + PrimitiveArray::from_option_iter((0..200000).map(|i| (i == 12345).then_some(1i64))); + assert_eq!(result.dtype(), expected.dtype()); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) +} + +#[test] +fn logical_null_indices_ignore_invalid_physical_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..128i64); + let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let indices = PrimitiveArray::new( + vec![i64::MIN, 4, i64::MAX], + Validity::from_iter([false, true, false]), + ) + .into_array(); + let expected = PrimitiveArray::from_option_iter([None, Some(4i64), None]); + assert_arrays_eq!(encoded.take(indices.clone())?, expected, &mut ctx); + // Reuse the same immutable request to check repeated execution. + assert_arrays_eq!( + encoded.take(indices)?, + PrimitiveArray::from_option_iter([None, Some(4i64), None]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn contiguous_take_preserves_nullable_index_dtype() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..32768i64); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_option_iter((1..32767u64).map(Some)).into_array(); + let expected = PrimitiveArray::from_option_iter((1..32767i64).map(Some)); + let result = encoded.take(indices)?.execute::(&mut ctx)?; + assert_eq!(result.dtype(), expected.dtype()); + assert_arrays_eq!(result, expected, &mut ctx); + let range = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([1u64]).into_array(), + PrimitiveArray::from_iter([32766u64]).into_array(), + PrimitiveArray::from_iter([1u64]).into_array(), + 32766, + )? + .into_array(); + assert_arrays_eq!( + encoded.take(range)?, + PrimitiveArray::from_iter(1..32767i64), + &mut ctx + ); + Ok(()) +} diff --git a/vortex-file/tests/pco_file_compatibility.rs b/vortex-file/tests/pco_file_compatibility.rs new file mode 100644 index 00000000000..9549459df89 --- /dev/null +++ b/vortex-file/tests/pco_file_compatibility.rs @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public file harness for an explicit bounded two-version compatibility check. + +#![expect(clippy::tests_outside_test_module)] + +use std::env::var; +use std::env::var_os; +use std::fs::create_dir_all; +use std::fs::read; +use std::fs::write; +use std::path::PathBuf; +use std::sync::Arc; + +use futures::TryStreamExt; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::assert_arrays_eq; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_file::register_default_encodings; +use vortex_io::session::RuntimeSession; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::session::LayoutSession; +use vortex_pco::Pco; +use vortex_session::VortexSession; + +mod common; + +#[tokio::test] +async fn public_pco_compatibility() -> VortexResult<()> { + let session = VortexSession::empty() + .with::() + .with::(); + register_default_encodings(&session); + common::enable_all_registered_array_encodings(&session); + let mut ctx = session.create_execution_ctx(); + let original = PrimitiveArray::from_option_iter( + (0..131_072i64).map(|i| (!(1024..=120_000).contains(&i) && i % 7 != 0).then_some(i * 19)), + ); + let mode = var("PCO_COMPAT_MODE").unwrap_or_else(|_| "roundtrip".to_owned()); + let directory = var_os("PCO_COMPAT_DIR").map(PathBuf::from); + for (name, range) in [("unsliced", 0..131_072), ("normalized_slice", 777..125_000)] { + let encoded = Pco::from_primitive(original.as_view(), 3, 127, &mut ctx)?.into_array(); + // Explicit writer-owned normalization, never a hidden serde fallback: execute selected + // pages and recompress only the chosen logical rows using the caller's configuration. + let selected = encoded + .slice(range.clone())? + .execute::(&mut ctx)?; + let normalized = Pco::from_primitive(selected.as_view(), 3, 127, &mut ctx)?.into_array(); + let expected = + StructArray::from_fields(&[("values", original.slice(range)?)])?.into_array(); + let data = StructArray::from_fields(&[("values", normalized)])?.into_array(); + let path = directory + .as_ref() + .map(|directory| directory.join(format!("{name}.vortex"))); + let bytes = if mode == "read" { + read( + path.as_ref() + .vortex_expect("read mode requires PCO_COMPAT_DIR"), + )? + } else { + let mut bytes = Vec::new(); + session + .write_options() + .with_strategy(Arc::new(FlatLayoutStrategy::default())) + .write(&mut bytes, data.to_array_stream()) + .await?; + if let Some(path) = &path { + create_dir_all(path.parent().vortex_expect("fixture path has a parent"))?; + write(path, &bytes)?; + } + bytes + }; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(bytes))?; + let chunks = file.scan()?.into_stream()?.try_collect::>().await?; + let actual = ChunkedArray::try_new(chunks, expected.dtype().clone())?.into_array(); + assert_arrays_eq!(actual, expected, &mut ctx); + } + Ok(()) +} From f56bcebfbeda930bd1619e09a79744282b12036c Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 17:05:24 -0400 Subject: [PATCH 8/8] Use checked integer conversion in the public large-input oracle Signed-off-by: Daniel King --- encodings/pco/src/range_semantics_tests.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/encodings/pco/src/range_semantics_tests.rs b/encodings/pco/src/range_semantics_tests.rs index 841fa9e76ba..d482f2b7474 100644 --- a/encodings/pco/src/range_semantics_tests.rs +++ b/encodings/pco/src/range_semantics_tests.rs @@ -109,7 +109,9 @@ fn large_public_compressor_selections() -> VortexResult<()> { ); assert_arrays_eq!( encoded.slice(len - 1000..len - 1)?, - PrimitiveArray::from_iter((len - 1000..len - 1).map(|row| row as u32)), + PrimitiveArray::from_iter( + (len - 1000..len - 1).map(|row| u32::try_from(row).vortex_expect("fixture fits u32")) + ), &mut ctx ); Ok(())