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

vortex_array/arrays/extension/
mod.rs
