diff --git a/Cargo.lock b/Cargo.lock index 8451ac1a3e4..d2729c9f008 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12520,6 +12520,8 @@ dependencies = [ "vortex-edition", "vortex-error", "vortex-expr", + "vortex-file", + "vortex-io", "vortex-layout", "vortex-mask", "vortex-session", diff --git a/Cargo.toml b/Cargo.toml index 7fc483deec5..658785c90f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -214,7 +214,7 @@ parquet = "59.2" parquet-variant = "59.2" parquet-variant-compute = "59.2" paste = "1.0.15" -pco = "1.0.1" +pco = "1.1.0" percent-encoding = "2.3.2" pin-project-lite = "0.2.15" primitive-types = { version = "0.14.0" } diff --git a/encodings/pco/src/array/mod.rs b/encodings/pco/src/array/mod.rs index c30c518cc93..2f11389543a 100644 --- a/encodings/pco/src/array/mod.rs +++ b/encodings/pco/src/array/mod.rs @@ -604,9 +604,6 @@ impl PcoData { let old_len = decompressed_values.len(); let new_len = old_len + page_n_values; decompressed_values.reserve(page_n_values); - unsafe { - decompressed_values.set_len(new_len); - } let page: &[u8] = self .pages .get(page_idx) @@ -626,8 +623,15 @@ impl PcoData { let mut pd = cd .page_decompressor(page, page_n_values) .map_err(vortex_err_from_pco)?; - pd.read(&mut decompressed_values[old_len..new_len]) + let progress = pd + .read_uninit(&mut decompressed_values.spare_capacity_mut()[..page_n_values]) .map_err(vortex_err_from_pco)?; + vortex_ensure_eq!(progress.n_processed, page_n_values); + // SAFETY: `read_uninit` initialized all `page_n_values` spare elements, + // as confirmed by the processed count above. + unsafe { + decompressed_values.set_len(new_len); + } chunk_decompressor = Some(cd); } else { diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs index 00aa801254a..f7a10c8844f 100644 --- a/encodings/pco/src/probe.rs +++ b/encodings/pco/src/probe.rs @@ -21,6 +21,7 @@ use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; use vortex_mask::Mask; use crate::Pco; @@ -156,10 +157,13 @@ fn decode_page( .page_decompressor(buffer, n_values) .map_err(vortex_err_from_pco)?; let mut values = BufferMut::::with_capacity_in(n_values, ctx.allocator().clone()); - // SAFETY: the buffer reserves `n_values` elements, and the page decompressor was built for - // exactly that count, so `read` writes every element before anything observes it. + let progress = decoder + .read_uninit(&mut values.spare_capacity_mut()[..n_values]) + .map_err(vortex_err_from_pco)?; + vortex_ensure_eq!(progress.n_processed, n_values); + // SAFETY: `read_uninit` initialized all `n_values` spare elements, as confirmed by the + // processed count above. unsafe { values.set_len(n_values) }; - decoder.read(&mut values).map_err(vortex_err_from_pco)?; Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable)) } diff --git a/encodings/uuid/src/arrow.rs b/encodings/uuid/src/arrow.rs index d69c35a7928..4a4011c5859 100644 --- a/encodings/uuid/src/arrow.rs +++ b/encodings/uuid/src/arrow.rs @@ -30,13 +30,12 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDType; -use vortex_array::dtype::extension::ExtVTable; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; use vortex_arrow::nulls; @@ -56,30 +55,30 @@ const UUID_BYTE_LEN: i32 = 16; /// The cached Id of Arrow's canonical `arrow.uuid` extension type. static ARROW_UUID: CachedId = CachedId::new(ArrowUuid::NAME); -impl ArrowExportVTable for Uuid { - fn arrow_ext_id(&self) -> Id { - *ARROW_UUID +impl ArrowExtVTable for Uuid { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::PreserveVortex } - fn vortex_id(&self) -> Id { - Uuid.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_UUID } - // Encode all of these. fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, _session: &ArrowSession, ) -> VortexResult> { let mut field = Field::new( - name.to_string(), + name, DataType::FixedSizeBinary(UUID_BYTE_LEN), - dtype.is_nullable(), + ext_dtype.storage_dtype().is_nullable(), ); field .try_with_extension_type(ArrowUuid) - .vortex_expect("FixedSizeBinary[16] is correct type for ArrowUuid"); + .vortex_expect("UUID field was constructed with FixedSizeBinary(16)"); + Ok(Some(field)) } @@ -90,28 +89,13 @@ impl ArrowExportVTable for Uuid { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_uuid = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_uuid { - return Ok(ArrowExport::Unsupported(array)); - } Ok(ArrowExport::Exported(try_fsl_to_fsb(array, options, ctx)?)) } -} - -impl ArrowImportVTable for Uuid { - fn arrow_ext_id(&self) -> Id { - *ARROW_UUID - } fn from_arrow_field( - &self, field: &Field, _session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { if !field.has_valid_extension_type::() { return Ok(None); } @@ -122,24 +106,20 @@ impl ArrowImportVTable for Uuid { field.is_nullable().into(), ); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Uuid, UuidMetadata::default(), storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + Uuid, + UuidMetadata::default(), + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, _field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, _session: &ArrowSession, ) -> VortexResult { - let DType::Extension(dtype) = dtype else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !matches!(array.data_type(), DataType::FixedSizeBinary(UUID_BYTE_LEN)) - || !dtype.is::() - { + if !matches!(array.data_type(), DataType::FixedSizeBinary(UUID_BYTE_LEN)) { return Ok(ArrowImport::Unsupported(array)); } @@ -150,7 +130,7 @@ impl ArrowImportVTable for Uuid { PType::U8, Validity::NonNullable, ); - let validity = nulls(fsb.nulls(), dtype.is_nullable())?; + let validity = nulls(fsb.nulls(), ext_dtype.storage_dtype().is_nullable())?; let storage = FixedSizeListArray::new( u8_array.into_array(), @@ -160,7 +140,7 @@ impl ArrowImportVTable for Uuid { ) .into_array(); Ok(ArrowImport::Imported( - ExtensionArray::new(dtype.clone(), storage).into_array(), + ExtensionArray::new(ext_dtype.clone().erased(), storage).into_array(), )) } } @@ -221,7 +201,9 @@ mod tests { use arrow_schema::Fields; use arrow_schema::Schema; use rstest::rstest; + use uuid::Version; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::DictArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::Struct; use vortex_array::arrays::StructArray; @@ -233,11 +215,44 @@ mod tests { use vortex_array::iter::ArrayIterator; use vortex_arrow::ArrowArrayStreamAdapter; use vortex_arrow::ArrowSessionExt; - use vortex_error::vortex_err; use vortex_session::VortexSession; use super::*; + #[track_caller] + fn check_version_round_trip( + version: Option, + nullable: bool, + dictionary: bool, + ) -> VortexResult<()> { + let session = uuid_session(); + let arrow = session.arrow(); + let mut ctx = session.create_execution_ctx(); + let array = versioned_uuid_array(version, nullable)?; + let mut field = arrow.to_arrow_field("id", array.dtype())?; + assert!(field.has_valid_extension_type::()); + assert_eq!(arrow.from_arrow_field(&field)?, *array.dtype()); + + let array = if dictionary { + field = field.clone().with_data_type(DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(field.data_type().clone()), + )); + DictArray::try_new(PrimitiveArray::from_iter([2u32, 1, 0]).into_array(), array)? + .into_array() + } else { + array + }; + + let exported = array.clone().execute_arrow(Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + + let imported = arrow.from_arrow_array(exported, &field)?; + vortex_array::assert_arrays_eq!(imported, array, &mut ctx); + + Ok(()) + } + /// A session with the UUID extension dtype and its Arrow plugins registered. fn uuid_session() -> VortexSession { let session = VortexSession::empty(); @@ -289,46 +304,6 @@ mod tests { Ok(()) } - #[test] - fn to_arrow_field_top_level_uuid_carries_extension_metadata() -> VortexResult<()> { - let vortex_session = uuid_session(); - let session = vortex_session.arrow(); - let field = session.to_arrow_field("id", &uuid_dtype(false))?; - assert!(field.has_valid_extension_type::()); - Ok(()) - } - - #[test] - fn to_arrow_field_struct_with_nested_uuid_preserves_metadata() -> VortexResult<()> { - let vortex_session = uuid_session(); - let session = vortex_session.arrow(); - let dtype = DType::Struct( - StructFields::from_iter([(FieldName::from("id"), uuid_dtype(false))]), - Nullability::NonNullable, - ); - let field = session.to_arrow_field("row", &dtype)?; - let DataType::Struct(inner) = field.data_type() else { - panic!("expected Struct, got {:?}", field.data_type()); - }; - assert_eq!(inner.len(), 1); - assert_eq!(inner[0].data_type(), &DataType::FixedSizeBinary(16)); - assert!(inner[0].has_valid_extension_type::()); - Ok(()) - } - - #[test] - fn to_arrow_field_list_of_uuid_preserves_metadata() -> VortexResult<()> { - let vortex_session = uuid_session(); - let session = vortex_session.arrow(); - let dtype = DType::List(Arc::new(uuid_dtype(true)), Nullability::NonNullable); - let field = session.to_arrow_field("ids", &dtype)?; - let DataType::List(elem) = field.data_type() else { - panic!("expected List, got {:?}", field.data_type()); - }; - assert!(elem.has_valid_extension_type::()); - Ok(()) - } - #[test] fn to_arrow_field_fixed_size_list_of_uuid_preserves_metadata() -> VortexResult<()> { let vortex_session = uuid_session(); @@ -379,7 +354,7 @@ mod tests { } #[test] - fn to_arrow_schema_struct_of_struct_uuid() -> VortexResult<()> { + fn schema_roundtrip_preserves_nested_uuid() -> VortexResult<()> { let vortex_session = uuid_session(); let session = vortex_session.arrow(); let inner = DType::Struct( @@ -387,7 +362,14 @@ mod tests { Nullability::NonNullable, ); let outer = DType::Struct( - StructFields::from_iter([(FieldName::from("payload"), inner)]), + StructFields::from_iter([ + (FieldName::from("payload"), inner), + (FieldName::from("id"), uuid_dtype(false)), + ( + FieldName::from("ids"), + DType::List(Arc::new(uuid_dtype(true)), Nullability::NonNullable), + ), + ]), Nullability::NonNullable, ); let schema = session.to_arrow_schema(&outer)?; @@ -396,6 +378,8 @@ mod tests { panic!("expected Struct, got {:?}", payload.data_type()); }; assert!(inner_fields[0].has_valid_extension_type::()); + assert_eq!(session.from_arrow_schema(&schema)?, outer); + Ok(()) } @@ -415,26 +399,6 @@ mod tests { Ok(()) } - #[test] - fn schema_roundtrip_preserves_nested_uuid() -> VortexResult<()> { - let vortex_session = uuid_session(); - let session = vortex_session.arrow(); - let dtype = DType::Struct( - StructFields::from_iter([ - (FieldName::from("id"), uuid_dtype(false)), - ( - FieldName::from("ids"), - DType::List(Arc::new(uuid_dtype(true)), Nullability::NonNullable), - ), - ]), - Nullability::NonNullable, - ); - let schema = session.to_arrow_schema(&dtype)?; - let roundtripped = session.from_arrow_schema(&schema)?; - assert_eq!(roundtripped, dtype); - Ok(()) - } - #[test] fn to_arrow_datatype_dispatches_plugins() -> VortexResult<()> { let vortex_session = uuid_session(); @@ -472,13 +436,7 @@ mod tests { let mut ctx = vortex_session.create_execution_ctx(); let session = vortex_session.arrow(); - let mut field = Field::new("id", DataType::FixedSizeBinary(16), false); - field.try_with_extension_type(ArrowUuid)?; - let arrow_array: ArrowArrayRef = Arc::new(FixedSizeBinaryArray::try_from_iter( - [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(), - )?); - - let vortex_array = session.from_arrow_array(arrow_array, &field)?; + let vortex_array = uuid_array(&session)?; let vortex_ext = vortex_array.dtype().as_extension(); assert!(vortex_ext.is::()); @@ -595,7 +553,11 @@ mod tests { &values, )?); - let expected = run_array.to_data(); + let expected = run_array + .to_data() + .into_builder() + .data_type(field.data_type().clone()) + .build()?; let vortex_array = session.from_arrow_array(run_array, &field)?; assert_eq!(vortex_array.len(), 5); // The array conversion must agree with the dtype conversion. @@ -603,20 +565,8 @@ mod tests { // And the values must round-trip back out through the export plugin. let exported = vortex_array.execute_arrow(Some(&field), &mut ctx)?; - assert_eq!(exported.len(), 5); - let ree = exported - .as_any() - .downcast_ref::>() - .ok_or_else(|| { - vortex_err!( - "expected an Int32 run-end array, got {}", - exported.data_type() - ) - })?; assert_eq!(exported.to_data(), expected); - let values = ree.values().as_fixed_size_binary(); - assert_eq!(values.value(0), b"0123456789abcdef"); - assert_eq!(values.value(1), b"fedcba9876543210"); + Ok(()) } @@ -665,12 +615,8 @@ mod tests { // extension survives the export as well. let mut ctx = vortex_session.create_execution_ctx(); let exported = array.execute_arrow(Some(&field), &mut ctx)?; - assert_eq!(exported.data_type(), field.data_type()); assert_eq!(exported.to_data(), expected); - let values = exported.as_any_dictionary().values().as_list::(); - let uuids = values.values().as_fixed_size_binary(); - assert_eq!(uuids.value(0), b"0123456789abcdef"); - assert_eq!(uuids.value(1), b"fedcba9876543210"); + Ok(()) } @@ -858,4 +804,78 @@ mod tests { Ok(()) } + + /// Builds three UUIDs with valid version bits, with a null middle row when nullable. + fn versioned_uuid_array(version: Option, nullable: bool) -> VortexResult { + let dtype = ExtDType::try_with_vtable( + Uuid, + UuidMetadata { version }, + uuid_dtype(nullable).as_extension().storage_dtype().clone(), + )? + .erased(); + + let bytes = match version { + Some(Version::Max) => [u8::MAX; 16], + Some(Version::Nil) | None => [0; 16], + Some(version) => uuid::Builder::from_bytes([0; 16]) + .with_version(version) + .with_variant(uuid::Variant::RFC4122) + .into_uuid() + .into_bytes(), + }; + + let validity = if nullable { + Validity::from_iter([ + true, // First UUID. + false, // Null middle UUID. + true, // Last UUID. + ]) + } else { + Validity::NonNullable + }; + let storage = FixedSizeListArray::try_new( + PrimitiveArray::from_iter((0..3).flat_map(|_| bytes)).into_array(), + 16, + validity, + 3, + )?; + + Ok(ExtensionArray::try_new(dtype, storage.into_array())?.into_array()) + } + + #[rstest] + #[case::unspecified(None, false, false)] + #[case::nil(Some(Version::Nil), false, false)] + #[case::mac(Some(Version::Mac), false, false)] + #[case::dce(Some(Version::Dce), false, false)] + #[case::md5(Some(Version::Md5), false, false)] + #[case::random(Some(Version::Random), false, false)] + #[case::sha1(Some(Version::Sha1), false, false)] + #[case::sort_mac(Some(Version::SortMac), false, false)] + #[case::sort_rand(Some(Version::SortRand), false, false)] + #[case::custom(Some(Version::Custom), false, false)] + #[case::max(Some(Version::Max), false, false)] + #[case::nullable(Some(Version::Random), true, false)] + #[case::dictionary(Some(Version::Random), false, true)] + #[case::nullable_dictionary(Some(Version::Random), true, true)] + fn version_round_trip( + #[case] version: Option, + #[case] nullable: bool, + #[case] dictionary: bool, + ) -> VortexResult<()> { + check_version_round_trip(version, nullable, dictionary) + } + + #[test] + fn mismatched_metadata_is_rejected() -> VortexResult<()> { + let session = uuid_session(); + let mut ctx = session.create_execution_ctx(); + let array = uuid_array(&session.arrow())?; + let target_array = versioned_uuid_array(Some(Version::Random), false)?; + let field = session.arrow().to_arrow_field("id", target_array.dtype())?; + + assert!(array.execute_arrow(Some(&field), &mut ctx).is_err()); + + Ok(()) + } } diff --git a/encodings/uuid/src/lib.rs b/encodings/uuid/src/lib.rs index b079120d88c..7db124c50a4 100644 --- a/encodings/uuid/src/lib.rs +++ b/encodings/uuid/src/lib.rs @@ -24,8 +24,6 @@ mod arrow; mod metadata; mod vtable; -use std::sync::Arc; - pub use metadata::UuidMetadata; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::session::DTypeSessionExt; @@ -44,8 +42,7 @@ pub struct Uuid; pub fn initialize(session: &VortexSession) { session.dtypes().register(Uuid); session.aggregate_fns().register_storage_ordered(Uuid); - session.arrow().register_exporter(Arc::new(Uuid)); - session.arrow().register_importer(Arc::new(Uuid)); + session.arrow().register_extension(Uuid); } #[cfg(test)] diff --git a/fuzz/src/gpu/mod.rs b/fuzz/src/gpu/mod.rs index 91e116e38cd..c972581d5f1 100644 --- a/fuzz/src/gpu/mod.rs +++ b/fuzz/src/gpu/mod.rs @@ -13,9 +13,17 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::arbitrary::ArbitraryArray; +use vortex_array::arrays::arbitrary::ArbitraryArrayConfig; +use vortex_array::arrays::arbitrary::ArbitraryWith; use vortex_array::arrays::dict::ArbitraryDictArray; +use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; use vortex_error::vortex_err; use crate::SESSION; @@ -28,13 +36,16 @@ use crate::error::VortexFuzzResult; pub enum GpuEncodingKind { /// Dictionary encoding with GPU take support. Dict, + /// Dictionary whose nullable codes hold arbitrary values, possibly outside the dictionary, + /// under null positions. + DictGarbageNullCodes, } impl<'a> Arbitrary<'a> for GpuEncodingKind { fn arbitrary(u: &mut Unstructured<'a>) -> Result { - // Currently only Dict is supported - match u.int_in_range(0..=0)? { + match u.int_in_range(0..=1)? { 0 => Ok(GpuEncodingKind::Dict), + 1 => Ok(GpuEncodingKind::DictGarbageNullCodes), _ => unreachable!(), } } @@ -58,14 +69,83 @@ impl<'a> Arbitrary<'a> for FuzzCompressGpu { .0 .into_array() } + GpuEncodingKind::DictGarbageNullCodes => arbitrary_dict_with_garbage_null_codes(u)?, }; Ok(FuzzCompressGpu { array }) } } +/// Generate a random UTF-8 or binary DType. +fn arbitrary_gpu_string_dtype(u: &mut Unstructured) -> Result { + let nullability: Nullability = u.arbitrary()?; + Ok(if u.arbitrary()? { + DType::Utf8(nullability) + } else { + DType::Binary(nullability) + }) +} + +/// Generate a dictionary whose codes are in range at valid positions and arbitrary, including far +/// outside the dictionary, at null positions. +fn arbitrary_dict_with_garbage_null_codes(u: &mut Unstructured) -> Result { + let dtype = if u.arbitrary()? { + arbitrary_gpu_string_dtype(u)? + } else { + arbitrary_gpu_primitive_dtype(u)? + }; + let values_len = u.int_in_range(1..=20)?; + let values = ArbitraryArray::arbitrary_with_config( + u, + &ArbitraryArrayConfig { + dtype: Some(dtype), + len: values_len..=values_len, + }, + )? + .0; + let len = u.int_in_range(0..=200)?; + let valid: Vec = (0..len).map(|_| u.arbitrary()).collect::>()?; + let values_len = u32::try_from(values_len).map_err(|_| arbitrary::Error::IncorrectFormat)?; + let validity: Validity = valid.iter().copied().collect(); + let codes = match u.int_in_range(0..=2)? { + 0 => PrimitiveArray::new( + Buffer::from(garbage_codes::(u, &valid, values_len)?), + validity, + ), + 1 => PrimitiveArray::new( + Buffer::from(garbage_codes::(u, &valid, values_len)?), + validity, + ), + _ => PrimitiveArray::new( + Buffer::from(garbage_codes::(u, &valid, values_len)?), + validity, + ), + }; + let dict = DictArray::try_new(codes.into_array(), values) + .map_err(|_| arbitrary::Error::IncorrectFormat)?; + Ok(dict.into_array()) +} + +/// Codes in `0..values_len` at valid positions and any value of `T` at null positions. +fn garbage_codes<'a, T>(u: &mut Unstructured<'a>, valid: &[bool], values_len: u32) -> Result> +where + T: Arbitrary<'a> + TryFrom, +{ + valid + .iter() + .map(|&is_valid| { + if is_valid { + T::try_from(u.int_in_range(0..=values_len - 1)?) + .map_err(|_| arbitrary::Error::IncorrectFormat) + } else { + u.arbitrary::() + } + }) + .collect() +} + /// Generate a random primitive DType suitable for GPU operations. -fn arbitrary_gpu_primitive_dtype(u: &mut Unstructured) -> Result { +fn arbitrary_gpu_primitive_dtype(u: &mut Unstructured) -> Result { let nullability: Nullability = u.arbitrary()?; let ptype = match u.int_in_range(0..=9)? { 0 => PType::U8, @@ -80,7 +160,7 @@ fn arbitrary_gpu_primitive_dtype(u: &mut Unstructured) -> Result PType::F64, _ => unreachable!(), }; - Ok(vortex_array::dtype::DType::Primitive(ptype, nullability)) + Ok(DType::Primitive(ptype, nullability)) } /// Run the GPU decompression fuzzer. diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java index 7d867bd2c91..b6731040a14 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java @@ -37,7 +37,7 @@ public final class VortexIo implements Serializable { private static final int DEFAULT_READ_CONCURRENCY = 0; /** - * Maximum gap between merged reads, in bytes. Defaults to 64 KiB; zero merges only adjacent or overlapping reads. + * Maximum gap between merged reads, in bytes. Defaults to 1 MiB; zero merges only adjacent or overlapping reads. */ public static final String COALESCE_DISTANCE_OPTION = "vortex.read.coalesce.distance"; diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/NestedSchemaPruning.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/NestedSchemaPruning.java index 80660ff2bc1..e8c21cf96e9 100644 --- a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/NestedSchemaPruning.java +++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/NestedSchemaPruning.java @@ -155,8 +155,7 @@ private static Node struct( boolean filled = false; for (int i = 0; i < fields.length; i++) { String fieldPath = path.isEmpty() ? fields[i].name() : path + "." + fields[i].name(); - int index = lookup.indexOf( - fields[i].name(), resolution.fileName(), path.isEmpty() ? "" : path + "."); + int index = lookup.indexOf(fields[i].name(), resolution.fileName(), path.isEmpty() ? "" : path + "."); if (index < 0) { if (!fields[i].nullable()) { throw new IllegalArgumentException(String.format( diff --git a/vortex-array/src/arrays/runend/array/mod.rs b/vortex-array/src/arrays/runend/array/mod.rs index ad3a293c258..000963a85da 100644 --- a/vortex-array/src/arrays/runend/array/mod.rs +++ b/vortex-array/src/arrays/runend/array/mod.rs @@ -1,6 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Run-end arrays and their execution into logical values. +//! +//! Each end is an exclusive logical index paired with the value for its preceding run. Decoding +//! **must** apply a slice's logical offset before expanding those values. Specialized decoders +//! handle common storage types. The generic fallback gathers values to retain their logical dtype. + use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; @@ -52,6 +58,7 @@ use crate::vtable::VTable; use crate::vtable::ValidityVTable; mod plugin; +mod to_dict; /// A [`RunEnd`]-encoded Vortex array. pub type RunEndArray = Array; @@ -482,7 +489,7 @@ pub(super) fn run_end_canonicalize( .execute_as::("values", ctx)?; runend_decode_listview(pends, values, array.offset(), array.len())?.into_array() } - _ => vortex_bail!("Unsupported RunEnd value type: {}", array.dtype()), + _ => array.to_dict(ctx)?.into_array(), }) } diff --git a/vortex-array/src/arrays/runend/array/to_dict.rs b/vortex-array/src/arrays/runend/array/to_dict.rs new file mode 100644 index 00000000000..511cb479bbd --- /dev/null +++ b/vortex-array/src/arrays/runend/array/to_dict.rs @@ -0,0 +1,281 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Express run-end decoding as a gather from the existing values child. +//! +//! Repeated run indices let the generic take machinery handle logical types without specialized +//! run-end decoders. Piecewise codes keep that mapping compact until a consumer needs flat indices. + +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::arrays::DictArray; +use crate::arrays::PiecewiseSequenceArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::runend::RunEndArray; +use crate::arrays::runend::RunEndArrayExt; +use crate::arrays::runend::RunEndArraySlotsExt; +use crate::arrays::runend::slice_relative_ends; +use crate::dtype::Nullability; +use crate::dtype::UnsignedPType; +use crate::match_each_unsigned_integer_ptype; + +impl RunEndArray { + /// Converts to a dictionary with the same logical dtype and slice. + /// + /// Non-empty arrays share the values child without executing it. Codes use one constant piece + /// per overlapping run, so construction needs no expanded index buffer. Executing the resulting + /// [`DictArray`] may still materialize those indices. Empty arrays retain no source children. + pub fn to_dict(&self, ctx: &mut ExecutionCtx) -> VortexResult { + if self.is_empty() { + return DictArray::try_new( + PrimitiveArray::empty::(Nullability::NonNullable).into_array(), + Canonical::empty(self.dtype()).into_array(), + ); + } + + let ends = self.ends().clone().execute::(ctx)?; + let codes = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { + run_indices(ends.as_slice::(), self.offset(), self.len())? + }); + + DictArray::try_new(codes.into_array(), self.values().clone()) + } +} + +/// Builds repeated physical value indices for a slice covered by validated run ends. +fn run_indices( + ends: &[E], + offset: usize, + len: usize, +) -> VortexResult { + let mut starts = BufferMut::::with_capacity(ends.len()); + let mut lengths = BufferMut::::with_capacity(ends.len()); + let mut run_start = 0; + + for (index, run_end) in slice_relative_ends(ends, offset, len).enumerate() { + let length = run_end - run_start; + if length != 0 { + starts.push(u64::try_from(index)?); + lengths.push(u64::try_from(length)?); + } + run_start = run_end; + + if run_end == len { + break; + } + } + + let multipliers = ConstantArray::new(0u64, starts.len()).into_array(); + + PiecewiseSequenceArray::try_new( + starts.freeze().into_array(), + lengths.freeze().into_array(), + multipliers, + len, + ) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::ArrayRef; + use crate::ExecutionCtx; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::ExtensionArray; + use crate::arrays::FixedSizeListArray; + use crate::arrays::PiecewiseSequence; + use crate::arrays::PrimitiveArray; + use crate::arrays::StructArray; + use crate::arrays::dict::DictArraySlotsExt; + use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; + use crate::arrays::runend::RunEnd; + use crate::arrays::runend::RunEndArray; + use crate::arrays::runend::RunEndArraySlotsExt; + use crate::arrays::runend::tests::SESSION; + use crate::assert_arrays_eq; + use crate::builtins::ArrayBuiltins; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::scalar::Scalar; + + #[track_caller] + fn check_conversion( + array: &RunEndArray, + expected: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let dict = array.to_dict(ctx)?; + + if array.is_empty() { + assert!(dict.codes().is_empty()); + assert!(dict.values().is_empty()); + } else { + assert!(ArrayRef::ptr_eq(dict.values(), array.values())); + let codes = dict.codes().as_::(); + assert_eq!(codes.multipliers().as_constant(), Some(&Scalar::from(0u64))); + } + + assert_arrays_eq!(dict.into_array(), expected, ctx); + + Ok(()) + } + + fn struct_values(storage: ArrayRef) -> VortexResult { + Ok(StructArray::from_fields(&[("value", storage)])?.into_array()) + } + + fn fixed_size_list_values(storage: ArrayRef) -> VortexResult { + Ok( + FixedSizeListArray::try_new(storage.clone(), 1, storage.validity()?, storage.len())? + .into_array(), + ) + } + + fn timestamp_values(storage: ArrayRef) -> VortexResult { + let ext_dtype = Timestamp::try_new_with_tz( + TimeUnit::Microseconds, + Some("UTC".into()), + storage.dtype().nullability(), + )?; + + Ok(ExtensionArray::try_new(ext_dtype.erased(), storage)?.into_array()) + } + + #[rstest] + #[case::full(PType::U32, 0, 8)] + #[case::clipped_u8(PType::U8, 2, 5)] + #[case::clipped_u16(PType::U16, 2, 5)] + #[case::clipped_u32(PType::U32, 2, 5)] + #[case::clipped_u64(PType::U64, 2, 5)] + #[case::run_boundary(PType::U32, 3, 2)] + #[case::single_row(PType::U32, 1, 1)] + #[case::empty_retained_runs(PType::U32, 8, 0)] + fn sliced_conversion( + #[case] ends_type: PType, + #[case] offset: usize, + #[case] len: usize, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let ends = buffer![3u64, 5, 8] + .into_array() + .cast(DType::Primitive(ends_type, Nullability::NonNullable))?; + let values = PrimitiveArray::from_option_iter([Some(10i64), None, Some(30)]).into_array(); + let array = RunEnd::try_new_offset_length(ends, values, offset, len, &mut ctx)?; + let expanded = [ + Some(10i64), + Some(10), + Some(10), // First run. + None, + None, // Null run. + Some(30), + Some(30), + Some(30), // Last run. + ]; + let expected = + PrimitiveArray::from_option_iter(expanded[offset..offset + len].iter().copied()); + + check_conversion(&array, expected.into_array(), &mut ctx) + } + + #[test] + fn codes_retain_only_overlapping_runs() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new_offset_length( + buffer![3u32, 5, 8].into_array(), + buffer![10i64, 20, 30].into_array(), + 3, + 4, + &mut ctx, + )?; + let dict = array.to_dict(&mut ctx)?; + let codes = dict.codes().as_::(); + + assert_arrays_eq!(codes.starts(), buffer![1u64, 2].into_array(), &mut ctx); + assert_arrays_eq!(codes.lengths(), buffer![2u64, 2].into_array(), &mut ctx); + + assert_arrays_eq!( + dict.into_array(), + buffer![20i64, 20, 30, 30].into_array(), + &mut ctx + ); + + Ok(()) + } + + #[test] + fn empty_conversion() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::empty::(Nullability::Nullable).into_array(); + let array = RunEnd::try_new( + PrimitiveArray::empty::(Nullability::NonNullable).into_array(), + values.clone(), + &mut ctx, + )?; + + check_conversion(&array, values, &mut ctx) + } + + #[test] + fn repeated_slicing() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new( + buffer![3u32, 5, 8].into_array(), + PrimitiveArray::from_option_iter([Some(10i64), None, Some(30)]).into_array(), + &mut ctx, + )?; + let sliced = array.slice(2..8)?.execute_until::(&mut ctx)?; + let sliced = sliced + .slice(1..5)? + .execute_until::(&mut ctx)? + .downcast::(); + let expected = PrimitiveArray::from_option_iter([None, None, Some(30i64), Some(30)]); + + check_conversion(&sliced, expected.into_array(), &mut ctx) + } + + #[rstest] + #[case::structure(struct_values)] + #[case::fixed_size_list(fixed_size_list_values)] + #[case::extension(timestamp_values)] + fn fallback_preserves_dtype( + #[case] wrap: fn(ArrayRef) -> VortexResult, + #[values(false, true)] empty: bool, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let storage = PrimitiveArray::from_option_iter([Some(10i64), None, Some(30)]).into_array(); + let values = wrap(storage)?; + let len = if empty { 0 } else { 5 }; + let expected = wrap( + PrimitiveArray::from_option_iter( + [Some(10i64), None, None, Some(30), Some(30)] + .into_iter() + .take(len), + ) + .into_array(), + )?; + let array = RunEnd::try_new_offset_length( + buffer![3u32, 5, 8].into_array(), + values, + 2, + len, + &mut ctx, + )?; + + check_conversion(&array, expected.clone(), &mut ctx)?; + assert_arrays_eq!(array.into_array(), expected, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-array/src/arrays/runend/compress.rs b/vortex-array/src/arrays/runend/compress.rs index 2c929eccb7f..9aa319c26bf 100644 --- a/vortex-array/src/arrays/runend/compress.rs +++ b/vortex-array/src/arrays/runend/compress.rs @@ -25,7 +25,7 @@ use crate::arrays::VarBinViewArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::decimal::DecimalArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::arrays::runend::iter::trimmed_ends_iter; +use crate::arrays::runend::iter::slice_relative_ends; use crate::buffer::BufferHandle; use crate::dtype::NativePType; use crate::dtype::Nullability; @@ -204,7 +204,7 @@ pub fn runend_decode_primitive( Ok(match_each_native_ptype!(values.ptype(), |P| { match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_typed_primitive( - trimmed_ends_iter(ends.as_slice::(), offset, length), + slice_relative_ends(ends.as_slice::(), offset, length), values.as_slice::

(), validity_mask, values.dtype().nullability(), @@ -381,7 +381,7 @@ pub fn runend_decode_decimal( Ok(match_each_decimal_value_type!(values.values_type(), |D| { let (decoded, validity) = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_slice( - trimmed_ends_iter(ends.as_slice::(), offset, length), + slice_relative_ends(ends.as_slice::(), offset, length), values.buffer::().as_slice(), validity_mask, values.dtype().nullability(), @@ -414,7 +414,7 @@ pub fn runend_decode_varbinview( let (decoded_views, validity) = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_slice( - trimmed_ends_iter(ends.as_slice::(), offset, length), + slice_relative_ends(ends.as_slice::(), offset, length), views, validity_mask, values.dtype().nullability(), diff --git a/vortex-array/src/arrays/runend/compute/take.rs b/vortex-array/src/arrays/runend/compute/take.rs index c7a568fba6f..cbe73efd869 100644 --- a/vortex-array/src/arrays/runend/compute/take.rs +++ b/vortex-array/src/arrays/runend/compute/take.rs @@ -20,7 +20,7 @@ use crate::arrays::dict::TakeExecute; use crate::arrays::runend::RunEnd; use crate::arrays::runend::array::RunEndArrayExt; use crate::arrays::runend::array::RunEndArraySlotsExt; -use crate::arrays::runend::iter::trimmed_ends_iter; +use crate::arrays::runend::iter::slice_relative_ends; use crate::dtype::UnsignedPType; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; @@ -378,7 +378,7 @@ where { let mut table = BufferMut::with_capacity(len); let mut run_start = 0; - for (run_idx, run_end) in trimmed_ends_iter(ends, offset, len).enumerate() { + for (run_idx, run_end) in slice_relative_ends(ends, offset, len).enumerate() { table.push_n(run_idx.as_(), run_end - run_start); run_start = run_end; } diff --git a/vortex-array/src/arrays/runend/decompress_bool.rs b/vortex-array/src/arrays/runend/decompress_bool.rs index 739fb968276..3136631de2d 100644 --- a/vortex-array/src/arrays/runend/decompress_bool.rs +++ b/vortex-array/src/arrays/runend/decompress_bool.rs @@ -18,7 +18,7 @@ use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; -use crate::arrays::runend::iter::trimmed_ends_iter; +use crate::arrays::runend::iter::slice_relative_ends; use crate::dtype::DType; use crate::dtype::Nullability; use crate::match_each_unsigned_integer_ptype; @@ -74,7 +74,7 @@ pub fn runend_decode_bools( Ok(match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_typed_bool( - trimmed_ends_iter(ends.as_slice::(), offset, length), + slice_relative_ends(ends.as_slice::(), offset, length), &values_buf, validity, nullability, @@ -112,7 +112,7 @@ pub fn runend_decode_typed_bool( } /// Fast path for few runs with no offset. Uses direct slice access to minimize overhead. -/// This avoids the `trimmed_ends_iter` iterator chain which adds significant overhead +/// This avoids the `slice_relative_ends` iterator chain which adds significant overhead /// for small numbers of runs. #[allow(clippy::inline_always)] #[inline(always)] diff --git a/vortex-array/src/arrays/runend/iter.rs b/vortex-array/src/arrays/runend/iter.rs index d41ee2da87c..925725959c7 100644 --- a/vortex-array/src/arrays/runend/iter.rs +++ b/vortex-array/src/arrays/runend/iter.rs @@ -7,7 +7,15 @@ use vortex_error::vortex_panic; use crate::dtype::IntegerPType; -pub fn trimmed_ends_iter( +/// Yields exclusive run ends relative to a logical slice. +/// +/// Subtracts `offset` from each end and clamps it to `length`. One end is yielded per input run, +/// including zero-length runs at the slice boundaries, so indices still match the values child. +/// +/// # Panics +/// +/// Panics if `offset` or `length` cannot be represented by `E`, or an input end is before `offset`. +pub fn slice_relative_ends( run_ends: &[E], offset: usize, length: usize, diff --git a/vortex-array/src/arrays/runend/mod.rs b/vortex-array/src/arrays/runend/mod.rs index 84914068f73..77a8f94335c 100644 --- a/vortex-array/src/arrays/runend/mod.rs +++ b/vortex-array/src/arrays/runend/mod.rs @@ -6,7 +6,7 @@ mod arbitrary; #[cfg(feature = "arbitrary")] pub use arbitrary::ArbitraryRunEndArray; pub use array::*; -pub use iter::trimmed_ends_iter; +pub use iter::slice_relative_ends; mod array; pub mod compress; diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 1df5dddb5f2..daff5f2a957 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -25,6 +25,7 @@ use vortex_utils::aliases::hash_map::HashTableEntry; use super::DictConstraints; use super::DictEncoder; +use super::codes_ptype_for_max_len; use crate::ArrayRef; use crate::ArrayView; use crate::IntoArray; @@ -59,18 +60,18 @@ pub fn bytes_dict_builder( constraints: &DictConstraints, allocator: BufferAllocatorRef, ) -> Box { - match constraints.max_len as u64 { - max if max <= u8::MAX as u64 => Box::new(BytesDictBuilder::::new_in( + match codes_ptype_for_max_len(constraints.max_len as u64) { + PType::U8 => Box::new(BytesDictBuilder::::new_in( dtype, constraints, allocator, )), - max if max <= u16::MAX as u64 => Box::new(BytesDictBuilder::::new_in( + PType::U16 => Box::new(BytesDictBuilder::::new_in( dtype, constraints, allocator, )), - max if max <= u32::MAX as u64 => Box::new(BytesDictBuilder::::new_in( + PType::U32 => Box::new(BytesDictBuilder::::new_in( dtype, constraints, allocator, @@ -363,6 +364,7 @@ mod test { use std::sync::Arc; use std::sync::LazyLock; + use rstest::rstest; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::ByteBuffer; @@ -370,6 +372,7 @@ mod test { use vortex_session::VortexSession; use super::BytesDictBuilder; + use super::bytes_dict_builder; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; @@ -379,11 +382,13 @@ mod test { use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; use crate::buffer::BufferHandle; + use crate::builders::dict::DictConstraints; use crate::builders::dict::UNCONSTRAINED; use crate::builders::dict::dict_encode; use crate::builders::dict::dict_encoder_in; use crate::dtype::DType; use crate::dtype::Nullability; + use crate::dtype::PType; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -510,6 +515,25 @@ mod test { Ok(()) } + #[rstest] + #[case(usize::from(u8::MAX) + 1, PType::U8)] + #[case(usize::from(u8::MAX) + 2, PType::U16)] + #[case(usize::from(u16::MAX) + 1, PType::U16)] + #[case(usize::from(u16::MAX) + 2, PType::U32)] + fn selects_narrowest_code_type(#[case] max_len: usize, #[case] expected: PType) { + let constraints = DictConstraints { + max_bytes: usize::MAX, + max_len, + }; + let encoder = bytes_dict_builder( + DType::Utf8(Nullability::NonNullable), + &constraints, + BufferAllocatorRef::statically_allocated(), + ); + + assert_eq!(encoder.codes_ptype(), expected); + } + #[test] fn max_dict_bytes_cannot_exceed_the_view_offset_range() { let builder = BytesDictBuilder::::new_in( diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index 8e75ea7a903..b2bc39db1e4 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -23,7 +23,9 @@ mod primitive; #[derive(Clone)] pub struct DictConstraints { + /// Limits the encoded storage for dictionary values. pub max_bytes: usize, + /// Limits the number of dictionary values. pub max_len: usize, } @@ -32,6 +34,19 @@ pub const UNCONSTRAINED: DictConstraints = DictConstraints { max_len: usize::MAX, }; +/// Returns the narrowest unsigned code type that can index `max_len` dictionary values. +/// +/// Codes count from zero, so a code type indexes one more value than its maximum: `u8` codes +/// index up to 256 values. +pub const fn codes_ptype_for_max_len(max_len: u64) -> PType { + match max_len { + len if len <= 1 << u8::BITS => PType::U8, + len if len <= 1 << u16::BITS => PType::U16, + len if len <= 1 << u32::BITS => PType::U32, + _ => PType::U64, + } +} + pub trait DictEncoder: Send { /// Assign dictionary codes to the given input array. fn encode(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index e4fc05bf4bc..30e89093216 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -10,9 +10,9 @@ use arrow_array::cast::AsArray; use arrow_array::new_null_array; use arrow_array::types::*; use arrow_schema::DataType; +use arrow_schema::Field; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Dict; @@ -24,6 +24,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::ArrowExportOptions; +use crate::ExecuteArrowExt; use crate::executor::execute_arrow_naive; /// Matches the encodings [`to_arrow_dictionary`] requires for export. @@ -40,27 +41,36 @@ impl Matcher for ArrowDictExportable { pub(super) fn to_arrow_dictionary( array: ArrayRef, codes_type: &DataType, - values_type: &DataType, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let array = array.execute_until::(ctx)?; + // Keep an existing empty dictionary: canonicalization can discard its codes, and Arrow + // cannot repack every extension's storage type into a dictionary. + let array = if array.is::() || array.is::() { + array + } else { + array.execute_until::(ctx)? + }; let array = match array.try_downcast::() { - Ok(dict) => return dict_to_dict(dict, codes_type, values_type, options, ctx), + Ok(dict) => return dict_to_dict(dict, codes_type, values_field, options, ctx), Err(array) => array, }; let array = match array.try_downcast::() { - Ok(constant) => return constant_to_dict(constant, codes_type, values_type, options, ctx), + Ok(constant) => return constant_to_dict(constant, codes_type, values_field, options, ctx), Err(array) => array, }; // Otherwise, we should try and build a dictionary. // Arrow hides this functionality inside the cast module! - let array = execute_arrow_naive(array, Some(values_type), options, ctx)?; + let array = array.execute_arrow_with_options(Some(values_field), options, ctx)?; arrow_cast::cast( &array, - &DataType::Dictionary(Box::new(codes_type.clone()), Box::new(values_type.clone())), + &DataType::Dictionary( + Box::new(codes_type.clone()), + Box::new(values_field.data_type().clone()), + ), ) .map_err(VortexError::from) } @@ -69,21 +79,22 @@ pub(super) fn to_arrow_dictionary( fn constant_to_dict( array: ConstantArray, codes_type: &DataType, - values_type: &DataType, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); let scalar = array.scalar(); if scalar.is_null() { - let dict_type = - DataType::Dictionary(Box::new(codes_type.clone()), Box::new(values_type.clone())); + let dict_type = DataType::Dictionary( + Box::new(codes_type.clone()), + Box::new(values_field.data_type().clone()), + ); return Ok(new_null_array(&dict_type, len)); } - let values = execute_arrow_naive( - ConstantArray::new(scalar.clone(), 1).into_array(), - Some(values_type), + let values = ConstantArray::new(scalar.clone(), 1).execute_arrow_with_options( + Some(values_field), options, ctx, )?; @@ -95,12 +106,17 @@ fn constant_to_dict( fn dict_to_dict( array: DictArray, codes_type: &DataType, - values_type: &DataType, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let codes = execute_arrow_naive(array.codes().clone(), Some(codes_type), options, ctx)?; - let values = execute_arrow_naive(array.values().clone(), Some(values_type), options, ctx)?; + let values = + array + .values() + .clone() + .execute_arrow_with_options(Some(values_field), options, ctx)?; + make_dict_array(codes_type, codes, values) } @@ -169,20 +185,54 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::Nullable; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::mask::Mask; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; + use crate::ArrowSessionExt; use crate::ExecuteArrowExt; use crate::executor::dictionary::ConstantArray; use crate::executor::dictionary::DictArray; + #[track_caller] + fn check_fixed_size_list_dictionary_round_trip(len: usize) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0..i32::try_from(len * 2)?).into_array(), + 2, + Validity::NonNullable, + len, + )? + .into_array(); + let source = DictArray::try_new( + PrimitiveArray::from_iter(0..u32::try_from(len)?).into_array(), + values, + )? + .into_array(); + let arrow = SESSION.arrow(); + let values_field = arrow.to_arrow_field("item", source.dtype())?; + let field = Field::new( + "dict", + dict_type(DataType::Int32, values_field.data_type().clone()), + false, + ); + + let exported = source.clone().execute_arrow(Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + let imported = arrow.from_arrow_array(exported, &field)?; + vortex_array::assert_arrays_eq!(imported, source, &mut ctx); + + Ok(()) + } + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); fn dict_type(codes: DataType, values: DataType) -> DataType { @@ -269,4 +319,11 @@ mod tests { assert_eq!(flat.as_ref(), &expected as &dyn arrow_array::Array); Ok(()) } + + #[rstest] + #[case::empty(0)] + #[case::nonempty(3)] + fn fixed_size_list_values_round_trip(#[case] len: usize) -> VortexResult<()> { + check_fixed_size_list_dictionary_round_trip(len) + } } diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index 803f038962f..386c8723dc9 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -173,6 +173,11 @@ fn execute_arrow_array( } }; + if let DataType::Dictionary(codes, values) = target_field.data_type() { + let values_field = target_field.clone().with_data_type(values.as_ref().clone()); + return to_arrow_dictionary(array, codes, &values_field, options, ctx); + } + if let Some(arrow_ext_name) = target_field.extension_type_name() { return execute_arrow_extension(array, arrow_ext_name, target_field, options, ctx); } @@ -297,9 +302,9 @@ fn execute_arrow_ext_storage( /// ignoring nullability, after every nested extension type is replaced by its storage type. The /// types at the top level are compared as they are. /// -/// An export compares the type that its target imports as with the storage type this way, because -/// an import can lose the metadata of a nested extension type, such as the version of a UUID, or -/// not know its Vortex ID. The export of each nested field checks that field when it reaches it. +/// An explicit target may omit nested extension metadata or name an unregistered Vortex ID, so +/// export compares its inferred type with the storage layout. Each nested field is checked when +/// it is exported. fn eq_storage_layout(lhs: &DType, rhs: &DType) -> bool { match (lhs, rhs) { (DType::List(lhs, _), DType::List(rhs, _)) => eq_nested_storage_layout(lhs, rhs), @@ -402,9 +407,17 @@ pub(crate) fn execute_arrow_naive( to_arrow_struct(array, fields, options, ctx) } // TODO(joe): pass down preferred - DataType::Dictionary(codes_type, values_type) => { - to_arrow_dictionary(array, codes_type, values_type, options, ctx) - } + DataType::Dictionary(codes_type, values_type) => to_arrow_dictionary( + array.clone(), + codes_type, + &Field::new( + "", + values_type.as_ref().clone(), + array.dtype().is_nullable(), + ), + options, + ctx, + ), dt @ DataType::Decimal32(..) => to_arrow_decimal(array, dt, ctx), dt @ DataType::Decimal64(..) => to_arrow_decimal(array, dt, ctx), dt @ DataType::Decimal128(..) => to_arrow_decimal(array, dt, ctx), diff --git a/vortex-arrow/src/executor/run_end.rs b/vortex-arrow/src/executor/run_end.rs index 04801861ea0..0b8e33a0760 100644 --- a/vortex-arrow/src/executor/run_end.rs +++ b/vortex-arrow/src/executor/run_end.rs @@ -1,11 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Export run-end arrays without expanding their logical values. +//! +//! Vortex slices retain absolute run ends. Adjusting before narrowing to Arrow's requested integer +//! width keeps small slices exportable even when their source indices are large. The values field +//! retains extension metadata so imports recover the same logical type. + use std::sync::Arc; +use arrow_array::Array as _; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::RunArray; use arrow_array::cast::AsArray; +use arrow_array::make_array; use arrow_array::new_null_array; use arrow_array::types::*; use arrow_buffer::ArrowNativeType; @@ -16,10 +24,12 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::RunEnd; use vortex_array::arrays::RunEndArray; use vortex_array::arrays::runend::RunEndArrayExt; use vortex_array::arrays::runend::RunEndArraySlotsExt; +use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::matcher::Matcher; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -44,7 +54,35 @@ impl Matcher for ArrowRunEndExportable { pub(super) fn to_arrow_run_end( array: ArrayRef, ends_type: &DataType, - values_type: &Field, + values_field: &Field, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let data_type = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", ends_type.clone(), false)), + Arc::new(values_field.clone()), + ); + if array.is_empty() { + return Ok(new_null_array(&data_type, 0)); + } + + // Arrow's RunArray constructors synthesize a values field without extension metadata. + // Retain the caller's field so the array remains consistent with its export schema. + let exported = build_run_end_array(array, ends_type, values_field, options, ctx)?; + + Ok(make_array( + exported + .to_data() + .into_builder() + .data_type(data_type) + .build()?, + )) +} + +fn build_run_end_array( + array: ArrayRef, + ends_type: &DataType, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -52,118 +90,108 @@ pub(super) fn to_arrow_run_end( let array = match array.try_downcast::() { Ok(constant) => { - return constant_to_run_end(constant, ends_type, values_type, options, ctx); + return constant_to_run_end(constant, ends_type, values_field, options, ctx); } Err(array) => array, }; let array = match array.try_downcast::() { - Ok(run_end) => return run_end_to_arrow(run_end, ends_type, values_type, options, ctx), + Ok(run_end) => return run_end_to_arrow(run_end, ends_type, values_field, options, ctx), Err(array) => array, }; // Fallback: canonicalize to flat Arrow, then cast to REE. - let flat = export_values(array, values_type, options, ctx)?; + let flat = array.execute_arrow_with_options(Some(values_field), options, ctx)?; let ree_type = DataType::RunEndEncoded( Arc::new(Field::new("run_ends", ends_type.clone(), false)), - Arc::new(values_type.clone()), + Arc::new(values_field.clone()), ); arrow_cast::cast(&flat, &ree_type).map_err(VortexError::from) } -/// Export the values of a run-end array through the session, so extension metadata on the target -/// values field dispatches to its export plugin instead of the naive Arrow mapping. -fn export_values( - values: ArrayRef, - values_type: &Field, - options: &ArrowExportOptions, - ctx: &mut ExecutionCtx, -) -> VortexResult { - values.execute_arrow_with_options(Some(values_type), options, ctx) -} - fn run_end_to_arrow( array: RunEndArray, ends_type: &DataType, - values_type: &Field, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let length = array.len(); - let offset = array.offset(); + let ends = array.ends().clone().execute::(ctx)?; + let (runs, ends) = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { + let run_ends = ends.as_slice::(); + let length = E::try_from(array.len())?; + + if array.offset() == 0 + && run_ends.first().is_some_and(|&end| end > 0) + && run_ends.last() == Some(&length) + { + (0..run_ends.len(), ends.clone().into_array()) + } else { + let offset = E::try_from(array.offset())?; + + // A run ending at the offset contributes no rows. Retaining it would give Arrow a + // zero end, and retaining trailing runs would give duplicate ends after clamping. + let start = run_ends.partition_point(|&end| end <= offset); + let end = start + run_ends[start..].partition_point(|&end| end - offset < length); + let end = (end + 1).min(run_ends.len()); + let adjusted = PrimitiveArray::from_iter( + run_ends[start..end] + .iter() + .map(|&end| (end - offset).min(length)), + ) + .into_array(); - let arrow_ends = execute_arrow_naive(array.ends().clone(), Some(ends_type), options, ctx)?; - let arrow_values = export_values(array.values().clone(), values_type, options, ctx)?; + (start..end, adjusted) + } + }); + let values = if runs.start == 0 && runs.end == array.values().len() { + array.values().clone() + } else { + array.values().slice(runs)? + }; + + // Narrow only the relative ends: absolute indices can exceed the target's range even when + // this slice fits. Unadjusted ends and values keep the existing buffer-sharing export path. + let arrow_ends = execute_arrow_naive(ends, Some(ends_type), options, ctx)?; + let arrow_values = values.execute_arrow_with_options(Some(values_field), options, ctx)?; match ends_type { - DataType::Int16 => build_run_array::(&arrow_ends, &arrow_values, offset, length), - DataType::Int32 => build_run_array::(&arrow_ends, &arrow_values, offset, length), - DataType::Int64 => build_run_array::(&arrow_ends, &arrow_values, offset, length), + DataType::Int16 => Ok(Arc::new(RunArray::::try_new( + arrow_ends.as_primitive::(), + &arrow_values, + )?)), + DataType::Int32 => Ok(Arc::new(RunArray::::try_new( + arrow_ends.as_primitive::(), + &arrow_values, + )?)), + DataType::Int64 => Ok(Arc::new(RunArray::::try_new( + arrow_ends.as_primitive::(), + &arrow_values, + )?)), _ => vortex_bail!("Unsupported run-end index type: {:?}", ends_type), } } -fn build_run_array( - ends: &ArrowArrayRef, - values: &ArrowArrayRef, - offset: usize, - length: usize, -) -> VortexResult -where - R::Native: std::ops::Sub + Ord, -{ - let offset_native = R::Native::from_usize(offset) - .ok_or_else(|| vortex_err!("Offset {offset} exceeds run-end index capacity"))?; - let length_native = R::Native::from_usize(length) - .ok_or_else(|| vortex_err!("Length {length} exceeds run-end index capacity"))?; - - let ends_prim = ends.as_primitive::(); - if offset == 0 && ends_prim.values().last() == Some(&length_native) { - // Fast path: no trimming or adjustment needed. - return Ok(Arc::new(RunArray::::try_new(ends_prim, values)?) as ArrowArrayRef); - } - - // Trim to only include runs covering the [offset, offset+length) range. - // Runs beyond this would produce duplicate adjusted ends, violating - // Arrow's strict-ordering requirement for RunArray. - // Run ends are strictly increasing, so we can binary search. - let num_runs = (ends_prim - .values() - .partition_point(|&e| e - offset_native < length_native) - + 1) - .min(ends_prim.len()); - - let trimmed_ends = ends.slice(0, num_runs); - let trimmed_values = values.slice(0, num_runs); - - let adjusted = trimmed_ends - .as_primitive::() - .unary(|end| (end - offset_native).min(length_native)); - - Ok(Arc::new(RunArray::::try_new(&adjusted, &trimmed_values)?) as ArrowArrayRef) -} - /// Convert a constant array to a run-end encoded array with a single run. fn constant_to_run_end( array: ConstantArray, ends_type: &DataType, - values_type: &Field, + values_field: &Field, options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); let scalar = array.scalar(); - if scalar.is_null() || len == 0 { + if scalar.is_null() { let ree_type = DataType::RunEndEncoded( Arc::new(Field::new("run_ends", ends_type.clone(), false)), - Arc::new(values_type.clone()), + Arc::new(values_field.clone()), ); return Ok(new_null_array(&ree_type, len)); } - let values = export_values( - ConstantArray::new(scalar.clone(), 1).into_array(), - values_type, + let values = ConstantArray::new(scalar.clone(), 1).execute_arrow_with_options( + Some(values_field), options, ctx, )?; @@ -201,13 +229,18 @@ mod tests { use arrow_schema::DataType; use arrow_schema::Field; use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SliceArray; use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; use vortex_array::dtype::Nullability::Nullable; use vortex_array::dtype::PType; + use vortex_array::extension::datetime::Date; + use vortex_array::extension::datetime::TimeUnit; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; @@ -215,10 +248,55 @@ mod tests { use vortex_error::vortex_err; use vortex_session::VortexSession; + use crate::ArrowSessionExt; use crate::ExecuteArrowExt; use crate::executor::run_end::ConstantArray; use crate::executor::run_end::RunEnd; + #[track_caller] + fn check_temporal_run_end_round_trip(len: usize, sliced: bool) -> VortexResult<()> { + let session = VortexSession::empty(); + let arrow = session.arrow(); + let mut ctx = session.create_execution_ctx(); + let dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); + let values = ExtensionArray::try_new( + dtype, + PrimitiveArray::from_iter((0..i32::try_from(len)?).map(|i| 7 + i * 4)).into_array(), + )? + .into_array(); + let source = RunEnd::try_new( + PrimitiveArray::from_iter((1..=u32::try_from(len)?).map(|i| i * 2)).into_array(), + values.clone(), + &mut ctx, + )? + .into_array(); + let source = if sliced && !source.is_empty() { + source.slice(1..source.len())? + } else { + source + }; + + let start = u32::from(sliced && len > 0); + let expected = values.take( + PrimitiveArray::from_iter((start..u32::try_from(len * 2)?).map(|i| i / 2)).into_array(), + )?; + let values_field = arrow.to_arrow_field("dates", source.dtype())?; + let field = Field::new( + "dates", + DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(values_field), + ), + false, + ); + let exported = source.execute_arrow(Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + let imported = arrow.from_arrow_array(exported, &field)?; + vortex_array::assert_arrays_eq!(imported, expected, &mut ctx); + + Ok(()) + } + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); fn ree_type(ends: DataType, values_dtype: DataType) -> DataType { @@ -228,10 +306,7 @@ mod tests { ) } - fn execute( - array: vortex_array::ArrayRef, - dt: &DataType, - ) -> VortexResult { + fn execute(array: ArrayRef, dt: &DataType) -> VortexResult { let field = Field::new("", dt.clone(), array.dtype().is_nullable()); array.execute_arrow(Some(&field), &mut SESSION.create_execution_ctx()) } @@ -284,7 +359,7 @@ mod tests { ), )] fn constant_to_ree( - #[case] input: vortex_array::ArrayRef, + #[case] input: ArrayRef, #[case] target_type: DataType, #[case] expected: arrow_array::ArrayRef, ) -> VortexResult<()> { @@ -344,37 +419,47 @@ mod tests { Ok(()) } - /// Regression: build_run_array must trim excess trailing runs and - /// respect the `length` parameter. This happens when a vortex - /// RunEndArray is sliced to fewer rows than the physical run_ends cover. #[rstest] - #[case::offset_zero(0, 5, &[3, 5], &[100, 200])] - #[case::nonzero_offset(2, 3, &[1, 3], &[100, 200])] - #[case::all_runs_needed_but_last_exceeds(0, 8, &[3, 5, 8], &[100, 200, 300])] - fn build_run_array_trims_excess_runs( + #[case::offset_zero( &[3, 5, 10], 0, 5, &[3, 5], &[100, 200])] + #[case::nonzero_offset( &[3, 5, 10], 2, 3, &[1, 3], &[100, 200])] + #[case::clipped_last_run( &[3, 5, 10], 0, 8, &[3, 5, 8], &[100, 200, 300])] + #[case::run_boundary( &[3, 5, 10], 3, 4, &[2, 4], &[200, 300])] + #[case::large_absolute_ends(&[40000, 40002], 39999, 3, &[1, 3], &[100, 200])] + #[case::empty_retained_runs(&[3, 5, 10], 11, 0, &[], &[])] + fn run_end_export_trims_runs( + #[case] ends: &[u32], #[case] offset: usize, #[case] length: usize, - #[case] expected_ends: &[i32], + #[case] expected_ends: &[i16], #[case] expected_values: &[i64], ) -> VortexResult<()> { - // 3 runs covering 10 rows: [0..3), [3..5), [5..10) - let ends: arrow_array::ArrayRef = Arc::new(Int32Array::from(vec![3i32, 5, 10])); - let values: arrow_array::ArrayRef = Arc::new(Int64Array::from(vec![100i64, 200, 300])); + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new_offset_length( + PrimitiveArray::from_iter(ends.iter().copied()).into_array(), + PrimitiveArray::from_iter((0..ends.len()).map(|i| (i as i64 + 1) * 100)).into_array(), + offset, + length, + &mut ctx, + )?; - let result = super::build_run_array::(&ends, &values, offset, length)?; + let result = execute( + array.into_array(), + &ree_type(DataType::Int16, DataType::Int64), + )?; assert_eq!(result.len(), length); + let expected = RunArray::::try_new( + &Int16Array::from(expected_ends.to_vec()), + &Int64Array::from(expected_values.to_vec()), + )?; + assert_eq!(result.as_ref(), &expected); - let ree = result - .as_any() - .downcast_ref::>() - .ok_or_else(|| vortex_err!("expected Int32 run-end array"))?; - assert_eq!(ree.run_ends().values(), expected_ends); - let values = ree - .values() - .as_any() - .downcast_ref::() - .ok_or_else(|| vortex_err!("expected Int64 values"))?; - assert_eq!(values.values(), expected_values); Ok(()) } + + #[rstest] + #[case::empty(0, false)] + #[case::sliced(3, true)] + fn temporal_run_end_round_trip(#[case] len: usize, #[case] sliced: bool) -> VortexResult<()> { + check_temporal_run_end_round_trip(len, sliced) + } } diff --git a/vortex-arrow/src/executor/struct_.rs b/vortex-arrow/src/executor/struct_.rs index 8559a0a327e..c48b119835d 100644 --- a/vortex-arrow/src/executor/struct_.rs +++ b/vortex-arrow/src/executor/struct_.rs @@ -170,14 +170,12 @@ fn create_from_fields( arrow_arrays.push(arrow_field); } - Ok(Arc::new(unsafe { - ArrowStructArray::new_unchecked_with_length( - fields.clone(), - arrow_arrays, - null_buffer, - len, - ) - })) + Ok(Arc::new(ArrowStructArray::try_new_with_length( + fields.clone(), + arrow_arrays, + null_buffer, + len, + )?)) } Err(names) => { // No target fields specified - use preferred types for each child. The inferred field @@ -213,14 +211,12 @@ fn create_from_fields( } let arrow_fields = Fields::from(arrow_fields); - Ok(Arc::new(unsafe { - ArrowStructArray::new_unchecked_with_length( - arrow_fields, - arrow_arrays, - null_buffer, - len, - ) - })) + Ok(Arc::new(ArrowStructArray::try_new_with_length( + arrow_fields, + arrow_arrays, + null_buffer, + len, + )?)) } } } diff --git a/vortex-arrow/src/ext_metadata.rs b/vortex-arrow/src/ext_metadata.rs index 85d7b47610d..37ae679684c 100644 --- a/vortex-arrow/src/ext_metadata.rs +++ b/vortex-arrow/src/ext_metadata.rs @@ -1,30 +1,121 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The Arrow field metadata that holds a Vortex extension type without an Arrow plugin. +//! Encoding Vortex extension types in Arrow field metadata. //! -//! Arrow field metadata must be UTF-8, but Vortex metadata is arbitrary bytes, so the metadata is -//! always base64, and [`EXTENSION_METADATA_ENCODING_KEY`] marks it as base64. The marker is needed -//! because other producers write the same Arrow extension keys with plain text metadata. +//! Typed plugins use [`crate::ArrowSession::register_extension`], whose adapter handles Vortex +//! metadata without exposing its keys or encoding to the callbacks. +//! +//! A custom extension with no Arrow [`ExtensionType`] implementation can use +//! [`set_arrow_extension_name`] to name its storage field. An extension without an Arrow plugin +//! uses [`ArrowSession::to_arrow_storage_field`] instead. +//! +//! Field metadata must be UTF-8, so Vortex payloads use marked base64. The Vortex ID path also +//! accepts unmarked UTF-8 metadata from older producers. A separate Vortex payload requires the +//! marker, and an absent payload imposes no constraint on a target's Vortex metadata. +//! +//! [`ArrowSession::to_arrow_storage_field`]: crate::ArrowSession::to_arrow_storage_field +//! [`ExtensionType`]: arrow_schema::extension::ExtensionType use arrow_schema::Field; use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY; use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; use base64::Engine; use base64::prelude::BASE64_STANDARD; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ExtDTypePlugin; use vortex_array::dtype::extension::ExtDTypeRef; +use vortex_array::dtype::extension::ExtVTable; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; -/// The Arrow field metadata key that names the encoding of `ARROW:extension:metadata`. +/// The Arrow field metadata key that names the encoding of a Vortex metadata payload. /// -/// The only encoding is `base64`. Without this key, the metadata is its own UTF-8 bytes. +/// This applies to `ARROW:extension:metadata` and [`VORTEX_EXTENSION_METADATA_KEY`]. The only +/// encoding is `base64`. Without this key, `ARROW:extension:metadata` is read as raw UTF-8 for +/// backward compatibility. pub const EXTENSION_METADATA_ENCODING_KEY: &str = "vortex:extension:metadata:encoding"; +/// Vortex metadata retained alongside an Arrow extension type. +/// +/// The payload is base64 and **must** carry [`EXTENSION_METADATA_ENCODING_KEY`]. It is separate +/// from the Arrow extension's own metadata, whose format is defined by that extension. +pub const VORTEX_EXTENSION_METADATA_KEY: &str = "vortex:extension:metadata"; + /// The value of [`EXTENSION_METADATA_ENCODING_KEY`] for standard padded base64. const BASE64_ENCODING: &str = "base64"; +/// Sets an Arrow extension name while preserving the field's other metadata. +/// +/// Prefer [`Field::try_with_extension_type`] when an Arrow extension has an +/// [`ExtensionType`](arrow_schema::extension::ExtensionType) implementation, which also validates +/// its storage and native metadata. +pub fn set_arrow_extension_name(field: &mut Field, name: impl Into) { + field + .metadata_mut() + .insert(EXTENSION_TYPE_NAME_KEY.to_owned(), name.into()); +} + +/// Writes Vortex metadata alongside an Arrow field's native metadata. +/// +/// Serialization, keys, and encoding are owned by this module. Callers decide whether the Arrow +/// representation needs a separate Vortex payload. A failed serialization leaves `field` unchanged. +pub(crate) fn write_vortex_metadata( + field: &mut Field, + vtable: &V, + metadata: &V::Metadata, +) -> VortexResult<()> { + let payload = BASE64_STANDARD.encode(vtable.serialize_metadata(metadata)?); + field.metadata_mut().extend([ + (VORTEX_EXTENSION_METADATA_KEY.to_owned(), payload), + ( + EXTENSION_METADATA_ENCODING_KEY.to_owned(), + BASE64_ENCODING.to_owned(), + ), + ]); + + Ok(()) +} + +/// Restores a Vortex dtype using metadata and the importer's inferred Vortex storage. +/// +/// Returns `None` when the field has no Vortex payload. A present payload **must** have a +/// supported encoding marker and metadata compatible with `storage_dtype`. Native Arrow metadata +/// is not read. Reconstruction goes through the dtype plugin so its storage checks also apply. +pub(crate) fn read_vortex_metadata( + field: &Field, + vtable: &V, + storage_dtype: &DType, +) -> VortexResult> { + let Some(payload) = field.metadata().get(VORTEX_EXTENSION_METADATA_KEY) else { + return Ok(None); + }; + + vortex_ensure_eq!( + field + .metadata() + .get(EXTENSION_METADATA_ENCODING_KEY) + .map(String::as_str) + .unwrap_or_default(), + BASE64_ENCODING, + "Arrow field {:?} must mark its Vortex metadata as base64", + field.name() + ); + let metadata = decode_base64_metadata(field, payload)?; + + vtable + .deserialize(&metadata, storage_dtype.clone()) + .map(Some) + .map_err(|err| { + err.with_context(format!( + "Invalid Vortex metadata on Arrow field {:?}", + field.name() + )) + }) +} + /// Returns the Arrow field metadata that holds `ext_dtype`. pub fn encode_arrow_ext_metadata(ext_dtype: &ExtDTypeRef) -> VortexResult<[(String, String); 3]> { let metadata = BASE64_STANDARD.encode(ext_dtype.serialize_metadata()?); @@ -57,15 +148,152 @@ pub fn decode_arrow_ext_metadata(field: &Field) -> VortexResult> { .map(String::as_str) { None => Ok(metadata.as_bytes().to_vec()), - Some(BASE64_ENCODING) => BASE64_STANDARD.decode(metadata).map_err(|err| { - vortex_err!( - "Arrow field {:?} must carry base64 extension metadata, got {metadata:?}: {err}", - field.name() - ) - }), + Some(BASE64_ENCODING) => decode_base64_metadata(field, metadata), Some(encoding) => vortex_bail!( "Arrow field {:?} must have the extension metadata encoding {BASE64_ENCODING:?}, got {encoding:?}", field.name() ), } } + +fn decode_base64_metadata(field: &Field, metadata: &str) -> VortexResult> { + BASE64_STANDARD.decode(metadata).map_err(|err| { + vortex_err!( + "Arrow field {:?} must carry base64 extension metadata, got {metadata:?}: {err}", + field.name() + ) + }) +} + +#[cfg(test)] +mod tests { + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY; + use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; + use rstest::rstest; + use vortex_array::dtype::PType; + use vortex_array::extension::datetime::Date; + use vortex_array::extension::datetime::TimeUnit; + use vortex_error::VortexResult; + + use super::*; + + #[track_caller] + fn check_metadata_round_trip(metadata: TimeUnit, native_type: DataType) -> VortexResult<()> { + let mut field = Field::new("date", native_type, true); + let inferred = Date::try_new(metadata, true.into())?.erased(); + field.metadata_mut().extend([ + ( + EXTENSION_TYPE_NAME_KEY.to_owned(), + "example.date".to_owned(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + "native metadata".to_owned(), + ), + ("application".to_owned(), "retained".to_owned()), // Unrelated application metadata. + ]); + let original = field.clone(); + + write_vortex_metadata(&mut field, &Date, &metadata)?; + + assert_eq!( + read_vortex_metadata(&field, &Date, inferred.storage_dtype())?, + Some(inferred) + ); + assert_eq!(field.data_type(), original.data_type()); + assert_eq!(field.name(), original.name()); + assert_eq!(field.is_nullable(), original.is_nullable()); + for (key, value) in original.metadata() { + assert_eq!(field.metadata().get(key), Some(value)); + } + + Ok(()) + } + + #[rstest] + #[case::days(TimeUnit::Days, DataType::Date32)] + #[case::milliseconds(TimeUnit::Milliseconds, DataType::Date64)] + fn typed_metadata_preserves_native_field( + #[case] metadata: TimeUnit, + #[case] native_type: DataType, + ) -> VortexResult<()> { + check_metadata_round_trip(metadata, native_type) + } + + #[test] + fn absent_vortex_metadata_ignores_native_metadata() -> VortexResult<()> { + let mut field = Field::new("date", DataType::Date32, false); + field.metadata_mut().extend([ + ( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + "native metadata".to_owned(), + ), + ( + EXTENSION_METADATA_ENCODING_KEY.to_owned(), + "native encoding".to_owned(), + ), + ]); + + assert_eq!( + read_vortex_metadata(&field, &Date, &DType::Primitive(PType::I32, false.into()))?, + None + ); + + Ok(()) + } + + #[rstest] + #[case::missing_encoding("AA==", None)] + #[case::unsupported_encoding("AA==", Some("hex"))] + #[case::invalid_base64("!", Some("base64"))] + #[case::invalid_typed_metadata("/w==", Some("base64"))] + #[case::empty_typed_metadata("", Some("base64"))] + fn malformed_metadata_is_rejected(#[case] payload: &str, #[case] encoding: Option<&str>) { + let mut field = Field::new("date", DataType::Date32, false); + field + .metadata_mut() + .insert(VORTEX_EXTENSION_METADATA_KEY.to_owned(), payload.to_owned()); + if let Some(encoding) = encoding { + field.metadata_mut().insert( + EXTENSION_METADATA_ENCODING_KEY.to_owned(), + encoding.to_owned(), + ); + } + + let error = + read_vortex_metadata(&field, &Date, &DType::Primitive(PType::I32, false.into())) + .expect_err("the metadata is malformed"); + assert!(error.to_string().contains("date")); + } + + #[test] + fn metadata_must_match_inferred_storage() -> VortexResult<()> { + let mut field = Field::new("date", DataType::Date32, false); + write_vortex_metadata(&mut field, &Date, &TimeUnit::Milliseconds)?; + + assert!( + read_vortex_metadata(&field, &Date, &DType::Primitive(PType::I32, false.into())) + .is_err() + ); + + Ok(()) + } + + #[test] + fn extension_name_preserves_other_metadata() { + let mut field = Field::new("value", DataType::Float32, false); + field + .metadata_mut() + .insert("application".to_owned(), "retained".to_owned()); + + set_arrow_extension_name(&mut field, "example.vector"); + + assert_eq!(field.extension_type_name(), Some("example.vector")); + assert_eq!( + field.metadata().get("application").map(String::as_str), + Some("retained") + ); + } +} diff --git a/vortex-arrow/src/extension/adapter.rs b/vortex-arrow/src/extension/adapter.rs new file mode 100644 index 00000000000..28a5310cb36 --- /dev/null +++ b/vortex-arrow/src/extension/adapter.rs @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Connects typed extension conversions to the Arrow session's importer and exporter registries. +//! +//! [`ArrowExtVTable`] defines conversion in both directions. This adapter provides the separate +//! interfaces that [`ArrowSession`] registers and applies the vtable's [`ArrowMetadataPolicy`]. +//! Keeping metadata handling here prevents plugins from forgetting it. The session checks that +//! completed fields and imported arrays preserve their dtypes. +//! +//! [`VORTEX_EXTENSION_METADATA_KEY`] names `vortex:extension:metadata`. This separate field entry +//! leaves `ARROW:extension:metadata` for the Arrow extension's own metadata. The codec in +//! [`crate::ext_metadata`] uses `vortex:extension:metadata:encoding` to mark the payload as base64 +//! and checks that marker before decoding. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::Field; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ExtDType; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::registry::Id; + +use super::ArrowExtVTable; +use super::ArrowMetadataPolicy; +use crate::ArrowExport; +use crate::ArrowExportOptions; +use crate::ArrowExportVTable; +use crate::ArrowImport; +use crate::ArrowImportVTable; +use crate::ArrowSession; +use crate::ArrowSessionExt; +use crate::ext_metadata::VORTEX_EXTENSION_METADATA_KEY; +use crate::ext_metadata::read_vortex_metadata; +use crate::ext_metadata::write_vortex_metadata; + +impl ArrowSession { + /// Registers both directions of a typed extension using the vtable's metadata policy. + /// + /// The adapter handles Vortex metadata encoding, decoding, and target checks. The session + /// verifies complete dtype identity for both typed and raw plugins. + pub fn register_extension(&self, vtable: V) { + let adapter = Arc::new(ExtensionPlugin { vtable }); + let exporter = Arc::clone(&adapter); + + self.register_exporter(exporter); + self.register_importer(adapter); + } +} + +/// Registers one typed extension as both an importer and an exporter with the same metadata policy. +#[derive(Debug)] +struct ExtensionPlugin { + vtable: V, +} + +impl ExtensionPlugin { + /// Rejects export targets whose Vortex metadata disagrees with `ext_dtype`. + /// + /// A target's metadata determines the meaning of exported values on import. Under + /// [`ArrowMetadataPolicy::PreserveVortex`], a target with [`VORTEX_EXTENSION_METADATA_KEY`] + /// **must** carry matching metadata. An absent payload adds no Vortex metadata constraint. + /// Comparison uses decoded metadata so it does not depend on serialization details. + fn validate_export_target( + &self, + ext_dtype: &ExtDType, + target: &Field, + ctx: &ExecutionCtx, + ) -> VortexResult<()> { + if V::metadata_policy() != ArrowMetadataPolicy::PreserveVortex + || !target + .metadata() + .contains_key(VORTEX_EXTENSION_METADATA_KEY) + { + return Ok(()); + } + + let native_dtype = + V::from_arrow_field(target, &ctx.session().arrow())?.ok_or_else(|| { + vortex_err!( + "Arrow field {:?} must have a supported native representation, got {}", + target.name(), + target.data_type() + ) + })?; + if let Some(actual) = + read_vortex_metadata(target, &self.vtable, native_dtype.storage_dtype())? + { + vortex_ensure_eq!( + actual.metadata::(), + ext_dtype.metadata(), + "Arrow field {:?} must carry the expected Vortex metadata", + target.name() + ); + } + + Ok(()) + } +} + +impl ArrowExportVTable for ExtensionPlugin { + fn arrow_ext_id(&self) -> Id { + self.vtable.arrow_ext_id() + } + + fn vortex_id(&self) -> Id { + self.vtable.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(None); + }; + let Ok(ext_dtype) = ext_dtype.clone().try_downcast::() else { + return Ok(None); + }; + + let Some(mut field) = self.vtable.to_arrow_field(name, &ext_dtype, session)? else { + return Ok(None); + }; + + // Import routing determines whether the native field needs a separate Vortex payload. + // The session checks the completed field after any payload has been added. + if V::metadata_policy() == ArrowMetadataPolicy::PreserveVortex + && session.from_arrow_field(&field)? != *dtype + { + write_vortex_metadata(&mut field, &self.vtable, ext_dtype.metadata())?; + } + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let Some(ext_dtype) = array.dtype().as_extension_opt() else { + return Ok(ArrowExport::Unsupported(array)); + }; + let Ok(ext_dtype) = ext_dtype.clone().try_downcast::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + + self.validate_export_target(&ext_dtype, target, ctx)?; + + self.vtable.execute_arrow(array, target, options, ctx) + } +} + +impl ArrowImportVTable for ExtensionPlugin { + fn arrow_ext_id(&self) -> Id { + self.vtable.arrow_ext_id() + } + + fn from_arrow_field( + &self, + field: &Field, + session: &ArrowSession, + ) -> VortexResult> { + let Some(native_dtype) = V::from_arrow_field(field, session)? else { + return Ok(None); + }; + + // The plugin can map Arrow's physical storage to a different Vortex storage dtype. Restore + // metadata against that inferred storage so the dtype plugin can validate the combination. + let restored_dtype = if V::metadata_policy() == ArrowMetadataPolicy::PreserveVortex { + read_vortex_metadata(field, &self.vtable, native_dtype.storage_dtype())? + } else { + None + }; + + Ok(Some(DType::Extension( + restored_dtype.unwrap_or_else(|| native_dtype.erased()), + ))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(ArrowImport::Unsupported(array)); + }; + let Ok(ext_dtype) = ext_dtype.clone().try_downcast::() else { + return Ok(ArrowImport::Unsupported(array)); + }; + + V::from_arrow_array(array, field, &ext_dtype, session) + } +} diff --git a/vortex-arrow/src/extension/mod.rs b/vortex-arrow/src/extension/mod.rs new file mode 100644 index 00000000000..928585dbd6f --- /dev/null +++ b/vortex-arrow/src/extension/mod.rs @@ -0,0 +1,142 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interoperability for Vortex extension types. +//! +//! Implement [`ArrowExtVTable`] and register it with [`ArrowSession::register_extension`] to enable +//! Arrow import and export. Registration wraps the callbacks in an adapter that handles Vortex +//! metadata, while the session validates schema round trips. Plugins only implement conversions. +//! +//! A Vortex extension does not need an equivalent standard Arrow type. It can use a custom Arrow +//! extension name, or decline native export so the session exports storage with the Vortex ID and +//! metadata. When a plugin supplies an Arrow representation, its metadata may not describe the full +//! Vortex dtype. Handling that gap in the adapter prevents plugins from silently losing metadata. +//! +//! Declare an [`ArrowMetadataPolicy`] in [`ArrowExtVTable::metadata_policy`] based on whether the +//! Arrow representation preserves the full Vortex dtype on its own. Callbacks leave transport to +//! the adapter and import arrays with the restored dtype they receive. Plugins remain responsible +//! for preserving values and lengths. + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::Field; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtVTable; +use vortex_error::VortexResult; +use vortex_session::registry::Id; + +use crate::ArrowExport; +use crate::ArrowExportOptions; +use crate::ArrowImport; +use crate::ArrowSession; + +/// Whether an Arrow representation needs separate Vortex metadata or not to restore the full dtype. +/// +/// On schema export, the session imports the generated field and checks that it reproduces the +/// original dtype, including metadata, storage, and nullability. On successful array import, the +/// session checks that the returned array has the dtype inferred from the field. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ArrowMetadataPolicy { + /// Rely on the Arrow representation alone to preserve the complete Vortex dtype. + /// + /// Choose this only when the plugin's Arrow field can reconstruct the dtype without additional + /// Vortex metadata. The adapter neither writes nor reads a separate Vortex metadata payload. + ArrowNative, + + /// Preserve Vortex metadata separately when the Arrow representation alone is insufficient. + /// + /// If the exported field does not reconstruct the original dtype, the adapter adds a Vortex + /// metadata payload before the session checks the completed field. The Arrow extension's own + /// metadata stays unchanged. On import, the payload restores metadata over the plugin's inferred + /// storage. + /// + /// An absent payload uses the native interpretation. A required empty payload is still + /// written, because a codec can distinguish it from absence. + PreserveVortex, +} + +/// Native Arrow conversion callbacks for one Vortex extension type. +/// +/// Register through [`ArrowSession::register_extension`], and use the session's conversion APIs. +/// Calling these callbacks directly bypasses metadata transport and schema validation. Callbacks +/// **must** describe only the native Arrow representation, leaving Vortex metadata transport to the +/// adapter. Register the owning dtype in the Vortex session as well to enable its Vortex ID path. +/// +/// Register the dtype and its Arrow conversions in the same session: +/// +/// ``` +/// use vortex_array::dtype::session::DTypeSessionExt; +/// use vortex_arrow::{ArrowExtVTable, ArrowSessionExt}; +/// use vortex_session::VortexSession; +/// +/// fn register(session: &VortexSession, vtable: V) { +/// session.dtypes().register(vtable.clone()); +/// session.arrow().register_extension(vtable); +/// } +/// ``` +pub trait ArrowExtVTable: ExtVTable { + /// The metadata policy used by the adapter for this extension type. + /// + /// The result **must** remain the same across calls and sessions. + fn metadata_policy() -> ArrowMetadataPolicy; + + /// The native Arrow extension ID used for both import and export. + fn arrow_ext_id(&self) -> Id; + + /// Builds the native field, or declines so that the session uses the Vortex ID path. + /// + /// Use [`Field::try_with_extension_type`] for standard Arrow extension metadata, or + /// [`crate::ext_metadata::set_arrow_extension_name`] for a custom name. The adapter adds any + /// Vortex metadata before the session validates the completed field against the source dtype. + fn to_arrow_field( + &self, + name: &str, + ext_dtype: &ExtDType, + session: &ArrowSession, + ) -> VortexResult>; + + /// Converts an array of this extension type to the native target representation. + /// + /// The adapter checks the source extension type before calling this method. Return + /// [`ArrowExport::Unsupported`] when the target is unsupported. Child exports **must** receive + /// the same `options` through [`crate::ExecuteArrowExt::execute_arrow_with_options`]. + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + options: &ArrowExportOptions, + ctx: &mut ExecutionCtx, + ) -> VortexResult; + + /// Interprets the native field, without applying Vortex side metadata. + /// + /// Infer the Vortex storage representation here, which can differ from Arrow's physical type. + /// Return `None` when the field is unsupported. + fn from_arrow_field( + field: &Field, + session: &ArrowSession, + ) -> VortexResult>>; + + /// Imports native values with the extension dtype restored by the adapter. + /// + /// `ext_dtype` includes any Vortex metadata restored by the adapter, so this callback does not + /// need to decode it. Inferring the dtype again from the native Arrow representation could lose + /// that metadata. + /// + /// A successful result **must** have exactly `ext_dtype` and the input length. Return + /// [`ArrowImport::Unsupported`] when the representation is unsupported. Import storage through + /// the session with a nullability rather than the extension field, to avoid dispatching this + /// plugin recursively. + fn from_arrow_array( + array: ArrowArrayRef, + field: &Field, + ext_dtype: &ExtDType, + session: &ArrowSession, + ) -> VortexResult; +} + +mod adapter; + +#[cfg(test)] +mod tests; diff --git a/vortex-arrow/src/extension/tests.rs b/vortex-arrow/src/extension/tests.rs new file mode 100644 index 00000000000..0490c13a685 --- /dev/null +++ b/vortex-arrow/src/extension/tests.rs @@ -0,0 +1,302 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_array::Int64Array; +use arrow_schema::DataType; +use rstest::rstest; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::dtype::session::DTypeSessionExt; +use vortex_array::scalar::ScalarValue; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::*; +use crate::ArrowExportVTable; +use crate::ArrowImportVTable; +use crate::ArrowSessionExt; +use crate::ext_metadata::VORTEX_EXTENSION_METADATA_KEY; +use crate::ext_metadata::set_arrow_extension_name; + +/// Exports a tagged field through its registered policy and checks its complete dtype on import. +#[track_caller] +fn round_trip_field(tag: u8, nullable: bool) -> VortexResult { + let session = tagged_session::(); + let arrow = session.arrow(); + let dtype = DType::Extension(Tagged::::dtype(tag, nullable)?.erased()); + + let field = arrow.to_arrow_field("tag", &dtype)?; + assert_eq!(arrow.from_arrow_field(&field)?, dtype); + + Ok(field) +} + +fn tagged_session() -> VortexSession { + let session = VortexSession::empty(); + session.dtypes().register(Tagged::); + session + .arrow() + .register_extension(Tagged::); + + session +} + +fn raw_tagged_session() -> VortexSession { + let session = VortexSession::empty(); + let plugin = Arc::new(RawTagged); + let exporter = Arc::clone(&plugin); + session.arrow().register_exporter(exporter); + session.arrow().register_importer(plugin); + + session +} + +/// The native field infers tag 255. Tag zero deliberately has an empty serialized payload. +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +struct Tagged; + +impl Tagged { + fn dtype(tag: u8, nullable: bool) -> VortexResult> { + ExtDType::try_with_vtable(Self, tag, DType::Primitive(PType::I64, nullable.into())) + } +} + +impl ExtVTable for Tagged { + type Metadata = u8; + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> Id { + static ID: CachedId = CachedId::new("test.typed_tag"); + *ID + } + + fn serialize_metadata(&self, tag: &u8) -> VortexResult> { + Ok(if *tag == 0 { vec![] } else { vec![*tag] }) + } + + fn deserialize_metadata(&self, bytes: &[u8]) -> VortexResult { + match bytes { + [] => Ok(0), + [tag] => Ok(*tag), + _ => vortex_bail!("tag must contain at most one byte"), + } + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + vortex_ensure_eq!( + *ext_dtype.storage_dtype(), + DType::Primitive(PType::I64, ext_dtype.storage_dtype().nullability()) + ); + + Ok(()) + } + + fn unpack_native<'a>( + _ext_dtype: &ExtDType, + value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(value) + } +} + +impl ArrowExtVTable for Tagged { + fn metadata_policy() -> ArrowMetadataPolicy { + if PRESERVE_VORTEX { + ArrowMetadataPolicy::PreserveVortex + } else { + ArrowMetadataPolicy::ArrowNative + } + } + + fn arrow_ext_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + ext_dtype: &ExtDType, + _session: &ArrowSession, + ) -> VortexResult> { + let mut field = Field::new( + name, + DataType::Int64, + ext_dtype.storage_dtype().is_nullable(), + ); + set_arrow_extension_name(&mut field, self.arrow_ext_id().to_string()); + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + _target: &Field, + _options: &ArrowExportOptions, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(ArrowExport::Unsupported(array)) + } + + fn from_arrow_field( + field: &Field, + _session: &ArrowSession, + ) -> VortexResult>> { + Ok(Some(Self::dtype(255, field.is_nullable())?)) + } + + fn from_arrow_array( + array: ArrowArrayRef, + _field: &Field, + _ext_dtype: &ExtDType, + _session: &ArrowSession, + ) -> VortexResult { + Ok(ArrowImport::Unsupported(array)) + } +} + +/// Bypasses the typed adapter and deliberately drops the extension when importing arrays. +#[derive(Debug)] +struct RawTagged; + +impl ArrowExportVTable for RawTagged { + fn arrow_ext_id(&self) -> Id { + Tagged::.id() + } + + fn vortex_id(&self) -> Id { + Tagged::.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + as ArrowExtVTable>::to_arrow_field( + &Tagged, + name, + &dtype.as_extension().clone().downcast::>(), + session, + ) + } + + fn execute_arrow( + &self, + array: ArrayRef, + _target: &Field, + _options: &ArrowExportOptions, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(ArrowExport::Unsupported(array)) + } +} + +impl ArrowImportVTable for RawTagged { + fn arrow_ext_id(&self) -> Id { + Tagged::.id() + } + + fn from_arrow_field( + &self, + field: &Field, + session: &ArrowSession, + ) -> VortexResult> { + Ok( + as ArrowExtVTable>::from_arrow_field(field, session)? + .map(|dtype| DType::Extension(dtype.erased())), + ) + } + + fn from_arrow_array( + &self, + _array: ArrowArrayRef, + _field: &Field, + _dtype: &DType, + _session: &ArrowSession, + ) -> VortexResult { + Ok(ArrowImport::Imported( + PrimitiveArray::from_iter([1i64]).into_array(), + )) + } +} + +#[rstest] +#[case::native(false, 255, false, None)] +#[case::nullable_native(false, 255, true, None)] +#[case::already_exact(true, 255, false, None)] +#[case::empty_metadata(true, 0, true, Some(""))] +#[case::nonempty_metadata(true, 7, false, Some("Bw=="))] +fn policy_preserves_dtype( + #[case] preserve_vortex: bool, + #[case] tag: u8, + #[case] nullable: bool, + #[case] expected_payload: Option<&str>, +) -> VortexResult<()> { + let field = if preserve_vortex { + round_trip_field::(tag, nullable)? + } else { + round_trip_field::(tag, nullable)? + }; + + assert_eq!( + field + .metadata() + .get(VORTEX_EXTENSION_METADATA_KEY) + .map(String::as_str), + expected_payload + ); + + Ok(()) +} + +#[test] +fn lossy_native_policy_is_rejected() -> VortexResult<()> { + let error = + round_trip_field::(7, false).expect_err("the native field cannot express tag 7"); + + assert!(error.to_string().contains("must preserve its Vortex dtype")); + + Ok(()) +} + +#[test] +fn raw_exporter_must_preserve_dtype() -> VortexResult<()> { + let session = raw_tagged_session(); + let dtype = DType::Extension(Tagged::::dtype(7, false)?.erased()); + let error = session + .arrow() + .to_arrow_field("tag", &dtype) + .expect_err("the raw exporter's field loses tag 7"); + + assert!(error.to_string().contains("must preserve its Vortex dtype")); + + Ok(()) +} + +#[test] +fn raw_importer_must_preserve_dtype() -> VortexResult<()> { + let session = raw_tagged_session(); + let dtype = DType::Extension(Tagged::::dtype(255, false)?.erased()); + let field = session.arrow().to_arrow_field("tag", &dtype)?; + let array: ArrowArrayRef = Arc::new(Int64Array::from(vec![1])); + let error = session + .arrow() + .from_arrow_array(array, &field) + .expect_err("the raw importer discards the inferred extension dtype"); + + assert!( + error + .to_string() + .contains("must preserve the dtype inferred from its field") + ); + + Ok(()) +} diff --git a/vortex-arrow/src/lib.rs b/vortex-arrow/src/lib.rs index 0ff749bcb08..671f2dfbfcb 100644 --- a/vortex-arrow/src/lib.rs +++ b/vortex-arrow/src/lib.rs @@ -20,6 +20,11 @@ pub mod convert; mod datum; pub mod dtype; mod executor; + +mod extension; +pub use extension::ArrowExtVTable; +pub use extension::ArrowMetadataPolicy; + pub mod ext_metadata; mod iter; mod null_buffer; diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index ef6aa4029e9..2a2a2d9dc3d 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -158,8 +158,12 @@ impl<'a> IntoArrowField<'a> for Nullability { /// Plugin layer for exporting a Vortex array to an Arrow extension type. /// +/// Prefer [`crate::ArrowExtVTable`] for automatic Vortex metadata transport. Both interfaces must +/// satisfy the dtype round-trip contract checked by [`ArrowSession`]. +/// /// This is purely an implementation trait, its methods should not be called directly. Instead, -/// use [`ExecuteArrowExt`] to export an array, and the methods on [`ArrowSession`] to convert a dtype. +/// use [`ExecuteArrowExt`] to export an array, and the methods on [`ArrowSession`] to convert a +/// dtype. /// /// [`ExecuteArrowExt`]: crate::ExecuteArrowExt pub trait ArrowExportVTable: 'static + Send + Sync + Debug { @@ -199,6 +203,9 @@ pub trait ArrowExportVTable: 'static + Send + Sync + Debug { /// Plugin layer for importing an Arrow extension-typed array into a Vortex array. /// +/// Prefer [`crate::ArrowExtVTable`] for automatic Vortex metadata transport. Both interfaces must +/// return arrays with the dtype inferred from their fields, as checked by [`ArrowSession`]. +/// /// Plugins are dispatched by `arrow_ext_id`. The session peels Arrow dictionaries before it /// dispatches, so a plugin never sees a [`DataType::Dictionary`] field or array. It sees the value /// field, with the dictionary field's metadata, and the values array. The session then wraps the @@ -258,6 +265,15 @@ type ArrowImporterRegistry = ArcSwapMap>; /// Session-scoped registry of Arrow extension plugins. /// +/// All plugins **must** preserve the complete dtype across [`Self::to_arrow_field`] and +/// [`Self::from_arrow_field`]. This includes extension IDs, metadata, storage, and nullability at +/// every nesting level. Exporting an array to that field and importing it with the same field must +/// preserve its dtype and logical values. The owning dtype and Arrow plugins must be registered in +/// the importing session. +/// +/// The session checks inferred fields and successful plugin imports, including raw plugins. +/// Explicit export targets may request a different dtype. Unknown Vortex IDs import as storage. +/// /// A plugin converts one Vortex extension type to and from an Arrow extension type. The crate that /// owns the type registers its plugin, as `vortex-uuid` does. Temporal types need no plugin, /// because Arrow has native types for them. @@ -485,6 +501,15 @@ impl ArrowSession { for plugin in self.exporters_by_vortex(&ext_dtype.id()).iter() { if let Some(field) = plugin.to_arrow_field(name, &dtype, self)? { + // Check the completed field through normal routing because another importer can + // take precedence over the exporter that inferred it. + vortex_ensure_eq!( + self.from_arrow_field(&field)?, + dtype, + "Arrow field {:?} must preserve its Vortex dtype", + name + ); + return Ok(Some(field)); } } @@ -804,6 +829,12 @@ impl ArrowSession { arr.len(), len ); + vortex_ensure_eq!( + arr.dtype(), + &dtype, + "Arrow importer must preserve the dtype inferred from its field" + ); + return Ok(arr); } ArrowImport::Unsupported(arr) => { diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index f1fa752fb92..badd7743984 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -62,3 +62,8 @@ test = false name = "compress_listview" harness = false test = false + +[[bench]] +name = "string_dict_selection" +harness = false +test = false diff --git a/vortex-btrblocks/benches/string_dict_selection.rs b/vortex-btrblocks/benches/string_dict_selection.rs new file mode 100644 index 00000000000..39bcd9aebb4 --- /dev/null +++ b/vortex-btrblocks/benches/string_dict_selection.rs @@ -0,0 +1,182 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(not(codspeed))] +mod benchmarks { + use std::sync::LazyLock; + + use divan::Bencher; + use divan::counter::ItemsCount; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::session::AggregateFnSession; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::memory::MemorySession; + use vortex_array::optimizer::kernels::KernelSession; + use vortex_array::session::ArraySession; + use vortex_btrblocks::BtrBlocksCompressorBuilder; + use vortex_btrblocks::CompressionSession; + use vortex_error::VortexExpect; + use vortex_session::VortexSession; + + static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::empty() + .with::() + .with::() + .with::() + .with::() + .with::() + }); + + #[derive(Clone, Copy)] + enum Distribution { + Uniform, + Clustered, + VariedPrefix, + } + + #[derive(Clone, Copy)] + struct Case { + name: &'static str, + rows: usize, + distinct_values: usize, + value_length: usize, + distribution: Distribution, + nullable: bool, + } + + impl std::fmt::Debug for Case { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(self.name) + } + } + + impl Case { + fn value_index(self, index: usize) -> usize { + match self.distribution { + Distribution::Uniform | Distribution::VariedPrefix => index % self.distinct_values, + Distribution::Clustered => index * self.distinct_values / self.rows, + } + } + + fn value(self, index: usize) -> String { + let value_index = self.value_index(index); + let mut value = match self.distribution { + Distribution::VariedPrefix => format!( + "{:016x}-value-{value_index:08x}", + value_index.wrapping_mul(0x9e3779b9) + ), + _ if self.value_length <= 12 => format!("{value_index:08x}"), + _ => format!("common-prefix-value-{value_index:08x}"), + }; + value.extend(std::iter::repeat_n( + 'x', + self.value_length.saturating_sub(value.len()), + )); + value + } + + fn make_array(self) -> ArrayRef { + let values = (0..self.rows) + .map(|index| self.value(index)) + .collect::>(); + let nullability = if self.nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }; + + VarBinViewArray::from_iter( + values.iter().enumerate().map(|(index, value)| { + (!self.nullable || index % 10 != 0).then_some(value.as_str()) + }), + DType::Utf8(nullability), + ) + .into_array() + } + } + + const fn case( + name: &'static str, + rows: usize, + distinct_values: usize, + value_length: usize, + distribution: Distribution, + ) -> Case { + Case { + name, + rows, + distinct_values, + value_length, + distribution, + nullable: false, + } + } + + const CASES: [Case; 9] = [ + case("Inline4096", 65_536, 4096, 8, Distribution::Uniform), + case("Outlined16", 65_536, 16, 28, Distribution::Uniform), + case("Outlined4096", 65_536, 4096, 28, Distribution::Uniform), + case("Outlined8192", 65_536, 8192, 28, Distribution::Uniform), + Case { + nullable: true, + ..case( + "NullableOutlined4096", + 65_536, + 4096, + 28, + Distribution::Uniform, + ) + }, + case("Clustered4096", 65_536, 4096, 28, Distribution::Clustered), + case("Long256", 65_536, 4096, 256, Distribution::Uniform), + case( + "UniqueCommonPrefix", + 65_536, + 65_536, + 28, + Distribution::Uniform, + ), + case( + "UniqueVariedPrefix", + 65_536, + 65_536, + 31, + Distribution::VariedPrefix, + ), + ]; + + #[divan::bench(args = CASES)] + fn compress(bencher: Bencher, case: Case) { + let array = case.make_array(); + let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .build(); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_refs(|(array, ctx)| compressor.compress(array, ctx)); + } + + #[divan::bench(args = CASES)] + fn decompress(bencher: Bencher, case: Case) { + let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .build(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor + .compress(&case.make_array(), &mut ctx) + .vortex_expect("benchmark input must compress"); + bencher + .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); + } +} + +fn main() { + divan::main(); +} diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 6d261788780..6d1ce2323f8 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -63,10 +63,12 @@ impl CompressionMode { excluded.push(binary::ZstdBuffersScheme.id()); } Self::Cuda => { - // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other - // string fragmentation and dictionary schemes still require unsupported decode - // paths. Delta has a CUDA decode kernel, but stays excluded until GPU delta - // decode is benchmarked against the schemes it would displace. + // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export, and + // string dictionaries, which the CUDA dictionary gather decodes into views over + // the decoded values. Other string fragmentation schemes and binary dictionaries + // still require unsupported decode paths. Delta has a CUDA decode kernel, but + // stays excluded until GPU delta decode is benchmarked against the schemes it + // would displace. excluded.extend([ integer::DeltaScheme::default().id(), integer::SparseScheme.id(), @@ -76,7 +78,6 @@ impl CompressionMode { float::NullDominatedSparseScheme.id(), string::NullDominatedSparseScheme.id(), fixed_size_list::FixedSizeListSparseScheme.id(), - string::StringDictScheme.id(), binary::BinaryDictScheme.id(), ]); // Both binary Zstd schemes are kept. Buffer-level compression preserves binary @@ -321,4 +322,14 @@ mod tests { assert!(has_delta(builder.clone().with_compact())); assert!(!has_delta(builder.only_cuda_compatible())); } + + #[test] + fn cuda_keeps_string_dictionaries_but_not_binary_ones() { + let cuda = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .only_cuda_compatible(); + let has = |id: SchemeId| cuda.allowed_schemes().iter().any(|s| s.id() == id); + assert!(has(string::StringDictScheme.id())); + assert!(!has(binary::BinaryDictScheme.id())); + } } diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 604018f98ba..976546ddba8 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -9,8 +9,11 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::dtype::PType; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; @@ -71,22 +74,29 @@ impl Scheme for NullDominatedSparseScheme { _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> CompressionEstimate { - let len = data.array_len() as f64; + let len = data.array_len(); let stats = data.varbinview_stats(exec_ctx); - let value_count = stats.value_count(); + let value_count = stats.value_count() as usize; // All-null arrays should be compressed as constant instead anyways. if value_count == 0 { return CompressionEstimate::Verdict(EstimateVerdict::Skip); } - // If the majority (90%) of values is null, this will compress well. - if stats.null_count() as f64 / len > 0.9 { - return CompressionEstimate::Verdict(EstimateVerdict::Ratio(len / value_count as f64)); + // Only null-dominated arrays (90% null) go this route. + if stats.null_count() as f64 / len as f64 <= 0.9 { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); } - // Otherwise we don't go this route. - CompressionEstimate::Verdict(EstimateVerdict::Skip) + // Sparse drops the views of null rows and adds an index per value. It keeps every value's + // bytes, so long values bound the ratio well below the fraction of rows that are null. + let Ok(before) = uncompressed_size_in_bytes(data.array(), exec_ctx) else { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + }; + let null_views = (len - value_count) * size_of::(); + let indices = value_count * PType::min_unsigned_ptype_for_value(len as u64).byte_width(); + let after = before.saturating_sub(null_views) + indices; + CompressionEstimate::Verdict(EstimateVerdict::Ratio(before as f64 / after as f64)) } fn compress( diff --git a/vortex-btrblocks/tests/string_dict_selection.rs b/vortex-btrblocks/tests/string_dict_selection.rs new file mode 100644 index 00000000000..30c748a5b71 --- /dev/null +++ b/vortex-btrblocks/tests/string_dict_selection.rs @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Regression tests for bounded variable-width dictionary selection. + +#![cfg(test)] +#![allow(clippy::tests_outside_test_module)] + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Dict; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::schemes::binary::BinaryDictScheme; +use vortex_btrblocks::schemes::binary::VarBinScheme; +use vortex_btrblocks::schemes::string::FSSTScheme; +use vortex_btrblocks::schemes::string::NullDominatedSparseScheme; +use vortex_btrblocks::schemes::string::StringDictScheme; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +#[derive(Clone, Copy, Debug)] +enum Distribution { + Uniform, + Clustered, +} + +#[rstest] +#[case::outlined_utf8(4096, 28, false, Distribution::Uniform)] +#[case::nullable_utf8(4096, 28, true, Distribution::Uniform)] +#[case::outlined_utf8_8192(8192, 28, false, Distribution::Uniform)] +#[case::clustered_utf8(4096, 28, false, Distribution::Clustered)] +fn test_dict_compressed_with_more_values_than_sample( + #[case] distinct_count: usize, + #[case] value_length: usize, + #[case] nullable: bool, + #[case] distribution: Distribution, +) -> VortexResult<()> { + let distinct_values = (0..distinct_count) + .map(|value| { + let mut string = format!("common-prefix-value-{value:08x}"); + string.extend(std::iter::repeat_n( + 'x', + value_length.saturating_sub(string.len()), + )); + string + }) + .collect::>(); + let values = (0..65_536) + .map(|index| { + let value_index = match distribution { + Distribution::Uniform => index % distinct_values.len(), + Distribution::Clustered => index * distinct_values.len() / 65_536, + }; + (!nullable || index % 10 != 0).then_some(distinct_values[value_index].as_bytes()) + }) + .collect::>(); + let nullability = if nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }; + let array = VarBinViewArray::from_iter(values, DType::Utf8(nullability)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .build(); + let compressed = compressor.compress(&array, &mut ctx)?; + + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +#[test] +fn test_unique_strings_with_common_prefix_not_dict_compressed() -> VortexResult<()> { + let values = (0usize..4096) + .map(|value| Some(format!("common-prefix-value-{value:08x}"))) + .collect::>(); + let array = + VarBinViewArray::from_iter(values, DType::Utf8(Nullability::NonNullable)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .build() + .compress(&array, &mut ctx)?; + + assert!( + !compressed.is::(), + "expected a non-Dict encoding, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +#[test] +fn test_sample_fallback_can_select_dict() -> VortexResult<()> { + let suffix = "x".repeat((1 << 20) + 1); + let distinct_values = [format!("first-{suffix}"), format!("second-{suffix}")]; + let values = (0..4) + .map(|index| distinct_values[index % distinct_values.len()].as_str()) + .collect::>(); + let array = VarBinViewArray::from_iter_str(values).into_array(); + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .build(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut ctx)?; + + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +#[test] +fn test_binary_dict_compressed_with_more_values_than_sample() -> VortexResult<()> { + let distinct_values = (0u32..4096) + .map(|value| { + format!( + "binary-value-{:08x}-{value:08x}", + value.wrapping_mul(0x9e37_79b9) + ) + }) + .collect::>(); + let values = (0..65_536) + .map(|index| Some(distinct_values[index % distinct_values.len()].as_bytes())) + .collect::>(); + let array = + VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&BinaryDictScheme) + .with_new_scheme(&VarBinScheme) + .build() + .compress(&array, &mut ctx)?; + + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +/// Sparse keeps every value's bytes, so it must not win a null-dominated column of long values +/// on the fraction of null rows alone. +#[test] +fn test_null_dominated_long_values_not_left_uncompressed() -> VortexResult<()> { + let values = (0..65_536) + .map(|index| { + (index % 25 == 0).then(|| { + format!( + "https://example.com/media/{index:08x}/{}", + "photo/".repeat(40) + ) + }) + }) + .collect::>(); + let array = VarBinViewArray::from_iter( + values.iter().map(|value| value.as_deref()), + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let value_bytes = values + .iter() + .flatten() + .map(|value| value.len() as u64) + .sum::(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .with_new_scheme(&NullDominatedSparseScheme) + .build() + .compress(&array, &mut ctx)?; + + assert!( + compressed.nbytes() < value_bytes / 2, + "{} compressed to {} bytes for {value_bytes} bytes of values", + compressed.encoding_id(), + compressed.nbytes() + ); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index c407d0251c6..83880cff102 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -10,30 +10,21 @@ use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::Dict; -use vortex_array::arrays::DictArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::dict::DictArrayExt; -use vortex_array::arrays::dict::DictArraySlotsExt; -use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::builders::dict::dict_encode; -use vortex_error::VortexExpect; use vortex_error::VortexResult; +use super::bytes_candidate::bytes_dictionary_estimate; +use super::bytes_candidate::compress_bytes_dictionary; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; use crate::scheme::DescendantExclusion; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; -use crate::stats::GenerateStatsOptions; /// Dictionary encoding for low-cardinality binary values. #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -52,12 +43,6 @@ impl Scheme for BinaryDictScheme { vec![Dict.id()] } - fn stats_options(&self) -> GenerateStatsOptions { - GenerateStatsOptions { - count_distinct_values: true, - } - } - /// Children: values=0, codes=1. fn num_children(&self) -> usize { 2 @@ -77,27 +62,11 @@ impl Scheme for BinaryDictScheme { fn expected_compression_ratio( &self, - data: &ArrayAndStats, + _data: &ArrayAndStats, _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, + _exec_ctx: &mut ExecutionCtx, ) -> CompressionEstimate { - let stats = data.varbinview_stats(exec_ctx); - - if stats.value_count() == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect( - "this must be present since `DictScheme` declared that we need distinct values", - ); - - // If > 50% of the values are distinct, skip dictionary scheme. - if estimated_distinct_values_count > stats.value_count() / 2 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + bytes_dictionary_estimate(&BinaryDictScheme) } fn compress( @@ -107,29 +76,6 @@ impl Scheme for BinaryDictScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let dict = dict_encode(data.array(), exec_ctx)?; - - // Values = child 0. - let compressed_values = - compressor.compress_child(dict.values(), &compress_ctx, self.id(), 0, exec_ctx)?; - - // Codes = child 1. - let narrowed_codes = dict - .codes() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)? - .into_array(); - let compressed_codes = - compressor.compress_child(&narrowed_codes, &compress_ctx, self.id(), 1, exec_ctx)?; - - // SAFETY: compressing codes or values does not alter the invariants. - unsafe { - Ok( - DictArray::new_unchecked(compressed_codes, compressed_values) - .set_all_values_referenced(dict.has_all_values_referenced()) - .into_array(), - ) - } + compress_bytes_dictionary(&BinaryDictScheme, compressor, data, compress_ctx, exec_ctx) } } diff --git a/vortex-compressor/src/builtins/dict/bytes_candidate.rs b/vortex-compressor/src/builtins/dict/bytes_candidate.rs new file mode 100644 index 00000000000..3ad83ec56a4 --- /dev/null +++ b/vortex-compressor/src/builtins/dict/bytes_candidate.rs @@ -0,0 +1,503 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Bounded trial compression for variable-width dictionary selection. +//! +//! Ordinary samples can miss repeated values when a column has thousands of distinct values, so +//! a sample-based estimate under-rates dictionaries for those columns. This module instead +//! builds one bounded dictionary candidate from the complete input. +//! +//! Before paying for that trial, a distributed probe of complete values estimates the number of +//! distinct values in the whole input. A cost model turns that estimate into an optimistic +//! compression ratio, and the trial runs only when that ratio could beat the best estimate so +//! far. A probe that cannot finish within its budget, and a dictionary too large for the trial, +//! fall back to the ordinary sample estimator. +//! +//! The trial's completed output is kept. The compression phase reuses it if the dictionary wins, +//! and the compressor falls back to it if another scheme wins but produces a larger output. + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::dict::DictArrayExt; +use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::builders::dict::DictConstraints; +use vortex_array::builders::dict::codes_ptype_for_max_len; +use vortex_array::builders::dict::dict_encode; +use vortex_array::builders::dict::dict_encode_with_constraints; +use vortex_error::VortexResult; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::CascadingCompressor; +use crate::compressor::TrialOutput; +use crate::compressor::estimate_compression_ratio_with_sampling; +use crate::compressor::exact_uncompressed_size; +use crate::compressor::sample_slices; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; +use crate::scheme::EstimateScore; +use crate::scheme::EstimateVerdict; +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::stats::ArrayAndStats; + +/// Limits the number of entries so that trial codes remain two bytes wide. +const MAX_DICTIONARY_ENTRIES: usize = 1 << u16::BITS; + +/// Holds the width of a trial dictionary's codes. +const TRIAL_CODE_BYTES: usize = codes_ptype_for_max_len(MAX_DICTIONARY_ENTRIES as u64).byte_width(); + +/// Limits the value and view storage in a trial dictionary. +/// +/// This is larger than the 1 MiB that `DictLayoutConstraints` gives a dictionary shared across a +/// column's chunks: one array's dictionary of 50,000 28-byte values still beats FSST, and a +/// smaller limit would hand it back to sampling, which cannot see that. +const MAX_DICTIONARY_BYTES: usize = 1 << 22; + +/// Limits code storage while the trial scans the complete input. +const MAX_DICTIONARY_CODE_BYTES: usize = 1 << 24; + +/// Limits complete input bytes hashed while constructing a trial dictionary. +const MAX_DICTIONARY_INPUT_BYTES: usize = 1 << 28; + +/// Limits complete value bytes hashed by the preliminary probe. +const MAX_PROBE_BYTES: usize = 1 << 20; + +/// Samples this many adjacent rows from each probe range. +const PROBE_RANGE_SIZE: u32 = 64; + +/// Distributes this many probe ranges across the input. +const PROBE_RANGE_COUNT: u32 = 128; + +/// Samples this many adjacent rows from each range in the head probe that screens out unique +/// inputs before the full probe. +const HEAD_PROBE_RANGE_SIZE: u32 = 16; + +/// Requires a head probe this large before it may reject an incomplete input. +/// +/// A distributed sample of 1024 values without a single repeat means the input holds hundreds of +/// thousands of distinct values, or nearly one per row. +const MIN_HEAD_PROBE_ROWS: usize = 1024; + +/// Bounds how well a string scheme compresses a dictionary's values. +/// +/// Dictionary values are distinct, so a scheme can only shrink them one value at a time. FSST +/// writes a 1-byte code for at most 8 bytes and OnPair a 2-byte code for at most 16, so neither +/// can exceed 8× on them, however well it compresses the repetitive input. +const MAX_VALUES_RATIO: f64 = 8.0; + +/// Returns a deferred estimate that can reuse a completed dictionary result. +pub(super) fn bytes_dictionary_estimate(scheme: &'static dyn Scheme) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + move |compressor, data, best_so_far, compress_ctx, exec_ctx| { + estimate_bytes_dictionary( + scheme, + compressor, + data, + best_so_far, + compress_ctx, + exec_ctx, + ) + }, + ))) +} + +/// Dictionary-encodes the input, reusing the trial result when the estimate cached one. +pub(super) fn compress_bytes_dictionary( + scheme: &'static dyn Scheme, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + if let Some(trial) = data.get::() + && trial.scheme == scheme.id() + { + return Ok(trial.array.clone()); + } + + let dict = dict_encode(data.array(), exec_ctx)?; + compress_dictionary(scheme, compressor, &dict, compress_ctx, exec_ctx) +} + +/// Compresses the value and code children of a dictionary. +fn compress_dictionary( + scheme: &'static dyn Scheme, + compressor: &CascadingCompressor, + dict: &DictArray, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + // Values = child 0. + let compressed_values = + compressor.compress_child(dict.values(), &compress_ctx, scheme.id(), 0, exec_ctx)?; + + // Codes = child 1. + let narrowed_codes = dict + .codes() + .clone() + .execute::(exec_ctx)? + .narrow(exec_ctx)? + .into_array(); + let compressed_codes = + compressor.compress_child(&narrowed_codes, &compress_ctx, scheme.id(), 1, exec_ctx)?; + + // SAFETY: compressing codes or values does not alter the invariants. + unsafe { + Ok( + DictArray::new_unchecked(compressed_codes, compressed_values) + .set_all_values_referenced(dict.has_all_values_referenced()) + .into_array(), + ) + } +} + +/// Resolves the deferred dictionary estimate. +fn estimate_bytes_dictionary( + scheme: &'static dyn Scheme, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + best_so_far: Option, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let profile = InputProfile::new(data, exec_ctx)?; + let array = data.array_as_varbinview(); + if profile.valid_rows == 0 { + return Ok(EstimateVerdict::Skip); + } + + if ProbeSample::new(&array, &profile.validity, HEAD_PROBE_RANGE_SIZE).is_some_and(|head| { + head.is_unique() && admit_trial(&profile, &head, best_so_far) == Admission::Reject + }) { + return Ok(EstimateVerdict::Skip); + } + + let Some(sample) = ProbeSample::new(&array, &profile.validity, PROBE_RANGE_SIZE) else { + return sample_verdict( + scheme, + compressor, + data, + best_so_far, + compress_ctx, + exec_ctx, + ); + }; + match admit_trial(&profile, &sample, best_so_far) { + Admission::Trial => {} + Admission::Reject => return Ok(EstimateVerdict::Skip), + Admission::TooLarge => { + return sample_verdict( + scheme, + compressor, + data, + best_so_far, + compress_ctx, + exec_ctx, + ); + } + } + + let Some(dictionary) = build_candidate(data.array(), exec_ctx)? else { + return sample_verdict( + scheme, + compressor, + data, + best_so_far, + compress_ctx, + exec_ctx, + ); + }; + let compressed = compress_dictionary(scheme, compressor, &dictionary, compress_ctx, exec_ctx)?; + let score = EstimateScore::from_sample_sizes(profile.uncompressed_nbytes, compressed.nbytes()); + // Keep the trial even when its ratio loses: the compressor prefers it to a winner whose + // actual output turns out larger. + data.get_or_insert_with(|| TrialOutput { + scheme: scheme.id(), + array: compressed, + }); + if !score_is_best(score, best_so_far) { + return Ok(EstimateVerdict::Skip); + } + Ok(score_verdict(score)) +} + +/// Runs the ordinary sample estimate and applies the current threshold. +fn sample_verdict( + scheme: &'static dyn Scheme, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + best_so_far: Option, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let score = estimate_compression_ratio_with_sampling( + compressor, + scheme, + data.array(), + compress_ctx, + exec_ctx, + )?; + if score_is_best(score, best_so_far) { + Ok(score_verdict(score)) + } else { + Ok(EstimateVerdict::Skip) + } +} + +/// Returns whether a score is valid and beats the current best score. +fn score_is_best(score: EstimateScore, best_so_far: Option) -> bool { + score.is_valid() && best_so_far.is_none_or(|best| score.beats(best)) +} + +/// Converts a score into a terminal estimate. +fn score_verdict(score: EstimateScore) -> EstimateVerdict { + match score { + EstimateScore::FiniteCompression(ratio) => EstimateVerdict::Ratio(ratio), + EstimateScore::ZeroBytes => EstimateVerdict::Skip, + } +} + +/// Describes the complete input with one pass over its views. +struct InputProfile { + /// Counts all rows, including nulls. + rows: usize, + /// Counts non-null rows. + valid_rows: usize, + /// Sums the lengths of non-null values. + value_bytes: usize, + /// Measures the bytes the rows reference, as every other estimate measures its input. + uncompressed_nbytes: u64, + /// Marks non-null rows. + validity: Mask, +} + +impl InputProfile { + /// Profiles the complete input. + fn new(data: &ArrayAndStats, exec_ctx: &mut ExecutionCtx) -> VortexResult { + let array = data.array_as_varbinview(); + let validity = array.validity()?.execute_mask(array.len(), exec_ctx)?; + let views = array.views(); + let value_bytes = match validity.bit_buffer() { + AllOr::All => views.iter().map(|view| view.len() as usize).sum(), + AllOr::None => 0, + AllOr::Some(bits) => bits + .set_indices() + .map(|index| views[index].len() as usize) + .sum(), + }; + + Ok(Self { + rows: array.len(), + valid_rows: validity.true_count(), + value_bytes, + uncompressed_nbytes: exact_uncompressed_size(data.array(), exec_ctx)?, + validity, + }) + } + + /// Returns the mean length of a non-null value. + fn mean_value_bytes(&self) -> f64 { + self.value_bytes as f64 / self.valid_rows as f64 + } +} + +/// Summarizes complete values from distributed ranges of the input. +#[derive(Debug, PartialEq)] +struct ProbeSample { + /// Counts probed non-null rows. + rows: usize, + /// Counts distinct probed non-null values. + distinct: usize, + /// Counts distinct values that occur exactly once in the probe. + singletons: usize, + /// Counts distinct values that occur exactly twice in the probe. + doubletons: usize, + /// Counts adjacent row pairs within a probe range. + adjacent_pairs: usize, + /// Counts adjacent row pairs within a probe range that hold different values. + adjacent_changes: usize, + /// Reports whether the probe covered every row of the input. + complete: bool, +} + +impl ProbeSample { + /// Probes `range_size` adjacent rows from each range, or returns `None` when the probed values + /// exceed the probe budget. + fn new(array: &ArrayView<'_, VarBinView>, validity: &Mask, range_size: u32) -> Option { + let views = array.views(); + let ranges = sample_slices(array.len(), range_size, PROBE_RANGE_COUNT); + let complete = matches!(ranges.as_slice(), [(0, end)] if *end == array.len()); + + let probe_bytes = ranges + .iter() + .flat_map(|&(start, end)| start..end) + .filter(|&index| validity.value(index)) + .map(|index| views[index].len() as usize) + .sum::(); + if probe_bytes > MAX_PROBE_BYTES { + return None; + } + + let mut counts = HashMap::<&[u8], u32>::new(); + let mut adjacent_pairs = 0; + let mut adjacent_changes = 0; + for (start, end) in ranges { + let mut previous = None; + for index in start..end { + let value = validity + .value(index) + .then(|| view_bytes(array, &views[index])); + if let Some(previous) = previous { + adjacent_pairs += 1; + adjacent_changes += usize::from(previous != value); + } + previous = Some(value); + if let Some(value) = value { + *counts.entry(value).or_default() += 1; + } + } + } + + let rows = counts.values().map(|&count| count as usize).sum::(); + if rows == 0 { + return None; + } + Some(Self { + rows, + distinct: counts.len(), + singletons: counts.values().filter(|&&count| count == 1).count(), + doubletons: counts.values().filter(|&&count| count == 2).count(), + adjacent_pairs, + adjacent_changes, + complete, + }) + } + + /// Reports whether every probed value is distinct, in a probe large enough to tell. + fn is_unique(&self) -> bool { + self.distinct == self.rows && (self.complete || self.rows >= MIN_HEAD_PROBE_ROWS) + } + + /// Estimates the number of distinct non-null values in the complete input. + /// + /// Uses the Chao1 estimator, which tends to under-count. An under-count only costs a trial + /// that the complete dictionary then rejects, while an over-count would skip a dictionary + /// that wins. + fn estimated_distinct(&self, valid_rows: usize) -> f64 { + let distinct = self.distinct as f64; + if self.complete { + return distinct; + } + let singletons = self.singletons as f64; + let doubletons = self.doubletons as f64; + let unseen = if self.doubletons > 0 { + singletons * singletons / (2.0 * doubletons) + } else { + singletons * (singletons - 1.0) / 2.0 + }; + (distinct + unseen).min(valid_rows as f64) + } + + /// Estimates the fraction of rows that start a new run of equal values. + fn run_start_fraction(&self) -> f64 { + if self.adjacent_pairs == 0 { + return 1.0; + } + (self.adjacent_changes as f64 / self.adjacent_pairs as f64) + .max(1.0 / f64::from(PROBE_RANGE_SIZE)) + } +} + +/// Decides whether to build a trial dictionary. +#[derive(Debug, PartialEq, Eq)] +enum Admission { + /// Builds the trial dictionary. + Trial, + /// Rejects Dictionary because it cannot beat the best estimate so far. + Reject, + /// Leaves the decision to sampling because the trial would exceed its bounds. + /// + /// A dictionary larger than the trial allows can still win with wider codes. + TooLarge, +} + +/// Applies the trial bounds and the cost model. +fn admit_trial( + profile: &InputProfile, + sample: &ProbeSample, + best_so_far: Option, +) -> Admission { + let best_ratio = match best_so_far { + Some(EstimateScore::FiniteCompression(ratio)) => ratio.max(1.0), + Some(EstimateScore::ZeroBytes) => return Admission::Reject, + None => 1.0, + }; + let distinct = sample.estimated_distinct(profile.valid_rows); + let has_nulls = profile.valid_rows < profile.rows; + let entries = distinct + f64::from(u8::from(has_nulls)); + let values_ratio = best_ratio.min(MAX_VALUES_RATIO); + if optimistic_ratio(profile, sample, entries, values_ratio) <= best_ratio { + return Admission::Reject; + } + + if entries > MAX_DICTIONARY_ENTRIES as f64 + || distinct * profile.mean_value_bytes() > MAX_DICTIONARY_BYTES as f64 + || profile.value_bytes > MAX_DICTIONARY_INPUT_BYTES + || profile.rows.saturating_mul(TRIAL_CODE_BYTES) > MAX_DICTIONARY_CODE_BYTES + { + return Admission::TooLarge; + } + Admission::Trial +} + +/// Estimates the compression ratio of a dictionary with `entries` values. +/// +/// The model assumes dictionary values compress by `values_ratio`, and that codes compress to their +/// bit width over the runs the probe saw. +fn optimistic_ratio( + profile: &InputProfile, + sample: &ProbeSample, + entries: f64, + values_ratio: f64, +) -> f64 { + let values_bytes = entries * (profile.mean_value_bytes() + size_of::() as f64); + let code_bits = entries.max(2.0).log2().ceil(); + let code_bytes = profile.rows as f64 * sample.run_start_fraction() * code_bits / 8.0; + profile.uncompressed_nbytes as f64 / (values_bytes / values_ratio + code_bytes).max(1.0) +} + +/// Builds a dictionary within the trial storage bounds, or returns `None` if it exceeds them. +fn build_candidate( + array: &ArrayRef, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult> { + let constraints = DictConstraints { + max_bytes: MAX_DICTIONARY_BYTES, + max_len: MAX_DICTIONARY_ENTRIES, + }; + let dictionary = dict_encode_with_constraints(array, &constraints, exec_ctx)?; + Ok((dictionary.len() == array.len()).then_some(dictionary)) +} + +/// Returns the complete value referenced by a binary view. +fn view_bytes<'a>(array: &'a ArrayView<'a, VarBinView>, view: &'a BinaryView) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let reference = view.as_view(); + &array.buffer(reference.buffer_index as usize)[reference.as_range()] + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-compressor/src/builtins/dict/bytes_candidate/tests.rs b/vortex-compressor/src/builtins/dict/bytes_candidate/tests.rs new file mode 100644 index 00000000000..07120149ae0 --- /dev/null +++ b/vortex-compressor/src/builtins/dict/bytes_candidate/tests.rs @@ -0,0 +1,305 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::Admission; +use super::HEAD_PROBE_RANGE_SIZE; +use super::InputProfile; +use super::MAX_DICTIONARY_ENTRIES; +use super::MAX_PROBE_BYTES; +use super::PROBE_RANGE_SIZE; +use super::ProbeSample; +use super::admit_trial; +use super::build_candidate; +use super::bytes_dictionary_estimate; +use crate::CascadingCompressor; +use crate::builtins::StringDictScheme; +use crate::compressor::TrialOutput; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; +use crate::scheme::EstimateScore; +use crate::scheme::EstimateVerdict; +use crate::scheme::Scheme; +use crate::stats::ArrayAndStats; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +/// Builds compression data for candidate tests. +fn string_data(values: &[Option]) -> ArrayAndStats { + let nullability = if values.iter().any(Option::is_none) { + Nullability::Nullable + } else { + Nullability::NonNullable + }; + let array = VarBinViewArray::from_iter( + values.iter().map(|value| value.as_deref()), + DType::Utf8(nullability), + ) + .into_array(); + ArrayAndStats::new(array, Default::default()) +} + +/// Returns `rows` values that cycle through `distinct` outlined strings. +fn cycled(rows: usize, distinct: usize) -> Vec> { + (0..rows) + .map(|index| { + let value = index % distinct; + Some(format!( + "{:016x}-value-{value:08x}", + value.wrapping_mul(0x9e37_79b9_7f4a_7c15) + )) + }) + .collect() +} + +/// Profiles and probes the input. +fn profile_and_probe( + data: &ArrayAndStats, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult<(InputProfile, Option)> { + let profile = InputProfile::new(data, exec_ctx)?; + let sample = ProbeSample::new( + &data.array_as_varbinview(), + &profile.validity, + PROBE_RANGE_SIZE, + ); + Ok((profile, sample)) +} + +/// Resolves the Dictionary estimate against an optional threshold. +fn estimate_dictionary( + compressor: &CascadingCompressor, + data: &ArrayAndStats, + best_so_far: Option, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let callback = match bytes_dictionary_estimate(&StringDictScheme) { + CompressionEstimate::Deferred(DeferredEstimate::Callback(callback)) => callback, + estimate => return Err(vortex_err!("dictionary estimation returned {estimate:?}")), + }; + callback( + compressor, + data, + best_so_far, + CompressorContext::new(), + exec_ctx, + ) +} + +#[test] +fn complete_probe_counts_distinct_exactly() -> VortexResult<()> { + let data = string_data(&cycled(4096, 100)); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert!(sample.complete); + assert_eq!(sample.estimated_distinct(profile.valid_rows), 100.0); + Ok(()) +} + +#[rstest] +#[case(20_000)] +#[case(50_000)] +fn partial_probe_estimates_distinct_beyond_probe_size(#[case] distinct: usize) -> VortexResult<()> { + let data = string_data(&cycled(1 << 20, distinct)); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + let estimate = sample.estimated_distinct(profile.valid_rows); + + assert!(!sample.complete); + // The probe sees at most 8192 rows, so its raw distinct count cannot reach `distinct`. + assert!(sample.distinct < distinct); + let error = (estimate - distinct as f64).abs() / distinct as f64; + assert!(error < 0.25, "estimated {estimate} for {distinct}"); + Ok(()) +} + +#[test] +fn probe_is_inconclusive_above_byte_budget() -> VortexResult<()> { + let data = string_data(&[Some("a".to_owned()), Some("x".repeat(MAX_PROBE_BYTES))]); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (_, sample) = profile_and_probe(&data, &mut exec_ctx)?; + + assert_eq!(sample, None); + Ok(()) +} + +#[test] +fn nulls_count_as_one_dictionary_entry() -> VortexResult<()> { + let values = (0..8192) + .map(|index| (index % 10 == 0).then(|| format!("unique-non-null-value-{index:08x}"))) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(profile.valid_rows, 820); + assert_eq!(sample.rows, 820); + assert_eq!(admit_trial(&profile, &sample, None), Admission::Trial); + Ok(()) +} + +#[test] +fn unique_values_are_not_admitted() -> VortexResult<()> { + let data = string_data(&cycled(8192, 8192)); + let compressor = CascadingCompressor::new(vec![&StringDictScheme]); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(admit_trial(&profile, &sample, None), Admission::Reject); + assert!(matches!( + estimate_dictionary(&compressor, &data, None, &mut exec_ctx)?, + EstimateVerdict::Skip + )); + Ok(()) +} + +#[test] +fn threshold_rejects_trial_before_building() -> VortexResult<()> { + let data = string_data(&cycled(65_536, 4096)); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(admit_trial(&profile, &sample, None), Admission::Trial); + assert_eq!( + admit_trial( + &profile, + &sample, + Some(EstimateScore::FiniteCompression(1000.0)) + ), + Admission::Reject + ); + Ok(()) +} + +#[test] +fn too_many_estimated_entries_defer_to_sampling() -> VortexResult<()> { + // Each value repeats 16 times, so a dictionary wins even though it exceeds the trial bounds. + let distinct = 2 * MAX_DICTIONARY_ENTRIES; + let values = (0..16 * distinct) + .map(|index| Some(format!("{:016x}", (index * 7919) % distinct))) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(admit_trial(&profile, &sample, None), Admission::TooLarge); + Ok(()) +} + +#[test] +fn candidate_stops_when_dictionary_storage_exceeds_limit() -> VortexResult<()> { + let suffix = "x".repeat(70_000); + let values = (0..128) + .map(|index| Some(format!("{:08x}-{suffix}", index % 64))) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = SESSION.create_execution_ctx(); + + assert!(build_candidate(data.array(), &mut exec_ctx)?.is_none()); + Ok(()) +} + +#[test] +fn candidate_uses_full_u16_code_range() -> VortexResult<()> { + let mut values = (0..=u16::MAX) + .map(|value| Some(format!("{value:08x}"))) + .collect::>(); + let mut exec_ctx = SESSION.create_execution_ctx(); + assert!(build_candidate(string_data(&values).array(), &mut exec_ctx)?.is_some()); + + values.push(Some("past-u16-code-range".to_owned())); + assert!(build_candidate(string_data(&values).array(), &mut exec_ctx)?.is_none()); + Ok(()) +} + +#[test] +fn completed_candidate_is_reused() -> VortexResult<()> { + let data = string_data(&cycled(65_536, 4096)); + let compressor = CascadingCompressor::new(vec![&StringDictScheme]); + let mut exec_ctx = SESSION.create_execution_ctx(); + let verdict = estimate_dictionary(&compressor, &data, None, &mut exec_ctx)?; + assert!(matches!(verdict, EstimateVerdict::Ratio(_))); + let cached = data + .get::() + .ok_or_else(|| vortex_err!("completed candidate was not cached"))?; + let compressed = + StringDictScheme.compress(&compressor, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(ArrayRef::ptr_eq(&cached.array, &compressed)); + Ok(()) +} + +/// Repeats far apart still reach the head probe: each value below recurs half the input away. +#[rstest] +#[case::unique(65_536, 65_536, true)] +#[case::two_repeats_per_value(65_536, 32_768, false)] +#[case::four_repeats_per_value(65_536, 16_384, false)] +#[case::two_repeats_in_a_shorter_input(16_384, 8_192, false)] +fn head_probe_screens_unique_inputs( + #[case] rows: usize, + #[case] distinct: usize, + #[case] expected: bool, +) -> VortexResult<()> { + let data = string_data(&cycled(rows, distinct)); + let mut exec_ctx = SESSION.create_execution_ctx(); + let profile = InputProfile::new(&data, &mut exec_ctx)?; + let head = ProbeSample::new( + &data.array_as_varbinview(), + &profile.validity, + HEAD_PROBE_RANGE_SIZE, + ) + .ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert!(!head.complete); + assert_eq!(head.is_unique(), expected); + if expected { + assert_eq!(admit_trial(&profile, &head, None), Admission::Reject); + } + Ok(()) +} + +/// A mostly-null column of distinct long values compresses far better as sparse values than as a +/// dictionary, whose distinct values no string scheme compresses more than eight-fold. +#[test] +fn values_ratio_bound_rejects_sparse_unique_values() -> VortexResult<()> { + let suffix = "x".repeat(240); + let values = (0..40_960) + .map(|index| (index % 23 == 0).then(|| format!("{index:016x}{suffix}"))) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = SESSION.create_execution_ctx(); + let (profile, sample) = profile_and_probe(&data, &mut exec_ctx)?; + let sample = sample.ok_or_else(|| vortex_err!("probe did not finish"))?; + + assert_eq!(admit_trial(&profile, &sample, None), Admission::Trial); + assert_eq!( + admit_trial( + &profile, + &sample, + Some(EstimateScore::FiniteCompression(22.0)) + ), + Admission::Reject + ); + Ok(()) +} diff --git a/vortex-compressor/src/builtins/dict/mod.rs b/vortex-compressor/src/builtins/dict/mod.rs index 4862df2b211..a49b9ea1a22 100644 --- a/vortex-compressor/src/builtins/dict/mod.rs +++ b/vortex-compressor/src/builtins/dict/mod.rs @@ -4,6 +4,7 @@ //! Dictionary encoding schemes for binary, integer, float, and string arrays. mod binary; +mod bytes_candidate; mod float; mod integer; mod string; diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index f5cbcd54d89..d339b60a9d6 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -10,30 +10,21 @@ use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::Dict; -use vortex_array::arrays::DictArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::dict::DictArrayExt; -use vortex_array::arrays::dict::DictArraySlotsExt; -use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::builders::dict::dict_encode; -use vortex_error::VortexExpect; use vortex_error::VortexResult; +use super::bytes_candidate::bytes_dictionary_estimate; +use super::bytes_candidate::compress_bytes_dictionary; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; use crate::scheme::DescendantExclusion; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; -use crate::stats::GenerateStatsOptions; /// Dictionary encoding for low-cardinality string values. #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -52,12 +43,6 @@ impl Scheme for StringDictScheme { vec![Dict.id()] } - fn stats_options(&self) -> GenerateStatsOptions { - GenerateStatsOptions { - count_distinct_values: true, - } - } - /// Children: values=0, codes=1. fn num_children(&self) -> usize { 2 @@ -77,27 +62,11 @@ impl Scheme for StringDictScheme { fn expected_compression_ratio( &self, - data: &ArrayAndStats, + _data: &ArrayAndStats, _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, + _exec_ctx: &mut ExecutionCtx, ) -> CompressionEstimate { - let stats = data.varbinview_stats(exec_ctx); - - if stats.value_count() == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect( - "this must be present since `DictScheme` declared that we need distinct values", - ); - - // If > 50% of the values are distinct, skip dictionary scheme. - if estimated_distinct_values_count > stats.value_count() / 2 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + bytes_dictionary_estimate(&StringDictScheme) } fn compress( @@ -107,29 +76,6 @@ impl Scheme for StringDictScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let dict = dict_encode(data.array(), exec_ctx)?; - - // Values = child 0. - let compressed_values = - compressor.compress_child(dict.values(), &compress_ctx, self.id(), 0, exec_ctx)?; - - // Codes = child 1. - let narrowed_codes = dict - .codes() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)? - .into_array(); - let compressed_codes = - compressor.compress_child(&narrowed_codes, &compress_ctx, self.id(), 1, exec_ctx)?; - - // SAFETY: compressing codes or values does not alter the invariants. - unsafe { - Ok( - DictArray::new_unchecked(compressed_codes, compressed_values) - .set_all_values_referenced(dict.has_all_values_referenced()) - .into_array(), - ) - } + compress_bytes_dictionary(&StringDictScheme, compressor, data, compress_ctx, exec_ctx) } } diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index de4599a9cc7..8686ad411b4 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -43,6 +43,19 @@ use crate::stats::ArrayAndStats; use crate::stats::GenerateStatsOptions; use crate::trace; +/// A complete compressed output that a scheme produced while estimating its ratio. +/// +/// An estimate callback may store one in its [`ArrayAndStats`]. Its size is a measurement, while +/// the winning scheme was picked on an estimate, so the compressor keeps the trial output +/// whenever it is smaller than what the winner actually produced. +#[derive(Debug)] +pub(crate) struct TrialOutput { + /// The scheme that produced the output. + pub(crate) scheme: SchemeId, + /// The compressed array. + pub(crate) array: ArrayRef, +} + /// The outcome of [`CascadingCompressor::choose_and_compress`]. enum Selection { /// A scheme, or the compressor's own empty, null or constant handling, produced a new array. @@ -357,6 +370,12 @@ impl CascadingCompressor { trace::scheme_compress_failed(winner.id(), before_nbytes, error_ctx.as_ref(), err); })?; + // The winner was picked on estimates. A trial output is measured, so it wins when smaller. + let compressed = match data.get::() { + Some(trial) if trial.array.nbytes() < compressed.nbytes() => trial.array.clone(), + _ => compressed, + }; + let after_nbytes = compressed.nbytes(); let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); @@ -380,7 +399,7 @@ impl CascadingCompressor { /// The bytes `array`'s rows use once canonical, read from the exact statistic the writer caches /// on each chunk and computed otherwise. A dtype that cannot be measured, such as variant, is /// sized by the bytes it holds. -pub(super) fn exact_uncompressed_size( +pub(crate) fn exact_uncompressed_size( array: &ArrayRef, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..8b2117e7b83 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,11 @@ mod sample; mod select; mod structural; +pub(crate) use cascade::TrialOutput; +pub(crate) use cascade::exact_uncompressed_size; +pub(crate) use sample::estimate_compression_ratio_with_sampling; +pub(crate) use sample::sample_slices; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index dcdce0dec1f..714af01e49a 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -43,12 +43,7 @@ pub(crate) fn sample(input: &ArrayRef, sample_size: u32, sample_count: u32) -> A return input.clone(); } - let slices = stratified_slices( - input.len(), - sample_size, - sample_count, - &mut StdRng::seed_from_u64(SAMPLE_SEED), - ); + let slices = sample_slices(input.len(), sample_size, sample_count); // For every slice, grab the relevant slice and repack into a new PrimitiveArray. let chunks = slices.into_iter().map(|(start, end)| { @@ -60,6 +55,20 @@ pub(crate) fn sample(input: &ArrayRef, sample_size: u32, sample_count: u32) -> A unsafe { ChunkedArray::new_unchecked(chunks, input.dtype().clone()) }.into_array() } +/// Returns deterministic stratified sample ranges for an array length. +pub(crate) fn sample_slices( + length: usize, + sample_size: u32, + sample_count: u32, +) -> Vec<(usize, usize)> { + stratified_slices( + length, + sample_size, + sample_count, + &mut StdRng::seed_from_u64(SAMPLE_SEED), + ) +} + /// Computes the number of sample chunks to cover approximately 1% of `len` elements, /// with a minimum of `SAMPLE_SIZE * SAMPLE_COUNT` (1024) values. pub(crate) fn sample_count_approx_one_percent(len: usize) -> u32 { @@ -147,7 +156,7 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> /// # Errors /// /// Returns an error if sample compression fails. -pub(super) fn estimate_compression_ratio_with_sampling( +pub(crate) fn estimate_compression_ratio_with_sampling( compressor: &CascadingCompressor, scheme: &S, array: &ArrayRef, @@ -163,6 +172,12 @@ pub(super) fn estimate_compression_ratio_with_sampling( canonical.into_array() }; + // A sample of only nulls compresses to almost nothing whatever the scheme, so its ratio says + // nothing about how the scheme handles the array's values. + if !compress_ctx.is_sample() && sample_array.all_invalid(exec_ctx)? { + return Ok(EstimateScore::ZeroBytes); + } + let sample_data = ArrayAndStats::new(sample_array, scheme.stats_options()); let error_ctx = trace::enabled_error_context(&compress_ctx); let sample_ctx = compress_ctx.with_sampling(); diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 3c73d2d4cdb..4afdcf48eea 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -55,11 +55,13 @@ impl CascadingCompressor { /// [`CompressionEstimate::Deferred`] are stashed for pass 2 so that we do not make any /// expensive computations if we don't have to. /// - /// Pass 2 evaluates the deferred work and, for each [`DeferredEstimate::Callback`], passes the - /// current best [`EstimateScore`] as an early-exit hint so the callback can return - /// [`EstimateVerdict::Skip`] without doing expensive work when it cannot beat the threshold. + /// Pass 2 evaluates the deferred work. Every [`DeferredEstimate::Sample`] runs before any + /// [`DeferredEstimate::Callback`], and each callback receives the current best + /// [`EstimateScore`] as an early-exit hint so it can return [`EstimateVerdict::Skip`] without + /// doing expensive work when it cannot beat the threshold. /// - /// Ties are broken by registration order within each pass. + /// Ties are broken by evaluation order: immediate verdicts, then samples, then callbacks, each + /// in registration order. /// /// [`expected_compression_ratio`]: Scheme::expected_compression_ratio pub(super) fn choose_best_scheme( @@ -96,7 +98,8 @@ impl CascadingCompressor { } // Pass 2: run deferred work. Callbacks receive the current best as a threshold so they can - // short-circuit with `Skip` when they cannot beat it. + // short-circuit with `Skip` when they cannot beat it, so samples run first to tighten it. + deferred.sort_by_key(|(_, estimate)| matches!(estimate, DeferredEstimate::Callback(_))); for (scheme, deferred_estimate) in deferred { let _span = trace::scheme_eval_span(scheme.id()).entered(); let threshold: Option = best.map(|(_, score)| score); diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 29e32b433c7..190b36ee520 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -15,6 +15,7 @@ use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_si use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Constant; +use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; @@ -28,6 +29,7 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbinview::BinaryView; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; +use vortex_array::builders::dict::dict_encode; use vortex_array::dtype::DType; use vortex_array::dtype::MapDType; use vortex_array::dtype::Nullability; @@ -47,6 +49,7 @@ use vortex_session::VortexSession; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; +use super::TrialOutput; use super::sample::estimate_compression_ratio_with_sampling; use super::select::WinnerEstimate; use super::structural; @@ -336,6 +339,126 @@ impl Scheme for ZeroBytesSamplingScheme { } } +/// Samples to a finite ratio by writing a single byte. +#[derive(Debug)] +struct FiniteSamplingScheme; + +impl Scheme for FiniteSamplingScheme { + fn scheme_name(&self) -> &'static str { + "test.finite_sampling" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Sample) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(PrimitiveArray::new(buffer![0u8], Validity::NonNullable).into_array()) + } +} + +/// Promises a ratio it does not deliver: its output is the input. +#[derive(Debug)] +struct OverpromisingScheme; + +impl Scheme for OverpromisingScheme { + fn scheme_name(&self) -> &'static str { + "test.overpromising" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Ratio(100.0)) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } +} + +/// Stores a dictionary as its trial output, then reports that it cannot beat the threshold. +#[derive(Debug)] +struct TrialStoringScheme; + +impl Scheme for TrialStoringScheme { + fn scheme_name(&self) -> &'static str { + "test.trial_storing" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, data, _best_so_far, _ctx, exec_ctx| { + let array = dict_encode(data.array(), exec_ctx)?.into_array(); + data.get_or_insert_with(|| TrialOutput { + scheme: TrialStoringScheme.id(), + array, + }); + Ok(EstimateVerdict::Skip) + }, + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + #[test] fn test_self_exclusion() { let c = compressor(); @@ -675,6 +798,61 @@ fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { Ok(()) } +#[test] +fn callbacks_observe_samples_registered_after_them() -> VortexResult<()> { + let _guard = OBSERVER_LOCK.lock(); + *OBSERVED_THRESHOLD.lock() = None; + + let compressor = + CascadingCompressor::new(vec![&ThresholdObservingScheme, &FiniteSamplingScheme]); + let schemes: [&'static dyn Scheme; 2] = [&ThresholdObservingScheme, &FiniteSamplingScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + let observed = *OBSERVED_THRESHOLD.lock(); + assert!(matches!( + observed, + Some(Some(EstimateScore::FiniteCompression(_))) + )); + Ok(()) +} + +#[test] +fn smaller_trial_output_replaces_an_overpromising_winner() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&OverpromisingScheme, &TrialStoringScheme]); + let array = PrimitiveArray::from_iter((0..1024i64).map(|value| value % 4)).into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressed = compressor.compress(&array, &mut exec_ctx)?; + + assert!(compressed.is::()); + assert_arrays_eq!(compressed, array, &mut exec_ctx); + Ok(()) +} + +#[test] +fn all_null_sample_gives_no_estimate() -> VortexResult<()> { + // Only the first row is valid, and the sampler does not land on it. + let array = + PrimitiveArray::from_option_iter((0..100_000).map(|index| (index == 0).then_some(index))) + .into_array(); + let compressor = CascadingCompressor::new(vec![&FiniteSamplingScheme]); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let score = estimate_compression_ratio_with_sampling( + &compressor, + &FiniteSamplingScheme, + &array, + CompressorContext::new(), + &mut exec_ctx, + )?; + + assert_eq!(score, EstimateScore::ZeroBytes); + Ok(()) +} + #[test] fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<()> { // Both schemes produce the same `Ratio(2.0)`, one from pass 1 (immediate) and one from diff --git a/vortex-compressor/src/scheme/estimate.rs b/vortex-compressor/src/scheme/estimate.rs index 20b9dc77e9b..13fe744816d 100644 --- a/vortex-compressor/src/scheme/estimate.rs +++ b/vortex-compressor/src/scheme/estimate.rs @@ -101,7 +101,7 @@ pub enum DeferredEstimate { pub enum EstimateScore { /// A finite compression ratio. Higher means a smaller amount of data, so it is better. FiniteCompression(f64), - /// Trial compression produced a 0-byte output. + /// Trial compression produced a 0-byte output, or sampled only null values. /// /// This has no finite ratio and is not eligible for scheme selection. /// diff --git a/vortex-compressor/src/stats/cache.rs b/vortex-compressor/src/stats/cache.rs index 2652403d3e6..07d0bf1cca1 100644 --- a/vortex-compressor/src/stats/cache.rs +++ b/vortex-compressor/src/stats/cache.rs @@ -67,6 +67,19 @@ impl StatsCache { new_arc } } + + /// Returns a cached value when one exists. + fn get(&self) -> Option> { + let type_id = TypeId::of::(); + let guard = self.entries.lock(); + let position = guard.iter().position(|(id, _)| *id == type_id)?; + Some( + Arc::clone(&guard[position].1) + .downcast::() + .ok() + .vortex_expect("we just checked the TypeID"), + ) + } } /// An array bundled with its lazily-computed statistics cache. @@ -207,4 +220,9 @@ impl ArrayAndStats { pub fn get_or_insert_with(&self, f: impl FnOnce() -> T) -> Arc { self.cache.get_or_insert_with::(f) } + + /// Returns a custom cached value when one exists. + pub(crate) fn get(&self) -> Option> { + self.cache.get::() + } } diff --git a/vortex-cuda/benches/dynamic_dispatch_cuda.rs b/vortex-cuda/benches/dynamic_dispatch_cuda.rs index bee5f2e3b77..9b4df35002c 100644 --- a/vortex-cuda/benches/dynamic_dispatch_cuda.rs +++ b/vortex-cuda/benches/dynamic_dispatch_cuda.rs @@ -625,9 +625,11 @@ mod standalone { { let mut launch = cuda_ctx.stream().launch_builder(&self.dict_fn); let codes_len = self.codes_len as u64; + let values_len = self.values_len as u64; launch.arg(&self.codes_buf); launch.arg(&codes_len); launch.arg(&self.values_alp_buf); + launch.arg(&values_len); launch.arg(&mut self.output_buf); unsafe { launch diff --git a/vortex-cuda/kernels/src/decimal_cast.cu b/vortex-cuda/kernels/src/decimal_cast.cu index 45c8c75e2a0..88618176a94 100644 --- a/vortex-cuda/kernels/src/decimal_cast.cu +++ b/vortex-cuda/kernels/src/decimal_cast.cu @@ -6,64 +6,43 @@ #include #include -// Arrow decimal schemas fix the physical values buffer width: -// - Decimal32: 4 bytes per value. -// - Decimal64: 8 bytes per value. -// - Decimal128: 16 bytes per value. -// - Decimal256: 32 bytes per value. +// Widen decimal values buffers to a wider physical width. Two callers use these kernels: +// - The Arrow export casts to the width that the Arrow schema fixes: 4 bytes for Decimal32, +// 8 bytes for Decimal64, 16 bytes for Decimal128, and 32 bytes for Decimal256. +// - The Widen executor casts a narrow decimal child to the canonical width of its decimal +// dtype, from 2 to 32 bytes. // -// Vortex storage width can differ, so export casts to the schema-implied width. -// Rust-side export rejects narrowing casts because detecting overflow on-device -// would require synchronizing an overflow flag back to the host. +// Only widening pairs are instantiated. A narrowing cast can overflow, and detecting that on the +// device would require synchronizing an overflow flag back to the host, so the Rust side rejects +// narrowing before it loads a kernel. -// Low 64-bit conversion for Decimal32/64 outputs. -template -__device__ __forceinline__ int64_t decimal_to_i64(Input value) { - if constexpr (std::is_same_v) { - return value.lo; - } else if constexpr (std::is_same_v) { - return value.parts[0]; - } else { - return static_cast(value); - } -} - -// 128-bit conversion for Decimal128 outputs. -template -__device__ __forceinline__ int128_t decimal_to_i128(Input value) { - if constexpr (std::is_same_v) { - return value; - } else if constexpr (std::is_same_v) { - return int128_t {value.parts[0], value.parts[1]}; - } else { - const int64_t lo = static_cast(value); - const int64_t hi = value < 0 ? -1 : 0; - return int128_t {lo, hi}; - } -} - -// Convert one value to the Arrow schema's physical width. +// Sign-extend one value to the wider output width in a single step, with no intermediate width. +// Native outputs use one `static_cast`. The 128-bit and 256-bit outputs keep the input's words +// and fill the upper words with its sign. template __device__ __forceinline__ Output decimal_cast_value(Input value) { - if constexpr (std::is_same_v) { - return static_cast(decimal_to_i64(value)); - } else if constexpr (std::is_same_v) { - return decimal_to_i64(value); - } else if constexpr (std::is_same_v) { - return decimal_to_i128(value); + static_assert(sizeof(Input) < sizeof(Output), "decimal casts on the GPU only widen"); + + if constexpr (std::is_same_v) { + static_assert(std::is_integral_v && std::is_signed_v); + const int64_t lo = static_cast(value); + return int128_t {lo, lo < 0 ? -1 : 0}; + } else if constexpr (std::is_same_v && std::is_same_v) { + const int64_t sign = value.hi < 0 ? -1 : 0; + return int256_t {{value.lo, value.hi, sign, sign}}; + } else if constexpr (std::is_same_v) { + static_assert(std::is_integral_v && std::is_signed_v); + const int64_t lo = static_cast(value); + const int64_t sign = lo < 0 ? -1 : 0; + return int256_t {{lo, sign, sign, sign}}; } else { - static_assert(std::is_same_v); - if constexpr (std::is_same_v) { - return value; - } else { - const int128_t value128 = decimal_to_i128(value); - const int64_t sign = value128.hi < 0 ? -1 : 0; - return int256_t {{value128.lo, value128.hi, sign, sign}}; - } + static_assert(std::is_integral_v && std::is_signed_v); + static_assert(std::is_integral_v && std::is_signed_v); + return static_cast(value); } } -// Cast a contiguous values buffer to the Arrow schema's physical width. +// Cast a contiguous values buffer to the output physical width. template __device__ void decimal_cast_device(const Input *__restrict input, Output *__restrict output, uint64_t array_len) { @@ -76,28 +55,26 @@ decimal_cast_device(const Input *__restrict input, Output *__restrict output, ui } } -// Generate Decimal32/64/128/256 cast kernels for one input storage type. -#define GENERATE_DECIMAL_CAST_KERNELS(input_suffix, InputType) \ - extern "C" __global__ void decimal_cast_##input_suffix##_i32(const InputType *__restrict input, \ - int32_t *__restrict output, \ - uint64_t array_len) { \ - decimal_cast_device(input, output, array_len); \ - } \ - extern "C" __global__ void decimal_cast_##input_suffix##_i64(const InputType *__restrict input, \ - int64_t *__restrict output, \ - uint64_t array_len) { \ - decimal_cast_device(input, output, array_len); \ - } \ - extern "C" __global__ void decimal_cast_##input_suffix##_i128(const InputType *__restrict input, \ - int128_t *__restrict output, \ - uint64_t array_len) { \ - decimal_cast_device(input, output, array_len); \ - } \ - extern "C" __global__ void decimal_cast_##input_suffix##_i256(const InputType *__restrict input, \ - int256_t *__restrict output, \ - uint64_t array_len) { \ +#define GENERATE_DECIMAL_CAST(input_suffix, InputType, output_suffix, OutputType) \ + extern "C" __global__ void decimal_cast_##input_suffix##_##output_suffix( \ + const InputType *__restrict input, \ + OutputType *__restrict output, \ + uint64_t array_len) { \ decimal_cast_device(input, output, array_len); \ } -FOR_EACH_SIGNED_INT(GENERATE_DECIMAL_CAST_KERNELS) -FOR_EACH_LARGE_DECIMAL(GENERATE_DECIMAL_CAST_KERNELS) +GENERATE_DECIMAL_CAST(i8, int8_t, i16, int16_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i32, int32_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i64, int64_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i32, int32_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i64, int64_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i32, int32_t, i64, int64_t) +GENERATE_DECIMAL_CAST(i32, int32_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i32, int32_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i64, int64_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i64, int64_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i128, int128_t, i256, int256_t) diff --git a/vortex-cuda/kernels/src/dict.cu b/vortex-cuda/kernels/src/dict.cu index a456b476821..db8b1caf249 100644 --- a/vortex-cuda/kernels/src/dict.cu +++ b/vortex-cuda/kernels/src/dict.cu @@ -8,10 +8,14 @@ #include "config.cuh" #include "types.cuh" +// Codes under null positions may hold any value, and so may corrupt input, so every gather checks +// the code against the dictionary length. A code outside the dictionary writes a zero value instead +// of reading past the end of `values`. template __device__ void dict_kernel(const IndexT *const __restrict codes, uint64_t codes_len, const ValueT *const __restrict values, + uint64_t values_len, ValueT *const __restrict output) { const uint32_t elements_per_block = blockDim.x * ELEMENTS_PER_THREAD; const uint64_t block_start = static_cast(blockIdx.x) * elements_per_block; @@ -20,7 +24,7 @@ __device__ void dict_kernel(const IndexT *const __restrict codes, for (uint64_t idx = block_start + threadIdx.x; idx < block_end; idx += blockDim.x) { const IndexT code = codes[idx]; - output[idx] = values[code]; + output[idx] = static_cast(code) < values_len ? values[code] : ValueT {}; } } @@ -39,6 +43,7 @@ template __device__ void dict_bool_kernel(const IndexT *const __restrict codes, uint64_t codes_len, const uint8_t *const __restrict values, + uint64_t values_len, uint64_t values_bit_offset, uint8_t *const __restrict output) { const uint64_t output_len = (codes_len + 7) / 8; @@ -53,7 +58,7 @@ __device__ void dict_bool_kernel(const IndexT *const __restrict codes, #pragma unroll for (uint32_t bit = 0; bit < 8; ++bit) { const uint64_t row = row_start + bit; - if (row < codes_len) { + if (row < codes_len && static_cast(codes[row]) < values_len) { const uint64_t value_idx = values_bit_offset + static_cast(codes[row]); const uint8_t value = (values[value_idx / 8] >> (value_idx % 8)) & 1; packed |= static_cast(value << bit); @@ -69,8 +74,9 @@ __device__ void dict_bool_kernel(const IndexT *const __restrict codes, const IndexType *const __restrict codes, \ uint64_t codes_len, \ const ValueType *const __restrict values, \ + uint64_t values_len, \ ValueType *const __restrict output) { \ - dict_kernel(codes, codes_len, values, output); \ + dict_kernel(codes, codes_len, values, values_len, output); \ } // Generate dict kernel for all index types (unsigned integers) for a given value type @@ -84,9 +90,10 @@ __device__ void dict_bool_kernel(const IndexT *const __restrict codes, extern "C" __global__ void dict_bool_##index_suffix(const IndexType *const __restrict codes, \ uint64_t codes_len, \ const uint8_t *const __restrict values, \ + uint64_t values_len, \ uint64_t values_bit_offset, \ uint8_t *const __restrict output) { \ - dict_bool_kernel(codes, codes_len, values, values_bit_offset, output); \ + dict_bool_kernel(codes, codes_len, values, values_len, values_bit_offset, output); \ } // Generate fixed-width kernels for all native ptypes and decimal values. diff --git a/vortex-cuda/kernels/src/widen.cu b/vortex-cuda/kernels/src/widen.cu new file mode 100644 index 00000000000..9654b9db3a0 --- /dev/null +++ b/vortex-cuda/kernels/src/widen.cu @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "scalar_kernel.cuh" +#include "types.cuh" +#include + +template +struct WidenOp { + static_assert(sizeof(InputT) < sizeof(OutputT)); + static_assert(std::is_signed_v == std::is_signed_v); + + __device__ OutputT operator()(InputT value) const { + return static_cast(value); + } +}; + +#define GENERATE_WIDEN(input_suffix, InputT, output_suffix, OutputT) \ + extern "C" __global__ void widen_##input_suffix##_##output_suffix(const InputT *__restrict input, \ + OutputT *__restrict output, \ + uint64_t array_len) { \ + scalar_kernel(input, output, array_len, WidenOp {}); \ + } + +GENERATE_WIDEN(i8, int8_t, i16, int16_t) +GENERATE_WIDEN(i8, int8_t, i32, int32_t) +GENERATE_WIDEN(i8, int8_t, i64, int64_t) +GENERATE_WIDEN(i16, int16_t, i32, int32_t) +GENERATE_WIDEN(i16, int16_t, i64, int64_t) +GENERATE_WIDEN(i32, int32_t, i64, int64_t) +GENERATE_WIDEN(u8, uint8_t, u16, uint16_t) +GENERATE_WIDEN(u8, uint8_t, u32, uint32_t) +GENERATE_WIDEN(u8, uint8_t, u64, uint64_t) +GENERATE_WIDEN(u16, uint16_t, u32, uint32_t) +GENERATE_WIDEN(u16, uint16_t, u64, uint64_t) +GENERATE_WIDEN(u32, uint32_t, u64, uint64_t) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index b3c87d1abbf..67dfe0cd9ae 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -44,16 +44,13 @@ use vortex::array::arrays::varbinview::VarBinViewDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::buffer::DeviceBuffer; use vortex::array::builtins::ArrayBuiltins; -use vortex::array::match_each_decimal_value_type; use vortex::array::optimizer::ArrayOptimizer; use vortex::array::validity::Validity; use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::DecimalType; -use vortex::dtype::NativeDecimalType; use vortex::dtype::Nullability; use vortex::dtype::PType; -use vortex::dtype::i256; use vortex::encodings::fsst::FSST; use vortex::error::VortexResult; use vortex::error::vortex_bail; @@ -62,7 +59,6 @@ use vortex::error::vortex_ensure_eq; use vortex::error::vortex_err; use vortex::extension::datetime::AnyTemporal; use vortex_cuda_driver::CudaSlice; -use vortex_cuda_driver::DeviceRepr; use vortex_cuda_driver::LaunchConfig; use vortex_cuda_driver::PushKernelArg; use vortex_onpair::OnPair; @@ -92,6 +88,7 @@ use crate::executor::execute_validity_cuda; use crate::kernel::DecodedVarBin; use crate::kernel::decode_fsst_varbin; use crate::kernel::decode_onpair_varbin; +use crate::kernel::widen_decimal_values; use crate::stream::zero_padding; const BITMAP_THREADS_PER_BLOCK: u32 = 256; @@ -475,73 +472,11 @@ async fn export_decimal_values( "cannot export decimal values from {values_type} storage to Arrow {target_type}: narrowing would require a device-to-host overflow check", ); } - let values = ctx.ensure_on_device(values).await?; if values_type == target_type { - return Ok(values); + return ctx.ensure_on_device(values).await; } - match_each_decimal_value_type!(values_type, |S| { - export_decimal_values_from::(values, target_type, len, ctx).await - }) -} - -/// Dispatch from a concrete storage type `S` to the Arrow-required output width. -async fn export_decimal_values_from( - values: BufferHandle, - target_type: DecimalType, - len: usize, - ctx: &mut CudaExecutionCtx, -) -> VortexResult -where - S: NativeDecimalType + DeviceRepr, -{ - match target_type { - DecimalType::I32 => decimal_cast::(values, len, ctx).await, - DecimalType::I64 => decimal_cast::(values, len, ctx).await, - DecimalType::I128 => decimal_cast::(values, len, ctx).await, - DecimalType::I256 => decimal_cast::(values, len, ctx).await, - target_type => { - vortex_bail!("cannot export DecimalArray as Arrow decimal value type {target_type}") - } - } -} - -/// Launches the CUDA kernel that casts from Vortex storage type `S` to Arrow output type `D`. -/// -/// The caller must ensure this is not a narrowing cast. -async fn decimal_cast( - values: BufferHandle, - len: usize, - ctx: &mut CudaExecutionCtx, -) -> VortexResult -where - S: NativeDecimalType + DeviceRepr, - D: NativeDecimalType + DeviceRepr, -{ - if len == 0 { - return ctx - .ensure_on_device(BufferHandle::new_host( - Buffer::::empty().into_byte_buffer(), - )) - .await; - } - - let mut output_buffer = ctx.device_alloc::(len)?; - - let values_view = values.cuda_view::()?; - let len_u64 = len as u64; - let cuda_function = ctx.load_function_with_suffixes( - "decimal_cast", - &[&S::DECIMAL_TYPE.to_string(), &D::DECIMAL_TYPE.to_string()], - )?; - - ctx.launch_kernel(&cuda_function, len, |args| { - args.arg(&values_view).arg(&mut output_buffer).arg(&len_u64); - })?; - - Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( - output_buffer, - )?))) + widen_decimal_values(values, values_type, target_type, len, ctx).await } /// Export Vortex binary views as an Arrow Device array with `Utf8View`/`BinaryView` layout. diff --git a/vortex-cuda/src/canonical.rs b/vortex-cuda/src/canonical.rs index 138033a08ea..836f2e7544b 100644 --- a/vortex-cuda/src/canonical.rs +++ b/vortex-cuda/src/canonical.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use async_trait::async_trait; +use futures::future::BoxFuture; use futures::future::try_join_all; use vortex::array::ArrayRef; use vortex::array::Canonical; @@ -29,11 +30,17 @@ use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; +use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_bail; -/// Copy a canonical child array to the host. +/// Copy a child array to the host, executing it there if CUDA execution left it encoded. async fn child_into_host(child: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let child = if child.is_canonical() { + child + } else { + buffers_into_host(child).await? + }; Ok(child .execute::(ctx)? .into_host(ctx) @@ -41,6 +48,46 @@ async fn child_into_host(child: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResul .into_array()) } +/// Copies every device buffer of `array` and of its children to the host, keeping the encoding. +/// +/// A child that CUDA execution left encoded, such as a validity mask that is still a lazy `take` +/// through device-resident dictionary codes or a run-end encoded mask, can only be executed on +/// the host once all of its buffers are there. +fn buffers_into_host(array: ArrayRef) -> BoxFuture<'static, VortexResult> { + Box::pin(async move { + if array.is_host() { + return Ok(array); + } + let mut array = array; + if array.children_iter().any(|child| !child.is_host()) { + let mut slots = Vec::with_capacity(array.slots().len()); + for slot in array.slots().iter() { + slots.push(match slot { + Some(child) => Some(buffers_into_host(child.clone()).await?), + None => None, + }); + } + // SAFETY: the slots are byte-for-byte copies; only their placement changes. + array = unsafe { array.with_slots(slots.into())? }; + } + let buffers = array.buffer_handles(); + // Some encodings, such as Sparse, reject buffer replacement, so only replace when a + // buffer actually moves. + if buffers.iter().all(BufferHandle::is_on_host) { + return Ok(array); + } + let buffers = try_join_all(buffers.into_iter().map(|buffer| async move { + if buffer.is_on_host() { + return Ok::<_, VortexError>(buffer); + } + Ok(BufferHandle::new_host(buffer.try_into_host()?.await?)) + })) + .await?; + // SAFETY: the buffers are byte-for-byte copies; only their placement changes. + unsafe { array.with_buffers(buffers) } + }) +} + /// Move all canonical data from to_host from device. #[async_trait] pub trait CanonicalCudaExt { @@ -196,11 +243,19 @@ mod tests { use std::sync::LazyLock; use vortex::array::VortexSessionExecute; + use vortex::array::arrays::DictArray; + use vortex::array::assert_arrays_eq; use vortex::dtype::PType; + use vortex::encodings::runend::RunEnd; + use vortex::encodings::sparse::Sparse; + use vortex::scalar::Scalar; use vortex::session::VortexSession; use super::*; use crate::CudaSession; + use crate::executor::CudaExecute; + use crate::executor::CudaExecutionCtx; + use crate::kernel::DictExecutor; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -225,4 +280,133 @@ mod tests { assert!(Canonical::List(list).into_host(&mut ctx).await.is_err()); Ok(()) } + + /// Moves every buffer of `array` and of its children to the device without decoding it. + fn encoded_on_device(array: ArrayRef, cuda_ctx: &CudaExecutionCtx) -> VortexResult { + let mut slots = Vec::new(); + for slot in array.slots().iter() { + slots.push(match slot { + Some(child) => Some(encoded_on_device(child.clone(), cuda_ctx)?), + None => None, + }); + } + // SAFETY: the slots are byte-for-byte copies; only their placement changes. + let array = unsafe { array.with_slots(slots.into())? }; + // Sparse keeps its fill value in a host buffer that it doesn't allow to be replaced. + if array.as_opt::().is_some() { + return Ok(array); + } + let mut buffers = Vec::new(); + for buffer in array.buffer_handles() { + buffers.push(cuda_ctx.ensure_on_device_sync(buffer)?); + } + // SAFETY: the buffers are byte-for-byte copies; only their placement changes. + unsafe { array.with_buffers(buffers) } + } + + /// A GPU decode can leave a validity mask encoded on the device, here run-end encoded. + /// Copying the values to the host used to execute that mask on the host while its buffers + /// were still on the device, which panicked with "expected host buffer". + #[crate::test] + async fn test_into_host_copies_encoded_device_validity() -> VortexResult<()> { + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut ctx = SESSION.create_execution_ctx(); + let mask = RunEnd::new( + PrimitiveArray::new(Buffer::from(vec![2u32, 5, 6]), Validity::NonNullable).into_array(), + BoolArray::from_iter([true, false, true]).into_array(), + &mut ctx, + ) + .into_array(); + let mask = encoded_on_device(mask, &cuda_ctx)?; + let values = cuda_ctx + .stream() + .copy_to_device(vec![1i32, 2, 3, 4, 5, 6])? + .await?; + let array = PrimitiveArray::from_buffer_handle(values, PType::I32, Validity::Array(mask)); + + let host = Canonical::Primitive(array) + .into_host(cuda_ctx.execution_ctx()) + .await? + .into_array(); + let expected = + PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, None, None, Some(6)]); + assert_arrays_eq!(host, expected, &mut ctx); + Ok(()) + } + + /// A sparse validity mask keeps its fill value in a host buffer, which Sparse doesn't allow to + /// be replaced, while its patch indices and values are on the device. Only the children move + /// to the host. + #[crate::test] + async fn test_into_host_copies_sparse_device_validity() -> VortexResult<()> { + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut ctx = SESSION.create_execution_ctx(); + let mask = Sparse::try_new( + PrimitiveArray::new(Buffer::from(vec![1u32, 4]), Validity::NonNullable).into_array(), + BoolArray::from_iter([false, false]).into_array(), + 6, + Scalar::from(true), + &mut ctx, + )? + .into_array(); + let mask = encoded_on_device(mask, &cuda_ctx)?; + let values = cuda_ctx + .stream() + .copy_to_device(vec![1i32, 2, 3, 4, 5, 6])? + .await?; + let array = PrimitiveArray::from_buffer_handle(values, PType::I32, Validity::Array(mask)); + + let host = Canonical::Primitive(array) + .into_host(cuda_ctx.execution_ctx()) + .await? + .into_array(); + let expected = + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None, Some(6)]); + assert_arrays_eq!(host, expected, &mut ctx); + Ok(()) + } + + /// A dictionary with nullable values returns its validity as a lazy gather through the codes. + /// With the codes on the device, as a file scan leaves them, copying the decoded dictionary to + /// the host used to panic on a host-buffer read. + async fn assert_device_dict_into_host( + codes: PrimitiveArray, + values: ArrayRef, + ) -> VortexResult<()> { + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut ctx = SESSION.create_execution_ctx(); + let expected = DictArray::try_new(codes.clone().into_array(), values.clone())?; + let device_codes = encoded_on_device(codes.into_array(), &cuda_ctx)?; + let dict = DictArray::try_new(device_codes, values)?; + let host = DictExecutor + .execute(dict.into_array(), &mut cuda_ctx) + .await? + .into_host(cuda_ctx.execution_ctx()) + .await? + .into_array(); + assert_arrays_eq!(host, expected, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_into_host_string_dict_with_nullable_values() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_nullable_str([ + Some("a"), + None, + Some("a string longer than twelve bytes"), + ]); + let codes = PrimitiveArray::new( + Buffer::from(vec![0u8, 1, 2, 1, 0, 2]), + Validity::NonNullable, + ); + assert_device_dict_into_host(codes, values.into_array()).await + } + + #[crate::test] + async fn test_into_host_primitive_dict_with_nullable_values() -> VortexResult<()> { + let values = PrimitiveArray::from_option_iter([Some(7i32), None, Some(-3)]); + let codes = + PrimitiveArray::new(Buffer::from(vec![2u32, 0, 1, 1, 0]), Validity::NonNullable); + assert_device_dict_into_host(codes, values.into_array()).await + } } diff --git a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index d0405f7ee41..9498f1c8419 100644 --- a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs +++ b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs @@ -133,6 +133,8 @@ fn is_dyn_dispatch_compatible(array: &ArrayRef) -> bool { _ => false, }; } + // TODO(connor): Add Widen fusion after fixing sign extension of signed narrow BitPacked and + // FoR children in dynamic_dispatch.cu. Those paths currently zero-extend to the output width. (id == FoR.id() && array.as_::().references().is_global()) || id == ZigZag.id() || id == Primitive.id() diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index aca22975bdb..b5b866ad70e 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -143,6 +143,7 @@ async fn execute_dict_bool_typed( let values_view = values_device.cuda_view::()?; let codes_view = codes_device.cuda_view::()?; let codes_len_u64 = codes_len as u64; + let values_len_u64 = values_len as u64; let values_offset_u64 = values_meta.offset() as u64; let codes_ptype = I::PTYPE.to_string(); @@ -151,6 +152,7 @@ async fn execute_dict_bool_typed( args.arg(&codes_view) .arg(&codes_len_u64) .arg(&values_view) + .arg(&values_len_u64) .arg(&values_offset_u64) .arg(&mut output_slice); })?; @@ -171,6 +173,7 @@ async fn execute_dict_prim_typed VortexResult { assert!(!codes.is_empty()); let codes_len = codes.len(); + let values_len_u64 = values.len() as u64; let PrimitiveDataParts { ptype: value_ptype, @@ -202,6 +205,7 @@ async fn execute_dict_prim_typed( + codes: Vec, + valid: &[bool], + sane: T, + ) -> (PrimitiveArray, PrimitiveArray) { + let validity: Validity = valid.iter().copied().collect(); + let sane_codes: Vec = codes + .iter() + .zip(valid) + .map(|(&code, &valid)| if valid { code } else { sane }) + .collect(); + ( + PrimitiveArray::new(Buffer::from(codes), validity.clone()), + PrimitiveArray::new(Buffer::from(sane_codes), validity), + ) + } + + async fn gpu_dict_to_host( + dict: DictArray, + cuda_ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let canonical = DictExecutor.execute(dict.into_array(), cuda_ctx).await?; + Ok(canonical + .into_host(cuda_ctx.execution_ctx()) + .await? + .into_array()) + } + + async fn assert_gpu_matches( + codes: PrimitiveArray, + sane: PrimitiveArray, + values: ArrayRef, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let dict = DictArray::try_new(codes.into_array(), values.clone())?; + let expected = DictArray::try_new(sane.into_array(), values)?; + let actual = gpu_dict_to_host(dict, &mut cuda_ctx).await?; + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_cuda_dict_string_out_of_range_codes_under_nulls() -> VortexResult<()> { + let values = + VarBinViewArray::from_iter_str(["a", "bb", "a string longer than twelve bytes"]); + // `u32::MAX` points 64 GiB past the views, so an unchecked gather faults. + let (codes, sane) = garbage_under_nulls( + vec![0u32, 200, 2, u32::MAX, 1], + &[true, false, true, false, true], + 0, + ); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_primitive_far_out_of_range_codes_under_nulls() -> VortexResult<()> { + // These codes point gigabytes past the dictionary, so an unchecked gather faults. + let values = PrimitiveArray::new(Buffer::from(vec![10i64, 20, 30]), NonNullable); + let (codes, sane) = garbage_under_nulls( + vec![0u32, u32::MAX, 2, 0x8000_0000, 1], + &[true, false, true, false, true], + 0, + ); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_decimal_out_of_range_codes_under_nulls() -> VortexResult<()> { + let values = DecimalArray::from_iter([1i128, 2, 3], DecimalDType::new(38, 0)); + let (codes, sane) = garbage_under_nulls(vec![1u32, u32::MAX, 2], &[true, false, true], 0); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_bool_out_of_range_codes_under_nulls() -> VortexResult<()> { + let values = BoolArray::from_iter([true, false, true]); + let (codes, sane) = garbage_under_nulls( + vec![2u64, u64::MAX, 0, 1_000_000, 1, 9, 2, 0, 1], + &[true, false, true, false, true, false, true, true, true], + 0, + ); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + /// The shape the `compress_gpu` fuzzer found: a one-value dictionary whose null code is + /// 3,960,089,894. An unchecked gather reads about 4 GB past the dictionary and fails with + /// `CUDA_ERROR_ILLEGAL_ADDRESS`. + #[crate::test] + async fn test_cuda_dict_fuzz_garbage_null_code() -> VortexResult<()> { + let values = PrimitiveArray::new(Buffer::from(vec![0u8]), NonNullable); + let (codes, sane) = garbage_under_nulls(vec![0u32, 3_960_089_894], &[true, false], 0); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_all_null_codes() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_str(["x", "a string longer than twelve bytes"]); + let (codes, sane) = + garbage_under_nulls(vec![5_000u16, 7, u16::MAX], &[false, false, false], 0); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_string_codes_sliced_across_block_boundaries() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_str( + (0..50).map(|i| format!("value-{i}-{}", "x".repeat(i % 20))), + ); + let len = 3_000u16; + let valid: Vec = (0..len).map(|i| i % 11 != 0).collect(); + let raw: Vec = (0..len) + .map(|i| { + if valid[usize::from(i)] { + (i * 7) % 50 + } else { + 60_000 + } + }) + .collect(); + let (codes, sane) = garbage_under_nulls(raw, &valid, 0); + for range in [0..1023, 1023..1025, 1000..2049, 2047..2048, 1024..3000] { + let mut ctx = SESSION.create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let sliced = codes.clone().into_array().slice(range.clone())?; + let sliced_sane = sane.clone().into_array().slice(range.clone())?; + let dict = DictArray::try_new(sliced, values.clone().into_array())?; + let expected = DictArray::try_new(sliced_sane, values.clone().into_array())?; + let actual = gpu_dict_to_host(dict, &mut cuda_ctx).await?; + assert_arrays_eq!(actual, expected, &mut ctx); + } + Ok(()) + } + + #[crate::test] + async fn test_cuda_dict_string_u16_codes_full_range() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_str((0..65_536u32).map(|i| format!("v{i:05}"))); + let codes = PrimitiveArray::new( + Buffer::from(vec![0u16, u16::MAX, 32_768, 1, u16::MAX - 1]), + NonNullable, + ); + assert_gpu_matches(codes.clone(), codes, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_string_u32_codes_beyond_u16() -> VortexResult<()> { + let values = + VarBinViewArray::from_iter_str((0..70_000u32).map(|i| format!("dictionary-entry-{i}"))); + let codes = PrimitiveArray::new( + Buffer::from(vec![0u32, 69_999, 65_536, 12_345]), + NonNullable, + ); + assert_gpu_matches(codes.clone(), codes, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_string_inline_boundary_and_long_values() -> VortexResult<()> { + let long = "z".repeat(4_097); + let values = + VarBinViewArray::from_iter_str(["", "twelve-bytes", "thirteen-byte", long.as_str()]); + let codes = PrimitiveArray::new(Buffer::from(vec![3u8, 0, 1, 2, 3, 2, 1, 0]), NonNullable); + assert_gpu_matches(codes.clone(), codes, values.into_array()).await + } + + #[crate::test] + async fn test_cuda_dict_binary_non_utf8_values() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_nullable_bin([ + Some(vec![0xffu8, 0xfe, 0x00]), + Some(vec![0x80]), + Some(vec![]), + None, + Some(vec![0xc3; 20]), + ]); + let (codes, sane) = garbage_under_nulls( + vec![0u8, 1, 2, 3, 4, 77, 4, 0], + &[true, true, true, true, true, false, true, true], + 0, + ); + assert_gpu_matches(codes, sane, values.into_array()).await + } + + /// Moves the codes' buffer to the device, as a file scan leaves it. + async fn codes_on_device( + codes: PrimitiveArray, + cuda_ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let handle = cuda_ctx + .ensure_on_device(codes.buffer_handle().clone()) + .await?; + Ok(PrimitiveArray::from_buffer_handle( + handle, + codes.ptype(), + codes.validity()?, + )) + } + + /// Decodes with device-resident codes, as a file scan leaves them, and copies the result back + /// to the host. + async fn assert_gpu_matches_with_device_codes( + codes: PrimitiveArray, + sane: PrimitiveArray, + values: ArrayRef, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let device_codes = codes_on_device(codes, &mut cuda_ctx).await?; + let dict = DictArray::try_new(device_codes.into_array(), values.clone())?; + let expected = DictArray::try_new(sane.into_array(), values)?; + let actual = gpu_dict_to_host(dict, &mut cuda_ctx).await?; + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_cuda_dict_string_nullable_values_and_codes_device_codes() -> VortexResult<()> { + let values = VarBinViewArray::from_iter_nullable_str([Some("x"), None, Some("zz")]); + let (codes, sane) = garbage_under_nulls( + vec![0u16, 9_999, 2, 1, u16::MAX, 0], + &[true, false, true, true, false, true], + 0, + ); + assert_gpu_matches_with_device_codes(codes, sane, values.into_array()).await + } } diff --git a/vortex-cuda/src/kernel/decimal_cast.rs b/vortex-cuda/src/kernel/decimal_cast.rs new file mode 100644 index 00000000000..8f1fb0d9f87 --- /dev/null +++ b/vortex-cuda/src/kernel/decimal_cast.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Widening casts of decimal values buffers on the device. +//! +//! The Arrow export widens decimal values to the width that the Arrow schema fixes, and the +//! Widen executor widens a narrow decimal child to the canonical width of its decimal dtype. Both +//! launch the `decimal_cast` kernels through this module. +//! +//! These casts only widen. A narrowing cast can overflow, and to detect that on the device, the +//! kernel would need an overflow flag that synchronizes with the host. So this module rejects a +//! narrowing cast, and each caller decides how to handle values that are wider than its target. + +use std::sync::Arc; + +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_decimal_value_type; +use vortex::buffer::Buffer; +use vortex::dtype::DecimalType; +use vortex::dtype::NativeDecimalType; +use vortex::dtype::i256; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_ensure_eq; +use vortex::error::vortex_panic; +use vortex_cuda_driver::DeviceRepr; +use vortex_cuda_driver::PushKernelArg; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::CudaExecutionCtx; + +/// Widen `len` decimal values from `values_type` to the wider `target_type` on the device. +/// +/// The values can be on the host or on the device, and the result is on the device. Each value is +/// sign-extended, so the result holds the same values as the input. +/// +/// # Errors +/// +/// Returns an error if `target_type` is not wider than `values_type`, if the buffer does not +/// contain `len` complete values, or if the copy to the device or the kernel launch fails. +pub(crate) async fn widen_decimal_values( + values: BufferHandle, + values_type: DecimalType, + target_type: DecimalType, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + vortex_ensure!( + values_type < target_type, + "a decimal cast on the GPU must widen to a type wider than {values_type}, got \ + {target_type}" + ); + vortex_ensure_eq!(values.len(), len * values_type.byte_width()); + + let values = ctx.ensure_on_device(values).await?; + + match_each_decimal_value_type!(values_type, |S| { + decimal_cast_to::(values, target_type, len, ctx).await + }) +} + +/// Dispatch [`decimal_cast`] from `S` on the runtime `target_type`. +async fn decimal_cast_to( + values: BufferHandle, + target_type: DecimalType, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativeDecimalType + DeviceRepr, +{ + match target_type { + DecimalType::I8 => vortex_panic!( + "widen_decimal_values ensures a target wider than the values type, and no decimal \ + type is narrower than i8, got {target_type}" + ), + DecimalType::I16 => decimal_cast::(values, len, ctx).await, + DecimalType::I32 => decimal_cast::(values, len, ctx).await, + DecimalType::I64 => decimal_cast::(values, len, ctx).await, + DecimalType::I128 => decimal_cast::(values, len, ctx).await, + DecimalType::I256 => decimal_cast::(values, len, ctx).await, + } +} + +/// Launch the `decimal_cast__` kernel, which sign-extends device values of type `S` to the +/// wider type `D`. +/// +/// `D` **must** be wider than `S`. The `decimal_cast.cu` kernels only instantiate widening pairs, so +/// loading any other pair fails. [`widen_decimal_values`] checks this before it dispatches. +async fn decimal_cast( + values: BufferHandle, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativeDecimalType + DeviceRepr, + D: NativeDecimalType + DeviceRepr, +{ + if len == 0 { + return ctx + .ensure_on_device(BufferHandle::new_host( + Buffer::::empty().into_byte_buffer(), + )) + .await; + } + + let mut output_buffer = ctx.device_alloc::(len)?; + + let values_view = values.cuda_view::()?; + let len_u64 = len as u64; + let cuda_function = ctx.load_function_with_suffixes( + "decimal_cast", + &[&S::DECIMAL_TYPE.to_string(), &D::DECIMAL_TYPE.to_string()], + )?; + + ctx.launch_kernel(&cuda_function, len, |args| { + args.arg(&values_view).arg(&mut output_buffer).arg(&len_u64); + })?; + + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + output_buffer, + )?))) +} diff --git a/vortex-cuda/src/kernel/encodings/mod.rs b/vortex-cuda/src/kernel/encodings/mod.rs index fb292bc445c..fbc2defb696 100644 --- a/vortex-cuda/src/kernel/encodings/mod.rs +++ b/vortex-cuda/src/kernel/encodings/mod.rs @@ -1,6 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! CUDA executors for encoded arrays. +//! +//! Executors decode children through the CUDA execution context and return canonical arrays. + mod alp; mod bitpacked; mod date_time_parts; @@ -11,6 +15,7 @@ mod fsst; mod onpair; mod runend; mod sequence; +mod widen; mod zigzag; mod zstd; mod zstd_buffers; @@ -29,6 +34,7 @@ pub(crate) use onpair::OnPairExecutor; pub(crate) use onpair::decode_onpair_varbin; pub(crate) use runend::RunEndExecutor; pub(crate) use sequence::SequenceExecutor; +pub(crate) use widen::WidenExecutor; pub(crate) use zigzag::ZigZagExecutor; pub(crate) use zstd::ZstdExecutor; pub use zstd::ZstdKernelPrep; diff --git a/vortex-cuda/src/kernel/encodings/widen/mod.rs b/vortex-cuda/src/kernel/encodings/widen/mod.rs new file mode 100644 index 00000000000..be324b868b7 --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/widen/mod.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Restores the logical width of integer and decimal [`Widen`] arrays on the device. +//! +//! Children execute at their own width before widening, so compressed signed values retain their +//! sign. Execution keeps the child's validity and allocates a wider values buffer, unless decimal +//! storage already has the output width. + +use std::sync::Arc; + +use async_trait::async_trait; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::Widen; +use vortex::array::arrays::widen::WidenArraySlotsExt; +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_integer_ptype; +use vortex::dtype::DType; +use vortex::dtype::NativePType; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex_cuda_driver::DeviceRepr; +use vortex_cuda_driver::PushKernelArg; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::CudaExecutionCtx; +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::kernel::widen_decimal_values; + +/// CUDA execution for integer and decimal arrays stored at a narrower width. +#[derive(Debug)] +pub(crate) struct WidenExecutor; + +#[async_trait] +impl CudaExecute for WidenExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let array = array.try_downcast::().map_err(|array| { + vortex_err!( + "WidenExecutor requires a Widen array, got {}", + array.encoding_id() + ) + })?; + let child = array.values().clone().execute_cuda(ctx).await?; + + match (array.dtype(), child) { + (DType::Primitive(output_ptype, _), Canonical::Primitive(child)) => { + let len = child.len(); + let parts = child.into_data_parts(); + let values = ctx.ensure_on_device(parts.buffer).await?; + let output = match_each_integer_ptype!(parts.ptype, |S| { + match_each_integer_ptype!(*output_ptype, |D| { + widen_integers::(values, len, ctx)? + }) + }); + + Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( + output, + *output_ptype, + parts.validity, + ))) + } + (DType::Decimal(decimal_dtype, _), Canonical::Decimal(child)) => { + let len = child.len(); + let parts = child.into_data_parts(); + let target_type = decimal_dtype.canonical_values_type(); + vortex_ensure!( + parts.values_type <= target_type, + "CUDA Widen requires decimal storage no wider than {target_type}, got {}", + parts.values_type + ); + + let values = if parts.values_type == target_type { + ctx.ensure_on_device(parts.values).await? + } else { + widen_decimal_values(parts.values, parts.values_type, target_type, len, ctx) + .await? + }; + + Ok(Canonical::Decimal(DecimalArray::try_new_handle( + values, + target_type, + *decimal_dtype, + parts.validity, + )?)) + } + (dtype, child) => vortex_bail!( + "CUDA Widen requires an integer or decimal child matching {dtype}, got {}", + child.dtype() + ), + } + } +} + +/// Widens `len` device values of type `S` into a new buffer of type `D`. +/// +/// `D` must be strictly wider than `S` with the same signedness, as the [`Widen`] invariant +/// guarantees. The `widen.cu` kernels only instantiate those pairs. The other monomorphized pairs +/// cannot run because a valid `Widen` never dispatches to them. +fn widen_integers( + values: BufferHandle, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativePType + DeviceRepr, + D: NativePType + DeviceRepr, +{ + let input = values.cuda_view::()?; + let mut output = ctx.device_alloc::(len)?; + let len_u64 = len as u64; + let function = ctx.load_function("widen", &[S::PTYPE, D::PTYPE])?; + ctx.launch_kernel(&function, len, |args| { + args.arg(&input).arg(&mut output).arg(&len_u64); + })?; + + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + output, + )?))) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-cuda/src/kernel/encodings/widen/tests.rs b/vortex-cuda/src/kernel/encodings/widen/tests.rs new file mode 100644 index 00000000000..9b0781eeb73 --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/widen/tests.rs @@ -0,0 +1,246 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Device execution and CPU agreement for Widen arrays. + +use num_traits::Bounded; +use rstest::rstest; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::WidenArray; +use vortex::array::arrays::decimal::DecimalArrayExt; +use vortex::array::assert_arrays_eq; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::dtype::DType; +use vortex::dtype::DecimalDType; +use vortex::dtype::NativeDecimalType; +use vortex::dtype::NativePType; +use vortex::dtype::Nullability; +use vortex::dtype::PType; +use vortex::encodings::fastlanes::FoR; +use vortex::error::VortexResult; + +use crate::CanonicalCudaExt; +use crate::CudaDispatchMode; +use crate::CudaExecutionCtx; +use crate::CudaSession; +use crate::cuda_session; +use crate::executor::CudaArrayExt; + +/// Copies fixture values to the device while preserving their dtype and validity. +async fn upload(child: ArrayRef, ctx: &mut CudaExecutionCtx) -> VortexResult { + match child.execute::(ctx.execution_ctx())? { + Canonical::Primitive(child) => { + let parts = child.into_data_parts(); + Ok(PrimitiveArray::from_buffer_handle( + ctx.ensure_on_device(parts.buffer).await?, + parts.ptype, + parts.validity, + ) + .into_array()) + } + Canonical::Decimal(child) => { + let parts = child.into_data_parts(); + Ok(DecimalArray::try_new_handle( + ctx.ensure_on_device(parts.values).await?, + parts.values_type, + parts.decimal_dtype, + parts.validity, + )? + .into_array()) + } + _ => unreachable!("Widen test fixtures are primitive or decimal"), + } +} + +async fn check_widen( + host_child: ArrayRef, + device_child: ArrayRef, + dtype: DType, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<()> { + let host = WidenArray::try_new(host_child, dtype.clone())?.into_array(); + let expected = host.execute::(ctx.execution_ctx())?.into_array(); + let device = WidenArray::try_new(device_child, dtype.clone())?.into_array(); + assert!(!device.is_host()); + + let result = device.execute_cuda(ctx).await?; + assert_eq!(result.dtype(), &dtype); + assert_eq!(result.len(), expected.len()); + match &result { + Canonical::Primitive(array) => assert!(!array.buffer_handle().is_on_host()), + Canonical::Decimal(array) => { + assert!(!array.buffer_handle().is_on_host()); + assert_eq!( + array.values_type(), + dtype + .as_decimal_opt() + .expect("the result is decimal") + .canonical_values_type() + ); + } + _ => unreachable!("Widen returns primitive or decimal arrays"), + } + + let actual = result.into_host(ctx.execution_ctx()).await?.into_array(); + assert_arrays_eq!(actual, expected, ctx.execution_ctx()); + + Ok(()) +} + +fn validity(nullable: bool, len: usize) -> Validity { + if nullable { + Validity::from_iter((0..len).map(|i| i % 3 != 0)) + } else { + Validity::NonNullable + } +} + +#[rstest] +#[case::i8_i16(Buffer::from_iter([i8::MIN, i8::MAX, -1, 0, 1]), PType::I16)] +#[case::i8_i32(Buffer::from_iter([i8::MIN, i8::MAX, -1, 0, 1]), PType::I32)] +#[case::i8_i64(Buffer::from_iter([i8::MIN, i8::MAX, -1, 0, 1]), PType::I64)] +#[case::i16_i32(Buffer::from_iter([i16::MIN, i16::MAX, -1, 0, 1]), PType::I32)] +#[case::i16_i64(Buffer::from_iter([i16::MIN, i16::MAX, -1, 0, 1]), PType::I64)] +#[case::i32_i64(Buffer::from_iter([i32::MIN, i32::MAX, -1, 0, 1]), PType::I64)] +#[case::u8_u16(Buffer::from_iter([u8::MIN, u8::MAX, 0, 1, 1 << 7]), PType::U16)] +#[case::u8_u32(Buffer::from_iter([u8::MIN, u8::MAX, 0, 1, 1 << 7]), PType::U32)] +#[case::u8_u64(Buffer::from_iter([u8::MIN, u8::MAX, 0, 1, 1 << 7]), PType::U64)] +#[case::u16_u32(Buffer::from_iter([u16::MIN, u16::MAX, 0, 1, 1 << 15]), PType::U32)] +#[case::u16_u64(Buffer::from_iter([u16::MIN, u16::MAX, 0, 1, 1 << 15]), PType::U64)] +#[case::u32_u64(Buffer::from_iter([u32::MIN, u32::MAX, 0, 1, 1 << 31]), PType::U64)] +#[crate::test] +async fn test_integer_pairs( + #[case] values: Buffer, + #[case] output: PType, + #[values(false, true)] nullable: bool, + #[values(CudaDispatchMode::Auto, CudaDispatchMode::StandaloneOnly)] mode: CudaDispatchMode, +) -> VortexResult<()> { + let validity = validity(nullable, values.len()); + let child = PrimitiveArray::new(values, validity).into_array(); + let dtype = DType::Primitive(output, child.dtype().nullability()); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?.with_dispatch_mode(mode); + let device_child = upload(child.clone(), &mut ctx).await?; + + check_widen(child, device_child, dtype, &mut ctx).await +} + +#[rstest] +#[case::i8_p4([-99i8, 99], 2, 4)] +#[case::i8_p9([-99i8, 99], 2, 9)] +#[case::i8_p18([-99i8, 99], 2, 18)] +#[case::i8_p38([-99i8, 99], 2, 38)] +#[case::i8_p76([-99i8, 99], 2, 76)] +#[case::i16_p9([-9_999i16, 9_999], 4, 9)] +#[case::i16_p18([-9_999i16, 9_999], 4, 18)] +#[case::i16_p38([-9_999i16, 9_999], 4, 38)] +#[case::i16_p76([-9_999i16, 9_999], 4, 76)] +#[case::i32_p18([-999_999_999i32, 999_999_999], 9, 18)] +#[case::i32_p38([-999_999_999i32, 999_999_999], 9, 38)] +#[case::i32_p76([-999_999_999i32, 999_999_999], 9, 76)] +#[case::i64_p38([-999_999_999_999_999_999i64, 999_999_999_999_999_999], 18, 38)] +#[case::i64_p76([-999_999_999_999_999_999i64, 999_999_999_999_999_999], 18, 76)] +#[case::i128_p76([-(1i128 << 100), (1i128 << 100)], 38, 76)] +#[crate::test] +async fn test_decimal_pairs( + #[case] bounds: [T; 2], + #[case] child_precision: u8, + #[case] precision: u8, + #[values(false, true)] nullable: bool, + #[values(CudaDispatchMode::Auto, CudaDispatchMode::StandaloneOnly)] mode: CudaDispatchMode, +) -> VortexResult<()> { + let mut values = [ + T::default(), + bounds[0], + T::default(), + T::default(), + bounds[1], + ]; + if nullable { + values[0] = T::min_value(); + values[3] = T::max_value(); + } + + let child = DecimalArray::new( + Buffer::from_iter(values), + DecimalDType::new(child_precision, 1), + validity(nullable, 5), + ) + .into_array(); + let dtype = DType::Decimal(DecimalDType::new(precision, 1), child.dtype().nullability()); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?.with_dispatch_mode(mode); + let device_child = upload(child.clone(), &mut ctx).await?; + + check_widen(child, device_child, dtype, &mut ctx).await +} + +#[rstest] +#[crate::test] +async fn test_compressed_child( + #[values(CudaDispatchMode::Auto, CudaDispatchMode::StandaloneOnly)] mode: CudaDispatchMode, +) -> VortexResult<()> { + let child = PrimitiveArray::from_iter([0i8, 1, 50, 99]).into_array(); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?.with_dispatch_mode(mode); + let device_child = upload(child.clone(), &mut ctx).await?; + let host = FoR::try_new(child, (-100i8).into())?.into_array(); + let device = FoR::try_new(device_child, (-100i8).into())?.into_array(); + + check_widen(host, device, PType::I64.into(), &mut ctx).await +} + +#[rstest] +#[case::one(1)] +#[case::before_vector(15)] +#[case::vector(16)] +#[case::after_vector(17)] +#[case::before_block(2047)] +#[case::block(2048)] +#[case::after_block(2049)] +#[crate::test] +async fn test_kernel_boundaries(#[case] len: usize) -> VortexResult<()> { + let child = PrimitiveArray::from_iter((0..len).map(|i| if i % 2 == 0 { -99i8 } else { 99 })) + .into_array(); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?; + let device_child = upload(child.clone(), &mut ctx).await?; + + check_widen(child, device_child, PType::I64.into(), &mut ctx).await +} + +#[rstest] +#[case::equal_width(false)] +#[case::narrowing(true)] +#[crate::test] +async fn test_legacy_decimal_storage(#[case] narrowing: bool) -> VortexResult<()> { + let dtype = DecimalDType::new(9, 1); + let child = if narrowing { + DecimalArray::from_iter([-99i128, 0, 99], dtype) + } else { + DecimalArray::from_iter([-99i64, 0, 99], dtype) + } + .into_array(); + let logical = DType::Decimal(DecimalDType::new(18, 1), Nullability::NonNullable); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?; + let device_child = upload(child.clone(), &mut ctx).await?; + + if narrowing { + let array = WidenArray::try_new(device_child, logical)?.into_array(); + let error = array + .execute_cuda(&mut ctx) + .await + .expect_err("i128 storage requires narrowing to i64"); + assert!( + error + .to_string() + .contains("CUDA Widen requires decimal storage no wider than i64, got i128"), + "{error}" + ); + } else { + check_widen(child, device_child, logical, &mut ctx).await?; + } + + Ok(()) +} diff --git a/vortex-cuda/src/kernel/mod.rs b/vortex-cuda/src/kernel/mod.rs index f1eb24b1fad..94d15ffa6d0 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -19,6 +19,10 @@ use vortex_cuda_driver::LaunchConfig; use vortex_cuda_driver::sys::CUevent_flags; mod arrays; + +mod decimal_cast; +pub(crate) use decimal_cast::widen_decimal_values; + // Fatbins let the CUDA driver select among the configured real and virtual architectures. mod embedded_kernels { include!(concat!(env!("OUT_DIR"), "/embedded_kernels.rs")); diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index bdf9af1421b..0e60e61b2c1 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -54,6 +54,7 @@ use kernel::OnPairExecutor; use kernel::RunEndExecutor; use kernel::SharedExecutor; pub use kernel::TracingLaunchStrategy; +use kernel::WidenExecutor; use kernel::ZigZagExecutor; use kernel::ZstdBuffersExecutor; use kernel::ZstdExecutor; @@ -80,6 +81,7 @@ use vortex::array::arrays::Filter; use vortex::array::arrays::List; use vortex::array::arrays::Masked; use vortex::array::arrays::Slice; +use vortex::array::arrays::Widen; use vortex::encodings::alp::ALP; use vortex::encodings::datetime_parts::DateTimeParts; use vortex::encodings::decimal_byte_parts::DecimalByteParts; @@ -124,6 +126,7 @@ pub fn initialize_cuda(session: &CudaSession) { session.register_kernel(OnPair.id(), &OnPairExecutor); session.register_kernel(RunEnd.id(), &RunEndExecutor); session.register_kernel(Sequence.id(), &SequenceExecutor); + session.register_kernel(Widen.id(), &WidenExecutor); session.register_kernel(ZigZag.id(), &ZigZagExecutor); session.register_kernel(Zstd.id(), &ZstdExecutor); session.register_kernel(ZstdBuffers.id(), &ZstdBuffersExecutor); diff --git a/vortex-duckdb/src/exporter/mod.rs b/vortex-duckdb/src/exporter/mod.rs index 59a0e7fe728..fe0333ef1cf 100644 --- a/vortex-duckdb/src/exporter/mod.rs +++ b/vortex-duckdb/src/exporter/mod.rs @@ -157,7 +157,7 @@ impl ArrayExporter { /// the offset, len and `WritableVector` as options. Not sure what it should return though? /// This would allow Vortex extension authors to plug into the DuckDB exporter system. pub trait ColumnExporter: 'static { - /// The dictionary parts of this exporter, if it exports DuckDB dictionary vectors. + /// The dictionary parts of this exporter, if it can export DuckDB dictionary vectors. /// /// A chunked exporter uses them to export a range across chunks as one dictionary vector. fn as_dictionary(&self) -> Option<&dyn DictionaryParts> { @@ -175,6 +175,9 @@ pub trait ColumnExporter: 'static { } /// The values and codes of an exporter that emits DuckDB dictionary vectors. +/// +/// Codes can be generated per request rather than stored for the whole array. This capability also +/// applies when the exporter uses a constant vector for a request contained in one run. pub trait DictionaryParts { /// The DuckDB vector of the dictionary values. fn values(&self) -> &ReusableDict; @@ -183,6 +186,9 @@ pub trait DictionaryParts { fn values_key(&self) -> usize; /// Write the codes of rows `offset..offset + codes.len()` into `codes`. + /// + /// The offset is relative to the exported array, including any slice. The range must lie within + /// that array, and every code must index [`Self::values`], including codes for null rows. fn write_codes(&self, offset: usize, codes: &mut [u32]); } diff --git a/vortex-duckdb/src/exporter/run_end.rs b/vortex-duckdb/src/exporter/run_end.rs index fbe0120850a..d455dc3cc1e 100644 --- a/vortex-duckdb/src/exporter/run_end.rs +++ b/vortex-duckdb/src/exporter/run_end.rs @@ -1,15 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Batch-sized run selections over a cached DuckDB values dictionary. +//! +//! [`DictionaryParts`] lets the chunked exporter join run-end and dictionary chunks that share +//! values without flattening them. Codes are expanded only for the requested batch. Routing through +//! [`RunEndArray::to_dict`] would make the generic dictionary exporter materialize all codes at setup. + use std::marker::PhantomData; use vortex::array::ArrayRef; +use vortex::array::Canonical; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::match_each_integer_ptype; -use vortex::array::search_sorted::SearchSorted; -use vortex::array::search_sorted::SearchSortedSide; use vortex::dtype::IntegerPType; use vortex::encodings::runend::RunEndArray; use vortex::encodings::runend::RunEndArrayExt; @@ -22,17 +27,18 @@ use crate::duckdb::ReusableDict; use crate::duckdb::SelectionVector; use crate::duckdb::VectorRef; use crate::exporter::ColumnExporter; +use crate::exporter::DictionaryParts; use crate::exporter::cache::ConversionCache; use crate::exporter::cached_values_dict; use crate::exporter::canonical; -/// We export run-end arrays to a DuckDB dictionary vector. Values are exported -/// into a ReusableDict with SelectionVector applied in export(). +/// Keeps the values dictionary cached while generating selections for each requested batch. struct RunEndExporter { ends: PrimitiveArray, ends_type: PhantomData, values: ArrayRef, values_dict: ReusableDict, + /// Added to request offsets because sliced run ends retain their source coordinates. run_end_offset: usize, } @@ -42,8 +48,11 @@ pub(crate) fn new_exporter_with_flatten( ctx: &mut ExecutionCtx, flatten: bool, ) -> VortexResult> { - // Our canonicalization is faster than creating a dictionary vector and - // letting duckdb flatten it for us. + if array.is_empty() { + return canonical::new_exporter(Canonical::empty(array.dtype()).into_array(), cache, ctx); + } + + // Canonicalize directly to avoid building a dictionary that DuckDB would immediately flatten. if flatten { return canonical::new_exporter(array.into_array(), cache, ctx); } @@ -65,7 +74,73 @@ pub(crate) fn new_exporter_with_flatten( }) } +impl RunEndExporter { + /// Finds the run containing an absolute row offset. Run ends are exclusive. + /// + /// The offset must refer to an existing row, rather than the end of an empty request. + fn first_run(&self, offset: usize) -> usize { + let offset = + E::from_usize(offset).vortex_expect("the request offset fits the run-end type"); + self.ends + .as_slice::() + .partition_point(|&end| end <= offset) + } + + /// Fills a non-empty request with indices into the original values child. + /// + /// `offset` is absolute, and `first_run` must be the index returned by [`Self::first_run`]. + fn write_run_codes(&self, offset: usize, first_run: usize, codes: &mut [u32]) { + let mut position = + E::from_usize(offset).vortex_expect("the request offset fits the run-end type"); + let end_offset = position + + E::from_usize(codes.len()).vortex_expect("the request length fits the run-end type"); + let mut remaining = codes; + + for (index, &end) in self.ends.as_slice::()[first_run..].iter().enumerate() { + let next_end = end.min(end_offset); + let run_len = (next_end - position) + .to_usize() + .vortex_expect("run length fits usize"); + let code = u32::try_from(first_run + index).vortex_expect("run index exceeds u32"); + remaining[..run_len].fill(code); + remaining = &mut remaining[run_len..]; + if remaining.is_empty() { + return; + } + position = next_end; + } + + assert!( + remaining.is_empty(), + "run ends must cover the requested range" + ); + } +} + +impl DictionaryParts for RunEndExporter { + fn values(&self) -> &ReusableDict { + &self.values_dict + } + + fn values_key(&self) -> usize { + self.values.addr() + } + + fn write_codes(&self, offset: usize, codes: &mut [u32]) { + if codes.is_empty() { + return; + } + + let offset = self.run_end_offset + offset; + self.write_run_codes(offset, self.first_run(offset), codes); + } +} + impl ColumnExporter for RunEndExporter { + fn as_dictionary(&self) -> Option<&dyn DictionaryParts> { + Some(self) + } + fn export( &self, offset: usize, @@ -73,76 +148,75 @@ impl ColumnExporter for RunEndExporter { vector: &mut VectorRef, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let ends_slice = self.ends.as_slice::(); - - // Adjust offset to account for the run-end offset. - let mut offset = E::from_usize(self.run_end_offset + offset) - .vortex_expect("RunEndExporter::export: offset is not a valid value"); - // Compute the final end offset. - let end_offset = offset + E::from_usize(len).vortex_expect("len is not end type"); - - // Find the run that contains the start offset. - let start_run_idx = ends_slice - .search_sorted(&offset, SearchSortedSide::Right)? - .to_ends_index(ends_slice.len()); - - // Find the final run in case we can short-circuit and return a constant vector. - let end_run_idx = ends_slice - .search_sorted(&end_offset, SearchSortedSide::Right)? - .to_ends_index(ends_slice.len()); - - if start_run_idx == end_run_idx { - // NOTE(ngates): would be great if we could just export and set type == CONSTANT - // self.values_exporter.export(start_run_idx, 1, vector, cache); - let constant = self.values.execute_scalar(start_run_idx, ctx)?; - let value = constant.try_to_duckdb_scalar()?; - vector.reference_value(&value); + if len == 0 { return Ok(()); } - // Build up a selection vector - let mut sel_vec = SelectionVector::with_capacity(len); - let mut sel_vec_slice = unsafe { sel_vec.as_slice_mut(len) }; - - for (run_idx, &next_end) in ends_slice[start_run_idx..=end_run_idx].iter().enumerate() { - let next_end = next_end.min(end_offset); - let run_len = (next_end - offset) - .to_usize() - .vortex_expect("run_len is usize"); - - let global_run_idx = - u32::try_from(start_run_idx + run_idx).vortex_expect("run index exceeds u32"); - sel_vec_slice[..run_len].fill(global_run_idx); - sel_vec_slice = &mut sel_vec_slice[run_len..]; + let offset = self.run_end_offset + offset; + let first_run = self.first_run(offset); + let end = + E::from_usize(offset + len).vortex_expect("the request end fits the run-end type"); - offset = next_end; + // The request end is exclusive, so equality still leaves every row inside this run. + if end <= self.ends.as_slice::()[first_run] { + let constant = self.values.execute_scalar(first_run, ctx)?; + let value = constant.try_to_duckdb_scalar()?; + vector.reference_value(&value); + return Ok(()); } - debug_assert!(sel_vec_slice.is_empty()); - vector.reuse_dictionary(&self.values_dict, &sel_vec, len); + let mut selection = SelectionVector::with_capacity(len); + // SAFETY: the selection has capacity for `len` entries. The validated run ends cover the + // requested range, and write_run_codes fills every entry before DuckDB reads the selection. + let codes = unsafe { selection.as_slice_mut(len) }; + self.write_run_codes(offset, first_run, codes); + vector.reuse_dictionary(&self.values_dict, &selection, len); + Ok(()) } } #[cfg(test)] mod tests { + use rstest::rstest; + use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; + use vortex::array::arrays::DictArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; + use vortex::buffer::Buffer; use vortex::buffer::buffer; use vortex::encodings::runend::RunEnd; use vortex::error::VortexResult; use crate::SESSION; - use crate::cpp::duckdb_type::DUCKDB_TYPE_INTEGER; use crate::duckdb::DataChunk; use crate::duckdb::LogicalType; use crate::exporter::ArrayExporter; use crate::exporter::ConversionCache; use crate::exporter::new_array_exporter; + #[track_caller] + fn check_exports(array: ArrayRef, requests: &[(usize, usize, &str)]) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let exporter = new_array_exporter(array, &ConversionCache::default(), &mut ctx)?; + let mut chunk = DataChunk::new([LogicalType::int32()]); + + for &(offset, len, expected) in requests { + chunk.reset(); + exporter.export(offset, len, chunk.get_vector_mut(0), &mut ctx)?; + chunk.set_len(len); + assert_eq!( + String::try_from(&*chunk)?, + format!("Chunk - [1 Columns]\n{expected}\n") + ); + } + + Ok(()) + } + #[test] fn test_one_chunk_null() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); @@ -150,22 +224,10 @@ mod tests { let array = PrimitiveArray::from_option_iter(source); let array = RunEnd::encode(array.into_array(), &mut ctx)?; - let mut chunk = DataChunk::new([LogicalType::new(DUCKDB_TYPE_INTEGER)]); - new_array_exporter(array.into_array(), &ConversionCache::default(), &mut ctx)?.export( - 0, - 5, - chunk.get_vector_mut(0), - &mut ctx, - )?; - chunk.set_len(5); - let chunk_str = String::try_from(&*chunk)?; - assert_eq!( - chunk_str, - r#"Chunk - [1 Columns] -- DICTIONARY INTEGER: 5 = [ 0, 1, NULL, 3, NULL] -"# - ); - Ok(()) + check_exports( + array.into_array(), + &[(0, 5, "- DICTIONARY INTEGER: 5 = [ 0, 1, NULL, 3, NULL]")], + ) } #[test] @@ -191,4 +253,85 @@ mod tests { assert!(!exporter.export(&mut chunk, None)?); Ok(()) } + + #[test] + fn sliced_batches_keep_run_boundaries() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = RunEnd::try_new_offset_length( + buffer![3u32, 5, 8].into_array(), + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array(), + 2, + 5, + &mut ctx, + )?; + + check_exports( + array.into_array(), + &[ + (0, 1, "- CONSTANT INTEGER: 1 = [ 10]"), // Exact first run end. + (1, 2, "- CONSTANT INTEGER: 2 = [ NULL]"), // Exact null run. + (2, 3, "- DICTIONARY INTEGER: 3 = [ NULL, 30, 30]"), // Clipped batch. + (5, 0, "- FLAT INTEGER: 0 = [ ]"), // Empty request at the end. + ], + ) + } + + #[test] + fn empty_arrays_do_not_cache_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = buffer![10i32].into_array(); + let retained = + RunEnd::try_new_offset_length(buffer![3u32].into_array(), values, 3, 0, &mut ctx)?; + let empty = RunEnd::try_new( + Buffer::::empty().into_array(), + Buffer::::empty().into_array(), + &mut ctx, + )?; + + for array in [empty, retained] { + let cache = ConversionCache::default(); + let exporter = new_array_exporter(array.into_array(), &cache, &mut ctx)?; + let mut chunk = DataChunk::new([LogicalType::int32()]); + exporter.export(0, 0, chunk.get_vector_mut(0), &mut ctx)?; + assert!(cache.dict_cache.is_empty()); + } + + Ok(()) + } + + #[rstest] + #[case::runend(false, true)] + #[case::dictionary(true, true)] + #[case::different_values(false, false)] + fn chunk_boundaries(#[case] dictionary: bool, #[case] shared: bool) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array(); + let first = RunEnd::try_new_offset_length( + buffer![2u32, 4, 6].into_array(), + values.clone(), + 2, + 4, + &mut ctx, + )? + .into_array(); + let next_values = if shared { + values + } else { + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array() + }; + let second = if dictionary { + DictArray::try_new(buffer![0u32, 1, 1, 2].into_array(), next_values)?.into_array() + } else { + RunEnd::try_new(buffer![1u32, 3, 4].into_array(), next_values, &mut ctx)?.into_array() + }; + let dtype = first.dtype().clone(); + let array = ChunkedArray::try_new(vec![first, second], dtype)?.into_array(); + let crossing = if shared { + "- DICTIONARY INTEGER: 6 = [ NULL, 30, 30, 10, NULL, NULL]" + } else { + "- FLAT INTEGER: 6 = [ NULL, 30, 30, 10, NULL, NULL]" + }; + + check_exports(array, &[(1, 6, crossing)]) + } } diff --git a/vortex-json/src/arrow.rs b/vortex-json/src/arrow.rs index 30c311a72d7..0fa7b31992b 100644 --- a/vortex-json/src/arrow.rs +++ b/vortex-json/src/arrow.rs @@ -15,12 +15,11 @@ use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; -use vortex_array::dtype::extension::ExtVTable; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexExpect; @@ -39,28 +38,21 @@ fn has_valid_json_extension(field: &Field) -> bool { && ArrowJson::try_new_from_field_metadata(field.data_type(), field.metadata()).is_ok() } -impl ArrowExportVTable for Json { - fn arrow_ext_id(&self) -> Id { - *ARROW_JSON +impl ArrowExtVTable for Json { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - Json.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_JSON } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let DType::Extension(ext_dtype) = dtype else { - return Ok(None); - }; - if !ext_dtype.is::() { - return Ok(None); - } - let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; field .try_with_extension_type(ArrowJson::default()) @@ -75,15 +67,6 @@ impl ArrowExportVTable for Json { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_json = array - .dtype() - .as_extension_opt() - .map(|ext_dtype| ext_dtype.is::()) - .unwrap_or(false); - if !is_json { - return Ok(ArrowExport::Unsupported(array)); - } - ArrowJson::try_new_from_field_metadata(target.data_type(), target.metadata())?; let executed = array.execute::(ctx)?; @@ -98,52 +81,40 @@ impl ArrowExportVTable for Json { Ok(ArrowExport::Exported(storage)) } -} - -impl ArrowImportVTable for Json { - fn arrow_ext_id(&self) -> Id { - *ARROW_JSON - } fn from_arrow_field( - &self, field: &Field, _session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { if !has_valid_json_extension(field) { return Ok(None); } - Ok(Some(DType::Extension( - ExtDType::::try_new(EmptyMetadata, DType::Utf8(field.is_nullable().into()))? - .erased(), - ))) + Ok(Some(ExtDType::::try_new( + EmptyMetadata, + DType::Utf8(field.is_nullable().into()), + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let DType::Extension(ext_dtype) = dtype else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() || !has_valid_json_extension(field) { + if !has_valid_json_extension(field) { return Ok(ArrowImport::Unsupported(array)); } let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::new(ext_dtype.clone(), storage).into_array(), + ExtensionArray::new(ext_dtype.clone().erased(), storage).into_array(), )) } } #[cfg(test)] mod tests { - use std::sync::Arc; use std::sync::LazyLock; diff --git a/vortex-json/src/lib.rs b/vortex-json/src/lib.rs index 15c06c82b9f..2759846d1e1 100644 --- a/vortex-json/src/lib.rs +++ b/vortex-json/src/lib.rs @@ -13,8 +13,6 @@ mod arrow; mod dtype; mod json_to_variant; -use std::sync::Arc; - pub mod editions; pub use dtype::Json; pub use json_to_variant::JsonToVariant; @@ -31,8 +29,7 @@ use vortex_session::VortexSession; /// Register JSON extension support with a session. pub fn initialize(session: &VortexSession) { session.dtypes().register(Json); - session.arrow().register_exporter(Arc::new(Json)); - session.arrow().register_importer(Arc::new(Json)); + session.arrow().register_extension(Json); session.scalar_fns().register(JsonToVariant); // JSON is an opt-in durable dtype, so it belongs to an independently enabled edition family. diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index b82e301a849..f40b790604c 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -73,7 +73,8 @@ pub struct DictLayoutConstraints { /// /// The codes dtype is chosen upfront and is at most [`PType::U16`]. Primitive dictionaries /// use [`PType::U8`] when `max_len <= 256` or the value domain fits in 256 entries including - /// any null entry. String and binary dictionaries use [`PType::U8`] when `max_len <= 255`. + /// any null entry. String and binary dictionaries use [`PType::U8`] when `max_len <= 256`, as + /// [`codes_ptype_for_max_len`](vortex_array::builders::dict::codes_ptype_for_max_len) picks. /// /// Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for external compatibility. pub max_len: u16, diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index b48831cf879..bfb82b4e652 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -39,6 +39,8 @@ divan = { workspace = true } mimalloc = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-file = { workspace = true } +vortex-io = { workspace = true } vortex-layout = { workspace = true } vortex-spatial = { path = ".", features = ["_test-harness"] } diff --git a/vortex-spatial/README.md b/vortex-spatial/README.md index c3af7eec152..8aeae021435 100644 --- a/vortex-spatial/README.md +++ b/vortex-spatial/README.md @@ -63,16 +63,22 @@ Avoid per-row `scalar_at` or validity lookups in either path. Materialize coordi The current rules can prune: - positive `ST_Intersects(geometry_column, constant_geometry)` filters when the boxes are disjoint; +- positive `ST_Contains(geometry_column, constant_geometry)` filters when the file or zone box cannot cover the constant's box; +- positive `ST_Contains(constant_geometry, geometry_column)` filters (equivalently `ST_Within(geometry_column, constant_geometry)`) when the boxes are disjoint; - `<`, `<=`, `>`, and `>=` comparisons on `ST_Distance(geometry_column, constant_geometry)` using box-distance bounds. These proofs are conservative. If the expression has a different shape, the constant is null or empty, or the file has no AABB statistic, the chunk is kept. Box overlap alone never proves that two geometries intersect. +Containment is directional. A file or zone box that extends outside a constant's box may still contain rows within that constant, so that alone cannot prune `ST_Within`. Box coverage does not prove geometry containment either; `NOT ST_Contains` and `NOT ST_Within` have no spatial pruning proof. + +The same rules work with file-level AABB statistics, including nested geometry fields. To record file bounds, use `include_nested_statistics().with_file_statistics(...)` with `GeometryAabb.bind(EmptyOptions)` in the aggregate list. This replaces the default file aggregate list. The [containment pruning tests](src/prune/contains.rs) exercise file pruning without zone maps and zone pruning without file statistics. + ## Interoperability and limits Arrow import and export use the standard GeoArrow extension names and preserve CRS metadata. Native import only supports separated coordinates; interleaved GeoArrow arrays are rejected. -DuckDB uses WKB for general geometry values. Its Vortex expression lowering converts supported WKB constants to native scalars and pushes `ST_Contains`, `ST_Distance`, and `ST_Intersects` into a Vortex scan. +DuckDB uses WKB for general geometry values. Its Vortex expression lowering converts supported WKB constants to native scalars and pushes `ST_Contains`, `ST_Within`, `ST_Distance`, and `ST_Intersects` into a Vortex scan. `ST_Within(a, b)` lowers to `ST_Contains(b, a)` and shares its compute and pruning implementation. Some important limits: diff --git a/vortex-spatial/src/extension/linestring.rs b/vortex-spatial/src/extension/linestring.rs index 7e7f2de49d9..78f48f9f33a 100644 --- a/vortex-spatial/src/extension/linestring.rs +++ b/vortex-spatial/src/extension/linestring.rs @@ -40,9 +40,9 @@ use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -255,27 +255,31 @@ impl LineStringData { } } -impl ArrowExportVTable for LineString { - fn arrow_ext_id(&self) -> Id { - *ARROW_LINESTRING +impl ArrowExtVTable for LineString { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_LINESTRING } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = linestring_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(linestring_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = linestring_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = linestring_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -287,15 +291,6 @@ impl ArrowExportVTable for LineString { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_linestring = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_linestring { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(linestring_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -320,22 +315,11 @@ impl ArrowExportVTable for LineString { Ok(ArrowExport::Exported(Arc::new(linestrings.into_arrow()))) } -} - -impl ArrowImportVTable for LineString { - fn arrow_ext_id(&self) -> Id { - *ARROW_LINESTRING - } - /// Import a `geoarrow.linestring` field as the [`LineString`] dtype. Keyed off the standard - /// GeoArrow name, so any producer resolves here. Accepts the full `LineStringType` extension, or - /// — for a metadata-less geometry literal — the name alone, inferring the dimension from the - /// coordinate field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(linestring_meta) = field.try_extension_type::() { vortex_ensure!( @@ -368,23 +352,20 @@ impl ArrowImportVTable for LineString { }; let storage_dtype = linestring_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(LineString, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + LineString, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -392,7 +373,7 @@ impl ArrowImportVTable for LineString { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index b31887bd83b..33e7eb61e31 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -1,6 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Spatial extension types and metadata shared by their Arrow plugins. +//! +//! Geometry types derive their coordinate dimension from storage and carry an optional coordinate +//! reference system in [`SpatialMetadata`]. Shared metadata handling keeps these properties +//! consistent across the GeoArrow fields used to import and export each geometry. + pub(crate) mod coordinate; mod linestring; mod multilinestring; @@ -352,13 +358,23 @@ mod tests { use prost::Message; use rstest::rstest; + use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; use vortex_array::dtype::proto::dtype as pb; use vortex_array::dtype::proto::dtype::d_type::DtypeType; + use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; + use vortex_arrow::ExecuteArrowExt; + use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -374,6 +390,55 @@ mod tests { use crate::extension::coordinate::coordinate_storage_dtype; use crate::test_harness::spatial_session; + #[track_caller] + fn check_nested_arrow_round_trip( + vtable: impl ExtVTable, + depth: usize, + nullable: bool, + fixed: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let mut ctx = session.create_execution_ctx(); + let mut storage = coordinate_storage_dtype(Dimension::Xy, Nullability::NonNullable); + for _ in 0..depth { + storage = DType::List(Arc::new(storage), Nullability::NonNullable); + } + + storage = storage.with_nullability(nullable.into()); + let dtype = ExtDType::try_with_vtable( + vtable, + SpatialMetadata::new(Some("EPSG:4326".into())), + storage, + )? + .erased(); + let mut builder = + builder_with_capacity_in(dtype.storage_dtype(), 2, BufferAllocatorRef::static_ref()); + builder.append_zeros(2); + let array = ExtensionArray::try_new(dtype, builder.finish())?.into_array(); + let array = if fixed { + FixedSizeListArray::try_new(array, 1, Validity::NonNullable, 2)?.into_array() + } else { + ListArray::try_new( + array, + PrimitiveArray::from_iter([0u32, 1, 2]).into_array(), + Validity::NonNullable, + &mut ctx, + )? + .into_array() + }; + + let arrow = session.arrow(); + let field = arrow.to_arrow_field("geometry", array.dtype())?; + assert_eq!(arrow.from_arrow_field(&field)?, *array.dtype()); + let exported = array.clone().execute_arrow(Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + exported.to_data().validate_full()?; + let imported = arrow.from_arrow_array(exported, &field)?; + vortex_array::assert_arrays_eq!(imported, array, &mut ctx); + + Ok(()) + } + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); /// Test shim: decode with an explicitly constructed session. @@ -607,4 +672,22 @@ mod tests { ))); Ok(()) } + + #[rstest] + #[case::line(LineString, 1, false, false)] + #[case::points(MultiPoint, 1, false, false)] + #[case::polygon(Polygon, 2, false, false)] + #[case::multiline(MultiLineString, 2, false, false)] + #[case::multipolygon(MultiPolygon, 3, false, false)] + #[case::nullable(LineString, 1, true, false)] + #[case::fixed_size_list(LineString, 1, false, true)] + #[case::nullable_fixed_list(LineString, 1, true, true)] + fn nested_arrow_round_trip( + #[case] vtable: impl ExtVTable, + #[case] depth: usize, + #[case] nullable: bool, + #[case] fixed: bool, + ) -> VortexResult<()> { + check_nested_arrow_round_trip(vtable, depth, nullable, fixed) + } } diff --git a/vortex-spatial/src/extension/multilinestring.rs b/vortex-spatial/src/extension/multilinestring.rs index a4ea09631ea..6d16cea1e60 100644 --- a/vortex-spatial/src/extension/multilinestring.rs +++ b/vortex-spatial/src/extension/multilinestring.rs @@ -33,9 +33,9 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -185,27 +185,31 @@ impl MultiLineStringData { } } -impl ArrowExportVTable for MultiLineString { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTILINESTRING +impl ArrowExtVTable for MultiLineString { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTILINESTRING } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = multilinestring_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(multilinestring_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = multilinestring_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = multilinestring_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -217,15 +221,6 @@ impl ArrowExportVTable for MultiLineString { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_multilinestring = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_multilinestring { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(multilinestring_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -252,20 +247,11 @@ impl ArrowExportVTable for MultiLineString { multilinestrings.into_arrow(), ))) } -} - -impl ArrowImportVTable for MultiLineString { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTILINESTRING - } - /// Import a `geoarrow.multilinestring` field (matched by GeoArrow name). Accepts the full - /// `MultiLineStringType`, or a metadata-less literal (name only), inferring the dimension. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(multilinestring_meta) = field.try_extension_type::() { vortex_ensure!( @@ -301,23 +287,20 @@ impl ArrowImportVTable for MultiLineString { }; let storage_dtype = multilinestring_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(MultiLineString, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + MultiLineString, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -325,7 +308,7 @@ impl ArrowImportVTable for MultiLineString { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/multipoint.rs b/vortex-spatial/src/extension/multipoint.rs index e57cc57c5be..67bc2f27986 100644 --- a/vortex-spatial/src/extension/multipoint.rs +++ b/vortex-spatial/src/extension/multipoint.rs @@ -33,9 +33,9 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -172,27 +172,31 @@ impl MultiPointData { } } -impl ArrowExportVTable for MultiPoint { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTIPOINT +impl ArrowExtVTable for MultiPoint { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOINT } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = multipoint_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(multipoint_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = multipoint_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = multipoint_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -204,15 +208,6 @@ impl ArrowExportVTable for MultiPoint { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_multipoint = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_multipoint { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(multipoint_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -237,21 +232,11 @@ impl ArrowExportVTable for MultiPoint { Ok(ArrowExport::Exported(Arc::new(multipoints.into_arrow()))) } -} - -impl ArrowImportVTable for MultiPoint { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTIPOINT - } - /// Import a `geoarrow.multipoint` field as the [`MultiPoint`] dtype (matched by GeoArrow name). - /// Accepts the full `MultiPointType`, or a metadata-less literal (name only), inferring the - /// dimension from the coordinate field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(multipoint_meta) = field.try_extension_type::() { vortex_ensure!( @@ -282,23 +267,20 @@ impl ArrowImportVTable for MultiPoint { }; let storage_dtype = multipoint_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(MultiPoint, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + MultiPoint, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -306,7 +288,7 @@ impl ArrowImportVTable for MultiPoint { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/multipolygon.rs b/vortex-spatial/src/extension/multipolygon.rs index 711bdb016d3..a1f8552f587 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -32,9 +32,9 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -182,27 +182,31 @@ impl MultiPolygonData { } } -impl ArrowExportVTable for MultiPolygon { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTIPOLYGON +impl ArrowExtVTable for MultiPolygon { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOLYGON } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = multipolygon_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(multipolygon_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = multipolygon_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = multipolygon_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -214,15 +218,6 @@ impl ArrowExportVTable for MultiPolygon { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_multipolygon = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_multipolygon { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(multipolygon_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -247,20 +242,11 @@ impl ArrowExportVTable for MultiPolygon { Ok(ArrowExport::Exported(Arc::new(multipolygons.into_arrow()))) } -} - -impl ArrowImportVTable for MultiPolygon { - fn arrow_ext_id(&self) -> Id { - *ARROW_MULTIPOLYGON - } - /// Import a `geoarrow.multipolygon` field (matched by GeoArrow name). Accepts the full - /// `MultiPolygonType`, or a metadata-less literal (name only), inferring the dimension. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(multipolygon_meta) = field.try_extension_type::() { vortex_ensure!( @@ -299,23 +285,20 @@ impl ArrowImportVTable for MultiPolygon { }; let storage_dtype = multipolygon_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(MultiPolygon, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + MultiPolygon, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -323,7 +306,7 @@ impl ArrowImportVTable for MultiPolygon { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index 46b465289ec..5fc9a31985b 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -31,9 +31,9 @@ use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -151,26 +151,25 @@ pub(crate) fn point_geometries( .collect() } -impl ArrowExportVTable for Point { - fn arrow_ext_id(&self) -> Id { - *ARROW_POINT +impl ArrowExtVTable for Point { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_POINT } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = coordinate_dimension(ext_type.storage_dtype())?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = coordinate_dimension(ext_dtype.storage_dtype())?; - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; field.try_with_extension_type(point_type(spatial_metadata, dimension))?; Ok(Some(field)) @@ -183,15 +182,6 @@ impl ArrowExportVTable for Point { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_point = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_point { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(point_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -217,22 +207,11 @@ impl ArrowExportVTable for Point { Ok(ArrowExport::Exported(points.into_arrow())) } -} - -impl ArrowImportVTable for Point { - fn arrow_ext_id(&self) -> Id { - *ARROW_POINT - } - /// Import a `geoarrow.point` field as the [`Point`] dtype. Keyed off the standard GeoArrow name, - /// so any producer (DataFusion, DuckDB, geoarrow-rs, …) resolves here. Accepts the full - /// `PointType` extension, or — for a metadata-less geometry literal — the name alone, inferring - /// the dimension from the coordinate field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(point_meta) = field.try_extension_type::() { vortex_ensure!( @@ -262,23 +241,20 @@ impl ArrowImportVTable for Point { }; let storage_dtype = coordinate_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Point, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + Point, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::Struct(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -286,7 +262,7 @@ impl ArrowImportVTable for Point { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index b43101799cd..08a6120adec 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -38,9 +38,9 @@ use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::ExecuteArrowExt; @@ -244,27 +244,31 @@ impl PolygonData { } } -impl ArrowExportVTable for Polygon { - fn arrow_ext_id(&self) -> Id { - *ARROW_POLYGON +impl ArrowExtVTable for Polygon { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_POLYGON } fn to_arrow_field( &self, name: &str, - dtype: &DType, - session: &ArrowSession, + ext_dtype: &ExtDType, + _session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = polygon_dimension(ext_type.storage_dtype())?; - - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(polygon_type(spatial_metadata, dimension))?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = polygon_dimension(ext_dtype.storage_dtype())?; + + let arrow_type = polygon_type(spatial_metadata, dimension); + let mut field = Field::new( + name, + arrow_type.data_type(), + ext_dtype.storage_dtype().is_nullable(), + ); + field.try_with_extension_type(arrow_type)?; Ok(Some(field)) } @@ -276,15 +280,6 @@ impl ArrowExportVTable for Polygon { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_polygon = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_polygon { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(polygon_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -309,22 +304,11 @@ impl ArrowExportVTable for Polygon { Ok(ArrowExport::Exported(Arc::new(polygons.into_arrow()))) } -} - -impl ArrowImportVTable for Polygon { - fn arrow_ext_id(&self) -> Id { - *ARROW_POLYGON - } - /// Import a `geoarrow.polygon` field as the [`Polygon`] dtype. Keyed off the standard GeoArrow - /// name, so any producer (DataFusion, DuckDB, geoarrow-rs, …) resolves here. Accepts the full - /// `PolygonType` extension, or — for a metadata-less geometry literal — the name alone, inferring - /// the dimension from the coordinate field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(polygon_meta) = field.try_extension_type::() { vortex_ensure!( @@ -361,23 +345,20 @@ impl ArrowImportVTable for Polygon { }; let storage_dtype = polygon_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Polygon, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + Polygon, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::List(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -385,7 +366,7 @@ impl ArrowImportVTable for Polygon { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/rect.rs b/vortex-spatial/src/extension/rect.rs index 9680595cbe5..bccda5917ec 100644 --- a/vortex-spatial/src/extension/rect.rs +++ b/vortex-spatial/src/extension/rect.rs @@ -40,9 +40,9 @@ use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexResult; @@ -205,26 +205,25 @@ pub(crate) fn rect_geometries( .collect() } -impl ArrowExportVTable for Rect { - fn arrow_ext_id(&self) -> Id { - *ARROW_BOX +impl ArrowExtVTable for Rect { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_BOX } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); - let dimension = box_dimension(ext_type.storage_dtype())?; + let spatial_metadata = ext_dtype.metadata(); + let dimension = box_dimension(ext_dtype.storage_dtype())?; - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; field.try_with_extension_type(box_type(spatial_metadata, dimension))?; Ok(Some(field)) @@ -237,15 +236,6 @@ impl ArrowExportVTable for Rect { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_box = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_box { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(box_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -267,21 +257,11 @@ impl ArrowExportVTable for Rect { Ok(ArrowExport::Exported(Arc::new(rects.into_arrow()))) } -} - -impl ArrowImportVTable for Rect { - fn arrow_ext_id(&self) -> Id { - *ARROW_BOX - } - /// Import a `geoarrow.box` field as the [`Rect`] dtype. Keyed off the standard GeoArrow name, - /// so any producer resolves here. Accepts the full `BoxType` extension, or — for a - /// metadata-less literal — the name alone, inferring the dimension from the field names. fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let (dimension, metadata) = if let Ok(box_meta) = field.try_extension_type::() { ( box_meta.dimension().into(), @@ -305,23 +285,20 @@ impl ArrowImportVTable for Rect { }; let storage_dtype = box_storage_dtype(dimension, field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Rect, metadata, storage_dtype)?.erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + Rect, + metadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!(array.data_type(), DataType::Struct(_)) { return Ok(ArrowImport::Unsupported(array)); @@ -329,7 +306,7 @@ impl ArrowImportVTable for Rect { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } diff --git a/vortex-spatial/src/extension/wkb.rs b/vortex-spatial/src/extension/wkb.rs index 47713736c7a..d40b31449a7 100644 --- a/vortex-spatial/src/extension/wkb.rs +++ b/vortex-spatial/src/extension/wkb.rs @@ -26,9 +26,9 @@ use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; use vortex_error::VortexError; @@ -168,25 +168,24 @@ impl ExtVTable for WellKnownBinary { static ARROW_WKB: CachedId = CachedId::new(WkbType::NAME); -impl ArrowExportVTable for WellKnownBinary { - fn arrow_ext_id(&self) -> Id { - *ARROW_WKB +impl ArrowExtVTable for WellKnownBinary { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - self.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_WKB } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let spatial_metadata = ext_dtype.metadata(); - let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; field.try_with_extension_type(wkb_type(spatial_metadata))?; Ok(Some(field)) @@ -199,15 +198,6 @@ impl ArrowExportVTable for WellKnownBinary { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let is_wkb = array - .dtype() - .as_extension_opt() - .map(|ext| ext.is::()) - .unwrap_or(false); - if !is_wkb { - return Ok(ArrowExport::Unsupported(array)); - } - let Ok(wkb_meta) = target.try_extension_type::() else { return Ok(ArrowExport::Unsupported(array)); }; @@ -247,41 +237,30 @@ impl ArrowExportVTable for WellKnownBinary { Ok(ArrowExport::Exported(arrow_ref)) } -} - -impl ArrowImportVTable for WellKnownBinary { - fn arrow_ext_id(&self) -> Id { - *ARROW_WKB - } fn from_arrow_field( - &self, field: &Field, _session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { let Ok(wkb_meta) = field.try_extension_type::() else { return Ok(None); }; let storage_dtype = DType::Binary(field.is_nullable().into()); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(WellKnownBinary, spatial_metadata(&wkb_meta), storage_dtype)? - .erased(), - ))) + Ok(Some(ExtDType::try_with_vtable( + WellKnownBinary, + spatial_metadata(&wkb_meta), + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let Some(ext_dtype) = dtype.as_extension_opt() else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !ext_dtype.is::() - || field.try_extension_type::().is_err() + if field.try_extension_type::().is_err() || !matches!( array.data_type(), DataType::Binary | DataType::LargeBinary | DataType::BinaryView @@ -292,7 +271,7 @@ impl ArrowImportVTable for WellKnownBinary { let storage = session.from_arrow_array(array, field.is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::new(ext_dtype.clone(), storage).into_array(), + ExtensionArray::new(ext_dtype.clone().erased(), storage).into_array(), )) } } diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index ecb3e1e55fd..5704e476915 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -3,8 +3,6 @@ #![doc = include_str!(concat!("../", env!("CARGO_PKG_README")))] -use std::sync::Arc; - use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; @@ -23,6 +21,7 @@ use crate::extension::Point; use crate::extension::Polygon; use crate::extension::Rect; use crate::extension::WellKnownBinary; +use crate::prune::SpatialContainsPrune; use crate::prune::SpatialDistancePrune; use crate::prune::SpatialIntersectsPrune; use crate::scalar_fn::area::SpatialArea; @@ -47,31 +46,25 @@ mod tests; /// Set up a session with spatial extension types, Arrow interop, functions, and pruning. pub fn initialize(session: &VortexSession) { - // Register the spatial extension types. - session.dtypes().register(WellKnownBinary); - session.arrow().register_exporter(Arc::new(WellKnownBinary)); - session.arrow().register_importer(Arc::new(WellKnownBinary)); - session.dtypes().register(Point); - session.arrow().register_exporter(Arc::new(Point)); - session.arrow().register_importer(Arc::new(Point)); - session.dtypes().register(LineString); - session.arrow().register_exporter(Arc::new(LineString)); - session.arrow().register_importer(Arc::new(LineString)); - session.dtypes().register(MultiPoint); - session.arrow().register_exporter(Arc::new(MultiPoint)); - session.arrow().register_importer(Arc::new(MultiPoint)); - session.dtypes().register(Polygon); - session.arrow().register_exporter(Arc::new(Polygon)); - session.arrow().register_importer(Arc::new(Polygon)); - session.dtypes().register(MultiLineString); - session.arrow().register_exporter(Arc::new(MultiLineString)); - session.arrow().register_importer(Arc::new(MultiLineString)); - session.dtypes().register(MultiPolygon); - session.arrow().register_exporter(Arc::new(MultiPolygon)); - session.arrow().register_importer(Arc::new(MultiPolygon)); - session.dtypes().register(Rect); - session.arrow().register_exporter(Arc::new(Rect)); - session.arrow().register_importer(Arc::new(Rect)); + let dtypes = session.dtypes(); + let arrow = session.arrow(); + + dtypes.register(WellKnownBinary); + arrow.register_extension(WellKnownBinary); + dtypes.register(Point); + arrow.register_extension(Point); + dtypes.register(LineString); + arrow.register_extension(LineString); + dtypes.register(MultiPoint); + arrow.register_extension(MultiPoint); + dtypes.register(Polygon); + arrow.register_extension(Polygon); + dtypes.register(MultiLineString); + arrow.register_extension(MultiLineString); + dtypes.register(MultiPolygon); + arrow.register_extension(MultiPolygon); + dtypes.register(Rect); + arrow.register_extension(Rect); // Register the geometry scalar functions. session.scalar_fns().register(SpatialArea); @@ -89,6 +82,7 @@ pub fn initialize(session: &VortexSession) { session.aggregate_fns().register(GeometryAabb); // Register the spatial pruning rules that use that AABB. + session.stats().register_rewrite(SpatialContainsPrune); session.stats().register_rewrite(SpatialDistancePrune); session.stats().register_rewrite(SpatialIntersectsPrune); diff --git a/vortex-spatial/src/prune/contains.rs b/vortex-spatial/src/prune/contains.rs new file mode 100644 index 00000000000..f4ff037c6c8 --- /dev/null +++ b/vortex-spatial/src/prune/contains.rs @@ -0,0 +1,416 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Contains` pruning with a geometry column and a constant in either operand order. + +use vortex_array::expr::BoundExpression; +use vortex_array::expr::bound::get_item; +use vortex_array::expr::bound::gt; +use vortex_array::expr::bound::lit; +use vortex_array::expr::bound::lt; +use vortex_array::expr::bound::or; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_error::VortexResult; +use vortex_expr::stats::rewrite::StatsRewriteRule; +use vortex_session::VortexSession; + +use super::aabb_stat; +use super::geometry_and_constant; +use super::is_column; +use super::query_aabb; +use crate::scalar_fn::contains::SpatialContains; + +/// Prunes files and zones for positive `ST_Contains` filters using their aggregate AABB. +/// +/// For `contains(column, constant)`, the aggregate box must cover the constant's box. For +/// `contains(constant, column)` (equivalently `within(column, constant)`), disjoint boxes prove +/// no row can match. Neither box relation proves containment, so this rule supplies no proof +/// for `NOT ST_Contains`. +#[derive(Debug)] +pub struct SpatialContainsPrune; + +impl StatsRewriteRule for SpatialContainsPrune { + fn scalar_fn_id(&self) -> ScalarFnId { + SpatialContains.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + session: &VortexSession, + ) -> VortexResult> { + let Some((geom, constant)) = geometry_and_constant(expr)? else { + return Ok(None); + }; + let Some(query) = query_aabb(constant, session)? else { + return Ok(None); + }; + let bounds = aabb_stat(geom); + let field = |name| get_item(name, bounds.clone()); + let outside = if is_column(expr.child(0)) { + // Every row's box lies inside the aggregate box. If even that union cannot cover + // the constant, no individual geometry can contain it. + [ + gt(field("xmin"), lit(query.min().x)), + gt(field("ymin"), lit(query.min().y)), + lt(field("xmax"), lit(query.max().x)), + lt(field("ymax"), lit(query.max().y)), + ] + } else { + // The union may extend outside the constant while some rows remain inside it. + // Only strict separation proves every row is outside; touching boxes must scan. + [ + gt(field("xmin"), lit(query.max().x)), + gt(field("ymin"), lit(query.max().y)), + lt(field("xmax"), lit(query.min().x)), + lt(field("ymax"), lit(query.min().y)), + ] + }; + let [xmin, ymin, xmax, ymax] = outside; + Ok(Some(or(or(xmin, ymin), or(xmax, ymax)))) + } +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroUsize; + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::EmptyOptions as AggregateOptions; + use vortex_array::arrays::StructArray; + use vortex_array::expr::Expression; + use vortex_array::expr::fill_null; + use vortex_array::expr::get_item; + use vortex_array::expr::lit; + use vortex_array::expr::not; + use vortex_array::expr::root; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_buffer::ByteBuffer; + use vortex_buffer::ByteBufferMut; + use vortex_edition::EDITION_DECLARATIONS; + use vortex_edition::EDITION_FAMILIES; + use vortex_edition::EditionSessionExt; + use vortex_edition::declarations::core::CORE_2026_08_3; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_expr::stats::rewrite::BoundExpressionPruneExt; + use vortex_file::OpenOptionsSessionExt; + use vortex_file::WriteOptionsSessionExt; + use vortex_io::runtime::Handle; + use vortex_io::runtime::single::block_on; + use vortex_io::session::RuntimeSessionExt; + use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; + use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; + use vortex_layout::layouts::repartition::RepartitionStrategy; + use vortex_layout::layouts::repartition::RepartitionWriterOptions; + use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; + use vortex_layout::layouts::zoned::writer::ZonedStrategy; + use vortex_layout::layouts::zoned::zone_map::ZoneMap; + use vortex_mask::Mask; + use vortex_session::VortexSession; + + use crate::aggregate_fn::GeometryAabb; + use crate::prune::test_harness::aabb_zone_map; + use crate::prune::test_harness::empty_zone_map; + use crate::scalar_fn::contains::SpatialContains; + use crate::scalar_fn::envelope::SpatialEnvelope; + use crate::test_harness::multipoint_column; + use crate::test_harness::nullable_point_column; + use crate::test_harness::nullable_rect_column; + use crate::test_harness::point_column; + use crate::test_harness::rect_column; + use crate::test_harness::spatial_session; + + fn contains(column: Expression, constant: Scalar, column_contains: bool) -> Expression { + let operands = if column_contains { + [column, lit(constant)] + } else { + [lit(constant), column] + }; + SpatialContains.new_expr(EmptyOptions, operands) + } + + #[rstest] + #[case::covers([-1.0, -1.0, 3.0, 3.0], false, false)] + #[case::equal([0.0, 0.0, 2.0, 2.0], false, false)] + #[case::inside([0.5, 0.5, 1.5, 1.5], true, false)] + #[case::overlaps_right([1.0, 0.0, 3.0, 2.0], true, false)] + #[case::overlaps_left([-1.0, 0.0, 1.0, 2.0], true, false)] + #[case::overlaps_above([0.0, 1.0, 2.0, 3.0], true, false)] + #[case::overlaps_below([0.0, -1.0, 2.0, 1.0], true, false)] + #[case::disjoint_right([3.0, 0.0, 4.0, 2.0], true, true)] + #[case::disjoint_left([-2.0, 0.0, -1.0, 2.0], true, true)] + #[case::disjoint_above([0.0, 3.0, 2.0, 4.0], true, true)] + #[case::disjoint_below([0.0, -2.0, 2.0, -1.0], true, true)] + #[case::touches([2.0, 0.0, 3.0, 2.0], true, false)] + fn box_relations_respect_operand_order( + #[case] bounds: [f64; 4], + #[case] prune_column_contains: bool, + #[case] prune_constant_contains: bool, + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let query = rect_column(vec![(0.0, 0.0, 2.0, 2.0)])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = query.dtype().clone(); + let proof = contains(root(), query, column_contains) + .bind(&dtype)? + .falsify(&session)? + .ok_or_else(|| vortex_err!("missing contains proof"))?; + let expected = if column_contains { + prune_column_contains + } else { + prune_constant_contains + }; + let zones = aabb_zone_map(&dtype, &[bounds])?; + assert_eq!(zones.prune(&proof, &session)?, Mask::from_iter([expected])); + Ok(()) + } + + #[rstest] + fn equal_points_are_not_pruned( + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let query = point_column(vec![0.0], vec![0.0])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = query.dtype().clone(); + let proof = contains(root(), query, column_contains) + .bind(&dtype)? + .falsify(&session)? + .ok_or_else(|| vortex_err!("missing contains proof"))?; + let zones = aabb_zone_map(&dtype, &[[0.0; 4]])?; + assert_eq!(zones.prune(&proof, &session)?, Mask::new_false(1)); + Ok(()) + } + + #[rstest] + fn missing_and_null_bounds_keep_zones( + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let query = rect_column(vec![(0.0, 0.0, 2.0, 2.0)])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = query.dtype().clone(); + let proof = contains(root(), query, column_contains) + .bind(&dtype)? + .falsify(&session)? + .ok_or_else(|| vortex_err!("missing contains proof"))?; + let aabb = GeometryAabb.bind(AggregateOptions); + let null_bounds = ZoneMap::try_new( + dtype.clone(), + StructArray::from_fields(&[( + aabb.to_string().as_str(), + nullable_rect_column(vec![None, None])?, + )])?, + Arc::new([aabb]), + 1, + 2, + )?; + for zones in [empty_zone_map(&dtype)?, null_bounds] { + assert_eq!(zones.prune(&proof, &session)?, Mask::new_false(2)); + } + Ok(()) + } + + #[rstest] + fn null_and_empty_constants_decline( + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let empty = multipoint_column(vec![vec![]])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = empty.dtype().clone(); + for query in [Scalar::null(dtype.as_nullable()), empty] { + assert!( + contains(root(), query, column_contains) + .bind(&dtype)? + .falsify(&session)? + .is_none() + ); + } + Ok(()) + } + + #[rstest] + fn negated_and_computed_predicates_decline( + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + let session = spatial_session(); + let query = rect_column(vec![(0.0, 0.0, 2.0, 2.0)])? + .execute_scalar(0, &mut session.create_execution_ctx())?; + let dtype = query.dtype().clone(); + for predicate in [ + not(contains(root(), query.clone(), column_contains)), + contains( + SpatialEnvelope.new_expr(EmptyOptions, [root()]), + query, + column_contains, + ), + SpatialContains.new_expr(EmptyOptions, [root(), root()]), + ] { + assert!(predicate.bind(&dtype)?.falsify(&session)?.is_none()); + } + Ok(()) + } + + fn file_session(handle: Handle) -> VortexResult { + let session = VortexSession::empty().with_handle(handle); + for family in EDITION_FAMILIES { + session.editions().declare_family(family)?; + } + for declaration in EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + session.enable_edition(CORE_2026_08_3)?; + vortex_file::register_default_encodings(&session); + crate::initialize(&session); + Ok(session) + } + + #[rstest] + fn prunes_files_without_zone_maps( + #[values(false, true)] nested: bool, + #[values(false, true)] record_bounds: bool, + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + block_on(|handle| async move { + let session = file_session(handle)?; + let points = nullable_point_column(vec![Some((0.0, 0.0)), None, Some((1.0, 1.0))])?; + let batch = StructArray::from_fields(&[("geom", points)])?.into_array(); + let (batch, column) = if nested { + ( + StructArray::from_fields(&[("location", batch)])?.into_array(), + get_item("geom", get_item("location", root())), + ) + } else { + (batch, get_item("geom", root())) + }; + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .with_strategy(Arc::new(ChunkedLayoutStrategy::new( + FlatLayoutStrategy::default(), + ))) + .include_nested_statistics() + .with_file_statistics(if record_bounds { + vec![GeometryAabb.bind(AggregateOptions)] + } else { + vec![] + }) + .write(&mut buffer, batch.to_array_stream()) + .await?; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(buffer))?; + let mut ctx = session.create_execution_ctx(); + let far = + nullable_point_column(vec![Some((10.0, 10.0))])?.execute_scalar(0, &mut ctx)?; + let near = + nullable_point_column(vec![Some((0.0, 0.0))])?.execute_scalar(0, &mut ctx)?; + let enclosing = + rect_column(vec![(-1.0, -1.0, 2.0, 2.0)])?.execute_scalar(0, &mut ctx)?; + + for (query, cannot_match) in [(far, true), (near, false), (enclosing, column_contains)] + { + let predicate = contains(column.clone(), query, column_contains); + let matches = batch + .clone() + .apply(&fill_null(predicate.clone(), lit(false)))? + .execute::(&mut ctx)?; + let expected = record_bounds && cannot_match; + assert_eq!(file.can_prune(&predicate)?, expected, "{predicate}"); + let mask = file + .layout_reader()? + .pruning_evaluation(&(0..3), &predicate.bind(file.dtype())?, Mask::new_true(3))? + .await?; + assert_eq!(mask, Mask::from_iter([!expected; 3]), "{predicate}"); + assert_eq!( + &mask & &matches, + matches, + "pruning dropped matches for {predicate}" + ); + } + Ok(()) + }) + } + + #[rstest] + fn prunes_zones_without_file_statistics( + #[values(false, true)] record_bounds: bool, + #[values(false, true)] column_contains: bool, + ) -> VortexResult<()> { + block_on(|handle| async move { + let session = file_session(handle)?; + let array = nullable_point_column(vec![ + Some((0.0, 0.0)), + Some((1.0, 1.0)), + Some((10.0, 10.0)), + Some((11.0, 11.0)), + ])?; + let aggregates = if record_bounds { + vec![GeometryAabb.bind(AggregateOptions)] + } else { + vec![] + }; + let strategy = RepartitionStrategy::new( + ZonedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + FlatLayoutStrategy::default(), + ZonedLayoutOptions { + block_size: NonZeroUsize::new(2).vortex_expect("non-zero zone size"), + aggregate_fns: Some(aggregates.into()), + ..Default::default() + }, + ), + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: 2, + block_size_target: None, + canonicalize: false, + }, + ); + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .with_strategy(Arc::new(strategy)) + .with_file_statistics(vec![]) + .write(&mut buffer, array.to_array_stream()) + .await?; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(buffer))?; + assert!(file.file_stats().is_none()); + let mut ctx = session.create_execution_ctx(); + let query = + nullable_point_column(vec![Some((10.0, 10.0))])?.execute_scalar(0, &mut ctx)?; + let predicate = contains(root(), query, column_contains); + let matches = array + .clone() + .apply(&fill_null(predicate.clone(), lit(false)))? + .execute::(&mut ctx)?; + let mask = file + .layout_reader()? + .pruning_evaluation(&(0..4), &predicate.bind(file.dtype())?, Mask::new_true(4))? + .await?; + assert_eq!( + mask, + Mask::from_iter([!record_bounds, !record_bounds, true, true]) + ); + assert_eq!( + &mask & &matches, + matches, + "pruning dropped matches for {predicate}" + ); + Ok(()) + }) + } +} diff --git a/vortex-spatial/src/prune/mod.rs b/vortex-spatial/src/prune/mod.rs index dd1a6951503..9ea202c8cbf 100644 --- a/vortex-spatial/src/prune/mod.rs +++ b/vortex-spatial/src/prune/mod.rs @@ -1,20 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Chunk pruning for spatial filters, using the per-chunk [`GeometryAabb`] axis-aligned -//! bounding box (AABB). +//! File and zone pruning for spatial filters, using the [`GeometryAabb`] axis-aligned +//! bounding box (AABB) statistic. //! //! One module per predicate. To prune a new one: recover the geometry column and the constant -//! (symmetric predicates can use `geometry_and_constant`), reduce the constant to its box with -//! `query_aabb`, build the proof over `aabb_stat` from the box-relation builders (`min_dist_sq`, -//! `max_dist_sq`), and register the rule in [`crate::initialize`]. No new statistic or -//! file-format change is needed. +//! with `geometry_and_constant` (preserving operand order for asymmetric predicates), reduce the +//! constant with `query_aabb`, build the proof over `aabb_stat`, and register the rule in +//! [`crate::initialize`]. No new statistic or file-format change is needed. +mod contains; mod distance; mod intersects; #[cfg(test)] mod test_harness; +pub use contains::SpatialContainsPrune; pub use distance::SpatialDistancePrune; use geo::BoundingRect; use geo::Rect as SpatialRect; @@ -34,6 +35,7 @@ use vortex_array::expr::bound::lit; use vortex_array::expr::bound::lt; use vortex_array::expr::bound::lt_eq; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::literal::Literal; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; @@ -44,22 +46,19 @@ use crate::aggregate_fn::GeometryAabb; use crate::extension::is_native_geometry; use crate::extension::single_geometry; -/// Splits a symmetric two-operand spatial predicate into the scope-rooted geometry column and the -/// constant operand's scalar. +/// Splits a two-operand spatial predicate into the geometry column and the constant operand's scalar. /// /// `None` means the rule must decline: the expression doesn't have the `f(column, constant)` -/// shape (in either operand order), or the column's dtype carries no [`GeometryAabb`] statistic. -/// An asymmetric predicate (e.g. a future contains) must recover which operand is the column -/// itself instead of calling this. +/// shape (in either operand order), or the column's dtype does not support [`GeometryAabb`]. +/// Asymmetric predicates must also check which operand is the column when choosing their proof. fn geometry_and_constant( expr: &BoundExpression, ) -> VortexResult> { - // The predicate is symmetric, so the column (scope root) and the constant may be on either - // side. + // Zone predicates use the column as their root; file predicates refer to it through fields. let (lhs, rhs) = (expr.child(0), expr.child(1)); - let (geom, constant) = if lhs.is_root() { + let (geom, constant) = if is_column(lhs) { (lhs, rhs) - } else if rhs.is_root() { + } else if is_column(rhs) { (rhs, lhs) } else { return Ok(None); @@ -74,6 +73,13 @@ fn geometry_and_constant( Ok(constant.as_opt::().map(|scalar| (geom, scalar))) } +fn is_column(mut expr: &BoundExpression) -> bool { + while expr.is::() { + expr = expr.child(0); + } + expr.is_root() +} + /// The 2D bounding box of a constant geometry of any type, or `None` for one without an extent /// (e.g. an empty geometry). /// @@ -93,10 +99,10 @@ fn query_aabb( Ok(single_geometry(constant, &mut exec)?.bounding_rect()) } -/// The chunk's AABB statistic, as the storage struct with `xmin`/`ymin`/`xmax`/`ymax` fields. +/// The scope's AABB statistic, as the storage struct with `xmin`/`ymin`/`xmax`/`ymax` fields. /// -/// A chunk written without the statistic reads as null here; every proof built on top must let -/// that null propagate to its root, where the zone map keeps the chunk. +/// A missing statistic reads as null here; every proof built on top must let that null +/// propagate to its root, keeping the file or zone. fn aabb_stat(geom: &BoundExpression) -> BoundExpression { // `ext_storage` unwraps the native `geoarrow.box` stat value to its backing struct, so // proofs can `get_item` the coordinate fields. diff --git a/vortex-tensor/src/lib.rs b/vortex-tensor/src/lib.rs index b2939c21d82..59bb0cd8f37 100644 --- a/vortex-tensor/src/lib.rs +++ b/vortex-tensor/src/lib.rs @@ -12,8 +12,6 @@ allow(clippy::unwrap_used, clippy::expect_used, clippy::unwrap_in_result) )] -use std::sync::Arc; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; @@ -56,8 +54,7 @@ pub fn initialize(session: &VortexSession) { session.dtypes().register(FixedShapeTensor); let arrow_session = session.arrow(); - arrow_session.register_exporter(Arc::new(Vector)); - arrow_session.register_importer(Arc::new(Vector)); + arrow_session.register_extension(Vector); let session_fns = session.scalar_fns(); diff --git a/vortex-tensor/src/types/fixed_shape_tensor/proto.rs b/vortex-tensor/src/types/fixed_shape_tensor/proto.rs index 2bd608d717d..7633a8842c4 100644 --- a/vortex-tensor/src/types/fixed_shape_tensor/proto.rs +++ b/vortex-tensor/src/types/fixed_shape_tensor/proto.rs @@ -4,7 +4,6 @@ //! Protobuf serialization for [`FixedShapeTensorMetadata`]. use prost::Message; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -36,12 +35,12 @@ struct FixedShapeTensorMetadataProto { } /// Serializes [`FixedShapeTensorMetadata`] to protobuf bytes. -pub(crate) fn serialize(metadata: &FixedShapeTensorMetadata) -> Vec { +pub(crate) fn serialize(metadata: &FixedShapeTensorMetadata) -> VortexResult> { let logical_shape = metadata .logical_shape() .iter() - .map(|&d| u32::try_from(d).vortex_expect("dimension size exceeds u32")) - .collect(); + .map(|&dimension| u32::try_from(dimension)) + .collect::, _>>()?; let dim_names = metadata.dim_names().map(|n| n.to_vec()).unwrap_or_default(); @@ -49,9 +48,10 @@ pub(crate) fn serialize(metadata: &FixedShapeTensorMetadata) -> Vec { .permutation() .map(|p| { p.iter() - .map(|&i| u32::try_from(i).vortex_expect("permutation index exceeds u32")) - .collect() + .map(|&index| u32::try_from(index)) + .collect::, _>>() }) + .transpose()? .unwrap_or_default(); let proto = FixedShapeTensorMetadataProto { @@ -59,7 +59,8 @@ pub(crate) fn serialize(metadata: &FixedShapeTensorMetadata) -> Vec { dim_names, permutation, }; - proto.encode_to_vec() + + Ok(proto.encode_to_vec()) } /// Deserializes [`FixedShapeTensorMetadata`] from protobuf bytes. diff --git a/vortex-tensor/src/types/fixed_shape_tensor/vtable.rs b/vortex-tensor/src/types/fixed_shape_tensor/vtable.rs index b98e3dd0acf..ad29b091479 100644 --- a/vortex-tensor/src/types/fixed_shape_tensor/vtable.rs +++ b/vortex-tensor/src/types/fixed_shape_tensor/vtable.rs @@ -28,7 +28,7 @@ impl ExtVTable for FixedShapeTensor { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(proto::serialize(metadata)) + proto::serialize(metadata) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { @@ -119,4 +119,13 @@ mod tests { ) -> VortexResult<()> { assert_roundtrip(&metadata?) } + + #[test] + #[cfg(target_pointer_width = "64")] + fn unrepresentable_dimension_returns_error() -> VortexResult<()> { + let metadata = FixedShapeTensorMetadata::new(vec![0, u32::MAX as usize + 1]); + assert!(FixedShapeTensor.serialize_metadata(&metadata).is_err()); + + Ok(()) + } } diff --git a/vortex-tensor/src/types/vector/arrow.rs b/vortex-tensor/src/types/vector/arrow.rs index fafe66297c3..d7fadcfc471 100644 --- a/vortex-tensor/src/types/vector/arrow.rs +++ b/vortex-tensor/src/types/vector/arrow.rs @@ -13,7 +13,6 @@ use arrow_array::Array; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_schema::DataType; use arrow_schema::Field; -use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; use vortex_array::ArrayRef; use vortex_array::EmptyMetadata; use vortex_array::ExecutionCtx; @@ -22,14 +21,14 @@ use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; -use vortex_array::dtype::extension::ExtVTable; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportOptions; -use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowExtVTable; use vortex_arrow::ArrowImport; -use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowMetadataPolicy; use vortex_arrow::ArrowSession; use vortex_arrow::ExecuteArrowExt; +use vortex_arrow::ext_metadata::set_arrow_extension_name; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use vortex_session::registry::Id; @@ -41,18 +40,6 @@ pub const ARROW_VECTOR_EXTENSION_NAME: &str = "vortex.tensor.vector"; static ARROW_VECTOR: CachedId = CachedId::new(ARROW_VECTOR_EXTENSION_NAME); -#[expect( - clippy::disallowed_types, - reason = "Arrow's Field::set_metadata requires std::collections::HashMap" -)] -fn vector_extension_metadata() -> std::collections::HashMap { - [( - EXTENSION_TYPE_NAME_KEY.to_string(), - ARROW_VECTOR_EXTENSION_NAME.to_string(), - )] - .into() -} - fn is_supported_float(data_type: &DataType) -> bool { matches!( data_type, @@ -60,32 +47,24 @@ fn is_supported_float(data_type: &DataType) -> bool { ) } -impl ArrowExportVTable for Vector { - fn arrow_ext_id(&self) -> Id { - *ARROW_VECTOR +impl ArrowExtVTable for Vector { + fn metadata_policy() -> ArrowMetadataPolicy { + ArrowMetadataPolicy::ArrowNative } - fn vortex_id(&self) -> Id { - Vector.id() + fn arrow_ext_id(&self) -> Id { + *ARROW_VECTOR } fn to_arrow_field( &self, name: &str, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult> { - let DType::Extension(dtype) = dtype else { - return Ok(None); - }; - - if !dtype.is::() { - return Ok(None); - } + let mut field = session.to_arrow_field(name, ext_dtype.storage_dtype())?; + set_arrow_extension_name(&mut field, ARROW_VECTOR_EXTENSION_NAME); - // Delegate to Arrow encoding of storage type. - let mut field = session.to_arrow_field(name, dtype.storage_dtype())?; - field.set_metadata(vector_extension_metadata()); Ok(Some(field)) } @@ -96,14 +75,6 @@ impl ArrowExportVTable for Vector { options: &ArrowExportOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - if !array - .dtype() - .as_extension_opt() - .is_some_and(|ext| ext.is::()) - { - return Ok(ArrowExport::Unsupported(array)); - } - let executed = array.execute::(ctx)?; let storage = executed.storage_array().clone(); @@ -111,21 +82,15 @@ impl ArrowExportVTable for Vector { Ok(ArrowExport::Exported(arrow_storage)) } -} - -impl ArrowImportVTable for Vector { - fn arrow_ext_id(&self) -> Id { - *ARROW_VECTOR - } fn from_arrow_field( - &self, field: &Field, session: &ArrowSession, - ) -> VortexResult> { + ) -> VortexResult>> { if field.extension_type_name() != Some(ARROW_VECTOR_EXTENSION_NAME) { return Ok(None); } + let DataType::FixedSizeList(elem, list_size) = field.data_type() else { return Ok(None); }; @@ -138,24 +103,20 @@ impl ArrowImportVTable for Vector { *list_size as u32, field.is_nullable().into(), ); - Ok(Some(DType::Extension( - ExtDType::try_with_vtable(Vector, EmptyMetadata, storage_dtype)?.erased(), - ))) + + Ok(Some(ExtDType::try_with_vtable( + Vector, + EmptyMetadata, + storage_dtype, + )?)) } fn from_arrow_array( - &self, array: ArrowArrayRef, _field: &Field, - dtype: &DType, + ext_dtype: &ExtDType, session: &ArrowSession, ) -> VortexResult { - let DType::Extension(dtype) = dtype else { - return Ok(ArrowImport::Unsupported(array)); - }; - if !dtype.is::() { - return Ok(ArrowImport::Unsupported(array)); - } let DataType::FixedSizeList(elem, _) = array.data_type() else { return Ok(ArrowImport::Unsupported(array)); }; @@ -163,9 +124,9 @@ impl ArrowImportVTable for Vector { return Ok(ArrowImport::Unsupported(array)); } - let storage = session.from_arrow_array(array, dtype.is_nullable())?; + let storage = session.from_arrow_array(array, ext_dtype.storage_dtype().is_nullable())?; Ok(ArrowImport::Imported( - ExtensionArray::try_new(dtype.clone(), storage)?.into_array(), + ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array(), )) } } @@ -190,7 +151,6 @@ mod tests { use vortex_array::dtype::StructFields; use vortex_array::dtype::extension::ExtDType; use vortex_array::validity::Validity; - use vortex_arrow::ArrowExport; use vortex_arrow::ArrowImport; use vortex_arrow::ArrowSessionExt; use vortex_buffer::buffer; @@ -235,8 +195,7 @@ mod tests { /// A session whose only Arrow plugins are the [`Vector`] exporter and importer. fn session_with_vector() -> VortexSession { let session = VortexSession::empty(); - session.arrow().register_exporter(Arc::new(Vector)); - session.arrow().register_importer(Arc::new(Vector)); + session.arrow().register_extension(Vector); session } @@ -264,15 +223,15 @@ mod tests { fn from_arrow_field_recovers_vector_dtype() -> VortexResult<()> { let vortex_session = session_with_vector(); let session = vortex_session.arrow(); - let arrow_field = Field::new( + let mut arrow_field = Field::new( "embedding", DataType::FixedSizeList( Arc::new(Field::new("item", DataType::Float32, false)), DIM as i32, ), true, - ) - .with_metadata(vector_extension_metadata()); + ); + set_arrow_extension_name(&mut arrow_field, ARROW_VECTOR_EXTENSION_NAME); let dtype = session.from_arrow_field(&arrow_field)?; assert_eq!(dtype, vector_dtype(true)); Ok(()) @@ -362,7 +321,7 @@ mod tests { ), false, ); - arrow_field.set_metadata(vector_extension_metadata()); + set_arrow_extension_name(&mut arrow_field, ARROW_VECTOR_EXTENSION_NAME); let arrow = arrow_fsl_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], DIM as i32); let imported = session.from_arrow_array(arrow, &arrow_field)?; assert!( @@ -381,20 +340,16 @@ mod tests { } #[test] - fn execute_arrow_returns_unsupported_for_non_vector() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); + fn execute_arrow_rejects_incompatible_non_vector() -> VortexResult<()> { + let session = session_with_vector(); + let mut ctx = session.create_execution_ctx(); let primitive = PrimitiveArray::from_iter([1_i32, 2, 3]).into_array(); - let target = Field::new("ints", DataType::Int32, false); - let result = ::execute_arrow( - &Vector, - primitive.clone(), - &target, - &ArrowExportOptions::default(), - &mut ctx, - )?; - assert!( - matches!(result, ArrowExport::Unsupported(arr) if arr.dtype() == primitive.dtype()) - ); + let target = session + .arrow() + .to_arrow_field("embedding", &vector_dtype(false))?; + + assert!(primitive.execute_arrow(Some(&target), &mut ctx).is_err()); + Ok(()) } @@ -404,36 +359,10 @@ mod tests { let field = Field::new("embedding", DataType::Int32, false); let int_array: ArrowArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let result = ::from_arrow_array( - &Vector, + let result = ::from_arrow_array( int_array, &field, - &dtype, - &VortexSession::empty().arrow(), - )?; - assert!(matches!(result, ArrowImport::Unsupported(_))); - Ok(()) - } - - #[test] - fn from_arrow_array_returns_unsupported_for_non_vector_dtype() -> VortexResult<()> { - use vortex_uuid::Uuid; - use vortex_uuid::UuidMetadata; - let uuid_storage = DType::FixedSizeList( - Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)), - 16, - Nullability::NonNullable, - ); - let uuid_ext = - ExtDType::try_with_vtable(Uuid, UuidMetadata::default(), uuid_storage)?.erased(); - - let fsl_arrow = arrow_fsl_f32(vec![1.0, 2.0, 3.0], DIM as i32); - let field = Field::new("embedding", fsl_arrow.data_type().clone(), false); - let result = ::from_arrow_array( - &Vector, - fsl_arrow, - &field, - &DType::Extension(uuid_ext), + &dtype.as_extension().clone().downcast::(), &VortexSession::empty().arrow(), )?; assert!(matches!(result, ArrowImport::Unsupported(_)));