simpler

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

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

@@ -14,18 +14,21 @@ use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::struct_::StructArrayExt;
 use vortex_array::dtype::DType;
-use vortex_array::dtype::FieldName;
-use vortex_array::dtype::FieldNames;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
 use vortex_array::dtype::StructFields;
 use vortex_array::expr::Expression;
+use vortex_array::expr::get_item;
+use vortex_array::expr::merge;
+use vortex_array::expr::pack;
 use vortex_array::expr::pruning::ROW_COUNT_FIELD;
+use vortex_array::expr::root;
 use vortex_array::expr::stats::Precision;
 use vortex_array::expr::stats::Stat;
 use vortex_array::expr::stats::StatsProvider;
 use vortex_array::stats::StatsSet;
 use vortex_array::validity::Validity;
+use vortex_buffer::buffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -188,7 +191,6 @@ impl ZoneMap {
     ) -> VortexResult<Mask> {
         let mut ctx = session.create_execution_ctx();
         self.with_row_count_field(zone_len, row_count, &mut ctx)?
-            .into_array()
             .apply(predicate)?
             .execute::<Mask>(&mut ctx)
     }
@@ -200,21 +202,22 @@ impl ZoneMap {
         zone_len: u64,
         row_count: u64,
         ctx: &mut ExecutionCtx,
-    ) -> VortexResult<StructArray> {
+    ) -> VortexResult<ArrayRef> {
         let num_zones = self.array.len();
         let row_count = self.row_count_array(zone_len, row_count, num_zones, ctx)?;
-
-        let mut names: Vec<FieldName> = self.array.names().iter().cloned().collect();
-        let mut fields: Vec<ArrayRef> = self.array.iter_unmasked_fields().cloned().collect();
-        names.push(FieldName::from(ZONE_ROW_COUNT_FIELD));
-        fields.push(row_count);
-
-        StructArray::try_new(
-            FieldNames::from_iter(names),
-            fields,
-            num_zones,
-            self.array.struct_validity(),
-        )
+        let wrapper = StructArray::from_fields(&[
+            ("zone_map", self.array.clone().into_array()),
+            ("row_count", row_count),
+        ])?
+        .into_array();
+        let expr = merge([
+            get_item("zone_map", root()),
+            pack(
+                [(ZONE_ROW_COUNT_FIELD, get_item("row_count", root()))],
+                Nullability::NonNullable,
+            ),
+        ]);
+        wrapper.apply(&expr)
     }
 
     fn row_count_array(
@@ -228,18 +231,27 @@ impl ZoneMap {
             return Ok(ConstantArray::new(0u64, 0).into_array());
         }
 
-        let last_zone_len =
-            row_count.saturating_sub(zone_len.saturating_mul((num_zones - 1) as u64));
+        let last_zone_len = row_count - (zone_len.saturating_mul((num_zones - 1) as u64));
 
         if num_zones == 1 || last_zone_len == zone_len {
-            let ends = PrimitiveArray::from_iter([num_zones as u64]).into_array();
-            let values = PrimitiveArray::from_iter([last_zone_len]).into_array();
-            return Ok(RunEnd::try_new(ends, values, ctx)?.into_array());
+            return Ok(ConstantArray::new(last_zone_len, num_zones).into_array());
+        }
+
+        let ends = unsafe {
+            PrimitiveArray::new_unchecked(
+                buffer![num_zones as u64 - 1, num_zones as u64],
+                Validity::NonNullable,
+            )
+        }
+        .into_array();
+        let values = unsafe {
+            PrimitiveArray::new_unchecked(buffer![zone_len, last_zone_len], Validity::NonNullable)
         }
+        .into_array();
 
-        let ends = PrimitiveArray::from_iter([num_zones as u64 - 1, num_zones as u64]).into_array();
-        let values = PrimitiveArray::from_iter([zone_len, last_zone_len]).into_array();
-        Ok(RunEnd::try_new(ends, values, ctx)?.into_array())
+        // SAFETY: `ends` are strictly increasing, terminate at `num_zones`, and align one-to-one
+        // with the non-null run values.
+        Ok(unsafe { RunEnd::new_unchecked(ends, values, 0, num_zones, ctx) }.into_array())
     }
 }