Array2: Migrate Roaring boolean and int

Cargo.toml

@@ -14,7 +14,7 @@ members = [
     "vortex-flatbuffers",
     "vortex-ipc",
     "vortex-ree",
-    #"vortex-roaring",
+    "vortex-roaring",
     "vortex-schema",
     #"vortex-zigzag",
 ]

bench-vortex/Cargo.toml

@@ -16,7 +16,7 @@ workspace = true
 
 [dependencies]
 #vortex-alp = { path = "../vortex-alp" }
-#vortex-roaring = { path = "../vortex-roaring" }
+vortex-roaring = { path = "../vortex-roaring" }
 #vortex-zigzag = { path = "../vortex-zigzag" }
 arrow = { workspace = true }
 arrow-array = { workspace = true }

bench-vortex/src/lib.rs

@@ -20,6 +20,7 @@ use vortex_datetime_parts::DateTimePartsEncoding;
 use vortex_dict::DictEncoding;
 use vortex_fastlanes::{BitPackedEncoding, FoREncoding};
 use vortex_ree::REEEncoding;
+use vortex_roaring::RoaringBoolEncoding;
 use vortex_schema::DType;
 
 use crate::data_downloads::FileType;
@@ -114,8 +115,8 @@ pub fn enumerate_arrays() -> Vec<EncodingRef> {
         &DateTimePartsEncoding,
         // &DeltaEncoding,  Blows up the search space too much.
         &REEEncoding,
-        //&RoaringBoolEncoding,
-        // RoaringIntEncoding,
+        &RoaringBoolEncoding,
+        // &RoaringIntEncoding,
         // Doesn't offer anything more than FoR really
         // ZigZagEncoding,
     ]

pyvortex/Cargo.toml

@@ -26,7 +26,7 @@ vortex-dict = { path = "../vortex-dict" }
 vortex-error = { path = "../vortex-error" }
 vortex-fastlanes = { path = "../vortex-fastlanes" }
 vortex-ree = { path = "../vortex-ree" }
-#vortex-roaring = { path = "../vortex-roaring" }
+vortex-roaring = { path = "../vortex-roaring" }
 vortex-schema = { path = "../vortex-schema" }
 #vortex-zigzag = { path = "../vortex-zigzag" }
 itertools = { workspace = true }

pyvortex/src/array.rs

@@ -23,6 +23,10 @@ use vortex_fastlanes::{
     FoREncoding, OwnedBitPackedArray, OwnedDeltaArray, OwnedFoRArray,
 };
 use vortex_ree::{OwnedREEArray, REEArray, REEEncoding, REE};
+use vortex_roaring::{
+    OwnedRoaringBoolArray, OwnedRoaringIntArray, RoaringBool, RoaringBoolArray,
+    RoaringBoolEncoding, RoaringInt, RoaringIntArray, RoaringIntEncoding,
+};
 
 use crate::dtype::PyDType;
 use crate::error::PyVortexError;
@@ -74,8 +78,8 @@ pyarray!(FoREncoding, FoRArray, "FoRArray");
 pyarray!(DeltaEncoding, DeltaArray, "DeltaArray");
 pyarray!(DictEncoding, DictArray, "DictArray");
 pyarray!(REEEncoding, REEArray, "REEArray");
-// pyarray!(RoaringBoolEncoding, RoaringBoolArray, "RoaringBoolArray");
-// pyarray!(RoaringIntEncoding, RoaringIntArray, "RoaringIntArray");
+pyarray!(RoaringBoolEncoding, RoaringBoolArray, "RoaringBoolArray");
+pyarray!(RoaringIntEncoding, RoaringIntArray, "RoaringIntArray");
 // pyarray!(ZigZagEncoding, ZigZagArray, "ZigZagArray");
 
 impl PyArray {
@@ -162,6 +166,18 @@ impl PyArray {
                 OwnedALPArray::try_from(inner.into_array()).map_err(PyVortexError::map_err)?,
             )?
             .extract(py),
+            RoaringBool::ID => PyBitPackedArray::wrap(
+                py,
+                OwnedBitPackedArray::try_from(inner.into_array())
+                    .map_err(PyVortexError::map_err)?,
+            )?
+            .extract(py),
+            RoaringInt::ID => PyBitPackedArray::wrap(
+                py,
+                OwnedBitPackedArray::try_from(inner.into_array())
+                    .map_err(PyVortexError::map_err)?,
+            )?
+            .extract(py),
             _ => Py::new(
                 py,
                 Self {

pyvortex/src/lib.rs

@@ -40,8 +40,8 @@ fn _lib(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PyFoRArray>()?;
     m.add_class::<PyPrimitiveArray>()?;
     m.add_class::<PyREEArray>()?;
-    // m.add_class::<PyRoaringBoolArray>()?;
-    // m.add_class::<PyRoaringIntArray>()?;
+    m.add_class::<PyRoaringBoolArray>()?;
+    m.add_class::<PyRoaringIntArray>()?;
     m.add_class::<PySparseArray>()?;
     m.add_class::<PyStructArray>()?;
     m.add_class::<PyVarBinArray>()?;

vortex-array/src/array/bool/compute/as_contiguous.rs

@@ -4,10 +4,10 @@ use vortex_error::VortexResult;
 use crate::array::bool::BoolArray;
 use crate::compute::as_contiguous::AsContiguousFn;
 use crate::validity::Validity;
-use crate::{Array, ArrayDType, IntoArray};
+use crate::{Array, ArrayDType, IntoArray, OwnedArray};
 
 impl AsContiguousFn for BoolArray<'_> {
-    fn as_contiguous(&self, arrays: &[Array]) -> VortexResult<Array<'static>> {
+    fn as_contiguous(&self, arrays: &[Array]) -> VortexResult<OwnedArray> {
         let validity = if self.dtype().is_nullable() {
             Validity::from_iter(arrays.iter().map(|a| a.with_dyn(|a| a.logical_validity())))
         } else {

vortex-array/src/array/bool/compute/fill.rs

@@ -4,10 +4,10 @@ use vortex_schema::Nullability;
 use crate::array::bool::BoolArray;
 use crate::compute::fill::FillForwardFn;
 use crate::validity::ArrayValidity;
-use crate::{Array, ArrayDType, IntoArray, ToArrayData};
+use crate::{ArrayDType, IntoArray, OwnedArray, ToArrayData};
 
 impl FillForwardFn for BoolArray<'_> {
-    fn fill_forward(&self) -> VortexResult<Array<'static>> {
+    fn fill_forward(&self) -> VortexResult<OwnedArray> {
         if self.dtype().nullability() == Nullability::NonNullable {
             return Ok(self.to_array_data().into_array());
         }

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

@@ -8,7 +8,7 @@ use crate::buffer::Buffer;
 use crate::ptype::{NativePType, PType};
 use crate::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
 use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
-use crate::{impl_encoding, ArrayDType};
+use crate::{impl_encoding, ArrayDType, OwnedArray};
 use crate::{match_each_native_ptype, ArrayFlatten};
 
 mod accessor;
@@ -129,7 +129,7 @@ impl<T: NativePType> From<Vec<T>> for PrimitiveArray<'_> {
 }
 
 impl<T: NativePType> IntoArray<'static> for Vec<T> {
-    fn into_array(self) -> Array<'static> {
+    fn into_array(self) -> OwnedArray {
         PrimitiveArray::from(self).into_array()
     }
 }

vortex-array/src/compute/as_contiguous.rs

@@ -1,13 +1,13 @@
 use itertools::Itertools;
 use vortex_error::{vortex_bail, vortex_err, VortexResult};
 
-use crate::{Array, ArrayDType};
+use crate::{Array, ArrayDType, OwnedArray};
 
 pub trait AsContiguousFn {
-    fn as_contiguous(&self, arrays: &[Array]) -> VortexResult<Array<'static>>;
+    fn as_contiguous(&self, arrays: &[Array]) -> VortexResult<OwnedArray>;
 }
 
-pub fn as_contiguous(arrays: &[Array]) -> VortexResult<Array<'static>> {
+pub fn as_contiguous(arrays: &[Array]) -> VortexResult<OwnedArray> {
     if arrays.is_empty() {
         vortex_bail!(ComputeError: "No arrays to concatenate");
     }

vortex-array/src/data.rs

@@ -9,7 +9,7 @@ use crate::encoding::EncodingRef;
 use crate::scalar::Scalar;
 use crate::stats::Stat;
 use crate::stats::Statistics;
-use crate::{Array, ArrayMetadata, IntoArray, ToArray};
+use crate::{Array, ArrayMetadata, IntoArray, OwnedArray, ToArray};
 
 #[derive(Clone, Debug)]
 pub struct ArrayData {
@@ -134,7 +134,7 @@ impl ToArray for ArrayData {
 }
 
 impl IntoArray<'static> for ArrayData {
-    fn into_array(self) -> Array<'static> {
+    fn into_array(self) -> OwnedArray {
         Array::Data(self)
     }
 }

vortex-array/src/ptype.rs

@@ -5,14 +5,15 @@ use arrow_array::types::*;
 use arrow_buffer::ArrowNativeType;
 use half::f16;
 use num_traits::{Num, NumCast};
+use serde::{Deserialize, Serialize};
 use vortex_error::{vortex_err, VortexError, VortexResult};
 use vortex_schema::DType::*;
 use vortex_schema::{DType, FloatWidth, IntWidth};
 
 use crate::scalar::{PScalar, Scalar};
 
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
-#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Hash, Serialize, Deserialize)]
 pub enum PType {
     U8,
     U16,

vortex-roaring/Cargo.toml

@@ -20,6 +20,8 @@ linkme = { workspace = true }
 croaring = { workspace = true }
 num-traits = { workspace = true }
 log = { workspace = true }
+serde = { workspace = true }
+paste = { workspace = true }
 
 [lints]
 workspace = true

vortex-roaring/src/boolean/compress.rs

@@ -1,22 +1,22 @@
 use croaring::Bitmap;
-use vortex::array::bool::{BoolArray, BoolEncoding};
-use vortex::array::downcast::DowncastArrayBuiltin;
-use vortex::array::{Array, ArrayRef};
+use vortex::array::bool::BoolArray;
 use vortex::compress::{CompressConfig, CompressCtx, EncodingCompression};
+use vortex::{Array, ArrayDType, ArrayDef, ArrayTrait, IntoArray, OwnedArray};
 use vortex_error::VortexResult;
 use vortex_schema::DType;
 use vortex_schema::Nullability::NonNullable;
 
-use crate::boolean::{RoaringBoolArray, RoaringBoolEncoding};
+use crate::boolean::RoaringBoolArray;
+use crate::{OwnedRoaringBoolArray, RoaringBool, RoaringBoolEncoding};
 
 impl EncodingCompression for RoaringBoolEncoding {
     fn can_compress(
         &self,
-        array: &dyn Array,
+        array: &Array,
         _config: &CompressConfig,
     ) -> Option<&dyn EncodingCompression> {
         // Only support bool enc arrays
-        if array.encoding().id() != BoolEncoding::ID {
+        if array.encoding().id() != RoaringBool::ID {
             return None;
         }
 
@@ -34,19 +34,19 @@ impl EncodingCompression for RoaringBoolEncoding {
 
     fn compress(
         &self,
-        array: &dyn Array,
-        _like: Option<&dyn Array>,
+        array: &Array,
+        _like: Option<&Array>,
         _ctx: CompressCtx,
-    ) -> VortexResult<ArrayRef> {
-        Ok(roaring_encode(array.as_bool()).into_array())
+    ) -> VortexResult<OwnedArray> {
+        roaring_encode(array.clone().flatten_bool()?).map(move |a| a.into_array())
     }
 }
 
-pub fn roaring_encode(bool_array: &BoolArray) -> RoaringBoolArray {
+pub fn roaring_encode(bool_array: BoolArray) -> VortexResult<OwnedRoaringBoolArray> {
     let mut bitmap = Bitmap::new();
     bitmap.extend(
         bool_array
-            .buffer()
+            .boolean_buffer()
             .iter()
             .enumerate()
             .filter(|(_, b)| *b)
@@ -55,5 +55,5 @@ pub fn roaring_encode(bool_array: &BoolArray) -> RoaringBoolArray {
     bitmap.run_optimize();
     bitmap.shrink_to_fit();
 
-    RoaringBoolArray::new(bitmap, bool_array.buffer().len())
+    RoaringBoolArray::try_new(bitmap, bool_array.len())
 }