Move PType into vortex-dtype (#274)

Next step is to change to `DType::Primitive(PType)`. This aligns our flattened arrays, our scalars and our dtypes to have essentially the same enum variants. We also vendor the half crate into vortex_dtype.
spiraldb · Apr 30, 2024 · 9d7b929 · 9d7b929
1 parent e845b4c
commit 9d7b929
Show file tree

Hide file tree

Showing 63 changed files with 179 additions and 202 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/vortex-alp/src/compress.rs b/vortex-alp/src/compress.rs
@@ -2,10 +2,10 @@ use itertools::Itertools;
 use vortex::array::primitive::PrimitiveArray;
 use vortex::array::sparse::{Sparse, SparseArray};
 use vortex::compress::{CompressConfig, CompressCtx, EncodingCompression};
-use vortex::ptype::{NativePType, PType};
 use vortex::scalar::Scalar;
 use vortex::validity::Validity;
 use vortex::{Array, ArrayDType, ArrayDef, AsArray, IntoArray, OwnedArray};
+use vortex_dtype::{NativePType, PType};
 use vortex_error::{vortex_bail, vortex_err, VortexResult};
 
 use crate::alp::ALPFloat;
@@ -16,8 +16,8 @@ use crate::{Exponents, OwnedALPArray};
 macro_rules! match_each_alp_float_ptype {
     ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
         macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
+        use vortex_dtype::PType;
         use vortex_error::vortex_err;
-        use vortex::ptype::PType;
         let ptype = $self;
         match ptype {
             PType::F32 => Ok(__with__! { f32 }),

diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml
@@ -25,7 +25,6 @@ arrow-buffer = { workspace = true }
 arrow-schema = { workspace = true }
 flatbuffers = { workspace = true }
 flexbuffers = { workspace = true }
-half = { workspace = true }
 humansize = { workspace = true }
 itertools = { workspace = true }
 leb128 = { workspace = true }

diff --git a/vortex-array/src/array/bool/compute/take.rs b/vortex-array/src/array/bool/compute/take.rs
@@ -1,10 +1,10 @@
 use arrow_buffer::BooleanBuffer;
 use num_traits::AsPrimitive;
+use vortex_dtype::match_each_integer_ptype;
 use vortex_error::VortexResult;
 
 use crate::array::bool::BoolArray;
 use crate::compute::take::TakeFn;
-use crate::match_each_integer_ptype;
 use crate::AsArray;
 use crate::IntoArray;
 use crate::{Array, OwnedArray};

diff --git a/vortex-array/src/array/chunked/compute/take.rs b/vortex-array/src/array/chunked/compute/take.rs
@@ -1,9 +1,9 @@
+use vortex_dtype::PType;
 use vortex_error::VortexResult;
 
 use crate::array::chunked::ChunkedArray;
 use crate::compute::cast::cast;
 use crate::compute::take::{take, TakeFn};
-use crate::ptype::PType;
 use crate::{Array, IntoArray, OwnedArray, ToArray, ToStatic};
 use crate::{ArrayDType, ArrayTrait};
 

diff --git a/vortex-array/src/array/chunked/mod.rs b/vortex-array/src/array/chunked/mod.rs
@@ -145,10 +145,10 @@ impl EncodingCompression for ChunkedEncoding {}
 
 #[cfg(test)]
 mod test {
+    use vortex_dtype::NativePType;
     use vortex_dtype::{DType, IntWidth, Nullability, Signedness};
 
     use crate::array::chunked::{ChunkedArray, OwnedChunkedArray};
-    use crate::ptype::NativePType;
     use crate::{Array, IntoArray};
 
     #[allow(dead_code)]

diff --git a/vortex-array/src/array/composite/array.rs b/vortex-array/src/array/composite/array.rs
@@ -6,7 +6,6 @@ use vortex_flatbuffers::{FlatBufferToBytes, ReadFlatBuffer};
 
 use crate::array::composite::{find_extension, CompositeExtensionRef, TypedCompositeArray};
 use crate::compute::ArrayCompute;
-use crate::scalar::AsBytes;
 use crate::stats::ArrayStatisticsCompute;
 use crate::validity::{ArrayValidity, LogicalValidity};
 use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
@@ -141,7 +140,7 @@ impl CompositeArray<'_> {
 impl<'a> CompositeArray<'a> {
     pub fn as_typed<M: UnderlyingMetadata>(&'a self) -> VortexResult<TypedCompositeArray<'a, M>> {
         Ok(TypedCompositeArray::new(
-            M::try_deserialize_metadata(Some(self.underlying_metadata().as_bytes()))?,
+            M::try_deserialize_metadata(Some(self.underlying_metadata()))?,
             self.underlying().clone(),
         ))
     }

diff --git a/vortex-array/src/array/constant/flatten.rs b/vortex-array/src/array/constant/flatten.rs
@@ -1,12 +1,12 @@
-use vortex_dtype::Nullability;
+use vortex_dtype::{match_each_native_ptype, Nullability};
 use vortex_error::VortexResult;
 
 use crate::array::bool::BoolArray;
 use crate::array::constant::ConstantArray;
 use crate::array::primitive::PrimitiveArray;
 use crate::scalar::Scalar;
 use crate::validity::Validity;
-use crate::{match_each_native_ptype, ArrayDType, ArrayTrait};
+use crate::{ArrayDType, ArrayTrait};
 use crate::{ArrayFlatten, Flattened};
 
 impl ArrayFlatten for ConstantArray<'_> {

diff --git a/vortex-array/src/array/datetime/localdatetime.rs b/vortex-array/src/array/datetime/localdatetime.rs
@@ -5,13 +5,13 @@ use arrow_array::{
     TimestampNanosecondArray, TimestampSecondArray,
 };
 use serde::{Deserialize, Serialize};
+use vortex_dtype::PType;
 use vortex_error::VortexResult;
 
 use crate::array::datetime::TimeUnit;
 use crate::compute::as_arrow::AsArrowArray;
 use crate::compute::cast::cast;
 use crate::impl_composite;
-use crate::ptype::PType;
 use crate::validity::ArrayValidity;
 
 impl_composite!("vortex.localdatetime", LocalDateTime);

diff --git a/vortex-array/src/array/primitive/accessor.rs b/vortex-array/src/array/primitive/accessor.rs
@@ -1,8 +1,8 @@
+use vortex_dtype::NativePType;
 use vortex_error::VortexResult;
 
 use crate::accessor::ArrayAccessor;
 use crate::array::primitive::PrimitiveArray;
-use crate::ptype::NativePType;
 use crate::validity::ArrayValidity;
 
 impl<T: NativePType> ArrayAccessor<T> for PrimitiveArray<'_> {

diff --git a/vortex-array/src/array/primitive/compute/as_arrow.rs b/vortex-array/src/array/primitive/compute/as_arrow.rs
@@ -4,11 +4,11 @@ use arrow_array::{
     ArrayRef as ArrowArrayRef, ArrowPrimitiveType, PrimitiveArray as ArrowPrimitiveArray,
 };
 use arrow_buffer::ScalarBuffer;
+use vortex_dtype::PType;
 use vortex_error::VortexResult;
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::as_arrow::AsArrowArray;
-use crate::ptype::PType;
 use crate::validity::ArrayValidity;
 use crate::ArrayTrait;
 

diff --git a/vortex-array/src/array/primitive/compute/as_contiguous.rs b/vortex-array/src/array/primitive/compute/as_contiguous.rs
@@ -1,10 +1,11 @@
 use arrow_buffer::{MutableBuffer, ScalarBuffer};
+use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::as_contiguous::AsContiguousFn;
 use crate::validity::Validity;
-use crate::{match_each_native_ptype, ArrayDType};
+use crate::ArrayDType;
 use crate::{Array, IntoArray, OwnedArray};
 
 impl AsContiguousFn for PrimitiveArray<'_> {

diff --git a/vortex-array/src/array/primitive/compute/cast.rs b/vortex-array/src/array/primitive/compute/cast.rs
@@ -1,11 +1,11 @@
-use vortex_dtype::DType;
+use vortex_dtype::{match_each_native_ptype, DType};
+use vortex_dtype::{NativePType, PType};
 use vortex_error::{vortex_err, VortexResult};
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::cast::CastFn;
-use crate::ptype::{NativePType, PType};
 use crate::validity::Validity;
-use crate::{match_each_native_ptype, ArrayDType};
+use crate::ArrayDType;
 use crate::{IntoArray, OwnedArray};
 
 impl CastFn for PrimitiveArray<'_> {
@@ -49,9 +49,9 @@ fn cast<T: NativePType>(array: &PrimitiveArray) -> VortexResult<Vec<T>> {
 
 #[cfg(test)]
 mod test {
+    use vortex_dtype::PType;
     use vortex_error::VortexError;
 
-    use crate::ptype::PType;
     use crate::{compute, IntoArray};
 
     #[test]

diff --git a/vortex-array/src/array/primitive/compute/fill.rs b/vortex-array/src/array/primitive/compute/fill.rs
@@ -1,8 +1,8 @@
+use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::fill::FillForwardFn;
-use crate::match_each_native_ptype;
 use crate::validity::ArrayValidity;
 use crate::{IntoArray, OwnedArray, ToArrayData};
 

diff --git a/vortex-array/src/array/primitive/compute/scalar_at.rs b/vortex-array/src/array/primitive/compute/scalar_at.rs
@@ -1,8 +1,8 @@
+use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::scalar_at::ScalarAtFn;
-use crate::match_each_native_ptype;
 use crate::scalar::PrimitiveScalar;
 use crate::scalar::Scalar;
 use crate::validity::ArrayValidity;

diff --git a/vortex-array/src/array/primitive/compute/search_sorted.rs b/vortex-array/src/array/primitive/compute/search_sorted.rs
@@ -1,9 +1,9 @@
+use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::search_sorted::{SearchResult, SearchSorted};
 use crate::compute::search_sorted::{SearchSortedFn, SearchSortedSide};
-use crate::match_each_native_ptype;
 use crate::scalar::Scalar;
 
 impl SearchSortedFn for PrimitiveArray<'_> {

diff --git a/vortex-array/src/array/primitive/compute/slice.rs b/vortex-array/src/array/primitive/compute/slice.rs
@@ -1,8 +1,8 @@
+use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::slice::SliceFn;
-use crate::match_each_native_ptype;
 use crate::IntoArray;
 use crate::OwnedArray;
 

diff --git a/vortex-array/src/array/primitive/compute/take.rs b/vortex-array/src/array/primitive/compute/take.rs
@@ -1,11 +1,11 @@
 use num_traits::PrimInt;
+use vortex_dtype::NativePType;
+use vortex_dtype::{match_each_integer_ptype, match_each_native_ptype};
 use vortex_error::VortexResult;
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::take::TakeFn;
-use crate::ptype::NativePType;
 use crate::IntoArray;
-use crate::{match_each_integer_ptype, match_each_native_ptype};
 use crate::{Array, OwnedArray};
 
 impl TakeFn for PrimitiveArray<'_> {

diff --git a/vortex-array/src/array/primitive/mod.rs b/vortex-array/src/array/primitive/mod.rs
@@ -2,14 +2,14 @@ use arrow_buffer::{ArrowNativeType, ScalarBuffer};
 use itertools::Itertools;
 use num_traits::AsPrimitive;
 use serde::{Deserialize, Serialize};
+use vortex_dtype::{match_each_native_ptype, NativePType, PType};
 use vortex_error::{vortex_bail, VortexResult};
 
 use crate::buffer::Buffer;
-use crate::ptype::{NativePType, PType};
 use crate::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
 use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
+use crate::ArrayFlatten;
 use crate::{impl_encoding, ArrayDType, OwnedArray};
-use crate::{match_each_native_ptype, ArrayFlatten};
 
 mod accessor;
 mod compute;
@@ -23,6 +23,7 @@ pub struct PrimitiveMetadata {
 }
 
 impl PrimitiveArray<'_> {
+    // TODO(ngates): remove the Arrow types from this API.
     pub fn try_new<T: NativePType + ArrowNativeType>(
         buffer: ScalarBuffer<T>,
         validity: Validity,
@@ -40,11 +41,15 @@ impl PrimitiveArray<'_> {
         })
     }
 
-    pub fn from_vec<T: NativePType + ArrowNativeType>(values: Vec<T>, validity: Validity) -> Self {
-        Self::try_new(ScalarBuffer::from(values), validity).unwrap()
+    pub fn from_vec<T: NativePType>(values: Vec<T>, validity: Validity) -> Self {
+        match_each_native_ptype!(T::PTYPE, |$P| {
+            Self::try_new(ScalarBuffer::<$P>::from(
+                unsafe { std::mem::transmute::<Vec<T>, Vec<$P>>(values) }
+            ), validity).unwrap()
+        })
     }
 
-    pub fn from_nullable_vec<T: NativePType + ArrowNativeType>(values: Vec<Option<T>>) -> Self {
+    pub fn from_nullable_vec<T: NativePType>(values: Vec<Option<T>>) -> Self {
         let elems: Vec<T> = values.iter().map(|v| v.unwrap_or_default()).collect();
         let validity = Validity::from(values.iter().map(|v| v.is_some()).collect::<Vec<_>>());
         Self::from_vec(elems, validity)
@@ -65,7 +70,8 @@ impl PrimitiveArray<'_> {
         self.array().buffer().expect("missing buffer")
     }
 
-    pub fn scalar_buffer<T: NativePType>(&self) -> ScalarBuffer<T> {
+    // TODO(ngates): deprecated, remove this.
+    pub fn scalar_buffer<T: NativePType + ArrowNativeType>(&self) -> ScalarBuffer<T> {
         assert_eq!(
             T::PTYPE,
             self.ptype(),
@@ -126,7 +132,7 @@ impl PrimitiveArray<'_> {
         for (idx, value) in positions.iter().zip_eq(values.iter()) {
             own_values[(*idx).as_()] = *value;
         }
-        Self::try_new(ScalarBuffer::from(own_values), validity)
+        Ok(Self::from_vec(own_values, validity))
     }
 }