Implement segment url udf in Rust

### Related

- tiny part of
https://linear.app/rerun/issue/RR-3568/udf-for-creating-links-is-incomplete

### What

This refactors the existing `segment_url()` stuff so that the UDF is
implement in rust side, because:
- the python implementation was really broken
- paves the way to add more features

Source-Ref: cddfb14e9a78cff092eb1cf9131fe060e51d87ab

Author: abey79

Cargo.lock

@@ -10955,6 +10955,7 @@ dependencies = [
  "re_server",
  "re_sorbet",
  "re_tuid",
+ "re_types_core",
  "re_uri",
  "re_video",
  "re_web_viewer_server",

docs/content/reference/migration/migration-0-30.md

@@ -5,13 +5,37 @@ order: 980
 
 ## 🐍 Python API
 
-### Deprecated UDF have been removed
+### `segment_url_udf` and `segment_url_with_timeref_udf` have been removed
 
-The deprecated `partition_url`, `partition_url_udf`, and `partition_url_with_timeref_udf` functions in
-`rerun.utilities.datafusion.functions.url_generation` have been removed. Use the `segment_url` equivalents instead:
+The `segment_url_udf()` and `segment_url_with_timeref_udf()` functions in
+`rerun.utilities.datafusion.functions.url_generation` have been removed. Use `segment_url()` instead,
+which covers both use cases:
 
-| Removed                            | Replacement                      |
-|------------------------------------|----------------------------------|
-| `partition_url()`                  | `segment_url()`                  |
-| `partition_url_udf()`              | `segment_url_udf()`              |
-| `partition_url_with_timeref_udf()` | `segment_url_with_timeref_udf()` |
+Before:
+
+```python
+from rerun.utilities.datafusion.functions.url_generation import segment_url_udf, segment_url_with_timeref_udf
+
+# Without timestamp
+udf = segment_url_udf(dataset)
+df.with_column("url", udf(col("rerun_segment_id")))
+
+# With timestamp
+udf = segment_url_with_timeref_udf(dataset, "my_timeline")
+df.with_column("url", udf(col("rerun_segment_id"), col("ts"), lit("my_timeline")))
+```
+
+After:
+
+```python
+from rerun.utilities.datafusion.functions.url_generation import segment_url
+
+# Without timestamp
+df.with_column("url", segment_url(dataset))
+
+# With timestamp
+df.with_column("url", segment_url(dataset, timestamp_col="ts", timeline_name="my_timeline"))
+```
+
+Also, the previously deprecated `partition_url()`, `partition_url_udf()`, and `partition_url_with_timeref_udf()`
+function have been removed.

rerun_py/Cargo.toml

@@ -71,6 +71,7 @@ re_memory.workspace = true
 re_sdk = { workspace = true, features = ["data_loaders"] }
 re_sorbet.workspace = true
 re_tuid.workspace = true
+re_types_core.workspace = true
 re_uri.workspace = true
 re_video.workspace = true
 re_web_viewer_server = { workspace = true, optional = true }

rerun_py/rerun_bindings/rerun_bindings.pyi

@@ -926,6 +926,9 @@ class EntryId:
     def __str__(self) -> str:
         """Return str(self)."""
 
+    def as_bytes(self) -> bytes:
+        """Return the raw 16-byte representation."""
+
 class EntryKind:
     """The kinds of entries that can be stored in the catalog."""
 
@@ -1057,6 +1060,12 @@ class DatasetEntryInternal:
         unsafe_allow_recent_cleanup: bool = False,
     ) -> None: ...
 
+class SegmentUrlUdfInternal:
+    """Rust-backed ScalarUDF for building segment URLs."""
+
+    def __datafusion_scalar_udf__(self) -> Any:
+        """Scalar UDF pycapsule."""
+
 class DatasetViewInternal:
     """Internal Rust implementation of DatasetView."""
 

rerun_py/rerun_sdk/rerun/utilities/datafusion/functions/url_generation.py

@@ -3,7 +3,6 @@
 from typing import TYPE_CHECKING
 
 import pyarrow as pa
-import pyarrow.compute
 
 from rerun.error_utils import RerunMissingDependencyError
 
@@ -12,7 +11,7 @@
 
 HAS_DATAFUSION = True
 try:
-    from datafusion import Expr, ScalarUDF, col, udf
+    from datafusion import Expr, ScalarUDF, col, lit
 except ModuleNotFoundError:
     HAS_DATAFUSION = False
 
@@ -51,73 +50,35 @@ def segment_url(
     """
     if not HAS_DATAFUSION:
         raise RerunMissingDependencyError("datafusion", "datafusion")
+
     if segment_id_col is None:
         segment_id_col = col("rerun_segment_id")
     if isinstance(segment_id_col, str):
         segment_id_col = col(segment_id_col)
 
+    rust_udf = _make_rust_udf()
+
+    origin_expr = lit(dataset.catalog.url)
+    entry_id_expr = lit(pa.scalar(dataset.id.as_bytes(), type=pa.binary(16)))
+
     if timestamp_col is not None:
         if timeline_name is None:
             timeline_name = str(timestamp_col)
 
         if isinstance(timestamp_col, str):
             timestamp_col = col(timestamp_col)
 
-        inner_udf = segment_url_with_timeref_udf(dataset, timeline_name)
-        return inner_udf(segment_id_col, timestamp_col).alias("segment_url_with_timestamp")
-
-    inner_udf = segment_url_udf(dataset)
-    return inner_udf(segment_id_col).alias("segment_url")
-
-
-def segment_url_udf(dataset: DatasetEntry) -> ScalarUDF:
-    """
-    Create a UDF to the URL for a segment within a Dataset.
-
-    This function will generate a UDF that expects one column of input,
-    a string containing the segment ID.
-    """
-    if not HAS_DATAFUSION:
-        raise RerunMissingDependencyError("datafusion", "datafusion")
-
-    def inner_udf(segment_id_arr: pa.Array) -> pa.Array:
-        return pa.compute.binary_join_element_wise(
-            dataset.segment_url(""),
-            segment_id_arr,
-            "",  # Required for join
+        return rust_udf(origin_expr, entry_id_expr, segment_id_col, timestamp_col, lit(timeline_name)).alias(
+            "segment_url"
         )
 
-    return udf(inner_udf, [pa.string()], pa.string(), "stable")
+    return rust_udf(origin_expr, entry_id_expr, segment_id_col, lit(None), lit(None)).alias("segment_url")
 
 
-def segment_url_with_timeref_udf(dataset: DatasetEntry, timeline_name: str) -> ScalarUDF:
-    """
-    Create a UDF to the URL for a segment within a Dataset with timestamp.
-
-    This function will generate a UDF that expects two columns of input,
-    a string containing the segment ID and the timestamp in nanoseconds.
-    """
+def _make_rust_udf() -> ScalarUDF:
     if not HAS_DATAFUSION:
         raise RerunMissingDependencyError("datafusion", "datafusion")
 
-    def inner_udf(segment_id_arr: pa.Array, timestamp_arr: pa.Array) -> pa.Array:
-        # The choice of `ceil_temporal` is important since this timestamp drives a cursor
-        # selection. Due to Rerun latest-at semantics, in order for data from the provided
-        # timestamp to be visible, the cursor must be set to a point in time which is
-        # greater than or equal to the target.
-        timestamp_us = pa.compute.ceil_temporal(timestamp_arr, unit="microsecond")
-
-        timestamp_us = pa.compute.strftime(
-            timestamp_us,
-            "%Y-%m-%dT%H:%M:%SZ",
-        )
-
-        return pa.compute.binary_join_element_wise(
-            dataset.segment_url(""),
-            segment_id_arr,
-            f"#when={timeline_name}@",
-            timestamp_us,
-            "",  # Required for join
-        )
+    from rerun_bindings import SegmentUrlUdfInternal  # type: ignore[attr-defined]
 
-    return udf(inner_udf, [pa.string(), pa.timestamp("ns")], pa.string(), "stable")
+    return ScalarUDF.from_pycapsule(SegmentUrlUdfInternal())

rerun_py/src/catalog/entry.rs

@@ -32,6 +32,11 @@ impl PyEntryId {
     pub fn __str__(&self) -> String {
         self.id.to_string()
     }
+
+    /// Return the raw 16-byte representation.
+    pub fn as_bytes<'py>(&self, py: Python<'py>) -> pyo3::Bound<'py, pyo3::types::PyBytes> {
+        pyo3::types::PyBytes::new(py, &self.id.id.as_bytes())
+    }
 }
 
 impl From<EntryId> for PyEntryId {

rerun_py/src/catalog/mod.rs

@@ -13,6 +13,7 @@ mod index_columns;
 mod indexes;
 mod registration_handle;
 mod schema;
+mod segment_url_udf;
 mod table_entry;
 mod table_provider_adapter;
 mod trace_context;
@@ -37,6 +38,7 @@ pub use self::indexes::{
 };
 pub use self::registration_handle::{PyRegistrationHandleInternal, PyRegistrationIterator};
 pub use self::schema::PySchemaInternal;
+pub use self::segment_url_udf::PySegmentUrlUdfInternal;
 pub use self::table_entry::{PyTableEntryInternal, PyTableInsertModeInternal};
 pub use self::table_provider_adapter::PyTableProviderAdapterInternal;
 pub use self::type_aliases::{AnyComponentColumn, IndexValuesLike, PyIndexValuesLikeInternal};
@@ -54,6 +56,7 @@ pub(crate) fn register(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
     m.add_class::<PyRegistrationHandleInternal>()?;
     m.add_class::<PyRegistrationIterator>()?;
     m.add_class::<PyTableProviderAdapterInternal>()?;
+    m.add_class::<PySegmentUrlUdfInternal>()?;
     m.add_class::<PyDatasetViewInternal>()?;
     m.add_class::<PyRerunHtmlTable>()?;