Bypass unknown aggregates in ZonedLayout

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::<AggregateFnSession>();
-        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 = <Zoned as VTable>::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 = <Zoned as VTable>::build(
+            &ZonedLayoutEncoding,
+            &dtype,
+            0,
+            &metadata,
+            vec![],
+            children.as_ref(),
+            &build_ctx,
+        );
+
+        assert!(result.is_err());
+    }
 }

vortex-layout/src/layouts/zoned/schema.rs

@@ -43,10 +43,20 @@ impl AggregateSpecProto {
         })
     }
 
-    pub(crate) fn to_aggregate_fn(&self, session: &VortexSession) -> VortexResult<AggregateFnRef> {
+    /// 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<Option<AggregateFnRef>> {
         #[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<Arc<[AggregateSpecProto]>> {
@@ -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<Arc<[AggregateFnRef]>> {
+) -> VortexResult<Option<Arc<[AggregateFnRef]>>> {
     aggregate_specs
         .iter()
-        .map(|aggregate_spec| aggregate_spec.to_aggregate_fn(session))
-        .collect::<VortexResult<Vec<_>>>()
-        .map(Into::into)
+        .map(|aggregate_spec| aggregate_spec.to_aggregate_fn_opt(session))
+        .collect::<VortexResult<Option<Vec<_>>>>()
+        .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(