Skip to content

Commit

Permalink
chore: hoist boolean -> string cast (#14122)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored Jan 31, 2024
1 parent ed672fa commit 65eb9ef
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 30 deletions.
28 changes: 15 additions & 13 deletions crates/polars-arrow/src/compute/cast/boolean_to.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use polars_error::PolarsResult;

use crate::array::{Array, BinaryArray, BooleanArray, PrimitiveArray, Utf8Array};
use crate::offset::Offset;
use super::{ArrayFromIter, BinaryViewArray, Utf8ViewArray};
use crate::array::{Array, BooleanArray, PrimitiveArray};
use crate::types::NativeType;

pub(super) fn boolean_to_primitive_dyn<T>(array: &dyn Array) -> PolarsResult<Box<dyn Array>>
Expand All @@ -26,24 +26,26 @@ where
PrimitiveArray::<T>::new(T::PRIMITIVE.into(), values.into(), from.validity().cloned())
}

/// Casts the [`BooleanArray`] to a [`Utf8Array`], casting trues to `"1"` and falses to `"0"`
pub fn boolean_to_utf8<O: Offset>(from: &BooleanArray) -> Utf8Array<O> {
let iter = from.values().iter().map(|x| if x { "1" } else { "0" });
Utf8Array::from_trusted_len_values_iter(iter)
pub fn boolean_to_utf8view(from: &BooleanArray) -> Utf8ViewArray {
unsafe { boolean_to_binaryview(from).to_utf8view_unchecked() }
}

pub(super) fn boolean_to_utf8_dyn<O: Offset>(array: &dyn Array) -> PolarsResult<Box<dyn Array>> {
pub(super) fn boolean_to_utf8view_dyn(array: &dyn Array) -> PolarsResult<Box<dyn Array>> {
let array = array.as_any().downcast_ref().unwrap();
Ok(Box::new(boolean_to_utf8::<O>(array)))
Ok(boolean_to_utf8view(array).boxed())
}

/// Casts the [`BooleanArray`] to a [`BinaryArray`], casting trues to `"1"` and falses to `"0"`
pub fn boolean_to_binary<O: Offset>(from: &BooleanArray) -> BinaryArray<O> {
let iter = from.values().iter().map(|x| if x { b"1" } else { b"0" });
BinaryArray::from_trusted_len_values_iter(iter)
pub fn boolean_to_binaryview(from: &BooleanArray) -> BinaryViewArray {
let iter = from.iter().map(|opt_b| match opt_b {
Some(true) => Some("true".as_bytes()),
Some(false) => Some("false".as_bytes()),
None => None,
});
BinaryViewArray::arr_from_iter_trusted(iter)
}

pub(super) fn boolean_to_binary_dyn<O: Offset>(array: &dyn Array) -> PolarsResult<Box<dyn Array>> {
pub(super) fn boolean_to_binaryview_dyn(array: &dyn Array) -> PolarsResult<Box<dyn Array>> {
let array = array.as_any().downcast_ref().unwrap();
Ok(Box::new(boolean_to_binary::<O>(array)))
Ok(boolean_to_binaryview(array).boxed())
}
4 changes: 2 additions & 2 deletions crates/polars-arrow/src/compute/cast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -462,8 +462,8 @@ pub fn cast(
Int64 => boolean_to_primitive_dyn::<i64>(array),
Float32 => boolean_to_primitive_dyn::<f32>(array),
Float64 => boolean_to_primitive_dyn::<f64>(array),
LargeUtf8 => boolean_to_utf8_dyn::<i64>(array),
LargeBinary => boolean_to_binary_dyn::<i64>(array),
Utf8View => boolean_to_utf8view_dyn(array),
BinaryView => boolean_to_binaryview_dyn(array),
_ => polars_bail!(InvalidOperation:
"casting from {from_type:?} to {to_type:?} not supported",
),
Expand Down
15 changes: 0 additions & 15 deletions crates/polars-core/src/chunked_array/cast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -360,24 +360,9 @@ impl ChunkCast for BinaryOffsetChunked {
}
}

fn boolean_to_string(ca: &BooleanChunked) -> StringChunked {
ca.into_iter()
.map(|opt_b| match opt_b {
Some(true) => Some("true"),
Some(false) => Some("false"),
None => None,
})
.collect()
}

impl ChunkCast for BooleanChunked {
fn cast(&self, data_type: &DataType) -> PolarsResult<Series> {
match data_type {
DataType::String => {
let mut ca = boolean_to_string(self);
ca.rename(self.name());
Ok(ca.into_series())
},
#[cfg(feature = "dtype-struct")]
DataType::Struct(fields) => cast_single_to_struct(self.name(), &self.chunks, fields),
_ => cast_impl(self.name(), &self.chunks, data_type),
Expand Down

0 comments on commit 65eb9ef

Please sign in to comment.