Scalar subtraction

vortex-array/Cargo.toml

@@ -52,3 +52,7 @@ criterion = { workspace = true }
 [[bench]]
 name = "search_sorted"
 harness = false
+
+[[bench]]
+name = "scalar_subtract"
+harness = false

vortex-array/benches/scalar_subtract.rs

@@ -0,0 +1,41 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use itertools::Itertools;
+use rand::distributions::Uniform;
+use rand::{thread_rng, Rng};
+use vortex::array::chunked::ChunkedArray;
+use vortex::IntoArray;
+use vortex_error::VortexError;
+
+fn scalar_subtract(c: &mut Criterion) {
+    let mut group = c.benchmark_group("scalar_subtract");
+
+    let mut rng = thread_rng();
+    let range = Uniform::new(0i64, 100_000_000);
+    let data1 = (0..10_000_000)
+        .map(|_| rng.sample(range))
+        .collect_vec()
+        .into_array();
+    let data2 = (0..10_000_000)
+        .map(|_| rng.sample(range))
+        .collect_vec()
+        .into_array();
+
+    let to_subtract = -1i64;
+
+    let chunked = ChunkedArray::from_iter([data1.clone(), data2]).into_array();
+
+    group.bench_function("vortex", |b| {
+        b.iter(|| {
+            let array =
+                vortex::compute::scalar_subtract::subtract_scalar(&chunked, &to_subtract.into())
+                    .unwrap();
+
+            let chunked = ChunkedArray::try_from(array).unwrap();
+            black_box(chunked);
+            Ok::<(), VortexError>(())
+        });
+    });
+}
+
+criterion_group!(benches, scalar_subtract);
+criterion_main!(benches);

vortex-array/src/array/chunked/compute/mod.rs

@@ -4,6 +4,7 @@ use vortex_scalar::Scalar;
 use crate::array::chunked::ChunkedArray;
 use crate::compute::as_contiguous::{as_contiguous, AsContiguousFn};
 use crate::compute::scalar_at::{scalar_at, ScalarAtFn};
+use crate::compute::scalar_subtract::SubtractScalarFn;
 use crate::compute::slice::SliceFn;
 use crate::compute::take::TakeFn;
 use crate::compute::ArrayCompute;
@@ -28,6 +29,10 @@ impl ArrayCompute for ChunkedArray<'_> {
     fn take(&self) -> Option<&dyn TakeFn> {
         Some(self)
     }
+
+    fn subtract_scalar(&self) -> Option<&dyn SubtractScalarFn> {
+        Some(self)
+    }
 }
 
 impl AsContiguousFn for ChunkedArray<'_> {

vortex-array/src/array/chunked/mod.rs

@@ -5,6 +5,7 @@ use vortex_error::{vortex_bail, VortexResult};
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::scalar_at::scalar_at;
+use crate::compute::scalar_subtract::{subtract_scalar, SubtractScalarFn};
 use crate::compute::search_sorted::{search_sorted, SearchSortedSide};
 use crate::validity::Validity::NonNullable;
 use crate::validity::{ArrayValidity, LogicalValidity};
@@ -139,14 +140,28 @@ impl ArrayValidity for ChunkedArray<'_> {
 
 impl EncodingCompression for ChunkedEncoding {}
 
+impl SubtractScalarFn for ChunkedArray<'_> {
+    fn subtract_scalar(&self, to_subtract: &Scalar) -> VortexResult<OwnedArray> {
+        self.chunks()
+            .map(|chunk| subtract_scalar(&chunk, to_subtract))
+            .collect::<VortexResult<Vec<_>>>()
+            .map(|chunks| {
+                ChunkedArray::try_new(chunks, self.dtype().clone())
+                    .expect("Subtraction on chunked array changed dtype")
+                    .into_array()
+            })
+    }
+}
+
 #[cfg(test)]
 mod test {
     use vortex_dtype::{DType, Nullability};
     use vortex_dtype::{NativePType, PType};
 
     use crate::array::chunked::{ChunkedArray, OwnedChunkedArray};
+    use crate::compute::scalar_subtract::subtract_scalar;
     use crate::compute::slice::slice;
-    use crate::{Array, IntoArray};
+    use crate::{Array, IntoArray, ToArray};
 
     fn chunked_array() -> OwnedChunkedArray {
         ChunkedArray::try_new(
@@ -197,4 +212,39 @@ mod test {
     pub fn slice_end() {
         assert_equal_slices(slice(chunked_array().array(), 7, 8).unwrap(), &[8u64]);
     }
+
+    #[test]
+    fn test_scalar_subtract() {
+        let chunked = chunked_array();
+        let to_subtract = 1u64;
+        let array = subtract_scalar(&chunked.to_array(), &to_subtract.into()).unwrap();
+
+        let chunked = ChunkedArray::try_from(array).unwrap();
+        let mut chunks_out = chunked.chunks();
+
+        let results = chunks_out
+            .next()
+            .unwrap()
+            .flatten_primitive()
+            .unwrap()
+            .typed_data::<u64>()
+            .to_vec();
+        assert_eq!(results, &[0u64, 1, 2]);
+        let results = chunks_out
+            .next()
+            .unwrap()
+            .flatten_primitive()
+            .unwrap()
+            .typed_data::<u64>()
+            .to_vec();
+        assert_eq!(results, &[3u64, 4, 5]);
+        let results = chunks_out
+            .next()
+            .unwrap()
+            .flatten_primitive()
+            .unwrap()
+            .typed_data::<u64>()
+            .to_vec();
+        assert_eq!(results, &[6u64, 7, 8]);
+    }
 }

vortex-array/src/array/primitive/compute/mod.rs

@@ -4,6 +4,7 @@ use crate::compute::as_contiguous::AsContiguousFn;
 use crate::compute::cast::CastFn;
 use crate::compute::fill::FillForwardFn;
 use crate::compute::scalar_at::ScalarAtFn;
+use crate::compute::scalar_subtract::SubtractScalarFn;
 use crate::compute::search_sorted::SearchSortedFn;
 use crate::compute::slice::SliceFn;
 use crate::compute::take::TakeFn;
@@ -16,6 +17,7 @@ mod fill;
 mod scalar_at;
 mod search_sorted;
 mod slice;
+mod subtract_scalar;
 mod take;
 
 impl ArrayCompute for PrimitiveArray<'_> {
@@ -39,6 +41,10 @@ impl ArrayCompute for PrimitiveArray<'_> {
         Some(self)
     }
 
+    fn subtract_scalar(&self) -> Option<&dyn SubtractScalarFn> {
+        Some(self)
+    }
+
     fn search_sorted(&self) -> Option<&dyn SearchSortedFn> {
         Some(self)
     }