diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index ce6d869acd8..d4e5d9350be 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -54,10 +54,10 @@ use crate::children::LayoutChildren; use crate::children::OwnedLayoutChildren; use crate::layouts::zoned::reader::ZonedReader; use crate::layouts::zoned::schema::AggregateSpecProto; -use crate::layouts::zoned::schema::aggregate_fns_from_specs; use crate::layouts::zoned::schema::aggregate_specs_from_fns; use crate::layouts::zoned::schema::aggregate_stats_table_dtype; use crate::layouts::zoned::schema::legacy_stats_table_dtype; +use crate::layouts::zoned::schema::try_aggregate_fns_from_specs; use crate::segments::SegmentId; use crate::segments::SegmentSource; use crate::vtable; @@ -157,9 +157,25 @@ impl VTable for Zoned { 2, "ZonedLayout expects exactly 2 children (data, zones)" ); - let aggregate_fns = aggregate_fns_from_specs(&metadata.aggregate_specs, build_ctx.session)?; + + // If any stored aggregate is unknown to this session (and unknown plugins are allowed), + // disable zoned pruning and fall back to a plain scan of the data child. + let Some(aggregate_fns) = + try_aggregate_fns_from_specs(&metadata.aggregate_specs, build_ctx.session)? + else { + return Ok(ZonedLayout { + dtype: dtype.clone(), + children: children.to_arc(), + zone_len: 0, + zone_map_schema: ZoneMapSchema::AggregateFns(Arc::new([])), + stats_table_dtype: aggregate_stats_table_dtype(dtype, &[]), + }); + }; + + // Verify that every aggregate is serializable. aggregate_specs_from_fns(&aggregate_fns)?; let stats_table_dtype = aggregate_stats_table_dtype(dtype, &aggregate_fns); + Ok(ZonedLayout { dtype: dtype.clone(), children: children.to_arc(), @@ -330,6 +346,8 @@ impl ZonedLayout { if zones.dtype() != &expected_dtype { vortex_bail!("Invalid zone map layout: zones dtype does not match expected dtype"); } + + // Verify that every aggregate is serializable. aggregate_specs_from_fns(&aggregate_fns)?; Ok(Self { @@ -551,7 +569,8 @@ mod tests { let deserialized = ZonedMetadata::deserialize(&metadata.serialize())?; let session = VortexSession::empty().with::(); - let aggregate_fns = aggregate_fns_from_specs(&deserialized.aggregate_specs, &session)?; + let aggregate_fns = try_aggregate_fns_from_specs(&deserialized.aggregate_specs, &session)? + .expect("known aggregates resolve"); assert_eq!(aggregate_fns.as_ref(), std::slice::from_ref(&aggregate_fn)); Ok(()) @@ -671,4 +690,74 @@ mod tests { assert!(result.is_err()); } + + #[test] + fn test_build_unknown_aggregate_disables_pruning_when_allowed() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let read_ctx = ReadContext::new([]); + let children = OwnedLayoutChildren::layout_children(vec![ + FlatLayout::new(0, dtype.clone(), SegmentId::from(0), read_ctx.clone()).into_layout(), + FlatLayout::new(0, dtype.clone(), SegmentId::from(1), read_ctx).into_layout(), + ]); + let session = vortex_array::array_session(); + session.allow_unknown(); + let build_read_ctx = ReadContext::new([]); + let build_ctx = LayoutBuildContext { + session: &session, + array_read_ctx: &build_read_ctx, + }; + + let metadata = ZonedMetadata { + zone_len: 8, + aggregate_specs: Arc::new([AggregateSpecProto::new_unknown("vortex.test.unknown")]), + }; + + let layout = ::build( + &ZonedLayoutEncoding, + &dtype, + 0, + &metadata, + vec![], + children.as_ref(), + &build_ctx, + )?; + + // An unknown aggregate disables zoned pruning (`zone_len == 0`) while leaving the data + // child readable, so the index is ignored rather than turned into a hard read error. + assert_eq!(layout.zone_len, 0); + Ok(()) + } + + #[test] + fn test_build_unknown_aggregate_errors_without_allow_unknown() { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let read_ctx = ReadContext::new([]); + let children = OwnedLayoutChildren::layout_children(vec![ + FlatLayout::new(0, dtype.clone(), SegmentId::from(0), read_ctx.clone()).into_layout(), + FlatLayout::new(0, dtype.clone(), SegmentId::from(1), read_ctx).into_layout(), + ]); + let session = vortex_array::array_session(); + let build_read_ctx = ReadContext::new([]); + let build_ctx = LayoutBuildContext { + session: &session, + array_read_ctx: &build_read_ctx, + }; + + let metadata = ZonedMetadata { + zone_len: 8, + aggregate_specs: Arc::new([AggregateSpecProto::new_unknown("vortex.test.unknown")]), + }; + + let result = ::build( + &ZonedLayoutEncoding, + &dtype, + 0, + &metadata, + vec![], + children.as_ref(), + &build_ctx, + ); + + assert!(result.is_err()); + } } diff --git a/vortex-layout/src/layouts/zoned/schema.rs b/vortex-layout/src/layouts/zoned/schema.rs index d05f5672640..f35eb4a59a9 100644 --- a/vortex-layout/src/layouts/zoned/schema.rs +++ b/vortex-layout/src/layouts/zoned/schema.rs @@ -43,10 +43,20 @@ impl AggregateSpecProto { }) } - pub(crate) fn to_aggregate_fn(&self, session: &VortexSession) -> VortexResult { + /// Resolve this spec to its aggregate function, returning `Ok(None)` for an unknown aggregate + /// ID when the session allows unknown plugins. + /// + /// A `None` result means the zone map cannot be reconstructed for this aggregate, so callers + /// should disable zoned pruning rather than fail the read. + fn to_aggregate_fn_opt(&self, session: &VortexSession) -> VortexResult> { #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let aggregate_fn_id = AggregateFnId::new(self.id.as_str()); + let Some(plugin) = session.aggregate_fns().find_plugin(&aggregate_fn_id) else { + if session.allows_unknown() { + return Ok(None); + } + vortex_bail!("unknown aggregate function id: {}", self.id); }; @@ -59,7 +69,17 @@ impl AggregateSpecProto { ); } - Ok(aggregate_fn) + Ok(Some(aggregate_fn)) + } + + /// Construct a spec referencing an aggregate ID that no session resolves, for exercising the + /// unknown-aggregate read paths. + #[cfg(test)] + pub(crate) fn new_unknown(id: &str) -> Self { + Self { + id: id.to_string(), + options: vec![], + } } } @@ -112,6 +132,10 @@ pub(crate) fn legacy_stats_table_dtype(column_dtype: &DType, present_stats: &[St ) } +/// Serializes each live [`AggregateFnRef`] into an [`AggregateSpecProto`] for storage in zoned +/// metadata. The inverse of [`try_aggregate_fns_from_specs`]. +/// +/// Errors if any aggregate is not serializable. pub(crate) fn aggregate_specs_from_fns( aggregate_fns: &[AggregateFnRef], ) -> VortexResult> { @@ -122,15 +146,27 @@ pub(crate) fn aggregate_specs_from_fns( .map(Into::into) } -pub(crate) fn aggregate_fns_from_specs( +/// Resolves each [`AggregateSpecProto`] into a live [`AggregateFnRef`], tolerating aggregates this +/// session does not recognize. Rebuilding the `zones` child needs every spec resolved, since its +/// dtype is derived from them. +/// +/// Returns: +/// +/// - `Ok(Some(fns))`: all resolved, so the zone map can be rebuilt. +/// - `Ok(None)`: an aggregate is unknown but unknown plugins are allowed (see +/// [`VortexSession::allows_unknown`]). The caller should disable pruning and scan only the `data` +/// child, which still returns correct results. +/// - `Err(..)`: any other failure (an unknown aggregate while disallowed, or undeserializable +/// options). +pub(crate) fn try_aggregate_fns_from_specs( aggregate_specs: &[AggregateSpecProto], session: &VortexSession, -) -> VortexResult> { +) -> VortexResult>> { aggregate_specs .iter() - .map(|aggregate_spec| aggregate_spec.to_aggregate_fn(session)) - .collect::>>() - .map(Into::into) + .map(|aggregate_spec| aggregate_spec.to_aggregate_fn_opt(session)) + .collect::>>>() + .map(|aggregate_fns| aggregate_fns.map(Arc::from)) } pub(crate) fn aggregate_state_dtype( @@ -166,6 +202,22 @@ mod tests { use super::*; + #[test] + fn unknown_aggregate_errors_without_allow_unknown() { + let session = vortex_array::array_session(); + let specs = [AggregateSpecProto::new_unknown("vortex.test.unknown")]; + assert!(try_aggregate_fns_from_specs(&specs, &session).is_err()); + } + + #[test] + fn unknown_aggregate_is_none_with_allow_unknown() -> VortexResult<()> { + let session = vortex_array::array_session(); + session.allow_unknown(); + let specs = [AggregateSpecProto::new_unknown("vortex.test.unknown")]; + assert!(try_aggregate_fns_from_specs(&specs, &session)?.is_none()); + Ok(()) + } + #[test] fn stats_table_dtype_adds_truncation_flags() { let dtype = legacy_stats_table_dtype(