diff --git a/vortex-array/public-api.lock b/vortex-array/public-api.lock index 55698904758..aff0c726750 100644 --- a/vortex-array/public-api.lock +++ b/vortex-array/public-api.lock @@ -17350,6 +17350,90 @@ pub fn vortex_array::scalar_fn::EmptyOptions::hash<__H: core::hash::Hasher>(&sel impl core::marker::StructuralPartialEq for vortex_array::scalar_fn::EmptyOptions +pub struct vortex_array::scalar_fn::ForeignScalarFnOptions + +impl vortex_array::scalar_fn::ForeignScalarFnOptions + +pub fn vortex_array::scalar_fn::ForeignScalarFnOptions::new(metadata: alloc::vec::Vec, arity: usize) -> Self + +impl core::clone::Clone for vortex_array::scalar_fn::ForeignScalarFnOptions + +pub fn vortex_array::scalar_fn::ForeignScalarFnOptions::clone(&self) -> vortex_array::scalar_fn::ForeignScalarFnOptions + +impl core::cmp::Eq for vortex_array::scalar_fn::ForeignScalarFnOptions + +impl core::cmp::PartialEq for vortex_array::scalar_fn::ForeignScalarFnOptions + +pub fn vortex_array::scalar_fn::ForeignScalarFnOptions::eq(&self, other: &vortex_array::scalar_fn::ForeignScalarFnOptions) -> bool + +impl core::fmt::Debug for vortex_array::scalar_fn::ForeignScalarFnOptions + +pub fn vortex_array::scalar_fn::ForeignScalarFnOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result + +impl core::fmt::Display for vortex_array::scalar_fn::ForeignScalarFnOptions + +pub fn vortex_array::scalar_fn::ForeignScalarFnOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result + +impl core::hash::Hash for vortex_array::scalar_fn::ForeignScalarFnOptions + +pub fn vortex_array::scalar_fn::ForeignScalarFnOptions::hash<__H: core::hash::Hasher>(&self, state: &mut __H) + +impl core::marker::StructuralPartialEq for vortex_array::scalar_fn::ForeignScalarFnOptions + +pub struct vortex_array::scalar_fn::ForeignScalarFnVTable + +impl vortex_array::scalar_fn::ForeignScalarFnVTable + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::make_scalar_fn(id: vortex_array::scalar_fn::ScalarFnId, metadata: alloc::vec::Vec, arity: usize) -> vortex_array::scalar_fn::ScalarFnRef + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::new(id: vortex_array::scalar_fn::ScalarFnId) -> Self + +impl core::clone::Clone for vortex_array::scalar_fn::ForeignScalarFnVTable + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::clone(&self) -> vortex_array::scalar_fn::ForeignScalarFnVTable + +impl core::fmt::Debug for vortex_array::scalar_fn::ForeignScalarFnVTable + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result + +impl vortex_array::scalar_fn::ScalarFnVTable for vortex_array::scalar_fn::ForeignScalarFnVTable + +pub type vortex_array::scalar_fn::ForeignScalarFnVTable::Options = vortex_array::scalar_fn::ForeignScalarFnOptions + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::arity(&self, options: &Self::Options) -> vortex_array::scalar_fn::Arity + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::child_name(&self, _options: &Self::Options, child_idx: usize) -> vortex_array::scalar_fn::ChildName + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::coerce_args(&self, options: &Self::Options, args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::deserialize(&self, metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::execute(&self, _options: &Self::Options, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::id(&self) -> vortex_array::scalar_fn::ScalarFnId + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::is_fallible(&self, options: &Self::Options) -> bool + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::is_null_sensitive(&self, options: &Self::Options) -> bool + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::reduce(&self, options: &Self::Options, node: &dyn vortex_array::scalar_fn::ReduceNode, ctx: &dyn vortex_array::scalar_fn::ReduceCtx) -> vortex_error::VortexResult> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::return_dtype(&self, _options: &Self::Options, _args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::serialize(&self, options: &Self::Options) -> vortex_error::VortexResult>> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::simplify(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, ctx: &dyn vortex_array::scalar_fn::SimplifyCtx) -> vortex_error::VortexResult> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::simplify_untyped(&self, options: &Self::Options, expr: &vortex_array::expr::Expression) -> vortex_error::VortexResult> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::stat_expression(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, stat: vortex_array::expr::stats::Stat, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::stat_falsification(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult> + pub struct vortex_array::scalar_fn::ScalarFn impl vortex_array::scalar_fn::ScalarFn @@ -17564,6 +17648,44 @@ pub fn vortex_array::scalar_fn::ScalarFnVTable::stat_falsification(&self, option pub fn vortex_array::scalar_fn::ScalarFnVTable::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult> +impl vortex_array::scalar_fn::ScalarFnVTable for vortex_array::scalar_fn::ForeignScalarFnVTable + +pub type vortex_array::scalar_fn::ForeignScalarFnVTable::Options = vortex_array::scalar_fn::ForeignScalarFnOptions + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::arity(&self, options: &Self::Options) -> vortex_array::scalar_fn::Arity + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::child_name(&self, _options: &Self::Options, child_idx: usize) -> vortex_array::scalar_fn::ChildName + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::coerce_args(&self, options: &Self::Options, args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::deserialize(&self, metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::execute(&self, _options: &Self::Options, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::id(&self) -> vortex_array::scalar_fn::ScalarFnId + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::is_fallible(&self, options: &Self::Options) -> bool + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::is_null_sensitive(&self, options: &Self::Options) -> bool + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::reduce(&self, options: &Self::Options, node: &dyn vortex_array::scalar_fn::ReduceNode, ctx: &dyn vortex_array::scalar_fn::ReduceCtx) -> vortex_error::VortexResult> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::return_dtype(&self, _options: &Self::Options, _args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::serialize(&self, options: &Self::Options) -> vortex_error::VortexResult>> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::simplify(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, ctx: &dyn vortex_array::scalar_fn::SimplifyCtx) -> vortex_error::VortexResult> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::simplify_untyped(&self, options: &Self::Options, expr: &vortex_array::expr::Expression) -> vortex_error::VortexResult> + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::stat_expression(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, stat: vortex_array::expr::stats::Stat, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::stat_falsification(&self, options: &Self::Options, expr: &vortex_array::expr::Expression, catalog: &dyn vortex_array::expr::pruning::StatsCatalog) -> core::option::Option + +pub fn vortex_array::scalar_fn::ForeignScalarFnVTable::validity(&self, options: &Self::Options, expression: &vortex_array::expr::Expression) -> vortex_error::VortexResult> + impl vortex_array::scalar_fn::ScalarFnVTable for vortex_array::scalar_fn::fns::between::Between pub type vortex_array::scalar_fn::fns::between::Between::Options = vortex_array::scalar_fn::fns::between::BetweenOptions diff --git a/vortex-array/src/aggregate_fn/foreign.rs b/vortex-array/src/aggregate_fn/foreign.rs new file mode 100644 index 00000000000..e79d2456493 --- /dev/null +++ b/vortex-array/src/aggregate_fn/foreign.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::Columnar; +use crate::ExecutionCtx; +use crate::aggregate_fn::AggregateFn; +use crate::aggregate_fn::AggregateFnId; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::AggregateFnVTable; +use crate::dtype::DType; +use crate::scalar::Scalar; + +/// Options payload for a foreign aggregate function. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct ForeignAggregateFnOptions { + metadata: Vec, +} + +impl ForeignAggregateFnOptions { + pub fn new(metadata: Vec) -> Self { + Self { metadata } + } +} + +impl Display for ForeignAggregateFnOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "foreign(metadata={}B)", self.metadata.len()) + } +} + +/// Aggregate-function placeholder used when deserializing an unknown aggregate function ID. +#[derive(Clone, Debug)] +pub struct ForeignAggregateFnVTable { + id: AggregateFnId, +} + +impl ForeignAggregateFnVTable { + pub fn new(id: AggregateFnId) -> Self { + Self { id } + } +} + +impl AggregateFnVTable for ForeignAggregateFnVTable { + type Options = ForeignAggregateFnOptions; + type Partial = (); + + fn id(&self) -> AggregateFnId { + self.id.clone() + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.metadata.clone())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(ForeignAggregateFnOptions::new(metadata.to_vec())) + } + + fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option { + None + } + + fn partial_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option { + None + } + + fn empty_partial( + &self, + _options: &Self::Options, + _input_dtype: &DType, + ) -> VortexResult { + vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) + } + + fn combine_partials(&self, _partial: &mut Self::Partial, _other: Scalar) -> VortexResult<()> { + vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) + } + + fn to_scalar(&self, _partial: &Self::Partial) -> VortexResult { + vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) + } + + fn reset(&self, _partial: &mut Self::Partial) {} + + fn is_saturated(&self, _state: &Self::Partial) -> bool { + false + } + + fn accumulate( + &self, + _state: &mut Self::Partial, + _batch: &Columnar, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) + } + + fn finalize(&self, _states: ArrayRef) -> VortexResult { + vortex_bail!("Cannot execute unknown aggregate function '{}'", self.id) + } +} + +pub fn new_foreign_aggregate_fn(id: AggregateFnId, metadata: Vec) -> AggregateFnRef { + AggregateFn::new( + ForeignAggregateFnVTable::new(id), + ForeignAggregateFnOptions::new(metadata), + ) + .erased() +} diff --git a/vortex-array/src/aggregate_fn/mod.rs b/vortex-array/src/aggregate_fn/mod.rs index 0dad56f3222..a22006cb981 100644 --- a/vortex-array/src/aggregate_fn/mod.rs +++ b/vortex-array/src/aggregate_fn/mod.rs @@ -20,6 +20,9 @@ pub use vtable::*; mod plugin; pub use plugin::*; +mod foreign; +pub(crate) use foreign::*; + mod typed; pub use typed::*; diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 18231993eae..dfb86ad0ed9 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -12,6 +12,7 @@ use vortex_session::VortexSession; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::new_foreign_aggregate_fn; use crate::aggregate_fn::session::AggregateFnSessionExt; impl AggregateFnRef { @@ -38,12 +39,13 @@ impl AggregateFnRef { /// Note: the serialization format is not stable and may change between versions. pub fn from_proto(proto: &pb::AggregateFn, session: &VortexSession) -> VortexResult { let agg_fn_id: AggregateFnId = ArcRef::new_arc(Arc::from(proto.id.as_str())); - let plugin = session - .aggregate_fns() - .registry() - .find(&agg_fn_id) - .ok_or_else(|| vortex_err!("unknown aggregate function id: {}", proto.id))?; - let agg_fn = plugin.deserialize(proto.metadata(), session)?; + let agg_fn = if let Some(plugin) = session.aggregate_fns().registry().find(&agg_fn_id) { + plugin.deserialize(proto.metadata(), session)? + } else if session.allows_unknown() { + new_foreign_aggregate_fn(agg_fn_id.clone(), proto.metadata().to_vec()) + } else { + return Err(vortex_err!("unknown aggregate function id: {}", proto.id)); + }; if agg_fn.id() != agg_fn_id { vortex_bail!( @@ -164,4 +166,23 @@ mod tests { assert_eq!(deserialized, agg_fn); } + + #[test] + fn unknown_aggregate_fn_id_allow_unknown() { + let session = VortexSession::empty() + .with::() + .allow_unknown(); + + let proto = pb::AggregateFn { + id: "vortex.test.foreign_aggregate".to_string(), + metadata: Some(vec![7, 8, 9]), + }; + + let agg_fn = AggregateFnRef::from_proto(&proto, &session).unwrap(); + assert_eq!(agg_fn.id().as_ref(), "vortex.test.foreign_aggregate"); + + let roundtrip = agg_fn.serialize_proto().unwrap(); + assert_eq!(roundtrip.id, proto.id); + assert_eq!(roundtrip.metadata(), proto.metadata()); + } } diff --git a/vortex-array/src/array/foreign.rs b/vortex-array/src/array/foreign.rs new file mode 100644 index 00000000000..c5141ef30f9 --- /dev/null +++ b/vortex-array/src/array/foreign.rs @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; + +use crate::Array; +use crate::ArrayRef; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::array::vtable::NotSupported; +use crate::array::vtable::ValidityVTable; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::executor::ExecutionCtx; +use crate::hash::ArrayEq; +use crate::hash::ArrayHash; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +#[derive(Clone, Debug)] +pub struct ForeignArrayData { + metadata: Vec, + buffers: Vec, +} + +impl ForeignArrayData { + pub fn new(metadata: Vec, buffers: Vec) -> Self { + Self { metadata, buffers } + } +} + +impl ArrayHash for ForeignArrayData { + fn array_hash(&self, state: &mut H, precision: crate::Precision) { + self.metadata.hash(state); + self.buffers.len().hash(state); + for buffer in &self.buffers { + buffer.array_hash(state, precision); + } + } +} + +impl ArrayEq for ForeignArrayData { + fn array_eq(&self, other: &Self, precision: crate::Precision) -> bool { + self.metadata == other.metadata + && self.buffers.len() == other.buffers.len() + && self + .buffers + .iter() + .zip(other.buffers.iter()) + .all(|(lhs, rhs)| lhs.array_eq(rhs, precision)) + } +} + +#[derive(Clone, Debug)] +pub struct ForeignArray { + id: ArrayId, +} + +impl ForeignArray { + pub fn new(id: ArrayId) -> Self { + Self { id } + } +} + +pub struct ForeignValidityVTable; + +impl ValidityVTable for ForeignValidityVTable { + fn validity(array: ArrayView<'_, ForeignArray>) -> VortexResult { + Ok(Validity::from(array.dtype().nullability())) + } +} + +impl VTable for ForeignArray { + type ArrayData = ForeignArrayData; + type OperationsVTable = NotSupported; + type ValidityVTable = ForeignValidityVTable; + + fn id(&self) -> ArrayId { + self.id.clone() + } + + fn validate( + &self, + _data: &Self::ArrayData, + _dtype: &DType, + _len: usize, + _slots: &[Option], + ) -> VortexResult<()> { + Ok(()) + } + + fn nbuffers(array: ArrayView<'_, Self>) -> usize { + array.buffers.len() + } + + fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + array.buffers[idx].clone() + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + Some(format!("buffer[{idx}]")) + } + + fn serialize(array: ArrayView<'_, Self>) -> VortexResult>> { + Ok(Some(array.metadata.clone())) + } + + fn fmt_metadata(array: ArrayView<'_, Self>, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "ForeignMetadata({}B)", array.metadata.len()) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + let child_arrays = (0..children.len()) + .map(|idx| children.get(idx, dtype, len)) + .collect::>>()?; + + Ok(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + ForeignArrayData::new(metadata.to_vec(), buffers.to_vec()), + ) + .with_slots(child_arrays.into_iter().map(Some).collect())) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + format!("child[{idx}]") + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + vortex_bail!( + "Cannot execute unknown array encoding '{}'", + array.encoding_id() + ) + } +} + +pub fn new_foreign_array( + id: ArrayId, + dtype: DType, + len: usize, + metadata: Vec, + buffers: Vec, + children: Vec, +) -> VortexResult { + Ok(Array::::try_from_parts( + ArrayParts::new( + ForeignArray::new(id), + dtype, + len, + ForeignArrayData::new(metadata, buffers), + ) + .with_slots(children.into_iter().map(Some).collect()), + )? + .into_array()) +} diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index f6ed110e172..ee05b55edf3 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -34,6 +34,9 @@ pub use erased::*; mod plugin; pub use plugin::*; +mod foreign; +pub(crate) use foreign::*; + mod typed; pub use typed::*; diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 516a8dbe1ad..78ef1de9a2b 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::hash::Hasher; +use std::sync::Arc; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -146,7 +147,7 @@ impl VTable for Filter { return Ok(ExecutionResult::done(canonical)); } let mask_values = match &array.mask { - Mask::Values(v) => v.clone(), + Mask::Values(v) => Arc::clone(v), _ => unreachable!("`execute_filter_fast_paths` handles AllTrue and AllFalse"), }; diff --git a/vortex-array/src/dtype/extension/foreign.rs b/vortex-array/src/dtype/extension/foreign.rs new file mode 100644 index 00000000000..7753a5b6b86 --- /dev/null +++ b/vortex-array/src/dtype/extension/foreign.rs @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + +use vortex_error::VortexResult; + +use crate::dtype::DType; +use crate::dtype::extension::ExtDType; +use crate::dtype::extension::ExtDTypeRef; +use crate::dtype::extension::ExtId; +use crate::dtype::extension::ExtVTable; +use crate::scalar::ScalarValue; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct ForeignExtMetadata(pub Vec); + +impl Display for ForeignExtMetadata { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{}B", self.0.len()) + } +} + +/// Placeholder extension dtype used when deserializing an unknown extension ID. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct ForeignExtDType { + id: ExtId, +} + +impl ForeignExtDType { + pub fn new(id: ExtId) -> Self { + Self { id } + } + + pub fn from_parts( + id: ExtId, + metadata: Vec, + storage_dtype: DType, + ) -> VortexResult { + Ok( + ExtDType::try_with_vtable(Self::new(id), ForeignExtMetadata(metadata), storage_dtype)? + .erased(), + ) + } +} + +impl ExtVTable for ForeignExtDType { + type Metadata = ForeignExtMetadata; + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + self.id.clone() + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.0.clone()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(ForeignExtMetadata(metadata.to_vec())) + } + + fn validate_dtype(_ext_dtype: &ExtDType) -> VortexResult<()> { + Ok(()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult> { + Ok(storage_value) + } +} diff --git a/vortex-array/src/dtype/extension/mod.rs b/vortex-array/src/dtype/extension/mod.rs index c6ecc86bbf7..0eae1587b40 100644 --- a/vortex-array/src/dtype/extension/mod.rs +++ b/vortex-array/src/dtype/extension/mod.rs @@ -19,6 +19,9 @@ pub use vtable::*; mod plugin; pub use plugin::*; +mod foreign; +pub(crate) use foreign::*; + mod typed; pub use typed::*; diff --git a/vortex-array/src/dtype/serde/flatbuffers.rs b/vortex-array/src/dtype/serde/flatbuffers.rs index 75eb32e4ad3..af7ed7b8af3 100644 --- a/vortex-array/src/dtype/serde/flatbuffers.rs +++ b/vortex-array/src/dtype/serde/flatbuffers.rs @@ -24,6 +24,7 @@ use crate::dtype::FieldDType; use crate::dtype::PType; use crate::dtype::StructFields; use crate::dtype::extension::ExtId; +use crate::dtype::extension::ForeignExtDType; use crate::dtype::flatbuffers as fb; use crate::dtype::session::DTypeSessionExt; @@ -213,21 +214,19 @@ impl TryFrom for DType { let storage_dtype = DType::try_from(storage_view) .map_err(|e| vortex_err!("failed to create DType from fbs message: {e}"))?; - let vtable = vfdt - .session - .dtypes() - .registry() - .find(&id) - .ok_or_else(|| vortex_err!("No such DType extension ID: {}", id))?; - let ext_dtype = vtable.deserialize( - fb_ext - .metadata() - .ok_or_else(|| { - vortex_err!("failed to parse extension metadata from flatbuffer") - })? - .bytes(), - storage_dtype, - )?; + let metadata = fb_ext + .metadata() + .ok_or_else(|| { + vortex_err!("failed to parse extension metadata from flatbuffer") + })? + .bytes(); + let ext_dtype = if let Some(vtable) = vfdt.session.dtypes().registry().find(&id) { + vtable.deserialize(metadata, storage_dtype)? + } else if vfdt.session.allows_unknown() { + ForeignExtDType::from_parts(id, metadata.to_vec(), storage_dtype)? + } else { + return Err(vortex_err!("No such DType extension ID: {}", id)); + }; Ok(Self::Extension(ext_dtype)) } diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index 508eb42d2b4..4ca8642120b 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -13,6 +13,7 @@ use crate::dtype::DecimalDType; use crate::dtype::PType; use crate::dtype::StructFields; use crate::dtype::extension::ExtId; +use crate::dtype::extension::ForeignExtDType; use crate::dtype::field::Field; use crate::dtype::field::FieldPath; use crate::dtype::proto::dtype as pb; @@ -87,16 +88,19 @@ impl DType { } DtypeType::Extension(e) => { let id = ExtId::new_arc(e.id.as_str().to_string().into()); - let vtable = session.dtypes().registry().find(&id).ok_or_else( - || vortex_err!(Serde: "Unregistered extension type ID: {}", e.id), - )?; let storage_dtype = DType::from_proto( e.storage_dtype .as_ref() .ok_or_else(|| vortex_err!("Extension DType missing storage proto"))?, session, )?; - let ext_dtype = vtable.deserialize(e.metadata(), storage_dtype)?; + let ext_dtype = if let Some(vtable) = session.dtypes().registry().find(&id) { + vtable.deserialize(e.metadata(), storage_dtype)? + } else if session.allows_unknown() { + ForeignExtDType::from_parts(id, e.metadata().to_vec(), storage_dtype)? + } else { + return Err(vortex_err!(Serde: "Unregistered extension type ID: {}", e.id)); + }; Ok(Self::Extension(ext_dtype)) } DtypeType::Variant(v) => Ok(Self::Variant(v.nullable.into())), @@ -221,6 +225,8 @@ impl TryFrom<&pb::FieldPath> for FieldPath { mod tests { use std::sync::Arc; + use vortex_session::VortexSession; + use super::*; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -490,4 +496,35 @@ mod tests { .contains("Extension DType missing storage proto") ); } + + #[test] + fn test_unknown_extension_allow_unknown() { + let session = VortexSession::empty().allow_unknown(); + let proto = pb::DType { + dtype_type: Some(DtypeType::Extension(Box::new(pb::Extension { + id: "vortex.test.foreign_ext".to_string(), + storage_dtype: Some(Box::new(pb::DType { + dtype_type: Some(DtypeType::Primitive(pb::Primitive { + r#type: pb::PType::I32.into(), + nullable: false, + })), + })), + metadata: Some(vec![1, 2, 3]), + }))), + }; + + let dtype = DType::from_proto(&proto, &session).unwrap(); + let DType::Extension(ext) = &dtype else { + panic!("Expected extension dtype"); + }; + assert_eq!(ext.id().as_ref(), "vortex.test.foreign_ext"); + assert_eq!(ext.serialize_metadata().unwrap(), vec![1, 2, 3]); + + let roundtrip = pb::DType::try_from(&dtype).unwrap(); + let DtypeType::Extension(roundtrip_ext) = roundtrip.dtype_type.unwrap() else { + panic!("Expected extension dtype"); + }; + assert_eq!(roundtrip_ext.id, "vortex.test.foreign_ext"); + assert_eq!(roundtrip_ext.metadata(), &[1, 2, 3]); + } } diff --git a/vortex-array/src/dtype/serde/serde.rs b/vortex-array/src/dtype/serde/serde.rs index 63044385548..e89a1a98658 100644 --- a/vortex-array/src/dtype/serde/serde.rs +++ b/vortex-array/src/dtype/serde/serde.rs @@ -30,6 +30,7 @@ use crate::dtype::StructFields; use crate::dtype::decimal::DecimalDType; use crate::dtype::extension::ExtDTypeRef; use crate::dtype::extension::ExtId; +use crate::dtype::extension::ForeignExtDType; use crate::dtype::session::DTypeSessionExt; /// Serialize Nullability as a boolean @@ -571,20 +572,30 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, ExtDTypeRef> { let id = id.ok_or_else(|| de::Error::missing_field("id"))?; let id = ExtId::new_arc(id); - let vtable = self.session.dtypes().registry().find(&id).ok_or_else(|| { - de::Error::custom(format!("unknown extension dtype id: {}", id)) - })?; - let storage_dtype = storage_dtype.ok_or_else(|| de::Error::missing_field("storage_dtype"))?; let metadata = metadata.ok_or_else(|| de::Error::missing_field("metadata"))?; - vtable.deserialize(&metadata, storage_dtype).map_err(|e| { - de::Error::custom(format!( - "failed to deserialize extension dtype {}: {}", - id, e - )) - }) + if let Some(vtable) = self.session.dtypes().registry().find(&id) { + vtable.deserialize(&metadata, storage_dtype).map_err(|e| { + de::Error::custom(format!( + "failed to deserialize extension dtype {}: {}", + id, e + )) + }) + } else if self.session.allows_unknown() { + ForeignExtDType::from_parts(id, metadata, storage_dtype).map_err(|e| { + de::Error::custom(format!( + "failed to deserialize unknown extension dtype: {}", + e + )) + }) + } else { + Err(de::Error::custom(format!( + "unknown extension dtype id: {}", + id + ))) + } } } diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 16b063ac831..22c5aa8b467 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -10,6 +10,7 @@ use vortex_proto::expr as pb; use vortex_session::VortexSession; use crate::expr::Expression; +use crate::scalar_fn::ForeignScalarFnVTable; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::session::ScalarFnSessionExt; @@ -41,19 +42,21 @@ impl ExprSerializeProtoExt for Expression { impl Expression { pub fn from_proto(expr: &pb::Expr, session: &VortexSession) -> VortexResult { let expr_id = ScalarFnId::new_arc(Arc::from(expr.id.to_string())); - let vtable = session - .scalar_fns() - .registry() - .find(&expr_id) - .ok_or_else(|| vortex_err!("unknown expression id: {}", expr_id))?; - let children = expr .children .iter() .map(|e| Expression::from_proto(e, session)) .collect::>>()?; - Expression::try_new(vtable.deserialize(expr.metadata(), session)?, children) + let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().find(&expr_id) { + vtable.deserialize(expr.metadata(), session)? + } else if session.allows_unknown() { + ForeignScalarFnVTable::make_scalar_fn(expr_id, expr.metadata().to_vec(), children.len()) + } else { + return Err(vortex_err!("unknown expression id: {}", expr_id)); + }; + + Expression::try_new(scalar_fn, children) } } @@ -70,6 +73,7 @@ pub fn deserialize_expr_proto( mod tests { use prost::Message; use vortex_proto::expr as pb; + use vortex_session::VortexSession; use super::ExprSerializeProtoExt; use crate::LEGACY_SESSION; @@ -83,6 +87,7 @@ mod tests { use crate::expr::root; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; + use crate::scalar_fn::session::ScalarFnSession; #[test] fn expression_serde() { @@ -109,4 +114,25 @@ mod tests { assert_eq!(&deser_expr, &expr); } + + #[test] + fn unknown_expression_id_allow_unknown() { + let session = VortexSession::empty() + .with::() + .allow_unknown(); + + let expr_proto = pb::Expr { + id: "vortex.test.foreign_scalar_fn".to_string(), + metadata: Some(vec![1, 2, 3, 4]), + children: vec![root().serialize_proto().unwrap()], + }; + + let expr = Expression::from_proto(&expr_proto, &session).unwrap(); + assert_eq!(expr.id().as_ref(), "vortex.test.foreign_scalar_fn"); + + let roundtrip = expr.serialize_proto().unwrap(); + assert_eq!(roundtrip.id, expr_proto.id); + assert_eq!(roundtrip.metadata(), expr_proto.metadata()); + assert_eq!(roundtrip.children.len(), 1); + } } diff --git a/vortex-array/src/scalar_fn/foreign.rs b/vortex-array/src/scalar_fn/foreign.rs new file mode 100644 index 00000000000..e7d3ee641cc --- /dev/null +++ b/vortex-array/src/scalar_fn/foreign.rs @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::expr::Expression; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFn; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::ScalarFnVTable; + +/// Options payload for a foreign scalar function. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct ForeignScalarFnOptions { + metadata: Vec, + arity: usize, +} + +impl ForeignScalarFnOptions { + pub fn new(metadata: Vec, arity: usize) -> Self { + Self { metadata, arity } + } +} + +impl Display for ForeignScalarFnOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!( + f, + "foreign(arity={}, metadata={}B)", + self.arity, + self.metadata.len() + ) + } +} + +/// Scalar function placeholder used when deserializing an unknown scalar function ID. +#[derive(Clone, Debug)] +pub struct ForeignScalarFnVTable { + id: ScalarFnId, +} + +impl ForeignScalarFnVTable { + pub fn new(id: ScalarFnId) -> Self { + Self { id } + } + + pub fn make_scalar_fn(id: ScalarFnId, metadata: Vec, arity: usize) -> ScalarFnRef { + ScalarFn::new(Self::new(id), ForeignScalarFnOptions::new(metadata, arity)).erased() + } +} + +impl ScalarFnVTable for ForeignScalarFnVTable { + type Options = ForeignScalarFnOptions; + + fn id(&self) -> ScalarFnId { + self.id.clone() + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.metadata.clone())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(ForeignScalarFnOptions::new(metadata.to_vec(), 0)) + } + + fn arity(&self, options: &Self::Options) -> Arity { + Arity::Exact(options.arity) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + ChildName::new_arc(format!("arg{child_idx}").into()) + } + + fn fmt_sql( + &self, + _options: &Self::Options, + expr: &Expression, + f: &mut Formatter<'_>, + ) -> fmt::Result { + write!(f, "{}(", self.id)?; + for i in 0..expr.children().len() { + if i > 0 { + write!(f, ", ")?; + } + expr.child(i).fmt_sql(f)?; + } + write!(f, ")") + } + + fn return_dtype(&self, _options: &Self::Options, _args: &[DType]) -> VortexResult { + Ok(DType::Variant(Nullability::Nullable)) + } + + fn execute( + &self, + _options: &Self::Options, + _args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + vortex_bail!("Cannot execute unknown scalar function '{}'", self.id); + } +} diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 945c37262de..39e45f00e08 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -15,6 +15,9 @@ pub use vtable::*; mod plugin; pub use plugin::*; +mod foreign; +pub use foreign::*; + mod typed; pub use typed::*; diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index 1dda5d01a48..4f8f9eed278 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -29,6 +29,7 @@ use vortex_utils::aliases::hash_map::HashMap; use crate::ArrayContext; use crate::ArrayRef; +use crate::array::new_foreign_array; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::TryFromBytes; @@ -323,11 +324,12 @@ impl SerializedArray { let encoding_id = ctx .resolve(encoding_idx) .ok_or_else(|| vortex_err!("Unknown encoding index: {}", encoding_idx))?; - let plugin = session - .arrays() - .registry() - .find(&encoding_id) - .ok_or_else(|| vortex_err!("Unknown encoding: {}", encoding_id))?; + let Some(plugin) = session.arrays().registry().find(&encoding_id) else { + if session.allows_unknown() { + return self.decode_foreign(encoding_id, dtype, len, ctx); + } + return Err(vortex_err!("Unknown encoding: {}", encoding_id)); + }; let children = SerializedArrayChildren { ser: self, @@ -374,6 +376,34 @@ impl SerializedArray { Ok(decoded) } + fn decode_foreign( + &self, + encoding_id: crate::array::ArrayId, + dtype: &DType, + len: usize, + ctx: &ReadContext, + ) -> VortexResult { + let children = (0..self.nchildren()) + .map(|idx| { + let child = self.child(idx); + let child_encoding_idx = child.flatbuffer().encoding(); + let child_encoding_id = ctx + .resolve(child_encoding_idx) + .ok_or_else(|| vortex_err!("Unknown encoding index: {}", child_encoding_idx))?; + child.decode_foreign(child_encoding_id, dtype, len, ctx) + }) + .collect::>>()?; + + new_foreign_array( + encoding_id, + dtype.clone(), + len, + self.metadata().to_vec(), + self.collect_buffers()?.into_owned(), + children, + ) + } + /// Returns the array encoding. pub fn encoding_id(&self) -> u16 { self.flatbuffer().encoding() @@ -662,3 +692,88 @@ impl TryFrom for SerializedArray { Self::try_from(value.try_to_host_sync()?) } } + +#[cfg(test)] +mod tests { + use flatbuffers::FlatBufferBuilder; + use vortex_session::VortexSession; + use vortex_session::registry::ReadContext; + + use super::SerializeOptions; + use super::SerializedArray; + use crate::ArrayContext; + use crate::array::ArrayId; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::flatbuffers as fba; + use crate::session::ArraySession; + + #[test] + fn unknown_array_encoding_allow_unknown() { + let mut fbb = FlatBufferBuilder::new(); + + let child_metadata = fbb.create_vector(&[9u8]); + let child = fba::ArrayNode::create( + &mut fbb, + &fba::ArrayNodeArgs { + encoding: 1, + metadata: Some(child_metadata), + children: None, + buffers: None, + stats: None, + }, + ); + + let children = fbb.create_vector(&[child]); + let metadata = fbb.create_vector(&[1u8, 2, 3]); + let root = fba::ArrayNode::create( + &mut fbb, + &fba::ArrayNodeArgs { + encoding: 0, + metadata: Some(metadata), + children: Some(children), + buffers: None, + stats: None, + }, + ); + let array = fba::Array::create( + &mut fbb, + &fba::ArrayArgs { + root: Some(root), + buffers: None, + }, + ); + fbb.finish_minimal(array); + let (buf, start) = fbb.collapse(); + let tree = vortex_buffer::ByteBuffer::from(buf).slice(start..); + + let ser = SerializedArray::from_array_tree(tree).unwrap(); + let ctx = ReadContext::new([ + ArrayId::new_ref("vortex.test.foreign_array"), + ArrayId::new_ref("vortex.test.foreign_child"), + ]); + let session = VortexSession::empty() + .with::() + .allow_unknown(); + + let decoded = ser + .decode(&DType::Variant(Nullability::Nullable), 5, &ctx, &session) + .unwrap(); + assert_eq!(decoded.encoding_id().as_ref(), "vortex.test.foreign_array"); + assert_eq!(decoded.nchildren(), 1); + assert_eq!( + decoded.nth_child(0).unwrap().encoding_id().as_ref(), + "vortex.test.foreign_child" + ); + assert_eq!(decoded.metadata().unwrap().unwrap(), vec![1, 2, 3]); + assert_eq!( + decoded.nth_child(0).unwrap().metadata().unwrap().unwrap(), + vec![9] + ); + + let serialized = decoded + .serialize(&ArrayContext::default(), &SerializeOptions::default()) + .unwrap(); + assert!(!serialized.is_empty()); + } +} diff --git a/vortex-file/src/footer/mod.rs b/vortex-file/src/footer/mod.rs index 056c8337ad7..e0d08568108 100644 --- a/vortex-file/src/footer/mod.rs +++ b/vortex-file/src/footer/mod.rs @@ -34,7 +34,7 @@ use vortex_flatbuffers::FlatBuffer; use vortex_flatbuffers::footer as fb; use vortex_layout::LayoutEncodingId; use vortex_layout::LayoutRef; -use vortex_layout::layout_from_flatbuffer; +use vortex_layout::layout_from_flatbuffer_with_options; use vortex_layout::session::LayoutSessionExt; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; @@ -101,12 +101,13 @@ impl Footer { .collect(); let array_read_ctx = ReadContext::new(array_ids); - let root_layout = layout_from_flatbuffer( + let root_layout = layout_from_flatbuffer_with_options( layout_bytes, &dtype, &layout_read_ctx, &array_read_ctx, session.layouts().registry(), + session.allows_unknown(), )?; let segments: Arc<[SegmentSpec]> = fb_footer diff --git a/vortex-layout/public-api.lock b/vortex-layout/public-api.lock index cc11334cb0c..657687c7de4 100644 --- a/vortex-layout/public-api.lock +++ b/vortex-layout/public-api.lock @@ -2124,6 +2124,8 @@ pub fn vortex_layout::layouts::zoned::Zoned::with_children(layout: &mut Self::La pub fn vortex_layout::layout_from_flatbuffer(flatbuffer: vortex_flatbuffers::FlatBuffer, dtype: &vortex_array::dtype::DType, layout_ctx: &vortex_session::registry::ReadContext, ctx: &vortex_session::registry::ReadContext, layouts: &vortex_layout::session::LayoutRegistry) -> vortex_error::VortexResult +pub fn vortex_layout::layout_from_flatbuffer_with_options(flatbuffer: vortex_flatbuffers::FlatBuffer, dtype: &vortex_array::dtype::DType, layout_ctx: &vortex_session::registry::ReadContext, ctx: &vortex_session::registry::ReadContext, layouts: &vortex_layout::session::LayoutRegistry, allow_unknown: bool) -> vortex_error::VortexResult + pub type vortex_layout::ArrayFuture = futures_core::future::BoxFuture<'static, vortex_error::VortexResult> pub type vortex_layout::LayoutContext = vortex_session::registry::Context diff --git a/vortex-layout/src/children.rs b/vortex-layout/src/children.rs index 2090aede5ee..916e216543e 100644 --- a/vortex-layout/src/children.rs +++ b/vortex-layout/src/children.rs @@ -103,6 +103,7 @@ pub(crate) struct ViewedLayoutChildren { array_read_ctx: ReadContext, layout_read_ctx: ReadContext, layouts: LayoutRegistry, + allow_unknown: bool, } impl ViewedLayoutChildren { @@ -117,6 +118,7 @@ impl ViewedLayoutChildren { array_read_ctx: ReadContext, layout_read_ctx: ReadContext, layouts: LayoutRegistry, + allow_unknown: bool, ) -> Self { Self { flatbuffer, @@ -124,6 +126,7 @@ impl ViewedLayoutChildren { array_read_ctx, layout_read_ctx, layouts, + allow_unknown, } } @@ -134,6 +137,41 @@ impl ViewedLayoutChildren { // The lifetime of the returned Layout is tied to self, ensuring the buffer remains valid. unsafe { fbl::Layout::follow(self.flatbuffer.as_ref(), self.flatbuffer_loc) } } + + fn foreign_layout_from_fb( + &self, + fb_layout: fbl::Layout<'_>, + dtype: &DType, + ) -> VortexResult { + let encoding_id = self + .layout_read_ctx + .resolve(fb_layout.encoding()) + .ok_or_else(|| vortex_err!("Encoding not found: {}", fb_layout.encoding()))?; + + let children = fb_layout + .children() + .unwrap_or_default() + .iter() + .map(|child| self.foreign_layout_from_fb(child, dtype)) + .collect::>>()?; + + Ok(crate::layouts::foreign::new_foreign_layout( + encoding_id, + dtype.clone(), + fb_layout.row_count(), + fb_layout + .metadata() + .map(|m| m.bytes().to_vec()) + .unwrap_or_default(), + fb_layout + .segments() + .unwrap_or_default() + .iter() + .map(SegmentId::from) + .collect_vec(), + children, + )) + } } impl LayoutChildren for ViewedLayoutChildren { @@ -153,15 +191,22 @@ impl LayoutChildren for ViewedLayoutChildren { array_read_ctx: self.array_read_ctx.clone(), layout_read_ctx: self.layout_read_ctx.clone(), layouts: self.layouts.clone(), + allow_unknown: self.allow_unknown, }; let encoding_id = self .layout_read_ctx .resolve(fb_child.encoding()) .ok_or_else(|| vortex_err!("Encoding not found: {}", fb_child.encoding()))?; - let encoding = self.layouts.find(&encoding_id).ok_or_else(|| { - vortex_err!("Encoding not found in registry: {}", fb_child.encoding()) - })?; + let Some(encoding) = self.layouts.find(&encoding_id) else { + if self.allow_unknown { + return viewed_children.foreign_layout_from_fb(fb_child, dtype); + } + return Err(vortex_err!( + "Encoding not found in registry: {}", + fb_child.encoding() + )); + }; encoding.build( dtype, diff --git a/vortex-layout/src/encoding.rs b/vortex-layout/src/encoding.rs index ce0086e2ffc..271a94f9059 100644 --- a/vortex-layout/src/encoding.rs +++ b/vortex-layout/src/encoding.rs @@ -128,8 +128,10 @@ impl dyn LayoutEncoding + '_ { mod private { use super::*; + use crate::layouts::foreign::ForeignLayoutEncoding; pub trait Sealed {} impl Sealed for LayoutEncodingAdapter {} + impl Sealed for ForeignLayoutEncoding {} } diff --git a/vortex-layout/src/flatbuffers.rs b/vortex-layout/src/flatbuffers.rs index f180ab9753c..02e05b199c8 100644 --- a/vortex-layout/src/flatbuffers.rs +++ b/vortex-layout/src/flatbuffers.rs @@ -48,14 +48,30 @@ pub fn layout_from_flatbuffer( layout_ctx: &ReadContext, ctx: &ReadContext, layouts: &LayoutRegistry, +) -> VortexResult { + layout_from_flatbuffer_with_options(flatbuffer, dtype, layout_ctx, ctx, layouts, false) +} + +/// Parse a [`LayoutRef`] from a layout flatbuffer with unknown-encoding behavior control. +pub fn layout_from_flatbuffer_with_options( + flatbuffer: FlatBuffer, + dtype: &DType, + layout_ctx: &ReadContext, + ctx: &ReadContext, + layouts: &LayoutRegistry, + allow_unknown: bool, ) -> VortexResult { let fb_layout = root_with_opts::(&LAYOUT_VERIFIER, &flatbuffer)?; let encoding_id = layout_ctx .resolve(fb_layout.encoding()) .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?; - let encoding = layouts - .find(&encoding_id) - .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?; + let encoding = layouts.find(&encoding_id); + + if encoding.is_none() && allow_unknown { + return foreign_layout_from_fb(fb_layout, dtype, layout_ctx); + } + let encoding = + encoding.ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?; // SAFETY: we validate the flatbuffer above in the `root` call, and extract a loc. let viewed_children = unsafe { @@ -65,6 +81,7 @@ pub fn layout_from_flatbuffer( ctx.clone(), layout_ctx.clone(), layouts.clone(), + allow_unknown, ) }; @@ -88,6 +105,40 @@ pub fn layout_from_flatbuffer( Ok(layout) } +fn foreign_layout_from_fb( + fb_layout: layout::Layout<'_>, + dtype: &DType, + layout_ctx: &ReadContext, +) -> VortexResult { + let encoding_id = layout_ctx + .resolve(fb_layout.encoding()) + .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?; + + let children = fb_layout + .children() + .unwrap_or_default() + .iter() + .map(|child| foreign_layout_from_fb(child, dtype, layout_ctx)) + .collect::>>()?; + + Ok(crate::layouts::foreign::new_foreign_layout( + encoding_id, + dtype.clone(), + fb_layout.row_count(), + fb_layout + .metadata() + .map(|m| m.bytes().to_vec()) + .unwrap_or_default(), + fb_layout + .segments() + .unwrap_or_default() + .iter() + .map(SegmentId::from) + .collect(), + children, + )) +} + impl dyn Layout + '_ { /// Serialize the layout into a [`FlatBufferBuilder`]. pub fn flatbuffer_writer<'a>( @@ -159,3 +210,83 @@ impl WriteFlatBuffer for LayoutFlatBufferWriter<'_> { )) } } + +#[cfg(test)] +mod tests { + use flatbuffers::FlatBufferBuilder; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_flatbuffers::layout as fbl; + use vortex_session::registry::ReadContext; + + use super::layout_from_flatbuffer_with_options; + use crate::LayoutEncodingId; + use crate::session::LayoutSession; + + #[test] + fn unknown_layout_encoding_allow_unknown() { + let mut fbb = FlatBufferBuilder::new(); + + let child_metadata = fbb.create_vector(&[9u8]); + let child = fbl::Layout::create( + &mut fbb, + &fbl::LayoutArgs { + encoding: 1, + row_count: 3, + metadata: Some(child_metadata), + children: None, + segments: None, + }, + ); + + let children = fbb.create_vector(&[child]); + let metadata = fbb.create_vector(&[1u8, 2, 3]); + let segments = fbb.create_vector(&[7u32]); + let root = fbl::Layout::create( + &mut fbb, + &fbl::LayoutArgs { + encoding: 0, + row_count: 10, + metadata: Some(metadata), + children: Some(children), + segments: Some(segments), + }, + ); + fbb.finish_minimal(root); + let (buf, start) = fbb.collapse(); + let layout_buffer = vortex_flatbuffers::FlatBuffer::align_from( + vortex_buffer::ByteBuffer::from(buf).slice(start..), + ); + + let layout_ctx = ReadContext::new([ + LayoutEncodingId::new_ref("vortex.test.foreign_layout"), + LayoutEncodingId::new_ref("vortex.test.foreign_child_layout"), + ]); + let array_ctx = ReadContext::new([]); + let layouts = LayoutSession::default().registry().clone(); + + let layout = layout_from_flatbuffer_with_options( + layout_buffer, + &DType::Variant(Nullability::Nullable), + &layout_ctx, + &array_ctx, + &layouts, + true, + ) + .unwrap(); + + assert_eq!(layout.encoding_id().as_ref(), "vortex.test.foreign_layout"); + assert_eq!(layout.row_count(), 10); + assert_eq!(layout.metadata(), vec![1, 2, 3]); + assert_eq!(layout.segment_ids().len(), 1); + assert_eq!(*layout.segment_ids()[0], 7); + assert_eq!(layout.nchildren(), 1); + + let child = layout.child(0).unwrap(); + assert_eq!( + child.encoding_id().as_ref(), + "vortex.test.foreign_child_layout" + ); + assert_eq!(child.metadata(), vec![9]); + } +} diff --git a/vortex-layout/src/layout.rs b/vortex-layout/src/layout.rs index 473d754df83..0d805623c0f 100644 --- a/vortex-layout/src/layout.rs +++ b/vortex-layout/src/layout.rs @@ -317,10 +317,12 @@ impl Layout for LayoutAdapter { mod private { use super::*; + use crate::layouts::foreign::ForeignLayout; pub trait Sealed {} impl Sealed for LayoutAdapter {} + impl Sealed for ForeignLayout {} } #[cfg(test)] diff --git a/vortex-layout/src/layouts/foreign/mod.rs b/vortex-layout/src/layouts/foreign/mod.rs new file mode 100644 index 00000000000..358e2a42be1 --- /dev/null +++ b/vortex-layout/src/layouts/foreign/mod.rs @@ -0,0 +1,180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; +use std::sync::Arc; + +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use crate::Layout; +use crate::LayoutChildType; +use crate::LayoutChildren; +use crate::LayoutEncoding; +use crate::LayoutEncodingId; +use crate::LayoutEncodingRef; +use crate::LayoutReaderRef; +use crate::LayoutRef; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +/// Placeholder layout encoding used when deserializing an unknown layout encoding ID. +#[derive(Clone, Debug)] +pub(crate) struct ForeignLayoutEncoding { + id: LayoutEncodingId, +} + +impl ForeignLayoutEncoding { + pub(crate) fn new(id: LayoutEncodingId) -> Self { + Self { id } + } +} + +impl LayoutEncoding for ForeignLayoutEncoding { + fn as_any(&self) -> &dyn Any { + self + } + + fn id(&self) -> LayoutEncodingId { + self.id.clone() + } + + fn build( + &self, + dtype: &DType, + row_count: u64, + metadata: &[u8], + segment_ids: Vec, + children: &dyn LayoutChildren, + _ctx: &ReadContext, + ) -> VortexResult { + let child_layouts = (0..children.nchildren()) + .map(|idx| children.child(idx, dtype)) + .collect::>>()?; + + Ok(new_foreign_layout( + self.id.clone(), + dtype.clone(), + row_count, + metadata.to_vec(), + segment_ids, + child_layouts, + )) + } +} + +/// Placeholder layout used when deserializing an unknown layout encoding ID. +#[derive(Clone, Debug)] +pub(crate) struct ForeignLayout { + encoding: LayoutEncodingRef, + dtype: DType, + row_count: u64, + metadata: Vec, + segment_ids: Vec, + children: Vec, +} + +impl ForeignLayout { + pub(crate) fn new( + encoding_id: LayoutEncodingId, + dtype: DType, + row_count: u64, + metadata: Vec, + segment_ids: Vec, + children: Vec, + ) -> Self { + let encoding = + LayoutEncodingRef::new_arc(Arc::new(ForeignLayoutEncoding::new(encoding_id))); + + Self { + encoding, + dtype, + row_count, + metadata, + segment_ids, + children, + } + } +} + +pub(crate) fn new_foreign_layout( + encoding_id: LayoutEncodingId, + dtype: DType, + row_count: u64, + metadata: Vec, + segment_ids: Vec, + children: Vec, +) -> LayoutRef { + Arc::new(ForeignLayout::new( + encoding_id, + dtype, + row_count, + metadata, + segment_ids, + children, + )) +} + +impl Layout for ForeignLayout { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn to_layout(&self) -> LayoutRef { + Arc::new(self.clone()) + } + + fn encoding(&self) -> LayoutEncodingRef { + self.encoding.clone() + } + + fn row_count(&self) -> u64 { + self.row_count + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn nchildren(&self) -> usize { + self.children.len() + } + + fn child(&self, idx: usize) -> VortexResult { + self.children.get(idx).cloned().ok_or_else(|| { + vortex_err!("Child index out of bounds: {} of {}", idx, self.nchildren()) + }) + } + + fn child_type(&self, idx: usize) -> LayoutChildType { + LayoutChildType::Auxiliary(format!("[{idx}]").into()) + } + + fn metadata(&self) -> Vec { + self.metadata.clone() + } + + fn segment_ids(&self) -> Vec { + self.segment_ids.clone() + } + + fn new_reader( + &self, + _name: Arc, + _segment_source: Arc, + _session: &VortexSession, + ) -> VortexResult { + vortex_bail!( + "Cannot read unknown layout encoding '{}'", + self.encoding.id() + ) + } +} diff --git a/vortex-layout/src/layouts/mod.rs b/vortex-layout/src/layouts/mod.rs index 03b29a97a8f..18df5b8f347 100644 --- a/vortex-layout/src/layouts/mod.rs +++ b/vortex-layout/src/layouts/mod.rs @@ -15,6 +15,7 @@ pub mod compressed; pub mod dict; pub mod file_stats; pub mod flat; +pub(crate) mod foreign; pub(crate) mod partitioned; pub mod repartition; pub mod row_idx; diff --git a/vortex-session/public-api.lock b/vortex-session/public-api.lock index ab635ab69e2..55413e768ef 100644 --- a/vortex-session/public-api.lock +++ b/vortex-session/public-api.lock @@ -110,6 +110,10 @@ pub struct vortex_session::VortexSession(_) impl vortex_session::VortexSession +pub fn vortex_session::VortexSession::allow_unknown(self) -> Self + +pub fn vortex_session::VortexSession::allows_unknown(&self) -> bool + pub fn vortex_session::VortexSession::empty() -> Self pub fn vortex_session::VortexSession::with(self) -> Self diff --git a/vortex-session/src/lib.rs b/vortex-session/src/lib.rs index 623192899e7..205acc9bef9 100644 --- a/vortex-session/src/lib.rs +++ b/vortex-session/src/lib.rs @@ -61,6 +61,24 @@ impl VortexSession { } self } + + /// Allow deserializing unknown plugin IDs as non-executable foreign placeholders. + pub fn allow_unknown(self) -> Self { + let mut policy = ::get_mut::(&self); + policy.allow_unknown = true; + drop(policy); + self + } + + /// Returns whether unknown plugins should deserialize as foreign placeholders. + pub fn allows_unknown(&self) -> bool { + ::get::(self).allow_unknown + } +} + +#[derive(Debug, Clone, Copy, Default)] +struct UnknownPluginPolicy { + allow_unknown: bool, } /// Trait for accessing and modifying the state of a Vortex session. @@ -245,3 +263,17 @@ impl<'a, T> RefMut<'a, T> { RefMut(self.0.map(f)) } } + +#[cfg(test)] +mod tests { + use super::VortexSession; + + #[test] + fn allow_unknown_flag_is_opt_in() { + let session = VortexSession::empty(); + assert!(!session.allows_unknown()); + + let session = session.allow_unknown(); + assert!(session.allows_unknown()); + } +} diff --git a/vortex-tui/src/main.rs b/vortex-tui/src/main.rs index d813ff4dc4f..af872073dd5 100644 --- a/vortex-tui/src/main.rs +++ b/vortex-tui/src/main.rs @@ -8,6 +8,6 @@ use vortex_tui::launch; #[tokio::main] async fn main() -> anyhow::Result<()> { - let session = VortexSession::default().with_tokio(); + let session = VortexSession::default().with_tokio().allow_unknown(); launch(&session).await } diff --git a/vortex-web/crate/src/lib.rs b/vortex-web/crate/src/lib.rs index 69ef6284c2d..03b36ee4515 100644 --- a/vortex-web/crate/src/lib.rs +++ b/vortex-web/crate/src/lib.rs @@ -7,4 +7,17 @@ //! //! Built with `wasm-pack build --target web` and consumed by the vortex-web frontend. +use std::sync::LazyLock; + +use vortex::VortexSessionDefault; +use vortex::io::runtime::wasm::WasmRuntime; +use vortex::io::session::RuntimeSessionExt; +use vortex::session::VortexSession; + mod wasm; + +static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::default() + .with_handle(WasmRuntime::handle()) + .allow_unknown() +}); diff --git a/vortex-web/crate/src/wasm.rs b/vortex-web/crate/src/wasm.rs index 1139907f47f..5463d6f37b4 100644 --- a/vortex-web/crate/src/wasm.rs +++ b/vortex-web/crate/src/wasm.rs @@ -18,7 +18,6 @@ use futures::FutureExt; use futures::TryStreamExt; use futures::future::BoxFuture; use serde::Serialize; -use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::LEGACY_SESSION; use vortex::array::VortexSessionExecute; @@ -35,8 +34,6 @@ use vortex::file::VERSION; use vortex::file::VortexFile; use vortex::io::CoalesceConfig; use vortex::io::VortexReadAt; -use vortex::io::runtime::wasm::WasmRuntime; -use vortex::io::session::RuntimeSessionExt; use vortex::layout::LayoutChildType; use vortex::layout::LayoutRef; use vortex::layout::layouts::flat::Flat; @@ -46,6 +43,8 @@ use vortex::session::registry::ReadContext; use wasm_bindgen::prelude::*; use wasm_bindgen_futures::JsFuture; +use crate::SESSION; + /// Initialize the WASM module (sets up panic hook for better error messages). #[wasm_bindgen(start)] pub fn init() { @@ -124,7 +123,6 @@ impl VortexReadAt for BlobReadAt { /// The `File` (a `Blob`) is read lazily — only the footer is read at open time. #[wasm_bindgen] pub async fn open_vortex_file(file: web_sys::File) -> Result { - let session = VortexSession::default().with_handle(WasmRuntime::handle()); let blob: &web_sys::Blob = file.as_ref(); let file_size = blob.size() as usize; let reader = Arc::new(BlobReadAt { @@ -132,7 +130,7 @@ pub async fn open_vortex_file(file: web_sys::File) -> Result Result