Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
[cast] make CastOptions private and introduce wrapping_dictionary_to_…
Browse files Browse the repository at this point in the history
…dictionary_keys wrapping_dictionary_to_dictionary_values
  • Loading branch information
sundy-li committed Aug 7, 2021
1 parent 591502b commit 647fddb
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 5 deletions.
38 changes: 35 additions & 3 deletions src/compute/cast/dictionary_to.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use super::{primitive_to_primitive, CastOptions};
use super::{primitive_as_primitive, primitive_to_primitive, CastOptions};
use crate::{
array::{Array, DictionaryArray, DictionaryKey, PrimitiveArray},
compute::{cast::cast_with_options, take::take},
Expand Down Expand Up @@ -28,12 +28,24 @@ macro_rules! key_cast {
pub fn dictionary_to_dictionary_values<K: DictionaryKey>(
from: &DictionaryArray<K>,
values_type: &DataType,
options: CastOptions,
) -> Result<DictionaryArray<K>> {
let keys = from.keys();
let values = from.values();

let values = cast_with_options(values.as_ref(), values_type, options)?.into();
let values = cast_with_options(values.as_ref(), values_type, CastOptions::default())?.into();
Ok(DictionaryArray::from_data(keys.clone(), values))
}

/// Similar to dictionary_to_dictionary_values, but overflowing cast is wrapped
pub fn wrapping_dictionary_to_dictionary_values<K: DictionaryKey>(
from: &DictionaryArray<K>,
values_type: &DataType,
) -> Result<DictionaryArray<K>> {
let keys = from.keys();
let values = from.values();

let values =
cast_with_options(values.as_ref(), values_type, CastOptions { wrapped: true })?.into();
Ok(DictionaryArray::from_data(keys.clone(), values))
}

Expand Down Expand Up @@ -61,6 +73,26 @@ where
}
}

/// Similar to dictionary_to_dictionary_keys, but overflowing cast is wrapped
pub fn wrapping_dictionary_to_dictionary_keys<K1, K2>(
from: &DictionaryArray<K1>,
) -> Result<DictionaryArray<K2>>
where
K1: DictionaryKey + num::traits::AsPrimitive<K2>,
K2: DictionaryKey,
{
let keys = from.keys();
let values = from.values();

let casted_keys = primitive_as_primitive::<K1, K2>(keys, &K2::DATA_TYPE);

if casted_keys.null_count() > keys.null_count() {
Err(ArrowError::KeyOverflowError)
} else {
Ok(DictionaryArray::from_data(casted_keys, values.clone()))
}
}

pub(super) fn dictionary_cast_dyn<K: DictionaryKey>(
array: &dyn Array,
to_type: &DataType,
Expand Down
4 changes: 2 additions & 2 deletions src/compute/cast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ pub use utf8_to::*;

/// options defining how Cast kernels behave
#[derive(Clone, Copy, Debug)]
pub struct CastOptions {
struct CastOptions {
/// default to false
/// whether an overflowing cast should be converted to `None` (default), or be wrapped (i.e. `256i16 as u8 = 0` vectorized).
/// Settings this to `true` is 5-6x faster for numeric types.
Expand Down Expand Up @@ -341,7 +341,7 @@ pub fn wrapping_cast(array: &dyn Array, to_type: &DataType) -> Result<Box<dyn Ar
}

#[inline]
pub fn cast_with_options(
fn cast_with_options(
array: &dyn Array,
to_type: &DataType,
options: CastOptions,
Expand Down

0 comments on commit 647fddb

Please sign in to comment.