Test public PCO selection contracts without depending on optimized kernels

Signed-off-by: Daniel King <dan@spiraldb.com>

Author: danking

encodings/pco/src/lib.rs

@@ -62,3 +62,6 @@ pub struct PcoMetadata {
 
 #[cfg(test)]
 mod tests;
+
+#[cfg(test)]
+mod selection_contract_tests;

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::<Vec<_>>(),
+        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<u64>,
+    #[case] lengths: Vec<u64>,
+    #[case] strides: Vec<u64>,
+) -> 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::<Vec<_>>();
+    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::<PrimitiveArray>(&mut ctx)
+            .is_err()
+    );
+    Ok(())
+}