vortex_array/arrays/extension/
mod.rs

1use std::sync::Arc;
2
3use vortex_dtype::{DType, ExtDType, ExtID};
4use vortex_error::VortexResult;
5use vortex_mask::Mask;
6
7use crate::array::{ArrayCanonicalImpl, ArrayValidityImpl};
8use crate::stats::{ArrayStats, StatsSetRef};
9use crate::variants::ExtensionArrayTrait;
10use crate::vtable::VTableRef;
11use crate::{
12    Array, ArrayImpl, ArrayRef, ArrayStatisticsImpl, ArrayVariantsImpl, Canonical, EmptyMetadata,
13    Encoding,
14};
15mod compute;
16mod serde;
17
18#[derive(Clone, Debug)]
19pub struct ExtensionArray {
20    dtype: DType,
21    storage: ArrayRef,
22    stats_set: ArrayStats,
23}
24
25pub struct ExtensionEncoding;
26impl Encoding for ExtensionEncoding {
27    type Array = ExtensionArray;
28    type Metadata = EmptyMetadata;
29}
30
31impl ExtensionArray {
32    pub fn new(ext_dtype: Arc<ExtDType>, storage: ArrayRef) -> Self {
33        assert_eq!(
34            ext_dtype.storage_dtype(),
35            storage.dtype(),
36            "ExtensionArray: storage_dtype must match storage array DType",
37        );
38        Self {
39            dtype: DType::Extension(ext_dtype),
40            storage,
41            stats_set: ArrayStats::default(),
42        }
43    }
44
45    pub fn storage(&self) -> &ArrayRef {
46        &self.storage
47    }
48
49    #[allow(dead_code)]
50    #[inline]
51    pub fn id(&self) -> &ExtID {
52        self.ext_dtype().id()
53    }
54}
55
56impl ArrayImpl for ExtensionArray {
57    type Encoding = ExtensionEncoding;
58
59    fn _len(&self) -> usize {
60        self.storage.len()
61    }
62
63    fn _dtype(&self) -> &DType {
64        &self.dtype
65    }
66
67    fn _vtable(&self) -> VTableRef {
68        VTableRef::new_ref(&ExtensionEncoding)
69    }
70
71    fn _with_children(&self, children: &[ArrayRef]) -> VortexResult<Self> {
72        Ok(Self::new(self.ext_dtype().clone(), children[0].clone()))
73    }
74}
75
76impl ArrayStatisticsImpl for ExtensionArray {
77    fn _stats_ref(&self) -> StatsSetRef<'_> {
78        self.stats_set.to_ref(self)
79    }
80}
81
82impl ArrayCanonicalImpl for ExtensionArray {
83    fn _to_canonical(&self) -> VortexResult<Canonical> {
84        Ok(Canonical::Extension(self.clone()))
85    }
86}
87
88impl ArrayValidityImpl for ExtensionArray {
89    fn _is_valid(&self, index: usize) -> VortexResult<bool> {
90        self.storage.is_valid(index)
91    }
92
93    fn _all_valid(&self) -> VortexResult<bool> {
94        self.storage.all_valid()
95    }
96
97    fn _all_invalid(&self) -> VortexResult<bool> {
98        self.storage.all_invalid()
99    }
100
101    fn _validity_mask(&self) -> VortexResult<Mask> {
102        self.storage.validity_mask()
103    }
104}
105
106impl ArrayVariantsImpl for ExtensionArray {
107    fn _as_extension_typed(&self) -> Option<&dyn ExtensionArrayTrait> {
108        Some(self)
109    }
110}
111
112impl ExtensionArrayTrait for ExtensionArray {
113    fn storage_data(&self) -> ArrayRef {
114        self.storage().clone()
115    }
116}
vortex_array/arrays/extension/mod.rs

vortex_array/arrays/extension/
mod.rs
