Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add dtype-i128 feature flag and Int128Type #6374

Merged
merged 6 commits into from
Feb 2, 2023
+147 −9
Merged

Add dtype-i128 feature flag and Int128Type #6374

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
296cf14
feat: add new dtype-i128 feature flag
plaflamme Jan 21, 2023
6c874de
feat: add DataType::Int128
plaflamme Jan 21, 2023
e910c22
feat: define ArrayArithmetics to abstract over basic vs decimal
plaflamme Jan 22, 2023
90681fd
Merge remote-tracking branch 'upstream/master' into dtype-i128
plaflamme Jan 31, 2023
b803b6d
tidy: fix clippy warning
plaflamme Jan 31, 2023
0d0a898
feat: add DataType::Decimal128
plaflamme Jan 31, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions polars/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,7 @@ dtype-full = [
"dtype-time",
"dtype-i8",
"dtype-i16",
"dtype-i128",
"dtype-u8",
"dtype-u16",
"dtype-categorical",
Expand Down Expand Up @@ -215,6 +216,7 @@ dtype-duration = [
dtype-time = ["polars-core/dtype-time", "polars-io/dtype-time", "polars-time/dtype-time", "polars-ops/dtype-time"]
dtype-i8 = ["polars-core/dtype-i8", "polars-lazy/dtype-i8", "polars-ops/dtype-i8"]
dtype-i16 = ["polars-core/dtype-i16", "polars-lazy/dtype-i16", "polars-ops/dtype-i16"]
dtype-i128 = ["polars-core/dtype-i128", "polars-lazy/dtype-i128", "polars-ops/dtype-i128"]
dtype-u8 = ["polars-core/dtype-u8", "polars-lazy/dtype-u8", "polars-ops/dtype-u8"]
dtype-u16 = ["polars-core/dtype-u16", "polars-lazy/dtype-u16", "polars-ops/dtype-u16"]
dtype-categorical = [
Expand Down
1 change: 1 addition & 0 deletions polars/polars-core/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ dtype-duration = ["temporal"]
dtype-time = ["temporal"]
dtype-i8 = []
dtype-i16 = []
dtype-i128 = []
dtype-u8 = []
dtype-u16 = []
dtype-categorical = []
Expand Down
121 changes: 113 additions & 8 deletions polars/polars-core/src/chunked_array/arithmetic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,94 @@ use std::borrow::Cow;
use std::ops::{Add, Div, Mul, Rem, Sub};

use arrow::array::PrimitiveArray;
use arrow::compute::arithmetics::basic;
use arrow::compute::arithmetics::{basic, decimal};
use arrow::compute::arity_assign;
use arrow::types::NativeType;
use num::{Num, NumCast, ToPrimitive};

use crate::prelude::*;
use crate::series::IsSorted;
use crate::utils::{align_chunks_binary, align_chunks_binary_owned};

pub trait ArrayArithmetics
where
Self: NativeType,
{
fn add(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self>;
fn sub(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self>;
fn mul(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self>;
fn div(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self>;
fn div_scalar(lhs: &PrimitiveArray<Self>, rhs: &Self) -> PrimitiveArray<Self>;
fn rem(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self>;
fn rem_scalar(lhs: &PrimitiveArray<Self>, rhs: &Self) -> PrimitiveArray<Self>;
}

macro_rules! native_array_arithmetics {
($ty: ty) => {
impl ArrayArithmetics for $ty
{
fn add(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
basic::add(lhs, rhs)
}
fn sub(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
basic::sub(lhs, rhs)
}
fn mul(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
basic::mul(lhs, rhs)
}
fn div(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
basic::div(lhs, rhs)
}
fn div_scalar(lhs: &PrimitiveArray<Self>, rhs: &Self) -> PrimitiveArray<Self> {
basic::div_scalar(lhs, rhs)
}
fn rem(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
basic::rem(lhs, rhs)
}
fn rem_scalar(lhs: &PrimitiveArray<Self>, rhs: &Self) -> PrimitiveArray<Self> {
basic::rem_scalar(lhs, rhs)
}
}
};
($($ty:ty),*) => {
$(native_array_arithmetics!($ty);)*
}
}

native_array_arithmetics!(u8, u16, u32, u64, i8, i16, i32, i64, f32, f64);

#[cfg(feature = "dtype-i128")]
impl ArrayArithmetics for i128 {
fn add(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
decimal::add(lhs, rhs)
}

fn sub(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
decimal::sub(lhs, rhs)
}

fn mul(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
decimal::mul(lhs, rhs)
}

fn div(lhs: &PrimitiveArray<Self>, rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
decimal::div(lhs, rhs)
}

fn div_scalar(_lhs: &PrimitiveArray<Self>, _rhs: &Self) -> PrimitiveArray<Self> {
// decimal::div_scalar(lhs, rhs)
todo!("decimal::div_scalar exists, but takes &PrimitiveScalar<i128>, not &i128");
}

fn rem(_lhs: &PrimitiveArray<Self>, _rhs: &PrimitiveArray<Self>) -> PrimitiveArray<Self> {
unimplemented!("requires support in arrow2 crate")
}

fn rem_scalar(_lhs: &PrimitiveArray<Self>, _rhs: &Self) -> PrimitiveArray<Self> {
unimplemented!("requires support in arrow2 crate")
}
}

macro_rules! apply_operand_on_chunkedarray_by_iter {

($self:ident, $rhs:ident, $operand:tt) => {
Expand Down Expand Up @@ -157,7 +237,12 @@ where
type Output = ChunkedArray<T>;

fn add(self, rhs: Self) -> Self::Output {
arithmetic_helper(self, rhs, basic::add, |lhs, rhs| lhs + rhs)
arithmetic_helper(
self,
rhs,
<T::Native as ArrayArithmetics>::add,
|lhs, rhs| lhs + rhs,
)
}
}

Expand All @@ -168,7 +253,12 @@ where
type Output = ChunkedArray<T>;

fn div(self, rhs: Self) -> Self::Output {
arithmetic_helper(self, rhs, basic::div, |lhs, rhs| lhs / rhs)
arithmetic_helper(
self,
rhs,
<T::Native as ArrayArithmetics>::div,
|lhs, rhs| lhs / rhs,
)
}
}

Expand All @@ -179,7 +269,12 @@ where
type Output = ChunkedArray<T>;

fn mul(self, rhs: Self) -> Self::Output {
arithmetic_helper(self, rhs, basic::mul, |lhs, rhs| lhs * rhs)
arithmetic_helper(
self,
rhs,
<T::Native as ArrayArithmetics>::mul,
|lhs, rhs| lhs * rhs,
)
}
}

Expand All @@ -190,7 +285,12 @@ where
type Output = ChunkedArray<T>;

fn rem(self, rhs: Self) -> Self::Output {
arithmetic_helper(self, rhs, basic::rem, |lhs, rhs| lhs % rhs)
arithmetic_helper(
self,
rhs,
<T::Native as ArrayArithmetics>::rem,
|lhs, rhs| lhs % rhs,
)
}
}

Expand All @@ -201,7 +301,12 @@ where
type Output = ChunkedArray<T>;

fn sub(self, rhs: Self) -> Self::Output {
arithmetic_helper(self, rhs, basic::sub, |lhs, rhs| lhs - rhs)
arithmetic_helper(
self,
rhs,
<T::Native as ArrayArithmetics>::sub,
|lhs, rhs| lhs - rhs,
)
}
}

Expand Down Expand Up @@ -317,7 +422,7 @@ where

fn div(self, rhs: N) -> Self::Output {
let rhs: T::Native = NumCast::from(rhs).expect("could not cast");
self.apply_kernel(&|arr| Box::new(basic::div_scalar(arr, &rhs)))
self.apply_kernel(&|arr| Box::new(<T::Native as ArrayArithmetics>::div_scalar(arr, &rhs)))
}
}

Expand All @@ -343,7 +448,7 @@ where

fn rem(self, rhs: N) -> Self::Output {
let rhs: T::Native = NumCast::from(rhs).expect("could not cast");
self.apply_kernel(&|arr| Box::new(basic::rem_scalar(arr, &rhs)))
self.apply_kernel(&|arr| Box::new(<T::Native as ArrayArithmetics>::rem_scalar(arr, &rhs)))
}
}

Expand Down
8 changes: 6 additions & 2 deletions polars/polars-core/src/datatypes/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ use std::ops::{Add, AddAssign, Div, Mul, Rem, Sub, SubAssign};
use ahash::RandomState;
pub use aliases::*;
pub use any_value::*;
use arrow::compute::arithmetics::basic::NativeArithmetics;
use arrow::compute::comparison::Simd8;
#[cfg(feature = "dtype-categorical")]
use arrow::datatypes::IntegerType;
Expand All @@ -41,6 +40,7 @@ use serde::{Deserialize, Serialize};
use serde::{Deserializer, Serializer};
pub use time_unit::*;

use crate::chunked_array::arithmetic::ArrayArithmetics;
pub use crate::chunked_array::logical::*;
#[cfg(feature = "object")]
use crate::chunked_array::object::PolarsObjectSafe;
Expand Down Expand Up @@ -82,6 +82,8 @@ impl_polars_datatype!(Int8Type, Int8, i8);
impl_polars_datatype!(Int16Type, Int16, i16);
impl_polars_datatype!(Int32Type, Int32, i32);
impl_polars_datatype!(Int64Type, Int64, i64);
#[cfg(feature = "dtype-i128")]
impl_polars_datatype!(Int128Type, Unknown, i128);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unknown should be Int128Type

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Presumably, you mean Int128 (which would be DataType::Int128)?

If so, I don't think that will work because I think DataType will eventually be something like DataType::Decimal(precision, scale), right?

Copy link
Member

@ritchie46 ritchie46 Jan 31, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, but then I think we should make the DataType something like DataType::Decimal(Option<precision, scale>). Then we can fill in the blanks later.

impl_polars_datatype!(Float32Type, Float32, f32);
impl_polars_datatype!(Float64Type, Float64, f64);
impl_polars_datatype!(DateType, Date, i32);
Expand Down Expand Up @@ -150,6 +152,8 @@ pub type Int8Chunked = ChunkedArray<Int8Type>;
pub type Int16Chunked = ChunkedArray<Int16Type>;
pub type Int32Chunked = ChunkedArray<Int32Type>;
pub type Int64Chunked = ChunkedArray<Int64Type>;
#[cfg(feature = "dtype-i128")]
pub type Int128Chunked = ChunkedArray<Int128Type>;
pub type Float32Chunked = ChunkedArray<Float32Type>;
pub type Float64Chunked = ChunkedArray<Float64Type>;
pub type Utf8Chunked = ChunkedArray<Utf8Type>;
Expand All @@ -175,7 +179,7 @@ pub trait NumericNative:
+ Bounded
+ FromPrimitive
+ IsFloat
+ NativeArithmetics
+ ArrayArithmetics
{
type POLARSTYPE: PolarsNumericType;
}
Expand Down
1 change: 1 addition & 0 deletions polars/polars-lazy/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ dtype-u8 = ["polars-plan/dtype-u8", "polars-pipe/dtype-u8"]
dtype-u16 = ["polars-plan/dtype-u16", "polars-pipe/dtype-u16"]
dtype-i8 = ["polars-plan/dtype-i8", "polars-pipe/dtype-i8"]
dtype-i16 = ["polars-plan/dtype-i16", "polars-pipe/dtype-i16"]
dtype-i128 = ["polars-plan/dtype-i128", "polars-pipe/dtype-i128"]
dtype-date = ["polars-plan/dtype-date", "polars-time/dtype-date", "temporal"]
dtype-datetime = ["polars-plan/dtype-datetime", "polars-time/dtype-datetime", "temporal"]
dtype-duration = ["polars-plan/dtype-duration", "polars-time/dtype-duration", "temporal"]
Expand Down
1 change: 1 addition & 0 deletions polars/polars-lazy/polars-pipe/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,4 +34,5 @@ dtype-u8 = ["polars-core/dtype-u8"]
dtype-u16 = ["polars-core/dtype-u16"]
dtype-i8 = ["polars-core/dtype-i8"]
dtype-i16 = ["polars-core/dtype-i16"]
dtype-i128 = ["polars-core/dtype-i128"]
dtype-categorical = ["polars-core/dtype-categorical"]
1 change: 1 addition & 0 deletions polars/polars-lazy/polars-plan/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ dtype-u8 = ["polars-core/dtype-u8"]
dtype-u16 = ["polars-core/dtype-u16"]
dtype-i8 = ["polars-core/dtype-i8"]
dtype-i16 = ["polars-core/dtype-i16"]
dtype-i128 = ["polars-core/dtype-i128"]
dtype-date = ["polars-core/dtype-date", "polars-time/dtype-date", "temporal"]
dtype-datetime = ["polars-core/dtype-datetime", "polars-time/dtype-datetime", "temporal"]
dtype-duration = ["polars-core/dtype-duration", "polars-time/dtype-duration", "temporal"]
Expand Down
1 change: 1 addition & 0 deletions polars/polars-ops/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ dtype-u8 = ["polars-core/dtype-u8"]
dtype-u16 = ["polars-core/dtype-u16"]
dtype-i8 = ["polars-core/dtype-i8"]
dtype-i16 = ["polars-core/dtype-i16"]
dtype-i128 = ["polars-core/dtype-i128"]
object = ["polars-core/object"]
propagate_nans = []
performant = ["polars-core/performant"]
Expand Down