vortex_datetime_parts/
ops.rs

1use vortex_array::vtable::OperationsVTable;
2use vortex_array::{Array, ArrayRef, IntoArray};
3use vortex_dtype::Nullability::{NonNullable, Nullable};
4use vortex_dtype::datetime::TemporalMetadata;
5use vortex_dtype::{DType, PType};
6use vortex_error::{VortexResult, vortex_bail};
7use vortex_scalar::Scalar;
8
9use crate::timestamp::TimestampParts;
10use crate::{DateTimePartsArray, DateTimePartsVTable, timestamp};
11
12impl OperationsVTable<DateTimePartsVTable> for DateTimePartsVTable {
13    fn slice(array: &DateTimePartsArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
14        Ok(DateTimePartsArray::try_new(
15            array.dtype().clone(),
16            array.days().slice(start, stop)?,
17            array.seconds().slice(start, stop)?,
18            array.subseconds().slice(start, stop)?,
19        )?
20        .into_array())
21    }
22
23    fn scalar_at(array: &DateTimePartsArray, index: usize) -> VortexResult<Scalar> {
24        let DType::Extension(ext) = array.dtype().clone() else {
25            vortex_bail!(
26                "DateTimePartsArray must have extension dtype, found {}",
27                array.dtype()
28            );
29        };
30
31        let Ok(temporal_metadata) = TemporalMetadata::try_from(ext.as_ref()) else {
32            vortex_bail!(ComputeError: "must decode TemporalMetadata from extension metadata");
33        };
34
35        if !array.is_valid(index)? {
36            return Ok(Scalar::null(DType::Extension(ext)));
37        }
38
39        let days: i64 = array
40            .days()
41            .scalar_at(index)?
42            .cast(&DType::Primitive(PType::I64, Nullable))?
43            .try_into()?;
44        let seconds: i64 = array
45            .seconds()
46            .scalar_at(index)?
47            .cast(&DType::Primitive(PType::I64, NonNullable))?
48            .try_into()?;
49        let subseconds: i64 = array
50            .subseconds()
51            .scalar_at(index)?
52            .cast(&DType::Primitive(PType::I64, NonNullable))?
53            .try_into()?;
54
55        let ts = timestamp::combine(
56            TimestampParts {
57                days,
58                seconds,
59                subseconds,
60            },
61            temporal_metadata.time_unit(),
62        )?;
63
64        Ok(Scalar::extension(ext, Scalar::from(ts)))
65    }
66}