Skip to content

Make sequence metadata in-memory + clean up #6759

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gatesn merged 2 commits into from
Mar 2, 2026
Merged

Make sequence metadata in-memory + clean up #6759

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fc0dd49
make sequence metadata in-memory + clean up
connortsui20 Mar 2, 2026
3af075e
fix cuda and public api
connortsui20 Mar 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions encodings/sequence/public-api.lock
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,11 @@ pub fn vortex_sequence::SequenceArray::last(&self) -> vortex_array::scalar::type

pub fn vortex_sequence::SequenceArray::multiplier(&self) -> vortex_array::scalar::typed_view::primitive::pvalue::PValue

pub fn vortex_sequence::SequenceArray::new(base: vortex_array::scalar::typed_view::primitive::pvalue::PValue, multiplier: vortex_array::scalar::typed_view::primitive::pvalue::PValue, ptype: vortex_array::dtype::ptype::PType, nullability: vortex_array::dtype::nullability::Nullability, length: usize) -> vortex_error::VortexResult<Self>

pub fn vortex_sequence::SequenceArray::ptype(&self) -> vortex_array::dtype::ptype::PType

pub fn vortex_sequence::SequenceArray::typed_new<T: vortex_array::dtype::ptype::NativePType + core::convert::Into<vortex_array::scalar::typed_view::primitive::pvalue::PValue>>(base: T, multiplier: T, nullability: vortex_array::dtype::nullability::Nullability, length: usize) -> vortex_error::VortexResult<Self>
pub fn vortex_sequence::SequenceArray::try_new(base: vortex_array::scalar::typed_view::primitive::pvalue::PValue, multiplier: vortex_array::scalar::typed_view::primitive::pvalue::PValue, ptype: vortex_array::dtype::ptype::PType, nullability: vortex_array::dtype::nullability::Nullability, length: usize) -> vortex_error::VortexResult<Self>

pub fn vortex_sequence::SequenceArray::try_new_typed<T: vortex_array::dtype::ptype::NativePType + core::convert::Into<vortex_array::scalar::typed_view::primitive::pvalue::PValue>>(base: T, multiplier: T, nullability: vortex_array::dtype::nullability::Nullability, length: usize) -> vortex_error::VortexResult<Self>

impl core::clone::Clone for vortex_sequence::SequenceArray

Expand Down Expand Up @@ -104,7 +104,7 @@ impl vortex_array::vtable::VTable for vortex_sequence::SequenceVTable

pub type vortex_sequence::SequenceVTable::Array = vortex_sequence::SequenceArray

pub type vortex_sequence::SequenceVTable::Metadata = vortex_array::metadata::ProstMetadata<vortex_sequence::array::SequenceMetadata>
pub type vortex_sequence::SequenceVTable::Metadata = vortex_sequence::array::SequenceMetadata

pub type vortex_sequence::SequenceVTable::OperationsVTable = vortex_sequence::SequenceVTable

Expand All @@ -124,7 +124,7 @@ pub fn vortex_sequence::SequenceVTable::child(_array: &vortex_sequence::Sequence

pub fn vortex_sequence::SequenceVTable::child_name(_array: &vortex_sequence::SequenceArray, idx: usize) -> alloc::string::String

pub fn vortex_sequence::SequenceVTable::deserialize(bytes: &[u8], _dtype: &vortex_array::dtype::DType, _len: usize, _buffers: &[vortex_array::buffer::BufferHandle], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Metadata>
pub fn vortex_sequence::SequenceVTable::deserialize(bytes: &[u8], dtype: &vortex_array::dtype::DType, _len: usize, _buffers: &[vortex_array::buffer::BufferHandle], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Metadata>

pub fn vortex_sequence::SequenceVTable::dtype(array: &vortex_sequence::SequenceArray) -> &vortex_array::dtype::DType

Expand Down
126 changes: 73 additions & 53 deletions encodings/sequence/src/array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,14 @@ use crate::rules::RULES;

vtable!(Sequence);

#[derive(Clone, prost::Message)]
#[derive(Debug, Clone, Copy)]
pub struct SequenceMetadata {
base: PValue,
multiplier: PValue,
}

#[derive(Clone, prost::Message)]
pub struct ProstSequenceMetadata {
#[prost(message, tag = "1")]
base: Option<vortex_proto::scalar::ScalarValue>,
#[prost(message, tag = "2")]
Expand All @@ -78,13 +84,13 @@ pub struct SequenceArray {
}

impl SequenceArray {
pub fn typed_new<T: NativePType + Into<PValue>>(
pub fn try_new_typed<T: NativePType + Into<PValue>>(
base: T,
multiplier: T,
nullability: Nullability,
length: usize,
) -> VortexResult<Self> {
Self::new(
Self::try_new(
base.into(),
multiplier.into(),
T::PTYPE,
Expand All @@ -94,7 +100,7 @@ impl SequenceArray {
}

/// Constructs a sequence array using two integer values (with the same ptype).
pub fn new(
pub fn try_new(
base: PValue,
multiplier: PValue,
ptype: PType,
Expand All @@ -111,16 +117,23 @@ impl SequenceArray {
))
})?;

Ok(Self::unchecked_new(
base,
multiplier,
ptype,
nullability,
length,
))
}

pub(crate) fn unchecked_new(
// SAFETY: we just validated that `ptype` is an integer and that the final
// element is representable via `try_last`.
Ok(unsafe { Self::new_unchecked(base, multiplier, ptype, nullability, length) })
}

/// Constructs a [`SequenceArray`] without validating that the `ptype` is an integer
/// type or that the final element is representable.
///
/// # Safety
///
/// The caller must ensure that:
/// - `ptype` is an integer type (i.e., `ptype.is_int()` returns `true`).
/// - `base + (length - 1) * multiplier` does not overflow the range of `ptype`.
///
/// Violating the first invariant will cause a panic. Violating the second will
/// cause silent wraparound when materializing elements, producing incorrect values.
pub(crate) unsafe fn new_unchecked(
base: PValue,
multiplier: PValue,
ptype: PType,
Expand Down Expand Up @@ -226,7 +239,7 @@ impl SequenceArray {
impl VTable for SequenceVTable {
type Array = SequenceArray;

type Metadata = ProstMetadata<SequenceMetadata>;
type Metadata = SequenceMetadata;
type OperationsVTable = Self;
type ValidityVTable = Self;

Expand Down Expand Up @@ -289,40 +302,37 @@ impl VTable for SequenceVTable {
}

fn metadata(array: &SequenceArray) -> VortexResult<Self::Metadata> {
Ok(ProstMetadata(SequenceMetadata {
base: Some((&array.base()).into()),
multiplier: Some((&array.multiplier()).into()),
}))
Ok(SequenceMetadata {
base: array.base(),
multiplier: array.multiplier(),
})
}

fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
Ok(Some(metadata.serialize()))
let prost = ProstMetadata(ProstSequenceMetadata {
base: Some((&metadata.base).into()),
multiplier: Some((&metadata.multiplier).into()),
});

Ok(Some(prost.serialize()))
}

fn deserialize(
bytes: &[u8],
_dtype: &DType,
dtype: &DType,
_len: usize,
_buffers: &[BufferHandle],
_session: &VortexSession,
) -> VortexResult<Self::Metadata> {
Ok(ProstMetadata(
<ProstMetadata<SequenceMetadata> as DeserializeMetadata>::deserialize(bytes)?,
))
}
let prost = ProstMetadata(
<ProstMetadata<ProstSequenceMetadata> as DeserializeMetadata>::deserialize(bytes)?,
);

fn build(
dtype: &DType,
len: usize,
metadata: &Self::Metadata,
_buffers: &[BufferHandle],
_children: &dyn ArrayChildren,
) -> VortexResult<SequenceArray> {
let ptype = dtype.as_ptype();

// We go via scalar to cast the scalar values into the correct PType
// We go via Scalar to validate that the value is valid for the ptype.
let base = Scalar::from_proto_value(
metadata
prost
.0
.base
.as_ref()
Expand All @@ -331,10 +341,10 @@ impl VTable for SequenceVTable {
)?
.as_primitive()
.pvalue()
.vortex_expect("non-nullable primitive");
.vortex_expect("sequence array base should be a non-nullable primitive");

let multiplier = Scalar::from_proto_value(
metadata
prost
.0
.multiplier
.as_ref()
Expand All @@ -343,15 +353,25 @@ impl VTable for SequenceVTable {
)?
.as_primitive()
.pvalue()
.vortex_expect("non-nullable primitive");
.vortex_expect("sequence array multiplier should be a non-nullable primitive");

Ok(SequenceArray::unchecked_new(
base,
multiplier,
ptype,
Ok(SequenceMetadata { base, multiplier })
}

fn build(
dtype: &DType,
len: usize,
metadata: &Self::Metadata,
_buffers: &[BufferHandle],
_children: &dyn ArrayChildren,
) -> VortexResult<SequenceArray> {
SequenceArray::try_new(
metadata.base,
metadata.multiplier,
dtype.as_ptype(),
dtype.nullability(),
len,
))
)
}

fn with_children(_array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
Expand Down Expand Up @@ -433,7 +453,7 @@ mod tests {

#[test]
fn test_sequence_canonical() {
let arr = SequenceArray::typed_new(2i64, 3, Nullability::NonNullable, 4).unwrap();
let arr = SequenceArray::try_new_typed(2i64, 3, Nullability::NonNullable, 4).unwrap();

let canon = PrimitiveArray::from_iter((0..4).map(|i| 2i64 + i * 3));

Expand All @@ -442,7 +462,7 @@ mod tests {

#[test]
fn test_sequence_slice_canonical() {
let arr = SequenceArray::typed_new(2i64, 3, Nullability::NonNullable, 4)
let arr = SequenceArray::try_new_typed(2i64, 3, Nullability::NonNullable, 4)
.unwrap()
.slice(2..3)
.unwrap();
Expand All @@ -454,7 +474,7 @@ mod tests {

#[test]
fn test_sequence_scalar_at() {
let scalar = SequenceArray::typed_new(2i64, 3, Nullability::NonNullable, 4)
let scalar = SequenceArray::try_new_typed(2i64, 3, Nullability::NonNullable, 4)
.unwrap()
.scalar_at(2)
.unwrap();
Expand All @@ -467,19 +487,19 @@ mod tests {

#[test]
fn test_sequence_min_max() {
assert!(SequenceArray::typed_new(-127i8, -1i8, Nullability::NonNullable, 2).is_ok());
assert!(SequenceArray::typed_new(126i8, -1i8, Nullability::NonNullable, 2).is_ok());
assert!(SequenceArray::try_new_typed(-127i8, -1i8, Nullability::NonNullable, 2).is_ok());
assert!(SequenceArray::try_new_typed(126i8, -1i8, Nullability::NonNullable, 2).is_ok());
}

#[test]
fn test_sequence_too_big() {
assert!(SequenceArray::typed_new(127i8, 1i8, Nullability::NonNullable, 2).is_err());
assert!(SequenceArray::typed_new(-128i8, -1i8, Nullability::NonNullable, 2).is_err());
assert!(SequenceArray::try_new_typed(127i8, 1i8, Nullability::NonNullable, 2).is_err());
assert!(SequenceArray::try_new_typed(-128i8, -1i8, Nullability::NonNullable, 2).is_err());
}

#[test]
fn positive_multiplier_is_strict_sorted() -> VortexResult<()> {
let arr = SequenceArray::typed_new(0i64, 3, Nullability::NonNullable, 4)?;
let arr = SequenceArray::try_new_typed(0i64, 3, Nullability::NonNullable, 4)?;

let is_sorted = arr
.statistics()
Expand All @@ -495,7 +515,7 @@ mod tests {

#[test]
fn zero_multiplier_is_sorted_not_strict() -> VortexResult<()> {
let arr = SequenceArray::typed_new(5i64, 0, Nullability::NonNullable, 4)?;
let arr = SequenceArray::try_new_typed(5i64, 0, Nullability::NonNullable, 4)?;

let is_sorted = arr
.statistics()
Expand All @@ -511,7 +531,7 @@ mod tests {

#[test]
fn negative_multiplier_not_sorted() -> VortexResult<()> {
let arr = SequenceArray::typed_new(10i64, -1, Nullability::NonNullable, 4)?;
let arr = SequenceArray::try_new_typed(10i64, -1, Nullability::NonNullable, 4)?;

let is_sorted = arr
.statistics()
Expand All @@ -530,7 +550,7 @@ mod tests {
#[test]
fn test_large_multiplier_sorted() -> VortexResult<()> {
let large_multiplier = (i64::MAX as u64) + 1;
let arr = SequenceArray::typed_new(0, large_multiplier, Nullability::NonNullable, 2)?;
let arr = SequenceArray::try_new_typed(0, large_multiplier, Nullability::NonNullable, 2)?;

let is_sorted = arr
.statistics()
Expand Down
4 changes: 2 additions & 2 deletions encodings/sequence/src/compress.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ fn encode_primitive_array<P: NativePType + Into<PValue> + CheckedAdd + CheckedSu
) -> VortexResult<Option<ArrayRef>> {
if slice.len() == 1 {
// The multiplier here can be any value, zero is chosen
return SequenceArray::typed_new(slice[0], P::zero(), nullability, 1)
return SequenceArray::try_new_typed(slice[0], P::zero(), nullability, 1)
.map(|a| Some(a.to_array()));
}
let base = slice[0];
Expand All @@ -69,7 +69,7 @@ fn encode_primitive_array<P: NativePType + Into<PValue> + CheckedAdd + CheckedSu
.windows(2)
.all(|w| Some(w[1]) == w[0].checked_add(&multiplier))
.then_some(
SequenceArray::typed_new(base, multiplier, nullability, slice.len())
SequenceArray::try_new_typed(base, multiplier, nullability, slice.len())
.map(|a| a.to_array()),
)
.transpose()
Expand Down
25 changes: 14 additions & 11 deletions encodings/sequence/src/compute/cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ impl CastReduce for SequenceVTable {
// For SequenceArray, we can just create a new one with the same parameters
// but different nullability
return Ok(Some(
SequenceArray::new(
SequenceArray::try_new(
array.base(),
array.multiplier(),
*target_ptype,
Expand Down Expand Up @@ -71,7 +71,7 @@ impl CastReduce for SequenceVTable {
.ok_or_else(|| vortex_err!("Cast resulted in null multiplier value"))?;

return Ok(Some(
SequenceArray::new(
SequenceArray::try_new(
new_base,
new_multiplier,
*target_ptype,
Expand Down Expand Up @@ -102,7 +102,8 @@ mod tests {

#[test]
fn test_cast_sequence_nullability() {
let sequence = SequenceArray::typed_new(0u32, 1u32, Nullability::NonNullable, 4).unwrap();
let sequence =
SequenceArray::try_new_typed(0u32, 1u32, Nullability::NonNullable, 4).unwrap();

// Cast to nullable
let casted = sequence
Expand All @@ -118,7 +119,7 @@ mod tests {
#[test]
fn test_cast_sequence_u32_to_i64() {
let sequence =
SequenceArray::typed_new(100u32, 10u32, Nullability::NonNullable, 4).unwrap();
SequenceArray::try_new_typed(100u32, 10u32, Nullability::NonNullable, 4).unwrap();

let casted = sequence
.to_array()
Expand All @@ -137,7 +138,8 @@ mod tests {
#[test]
fn test_cast_sequence_i16_to_i32_nullable() {
// Test ptype change AND nullability change in one cast
let sequence = SequenceArray::typed_new(5i16, 3i16, Nullability::NonNullable, 3).unwrap();
let sequence =
SequenceArray::try_new_typed(5i16, 3i16, Nullability::NonNullable, 3).unwrap();

let casted = sequence
.to_array()
Expand All @@ -158,7 +160,8 @@ mod tests {

#[test]
fn test_cast_sequence_to_float_delegates_to_canonical() {
let sequence = SequenceArray::typed_new(0i32, 1i32, Nullability::NonNullable, 5).unwrap();
let sequence =
SequenceArray::try_new_typed(0i32, 1i32, Nullability::NonNullable, 5).unwrap();

// Cast to float should delegate to canonical (SequenceArray doesn't support float)
let casted = sequence
Expand All @@ -180,15 +183,15 @@ mod tests {
}

#[rstest]
#[case::i32(SequenceArray::typed_new(0i32, 1i32, Nullability::NonNullable, 5).unwrap())]
#[case::u64(SequenceArray::typed_new(1000u64, 100u64, Nullability::NonNullable, 4).unwrap())]
#[case::i32(SequenceArray::try_new_typed(0i32, 1i32, Nullability::NonNullable, 5).unwrap())]
#[case::u64(SequenceArray::try_new_typed(1000u64, 100u64, Nullability::NonNullable, 4).unwrap())]
// TODO(DK): SequenceArray does not actually conform. You cannot cast this array to u8 even
// though all its values are representable therein.
//
// #[case::negative_step(SequenceArray::typed_new(100i32, -10i32, Nullability::NonNullable,
// #[case::negative_step(SequenceArray::try_new_typed(100i32, -10i32, Nullability::NonNullable,
// 5).unwrap())]
#[case::single(SequenceArray::typed_new(42i64, 0i64, Nullability::NonNullable, 1).unwrap())]
#[case::constant(SequenceArray::typed_new(
#[case::single(SequenceArray::try_new_typed(42i64, 0i64, Nullability::NonNullable, 1).unwrap())]
#[case::constant(SequenceArray::try_new_typed(
100i32,
0i32, // multiplier of 0 means constant array
Nullability::NonNullable,
Expand Down
Loading
Loading