vortex_scalar/
extension.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::{Display, Formatter};
5use std::hash::Hash;
6use std::sync::Arc;
7
8use vortex_dtype::datetime::{TemporalMetadata, is_temporal_ext_type};
9use vortex_dtype::{DType, ExtDType};
10use vortex_error::{VortexError, VortexResult, vortex_bail};
11
12use crate::{Scalar, ScalarValue};
13
14/// A scalar value representing an extension type.
15///
16/// Extension types allow wrapping a storage type with custom semantics.
17pub struct ExtScalar<'a> {
18    ext_dtype: &'a ExtDType,
19    value: &'a ScalarValue,
20}
21
22impl Display for ExtScalar<'_> {
23    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
24        // Specialized handling for date/time/timestamp builtin extension types.
25        if is_temporal_ext_type(self.ext_dtype().id()) {
26            let metadata =
27                TemporalMetadata::try_from(self.ext_dtype()).map_err(|_| std::fmt::Error)?;
28
29            let maybe_timestamp = self
30                .storage()
31                .as_primitive()
32                .as_::<i64>()
33                .and_then(|maybe_timestamp| {
34                    maybe_timestamp.map(|v| metadata.to_jiff(v)).transpose()
35                })
36                .map_err(|_| std::fmt::Error)?;
37
38            match maybe_timestamp {
39                None => write!(f, "null"),
40                Some(v) => write!(f, "{v}"),
41            }
42        } else {
43            write!(f, "{}({})", self.ext_dtype().id(), self.storage())
44        }
45    }
46}
47
48impl PartialEq for ExtScalar<'_> {
49    fn eq(&self, other: &Self) -> bool {
50        self.ext_dtype.eq_ignore_nullability(other.ext_dtype) && self.storage() == other.storage()
51    }
52}
53
54impl Eq for ExtScalar<'_> {}
55
56// Ord is not implemented since it's undefined for different Extension DTypes
57impl PartialOrd for ExtScalar<'_> {
58    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
59        if !self.ext_dtype.eq_ignore_nullability(other.ext_dtype) {
60            return None;
61        }
62        self.storage().partial_cmp(&other.storage())
63    }
64}
65
66impl Hash for ExtScalar<'_> {
67    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
68        self.ext_dtype.hash(state);
69        self.storage().hash(state);
70    }
71}
72
73impl<'a> ExtScalar<'a> {
74    /// Creates a new extension scalar from a data type and scalar value.
75    ///
76    /// # Errors
77    ///
78    /// Returns an error if the data type is not an extension type.
79    pub fn try_new(dtype: &'a DType, value: &'a ScalarValue) -> VortexResult<Self> {
80        let DType::Extension(ext_dtype) = dtype else {
81            vortex_bail!("Expected extension scalar, found {}", dtype)
82        };
83
84        Ok(Self { ext_dtype, value })
85    }
86
87    /// Returns the storage scalar of the extension scalar.
88    pub fn storage(&self) -> Scalar {
89        Scalar::new(self.ext_dtype.storage_dtype().clone(), self.value.clone())
90    }
91
92    /// Returns the extension data type.
93    pub fn ext_dtype(&self) -> &'a ExtDType {
94        self.ext_dtype
95    }
96
97    pub(crate) fn cast(&self, dtype: &DType) -> VortexResult<Scalar> {
98        if self.value.is_null() && !dtype.is_nullable() {
99            vortex_bail!(
100                "cannot cast extension dtype with id {} and storage type {} to {}",
101                self.ext_dtype.id(),
102                self.ext_dtype.storage_dtype(),
103                dtype
104            );
105        }
106
107        if self.ext_dtype.storage_dtype().eq_ignore_nullability(dtype) {
108            // Casting from an extension type to the underlying storage type is OK.
109            return Ok(Scalar::new(dtype.clone(), self.value.clone()));
110        }
111
112        if let DType::Extension(ext_dtype) = dtype {
113            if self.ext_dtype.eq_ignore_nullability(ext_dtype) {
114                return Ok(Scalar::new(dtype.clone(), self.value.clone()));
115            }
116        }
117
118        vortex_bail!(
119            "cannot cast extension dtype with id {} and storage type {} to {}",
120            self.ext_dtype.id(),
121            self.ext_dtype.storage_dtype(),
122            dtype
123        );
124    }
125}
126
127impl<'a> TryFrom<&'a Scalar> for ExtScalar<'a> {
128    type Error = VortexError;
129
130    fn try_from(value: &'a Scalar) -> Result<Self, Self::Error> {
131        ExtScalar::try_new(value.dtype(), &value.value)
132    }
133}
134
135impl Scalar {
136    /// Creates a new extension scalar wrapping the given storage value.
137    pub fn extension(ext_dtype: Arc<ExtDType>, value: Scalar) -> Self {
138        // TODO(joe): enable once we use rust duckdb
139        // assert_eq!(ext_dtype.storage_dtype(), value.dtype());
140        Self {
141            dtype: DType::Extension(ext_dtype),
142            value: value.value().clone(),
143        }
144    }
145}
vortex_scalar/extension.rs

vortex_scalar/
extension.rs
