EXPERIMENT: Daft oopsies

daft/daft/__init__.pyi

@@ -1168,6 +1168,7 @@ def udf(
     expressions: list[PyExpr],
     return_dtype: PyDataType,
     init_args: InitArgsType,
+    scalar_udf: bool,
     resource_request: ResourceRequest | None,
     batch_size: int | None,
     concurrency: int | None,

daft/expressions/expressions.py

@@ -374,6 +374,7 @@ def udf(
         expressions: builtins.list[Expression],
         return_dtype: DataType,
         init_args: InitArgsType,
+        scalar_udf: bool,
         resource_request: ResourceRequest | None,
         batch_size: int | None,
         concurrency: int | None,
@@ -386,6 +387,7 @@ def udf(
                 [e._expr for e in expressions],
                 return_dtype._dtype,
                 init_args,
+                scalar_udf,
                 resource_request,
                 batch_size,
                 concurrency,

daft/udf.py

@@ -10,6 +10,7 @@
 from daft.datatype import DataType, DataTypeLike
 from daft.dependencies import np, pa
 from daft.expressions import Expression
+from daft.scalar_udf import run_scalar_udf
 from daft.series import Series
 
 InitArgsType = Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]]
@@ -247,6 +248,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Expression:
             expressions=expressions,
             return_dtype=self.return_dtype,
             init_args=self.init_args,
+            scalar_udf=False,
             resource_request=self.resource_request,
             batch_size=self.batch_size,
             concurrency=self.concurrency,

src/daft-csv/src/local.rs

@@ -898,7 +898,7 @@ where
             })
             .collect::<DaftResult<Vec<Series>>>()?;
         let num_rows = chunk.first().map(|s| s.len()).unwrap_or(0);
-        let table = RecordBatch::new_unchecked(read_schema.clone(), chunk, num_rows);
+        let table = RecordBatch::new_unchecked(read_schema.clone(), chunk, num_rows, None);
         let table = if let Some(predicate) = &predicate {
             let predicate = BoundExpr::try_new(predicate.clone(), &read_schema)?;
             let filtered = table.filter(&[predicate.clone()])?;

src/daft-csv/src/read.rs

@@ -687,7 +687,7 @@ fn parse_into_column_array_chunk_stream(
                         })
                         .collect::<DaftResult<Vec<Series>>>()?;
                     let num_rows = chunk.first().map_or(0, daft_core::series::Series::len);
-                    Ok(RecordBatch::new_unchecked(read_schema, chunk, num_rows))
+                    Ok(RecordBatch::new_unchecked(read_schema, chunk, num_rows, None))
                 })();
                 let _ = send.send(result);
             });

src/daft-dsl/src/expr/bound_expr.rs

@@ -9,7 +9,7 @@ use super::{
     bound_col, AggExpr, Column, Expr, ExprRef, ResolvedColumn, UnresolvedColumn, WindowExpr,
 };
 
-#[derive(Clone, Display)]
+#[derive(Clone, Display, Debug)]
 /// A simple newtype around ExprRef that ensures that all of the columns in the held expression are bound.
 ///
 /// We have several column variants: unresolved, resolved, and bound.

src/daft-dsl/src/functions/map/get.rs

@@ -2,7 +2,10 @@ use common_error::{DaftError, DaftResult};
 use daft_core::prelude::*;
 
 use super::super::FunctionEvaluator;
-use crate::{functions::FunctionExpr, ExprRef};
+use crate::{
+    functions::{FunctionExpr, FunctionResult},
+    ExprRef,
+};
 
 pub(super) struct GetEvaluator {}
 
@@ -34,14 +37,14 @@ impl FunctionEvaluator for GetEvaluator {
         Ok(field)
     }
 
-    fn evaluate(&self, inputs: &[Series], _: &FunctionExpr) -> DaftResult<Series> {
+    fn evaluate(&self, inputs: &[Series], _: &FunctionExpr) -> DaftResult<FunctionResult> {
         let [input, key] = inputs else {
             return Err(DaftError::ValueError(format!(
                 "Expected 2 input args, got {}",
                 inputs.len()
             )));
         };
 
-        input.map_get(key)
+        Ok(input.map_get(key)?.into())
     }
 }

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

@@ -35,6 +35,25 @@
     Partitioning(PartitioningExpr),
 }
 
+pub struct FunctionResult {
+    pub ok: Series,
+    pub error: Option<Series>,
+}
+
+impl FunctionResult {
+    pub fn into_inner(self) -> (Series, Option<Series>) {
+        (self.ok, self.error)
+    }
+}
+
+impl From<Series> for FunctionResult {
+    fn from(series: Series) -> Self {
+        FunctionResult { ok: series, error: None }
+    }
+}
+
+
+
 pub trait FunctionEvaluator {
     fn fn_name(&self) -> &'static str;
     fn to_field(
@@ -43,7 +62,7 @@
         schema: &Schema,
         expr: &FunctionExpr,
     ) -> DaftResult<Field>;
-    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<Series>;
+    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<FunctionResult>;
 }
 
 impl FunctionExpr {
@@ -79,7 +98,7 @@
         self.get_evaluator().to_field(inputs, schema, expr)
     }
 
-    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<Series> {
+    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<FunctionResult> {
         self.get_evaluator().evaluate(inputs, expr)
     }
 }

src/daft-dsl/src/functions/partitioning/evaluators.rs

@@ -1,7 +1,7 @@
 use common_error::{DaftError, DaftResult};
 use daft_core::prelude::*;
 
-use super::super::FunctionEvaluator;
+use super::super::{FunctionEvaluator, FunctionResult};
 use crate::{functions::partitioning::PartitioningExpr, ExprRef};
 
 macro_rules! impl_func_evaluator_for_partitioning {
@@ -39,9 +39,9 @@ macro_rules! impl_func_evaluator_for_partitioning {
                 }
             }
 
-            fn evaluate(&self, inputs: &[Series], _: &FunctionExpr) -> DaftResult<Series> {
+            fn evaluate(&self, inputs: &[Series], _: &FunctionExpr) -> DaftResult<FunctionResult> {
                 match inputs {
-                    [input] => input.$kernel(),
+                    [input] => Ok(input.$kernel()?.into()),
                     _ => Err(DaftError::ValueError(format!(
                         "Expected 1 input arg for {}, got {}",
                         stringify!($op),
@@ -98,14 +98,14 @@ impl FunctionEvaluator for IcebergBucketEvaluator {
         }
     }
 
-    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<Series> {
+    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<FunctionResult> {
         let n = match expr {
             FunctionExpr::Partitioning(PartitioningExpr::IcebergBucket(n)) => n,
             _ => panic!("Expected PartitioningExpr::IcebergBucket Expr, got {expr}"),
         };
 
         match inputs {
-            [input] => input.partitioning_iceberg_bucket(*n),
+            [input] => Ok(input.partitioning_iceberg_bucket(*n)?.into()),
             _ => Err(DaftError::ValueError(format!(
                 "Expected 1 input arg, got {}",
                 inputs.len()
@@ -142,14 +142,14 @@ impl FunctionEvaluator for IcebergTruncateEvaluator {
         }
     }
 
-    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<Series> {
+    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<FunctionResult> {
         let w = match expr {
             FunctionExpr::Partitioning(PartitioningExpr::IcebergTruncate(w)) => w,
             _ => panic!("Expected PartitioningExpr::IcebergTruncate Expr, got {expr}"),
         };
 
         match inputs {
-            [input] => input.partitioning_iceberg_truncate(*w),
+            [input] => Ok(input.partitioning_iceberg_truncate(*w)?.into()),
             _ => Err(DaftError::ValueError(format!(
                 "Expected 1 input arg, got {}",
                 inputs.len()

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

@@ -74,6 +74,7 @@ pub struct PythonUDF {
     pub bound_args: RuntimePyObject,
     pub num_expressions: usize,
     pub return_dtype: DataType,
+    pub scalar_udf: bool,
     pub resource_request: Option<ResourceRequest>,
     pub batch_size: Option<usize>,
     pub concurrency: Option<usize>,
@@ -87,6 +88,7 @@ pub fn udf(
     expressions: &[ExprRef],
     return_dtype: DataType,
     init_args: RuntimePyObject,
+    scalar_udf: bool,
     resource_request: Option<ResourceRequest>,
     batch_size: Option<usize>,
     concurrency: Option<usize>,
@@ -101,6 +103,7 @@ pub fn udf(
             resource_request,
             batch_size,
             concurrency,
+            scalar_udf,
         }),
         inputs: expressions.into(),
     })

src/daft-dsl/src/functions/python/udf.rs

@@ -7,7 +7,10 @@
 };
 
 use super::{super::FunctionEvaluator, PythonUDF};
-use crate::{functions::FunctionExpr, ExprRef};
+use crate::{
+    functions::{FunctionExpr, FunctionResult},
+    ExprRef,
+};
 
 #[cfg(feature = "python")]
 fn run_udf(
@@ -58,9 +61,56 @@
     }
 }
 
+#[cfg(feature = "python")]
+fn run_scalar_udf(
+    py: pyo3::Python,
+    inputs: &[Series],
+    func: pyo3::Py<PyAny>,
+    bound_args: pyo3::Py<PyAny>,
+    return_dtype: &DataType,
+) -> DaftResult<(Series, Series)> {
+    use daft_core::python::{PyDataType, PySeries};
+
+    // Convert input Rust &[Series] to wrapped Python Vec<Bound<PyAny>>
+    let py_series_module = PyModule::import(py, pyo3::intern!(py, "daft.series"))?;
+    let py_series_class = py_series_module.getattr(pyo3::intern!(py, "Series"))?;
+    let pyseries: PyResult<Vec<Bound<PyAny>>> = inputs
+        .iter()
+        .map(|s| {
+            py_series_class.call_method(
+                pyo3::intern!(py, "_from_pyseries"),
+                (PySeries { series: s.clone() },),
+                None,
+            )
+        })
+        .collect();
+    let pyseries = pyseries?;
+
+    // Run the function on the converted Vec<Bound<PyAny>>
+    let py_udf_module = PyModule::import(py, pyo3::intern!(py, "daft.scalar_udf"))?;
+    let run_udf_func = py_udf_module.getattr(pyo3::intern!(py, "run_scalar_udf"))?;
+    let result = run_udf_func.call1((
+        func,                                   // Function to run
+        bound_args,                             // Arguments bound to the function
+        pyseries,                               // evaluated_expressions
+        PyDataType::from(return_dtype.clone()), // Returned datatype
+    ));
+
+    match result {
+        Ok(pyany) => {
+            let ok_and_error_series = pyany.extract::<(PySeries, PySeries)>();
+            match ok_and_error_series {
+                Ok((ok_series, error_series)) => Ok((ok_series.series, error_series.series)),
+                Err(e) => Err(DaftError::ValueError(format!("Internal error occurred when coercing the results of running UDF to Series:\n\n{e}"))),
+            }
+        }
+        Err(e) => Err(e.into()),
+    }
+}
+
 impl PythonUDF {
     #[cfg(feature = "python")]
-    pub fn call_udf(&self, inputs: &[Series]) -> DaftResult<Series> {
+    pub fn call_udf(&self, inputs: &[Series]) -> DaftResult<FunctionResult> {
         use pyo3::Python;
 
         use crate::functions::python::{py_udf_initialize, MaybeInitializedUDF};
@@ -83,14 +133,29 @@
                 }
             };
 
-            run_udf(
-                py,
-                inputs,
-                func,
-                self.bound_args.clone().unwrap().clone_ref(py),
-                &self.return_dtype,
-                self.batch_size,
-            )
+            if self.scalar_udf {
+                let (ok, error) = run_scalar_udf(
+                    py,
+                    inputs,
+                    func,
+                    self.bound_args.clone().unwrap().clone_ref(py),
+                    &self.return_dtype,
+                )?;
+                Ok(FunctionResult {
+                    ok,
+                    error: Some(error),
+                })
+            } else {
+                let result = run_udf(
+                    py,
+                    inputs,
+                    func,
+                    self.bound_args.clone().unwrap().clone_ref(py),
+                    &self.return_dtype,
+                    self.batch_size,
+                )?;
+                Ok(result.into())
+            }
         })
     }
 }
@@ -116,7 +181,7 @@
         }
     }
 
-    fn evaluate(&self, inputs: &[Series], _: &FunctionExpr) -> DaftResult<Series> {
+    fn evaluate(&self, inputs: &[Series], _: &FunctionExpr) -> DaftResult<FunctionResult> {
         #[cfg(not(feature = "python"))]
         {
             panic!("Cannot evaluate a PythonUDF without compiling for Python");

src/daft-dsl/src/functions/sketch/percentile.rs

@@ -2,7 +2,7 @@ use common_error::{DaftError, DaftResult};
 use daft_core::prelude::*;
 
 use super::{super::FunctionEvaluator, SketchExpr};
-use crate::{functions::FunctionExpr, ExprRef};
+use crate::{functions::{FunctionExpr, FunctionResult}, ExprRef};
 
 pub(super) struct PercentileEvaluator {}
 
@@ -57,13 +57,13 @@ impl FunctionEvaluator for PercentileEvaluator {
         }
     }
 
-    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<Series> {
+    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<FunctionResult> {
         match inputs {
             [input] => match expr {
                 FunctionExpr::Sketch(SketchExpr::Percentile {
                     percentiles,
                     force_list_output,
-                }) => input.sketch_percentile(percentiles.0.as_slice(), *force_list_output),
+                }) => Ok(input.sketch_percentile(percentiles.0.as_slice(), *force_list_output)?.into()),
                 _ => unreachable!(
                     "PercentileEvaluator must evaluate a SketchExpr::Percentile expression"
                 ),

src/daft-dsl/src/functions/struct_/get.rs

@@ -2,7 +2,7 @@ use common_error::{DaftError, DaftResult};
 use daft_core::prelude::*;
 
 use super::{super::FunctionEvaluator, StructExpr};
-use crate::{functions::FunctionExpr, ExprRef};
+use crate::{functions::{FunctionExpr, FunctionResult}, ExprRef};
 
 pub(super) struct GetEvaluator {}
 
@@ -56,15 +56,15 @@ impl FunctionEvaluator for GetEvaluator {
         }
     }
 
-    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<Series> {
+    fn evaluate(&self, inputs: &[Series], expr: &FunctionExpr) -> DaftResult<FunctionResult> {
         match inputs {
             [input] => {
                 let name = match expr {
                     FunctionExpr::Struct(StructExpr::Get(name)) => name,
                     _ => panic!("Expected Struct Get Expr, got {expr}"),
                 };
 
-                input.struct_get(name)
+                Ok(input.struct_get(name)?.into())
             }
             _ => Err(DaftError::ValueError(format!(
                 "Expected 1 input arg, got {}",