Skip to content

Commit

Permalink
Allow mixed integers
Browse files Browse the repository at this point in the history
  • Loading branch information
@stinodego
stinodego committed Feb 19, 2024
1 parent 8319ba9 commit d8c6051
Showing 1 changed file with 90 additions and 20 deletions.
110 changes: 90 additions & 20 deletions crates/polars-core/src/series/any_value.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,6 @@ use std::fmt::Write;
use crate::prelude::*;
use crate::utils::get_supertype;

fn any_values_to_primitive<T: PolarsNumericType>(avs: &[AnyValue]) -> ChunkedArray<T> {
avs.iter()
.map(|av| av.extract::<T::Native>())
.collect_trusted()
}

#[cfg(feature = "dtype-decimal")]
fn any_values_to_decimal(
avs: &[AnyValue],
Expand Down Expand Up @@ -225,36 +219,36 @@ impl Series {
) -> PolarsResult<Series> {
let mut s = match dtype {
#[cfg(feature = "dtype-i8")]
DataType::Int8 => any_values_to_primitive::<Int8Type>(av).into_series(),
DataType::Int8 => any_values_to_integer::<Int8Type>(av, strict)?.into_series(),
#[cfg(feature = "dtype-i16")]
DataType::Int16 => any_values_to_primitive::<Int16Type>(av).into_series(),
DataType::Int32 => any_values_to_primitive::<Int32Type>(av).into_series(),
DataType::Int64 => any_values_to_primitive::<Int64Type>(av).into_series(),
DataType::Int16 => any_values_to_integer::<Int16Type>(av, strict)?.into_series(),
DataType::Int32 => any_values_to_integer::<Int32Type>(av, strict)?.into_series(),
DataType::Int64 => any_values_to_integer::<Int64Type>(av, strict)?.into_series(),
#[cfg(feature = "dtype-u8")]
DataType::UInt8 => any_values_to_primitive::<UInt8Type>(av).into_series(),
DataType::UInt8 => any_values_to_integer::<UInt8Type>(av, strict)?.into_series(),
#[cfg(feature = "dtype-u16")]
DataType::UInt16 => any_values_to_primitive::<UInt16Type>(av).into_series(),
DataType::UInt32 => any_values_to_primitive::<UInt32Type>(av).into_series(),
DataType::UInt64 => any_values_to_primitive::<UInt64Type>(av).into_series(),
DataType::Float32 => any_values_to_primitive::<Float32Type>(av).into_series(),
DataType::Float64 => any_values_to_primitive::<Float64Type>(av).into_series(),
DataType::UInt16 => any_values_to_integer::<UInt16Type>(av, strict)?.into_series(),
DataType::UInt32 => any_values_to_integer::<UInt32Type>(av, strict)?.into_series(),
DataType::UInt64 => any_values_to_integer::<UInt64Type>(av, strict)?.into_series(),
DataType::Float32 => any_values_to_f32(av, strict)?.into_series(),
DataType::Float64 => any_values_to_f64(av, strict)?.into_series(),
DataType::String => any_values_to_string(av, strict)?.into_series(),
DataType::Binary => any_values_to_binary(av, strict)?.into_series(),
DataType::Boolean => any_values_to_bool(av, strict)?.into_series(),
#[cfg(feature = "dtype-date")]
DataType::Date => any_values_to_primitive::<Int32Type>(av)
DataType::Date => any_values_to_primitive_nonstrict::<Int32Type>(av)
.into_date()
.into_series(),
#[cfg(feature = "dtype-datetime")]
DataType::Datetime(tu, tz) => any_values_to_primitive::<Int64Type>(av)
DataType::Datetime(tu, tz) => any_values_to_primitive_nonstrict::<Int64Type>(av)
.into_datetime(*tu, (*tz).clone())
.into_series(),
#[cfg(feature = "dtype-time")]
DataType::Time => any_values_to_primitive::<Int64Type>(av)
DataType::Time => any_values_to_primitive_nonstrict::<Int64Type>(av)
.into_time()
.into_series(),
#[cfg(feature = "dtype-duration")]
DataType::Duration(tu) => any_values_to_primitive::<Int64Type>(av)
DataType::Duration(tu) => any_values_to_primitive_nonstrict::<Int64Type>(av)
.into_duration(*tu)
.into_series(),
#[cfg(feature = "dtype-decimal")]
Expand Down Expand Up @@ -460,6 +454,82 @@ impl Series {
}
}

fn any_values_to_primitive_nonstrict<T: PolarsNumericType>(values: &[AnyValue]) -> ChunkedArray<T> {
values
.iter()
.map(|av| av.extract::<T::Native>())
.collect_trusted()
}

fn any_values_to_integer<T: PolarsIntegerType>(
values: &[AnyValue],
strict: bool,
) -> PolarsResult<ChunkedArray<T>> {
fn any_values_to_integer_strict<T: PolarsIntegerType>(
values: &[AnyValue],
) -> PolarsResult<ChunkedArray<T>> {
let mut builder = PrimitiveChunkedBuilder::<T>::new("", values.len());
for av in values {
match av {
av if av.is_integer() => {
let opt_val = av.extract::<T::Native>();
let val = match opt_val {
Some(v) => v,
None => return Err(invalid_value_error(&T::get_dtype(), av)),
};
builder.append_value(val)
},
AnyValue::Null => builder.append_null(),
av => return Err(invalid_value_error(&T::get_dtype(), av)),
}
}
Ok(builder.finish())
}
if strict {
any_values_to_integer_strict::<T>(values)
} else {
Ok(any_values_to_primitive_nonstrict::<T>(values))
}
}

fn any_values_to_f32(values: &[AnyValue], strict: bool) -> PolarsResult<Float32Chunked> {
fn any_values_to_f32_strict(values: &[AnyValue]) -> PolarsResult<Float32Chunked> {
let mut builder = PrimitiveChunkedBuilder::<Float32Type>::new("", values.len());
for av in values {
match av {
AnyValue::Float32(i) => builder.append_value(*i),
AnyValue::Null => builder.append_null(),
av => return Err(invalid_value_error(&DataType::Float32, av)),
}
}
Ok(builder.finish())
}
if strict {
any_values_to_f32_strict(values)
} else {
Ok(any_values_to_primitive_nonstrict::<Float32Type>(values))
}
}
fn any_values_to_f64(values: &[AnyValue], strict: bool) -> PolarsResult<Float64Chunked> {
fn any_values_to_f64_strict(values: &[AnyValue]) -> PolarsResult<Float64Chunked> {
let mut builder = PrimitiveChunkedBuilder::<Float64Type>::new("", values.len());
for av in values {
match av {
AnyValue::Float64(i) => builder.append_value(*i),
AnyValue::Float32(i) => builder.append_value(*i as f64),
AnyValue::Null => builder.append_null(),
av => return Err(invalid_value_error(&DataType::Float64, av)),
}
}
Ok(builder.finish())
}
if strict {
any_values_to_f64_strict(values)
} else {
Ok(any_values_to_primitive_nonstrict::<Float64Type>(values))
}
}

fn any_values_to_bool(values: &[AnyValue], strict: bool) -> PolarsResult<BooleanChunked> {
if strict {
any_values_to_bool_strict(values)
Expand Down

0 comments on commit d8c6051

Please sign in to comment.