Switching over to creating a getter to return a scalar udf

Author: timsaucer

rerun_py/src/catalog/dataset.rs

@@ -4,8 +4,10 @@ use arrow::array::{ArrayData, ArrayRef, RecordBatch, StringArray, StringViewArra
 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
 use arrow::pyarrow::{FromPyArrow as _, PyArrowType, ToPyArrow as _};
 use datafusion::common::exec_err;
-use pyo3::Bound;
+use pyo3::types::{PyAnyMethods, PyString, PyTuple};
+use pyo3::{Bound, IntoPyObject};
 use pyo3::{exceptions::PyRuntimeError, pyclass, pymethods, Py, PyAny, PyRef, PyResult, Python};
+use re_tuid::Tuid;
 use tokio_stream::StreamExt as _;
 
 use re_chunk_store::{ChunkStore, ChunkStoreHandle};
@@ -26,6 +28,7 @@ use re_protos::manifest_registry::v1alpha1::{
 };
 use re_sdk::{ComponentDescriptor, ComponentName};
 
+use crate::catalog::ConnectionHandle;
 use crate::catalog::{
     dataframe_query::PyDataframeQueryView, to_py_err, PyEntry, VectorDistanceMetricLike, VectorLike,
 };
@@ -98,54 +101,83 @@ impl PyDataset {
         .to_string()
     }
 
+    #[getter]
     fn partition_url_udf(
-        self_: PyRef<'_, Self>,
-        partition_id_expr: &Bound<'_, PyAny>,
+        self_: PyRef<'_, Self>
     ) -> PyResult<Py<PyAny>> {
+
         let super_ = self_.as_super();
         let connection = super_.client.borrow(self_.py()).connection().clone();
+        let py = self_.py();
 
-        let mut url = re_uri::DatasetDataUri {
-            origin: connection.origin().clone(),
-            dataset_id: super_.details.id.id,
-            partition_id: "default".to_owned(), // to be replaced during loop
-
-            //TODO(ab): add support for these two
-            time_range: None,
-            fragment: Default::default(),
-        };
+        #[pyclass]
+        struct PartitionUrlInner {
+            pub connection: ConnectionHandle,
+            pub dataset_id: Tuid,
+        }
 
-        let array_data = ArrayData::from_pyarrow_bound(partition_id_expr)?;
-
-        match array_data.data_type() {
-            DataType::Utf8 => {
-                let str_array = StringArray::from(array_data);
-                let str_iter = str_array.iter().map(|maybe_id| {
-                    maybe_id.map(|id| {
-                        url.partition_id = id.to_owned();
-                        url.to_string()
-                    })
-                });
-                let output_array: ArrayRef = Arc::new(str_iter.collect::<StringArray>());
-                output_array.to_data().to_pyarrow(super_.py())
+        #[pymethods]
+        impl PartitionUrlInner {
+            pub fn __call__(&self, py: Python<'_>, partition_id_expr: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
+
+                let mut url = re_uri::DatasetDataUri {
+                    origin: self.connection.origin().clone(),
+                    dataset_id: self.dataset_id,
+                    partition_id: "default".to_owned(), // to be replaced during loop
+        
+                    //TODO(ab): add support for these two
+                    time_range: None,
+                    fragment: Default::default(),
+                };
+        
+                let array_data = ArrayData::from_pyarrow_bound(partition_id_expr)?;
+        
+                match array_data.data_type() {
+                    DataType::Utf8 => {
+                        let str_array = StringArray::from(array_data);
+                        let str_iter = str_array.iter().map(|maybe_id| {
+                            maybe_id.map(|id| {
+                                url.partition_id = id.to_owned();
+                                url.to_string()
+                            })
+                        });
+                        let output_array: ArrayRef = Arc::new(str_iter.collect::<StringArray>());
+                        output_array.to_data().to_pyarrow(py)
+                    }
+                    DataType::Utf8View => {
+                        let str_array = StringViewArray::from(array_data);
+                        let str_iter = str_array.iter().map(|maybe_id| {
+                            maybe_id.map(|id| {
+                                url.partition_id = id.to_owned();
+                                url.to_string()
+                            })
+                        });
+                        let output_array: ArrayRef = Arc::new(str_iter.collect::<StringViewArray>());
+                        output_array.to_data().to_pyarrow(py)
+                    }
+                    _ => exec_err!(
+                        "Incorrect data type for partition_url_udf. Expected utf8 or utf8view. Received {}",
+                        array_data.data_type()
+                    )
+                    .map_err(to_py_err),
+                }
             }
-            DataType::Utf8View => {
-                let str_array = StringViewArray::from(array_data);
-                let str_iter = str_array.iter().map(|maybe_id| {
-                    maybe_id.map(|id| {
-                        url.partition_id = id.to_owned();
-                        url.to_string()
-                    })
-                });
-                let output_array: ArrayRef = Arc::new(str_iter.collect::<StringViewArray>());
-                output_array.to_data().to_pyarrow(super_.py())
-            }
-            _ => exec_err!(
-                "Incorrect data type for partition_url_udf. Expected utf8 or utf8view. Received {}",
-                array_data.data_type()
-            )
-            .map_err(to_py_err),
         }
+
+        let udf_factory = py.import("datafusion").and_then(|datafusion| datafusion.getattr("udf"))?;
+        let pa_utf8 = py.import("pyarrow").and_then(|pa| pa.getattr("utf8")?.call0())?;
+
+        let inner = PartitionUrlInner {
+            connection,
+            dataset_id: super_.details.id.id,
+        };
+        let bound_inner = inner.into_pyobject(py)?;
+        let py_stable = PyString::new(py, "stable");
+        
+        // df.udf(dataset.partition_url_udf,  pa.utf8(), pa.utf8(), 'stable')
+        let args = PyTuple::new(py, vec![bound_inner.as_any(), pa_utf8.as_any(), pa_utf8.as_any(), py_stable.as_any()])?;
+
+        Ok(udf_factory.call1(args)?.unbind())
     }
 
     /// Register a RRD URI to the dataset.