diff --git a/src/dask_awkward/layers/layers.py b/src/dask_awkward/layers/layers.py index eef1a81a..c33a54a6 100644 --- a/src/dask_awkward/layers/layers.py +++ b/src/dask_awkward/layers/layers.py @@ -112,7 +112,11 @@ def io_func_implements_projection(func: ImplementsIOFunction) -> bool: def io_func_implements_mocking(func: ImplementsIOFunction) -> bool: - return hasattr(func, "mock") and hasattr(func, "mock_empty") + return hasattr(func, "mock") + + +def io_func_implements_mock_empty(func: ImplementsIOFunction) -> bool: + return hasattr(func, "mock_empty") def io_func_implements_columnar(func: ImplementsIOFunction) -> bool: diff --git a/src/dask_awkward/lib/io/columnar.py b/src/dask_awkward/lib/io/columnar.py index 2fee74f9..6644ea3f 100644 --- a/src/dask_awkward/lib/io/columnar.py +++ b/src/dask_awkward/lib/io/columnar.py @@ -64,7 +64,7 @@ def mock(self: S) -> AwkwardArray: def mock_empty(self: S, backend: BackendT = "cpu") -> AwkwardArray: return ak.to_backend( - self.form.length_zero_array(highlevel=False), + self.form.length_zero_array(highlevel=False, behavior=self.behavior), backend, highlevel=True, ) diff --git a/src/dask_awkward/lib/io/io.py b/src/dask_awkward/lib/io/io.py index 861cff4a..eaa26729 100644 --- a/src/dask_awkward/lib/io/io.py +++ b/src/dask_awkward/lib/io/io.py @@ -31,6 +31,7 @@ BackendT, ImplementsMocking, IOFunctionWithMocking, + io_func_implements_mock_empty, io_func_implements_mocking, ) from dask_awkward.lib.core import ( @@ -622,6 +623,8 @@ def from_map( raise ValueError("empty_on_raise and empty_backend must be used together.") if empty_on_raise and empty_backend: + if not io_func_implements_mock_empty(io_func): + raise ValueError("io_func must implement mock_empty method.") io_func = return_empty_on_raise( io_func, allowed_exceptions=empty_on_raise, diff --git a/tests/test_io.py b/tests/test_io.py index 3ee3cf2b..912f02d6 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -398,3 +398,24 @@ def test_random_fail_from_lists(): label="from-lists", empty_backend="cpu", ) + + class NoMockEmpty: + def __init__(self, x): + self.x = x + + def mock(self): + return 5 + + def __call__(self, *args): + return self.x * args[0] + + with pytest.raises(ValueError, match="must implement"): + array = from_map( + NoMockEmpty(5), + many, + meta=typetracer_array(ak.Array(many[0])), + divisions=divs, + label="from-lists", + empty_on_raise=(RuntimeError,), + empty_backend="cpu", + )