diff --git a/development_docs/architecture/execution-and-caching.md b/development_docs/architecture/execution-and-caching.md index 33f99aa..8168b59 100644 --- a/development_docs/architecture/execution-and-caching.md +++ b/development_docs/architecture/execution-and-caching.md @@ -160,6 +160,10 @@ unchanged. For an owned export graph, the adapter can: - add one frozen producer-dependency input to native cell hashing +- retry a failed native hash with a digest of the type and Arrow IPC stream of + each referenced PyArrow value, because Marimo's NumPy view rejects + object-typed columns such as strings. Cells whose native hash succeeds keep + their native keys. - record the effective hit or miss for authored and projection cells - run complete-cell owners and selected exporter leaves live when their output contract includes uncached side effects or session-bound resources diff --git a/docs/guide/choose-states.md b/docs/guide/choose-states.md index 548bf5d..340edd6 100644 --- a/docs/guide/choose-states.md +++ b/docs/guide/choose-states.md @@ -167,7 +167,7 @@ Each output has one source kind: | -------------- | --------------------------------------------------------------------- | ---------------------------------------------- | | `kind: json` | Canonical portable JSON | Python, browser, agent, or custom client | | `kind: native` | marimo scalar, JSON, NumPy, Arrow, or `BlobAsset` representation | Typed Python or browser loader | -| `kind: export` | `BlobAsset` returned by one declared exporter | Chart, table, media, or domain-specific loader | +| `kind: export` | `BlobAsset` or canonical JSON returned by one declared exporter | Chart, table, media, or domain-specific loader | | `kind: output` | Formatted marimo output and replay resources | marimo-aware browser application | | `kind: cell` | Cell identity, terminal output, console records, and replay resources | marimo-aware browser application or agent | diff --git a/docs/reference/export-spec.md b/docs/reference/export-spec.md index aafaeff..b6401a2 100644 --- a/docs/reference/export-spec.md +++ b/docs/reference/export-spec.md @@ -184,7 +184,8 @@ exporter: altair.vegalite ``` An export source passes the selected value to one declared exporter. The -exporter returns a `BlobAsset` with bytes, media type, filename, and metadata. +exporter returns a `BlobAsset` with bytes, media type, filename, and metadata, +or a JSON value that the export stores as canonical `marimo.json.v1` JSON. ### Rendered-output source diff --git a/docs/reference/python/produce.md b/docs/reference/python/produce.md index 7134d4a..631b4d7 100644 --- a/docs/reference/python/produce.md +++ b/docs/reference/python/produce.md @@ -145,7 +145,7 @@ OutputSpec.cell(name: str | None = None, *, id: str | None = None) -> OutputSpec | ---------- | ---------------------------------------------------------------------------------------------- | | `json()` | Canonical portable JSON selected from a notebook definition | | `native()` | marimo cache representation for a scalar, JSON value, NumPy array, Arrow table, or `BlobAsset` | -| `export()` | `BlobAsset` returned by an explicit exporter | +| `export()` | `BlobAsset` or canonical JSON returned by an explicit exporter | | `output()` | Formatted `marimo.output.v1` snapshot and replay resources | | `cell()` | Complete `marimo.cell.v1` snapshot selected by authored cell name or inspected runtime ID | @@ -222,7 +222,7 @@ modules, so restart it after changing custom exporter source. ### `BlobAsset` -Custom exporters return `marimo_export.outputs.BlobAsset`: +Custom exporters return `marimo_export.outputs.BlobAsset` or a JSON value: ```python from marimo_export.outputs import BlobAsset @@ -253,6 +253,20 @@ canonical encoding is limited to 256 KiB. Supply `media_type` for a value that will enter a notebook export. Export production rejects a `BlobAsset` whose media type is absent or invalid. +An exporter that returns a JSON value produces the canonical `marimo.json.v1` +representation of `OutputSpec.json()`, and readers return it through +`output.json()`: + +```python +def summarize(value) -> dict[str, object]: + return {"rows": value.num_rows, "columns": value.column_names} +``` + +One output keeps one codec across states, so an exporter returns a `BlobAsset` +for every state or a JSON value for every state. Any other result raises +`OutputError` with code `output_execution_failed` and `exception_type` +`TypeError`. + ## `plan()` ```python diff --git a/docs/reference/terminology.md b/docs/reference/terminology.md index 0b69be3..1ab0637 100644 --- a/docs/reference/terminology.md +++ b/docs/reference/terminology.md @@ -50,7 +50,7 @@ The concept pages introduce them through worked examples. | Output | One published name and representation available for every exported state. | | Output source | The `json`, `native`, `export`, `output`, or `cell` selection declared by an output spec. | | Selector | A path from one Python definition through supported attribute or item steps to a selected notebook result. | -| Exporter | A producer-side converter that returns a `BlobAsset` for one selected value. | +| Exporter | A producer-side converter that returns a `BlobAsset` or a JSON value for one selected value. | | Output plan | The complete set of authored output declarations. Its identity changes when an output source, exporter, option, or declared dependency changes. | | Output representation | The codec and media type that define how one output is stored and decoded. One output name keeps the same representation across every state. | | Codec | A versioned identifier for the native storage envelope, such as `marimo.json.v1` or `numpy.npy.v1`. | diff --git a/packages/python/src/marimo_export/_marimo/blob.py b/packages/python/src/marimo_export/_marimo/blob.py index 01d1850..d2a41e1 100644 --- a/packages/python/src/marimo_export/_marimo/blob.py +++ b/packages/python/src/marimo_export/_marimo/blob.py @@ -2,10 +2,8 @@ from __future__ import annotations -from marimo_export.outputs import BlobAsset - -def to_native_blob_asset(value: BlobAsset) -> object: +def to_native_blob_asset(value: object) -> object: from marimo_export._marimo.compat.blob import to_native_blob_asset as convert return convert(value) diff --git a/packages/python/src/marimo_export/_marimo/compat/blob.py b/packages/python/src/marimo_export/_marimo/compat/blob.py index 86e6b32..21867d2 100644 --- a/packages/python/src/marimo_export/_marimo/compat/blob.py +++ b/packages/python/src/marimo_export/_marimo/compat/blob.py @@ -2,17 +2,32 @@ from __future__ import annotations -from marimo_export._json import portable_json_object +from marimo_export._json import canonical_bytes, json_value, portable_json_object +from marimo_export.descriptors import JSON_CODEC, JSON_MEDIA_TYPE from marimo_export.outputs import BlobAsset -def to_native_blob_asset(value: BlobAsset) -> object: - """Return the native value required by Marimo's lazy ``.bin`` codec.""" +def to_native_blob_asset(value: object) -> object: + """Return the native value required by Marimo's lazy ``.bin`` codec. + + An exporter returns a ``BlobAsset`` or a JSON value. A JSON value uses the + canonical JSON representation of a JSON source. + """ - if not isinstance(value, BlobAsset): - raise TypeError("output exporter must return marimo_export.outputs.BlobAsset") from marimo._save.stubs import BlobAsset as NativeBlobAsset + if not isinstance(value, BlobAsset): + try: + portable = json_value(value, "exporter result") + except (TypeError, ValueError) as error: + raise TypeError( + "output exporter must return marimo_export.outputs.BlobAsset or a JSON value" + ) from error + return NativeBlobAsset( + data=canonical_bytes(portable), + media_type=JSON_MEDIA_TYPE, + metadata={"schema": JSON_CODEC}, + ) return NativeBlobAsset( data=value.data, media_type=value.media_type, diff --git a/packages/python/src/marimo_export/_marimo/compat/cache/attempts.py b/packages/python/src/marimo_export/_marimo/compat/cache/attempts.py index 4615a64..1098abb 100644 --- a/packages/python/src/marimo_export/_marimo/compat/cache/attempts.py +++ b/packages/python/src/marimo_export/_marimo/compat/cache/attempts.py @@ -4,6 +4,7 @@ import ast import copy +import hashlib import threading from collections.abc import Callable, Iterator from contextlib import contextmanager @@ -132,17 +133,17 @@ def tracked( ) -> Cache: with _SCOPES_LOCK: scope = _SCOPES.get(id(graph)) + tracked_graph = scope is not None and scope.graph is graph environment = scope.environment if scope is not None and scope.graph is graph else None if environment is not None: module, scope_values = _with_environment(module, scope_values, environment) - attempt = native( - module, - graph, - cell_id, - scope_values, - *args, - **kwargs, - ) + try: + attempt = native(module, graph, cell_id, scope_values, *args, **kwargs) + except TypeError: + hashable = _with_arrow_digests(graph, cell_id, scope_values) if tracked_graph else None + if hashable is None: + raise + attempt = native(module, graph, cell_id, hashable, *args, **kwargs) with _SCOPES_LOCK: scope = _SCOPES.get(id(graph)) if scope is None or scope.graph is not graph: @@ -194,6 +195,61 @@ def _with_environment( return ast.fix_missing_locations(lookup), {**scope, name: environment} +def _with_arrow_digests( + graph: Any, + cell_id: Any, + scope: dict[str, Any], +) -> dict[str, Any] | None: + # Marimo hashes Arrow data through NumPy, which cannot view object-typed + # columns such as strings as bytes. After that failure the native hasher + # receives a digest of each referenced Arrow value's type and IPC stream, + # and execution keeps the value. Cells whose native hash succeeds keep + # their native keys. + cell = graph.cells.get(cell_id) + if cell is None: + return None + digests = { + name: digest + for name in cell.refs + if name in scope and (digest := _arrow_digest(scope[name])) is not None + } + return {**scope, **digests} if digests else None + + +def _arrow_digest(value: object) -> str | None: + if not type(value).__module__.startswith("pyarrow"): + return None + import pyarrow as pa + + kind = type(value).__qualname__ + if isinstance(value, (pa.Array, pa.ChunkedArray)): + value = pa.Table.from_arrays([value], names=["value"]) + elif not isinstance(value, (pa.Table, pa.RecordBatch)): + return None + sink = _DigestSink() + with pa.ipc.new_stream(sink, value.schema) as writer: + writer.write(value) + return f"arrow-ipc-sha256:{kind}:{sink.digest.hexdigest()}" + + +class _DigestSink: + """A write-only file that hashes the Arrow IPC stream written to it.""" + + def __init__(self) -> None: + self.digest = hashlib.sha256() + self.closed = False + + def write(self, data: bytes) -> int: + self.digest.update(data) + return len(data) + + def flush(self) -> None: + pass + + def close(self) -> None: + self.closed = True + + def record_cache_miss(graph: Any, cell_id: Any) -> None: """Record a live run chosen after native restoration.""" diff --git a/packages/python/src/marimo_export/_marimo/compat/cache/lifecycle.py b/packages/python/src/marimo_export/_marimo/compat/cache/lifecycle.py index 90c43fd..05941db 100644 --- a/packages/python/src/marimo_export/_marimo/compat/cache/lifecycle.py +++ b/packages/python/src/marimo_export/_marimo/compat/cache/lifecycle.py @@ -5,6 +5,7 @@ import time from typing import Any +from marimo._runtime.exceptions import MarimoRuntimeException from marimo._runtime.executor.lifecycles import Skip from marimo._runtime.executor.lifecycles.cached import CachedLifecycle @@ -20,7 +21,15 @@ class CompleteCachedLifecycle(CachedLifecycle): """Rerun a hit when its restored values cannot serve the live session.""" def setup(self, cell: Any, glbls: Any) -> Any: - decision = super().setup(cell, glbls) + try: + decision = super().setup(cell, glbls) + except Exception as error: + if not has_cache_scope(self._graph): + raise + # Marimo records an exception that escapes a lifecycle as an + # untyped error. A runtime exception keeps the exception on the + # cell, so export failures name its type. + raise MarimoRuntimeException from error if not has_cache_scope(self._graph): return decision if not isinstance(decision, Skip): diff --git a/packages/python/src/marimo_export/_marimo/compat/projections.py b/packages/python/src/marimo_export/_marimo/compat/projections.py index d18ed4b..9d00b29 100644 --- a/packages/python/src/marimo_export/_marimo/compat/projections.py +++ b/packages/python/src/marimo_export/_marimo/compat/projections.py @@ -50,6 +50,7 @@ class ProjectionRecording: "polars.series.series.Series", } ) +_PYARROW_TYPES = frozenset({"pyarrow.lib.RecordBatch", "pyarrow.lib.Table"}) class _RecordingPipe: @@ -220,16 +221,25 @@ def capture_native_value( def _native_arrow_value(value: object) -> object | None: python_type = f"{type(value).__module__}.{type(value).__qualname__}" - if python_type not in _POLARS_TYPES: + if python_type in _POLARS_TYPES: + frame = cast(Any, value).to_frame() if python_type.endswith(".Series") else value + buffer = BytesIO() + cast(Any, frame).write_ipc_stream(buffer, compression="uncompressed") + data = buffer.getvalue() + elif python_type in _PYARROW_TYPES: + import pyarrow as pa + + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, cast(Any, value).schema) as writer: + writer.write(value) + data = sink.getvalue().to_pybytes() + else: return None - frame = cast(Any, value).to_frame() if python_type.endswith(".Series") else value - buffer = BytesIO() - cast(Any, frame).write_ipc_stream(buffer, compression="uncompressed") from marimo._save.stubs import BlobAsset return BlobAsset( - data=buffer.getvalue(), + data=data, media_type=ARROW_MEDIA_TYPE, metadata={ "python_type": python_type, diff --git a/packages/python/src/marimo_export/_marimo/compat/receipts.py b/packages/python/src/marimo_export/_marimo/compat/receipts.py index 9ec93c8..360773f 100644 --- a/packages/python/src/marimo_export/_marimo/compat/receipts.py +++ b/packages/python/src/marimo_export/_marimo/compat/receipts.py @@ -22,6 +22,7 @@ from marimo_export._marimo.compat.projections import _NATIVE_ARROW_SCHEMA from marimo_export.descriptors import ( ARROW_MEDIA_TYPE, + JSON_CODEC, JSON_MEDIA_TYPE, MARIMO_CELL_MEDIA_TYPE, MARIMO_OUTPUT_MEDIA_TYPE, @@ -38,7 +39,13 @@ ) from marimo_export.errors import CodecError, OutputError from marimo_export.outputs import BlobAsset -from marimo_export.spec import CellSource, JsonSource, NativeSource, RenderedOutputSource +from marimo_export.spec import ( + CellSource, + ExportSource, + JsonSource, + NativeSource, + RenderedOutputSource, +) _BLOB_ASSET_PYTHON_TYPE = f"{BlobAsset.__module__}.{BlobAsset.__qualname__}" @@ -156,7 +163,14 @@ def native_receipt( payload=payload, disposition=disposition, ) - if isinstance(source, (JsonSource, NativeSource)) and cached.media_type == JSON_MEDIA_TYPE: + if cached.media_type == JSON_MEDIA_TYPE and ( + isinstance(source, (JsonSource, NativeSource)) + or ( + isinstance(source, ExportSource) + and cached.metadata == {"schema": JSON_CODEC} + and cached.filename is None + ) + ): try: value = decode_json(data, f"output {output!r} JSON projection") except (TypeError, ValueError) as error: diff --git a/packages/python/tests/test_custom_exporter_integration.py b/packages/python/tests/test_custom_exporter_integration.py index 1d3ed96..32f67cc 100644 --- a/packages/python/tests/test_custom_exporter_integration.py +++ b/packages/python/tests/test_custom_exporter_integration.py @@ -181,6 +181,54 @@ def test_capture_sideloads_an_importable_callable( assert notebook.read_bytes() == source +def test_custom_exporter_returns_a_json_value( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + notebook = tmp_path / "notebook.py" + _write_notebook(notebook) + (tmp_path / "export_exports.py").write_text( + "def describe(value):\n return {'answer': value, 'labels': ['forty', 'one']}\n", + encoding="utf-8", + ) + monkeypatch.setenv("PYTHONPATH", str(tmp_path)) + spec = ExportSpec( + default_state="baseline", + states={"baseline": {}}, + outputs={"summary": OutputSpec.export("answer", importable("export_exports:describe"))}, + ) + + _capture(notebook, spec, tmp_path / "export") + output = open_export(tmp_path / "export").state("baseline").output("summary") + + assert output.descriptor.codec == "marimo.json.v1" + assert output.json() == {"answer": 41, "labels": ("forty", "one")} + + +def test_custom_exporter_rejects_a_result_without_a_portable_form( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + notebook = tmp_path / "notebook.py" + _write_notebook(notebook) + (tmp_path / "export_exports.py").write_text( + "def describe(value):\n return object()\n", + encoding="utf-8", + ) + monkeypatch.setenv("PYTHONPATH", str(tmp_path)) + spec = ExportSpec( + default_state="baseline", + states={"baseline": {}}, + outputs={"summary": OutputSpec.export("answer", importable("export_exports:describe"))}, + ) + + with pytest.raises(OutputError) as raised: + _capture(notebook, spec, tmp_path / "export") + + assert raised.value.code == "output_execution_failed" + assert raised.value.details["exception_type"] == "TypeError" + + def test_custom_exporter_builds_are_deterministic_while_both_execute( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, diff --git a/packages/python/tests/test_marimo_cache_patch.py b/packages/python/tests/test_marimo_cache_patch.py index 3cac8db..f51febf 100644 --- a/packages/python/tests/test_marimo_cache_patch.py +++ b/packages/python/tests/test_marimo_cache_patch.py @@ -18,6 +18,7 @@ from marimo._save.loaders import PERSISTENT_LOADERS from marimo._save.stubs.lazy_stub import UnhashableStub from marimo_export._marimo.compat.cache.attempts import ( + _arrow_digest, cache_attempt_wrapper, track_managed_parent_cache, track_notebook_cache, @@ -158,6 +159,21 @@ def test_complete_lifecycle_leaves_untracked_unavailable_hits_native( assert lifecycle._attempts["cell"] is attempt +def test_complete_lifecycle_leaves_untracked_setup_failures_native( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle = cast(Any, object.__new__(CompleteCachedLifecycle)) + lifecycle._graph = object() + + def fail(self: Any, cell: Any, glbls: Any) -> None: + raise TypeError("unhashable") + + monkeypatch.setattr(CachedLifecycle, "setup", fail) + + with pytest.raises(TypeError, match="unhashable"): + lifecycle.setup(SimpleNamespace(cell_id="cell"), {}) + + def test_recreated_session_state_counts_as_an_authored_miss( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -216,3 +232,18 @@ def test_complete_lifecycle_reruns_unavailable_hits_in_managed_parent_scope( assert observed is None assert not lifecycle._attempts["cell"].hit + + +def test_arrow_hash_digests_distinguish_arrow_types_with_equal_rows() -> None: + pa = pytest.importorskip("pyarrow") + rows = {"carrier": ["AA", "DL"]} + digests = { + _arrow_digest(pa.table(rows)), + _arrow_digest(pa.RecordBatch.from_pydict(rows)), + _arrow_digest(pa.array(rows["carrier"])), + _arrow_digest(pa.chunked_array([rows["carrier"]])), + } + + assert len(digests) == 4 + assert _arrow_digest(pa.table(rows)) == _arrow_digest(pa.table(dict(rows))) + assert _arrow_digest(rows) is None diff --git a/packages/python/tests/test_native_output_integration.py b/packages/python/tests/test_native_output_integration.py index 6145b8e..77b0ca6 100644 --- a/packages/python/tests/test_native_output_integration.py +++ b/packages/python/tests/test_native_output_integration.py @@ -11,6 +11,7 @@ from marimo_export import ExportSpec, OutputSpec, open_export from marimo_export.descriptors import ArrowDescriptor, NumpyDescriptor from marimo_export.errors import OutputError +from marimo_export.exporters import importable def _write_notebook(notebook: Path, counter: Path) -> None: @@ -304,3 +305,107 @@ def _(): assert raised.value.code == "output_execution_failed" assert not (tmp_path / "export").exists() + + +def test_arrow_tables_with_string_columns_export_through_every_value_source( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + notebook = tmp_path / "notebook.py" + notebook.write_text( + """ +import marimo + +app = marimo.App() + + +@app.cell +def _(): + import pyarrow as pa + + flights = pa.table( + { + "carrier": ["AA", "DL"], + "origin": pa.array(["JFK", "LAX"]).dictionary_encode(), + "delay": [5, 25], + } + ) + return (flights,) + + +if __name__ == "__main__": + app.run() +""".lstrip(), + encoding="utf-8", + ) + (tmp_path / "flight_exports.py").write_text( + "def columns(value):\n return value.column_names\n", + encoding="utf-8", + ) + monkeypatch.setenv("PYTHONPATH", str(tmp_path)) + spec = ExportSpec( + default_state="baseline", + states={"baseline": {}}, + outputs={ + "flights": OutputSpec.native("flights"), + "columns": OutputSpec.export("flights", importable("flight_exports:columns")), + }, + ) + + result = build(notebook, spec=spec, output=tmp_path / "export", timeout=30) + warm = build(notebook, spec=spec, output=tmp_path / "warm", timeout=30) + state = open_export(result.path).state("baseline") + + # Custom exporter outputs run for every prepared state, so only the native + # output can restore from the marimo cache. + assert (warm.cache_activity.projection_hits, warm.cache_activity.projection_misses) == (1, 1) + + assert isinstance(state.output("flights").descriptor, ArrowDescriptor) + assert state.output("flights").descriptor.provenance.python_type == "pyarrow.lib.Table" + table = pa.ipc.open_stream(state.output("flights").asset_bytes()).read_all() + assert table.to_pylist() == [ + {"carrier": "AA", "origin": "JFK", "delay": 5}, + {"carrier": "DL", "origin": "LAX", "delay": 25}, + ] + assert pa.types.is_dictionary(table.schema.field("origin").type) + assert state.output("columns").json() == ("carrier", "origin", "delay") + + +def test_a_value_the_cell_cache_cannot_hash_names_its_exception(tmp_path: Path) -> None: + notebook = tmp_path / "notebook.py" + notebook.write_text( + """ +import marimo + +app = marimo.App() + + +@app.cell +def _(): + import numpy as np + + class Labels: + def __array__(self, dtype=None, copy=None): + return np.array(["alpha", "beta"], dtype=object) + + labels = Labels() + return (labels,) + + +if __name__ == "__main__": + app.run() +""".lstrip(), + encoding="utf-8", + ) + spec = ExportSpec( + default_state="baseline", + states={"baseline": {}}, + outputs={"labels": OutputSpec.json("labels")}, + ) + + with pytest.raises(OutputError) as raised: + build(notebook, spec=spec, output=tmp_path / "export", timeout=30) + + assert raised.value.code == "output_execution_failed" + assert raised.value.details["exception_type"] == "TypeError" + assert str(raised.value).endswith("with TypeError")