Skip to content

Commit

Permalink
remove obsolete code & improve error message
Browse files Browse the repository at this point in the history
  • Loading branch information
pfackeldey committed Nov 30, 2024
1 parent cfa79f5 commit bae0915
Showing 1 changed file with 13 additions and 37 deletions.
50 changes: 13 additions & 37 deletions src/dask_awkward/lib/mapfilter.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,28 +118,6 @@ def _multi_return_map_partitions(
return tuple(ret)


def _compare_return_vals(left: tp.Any, right: tp.Any) -> None:
def cmp(left, right):
msg = (
"The provided 'meta' does not match "
"the output type inferred from the pre-run step; "
"got {right}, but expected {left}.".format(left=left, right=right)
)
if isinstance(left, ak.Array):
if left.layout.form != right.layout.form:
raise ValueError(msg)

else:
if left != right:
raise ValueError(msg)

if isinstance(left, tuple) and isinstance(right, tuple):
for left_, right_ in zip(left, right):
cmp(left_, right_)
else:
cmp(left, right)


class UntraceableFunctionError(Exception): ...


Expand Down Expand Up @@ -188,7 +166,7 @@ def prerun(
fn: tp.Callable, *args: tp.Any, **kwargs: tp.Any
) -> tuple[tp.Any, tp.Mapping]:
"""
Pre-runs the provided function with typetracer arrays to determine the necessary columns/slices
Pre-runs the provided function with typetracer arrays to determine the necessary columns
that should be touched explicitly and to infer the metadata of the function's output.
Parameters
Expand Down Expand Up @@ -264,17 +242,15 @@ def _render_buffer_key(
needs = dict(_reports2needs(reports=reports))

msg = (
f"This wrapped function '{fn}' is not traceable. "
f"An error occurred at line {line_number}.\n"
"'mapfilter' can circumvent this by providing the 'needs' and "
"'meta' arguments to the decorator.\n"
"\n- 'needs': mapping where the keys point to input argument "
"dask_awkward arrays and the values to columns/slices that "
"should be touched explicitly. The typetracing step could "
"determine the following necessary columns/slices.\n\n"
f"Pre-running reported the following 'needs':\n"
f"\n{needs}\n"
"\n- 'meta': value(s) of what the wrapped function would "
f"'{fn}' is not traceable, an error occurred at line {line_number}. "
"'dak.mapfilter' can circumvent this by providing 'needs' and "
"'meta' arguments to it.\n\n"
"- 'needs': mapping where the keys point to input argument "
"dask_awkward arrays and the values to columns that should "
"be touched explicitly. The typetracing step could determine "
"the following necessary columns until the exception occurred:\n\n"
f"{needs=}\n\n"
"- 'meta': value(s) of what the wrapped function would "
"return. For arrays, only the shape and type matter."
)
raise UntraceableFunctionError(msg) from err
Expand Down Expand Up @@ -310,7 +286,7 @@ class mapfilter:
needs : dict, optional
If ``None`` (the default), nothing is touched in addition to the standard typetracer report.
In certain cases, it is necessary to touch additional objects **explicitly** to get the correct typetracer report.
For this, provide a dictionary that maps input arguments that are arrays to the columns/slices of that array that should be touched.
For this, provide a dictionary that maps input arguments that are arrays to the columns of that array that should be touched.
If ``needs`` is used together with ``meta``, **only** the columns provided by the ``needs`` argument will be touched explicitly.
"""

Expand All @@ -326,8 +302,8 @@ def __post_init__(self) -> None:
msg = ( # type: ignore[unreachable]
"'needs' argument must be a mapping where the keys "
"point to input argument dask_awkward arrays and the values "
"to columns/slices that should be touched explicitly, "
f"got '{self.needs!r}' instead.\n\n"
"to columns that should be touched explicitly, "
f"got {self.needs!r} instead.\n\n"
"Exemplary usage:\n"
"\n@partial(mapfilter, needs={'array': ['col1', 'col2']})"
"\ndef process(array: ak.Array) -> ak.Array:"
Expand Down

0 comments on commit bae0915

Please sign in to comment.