diff --git a/vortex-compressor/src/compressor/constant/mod.rs b/vortex-compressor/src/compressor/constant/mod.rs index ca28f72a094..ab835a11759 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; @@ -164,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 @@ -182,7 +203,8 @@ fn compress_constant(source: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult< ); }; - let scalar = source.execute_scalar(first_valid, ctx)?; + 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 8a160f3d533..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; @@ -86,6 +87,59 @@ fn constant_string_compresses_without_schemes() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::utf8(DType::Utf8 as fn(Nullability) -> DType, b'a')] +#[case::binary(DType::Binary, 0xff)] +fn constant_bytes_release_source_buffer( + #[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(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, +) -> VortexResult<()> { + let value = vec![byte; width]; + let source = VarBinViewArray::from_iter( + (0..1024).map(|i| (!nullability.is_nullable() || i % 5 != 0).then_some(value.as_slice())), + dtype(nullability), + ); + 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)?; + 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_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(); @@ -264,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)]