From c2273d6cff0484d06e7e3e7f9f18e594c06dab68 Mon Sep 17 00:00:00 2001 From: Christopher Wirt Date: Wed, 4 Nov 2026 11:38:58 +0000 Subject: [PATCH 1/2] fix(compressor): detach constant byte scalars from source buffers Signed-off-by: Christopher Wirt --- .../src/compressor/constant/mod.rs | 16 +++++++++ .../src/compressor/constant/tests.rs | 34 +++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/vortex-compressor/src/compressor/constant/mod.rs b/vortex-compressor/src/compressor/constant/mod.rs index ca28f72a094..eb517f64cc0 100644 --- a/vortex-compressor/src/compressor/constant/mod.rs +++ b/vortex-compressor/src/compressor/constant/mod.rs @@ -27,6 +27,9 @@ use vortex_array::dtype::NativePType; use vortex_array::match_each_decimal_value_type; use vortex_array::match_each_native_ptype; use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_buffer::BufferString; +use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -183,6 +186,19 @@ fn compress_constant(source: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult< }; let scalar = source.execute_scalar(first_valid, ctx)?; + // A string scalar can be a slice of the source's data or views buffer. Constants may + // remain buffered until EOF, so keeping that slice would retain the whole input chunk. + let scalar = match scalar.value() { + Some(ScalarValue::Utf8(value)) => Scalar::try_new( + scalar.dtype().clone(), + Some(ScalarValue::Utf8(BufferString::from(value.as_str()))), + )?, + Some(ScalarValue::Binary(value)) => Scalar::try_new( + scalar.dtype().clone(), + Some(ScalarValue::Binary(ByteBuffer::copy_from(value))), + )?, + _ => scalar, + }; let const_arr = ConstantArray::new(scalar, source.len()).into_array(); if mask.all_true() { diff --git a/vortex-compressor/src/compressor/constant/tests.rs b/vortex-compressor/src/compressor/constant/tests.rs index 8a160f3d533..3d7fa598557 100644 --- a/vortex-compressor/src/compressor/constant/tests.rs +++ b/vortex-compressor/src/compressor/constant/tests.rs @@ -86,6 +86,40 @@ fn constant_string_compresses_without_schemes() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::utf8(false)] +#[case::binary(true)] +fn constant_bytes_release_source_buffer( + #[case] binary: bool, + #[values(8, 128)] width: usize, + #[values(false, true)] nullable: bool, +) -> VortexResult<()> { + let nullability = Nullability::from(nullable); + let dtype = if binary { + DType::Binary(nullability) + } else { + DType::Utf8(nullability) + }; + let value = vec![if binary { 255 } else { b'a' }; width]; + let source = VarBinViewArray::from_iter( + (0..8192).map(|i| (!nullable || i % 5 != 0).then_some(value.as_slice())), + dtype, + ); + // Short values slice the views buffer; long values slice a data buffer. + let backing = source.bytes_at(usize::from(nullable)); + let source = source.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = empty_compressor().compress(&source, &mut ctx)?; + assert_arrays_eq!(compressed, source, &mut ctx); + drop(source); + + assert!( + backing.try_into_mut().is_ok(), + "constant retained the source buffer" + ); + Ok(()) +} + #[test] fn constant_bool_compresses_without_schemes() -> VortexResult<()> { let array = BoolArray::from_iter(std::iter::repeat_n(true, 100)).into_array(); From 8a87d0ea4259d80b2b1fd2e630b770a7cd3a1a5b Mon Sep 17 00:00:00 2001 From: Christopher Wirt Date: Wed, 4 Nov 2026 12:03:48 +0000 Subject: [PATCH 2/2] refactor(compressor): detach nested constant bytes and simplify tests Copy string and binary bytes nested in tuple values too, so constant struct- or list-backed extension columns also release the source chunk. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Christopher Wirt --- .../src/compressor/constant/mod.rs | 34 ++++---- .../src/compressor/constant/tests.rs | 84 +++++++++++++++---- 2 files changed, 89 insertions(+), 29 deletions(-) diff --git a/vortex-compressor/src/compressor/constant/mod.rs b/vortex-compressor/src/compressor/constant/mod.rs index eb517f64cc0..ab835a11759 100644 --- a/vortex-compressor/src/compressor/constant/mod.rs +++ b/vortex-compressor/src/compressor/constant/mod.rs @@ -167,6 +167,24 @@ pub(crate) fn compress_as_constant( Ok(if accepted { compressed } else { array }) } +/// Copies string and binary bytes, including nested ones, into their own allocations. +/// +/// Scalars extracted from an array can slice its buffers, so a constant built from one would +/// otherwise keep the whole source array alive. +fn detach_bytes(value: ScalarValue) -> ScalarValue { + match value { + ScalarValue::Utf8(s) => ScalarValue::Utf8(BufferString::from(s.as_str())), + ScalarValue::Binary(b) => ScalarValue::Binary(ByteBuffer::copy_from(&b)), + ScalarValue::Tuple(children) => ScalarValue::Tuple( + children + .into_iter() + .map(|child| child.map(detach_bytes)) + .collect(), + ), + value => value, + } +} + /// Encodes an array whose valid values are all equal. /// /// Returns a [`ConstantArray`], wrapped in a [`MaskedArray`] when the array has some nulls, or a @@ -185,20 +203,8 @@ fn compress_constant(source: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult< ); }; - let scalar = source.execute_scalar(first_valid, ctx)?; - // A string scalar can be a slice of the source's data or views buffer. Constants may - // remain buffered until EOF, so keeping that slice would retain the whole input chunk. - let scalar = match scalar.value() { - Some(ScalarValue::Utf8(value)) => Scalar::try_new( - scalar.dtype().clone(), - Some(ScalarValue::Utf8(BufferString::from(value.as_str()))), - )?, - Some(ScalarValue::Binary(value)) => Scalar::try_new( - scalar.dtype().clone(), - Some(ScalarValue::Binary(ByteBuffer::copy_from(value))), - )?, - _ => scalar, - }; + let (dtype, value) = source.execute_scalar(first_valid, ctx)?.into_parts(); + let scalar = Scalar::try_new(dtype, value.map(detach_bytes))?; let const_arr = ConstantArray::new(scalar, source.len()).into_array(); if mask.all_true() { diff --git a/vortex-compressor/src/compressor/constant/tests.rs b/vortex-compressor/src/compressor/constant/tests.rs index 3d7fa598557..29284f4544a 100644 --- a/vortex-compressor/src/compressor/constant/tests.rs +++ b/vortex-compressor/src/compressor/constant/tests.rs @@ -16,6 +16,7 @@ use vortex_array::arrays::DictArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::Masked; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::TemporalArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbinview::BinaryView; @@ -87,26 +88,21 @@ fn constant_string_compresses_without_schemes() -> VortexResult<()> { } #[rstest] -#[case::utf8(false)] -#[case::binary(true)] +#[case::utf8(DType::Utf8 as fn(Nullability) -> DType, b'a')] +#[case::binary(DType::Binary, 0xff)] fn constant_bytes_release_source_buffer( - #[case] binary: bool, + #[case] dtype: fn(Nullability) -> DType, + #[case] byte: u8, + // 8 bytes is inlined in the views buffer; 128 bytes lives in a data buffer. #[values(8, 128)] width: usize, - #[values(false, true)] nullable: bool, + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, ) -> VortexResult<()> { - let nullability = Nullability::from(nullable); - let dtype = if binary { - DType::Binary(nullability) - } else { - DType::Utf8(nullability) - }; - let value = vec![if binary { 255 } else { b'a' }; width]; + let value = vec![byte; width]; let source = VarBinViewArray::from_iter( - (0..8192).map(|i| (!nullable || i % 5 != 0).then_some(value.as_slice())), - dtype, + (0..1024).map(|i| (!nullability.is_nullable() || i % 5 != 0).then_some(value.as_slice())), + dtype(nullability), ); - // Short values slice the views buffer; long values slice a data buffer. - let backing = source.bytes_at(usize::from(nullable)); + let backing = source.bytes_at(1); let source = source.into_array(); let mut ctx = SESSION.create_execution_ctx(); let compressed = empty_compressor().compress(&source, &mut ctx)?; @@ -120,6 +116,30 @@ fn constant_bytes_release_source_buffer( Ok(()) } +#[test] +fn constant_nested_bytes_release_source_buffer() -> VortexResult<()> { + let value = vec![b'a'; 128]; + let field = VarBinViewArray::from_iter( + (0..1024).map(|_| Some(value.as_slice())), + DType::Utf8(Nullability::NonNullable), + ); + let backing = field.bytes_at(1); + let storage = StructArray::from_fields(&[("name", field.into_array())])?.into_array(); + let dtype = ExtDType::::try_new(0, storage.dtype().clone())?; + let source = ExtensionArray::new(dtype.erased(), storage).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = empty_compressor().compress(&source, &mut ctx)?; + assert!(compressed.is::()); + assert_arrays_eq!(compressed, source, &mut ctx); + drop(source); + + assert!( + backing.try_into_mut().is_ok(), + "constant retained the source buffer" + ); + Ok(()) +} + #[test] fn constant_bool_compresses_without_schemes() -> VortexResult<()> { let array = BoolArray::from_iter(std::iter::repeat_n(true, 100)).into_array(); @@ -298,6 +318,40 @@ impl ExtVTable for FloatExtension { } } +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] +struct StructExtension; + +impl ExtVTable for StructExtension { + type Metadata = u8; + type NativeValue<'a> = &'a ScalarValue; + + #[expect(clippy::disallowed_methods, reason = "test-only id")] + fn id(&self) -> ExtId { + ExtId::new("test.struct") + } + + fn serialize_metadata(&self, metadata: &u8) -> VortexResult> { + Ok(vec![*metadata]) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + vortex_ensure!(metadata.len() == 1, "expected one metadata byte"); + Ok(metadata[0]) + } + + fn validate_dtype(dtype: &ExtDType) -> VortexResult<()> { + vortex_ensure!(dtype.storage_dtype().is_struct(), "expected struct storage"); + Ok(()) + } + + fn unpack_native<'a>( + _dtype: &'a ExtDType, + value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(value) + } +} + #[rstest] #[case::equal(1.5, 1.5, true)] #[case::nan_among_equal_values(1.5, f64::NAN, false)]