spiceai-54-vortex-0.85.0 : cherry-picks

encodings/datetime-parts/src/ops.rs

@@ -65,9 +65,11 @@ impl OperationsVTable<DateTimeParts> for DateTimeParts {
             options.unit,
         );
 
-        Ok(Scalar::extension::<Timestamp>(
-            options.clone(),
-            Scalar::primitive(ts, ext.storage_dtype().nullability()),
-        ))
+        // Build the extension scalar fallibly: an extension type whose metadata the storage value
+        // cannot satisfy (an unresolvable timezone, say) must surface as an error rather than abort
+        // the process, since `execute_scalar` dispatches here when reading a value out of a
+        // timestamp column stored under this encoding.
+        let storage = Scalar::primitive(ts, ext.storage_dtype().nullability());
+        Scalar::try_extension_ref(ext, storage)
     }
 }

vortex-array/src/aggregate_fn/fns/min_max/extension.rs

@@ -19,15 +19,21 @@ pub(super) fn accumulate_extension(
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<()> {
     let non_nullable_ext_dtype = array.ext_dtype().with_nullability(Nullability::NonNullable);
+    // Build the extension scalars fallibly: an extension type whose metadata the storage value
+    // cannot satisfy (an unresolvable timezone, say) must fail the aggregation rather than abort
+    // the process, since this runs on the write path when building zone maps.
     let local = min_max(
         array.storage_array(),
         ctx,
         NumericalAggregateOpts::default(),
     )?
-    .map(|MinMaxResult { min, max }| MinMaxResult {
-        min: Scalar::extension_ref(non_nullable_ext_dtype.clone(), min),
-        max: Scalar::extension_ref(non_nullable_ext_dtype, max),
-    });
+    .map(|MinMaxResult { min, max }| -> VortexResult<MinMaxResult> {
+        Ok(MinMaxResult {
+            min: Scalar::try_extension_ref(non_nullable_ext_dtype.clone(), min)?,
+            max: Scalar::try_extension_ref(non_nullable_ext_dtype, max)?,
+        })
+    })
+    .transpose()?;
     partial.merge(local);
     Ok(())
 }

vortex-array/src/arrays/dict/take.rs

@@ -156,17 +156,97 @@ pub(crate) fn propagate_take_stats(
                 st.set(Stat::IsConstant, Precision::exact(true));
             }
         }
-        let inexact_min_max = [Stat::Min, Stat::Max]
-            .into_iter()
-            .filter_map(|stat| match source.statistics().get(stat).into_inexact() {
-                Precision::Exact(scalar) | Precision::Inexact(scalar) => {
-                    scalar.into_value().map(|sv| (stat, Precision::Inexact(sv)))
-                }
-                Precision::Absent => None,
-            })
-            .collect::<SmallVec<_>>();
+        let inexact_min_max = [
+            Stat::Min,
+            Stat::Max,
+            Stat::UncompressedSizeInBytes,
+            Stat::IsConstant,
+        ]
+        .into_iter()
+        .filter_map(|stat| match source.statistics().get(stat).into_inexact() {
+            Precision::Exact(scalar) | Precision::Inexact(scalar) => {
+                scalar.into_value().map(|sv| (stat, Precision::Inexact(sv)))
+            }
+            Precision::Absent => None,
+        })
+        .collect::<SmallVec<_>>();
         st.combine_sets(
             &(unsafe { StatsSet::new_unchecked(inexact_min_max) }).as_typed_ref(source.dtype()),
         )
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::buffer;
+
+    use vortex_error::VortexExpect;
+
+    use super::*;
+    use crate::dtype::Nullability;
+    use crate::scalar::Scalar;
+
+    /// `UncompressedSizeInBytes` has to survive a dict take, inexactly, the same way `Min`/`Max`
+    /// already do: the optimizer sizes joins from it, and a value's absence after a take is read
+    /// as "unknown" rather than "zero", so losing it silently degrades plans over dictionary-
+    /// encoded columns rather than failing loudly.
+    #[test]
+    fn uncompressed_size_in_bytes_propagates_through_take_as_inexact() -> VortexResult<()> {
+        let source = buffer![1i32, 2, 3, 4].into_array();
+        source.statistics().set(
+            Stat::UncompressedSizeInBytes,
+            Precision::exact(
+                Scalar::primitive(16u64, Nullability::NonNullable)
+                    .into_value()
+                    .vortex_expect("non-null size"),
+            ),
+        );
+
+        let target = buffer![1i32, 3].into_array();
+        let indices = buffer![0u64, 2].into_array();
+        propagate_take_stats(&source, &target, &indices)?;
+
+        assert_eq!(
+            target
+                .statistics()
+                .get_as::<u64>(Stat::UncompressedSizeInBytes),
+            Precision::Inexact(16),
+            "the source's size should carry forward, downgraded to inexact"
+        );
+        Ok(())
+    }
+
+    /// `IsConstant` is handled twice: exactly when every index is valid (the block above the
+    /// one this row extends), and inexactly here otherwise. A null index breaks the "every
+    /// output value equals the source's constant" guarantee, so the exact path must not fire,
+    /// but the source's own `IsConstant` reading (whatever it is) is still a reasonable inexact
+    /// carry-forward rather than dropping the stat entirely.
+    #[test]
+    fn is_constant_propagates_inexactly_through_take_with_a_null_index() -> VortexResult<()> {
+        let source = ConstantArray::new(Scalar::primitive(7i32, Nullability::NonNullable), 4)
+            .into_array();
+        source.statistics().set(
+            Stat::IsConstant,
+            Precision::exact(
+                Scalar::from(true)
+                    .into_value()
+                    .vortex_expect("non-null bool"),
+            ),
+        );
+
+        let target = ConstantArray::new(Scalar::primitive(7i32, Nullability::NonNullable), 2)
+            .into_array();
+        // A nullable index array with one null: `indices_all_valid` is false, so the exact
+        // IsConstant path above must not run.
+        let indices = crate::arrays::PrimitiveArray::from_option_iter([Some(0u64), None])
+            .into_array();
+        propagate_take_stats(&source, &target, &indices)?;
+
+        assert_eq!(
+            target.statistics().get_as::<bool>(Stat::IsConstant),
+            Precision::Inexact(true),
+            "IsConstant should still carry forward inexactly when it can't be proven exact"
+        );
+        Ok(())
+    }
+}

vortex-array/src/arrays/extension/compute/cast.rs

@@ -1,47 +1,166 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::Extension;
 use crate::arrays::ExtensionArray;
+use crate::arrays::PrimitiveArray;
 use crate::arrays::extension::ExtensionArrayExt;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
+use crate::dtype::PType;
+use crate::extension::datetime::AnyTemporal;
+use crate::extension::datetime::DateToTimestamp;
+use crate::extension::datetime::TemporalMetadata;
+use crate::extension::datetime::TimeUnit;
+use crate::scalar_fn::fns::cast::CastKernel;
 use crate::scalar_fn::fns::cast::CastReduce;
+use crate::validity::Validity;
 
 impl CastReduce for Extension {
     fn cast(array: ArrayView<'_, Extension>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
-        if !array.dtype().eq_ignore_nullability(dtype) {
-            // Target is not the same extension type.
+        let DType::Extension(ext_dtype) = dtype else {
+            // Target is not an extension type.
             // Delegate to the storage array's cast.
             return Ok(Some(array.storage_array().cast(dtype.clone())?));
+        };
+
+        if array.ext_dtype().eq_ignore_nullability(ext_dtype) {
+            // Same extension type: restructure by casting the storage array. This is buffer-free
+            // (`cast` returns a lazy cast expression), so it stays in the reduce phase.
+            let new_storage = match array
+                .storage_array()
+                .cast(ext_dtype.storage_dtype().clone())
+            {
+                Ok(arr) => arr,
+                Err(e) => {
+                    tracing::warn!("Failed to cast storage array: {e}");
+                    return Ok(None);
+                }
+            };
+
+            return Ok(Some(
+                ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(),
+            ));
         }
 
+        // A different extension type (e.g. Date -> Timestamp) requires per-value conversion that
+        // reads buffers, so it is deferred to the execute phase in `CastKernel`, where an
+        // `ExecutionCtx` is available.
+        Ok(None)
+    }
+}
+
+impl CastKernel for Extension {
+    fn cast(
+        array: ArrayView<'_, Extension>,
+        dtype: &DType,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
         let DType::Extension(ext_dtype) = dtype else {
-            unreachable!("Already verified we have an extension dtype");
+            // Non-extension targets are restructured buffer-free by `CastReduce`.
+            return Ok(None);
         };
 
-        let new_storage = match array
-            .storage_array()
-            .cast(ext_dtype.storage_dtype().clone())
-        {
-            Ok(arr) => arr,
-            Err(e) => {
-                tracing::warn!("Failed to cast storage array: {e}");
-                return Ok(None);
-            }
-        };
+        if let Some(new_storage) = cast_temporal_date_to_timestamp(&array, dtype, ctx)? {
+            return Ok(Some(
+                ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(),
+            ));
+        }
 
-        Ok(Some(
-            ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(),
-        ))
+        Ok(None)
     }
 }
 
+fn cast_temporal_date_to_timestamp(
+    array: &ArrayView<'_, Extension>,
+    target_dtype: &DType,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Option<ArrayRef>> {
+    let DType::Extension(target_ext_dtype) = target_dtype else {
+        return Ok(None);
+    };
+
+    let Some(source_temporal) = array.ext_dtype().metadata_opt::<AnyTemporal>() else {
+        return Ok(None);
+    };
+    let Some(target_temporal) = target_ext_dtype.metadata_opt::<AnyTemporal>() else {
+        return Ok(None);
+    };
+
+    let (TemporalMetadata::Date(source_unit), TemporalMetadata::Timestamp(target_unit, _)) =
+        (source_temporal, target_temporal)
+    else {
+        return Ok(None);
+    };
+
+    let source_i64 = array
+        .storage_array()
+        .cast(DType::Primitive(PType::I64, array.dtype().nullability()))?
+        .execute::<PrimitiveArray>(ctx)?;
+
+    let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit, ctx)?;
+
+    let new_storage = converted
+        .into_array()
+        .cast(target_ext_dtype.storage_dtype().clone())?
+        .execute::<PrimitiveArray>(ctx)?
+        .into_array();
+
+    Ok(Some(new_storage))
+}
+
+fn cast_date_values_to_timestamp(
+    values: &PrimitiveArray,
+    source_unit: TimeUnit,
+    target_unit: TimeUnit,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<PrimitiveArray> {
+    let conversion = DateToTimestamp::new(source_unit, target_unit)?;
+
+    let input = values.as_slice::<i64>();
+    let validity = values.validity()?;
+    let mut output = BufferMut::with_capacity(input.len());
+
+    match &validity {
+        Validity::NonNullable | Validity::AllValid => {
+            for &value in input {
+                // SAFETY: output has sufficient capacity for all pushed values.
+                unsafe { output.push_unchecked(conversion.convert(value)?) };
+            }
+        }
+        Validity::AllInvalid => {
+            for _ in 0..input.len() {
+                // SAFETY: output has sufficient capacity for all pushed values.
+                unsafe { output.push_unchecked(0i64) };
+            }
+        }
+        Validity::Array(_) => {
+            // Resolve validity to a boolean mask once. Null slots keep a placeholder 0 so a garbage
+            // source value in a null slot cannot trip the overflow check in
+            // `DateToTimestamp::convert`; the output re-uses `validity`, so those slots stay null.
+            let mask = validity.execute_mask(input.len(), ctx)?;
+            for (i, &value) in input.iter().enumerate() {
+                let converted = if mask.value(i) {
+                    conversion.convert(value)?
+                } else {
+                    0i64
+                };
+                // SAFETY: output has sufficient capacity for all pushed values.
+                unsafe { output.push_unchecked(converted) };
+            }
+        }
+    }
+
+    Ok(PrimitiveArray::new(output.freeze(), validity))
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::LazyLock;
@@ -61,6 +180,7 @@ mod tests {
     use crate::dtype::Nullability;
     use crate::dtype::PType;
     use crate::executor::VortexSessionExecute;
+    use crate::extension::datetime::Date;
     use crate::extension::datetime::TimeUnit;
     use crate::extension::datetime::Timestamp;
 
@@ -99,6 +219,50 @@ mod tests {
         assert_eq!(output.dtype(), &new_dtype);
     }
 
+    #[test]
+    fn cast_date_days_to_timestamp_nanoseconds() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let source_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased();
+        let target_dtype = Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased();
+
+        let arr = ExtensionArray::new(source_dtype, buffer![0i32, 1, -1].into_array());
+        let output = arr
+            .into_array()
+            .cast(DType::Extension(target_dtype.clone()))?;
+
+        let expected = ExtensionArray::new(
+            target_dtype,
+            buffer![0i64, 86_400_000_000_000, -86_400_000_000_000].into_array(),
+        )
+        .into_array();
+        assert_arrays_eq!(output, expected, &mut ctx);
+        Ok(())
+    }
+
+    #[test]
+    fn cast_date_days_to_timestamp_seconds_nullable() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let source_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased();
+        let target_dtype = Timestamp::new(TimeUnit::Seconds, Nullability::Nullable).erased();
+
+        let arr = ExtensionArray::new(
+            source_dtype,
+            PrimitiveArray::from_option_iter([Some(0i32), None, Some(2)]).into_array(),
+        );
+
+        let output = arr
+            .into_array()
+            .cast(DType::Extension(target_dtype.clone()))?;
+
+        let expected = ExtensionArray::new(
+            target_dtype,
+            PrimitiveArray::from_option_iter([Some(0i64), None, Some(172_800)]).into_array(),
+        )
+        .into_array();
+        assert_arrays_eq!(output, expected, &mut ctx);
+        Ok(())
+    }
+
     #[test]
     fn cast_different_ext_dtype() {
         let original_dtype =

vortex-array/src/arrays/extension/compute/rules.rs

@@ -36,7 +36,7 @@ impl ArrayReduceRule<Extension> for ExtensionConstantRule {
         };
 
         let storage_scalar = const_array.scalar().clone();
-        let ext_scalar = Scalar::extension_ref(array.ext_dtype().clone(), storage_scalar);
+        let ext_scalar = Scalar::try_extension_ref(array.ext_dtype().clone(), storage_scalar)?;
 
         let constant_with_extension_scalar =
             ConstantArray::new(ext_scalar, array.len()).into_array();