From d012fe75d84a5fca4a2bb125f19939681bb23533 Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Mon, 13 Feb 2023 13:49:07 +0100 Subject: [PATCH] fix(rust, python): stabilize integer operation to minimal required dtype (#6841) --- .../polars-plan/src/dsl/arithmetic.rs | 233 ++++++++++++++++++ polars/polars-lazy/polars-plan/src/dsl/mod.rs | 225 +---------------- .../src/logical_plan/aexpr/schema.rs | 5 +- .../polars-plan/src/logical_plan/lit.rs | 4 + .../optimizer/type_coercion/mod.rs | 176 +++++++++++-- polars/polars-lazy/src/tests/mod.rs | 18 ++ py-polars/tests/unit/test_schema.py | 10 + 7 files changed, 429 insertions(+), 242 deletions(-) create mode 100644 polars/polars-lazy/polars-plan/src/dsl/arithmetic.rs diff --git a/polars/polars-lazy/polars-plan/src/dsl/arithmetic.rs b/polars/polars-lazy/polars-plan/src/dsl/arithmetic.rs new file mode 100644 index 000000000000..6f8cb17b11b2 --- /dev/null +++ b/polars/polars-lazy/polars-plan/src/dsl/arithmetic.rs @@ -0,0 +1,233 @@ +use std::ops::{Add, Div, Mul, Rem, Sub}; + +use super::*; + +// Arithmetic ops +impl Add for Expr { + type Output = Expr; + + fn add(self, rhs: Self) -> Self::Output { + binary_expr(self, Operator::Plus, rhs) + } +} + +impl Sub for Expr { + type Output = Expr; + + fn sub(self, rhs: Self) -> Self::Output { + binary_expr(self, Operator::Minus, rhs) + } +} + +impl Div for Expr { + type Output = Expr; + + fn div(self, rhs: Self) -> Self::Output { + binary_expr(self, Operator::Divide, rhs) + } +} + +impl Mul for Expr { + type Output = Expr; + + fn mul(self, rhs: Self) -> Self::Output { + binary_expr(self, Operator::Multiply, rhs) + } +} + +impl Rem for Expr { + type Output = Expr; + + fn rem(self, rhs: Self) -> Self::Output { + binary_expr(self, Operator::Modulus, rhs) + } +} + +impl Expr { + /// Floor divide `self` by `rhs`. + pub fn floor_div(self, rhs: Self) -> Self { + binary_expr(self, Operator::FloorDivide, rhs) + } + + /// Raise expression to the power `exponent` + pub fn pow>(self, exponent: E) -> Self { + Expr::Function { + input: vec![self, exponent.into()], + function: FunctionExpr::Pow, + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the sine of the given expression + #[cfg(feature = "trigonometry")] + pub fn sin(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::Sin), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the cosine of the given expression + #[cfg(feature = "trigonometry")] + pub fn cos(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::Cos), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the tangent of the given expression + #[cfg(feature = "trigonometry")] + pub fn tan(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::Tan), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the inverse sine of the given expression + #[cfg(feature = "trigonometry")] + pub fn arcsin(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcSin), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the inverse cosine of the given expression + #[cfg(feature = "trigonometry")] + pub fn arccos(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcCos), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the inverse tangent of the given expression + #[cfg(feature = "trigonometry")] + pub fn arctan(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcTan), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the hyperbolic sine of the given expression + #[cfg(feature = "trigonometry")] + pub fn sinh(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::Sinh), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the hyperbolic cosine of the given expression + #[cfg(feature = "trigonometry")] + pub fn cosh(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::Cosh), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the hyperbolic tangent of the given expression + #[cfg(feature = "trigonometry")] + pub fn tanh(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::Tanh), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the inverse hyperbolic sine of the given expression + #[cfg(feature = "trigonometry")] + pub fn arcsinh(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcSinh), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the inverse hyperbolic cosine of the given expression + #[cfg(feature = "trigonometry")] + pub fn arccosh(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcCosh), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } + + /// Compute the inverse hyperbolic tangent of the given expression + #[cfg(feature = "trigonometry")] + pub fn arctanh(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcTanh), + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + fmt_str: "arctanh", + ..Default::default() + }, + } + } + + /// Compute the sign of the given expression + #[cfg(feature = "sign")] + pub fn sign(self) -> Self { + Expr::Function { + input: vec![self], + function: FunctionExpr::Sign, + options: FunctionOptions { + collect_groups: ApplyOptions::ApplyFlat, + ..Default::default() + }, + } + } +} diff --git a/polars/polars-lazy/polars-plan/src/dsl/mod.rs b/polars/polars-lazy/polars-plan/src/dsl/mod.rs index 80b343925585..2b9fa4839f61 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/mod.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/mod.rs @@ -3,6 +3,7 @@ pub mod cat; #[cfg(feature = "dtype-categorical")] pub use cat::*; +mod arithmetic; #[cfg(feature = "dtype-binary")] pub mod binary; #[cfg(feature = "temporal")] @@ -23,7 +24,6 @@ pub mod string; mod struct_; use std::fmt::Debug; -use std::ops::{Add, Div, Mul, Rem, Sub}; use std::sync::Arc; pub use expr::*; @@ -1186,188 +1186,6 @@ impl Expr { binary_expr(self, Operator::Or, expr.into()) } - /// Raise expression to the power `exponent` - pub fn pow>(self, exponent: E) -> Self { - Expr::Function { - input: vec![self, exponent.into()], - function: FunctionExpr::Pow, - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the sine of the given expression - #[cfg(feature = "trigonometry")] - pub fn sin(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::Sin), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the cosine of the given expression - #[cfg(feature = "trigonometry")] - pub fn cos(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::Cos), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the tangent of the given expression - #[cfg(feature = "trigonometry")] - pub fn tan(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::Tan), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the inverse sine of the given expression - #[cfg(feature = "trigonometry")] - pub fn arcsin(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcSin), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the inverse cosine of the given expression - #[cfg(feature = "trigonometry")] - pub fn arccos(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcCos), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the inverse tangent of the given expression - #[cfg(feature = "trigonometry")] - pub fn arctan(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcTan), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the hyperbolic sine of the given expression - #[cfg(feature = "trigonometry")] - pub fn sinh(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::Sinh), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the hyperbolic cosine of the given expression - #[cfg(feature = "trigonometry")] - pub fn cosh(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::Cosh), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the hyperbolic tangent of the given expression - #[cfg(feature = "trigonometry")] - pub fn tanh(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::Tanh), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the inverse hyperbolic sine of the given expression - #[cfg(feature = "trigonometry")] - pub fn arcsinh(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcSinh), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the inverse hyperbolic cosine of the given expression - #[cfg(feature = "trigonometry")] - pub fn arccosh(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcCosh), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - - /// Compute the inverse hyperbolic tangent of the given expression - #[cfg(feature = "trigonometry")] - pub fn arctanh(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Trigonometry(TrigonometricFunction::ArcTanh), - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - fmt_str: "arctanh", - ..Default::default() - }, - } - } - - /// Compute the sign of the given expression - #[cfg(feature = "sign")] - pub fn sign(self) -> Self { - Expr::Function { - input: vec![self], - function: FunctionExpr::Sign, - options: FunctionOptions { - collect_groups: ApplyOptions::ApplyFlat, - ..Default::default() - }, - } - } - /// Filter a single column /// Should be used in aggregation context. If you want to filter on a DataFrame level, use /// [LazyFrame::filter](LazyFrame::filter) @@ -2249,47 +2067,6 @@ impl Expr { } } -// Arithmetic ops -impl Add for Expr { - type Output = Expr; - - fn add(self, rhs: Self) -> Self::Output { - binary_expr(self, Operator::Plus, rhs) - } -} - -impl Sub for Expr { - type Output = Expr; - - fn sub(self, rhs: Self) -> Self::Output { - binary_expr(self, Operator::Minus, rhs) - } -} - -impl Div for Expr { - type Output = Expr; - - fn div(self, rhs: Self) -> Self::Output { - binary_expr(self, Operator::Divide, rhs) - } -} - -impl Mul for Expr { - type Output = Expr; - - fn mul(self, rhs: Self) -> Self::Output { - binary_expr(self, Operator::Multiply, rhs) - } -} - -impl Rem for Expr { - type Output = Expr; - - fn rem(self, rhs: Self) -> Self::Output { - binary_expr(self, Operator::Modulus, rhs) - } -} - /// Apply a function/closure over multiple columns once the logical plan get executed. /// /// This function is very similar to `[apply_mul]`, but differs in how it handles aggregations. diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/aexpr/schema.rs b/polars/polars-lazy/polars-plan/src/logical_plan/aexpr/schema.rs index 08a5775f17a8..296af3ffa3f8 100644 --- a/polars/polars-lazy/polars-plan/src/logical_plan/aexpr/schema.rs +++ b/polars/polars-lazy/polars-plan/src/logical_plan/aexpr/schema.rs @@ -43,7 +43,10 @@ impl AExpr { }), } } - Literal(sv) => Ok(Field::new("literal", sv.get_datatype())), + Literal(sv) => Ok(match sv { + LiteralValue::Series(s) => s.field().into_owned(), + _ => Field::new("literal", sv.get_datatype()), + }), BinaryExpr { left, right, op } => { use DataType::*; diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/lit.rs b/polars/polars-lazy/polars-plan/src/logical_plan/lit.rs index c1fccff44899..a30e84f87389 100644 --- a/polars/polars-lazy/polars-plan/src/logical_plan/lit.rs +++ b/polars/polars-lazy/polars-plan/src/logical_plan/lit.rs @@ -57,6 +57,10 @@ pub enum LiteralValue { } impl LiteralValue { + pub(crate) fn is_float(&self) -> bool { + matches!(self, LiteralValue::Float32(_) | LiteralValue::Float64(_)) + } + pub fn to_anyvalue(&self) -> Option { use LiteralValue::*; let av = match self { diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/type_coercion/mod.rs b/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/type_coercion/mod.rs index f48ba049226d..ff0aae0bb78e 100644 --- a/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/type_coercion/mod.rs +++ b/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/type_coercion/mod.rs @@ -22,6 +22,136 @@ macro_rules! unpack { }}; } +// `dtype_other` comes from a column +// so we shrink literal so it fits into that column dtype. +fn shrink_literal(dtype_other: &DataType, literal: &LiteralValue) -> Option { + match (dtype_other, literal) { + (DataType::UInt64, LiteralValue::Int64(v)) => { + if *v > 0 { + return Some(DataType::UInt64); + } + } + (DataType::UInt64, LiteralValue::Int32(v)) => { + if *v > 0 { + return Some(DataType::UInt64); + } + } + #[cfg(feature = "dtype-i16")] + (DataType::UInt64, LiteralValue::Int16(v)) => { + if *v > 0 { + return Some(DataType::UInt64); + } + } + #[cfg(feature = "dtype-i8")] + (DataType::UInt64, LiteralValue::Int8(v)) => { + if *v > 0 { + return Some(DataType::UInt64); + } + } + (DataType::UInt32, LiteralValue::Int64(v)) => { + if *v > 0 && *v < u32::MAX as i64 { + return Some(DataType::UInt32); + } + } + (DataType::UInt32, LiteralValue::Int32(v)) => { + if *v > 0 { + return Some(DataType::UInt32); + } + } + #[cfg(feature = "dtype-i16")] + (DataType::UInt32, LiteralValue::Int16(v)) => { + if *v > 0 { + return Some(DataType::UInt32); + } + } + #[cfg(feature = "dtype-i8")] + (DataType::UInt32, LiteralValue::Int8(v)) => { + if *v > 0 { + return Some(DataType::UInt32); + } + } + (DataType::UInt16, LiteralValue::Int64(v)) => { + if *v > 0 && *v < u16::MAX as i64 { + return Some(DataType::UInt16); + } + } + (DataType::UInt16, LiteralValue::Int32(v)) => { + if *v > 0 && *v < u16::MAX as i32 { + return Some(DataType::UInt16); + } + } + #[cfg(feature = "dtype-i16")] + (DataType::UInt16, LiteralValue::Int16(v)) => { + if *v > 0 { + return Some(DataType::UInt16); + } + } + #[cfg(feature = "dtype-i8")] + (DataType::UInt16, LiteralValue::Int8(v)) => { + if *v > 0 { + return Some(DataType::UInt16); + } + } + (DataType::UInt8, LiteralValue::Int64(v)) => { + if *v > 0 && *v < u8::MAX as i64 { + return Some(DataType::UInt8); + } + } + (DataType::UInt8, LiteralValue::Int32(v)) => { + if *v > 0 && *v < u8::MAX as i32 { + return Some(DataType::UInt8); + } + } + #[cfg(feature = "dtype-i16")] + (DataType::UInt8, LiteralValue::Int16(v)) => { + if *v > 0 && *v < u8::MAX as i16 { + return Some(DataType::UInt8); + } + } + #[cfg(feature = "dtype-i8")] + (DataType::UInt8, LiteralValue::Int8(v)) => { + if *v > 0 && *v < u8::MAX as i8 { + return Some(DataType::UInt8); + } + } + (DataType::Int32, LiteralValue::Int64(v)) => { + if *v <= i32::MAX as i64 { + return Some(DataType::Int32); + } + } + (DataType::Int16, LiteralValue::Int64(v)) => { + if *v <= i16::MAX as i64 { + return Some(DataType::Int16); + } + } + (DataType::Int16, LiteralValue::Int32(v)) => { + if *v <= i16::MAX as i32 { + return Some(DataType::Int16); + } + } + (DataType::Int8, LiteralValue::Int64(v)) => { + if *v <= i8::MAX as i64 { + return Some(DataType::Int8); + } + } + (DataType::Int8, LiteralValue::Int32(v)) => { + if *v <= i8::MAX as i32 { + return Some(DataType::Int8); + } + } + #[cfg(feature = "dtype-i16")] + (DataType::Int8, LiteralValue::Int16(v)) => { + if *v <= i8::MAX as i16 { + return Some(DataType::Int8); + } + } + _ => { + // the rest is done by supertypes. + } + } + None +} + /// determine if we use the supertype or not. For instance when we have a column Int64 and we compare with literal UInt32 /// it would be wasteful to cast the column instead of the literal. fn modify_supertype( @@ -34,34 +164,46 @@ fn modify_supertype( // only interesting on numerical types // other types will always use the supertype. if type_left.is_numeric() && type_right.is_numeric() { + use AExpr::*; match (left, right) { // don't let the literal f64 coerce the f32 column - (AExpr::Literal(LiteralValue::Float64(_) | LiteralValue::Int32(_) | LiteralValue::Int64(_)), _) if matches!(type_right, DataType::Float32) => { - st = DataType::Float32 + ( + Literal(LiteralValue::Float64(_) | LiteralValue::Int32(_) | LiteralValue::Int64(_)), + _, + ) if matches!(type_right, DataType::Float32) => st = DataType::Float32, + ( + _, + Literal(LiteralValue::Float64(_) | LiteralValue::Int32(_) | LiteralValue::Int64(_)), + ) if matches!(type_left, DataType::Float32) => st = DataType::Float32, + // always make sure that we cast to floats if one of the operands is float + (Literal(lv), _) | (_, Literal(lv)) if lv.is_float() => {} + + // TODO: see if we can activate this for columns as well. + // shrink the literal value if it fits in the column dtype + (Literal(LiteralValue::Series(_)), Literal(lv)) => { + if let Some(dtype) = shrink_literal(type_left, lv) { + st = dtype; + } } - (_, AExpr::Literal(LiteralValue::Float64(_) | LiteralValue::Int32(_) | LiteralValue::Int64(_))) if matches!(type_left, DataType::Float32) => { - st = DataType::Float32 + // shrink the literal value if it fits in the column dtype + (Literal(lv), Literal(LiteralValue::Series(_))) => { + if let Some(dtype) = shrink_literal(type_right, lv) { + st = dtype; + } } - // do nothing and use supertype - (AExpr::Literal(_), AExpr::Literal(_)) - // always make sure that we cast to floats if one of the operands is float - // and the left type is integer - |(AExpr::Literal(LiteralValue::Float32(_) | LiteralValue::Float64(_)), _) - |(_, AExpr::Literal(LiteralValue::Float32(_) | LiteralValue::Float64(_))) - => {} + (Literal(_), Literal(_)) => {} // cast literal to right type if they fit in the range - (AExpr::Literal(value), _) => { + (Literal(value), _) => { if let Some(lit_val) = value.to_anyvalue() { - if type_right.value_within_range(lit_val) { - st = type_right.clone(); - } + if type_right.value_within_range(lit_val) { + st = type_right.clone(); + } } } // cast literal to left type - (_, AExpr::Literal(value)) => { - + (_, Literal(value)) => { if let Some(lit_val) = value.to_anyvalue() { if type_left.value_within_range(lit_val) { st = type_left.clone(); diff --git a/polars/polars-lazy/src/tests/mod.rs b/polars/polars-lazy/src/tests/mod.rs index 7cadab2d543f..598a3bf05fd9 100644 --- a/polars/polars-lazy/src/tests/mod.rs +++ b/polars/polars-lazy/src/tests/mod.rs @@ -150,3 +150,21 @@ pub(crate) fn get_df() -> DataFrame { .unwrap(); df } + +#[test] +fn test_foo() { + let df: DataFrame = df![ + "a" => [1u64] + ] + .unwrap(); + + let s = df.column("a").unwrap().clone(); + + let df = df + .lazy() + .select([lit(s).floor_div(lit(1i64))]) + .collect() + .unwrap(); + + dbg!(df); +} diff --git a/py-polars/tests/unit/test_schema.py b/py-polars/tests/unit/test_schema.py index a53beeef3fbf..ca42a33d4ab3 100644 --- a/py-polars/tests/unit/test_schema.py +++ b/py-polars/tests/unit/test_schema.py @@ -353,3 +353,13 @@ def test_duration_divison_schema() -> None: assert q.schema == {"a": pl.Float64} assert q.collect().to_dict(False) == {"a": [1.0]} + + +def test_int_operator_stability() -> None: + for dt in pl.datatypes.INTEGER_DTYPES: + s = pl.Series(values=[10], dtype=dt) + assert pl.select(pl.lit(s) // 2).dtypes == [dt] + assert pl.select(pl.lit(s) + 2).dtypes == [dt] + assert pl.select(pl.lit(s) - 2).dtypes == [dt] + assert pl.select(pl.lit(s) * 2).dtypes == [dt] + assert pl.select(pl.lit(s) / 2).dtypes == [pl.Float64]