allow returning scalar values as dataframes

basf · Feb 15, 2024 · 5cdf04d · 5cdf04d
1 parent a32b505
commit 5cdf04d
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/rormula/rormula/_rormula.pyi → rormula/rormula/rormula.pyi b/rormula/rormula/_rormula.pyi → rormula/rormula/rormula.pyi
diff --git a/rormula/src/lib.rs b/rormula/src/lib.rs
@@ -74,10 +74,8 @@ fn eval_arithmetic<'py>(
                 let res = pya.into_pyarray(py);
                 Ok(res)
             }
+            Value::Scalar(s) => Ok(Array2::<f64>::from_elem((1, 1), s).into_pyarray(py)),
             Value::Cats(_) => Err(PyValueError::new_err("result cannot be cat".to_string())),
-            Value::Scalar(s) => Err(PyValueError::new_err(format!(
-                "result cannot be skalar but got {s}"
-            ))),
             Value::Error(e) => Err(PyValueError::new_err(format!("computation failed, {e:?}"))),
         }
     }

diff --git a/rormula/test/test_arithmetic.py b/rormula/test/test_arithmetic.py
@@ -72,6 +72,11 @@ def test_scalar_scalar():
     res = rormula.eval_asdf(df)
     ref = df.eval(s)
     np.allclose(res[name].to_numpy(), ref.values)
+    s = "5/3"
+    rormula = Arithmetic(s, name)
+    res = rormula.eval_asdf(df)
+    ref = df.eval(s)    
+    np.allclose(res[name].to_numpy(), ref)
 
 
 if __name__ == "__main__":