Make Optional Dependencies Clearer (#10696)

### Related
Closes #10695 

### What
~This add the datafusion version we specify in our examples as a lower
bound to make sure datafusion gets install with rerun before we use it.
Added a simple smoke test to repro the issue then verified it passed
with the install specification.~

Right now we have some optional dependencies specified with our package.
However we don't handle them very carefully. So importing different
parts of the package could blow up. Taking inspiration from pandas I
deferred imports so we should be able to import anything and not hit
errors unless we ACTUALLY use the thing with the dependency. It also
makes it clearer how to resolve things since it wasn't obvious to me
that our notebook was already incorporated as an optional dependency.
`python -c "import rerun.notebook"` just errors.

Adds `rerun[datafusion]` for the datafusion dependencies and
`rerun[all]` to get all of the non-testing features.

If we like this approach we should be able to update the couple of
different [internal
places](https://github.com/rerun-io/rerun/blob/7484e03f9a98341114c30abad49895258288df76/rerun_py/rerun_sdk/rerun/recording_stream.py#L900)
where we duplicate notebook checking to use this global check approach.

Author: ntjohnson1

docs/snippets/all/archetypes/image_advanced.py

@@ -22,14 +22,14 @@
 rr.log("from_file", rr.EncodedImage(path=file_path))
 
 # Read with Pillow and NumPy, and log the image.
-image = np.array(PILImage.open(file_path))
-rr.log("from_pillow_rgba", rr.Image(image))
+image_data = np.array(PILImage.open(file_path))
+rr.log("from_pillow_rgba", rr.Image(image_data))
 
 # Drop the alpha channel from the image.
-image_rgb = image[..., :3]
+image_rgb = image_data[..., :3]
 rr.log("from_pillow_rgb", rr.Image(image_rgb))
 
 # Read with OpenCV.
-image = cv2.imread(file_path)
+image_cv = cv2.imread(file_path)
 # OpenCV uses BGR ordering, we need to make this known to Rerun.
-rr.log("from_opencv", rr.Image(image, color_model="BGR"))
+rr.log("from_opencv", rr.Image(image_cv, color_model="BGR"))

pixi.lock

@@ -685,7 +685,7 @@ numpy = ">=2"   # Rerun still needs numpy <2. Enforce this outside of the pypi d
 # This is very similar to `pixi run py-build`, and dispatches to maturin by way of PEP621.
 # However, pixi doesn't know how to track the rust dependencies of the python package, so
 # you still need to `pixi run py-build` in the correct environment if you change the rust code.
-rerun-sdk = { path = "rerun_py", editable = true }
+rerun-sdk = { path = "rerun_py", editable = true, extras = ["all"] }
 
 # The same applies to the notebook.
 # However, in order to build rerun-notebook and thus to even activate environments with the
@@ -735,7 +735,7 @@ numpy = ">=2"   # Rerun still needs numpy <2. Enforce this outside of the pypi d
 
 
 [feature.python-pypi.pypi-dependencies]
-rerun-sdk = "==0.24.0a6"
+rerun-sdk = { version = "==0.24.0a6", extras = ["all"] }
 rerun-notebook = "==0.24.0a6"
 
 # EXAMPLES ENVIRONMENT
@@ -759,7 +759,7 @@ umap-learn = "==0.5.7"
 [feature.examples-common.pypi-dependencies]
 # External deps
 av = ">=14.2.0"
-datafusion = "==45.2.0"
+datafusion = "==47.0.0"
 jupyter = ">=1.0"
 polars = ">=0.12.0"
 

pixi.toml

@@ -37,6 +37,8 @@ text = "MIT OR Apache-2.0"
 [project.optional-dependencies]
 tests = ["pytest==7.1.2"]
 notebook = ["rerun-notebook==0.25.0-alpha.1+dev"]
+datafusion = ["datafusion==47.0.0"]
+all = ["notebook", "datafusion"]
 
 [project.urls]
 documentation = "https://www.rerun.io/docs"

rerun_py/pyproject.toml

@@ -267,3 +267,14 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
         return cast(T, wrapper)
 
     return decorator
+
+
+class RerunOptionalDependencyError(ImportError):
+    """Raised when an optional dependency is not installed."""
+
+    def __init__(self, package: str, optional_dep: str) -> None:
+        super().__init__(
+            f"'{package}' could not be imported. "
+            f"Please install it, or install rerun as rerun[{optional_dep}]/rerun[all] "
+            "to use this functionality."
+        )

rerun_py/rerun_sdk/rerun/error_utils.py

@@ -19,14 +19,20 @@
 
 
 from rerun import bindings
-from rerun_notebook import ErrorWidget as _ErrorWidget, Viewer as _Viewer
+from rerun.error_utils import RerunOptionalDependencyError
 
 from .event import (
     ViewerEvent as ViewerEvent,
     _viewer_event_from_json_str,
 )
 from .recording_stream import RecordingStream, get_data_recording
 
+HAS_NOTEBOOK = True
+try:
+    from rerun_notebook import ErrorWidget as _ErrorWidget, Viewer as _Viewer
+except ModuleNotFoundError:
+    HAS_NOTEBOOK = False
+
 _default_width = 640
 _default_height = 480
 
@@ -111,7 +117,8 @@ def __init__(
             Defaults to `False` if `url` is provided, and `True` otherwise.
 
         """
-
+        if not HAS_NOTEBOOK:
+            raise RerunOptionalDependencyError("rerun-notebook", "notebook")
         self._error_widget = _ErrorWidget()
         self._viewer = _Viewer(
             width=width if width is not None else _default_width,

rerun_py/rerun_sdk/rerun/notebook.py

@@ -0,0 +1 @@
+"""DataFusion utilities."""

rerun_py/rerun_sdk/rerun/utilities/datafusion/__init__.py

@@ -2,9 +2,16 @@
 
 import pyarrow as pa
 import pyarrow.compute
-from datafusion import Expr, ScalarUDF, col, udf
 from rerun_bindings import DatasetEntry
 
+from rerun.error_utils import RerunOptionalDependencyError
+
+HAS_DATAFUSION = True
+try:
+    from datafusion import Expr, ScalarUDF, col, udf
+except ModuleNotFoundError:
+    HAS_DATAFUSION = False
+
 
 def partition_url(
     dataset: DatasetEntry,
@@ -37,6 +44,8 @@ def partition_url(
         to seek along. By default this will use the same string as timestamp_col.
 
     """
+    if not HAS_DATAFUSION:
+        raise RerunOptionalDependencyError("datafusion", "datafusion")
     if partition_id_col is None:
         partition_id_col = col("rerun_partition_id")
     if isinstance(partition_id_col, str):
@@ -63,6 +72,8 @@ def partition_url_udf(dataset: DatasetEntry) -> ScalarUDF:
     This function will generate a UDF that expects one column of input,
     a string containing the Partition ID.
     """
+    if not HAS_DATAFUSION:
+        raise RerunOptionalDependencyError("datafusion", "datafusion")
 
     def inner_udf(partition_id_arr: pa.Array) -> pa.Array:
         return pa.compute.binary_join_element_wise(
@@ -81,6 +92,8 @@ def partition_url_with_timeref_udf(dataset: DatasetEntry, timeline_name: str) ->
     This function will generate a UDF that expects two columns of input,
     a string containing the Partition ID and the timestamp in nanoseconds.
     """
+    if not HAS_DATAFUSION:
+        raise RerunOptionalDependencyError("datafusion", "datafusion")
 
     def inner_udf(partition_id_arr: pa.Array, timestamp_arr: pa.Array) -> pa.Array:
         timestamp_us = pa.compute.cast(timestamp_arr, pa.timestamp("us"))

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

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from unittest.mock import Mock, patch
+
+import pytest
+from rerun.error_utils import RerunOptionalDependencyError
+
+
+def test_smoke() -> None:
+    """Just check that we can import the module."""
+
+
+def test_partition_url_import_normal() -> None:
+    """Check that we can import partition_url when datafusion is available."""
+    from rerun.utilities.datafusion.functions.url_generation import partition_url
+
+    assert partition_url is not None
+
+
+def test_partition_url_without_datafusion() -> None:
+    """Check that calling partition_url raises RerunOptionalDependencyError when datafusion is unavailable."""
+    # Mock the import to make datafusion unavailable
+    with patch.dict("sys.modules", {"datafusion": None}):
+        # Import the module - this should work
+        from rerun.utilities.datafusion.functions.url_generation import partition_url
+
+        # Create a mock dataset for testing
+        mock_dataset = Mock()
+
+        # But calling the function should raise an error
+        with pytest.raises(RerunOptionalDependencyError, match="'datafusion' could not be imported"):
+            partition_url(mock_dataset)

rerun_py/tests/unit/test_utilities_datafusion.py

@@ -86,7 +86,7 @@ def build_image_stack(image: Image) -> list[tuple[int | None, Image]]:
     return image_stack
 
 
-def image_from_clipboard() -> Image:
+def image_from_clipboard() -> Image | None:
     """
     Get image from the clipboard.
 
@@ -114,7 +114,12 @@ def image_from_clipboard() -> Image:
         os.unlink(filepath)
         return im
     else:
-        return PIL.ImageGrab.grabclipboard()
+        # On windows might return a list, of files,
+        # so return None signaling no image found.
+        content = PIL.ImageGrab.grabclipboard()
+        if isinstance(content, list):
+            return None
+        return content
 
 
 class Uploader: