Tsaucer/dataset url udf

rerun_py/src/catalog/dataset.rs

@@ -1,9 +1,13 @@
 use std::sync::Arc;
 
-use arrow::array::{RecordBatch, StringArray};
-use arrow::datatypes::{Field, Schema as ArrowSchema};
-use arrow::pyarrow::PyArrowType;
+use arrow::array::{ArrayData, ArrayRef, RecordBatch, StringArray, StringViewArray};
+use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
+use arrow::pyarrow::{FromPyArrow as _, PyArrowType, ToPyArrow};
+use datafusion::common::exec_err;
+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};
@@ -24,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,
 };
@@ -94,6 +99,85 @@ impl PyDataset {
         url.to_string()
     }
 
+    #[getter]
+    fn partition_url_udf(
+        self_: PyRef<'_, Self>
+    ) -> PyResult<Py<PyAny>> {
+
+        let super_ = self_.as_super();
+        let connection = super_.client.borrow(self_.py()).connection().clone();
+        let py = self_.py();
+
+        #[pyclass]
+        struct PartitionUrlInner {
+            pub connection: ConnectionHandle,
+            pub dataset_id: Tuid,
+        }
+
+        #[pymethods]
+        impl PartitionUrlInner {
+            pub fn __call__(&self, py: Python<'_>, partition_id_expr: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
+
+                let mut url = re_uri::DatasetDataEndpoint {
+                    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),
+                }
+            }
+        }
+
+        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.
     fn register(self_: PyRef<'_, Self>, recording_uri: String) -> PyResult<()> {
         let super_ = self_.as_super();