From e1fe9b2e3bc1c0b0832ac36f60d7f06b1eca8051 Mon Sep 17 00:00:00 2001 From: Doug Davis Date: Thu, 2 Nov 2023 07:24:27 -0500 Subject: [PATCH 1/3] use a duck typing check --- src/dask_awkward/lib/core.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index 6f8924dc..16033486 100644 --- a/src/dask_awkward/lib/core.py +++ b/src/dask_awkward/lib/core.py @@ -534,10 +534,21 @@ def new_record_object(dsk: HighLevelGraph, name: str, *, meta: Any) -> Record: return Record(dsk, name, meta) +def _is_numpy_or_cupy_like(arr: Any) -> bool: + return ( + hasattr(arr, "ndim") + and hasattr(arr, "shape") + and isinstance(arr.shape, tuple) + and hasattr(arr, "dtype") + ) + + def _finalize_array(results: Sequence[Any]) -> Any: # special cases for length 1 results if len(results) == 1: - if isinstance(results[0], (int, ak.Array, np.ndarray)): + if isinstance(results[0], (int, ak.Array)) or _is_numpy_or_cupy_like( + results[0] + ): return results[0] # a sequence of arrays that need to be concatenated. From c6f6e0b33236533958e9c17c1351fdfd3c56c3b8 Mon Sep 17 00:00:00 2001 From: Doug Davis Date: Thu, 2 Nov 2023 07:46:12 -0500 Subject: [PATCH 2/3] more duck typing, remove redundant branch --- src/dask_awkward/lib/core.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index 16033486..102961a2 100644 --- a/src/dask_awkward/lib/core.py +++ b/src/dask_awkward/lib/core.py @@ -557,18 +557,11 @@ def _finalize_array(results: Sequence[Any]) -> Any: # a sequence of scalars that are stored as np.ndarray(N) where N # is a number (i.e. shapeless numpy array) - elif any(isinstance(r, np.ndarray) for r in results) and any( + elif any(_is_numpy_or_cupy_like(r) for r in results) and any( r.shape == () for r in results ): return ak.Array(list(results)) - # sometimes we just check the length of partitions so all results - # will be integers, just make an array out of that. - elif isinstance(results, (tuple, list)) and all( - isinstance(r, (int, np.integer)) for r in results - ): - return ak.Array(list(results)) - # sometimes all partition results will be None (some write-to-disk # operations) elif all(r is None for r in results): From c353ac667996a80b4d83620ad80ad2cc27187f6d Mon Sep 17 00:00:00 2001 From: Doug Davis Date: Thu, 2 Nov 2023 10:52:16 -0500 Subject: [PATCH 3/3] bring back old branch necessary for awkward < 2.5 --- src/dask_awkward/lib/core.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index 102961a2..a7a6458f 100644 --- a/src/dask_awkward/lib/core.py +++ b/src/dask_awkward/lib/core.py @@ -546,9 +546,8 @@ def _is_numpy_or_cupy_like(arr: Any) -> bool: def _finalize_array(results: Sequence[Any]) -> Any: # special cases for length 1 results if len(results) == 1: - if isinstance(results[0], (int, ak.Array)) or _is_numpy_or_cupy_like( - results[0] - ): + np_like = _is_numpy_or_cupy_like(results[0]) + if isinstance(results[0], (int, ak.Array)) or np_like: # type: ignore[unreachable] return results[0] # a sequence of arrays that need to be concatenated. @@ -562,6 +561,12 @@ def _finalize_array(results: Sequence[Any]) -> Any: ): return ak.Array(list(results)) + # in awkward < 2.5 we can get integers instead of np.array scalars + elif isinstance(results, (tuple, list)) and all( + isinstance(r, (int, np.integer)) for r in results + ): + return ak.Array(list(results)) + # sometimes all partition results will be None (some write-to-disk # operations) elif all(r is None for r in results):