[FEAT] adding floor function (#1960)

Adding `floor` function. Resolves - #1910

daft/daft.pyi

@@ -871,6 +871,7 @@ class PyExpr:
     def alias(self, name: str) -> PyExpr: ...
     def cast(self, dtype: PyDataType) -> PyExpr: ...
     def ceil(self) -> PyExpr: ...
+    def floor(self) -> PyExpr: ...
     def if_else(self, if_true: PyExpr, if_false: PyExpr) -> PyExpr: ...
     def count(self, mode: CountMode) -> PyExpr: ...
     def sum(self) -> PyExpr: ...
@@ -986,6 +987,7 @@ class PySeries:
     def _agg_list(self) -> PySeries: ...
     def cast(self, dtype: PyDataType) -> PySeries: ...
     def ceil(self) -> PySeries: ...
+    def floor(self) -> PySeries: ...
     @staticmethod
     def concat(series: list[PySeries]) -> PySeries: ...
     def __len__(self) -> int: ...

daft/expressions/expressions.py

@@ -327,6 +327,11 @@ def ceil(self) -> Expression:
         expr = self._expr.ceil()
         return Expression._from_pyexpr(expr)
 
+    def floor(self) -> Expression:
+        """The floor of a numeric expression (``expr.floor()``)"""
+        expr = self._expr.floor()
+        return Expression._from_pyexpr(expr)
+
     def _count(self, mode: CountMode = CountMode.Valid) -> Expression:
         expr = self._expr.count(mode)
         return Expression._from_pyexpr(expr)

daft/series.py

@@ -352,6 +352,9 @@ def __abs__(self) -> Series:
     def ceil(self) -> Series:
         return Series._from_pyseries(self._series.ceil())
 
+    def floor(self) -> Series:
+        return Series._from_pyseries(self._series.floor())
+
     def __add__(self, other: object) -> Series:
         if not isinstance(other, Series):
             raise TypeError(f"expected another Series but got {type(other)}")

docs/source/api_docs/expressions.rst

@@ -47,6 +47,7 @@ Numeric
     Expression.__truediv__
     Expression.__mod__
     Expression.ceil
+    Expression.floor
 
 .. _api-comparison-expression:
 

src/daft-core/src/array/ops/floor.rs

@@ -0,0 +1,18 @@
+use num_traits::Float;
+
+use crate::{
+    array::DataArray,
+    datatypes::{DaftFloatType, DaftNumericType},
+};
+
+use common_error::DaftResult;
+
+impl<T: DaftFloatType> DataArray<T>
+where
+    T: DaftNumericType,
+    T::Native: Float,
+{
+    pub fn floor(&self) -> DaftResult<Self> {
+        self.apply(|v| v.floor())
+    }
+}

src/daft-core/src/array/ops/mod.rs

@@ -15,6 +15,7 @@ mod count;
 mod date;
 mod filter;
 mod float;
+mod floor;
 pub mod from_arrow;
 pub mod full;
 mod get;

src/daft-core/src/python/series.rs

@@ -112,6 +112,10 @@ impl PySeries {
         Ok(self.series.ceil()?.into())
     }
 
+    pub fn floor(&self) -> PyResult<Self> {
+        Ok(self.series.floor()?.into())
+    }
+
     pub fn take(&self, idx: &Self) -> PyResult<Self> {
         Ok(self.series.take(&idx.series)?.into())
     }

src/daft-core/src/series/ops/floor.rs

@@ -0,0 +1,20 @@
+use crate::datatypes::DataType;
+use crate::series::Series;
+use common_error::DaftError;
+use common_error::DaftResult;
+
+impl Series {
+    pub fn floor(&self) -> DaftResult<Series> {
+        use crate::series::array_impl::IntoSeries;
+        use DataType::*;
+        match self.data_type() {
+            Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 => Ok(self.clone()),
+            Float32 => Ok(self.f32().unwrap().floor()?.into_series()),
+            Float64 => Ok(self.f64().unwrap().floor()?.into_series()),
+            dt => Err(DaftError::TypeError(format!(
+                "floor not implemented for {}",
+                dt
+            ))),
+        }
+    }
+}

src/daft-core/src/series/ops/mod.rs

@@ -15,6 +15,7 @@ pub mod date;
 pub mod downcast;
 pub mod filter;
 pub mod float;
+pub mod floor;
 pub mod groups;
 pub mod hash;
 pub mod if_else;

src/daft-dsl/src/functions/numeric/floor.rs

@@ -0,0 +1,40 @@
+use common_error::{DaftError, DaftResult};
+use daft_core::{datatypes::Field, schema::Schema, series::Series};
+
+use super::super::FunctionEvaluator;
+use crate::Expr;
+
+pub(super) struct FloorEvaluator {}
+
+impl FunctionEvaluator for FloorEvaluator {
+    fn fn_name(&self) -> &'static str {
+        "floor"
+    }
+
+    fn to_field(&self, inputs: &[Expr], schema: &Schema, _: &Expr) -> DaftResult<Field> {
+        if inputs.len() != 1 {
+            return Err(DaftError::SchemaMismatch(format!(
+                "Expected 1 input arg, got {}",
+                inputs.len()
+            )));
+        }
+        let field = inputs.first().unwrap().to_field(schema)?;
+        if !field.dtype.is_numeric() {
+            return Err(DaftError::TypeError(format!(
+                "Expected input to floor to be numeric, got {}",
+                field.dtype
+            )));
+        }
+        Ok(field)
+    }
+
+    fn evaluate(&self, inputs: &[Series], _: &Expr) -> DaftResult<Series> {
+        if inputs.len() != 1 {
+            return Err(DaftError::SchemaMismatch(format!(
+                "Expected 1 input arg, got {}",
+                inputs.len()
+            )));
+        }
+        inputs.first().unwrap().floor()
+    }
+}

src/daft-dsl/src/functions/numeric/mod.rs

@@ -1,8 +1,10 @@
 mod abs;
 mod ceil;
+mod floor;
 
 use abs::AbsEvaluator;
 use ceil::CeilEvaluator;
+use floor::FloorEvaluator;
 
 use serde::{Deserialize, Serialize};
 
@@ -14,6 +16,7 @@ use super::FunctionEvaluator;
 pub enum NumericExpr {
     Abs,
     Ceil,
+    Floor,
 }
 
 impl NumericExpr {
@@ -23,6 +26,7 @@ impl NumericExpr {
         match self {
             Abs => &AbsEvaluator {},
             Ceil => &CeilEvaluator {},
+            Floor => &FloorEvaluator {},
         }
     }
 }
@@ -40,3 +44,10 @@ pub fn ceil(input: &Expr) -> Expr {
         inputs: vec![input.clone()],
     }
 }
+
+pub fn floor(input: &Expr) -> Expr {
+    Expr::Function {
+        func: super::FunctionExpr::Numeric(NumericExpr::Floor),
+        inputs: vec![input.clone()],
+    }
+}

src/daft-dsl/src/python.rs

@@ -155,6 +155,11 @@ impl PyExpr {
         Ok(ceil(&self.expr).into())
     }
 
+    pub fn floor(&self) -> PyResult<Self> {
+        use functions::numeric::floor;
+        Ok(floor(&self.expr).into())
+    }
+
     pub fn if_else(&self, if_true: &Self, if_false: &Self) -> PyResult<Self> {
         Ok(self.expr.if_else(&if_true.expr, &if_false.expr).into())
     }

tests/expressions/test_expressions.py

@@ -96,6 +96,15 @@ def test_repr_functions_ceil() -> None:
     assert repr_out == repr(copied)
 
 
+def test_repr_functions_floor() -> None:
+    a = col("a")
+    y = a.floor()
+    repr_out = repr(y)
+    assert repr_out == "floor(col(a))"
+    copied = copy.deepcopy(y)
+    assert repr_out == repr(copied)
+
+
 def test_repr_functions_day() -> None:
     a = col("a")
     y = a.dt.day()

tests/expressions/typing/test_arithmetic.py

@@ -82,3 +82,13 @@ def test_ceil(unary_data_fixture):
         run_kernel=lambda: arg.ceil(),
         resolvable=is_numeric(arg.datatype()),
     )
+
+
+def test_floor(unary_data_fixture):
+    arg = unary_data_fixture
+    assert_typing_resolve_vs_runtime_behavior(
+        data=(unary_data_fixture,),
+        expr=col(arg.name()).floor(),
+        run_kernel=lambda: arg.floor(),
+        resolvable=is_numeric(arg.datatype()),
+    )

tests/table/test_eval.py

@@ -180,3 +180,25 @@ def test_table_ceil_bad_input() -> None:
 
     with pytest.raises(ValueError, match="Expected input to ceil to be numeric"):
         table.eval_expression_list([col("a").ceil()])
+
+
+def test_table_numeric_floor() -> None:
+    table = MicroPartition.from_pydict(
+        {"a": [None, -1.0, -0.5, 0.0, 0.5, 2, None], "b": [-1.7, -1.5, -1.3, 0.3, 0.7, None, None]}
+    )
+
+    floor_table = table.eval_expression_list([col("a").floor(), col("b").floor()])
+
+    assert [math.floor(v) if v is not None else v for v in table.get_column("a").to_pylist()] == floor_table.get_column(
+        "a"
+    ).to_pylist()
+    assert [math.floor(v) if v is not None else v for v in table.get_column("b").to_pylist()] == floor_table.get_column(
+        "b"
+    ).to_pylist()
+
+
+def test_table_floor_bad_input() -> None:
+    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})
+
+    with pytest.raises(ValueError, match="Expected input to floor to be numeric"):
+        table.eval_expression_list([col("a").floor()])