Skip to content

Commit af5e659

Browse files
gatesngatesn
authored
Remove ExtDTypeInner (#6874)
Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent c3d57a1 commit af5e659

20 files changed

Lines changed: 196 additions & 274 deletions

File tree

‎docs/developer-guide/internals/vtables.md‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -63,21 +63,21 @@ registered in the session by their ID.
6363

6464
For a concept `Foo`, the components are organized into these files:
6565

66-
| File | Contains |
67-
|---------------|-----------------------------------------------------------------------|
68-
| `vtable.rs` | `FooVTable` trait definition |
69-
| `typed.rs` | `Foo<V>` data struct, inherent methods, `Deref` impl |
70-
| `erased.rs` | `FooRef` struct, `DynFoo` sealed trait, blanket impl |
71-
| `plugin.rs` | `FooPlugin` trait, registration |
72-
| `matcher.rs` | Downcasting helpers (`is`, `as_`, `as_opt`, pattern matching traits) |
66+
| File | Contains |
67+
|--------------|----------------------------------------------------------------------|
68+
| `vtable.rs` | `FooVTable` trait definition |
69+
| `typed.rs` | `Foo<V>` data struct, inherent methods, `Deref` impl |
70+
| `erased.rs` | `FooRef` struct, `DynFoo` sealed trait, blanket impl |
71+
| `plugin.rs` | `FooPlugin` trait, registration |
72+
| `matcher.rs` | Downcasting helpers (`is`, `as_`, `as_opt`, pattern matching traits) |
7373

7474
For Array encodings, each encoding has its own module (e.g. `arrays/primitive/`):
7575

76-
| File | Contains |
77-
|-------------------------|-------------------------------------------------------------|
78-
| `arrays/foo/mod.rs` | `V::Array` associated type, encoding-specific methods on it |
79-
| `arrays/foo/vtable.rs` | `ArrayVTable` impl for this encoding |
80-
| `arrays/foo/compute/` | Compute kernel implementations |
76+
| File | Contains |
77+
|------------------------|-------------------------------------------------------------|
78+
| `arrays/foo/mod.rs` | `V::Array` associated type, encoding-specific methods on it |
79+
| `arrays/foo/vtable.rs` | `ArrayVTable` impl for this encoding |
80+
| `arrays/foo/compute/` | Compute kernel implementations |
8181

8282
## Example: ExtDType
8383

‎vortex-array/public-api.lock‎

Lines changed: 35 additions & 35 deletions
Large diffs are not rendered by default.

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

Lines changed: 4 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -95,18 +95,13 @@ mod tests {
9595
Ok(EmptyMetadata)
9696
}
9797

98-
fn validate_dtype(
99-
&self,
100-
_options: &Self::Metadata,
101-
_storage_dtype: &DType,
102-
) -> VortexResult<()> {
98+
fn validate_dtype(&self, _extension_dtype: &ExtDType<Self>) -> VortexResult<()> {
10399
Ok(())
104100
}
105101

106102
fn unpack_native<'a>(
107103
&self,
108-
_metadata: &'a Self::Metadata,
109-
_storage_dtype: &'a DType,
104+
_extension_dtype: &'a ExtDType<Self>,
110105
_storage_value: &'a ScalarValue,
111106
) -> VortexResult<Self::NativeValue<'a>> {
112107
Ok("")
@@ -197,18 +192,13 @@ mod tests {
197192
Ok(EmptyMetadata)
198193
}
199194

200-
fn validate_dtype(
201-
&self,
202-
_options: &Self::Metadata,
203-
_storage_dtype: &DType,
204-
) -> VortexResult<()> {
195+
fn validate_dtype(&self, _extension_dtype: &ExtDType<Self>) -> VortexResult<()> {
205196
Ok(())
206197
}
207198

208199
fn unpack_native<'a>(
209200
&self,
210-
_metadata: &'a Self::Metadata,
211-
_storage_dtype: &'a DType,
201+
_extension_dtype: &'a ExtDType<Self>,
212202
_storage_value: &'a ScalarValue,
213203
) -> VortexResult<Self::NativeValue<'a>> {
214204
Ok("")

‎vortex-array/src/dtype/extension/erased.rs‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ use crate::dtype::extension::ExtId;
2020
use crate::dtype::extension::ExtVTable;
2121
use crate::dtype::extension::Matcher;
2222
use crate::dtype::extension::typed::DynExtDType;
23-
use crate::dtype::extension::typed::ExtDTypeInner;
2423
use crate::scalar::ScalarValue;
2524

2625
/// A type-erased extension dtype.
@@ -40,12 +39,12 @@ pub struct ExtDTypeRef(pub(super) Arc<dyn DynExtDType>);
4039
impl ExtDTypeRef {
4140
/// Returns the [`ExtId`] identifying this extension type.
4241
pub fn id(&self) -> ExtId {
43-
self.0.id()
42+
self.0.ext_id()
4443
}
4544

4645
/// Returns the storage dtype of the extension type.
4746
pub fn storage_dtype(&self) -> &DType {
48-
self.0.storage_dtype()
47+
self.0.ext_storage_dtype()
4948
}
5049

5150
/// Returns the nullability of the storage dtype.
@@ -127,12 +126,10 @@ impl ExtDTypeRef {
127126
/// Downcast to the concrete [`ExtDType`].
128127
///
129128
/// Returns `Err(self)` if the downcast fails.
130-
pub fn try_downcast<V: ExtVTable>(self) -> Result<ExtDType<V>, ExtDTypeRef> {
131-
if self.0.as_any().is::<ExtDTypeInner<V>>() {
132-
// SAFETY: type matches and ExtDTypeInner<V> is the only implementor
133-
let ptr = Arc::into_raw(self.0) as *const ExtDTypeInner<V>;
134-
let inner = unsafe { Arc::from_raw(ptr) };
135-
Ok(ExtDType(inner))
129+
pub fn try_downcast<V: ExtVTable>(self) -> Result<Arc<ExtDType<V>>, ExtDTypeRef> {
130+
if self.0.as_any().is::<ExtDType<V>>() {
131+
let ptr = Arc::into_raw(self.0) as *const ExtDType<V>;
132+
Ok(unsafe { Arc::from_raw(ptr) })
136133
} else {
137134
Err(self)
138135
}
@@ -143,12 +140,12 @@ impl ExtDTypeRef {
143140
/// # Panics
144141
///
145142
/// Panics if the downcast fails.
146-
pub fn downcast<V: ExtVTable>(self) -> ExtDType<V> {
143+
pub fn downcast<V: ExtVTable>(self) -> Arc<ExtDType<V>> {
147144
self.try_downcast::<V>()
148145
.map_err(|this| {
149146
vortex_err!(
150147
"Failed to downcast ExtDTypeRef {} to {}",
151-
this.0.id(),
148+
this.0.ext_id(),
152149
type_name::<V>(),
153150
)
154151
})

‎vortex-array/src/dtype/extension/matcher.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4+
use crate::dtype::extension::ExtDType;
45
use crate::dtype::extension::ExtDTypeRef;
56
use crate::dtype::extension::ExtVTable;
6-
use crate::dtype::extension::typed::ExtDTypeInner;
77

88
/// A trait for matching extension dtypes.
99
pub trait Matcher {
@@ -23,13 +23,13 @@ impl<V: ExtVTable> Matcher for V {
2323
type Match<'a> = &'a V::Metadata;
2424

2525
fn matches(item: &ExtDTypeRef) -> bool {
26-
item.0.as_any().is::<ExtDTypeInner<V>>()
26+
item.0.as_any().is::<ExtDType<V>>()
2727
}
2828

2929
fn try_match<'a>(item: &'a ExtDTypeRef) -> Option<Self::Match<'a>> {
3030
item.0
3131
.as_any()
32-
.downcast_ref::<ExtDTypeInner<V>>()
33-
.map(|inner| &inner.metadata)
32+
.downcast_ref::<ExtDType<V>>()
33+
.map(|inner| inner.metadata())
3434
}
3535
}

‎vortex-array/src/dtype/extension/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,11 @@ pub type ExtId = arcref::ArcRef<str>;
3434
/// Private module to seal [`typed::DynExtDType`].
3535
mod sealed {
3636
use crate::dtype::extension::ExtVTable;
37-
use crate::dtype::extension::typed::ExtDTypeInner;
37+
use crate::dtype::extension::typed::ExtDType;
3838

3939
/// Marker trait to prevent external implementations of [`super::typed::DynExtDType`].
4040
pub(crate) trait Sealed {}
4141

4242
/// This can be the **only** implementor for [`super::typed::DynExtDType`].
43-
impl<V: ExtVTable> Sealed for ExtDTypeInner<V> {}
43+
impl<V: ExtVTable> Sealed for ExtDType<V> {}
4444
}

‎vortex-array/src/dtype/extension/typed.rs‎

Lines changed: 26 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,14 @@ use crate::scalar::ScalarValue;
3232
/// [`try_with_vtable()`]: ExtDType::try_with_vtable
3333
/// [`erased()`]: ExtDType::erased
3434
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
35-
pub struct ExtDType<V: ExtVTable>(pub(super) Arc<ExtDTypeInner<V>>);
35+
pub struct ExtDType<V: ExtVTable> {
36+
/// The extension dtype vtable.
37+
vtable: V,
38+
/// The extension dtype metadata.
39+
metadata: V::Metadata,
40+
/// The underlying storage dtype.
41+
storage_dtype: DType,
42+
}
3643

3744
/// Convenience implementation for zero-sized VTables (or VTables that implement `Default`).
3845
impl<V: ExtVTable + Default> ExtDType<V> {
@@ -49,60 +56,43 @@ impl<V: ExtVTable> ExtDType<V> {
4956
metadata: V::Metadata,
5057
storage_dtype: DType,
5158
) -> VortexResult<Self> {
52-
vtable.validate_dtype(&metadata, &storage_dtype)?;
53-
54-
Ok(Self(Arc::new(ExtDTypeInner::<V> {
59+
let this = Self {
5560
vtable,
5661
metadata,
5762
storage_dtype,
58-
})))
63+
};
64+
65+
this.vtable.validate_dtype(&this)?;
66+
67+
Ok(this)
5968
}
6069

6170
/// Returns the identifier of the extension type.
6271
pub fn id(&self) -> ExtId {
63-
self.0.vtable.id()
72+
self.vtable.id()
6473
}
6574

6675
/// Returns the vtable of the extension type.
6776
pub fn vtable(&self) -> &V {
68-
&self.0.vtable
77+
&self.vtable
6978
}
7079

7180
/// Returns the metadata of the extension type.
7281
pub fn metadata(&self) -> &V::Metadata {
73-
&self.0.metadata
82+
&self.metadata
7483
}
7584

7685
/// Returns the storage dtype of the extension type.
7786
pub fn storage_dtype(&self) -> &DType {
78-
&self.0.storage_dtype
87+
&self.storage_dtype
7988
}
8089

8190
/// Erase the concrete type information, returning a type-erased extension dtype.
8291
pub fn erased(self) -> ExtDTypeRef {
83-
ExtDTypeRef(self.0)
92+
ExtDTypeRef(Arc::new(self))
8493
}
8594
}
8695

87-
// ---------------------------------------------------------------------------
88-
// Private inner struct + sealed trait
89-
// ---------------------------------------------------------------------------
90-
91-
/// The private inner representation of an extension dtype, pairing a vtable with its metadata
92-
/// and storage dtype.
93-
///
94-
/// This is the sole implementor of [`DynExtDType`], enabling [`ExtDTypeRef`] to safely downcast
95-
/// back to the concrete vtable type via [`Any`].
96-
#[derive(Debug, PartialEq, Eq, Hash)]
97-
pub(super) struct ExtDTypeInner<V: ExtVTable> {
98-
/// The extension dtype vtable.
99-
pub(super) vtable: V,
100-
/// The extension dtype metadata.
101-
pub(super) metadata: V::Metadata,
102-
/// The underlying storage dtype.
103-
pub(super) storage_dtype: DType,
104-
}
105-
10696
/// An object-safe, sealed trait encapsulating the behavior for extension dtypes.
10797
///
10898
/// This provides type-erased access to the extension dtype's identity, storage dtype, and
@@ -111,9 +101,9 @@ pub(super) trait DynExtDType: 'static + Send + Sync + super::sealed::Sealed {
111101
/// Returns `self` as a trait object for downcasting.
112102
fn as_any(&self) -> &dyn Any;
113103
/// Returns the [`ExtId`] identifying this extension type.
114-
fn id(&self) -> ExtId;
104+
fn ext_id(&self) -> ExtId;
115105
/// Returns a reference to the storage [`DType`].
116-
fn storage_dtype(&self) -> &DType;
106+
fn ext_storage_dtype(&self) -> &DType;
117107
/// Returns the metadata as a trait object for downcasting.
118108
fn metadata_any(&self) -> &dyn Any;
119109
/// Formats the metadata using [`Debug`].
@@ -135,16 +125,16 @@ pub(super) trait DynExtDType: 'static + Send + Sync + super::sealed::Sealed {
135125
-> fmt::Result;
136126
}
137127

138-
impl<V: ExtVTable> DynExtDType for ExtDTypeInner<V> {
128+
impl<V: ExtVTable> DynExtDType for ExtDType<V> {
139129
fn as_any(&self) -> &dyn Any {
140130
self
141131
}
142132

143-
fn id(&self) -> ExtId {
133+
fn ext_id(&self) -> ExtId {
144134
self.vtable.id()
145135
}
146136

147-
fn storage_dtype(&self) -> &DType {
137+
fn ext_storage_dtype(&self) -> &DType {
148138
&self.storage_dtype
149139
}
150140

@@ -186,19 +176,15 @@ impl<V: ExtVTable> DynExtDType for ExtDTypeInner<V> {
186176
}
187177

188178
fn value_validate(&self, storage_value: &ScalarValue) -> VortexResult<()> {
189-
self.vtable
190-
.validate_scalar_value(&self.metadata, &self.storage_dtype, storage_value)
179+
self.vtable.validate_scalar_value(self, storage_value)
191180
}
192181

193182
fn value_display(
194183
&self,
195184
f: &mut fmt::Formatter<'_>,
196185
storage_value: &ScalarValue,
197186
) -> fmt::Result {
198-
match self
199-
.vtable
200-
.unpack_native(&self.metadata, &self.storage_dtype, storage_value)
201-
{
187+
match self.vtable.unpack_native(self, storage_value) {
202188
Ok(native) => fmt::Display::fmt(&native, f),
203189
Err(_) => write!(
204190
f,

‎vortex-array/src/dtype/extension/vtable.rs‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use std::hash::Hash;
77

88
use vortex_error::VortexResult;
99

10-
use crate::dtype::DType;
10+
use crate::dtype::extension::ExtDType;
1111
use crate::dtype::extension::ExtId;
1212
use crate::scalar::ScalarValue;
1313

@@ -36,7 +36,7 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
3636
fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<Self::Metadata>;
3737

3838
/// Validate that the given storage type is compatible with this extension type.
39-
fn validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &DType) -> VortexResult<()>;
39+
fn validate_dtype(&self, ext_dtype: &ExtDType<Self>) -> VortexResult<()>;
4040

4141
// Methods related to the extension scalar values.
4242

@@ -49,27 +49,24 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash {
4949
/// Returns an error if the storage [`ScalarValue`] is not compatible with the extension type.
5050
fn validate_scalar_value(
5151
&self,
52-
metadata: &Self::Metadata,
53-
storage_dtype: &DType,
52+
ext_dtype: &ExtDType<Self>,
5453
storage_value: &ScalarValue,
5554
) -> VortexResult<()> {
56-
self.unpack_native(metadata, storage_dtype, storage_value)
57-
.map(|_| ())
55+
self.unpack_native(ext_dtype, storage_value).map(|_| ())
5856
}
5957

6058
/// Validate and unpack a native value from the storage [`ScalarValue`].
6159
///
6260
/// Note that [`ExtVTable::validate_dtype()`] is always called first to validate the storage
63-
/// [`DType`], and the [`Scalar`](crate::scalar::Scalar) implementation will verify that the
64-
/// storage value is compatible with the storage dtype on construction.
61+
/// [`crate::dtype::DType`], and the [`Scalar`](crate::scalar::Scalar) implementation will
62+
/// verify that the storage value is compatible with the storage dtype on construction.
6563
///
6664
/// # Errors
6765
///
6866
/// Returns an error if the storage [`ScalarValue`] is not compatible with the extension type.
6967
fn unpack_native<'a>(
7068
&self,
71-
metadata: &'a Self::Metadata,
72-
storage_dtype: &'a DType,
69+
ext_dtype: &'a ExtDType<Self>,
7370
storage_value: &'a ScalarValue,
7471
) -> VortexResult<Self::NativeValue<'a>>;
7572
}

‎vortex-array/src/extension/datetime/date.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -91,12 +91,13 @@ impl ExtVTable for Date {
9191
TimeUnit::try_from(tag)
9292
}
9393

94-
fn validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &DType) -> VortexResult<()> {
94+
fn validate_dtype(&self, ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
95+
let metadata = ext_dtype.metadata();
9596
let ptype = date_ptype(metadata)
9697
.ok_or_else(|| vortex_err!("Date type does not support time unit {}", metadata))?;
9798

9899
vortex_ensure!(
99-
storage_dtype.as_ptype() == ptype,
100+
ext_dtype.storage_dtype().as_ptype() == ptype,
100101
"Date storage dtype for {} must be {}",
101102
metadata,
102103
ptype
@@ -107,10 +108,10 @@ impl ExtVTable for Date {
107108

108109
fn unpack_native(
109110
&self,
110-
metadata: &Self::Metadata,
111-
_storage_dtype: &DType,
111+
ext_dtype: &ExtDType<Self>,
112112
storage_value: &ScalarValue,
113113
) -> VortexResult<Self::NativeValue<'_>> {
114+
let metadata = ext_dtype.metadata();
114115
match metadata {
115116
TimeUnit::Milliseconds => Ok(DateValue::Milliseconds(
116117
storage_value.as_primitive().cast::<i64>()?,

0 commit comments

Comments
 (0)