From 5e41e6c7790e98146189a62e56ea6d75881e51a8 Mon Sep 17 00:00:00 2001 From: Henry Schreiner Date: Wed, 10 Jun 2026 16:34:06 -0400 Subject: [PATCH 1/4] fix: assorted operation bugs (softmax, enforce_type, merge_union_of_records, and more) - ak.softmax: default axis was None, which always raised TypeError; default to axis=-1, update the docstring, and raise a clear error for axis=None. - ak.enforce_type: fix inverted regular->regular enforceability check in _type_is_enforceable (equal sizes now recurse, unequal sizes are rejected), matching _enforce_type. - ak.merge_union_of_records: use the outer union index (layout.index.data) when rewriting indexed/non-indexed union contents, fixing an IndexError on unions containing categorical/indexed records. - ak.argcartesian: honour the requested axis for list inputs (was dropped), matching the dict-input branch. - ak.from_raggedtensor: recurse with count + 1 to avoid RecursionError for tensors with three or more ragged dimensions. - ak.to_parquet_dataset: drop set-literal concatenation in the not-a-directory error (was raising TypeError instead of ValueError). - ak.to_layout: report the invalid string_policy value, not primitive_policy. - ak.count_nonzero: remove duplicated HighLevelContext unwrap block. - ak.to_json: remove no-op 'except Exception as err: raise err from err'. - ak.round / ak.from_safetensors: fix copy-pasted docstring / error message. - ak.unflatten / ak.nan_to_num: yield array-like secondary arguments to the __awkward_function__ dispatch protocol. - ak.from_rdataframe: thread attrs through to the output array. Assisted-by: ClaudeCode:claude-opus-4.8 --- .../_connect/rdataframe/from_rdataframe.py | 11 +- src/awkward/operations/ak_argcartesian.py | 2 +- src/awkward/operations/ak_count_nonzero.py | 2 - src/awkward/operations/ak_enforce_type.py | 6 +- .../operations/ak_from_raggedtensor.py | 2 +- src/awkward/operations/ak_from_rdataframe.py | 14 ++- src/awkward/operations/ak_from_safetensors.py | 2 +- .../operations/ak_merge_union_of_records.py | 4 +- src/awkward/operations/ak_nan_to_num.py | 2 +- src/awkward/operations/ak_round.py | 5 +- src/awkward/operations/ak_softmax.py | 20 +++- src/awkward/operations/ak_to_json.py | 110 +++++++++--------- src/awkward/operations/ak_to_layout.py | 2 +- .../operations/ak_to_parquet_dataset.py | 2 +- src/awkward/operations/ak_unflatten.py | 2 +- tests/test_2596_named_axis.py | 13 ++- 16 files changed, 111 insertions(+), 88 deletions(-) diff --git a/src/awkward/_connect/rdataframe/from_rdataframe.py b/src/awkward/_connect/rdataframe/from_rdataframe.py index ff5341ce48..97eb477363 100644 --- a/src/awkward/_connect/rdataframe/from_rdataframe.py +++ b/src/awkward/_connect/rdataframe/from_rdataframe.py @@ -63,7 +63,14 @@ def from_rdataframe( - data_frame, columns, highlevel, behavior, with_name, offsets_type, keep_order + data_frame, + columns, + highlevel, + behavior, + attrs, + with_name, + offsets_type, + keep_order, ): if hasattr(data_frame, "proxied_node"): raise NotImplementedError("Distributed RDataFrame is not yet supported") @@ -269,6 +276,7 @@ def form_dtype(form): depth_limit=1, highlevel=True, behavior=behavior, + attrs=attrs, with_name=with_name, ) @@ -278,6 +286,7 @@ def form_dtype(form): ak.contents.IndexedArray(sorted, out.layout), highlevel=True, behavior=behavior, + attrs=attrs, ) return wrap_layout(out.layout, highlevel=highlevel, behavior=behavior) diff --git a/src/awkward/operations/ak_argcartesian.py b/src/awkward/operations/ak_argcartesian.py index 82937d4f19..07f2376da1 100644 --- a/src/awkward/operations/ak_argcartesian.py +++ b/src/awkward/operations/ak_argcartesian.py @@ -118,7 +118,7 @@ def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior, attr if isinstance(arrays, Mapping): index_arrays = {n: ak.local_index(x, axis) for n, x in arrays.items()} else: - index_arrays = [ak.local_index(x) for x in arrays] + index_arrays = [ak.local_index(x, axis) for x in arrays] if with_name is not None: if parameters is None: diff --git a/src/awkward/operations/ak_count_nonzero.py b/src/awkward/operations/ak_count_nonzero.py index 409e9b82d9..16324a150a 100644 --- a/src/awkward/operations/ak_count_nonzero.py +++ b/src/awkward/operations/ak_count_nonzero.py @@ -98,8 +98,6 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): axis = regularize_axis(axis, none_allowed=True) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: - layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") reducer = ak._reducers.CountNonzero() out = ak._do.reduce( diff --git a/src/awkward/operations/ak_enforce_type.py b/src/awkward/operations/ak_enforce_type.py index a98c24e87c..fcec91a289 100644 --- a/src/awkward/operations/ak_enforce_type.py +++ b/src/awkward/operations/ak_enforce_type.py @@ -487,10 +487,8 @@ def _type_is_enforceable( elif layout.is_regular: if isinstance(type_, ak.types.RegularType): if layout.size == type_.size: - return _TypeEnforceableResult( - is_enforceable=False, requires_packing=False - ) - return _type_is_enforceable(layout.content, type_.content) + return _type_is_enforceable(layout.content, type_.content) + return _TypeEnforceableResult(is_enforceable=False, requires_packing=False) elif isinstance(type_, ak.types.ListType): return _type_is_enforceable(layout.content, type_.content) diff --git a/src/awkward/operations/ak_from_raggedtensor.py b/src/awkward/operations/ak_from_raggedtensor.py index cadf1edb12..b2a96ff168 100644 --- a/src/awkward/operations/ak_from_raggedtensor.py +++ b/src/awkward/operations/ak_from_raggedtensor.py @@ -99,5 +99,5 @@ def _recursive_call(content, offsets_arr, count): ) else: return ak.contents.ListOffsetArray( - offsets_arr[count], _recursive_call(content, offsets_arr, count) + offsets_arr[count], _recursive_call(content, offsets_arr, count + 1) ) diff --git a/src/awkward/operations/ak_from_rdataframe.py b/src/awkward/operations/ak_from_rdataframe.py index 37bae96117..7b3f85b6f1 100644 --- a/src/awkward/operations/ak_from_rdataframe.py +++ b/src/awkward/operations/ak_from_rdataframe.py @@ -53,11 +53,20 @@ def from_rdataframe( See also #ak.to_rdataframe. """ - return _impl(rdf, columns, highlevel, behavior, with_name, offsets_type, keep_order) + return _impl( + rdf, + columns, + highlevel, + behavior, + attrs, + with_name, + offsets_type, + keep_order, + ) def _impl( - data_frame, columns, highlevel, behavior, with_name, offsets_type, keep_order + data_frame, columns, highlevel, behavior, attrs, with_name, offsets_type, keep_order ): import awkward._connect.rdataframe.from_rdataframe # noqa: F401 @@ -90,6 +99,7 @@ def _impl( columns, highlevel=highlevel, behavior=behavior, + attrs=attrs, with_name=with_name, offsets_type=offsets_type, keep_order=keep_order, diff --git a/src/awkward/operations/ak_from_safetensors.py b/src/awkward/operations/ak_from_safetensors.py index f161fe8e06..57364a3f80 100644 --- a/src/awkward/operations/ak_from_safetensors.py +++ b/src/awkward/operations/ak_from_safetensors.py @@ -109,7 +109,7 @@ def _impl( from safetensors import _safe_open_handle except ImportError as err: raise ImportError( - """to use ak.from_tensorflow, you must install the 'safetensors' package with: + """to use ak.from_safetensors, you must install the 'safetensors' package with: pip install safetensors or diff --git a/src/awkward/operations/ak_merge_union_of_records.py b/src/awkward/operations/ak_merge_union_of_records.py index c0ab4f17a2..6f92439b8e 100644 --- a/src/awkward/operations/ak_merge_union_of_records.py +++ b/src/awkward/operations/ak_merge_union_of_records.py @@ -291,12 +291,12 @@ def apply(layout, depth, backend, **kwargs): # Rewrite union index of indexed types if content.is_indexed: next_index_data[is_this_tag] = content.index.data[ - content.index.data[is_this_tag] + layout.index.data[is_this_tag] ] next_contents.append(content.content) else: - next_index_data[is_this_tag] = content.index.data[is_this_tag] + next_index_data[is_this_tag] = layout.index.data[is_this_tag] next_contents.append(content) return invert_record_union( diff --git a/src/awkward/operations/ak_nan_to_num.py b/src/awkward/operations/ak_nan_to_num.py index 69e2617c00..6bc9e01126 100644 --- a/src/awkward/operations/ak_nan_to_num.py +++ b/src/awkward/operations/ak_nan_to_num.py @@ -50,7 +50,7 @@ def nan_to_num( See also #ak.nan_to_none to convert NaN to None, i.e. missing values with option-type. """ # Dispatch - yield (array,) + yield array, nan, posinf, neginf # Implementation return _impl(array, copy, nan, posinf, neginf, highlevel, behavior, attrs) diff --git a/src/awkward/operations/ak_round.py b/src/awkward/operations/ak_round.py index 6d6e0e0471..3011726dd0 100644 --- a/src/awkward/operations/ak_round.py +++ b/src/awkward/operations/ak_round.py @@ -40,8 +40,9 @@ def round( attrs (None or dict): Custom attributes for the output array, if high-level. - Returns the real components of the given array elements. - If the arrays have complex elements, the returned arrays are floats. + Rounds the given array elements to the given number of `decimals`. + Implements [np.round](https://numpy.org/doc/stable/reference/generated/numpy.round.html) + for Awkward Arrays. """ # Dispatch yield (array,) diff --git a/src/awkward/operations/ak_softmax.py b/src/awkward/operations/ak_softmax.py index 67bac6407b..1d4809db42 100644 --- a/src/awkward/operations/ak_softmax.py +++ b/src/awkward/operations/ak_softmax.py @@ -26,7 +26,7 @@ @high_level_function() def softmax( x, - axis=None, + axis=-1, *, keepdims=False, mask_identity=False, @@ -37,11 +37,12 @@ def softmax( """ Args: x: The data on which to compute the softmax (anything #ak.to_layout recognizes). - axis (None or int or str): The dimension over which the softmax is - computed. If None, the softmax is computed over the innermost - dimension. If an int, `0` is the outermost dimension, `1` is the - first level of nested lists, etc., and negative values count from - the innermost: `-1` is the innermost, `-2` is the next level up, etc. + axis (int or str): The dimension over which the softmax is + computed (default is `-1`, the innermost dimension). ak.softmax is + currently only defined for the innermost axis. If an int, `0` is the + outermost dimension, `1` is the first level of nested lists, etc., + and negative values count from the innermost: `-1` is the innermost, + `-2` is the next level up, etc. If a str, it is interpreted as the name of the axis which maps to an int if named axes are present. Named axes are attached to an array using #ak.with_named_axis and @@ -96,6 +97,13 @@ def _impl(x, axis, keepdims, mask_identity, highlevel, behavior, attrs): axis = _named_axis_to_positional_axis(named_axis, axis) axis = regularize_axis(axis, none_allowed=True) + if axis is None: + raise ValueError( + "ak.softmax requires an explicit axis (it is only defined for " + "axis=-1); axis=None is not supported. " + "See https://github.com/scikit-hep/awkward/issues/2760" + ) + x = ctx.wrap(x_layout) if maybe_posaxis(x_layout, axis, 1) != maybe_posaxis(x_layout, -1, 1): diff --git a/src/awkward/operations/ak_to_json.py b/src/awkward/operations/ak_to_json.py index 4060f8aa15..f8cfbdda1d 100644 --- a/src/awkward/operations/ak_to_json.py +++ b/src/awkward/operations/ak_to_json.py @@ -217,51 +217,64 @@ def opener(): def opener(): return _NoContextManager(file) - try: - if line_delimited: - if file is None: - out = [] - for datum in jsondata: - out.append( - json.dumps( - datum, - skipkeys=True, - ensure_ascii=True, - check_circular=False, - allow_nan=False, - indent=None, - separators=separators, - default=convert_other, - sort_keys=False, - ) + if line_delimited: + if file is None: + out = [] + for datum in jsondata: + out.append( + json.dumps( + datum, + skipkeys=True, + ensure_ascii=True, + check_circular=False, + allow_nan=False, + indent=None, + separators=separators, + default=convert_other, + sort_keys=False, ) - out.append(line_delimited) - return "".join(out) - - else: - with opener() as openfile: - for datum in jsondata: - json.dump( - datum, - openfile, - skipkeys=True, - ensure_ascii=True, - check_circular=False, - allow_nan=False, - indent=None, - separators=separators, - default=convert_other, - sort_keys=False, - ) - openfile.write(line_delimited) + ) + out.append(line_delimited) + return "".join(out) else: - if isinstance(array, (ak.highlevel.Record, ak.record.Record)): - jsondata = jsondata[0] + with opener() as openfile: + for datum in jsondata: + json.dump( + datum, + openfile, + skipkeys=True, + ensure_ascii=True, + check_circular=False, + allow_nan=False, + indent=None, + separators=separators, + default=convert_other, + sort_keys=False, + ) + openfile.write(line_delimited) - if file is None: - return json.dumps( + else: + if isinstance(array, (ak.highlevel.Record, ak.record.Record)): + jsondata = jsondata[0] + + if file is None: + return json.dumps( + jsondata, + skipkeys=True, + ensure_ascii=True, + check_circular=False, + allow_nan=False, + indent=num_indent_spaces, + separators=separators, + default=convert_other, + sort_keys=False, + ) + else: + with opener() as openfile: + return json.dump( jsondata, + openfile, skipkeys=True, ensure_ascii=True, check_circular=False, @@ -271,23 +284,6 @@ def opener(): default=convert_other, sort_keys=False, ) - else: - with opener() as openfile: - return json.dump( - jsondata, - openfile, - skipkeys=True, - ensure_ascii=True, - check_circular=False, - allow_nan=False, - indent=num_indent_spaces, - separators=separators, - default=convert_other, - sort_keys=False, - ) - - except Exception as err: - raise err from err class _NoContextManager: diff --git a/src/awkward/operations/ak_to_layout.py b/src/awkward/operations/ak_to_layout.py index 99b22ff0ce..74b33d5f9a 100644 --- a/src/awkward/operations/ak_to_layout.py +++ b/src/awkward/operations/ak_to_layout.py @@ -259,7 +259,7 @@ def _impl( ) else: raise ValueError( - f"Encountered an invalid string policy value {primitive_policy!r}. " + f"Encountered an invalid string policy value {string_policy!r}. " f'The permitted values are "pass-through", "as-characters", "promote", and "error".' ) elif isinstance(obj, (datetime, date, time, Number, bool)): diff --git a/src/awkward/operations/ak_to_parquet_dataset.py b/src/awkward/operations/ak_to_parquet_dataset.py index af7a4053e9..836ecfce7d 100644 --- a/src/awkward/operations/ak_to_parquet_dataset.py +++ b/src/awkward/operations/ak_to_parquet_dataset.py @@ -58,7 +58,7 @@ def _impl(directory, filenames, storage_options): ) from None fs, destination = fsspec.core.url_to_fs(directory, **(storage_options or {})) if not fs.isdir(destination): - raise ValueError(f"{destination!r} is not a directory" + {__file__}) + raise ValueError(f"{destination!r} is not a directory") filepaths = get_filepaths(filenames, fs, destination) diff --git a/src/awkward/operations/ak_unflatten.py b/src/awkward/operations/ak_unflatten.py index 42df503fe8..221ef0ad7c 100644 --- a/src/awkward/operations/ak_unflatten.py +++ b/src/awkward/operations/ak_unflatten.py @@ -97,7 +97,7 @@ def unflatten(array, counts, axis=0, *, highlevel=True, behavior=None, attrs=Non See also #ak.num and #ak.flatten. """ # Dispatch - yield (array,) + yield array, counts # Implementation return _impl(array, counts, axis, highlevel, behavior, attrs) diff --git a/tests/test_2596_named_axis.py b/tests/test_2596_named_axis.py index 59b78cf233..db8d33db80 100644 --- a/tests/test_2596_named_axis.py +++ b/tests/test_2596_named_axis.py @@ -689,21 +689,24 @@ def test_named_axis_ak_argcartesian(): named_two = ak.with_named_axis(two, named_axis=("x", "y")) named_three = ak.with_named_axis(three, named_axis=("x", "y")) + # The list input now honours the requested axis (matching the dict input + # branch), so the named-axis output agrees with ak.cartesian/argcartesian + # of an equivalent mapping. assert ak.argcartesian( [named_one, named_two, named_three], axis="x", nested=False - ).named_axis == {"x": 0, "y": 1} + ).named_axis == {"x": 0} assert ak.argcartesian( [named_one, named_two, named_three], axis="x", nested=True - ).named_axis == {"x": 1, "y": 2} + ).named_axis == {"x": 1} assert ak.argcartesian( [named_one, named_two, named_three], axis="x", nested=[0] - ).named_axis == {"x": 1, "y": 2} + ).named_axis == {"x": 1} assert ak.argcartesian( [named_one, named_two, named_three], axis="x", nested=[1] - ).named_axis == {"x": 0, "y": 2} + ).named_axis == {"x": 0} assert ak.argcartesian( [named_one, named_two, named_three], axis="x", nested=[0, 1] - ).named_axis == {"x": 2, "y": 3} + ).named_axis == {"x": 2} def test_negative_named_axis_ak_argcartesian(): From 3e125f1af4d0b7354d838140f3e7439bf34a0ccf Mon Sep 17 00:00:00 2001 From: Henry Schreiner Date: Wed, 10 Jun 2026 16:35:00 -0400 Subject: [PATCH 2/4] test: add regression tests for assorted operation fixes Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_4107_operations_batch.py | 162 ++++++++++++++++++++++++++++ 1 file changed, 162 insertions(+) create mode 100644 tests/test_4107_operations_batch.py diff --git a/tests/test_4107_operations_batch.py b/tests/test_4107_operations_batch.py new file mode 100644 index 0000000000..f9edeaa16c --- /dev/null +++ b/tests/test_4107_operations_batch.py @@ -0,0 +1,162 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +from __future__ import annotations + +import numpy as np +import pytest + +import awkward as ak + + +def test_softmax_default_axis(): + # Previously ak.softmax with the default axis raised TypeError because the + # default was axis=None, which maybe_posaxis cannot handle. + array = ak.Array([[1.0, 2.0, 3.0], [4.0], [], [5.0, 6.0]]) + result = ak.softmax(array) + + # Compare against a numpy/scipy-style softmax computed per inner list. + expected = [] + for row in array.to_list(): + if len(row) == 0: + expected.append([]) + continue + arr = np.array(row) + e = np.exp(arr) + expected.append((e / e.sum()).tolist()) + + result_list = result.to_list() + assert len(result_list) == len(expected) + for got, exp in zip(result_list, expected, strict=True): + assert np.allclose(got, exp) + + # axis=-1 is the same as the default + assert ak.softmax(array, axis=-1).to_list() == result_list + + +def test_softmax_axis_none_raises(): + array = ak.Array([[1.0, 2.0, 3.0]]) + with pytest.raises(ValueError): + ak.softmax(array, axis=None) + + +def test_enforce_type_regular_size_equal(): + # 2 * int32 -> 2 * int64 should be enforceable (recurse into content) + layout = ak.to_regular(ak.Array(np.arange(6, dtype=np.int32).reshape(3, 2))).layout + from awkward.operations.ak_enforce_type import _type_is_enforceable + + target = ak.types.from_datashape("2 * int64", highlevel=False) + result = _type_is_enforceable(layout, target) + assert result.is_enforceable + + # The full operation must agree with _type_is_enforceable + out = ak.enforce_type(ak.Array(layout), "2 * int64") + assert str(out.type) == "3 * 2 * int64" + + +def test_enforce_type_regular_size_mismatch(): + layout = ak.to_regular(ak.Array(np.arange(6, dtype=np.int32).reshape(3, 2))).layout + from awkward.operations.ak_enforce_type import _type_is_enforceable + + target = ak.types.from_datashape("3 * int64", highlevel=False) + result = _type_is_enforceable(layout, target) + assert not result.is_enforceable + + with pytest.raises(ValueError): + ak.enforce_type(ak.Array(layout), "3 * int64") + + +def test_merge_union_of_records_categorical_indexed(): + # A union where one branch is a categorical (IndexedArray) of records. + # Build a categorical (IndexedArray) record array directly. + index = ak.index.Index64(np.array([1, 0, 1], dtype=np.int64)) + inner_records = ak.contents.RecordArray( + [ak.contents.NumpyArray(np.array([10.0, 20.0]))], ["c"] + ) + indexed = ak.contents.IndexedArray( + index, inner_records, parameters={"__array__": "categorical"} + ) + + other = ak.contents.RecordArray( + [ak.contents.NumpyArray(np.array([1, 2], dtype=np.int64))], ["a"] + ) + + tags = ak.index.Index8(np.array([0, 1, 0, 1, 0], dtype=np.int8)) + union_index = ak.index.Index64(np.array([0, 0, 1, 1, 2], dtype=np.int64)) + union = ak.contents.UnionArray(tags, union_index, [other, indexed]) + + arr = ak.Array(union) + # Previously raised IndexError due to self-indexing the inner index. + out = ak.merge_union_of_records(arr) + out_list = out.to_list() + + # tag 0 -> other (field a); tag 1 -> categorical indexed records (field c). + # The inner categorical index [1, 0, 1] selects content [20.0, 10.0, 20.0], + # consumed in order by the two tag-1 slots (union_index 0 and 1). + assert out_list == [ + {"a": 1, "c": None}, + {"a": None, "c": 20.0}, + {"a": 2, "c": None}, + {"a": None, "c": 10.0}, + {"a": None, "c": None}, + ] + assert str(out.type) == "5 * {a: ?int64, c: ?float64}" + + +def test_argcartesian_list_vs_dict_parity_axis1(): + x = ak.Array([[1, 2, 3], [4]]) + y = ak.Array([[10, 20], [30]]) + + as_dict = ak.argcartesian({"x": x, "y": y}, axis=1) + as_list = ak.argcartesian([x, y], axis=1) + + # Field names differ ("x"/"y" vs "0"/"1") but the index values must match. + dict_vals = [[(d["x"], d["y"]) for d in row] for row in as_dict.to_list()] + # The list branch produces tuples, indexed positionally. + list_vals = [[tuple(t) for t in row] for row in as_list.to_list()] + assert dict_vals == list_vals + + +def test_from_raggedtensor_recursive_deep(): + # Drive _recursive_call directly (no TensorFlow needed) for a structure with + # three ragged dimensions (3 offset arrays). Before the fix the recursive + # else-branch reused the same `count`, causing a RecursionError for >= 3 + # ragged dimensions. + from awkward.operations.ak_from_raggedtensor import _recursive_call + + content = ak.contents.NumpyArray(np.arange(8, dtype=np.float64)) + offsets = [ + ak.index.Index64(np.array([0, 2], dtype=np.int64)), + ak.index.Index64(np.array([0, 2, 4], dtype=np.int64)), + ak.index.Index64(np.array([0, 2, 4, 6, 8], dtype=np.int64)), + ] + result = ak.Array(_recursive_call(content, offsets, 0)) + assert result.to_list() == [[[[0.0, 1.0], [2.0, 3.0]], [[4.0, 5.0], [6.0, 7.0]]]] + + +class _Dispatcher: + """Dummy array-like recording __awkward_function__ invocations.""" + + def __init__(self): + self.captured = None + + def __awkward_function__(self, func, array_likes, args, kwargs): + self.captured = list(array_likes) + return "intercepted" + + +def test_unflatten_dispatch_sees_counts(): + array = ak.Array([1, 2, 3, 4]) + sentinel = _Dispatcher() + # counts is array-like and must be offered to dispatch. + result = ak.unflatten(array, sentinel) + assert result == "intercepted" + assert any(x is sentinel for x in sentinel.captured) + + +def test_nan_to_num_dispatch_sees_extra_args(): + sentinel = _Dispatcher() + array = ak.Array([1.0, 2.0]) + # Pass the dispatcher as posinf; dispatch must offer it. + result = ak.nan_to_num(array, posinf=sentinel) + assert result == "intercepted" + assert any(x is sentinel for x in sentinel.captured) From 82f6ce953fdd0d376e507b5ff6d58f82a2aacd09 Mon Sep 17 00:00:00 2001 From: Henry Schreiner Date: Thu, 11 Jun 2026 15:46:28 -0400 Subject: [PATCH 3/4] chore: slim down test file and drop _Dispatcher class - Convert _Dispatcher helper class to _make_dispatch_sentinel() factory function using types.SimpleNamespace + closure - Remove comments that merely restate the surrounding code - Rebase was clean (no conflicts) Assisted-by: ClaudeCode:claude-sonnet-4-6 --- tests/test_4107_operations_batch.py | 33 +++++++++++++++-------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/tests/test_4107_operations_batch.py b/tests/test_4107_operations_batch.py index f9edeaa16c..1994d31442 100644 --- a/tests/test_4107_operations_batch.py +++ b/tests/test_4107_operations_batch.py @@ -2,6 +2,8 @@ from __future__ import annotations +import types + import numpy as np import pytest @@ -48,7 +50,6 @@ def test_enforce_type_regular_size_equal(): result = _type_is_enforceable(layout, target) assert result.is_enforceable - # The full operation must agree with _type_is_enforceable out = ak.enforce_type(ak.Array(layout), "2 * int64") assert str(out.type) == "3 * 2 * int64" @@ -66,8 +67,6 @@ def test_enforce_type_regular_size_mismatch(): def test_merge_union_of_records_categorical_indexed(): - # A union where one branch is a categorical (IndexedArray) of records. - # Build a categorical (IndexedArray) record array directly. index = ak.index.Index64(np.array([1, 0, 1], dtype=np.int64)) inner_records = ak.contents.RecordArray( [ak.contents.NumpyArray(np.array([10.0, 20.0]))], ["c"] @@ -111,16 +110,14 @@ def test_argcartesian_list_vs_dict_parity_axis1(): # Field names differ ("x"/"y" vs "0"/"1") but the index values must match. dict_vals = [[(d["x"], d["y"]) for d in row] for row in as_dict.to_list()] - # The list branch produces tuples, indexed positionally. list_vals = [[tuple(t) for t in row] for row in as_list.to_list()] assert dict_vals == list_vals def test_from_raggedtensor_recursive_deep(): # Drive _recursive_call directly (no TensorFlow needed) for a structure with - # three ragged dimensions (3 offset arrays). Before the fix the recursive - # else-branch reused the same `count`, causing a RecursionError for >= 3 - # ragged dimensions. + # three ragged dimensions. Before the fix the recursive else-branch reused + # the same `count`, causing a RecursionError for >= 3 ragged dimensions. from awkward.operations.ak_from_raggedtensor import _recursive_call content = ak.contents.NumpyArray(np.arange(8, dtype=np.float64)) @@ -133,20 +130,24 @@ def test_from_raggedtensor_recursive_deep(): assert result.to_list() == [[[[0.0, 1.0], [2.0, 3.0]], [[4.0, 5.0], [6.0, 7.0]]]] -class _Dispatcher: - """Dummy array-like recording __awkward_function__ invocations.""" - - def __init__(self): - self.captured = None +def _make_dispatch_sentinel(): + """Array-like that intercepts __awkward_function__ dispatch and records arguments.""" + captured = [] - def __awkward_function__(self, func, array_likes, args, kwargs): - self.captured = list(array_likes) + def __awkward_function__(func, array_likes, args, kwargs): + captured[:] = list(array_likes) return "intercepted" + sentinel = types.SimpleNamespace( + captured=captured, + __awkward_function__=__awkward_function__, + ) + return sentinel + def test_unflatten_dispatch_sees_counts(): array = ak.Array([1, 2, 3, 4]) - sentinel = _Dispatcher() + sentinel = _make_dispatch_sentinel() # counts is array-like and must be offered to dispatch. result = ak.unflatten(array, sentinel) assert result == "intercepted" @@ -154,8 +155,8 @@ def test_unflatten_dispatch_sees_counts(): def test_nan_to_num_dispatch_sees_extra_args(): - sentinel = _Dispatcher() array = ak.Array([1.0, 2.0]) + sentinel = _make_dispatch_sentinel() # Pass the dispatcher as posinf; dispatch must offer it. result = ak.nan_to_num(array, posinf=sentinel) assert result == "intercepted" From 4edf9f70479cc46c406eaab58821c970faf5c787 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Sat, 11 Jul 2026 22:02:35 -0500 Subject: [PATCH 4/4] not needed --- tests/test_2596_named_axis.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_2596_named_axis.py b/tests/test_2596_named_axis.py index db8d33db80..b446eaa9eb 100644 --- a/tests/test_2596_named_axis.py +++ b/tests/test_2596_named_axis.py @@ -689,9 +689,6 @@ def test_named_axis_ak_argcartesian(): named_two = ak.with_named_axis(two, named_axis=("x", "y")) named_three = ak.with_named_axis(three, named_axis=("x", "y")) - # The list input now honours the requested axis (matching the dict input - # branch), so the named-axis output agrees with ak.cartesian/argcartesian - # of an equivalent mapping. assert ak.argcartesian( [named_one, named_two, named_three], axis="x", nested=False ).named_axis == {"x": 0}