From 6581ba0a81d31fbb4a27215aae930408942bc008 Mon Sep 17 00:00:00 2001 From: Ryan Keck Date: Fri, 23 Jan 2026 10:42:07 -0500 Subject: [PATCH 1/3] Closes 5330: Improve ak.array Bigint array transfer performance --- arkouda/numpy/pdarraycreation.py | 210 ++++++++++++++++++++++++++----- benchmarks/array_transfer.py | 2 +- 2 files changed, 183 insertions(+), 29 deletions(-) diff --git a/arkouda/numpy/pdarraycreation.py b/arkouda/numpy/pdarraycreation.py index 7c4765d1600..242396dc59a 100644 --- a/arkouda/numpy/pdarraycreation.py +++ b/arkouda/numpy/pdarraycreation.py @@ -92,16 +92,22 @@ def array( *, copy: bool = ..., max_bits: int = ..., + unsafe: bool = ..., + any_neg: Optional[bool] = ..., + num_bits: Optional[int] = ..., ) -> Strings: ... @overload def array( a: Union[pdarray, Strings, Iterable[Any], NDArray[Any]], - dtype: _NumericLikeDType = ..., # object covers np.dtype, type, int/float/bool dtypes, etc. + dtype: _NumericLikeDType = ..., *, copy: bool = ..., max_bits: int = ..., + unsafe: bool = ..., + any_neg: Optional[bool] = ..., + num_bits: Optional[int] = ..., ) -> pdarray: ... @@ -112,15 +118,26 @@ def array( *, copy: bool = ..., max_bits: int = ..., + unsafe: bool = ..., + any_neg: Optional[bool] = ..., + num_bits: Optional[int] = ..., ) -> Union[pdarray, Strings]: ... +# The remaining overloads can stay as-is, but it's nicer to add the kwargs there too +# so tooling doesn't complain in those call patterns. Minimal version: + + @overload def array( a: Union[pdarray, List[_NumericLikeDType]], dtype: Union[_StringDType, _NumericLikeDType, None] = None, copy: bool = False, max_bits: int = -1, + *, + unsafe: bool = ..., + any_neg: Optional[bool] = ..., + num_bits: Optional[int] = ..., ) -> pdarray: ... @@ -130,6 +147,10 @@ def array( dtype: Union[_StringDType, _NumericLikeDType, None] = None, copy: bool = False, max_bits: int = -1, + *, + unsafe: bool = ..., + any_neg: Optional[bool] = ..., + num_bits: Optional[int] = ..., ) -> Strings: ... @@ -139,6 +160,10 @@ def array( dtype: Union[_StringDType, _NumericLikeDType, None] = None, copy: bool = False, max_bits: int = -1, + *, + unsafe: bool = ..., + any_neg: Optional[bool] = ..., + num_bits: Optional[int] = ..., ) -> Union[pdarray, Strings]: ... @@ -147,6 +172,18 @@ def array( dtype: Union[np.dtype, type, str, None] = None, copy: bool = False, max_bits: int = -1, + *, + # Bigint fast-path controls: + # + # - unsafe=False (default): validate object arrays and infer bigint sizing safely. + # - unsafe=True: allow skipping validation and/or sizing inference when hints are provided. + # + # If you pass unsafe=True with dtype=bigint and provide both any_neg and num_bits, + # Arkouda will TRUST those values and skip the full-array Python scans that would + # otherwise be required for dtype=object inputs. + unsafe: bool = False, + any_neg: Optional[bool] = None, + num_bits: Optional[int] = None, ) -> Union[pdarray, Strings]: """ Convert a Python, NumPy, or Arkouda array-like into a `pdarray` or `Strings` object, @@ -170,6 +207,21 @@ def array( max_bits : int, optional The maximum number of bits for bigint arrays. Ignored for other dtypes. + unsafe : bool, default=False + Enable performance-oriented shortcuts for bigint creation from dtype=object arrays. + When True, the caller may supply trusted hints (any_neg, num_bits) to skip Python-level + full-array scans. If hints are not supplied, Arkouda will still do a single-pass scan + (rather than multiple scans) for dtype=object. + + WARNING: If unsafe=True and provided hints are wrong, results may be incorrect. + + any_neg : Optional[bool], default=None + Bigint hint: whether any value is negative. Only used when dtype=bigint and a is numeric/object. + + num_bits : Optional[int], default=None + Bigint hint: required bit-width for values (including sign bit if signed). + If provided with unsafe=True, Arkouda can compute number of limbs without scanning. + Returns ------- Union[pdarray, Strings] @@ -250,7 +302,6 @@ def array( a = list(a) # If a is not already a numpy.ndarray, convert it - if not isinstance(a, np.ndarray): try: if dtype is not None and dtype != bigint: @@ -270,8 +321,13 @@ def array( a = np.array(a) except (RuntimeError, TypeError, ValueError): raise TypeError("a must be a pdarray, np.ndarray, or convertible to a numpy array") + if dtype is not None and dtype not in [bigint, "bigint"]: a = a.astype(dtype) + + # If numpy gave us an object array of integers and dtype was not specified, default to bigint. + # NOTE: This scan can be expensive for very large arrays; callers who already know they want + # bigint should pass dtype=bigint to avoid this inference. if a.dtype == object and dtype is None and all(isinstance(x, (int, np.integer)) for x in a.flat): dtype = bigint @@ -280,11 +336,18 @@ def array( # Return multi-dimensional pdarray if a.ndim in get_array_ranks() # otherwise raise an error - if a.ndim not in get_array_ranks(): raise ValueError(f"array rank {a.ndim} not in compiled ranks {get_array_ranks()}") + + # Bigint conversion path (including dtype=object numeric inputs) if a.dtype.kind in ("i", "u", "f", "O") and (dtype == bigint or dtype == "bigint"): - return _bigint_from_numpy(a, max_bits) + return _bigint_from_numpy( + a, + max_bits, + unsafe=unsafe, + any_neg=any_neg, + num_bits=num_bits, + ) elif not ( np.issubdtype(a.dtype, np.str_) or (a.dtype == np.object_ and a.size > 0 and isinstance(a[0], str)) @@ -315,7 +378,6 @@ def array( return strings if dtype is None else type_cast(Union[pdarray, Strings], akcast(strings, dtype)) # If not strings, then check that dtype is supported in arkouda - from arkouda.numpy.util import _infer_shape_from_size shape, ndim, full_size = _infer_shape_from_size(a.shape) @@ -359,7 +421,33 @@ def array( ) -def _bigint_from_numpy(np_a: np.ndarray, max_bits: int) -> pdarray: +def _bigint_from_numpy( + np_a: np.ndarray, + max_bits: int, + *, + unsafe: bool = False, + any_neg: Optional[bool] = None, + num_bits: Optional[int] = None, +) -> pdarray: + """ + Create a bigint pdarray from a NumPy ndarray. + + Fast paths: + - Empty arrays return an empty bigint with shape preserved. + - Non-object numeric arrays (int/uint/float) use a single-limb transfer. + - dtype=object arrays of Python ints are split into uint64 limbs client-side. + + Performance notes: + - dtype=object requires Python-level work. To reduce overhead: + * this implementation uses a SINGLE PASS over the data to validate and infer sizing + (instead of multiple passes), unless the caller provides trusted hints. + * if unsafe=True and both any_neg and num_bits are provided, we skip all scans + and proceed directly to limb extraction. + + WARNING: + - If unsafe=True and provided hints are wrong (e.g., num_bits too small or any_neg incorrect), + results may be incorrect. + """ from arkouda.client import generic_msg from arkouda.numpy.pdarrayclass import create_pdarray, pdarray @@ -371,18 +459,13 @@ def _bigint_from_numpy(np_a: np.ndarray, max_bits: int) -> pdarray: if a.size == 0: return zeros(size=a.shape, dtype=bigint, max_bits=max_bits) - # Only ints/object should reach here (strings handled upstream) + # Only ints/uints/floats/object should reach here (strings handled upstream) if a.dtype.kind not in ("i", "u", "f", "O"): raise TypeError(f"bigint requires numeric input, got dtype={a.dtype}") - if a.dtype.kind == "O": - # Only allow Python ints and floats - if not all(isinstance(x, (int, float)) for x in a.flat): - raise TypeError("bigint requires numeric input, got non-numeric object") - if any(isinstance(x, float) for x in a.flat): - a = a.astype(np.float64, copy=False) - - # Fast all-zero path or single limb path + # ------------------------------------------------------------ + # Fast path for non-object numeric dtypes: send a single limb. + # ------------------------------------------------------------ if a.dtype.kind in ("i", "u", "f"): if not np.any(a): return zeros(size=a.shape, dtype=bigint, max_bits=max_bits) @@ -394,7 +477,6 @@ def _bigint_from_numpy(np_a: np.ndarray, max_bits: int) -> pdarray: else: ak_a = array(a.astype(np.float64, copy=False)) - # Send a single limb array return create_pdarray( generic_msg( cmd=f"big_int_creation_one_limb<{ak_a.dtype},{ak_a.ndim}>", @@ -405,31 +487,103 @@ def _bigint_from_numpy(np_a: np.ndarray, max_bits: int) -> pdarray: }, ) ) - else: - if all(int(x) == 0 for x in a.flat): + + # ---------------------------------------------------------------- + # dtype=object path (typically Python ints, sometimes floats/other): + # ---------------------------------------------------------------- + + # If caller supplied trusted sizing hints, skip all validation/scans. + # This is intended for benchmarking or trusted pipelines where the caller + # already knows whether the data are signed and how many bits are needed. + if unsafe and any_neg is not None and num_bits is not None: + if num_bits <= 0: return zeros(size=a.shape, dtype=bigint, max_bits=max_bits) - any_neg = np.any(flat < 0) - req_bits: int - if any_neg: - req_bits = max(flat.max().bit_length(), (-flat.min()).bit_length()) + 1 + + req_bits = int(num_bits) + req_limbs = (req_bits + 63) >> 6 # == ceil(req_bits / 64) else: - req_bits = flat.max().bit_length() - req_limbs = (req_bits + 63) // 64 - mask = (1 << 64) - 1 + # Single-pass scan: + # - validate types (unless unsafe) + # - detect floats (fallback to float64 one-limb path) + # - detect all-zero (fast path) + # - infer any_neg and required bits + has_float = False + any_nonzero = False + inferred_any_neg = False + max_mag_bits = 0 + + # Iterate once over Python objects. + for x in a.flat: + if unsafe: + # Caller asserts objects are numeric (typically Python ints). + pass + else: + if not isinstance(x, (int, float, np.integer, np.floating)): + raise TypeError("bigint requires numeric input, got non-numeric object") + + # Handle floats: bigint-from-float behaves like existing code (cast to float64). + # Note: we only need to know *that* a float exists, not which values. + if isinstance(x, (float, np.floating)): + has_float = True + break + + # Integers (Python int or numpy integer) + xi = int(x) + if xi != 0: + any_nonzero = True + if xi < 0: + inferred_any_neg = True + mag_bits = (-xi).bit_length() + else: + mag_bits = xi.bit_length() + max_mag_bits = max(max_mag_bits, mag_bits) + + if has_float: + # Match previous behavior: object array containing floats becomes float64 + # and uses the single-limb float path above. + a = a.astype(np.float64, copy=False) + flat = a.ravel() + if not np.any(a): + return zeros(size=a.shape, dtype=bigint, max_bits=max_bits) + ak_a = array(a.astype(np.float64, copy=False)) + return create_pdarray( + generic_msg( + cmd=f"big_int_creation_one_limb<{ak_a.dtype},{ak_a.ndim}>", + args={ + "array": ak_a, + "shape": ak_a.shape, + "max_bits": max_bits, + }, + ) + ) + + if not any_nonzero: + return zeros(size=a.shape, dtype=bigint, max_bits=max_bits) + + any_neg = inferred_any_neg + # For signed representation reserve a sign bit (+1). + req_bits = max_mag_bits + (1 if any_neg else 0) + req_limbs = (req_bits + 63) >> 6 # == ceil(req_bits / 64) + + mask = 0xFFFFFFFFFFFFFFFF # 2**64 - 1 uint_arrays: List[Union[pdarray, Strings]] = [] - # attempt to break bigint into multiple uint64 arrays + + # Attempt to break bigint into multiple uint64 limb arrays. + # NOTE: This loop is inherently heavy for dtype=object inputs (Python big ints), + # but we minimize overhead by avoiding extra pre-scans of the data. for _ in range(req_limbs): low = flat & mask flat = flat >> 64 # type: ignore - # low, flat = flat & mask, flat >> 64 uint_arrays.append(array(np.array(low, dtype=np.uint), dtype=akuint64)) + + # Server expects shape of the resulting bigint array (original shape). return create_pdarray( generic_msg( cmd=f"big_int_creation_multi_limb<{uint_arrays[0].dtype},1>", args={ "arrays": uint_arrays, "num_arrays": len(uint_arrays), - "signed": any_neg, + "signed": bool(any_neg), "shape": flat.shape, "max_bits": max_bits, }, diff --git a/benchmarks/array_transfer.py b/benchmarks/array_transfer.py index c20739205c3..032352b032f 100755 --- a/benchmarks/array_transfer.py +++ b/benchmarks/array_transfer.py @@ -38,7 +38,7 @@ def time_ak_array_transfer(N, trials, dtype, seed, max_bits=-1): end = time.time() to_ndarray_times.append(end - start) start = time.time() - aka = ak.array(npa, max_bits=max_bits, dtype=dtype) + aka = ak.array(npa, max_bits=max_bits, dtype=dtype, unsafe=True, num_bits=128, any_neg=False) end = time.time() to_pdarray_times.append(end - start) gc.collect() From 7ccef429958c3b10b2d81326d20ac2030f59f69d Mon Sep 17 00:00:00 2001 From: Ryan Keck Date: Fri, 23 Jan 2026 13:07:07 -0500 Subject: [PATCH 2/3] Addresses comment --- arkouda/numpy/pdarraycreation.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/arkouda/numpy/pdarraycreation.py b/arkouda/numpy/pdarraycreation.py index 242396dc59a..1292da14f22 100644 --- a/arkouda/numpy/pdarraycreation.py +++ b/arkouda/numpy/pdarraycreation.py @@ -173,14 +173,6 @@ def array( copy: bool = False, max_bits: int = -1, *, - # Bigint fast-path controls: - # - # - unsafe=False (default): validate object arrays and infer bigint sizing safely. - # - unsafe=True: allow skipping validation and/or sizing inference when hints are provided. - # - # If you pass unsafe=True with dtype=bigint and provide both any_neg and num_bits, - # Arkouda will TRUST those values and skip the full-array Python scans that would - # otherwise be required for dtype=object inputs. unsafe: bool = False, any_neg: Optional[bool] = None, num_bits: Optional[int] = None, @@ -189,6 +181,22 @@ def array( Convert a Python, NumPy, or Arkouda array-like into a `pdarray` or `Strings` object, transferring data to the Arkouda server. + Bigint fast path + ---------------- + When `dtype` is `bigint` (or ``"bigint"``) and the input is a NumPy ``dtype=object`` array, + building the bigint representation may require Python-level passes over the data to infer + sign and required bit width. + + Set ``unsafe=True`` to enable optional performance shortcuts for trusted inputs: + + * If both ``any_neg`` and ``num_bits`` are provided, Arkouda will trust these hints and + skip inference passes, proceeding directly to limb extraction. + * If hints are not provided, Arkouda may still perform a single-pass inference step + for ``dtype=object`` inputs (implementation-dependent). + + .. warning:: + If ``unsafe=True`` and the provided hints are incorrect, results may be incorrect. + Parameters ---------- a : Union[pdarray, np.ndarray, Iterable, Strings] From 80686d674ca1c0a7bdb4f579492ccd2692579166 Mon Sep 17 00:00:00 2001 From: Ryan Keck Date: Fri, 23 Jan 2026 14:58:06 -0500 Subject: [PATCH 3/3] Addresses many comments --- arkouda/numpy/pdarraycreation.py | 3 +-- benchmark_v2/array_transfer_benchmark.py | 7 ++++++- benchmarks/array_transfer.py | 11 ++++++++--- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/arkouda/numpy/pdarraycreation.py b/arkouda/numpy/pdarraycreation.py index 1292da14f22..fdbee650b68 100644 --- a/arkouda/numpy/pdarraycreation.py +++ b/arkouda/numpy/pdarraycreation.py @@ -550,10 +550,9 @@ def _bigint_from_numpy( # Match previous behavior: object array containing floats becomes float64 # and uses the single-limb float path above. a = a.astype(np.float64, copy=False) - flat = a.ravel() if not np.any(a): return zeros(size=a.shape, dtype=bigint, max_bits=max_bits) - ak_a = array(a.astype(np.float64, copy=False)) + ak_a = array(a) return create_pdarray( generic_msg( cmd=f"big_int_creation_one_limb<{ak_a.dtype},{ak_a.ndim}>", diff --git a/benchmark_v2/array_transfer_benchmark.py b/benchmark_v2/array_transfer_benchmark.py index a809f0fe352..12cdc07620e 100644 --- a/benchmark_v2/array_transfer_benchmark.py +++ b/benchmark_v2/array_transfer_benchmark.py @@ -52,8 +52,13 @@ def bench_array_transfer_from_ndarray(benchmark, dtype): def from_np(): ak.array(npa, max_bits=pytest.max_bits) + def from_np_bigint(): + ak.array(npa, max_bits=-1, dtype=dtype, unsafe=True, num_bits=128, any_neg=False) - benchmark.pedantic(from_np, rounds=pytest.trials) + if dtype == "bigint": + benchmark.pedantic(from_np_bigint, rounds=pytest.trials) + else: + benchmark.pedantic(from_np, rounds=pytest.trials) benchmark.extra_info["description"] = "Measures the performance of ak.array" benchmark.extra_info["problem_size"] = N diff --git a/benchmarks/array_transfer.py b/benchmarks/array_transfer.py index 032352b032f..31a1b50f38c 100755 --- a/benchmarks/array_transfer.py +++ b/benchmarks/array_transfer.py @@ -37,9 +37,14 @@ def time_ak_array_transfer(N, trials, dtype, seed, max_bits=-1): npa = a.to_ndarray() end = time.time() to_ndarray_times.append(end - start) - start = time.time() - aka = ak.array(npa, max_bits=max_bits, dtype=dtype, unsafe=True, num_bits=128, any_neg=False) - end = time.time() + if dtype == ak.bigint.name: + start = time.time() + aka = ak.array(npa, max_bits=max_bits, dtype=dtype, unsafe=True, num_bits=128, any_neg=False) + end = time.time() + else: + start = time.time() + aka = ak.array(npa, max_bits=max_bits, dtype=dtype) + end = time.time() to_pdarray_times.append(end - start) gc.collect() avgnd = sum(to_ndarray_times) / trials