merge

vortex/src/array/primitive/compute/cast.rs

@@ -1,17 +1,3 @@
-// (c) Copyright 2024 Fulcrum Technologies, Inc. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
 use crate::array::primitive::PrimitiveArray;
 use crate::array::CloneOptionalArray;
 use crate::compute::cast::CastPrimitiveFn;
@@ -35,22 +21,25 @@ impl CastPrimitiveFn for PrimitiveArray {
 }
 
 fn cast<T: NativePType>(array: &PrimitiveArray) -> VortexResult<Vec<T>> {
-    array
-        .typed_data::<u32>()
-        .iter()
-        // TODO(ngates): allow configurable checked/unchecked casting
-        .map(|v| {
-            T::from(*v).ok_or_else(|| {
-                VortexError::ComputeError(format!("Failed to cast {} to {:?}", v, T::PTYPE).into())
+    match_each_native_ptype!(array.ptype(), |$E| {
+        array
+            .typed_data::<$E>()
+            .iter()
+            // TODO(ngates): allow configurable checked/unchecked casting
+            .map(|v| {
+                T::from(*v).ok_or_else(|| {
+                    VortexError::ComputeError(format!("Failed to cast {} to {:?}", v, T::PTYPE).into())
+                })
             })
-        })
-        .collect()
+            .collect()
+    })
 }
 
 #[cfg(test)]
 mod test {
     use crate::array::primitive::PrimitiveArray;
     use crate::compute;
+    use crate::error::VortexError;
     use crate::ptype::PType;
 
     #[test]
@@ -70,7 +59,11 @@ mod test {
     #[test]
     fn cast_i32_u32() {
         let arr = PrimitiveArray::from_vec(vec![-1i32]);
-        let u8arr = compute::cast::cast_primitive(&arr, &PType::U32).unwrap();
-        assert_eq!(u8arr.typed_data::<f32>(), vec![0.0f32, 10., 200.]);
+        assert_eq!(
+            compute::cast::cast_primitive(&arr, &PType::U32)
+                .err()
+                .unwrap(),
+            VortexError::ComputeError("Failed to cast -1 to U32".into(),)
+        )
     }
 }

vortex/src/array/primitive/compute/mod.rs

@@ -1,20 +1,7 @@
-// (c) Copyright 2024 Fulcrum Technologies, Inc. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::cast::CastPrimitiveFn;
 use crate::compute::patch::PatchFn;
+use crate::compute::scalar_at::ScalarAtFn;
 use crate::compute::ArrayCompute;
 
 mod cast;
@@ -29,4 +16,8 @@ impl ArrayCompute for PrimitiveArray {
     fn patch(&self) -> Option<&dyn PatchFn> {
         Some(self)
     }
+
+    fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
+        Some(self)
+    }
 }

vortex/src/array/primitive/compute/patch.rs

@@ -1,17 +1,3 @@
-// (c) Copyright 2024 Fulcrum Technologies, Inc. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
 use itertools::Itertools;
 
 use crate::array::downcast::DowncastArrayBuiltin;

vortex/src/array/primitive/compute/scalar_at.rs

@@ -1,3 +1,10 @@
+use crate::array::primitive::PrimitiveArray;
+use crate::array::Array;
+use crate::compute::scalar_at::ScalarAtFn;
+use crate::error::VortexResult;
+use crate::match_each_native_ptype;
+use crate::scalar::{NullableScalar, Scalar};
+
 impl ScalarAtFn for PrimitiveArray {
     fn scalar_at(&self, index: usize) -> VortexResult<Box<dyn Scalar>> {
         if self.is_valid(index) {

vortex/src/array/primitive/mod.rs

@@ -21,7 +21,6 @@ use crate::array::{
 use crate::arrow::CombineChunks;
 use crate::compress::EncodingCompression;
 use crate::compute::scalar_at::scalar_at;
-use crate::compute::ArrayCompute;
 use crate::dtype::DType;
 use crate::error::VortexResult;
 use crate::formatter::{ArrayDisplay, ArrayFormatter};

vortex/src/compute/cast.rs

@@ -1,22 +1,7 @@
-// (c) Copyright 2024 Fulcrum Technologies, Inc. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
 use crate::array::primitive::PrimitiveArray;
 use crate::array::Array;
 use crate::error::{VortexError, VortexResult};
 use crate::ptype::PType;
-use crate::scalar::Scalar;
 
 pub trait CastPrimitiveFn {
     fn cast_primitive(&self, ptype: &PType) -> VortexResult<PrimitiveArray>;

vortex/src/compute/mod.rs

@@ -1,17 +1,3 @@
-// (c) Copyright 2024 Fulcrum Technologies, Inc. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
 use cast::CastPrimitiveFn;
 use patch::PatchFn;
 use scalar_at::ScalarAtFn;

vortex/src/compute/patch.rs

@@ -1,17 +1,3 @@
-// (c) Copyright 2024 Fulcrum Technologies, Inc. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
 use crate::array::{Array, ArrayRef};
 use crate::error::{VortexError, VortexResult};
 
@@ -30,8 +16,7 @@ pub fn patch(array: &dyn Array, patch: &dyn Array) -> VortexResult<ArrayRef> {
     // TODO(ngates): check the dtype matches
 
     array
-        .compute()
-        .and_then(|c| c.patch())
+        .patch()
         .map(|t| t.patch(patch))
         .unwrap_or_else(|| Err(VortexError::NotImplemented("take", array.encoding().id())))
 }