Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

C++ refactoring: reducers #1099

Merged
merged 26 commits into from
Sep 30, 2021
Merged

C++ refactoring: reducers #1099

Show file tree
Hide file tree
Changes from 22 commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
d1b2462
start refactoring reducers
ianna Sep 17, 2021
fe58dfd
save for now
ianna Sep 21, 2021
a5b96a1
need more tests
ianna Sep 22, 2021
b87f1f7
all current tests pass
ianna Sep 23, 2021
7958016
reverse 3.8 syntax
ianna Sep 23, 2021
279e915
syntax cleanup
ianna Sep 23, 2021
2675681
windows fix
ianna Sep 23, 2021
d44fdb7
Don't hide unexecuted tests. UnionArray, keepdims, and sum/prod bool …
jpivarski Sep 23, 2021
235f621
Merge branch 'main' into ianna/refactoring-reducers
ianna Sep 23, 2021
078aa26
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 23, 2021
a8510db
Replace reducer name strings with singleton objects.
jpivarski Sep 23, 2021
06576d4
Set unnecessary 'shifts' to None, rather than length 0.
jpivarski Sep 23, 2021
4175873
Clean usage of 'starts' and all programming errors should be Assertio…
jpivarski Sep 23, 2021
773adc7
Included stub implementations of all the untested Content types.
jpivarski Sep 23, 2021
51f0cea
Merge branch 'ianna/refactoring-reducers' of https://github.com/sciki…
jpivarski Sep 23, 2021
5f45807
Still merging in main.
jpivarski Sep 23, 2021
62ee3a3
bytemaskedarray reducer and test
ianna Sep 24, 2021
85e2785
implement preferred_dtype
ianna Sep 24, 2021
31c17fd
fix keepdims
ianna Sep 24, 2021
3f714dd
promote boolean type to int for sum and prod reducers
ianna Sep 24, 2021
3af12de
unmaskedarray and windows fix
ianna Sep 24, 2021
752e699
implement dimension_optiontype
ianna Sep 26, 2021
29a1844
add initial for min and max and more tests
ianna Sep 28, 2021
3482908
indexedarray type bug fix
ianna Sep 29, 2021
07712cb
toregulararray for listarray and listoffsetarray
ianna Sep 29, 2021
c90b33f
ignore B904
ianna Sep 30, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
376 changes: 376 additions & 0 deletions src/awkward/_v2/_reducers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,376 @@
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE

from __future__ import absolute_import

import awkward as ak

np = ak.nplike.NumpyMetadata.instance()


class Reducer(object):
needs_position = False

@classmethod
def return_dtype(cls, given_dtype):
if (
given_dtype == np.bool_
or given_dtype == np.int8
or given_dtype == np.int16
or given_dtype == np.int32
):
if ak._util.win or ak._util.bits32:
return np.int32
else:
return np.int64

if (
given_dtype == np.uint8
or given_dtype == np.uint16
or given_dtype == np.uint32
):
if ak._util.win or ak._util.bits32:
return np.uint32
else:
return np.uint64

return given_dtype


class ArgMin(Reducer):
name = "argmin"

needs_position = True
preferred_dtype = np.int64

@classmethod
def apply(cls, array, parents, outlength):
result = ak._v2.index.Index64.empty(
outlength, array.nplike, dtype=cls.preferred_dtype
)
array._handle_error(
array.nplike[
"awkward_reduce_argmin",
result.dtype.type,
array.data.dtype.type,
parents.dtype.type,
](
result.to(array.nplike),
array.data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
return ak._v2.contents.NumpyArray(result)


class ArgMax(Reducer):
name = "argmax"

needs_position = True
preferred_dtype = np.int64

@classmethod
def apply(cls, array, parents, outlength):
result = ak._v2.index.Index64.empty(
outlength, array.nplike, dtype=cls.preferred_dtype
)
array._handle_error(
array.nplike[
"awkward_reduce_argmax",
result.dtype.type,
array._data.dtype.type,
parents.dtype.type,
](
result.to(array.nplike),
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
return ak._v2.contents.NumpyArray(result)


class Count(Reducer):
name = "count"
preferred_dtype = np.float64

@classmethod
def apply(cls, array, parents, outlength):
result = ak._v2.index.Index64.empty(outlength, array.nplike)
array._handle_error(
array.nplike[
"awkward_reduce_count_64", result.dtype.type, parents.dtype.type
](
result.to(array.nplike),
parents.to(array.nplike),
len(parents),
outlength,
)
)
return ak._v2.contents.NumpyArray(result)


class CountNonzero(Reducer):
name = "count_nonzero"
preferred_dtype = np.float64

@classmethod
def apply(cls, array, parents, outlength):
result = ak._v2.index.Index64.empty(outlength, array.nplike)
array._handle_error(
array.nplike[
"awkward_reduce_countnonzero",
result.dtype.type,
array._data.dtype.type,
parents.dtype.type,
](
result.to(array.nplike),
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
return ak._v2.contents.NumpyArray(result)


class Sum(Reducer):
name = "sum"
preferred_dtype = np.float64

@classmethod
def apply(cls, array, parents, outlength):
result = ak._v2.contents.NumpyArray(
array.nplike.empty(outlength, dtype=cls.return_dtype(array.dtype))
)
if array.dtype == np.bool_:
if result.dtype == np.int64 or result.dtype == np.uint64:
array._handle_error(
array.nplike[
"awkward_reduce_sum_int64_bool_64",
np.int64,
array.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
elif result.dtype == np.int32 or result.dtype == np.uint32:
array._handle_error(
array.nplike[
"awkward_reduce_sum_int32_bool_64",
np.int32,
array.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
else:
raise NotImplementedError
else:
array._handle_error(
array.nplike[
"awkward_reduce_sum",
result.dtype.type,
array.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
return ak._v2.contents.NumpyArray(result)


class Prod(Reducer):
name = "prod"
preferred_dtype = np.int64

@classmethod
def apply(cls, array, parents, outlength):
result = ak._v2.contents.NumpyArray(
array.nplike.empty(outlength, dtype=cls.return_dtype(array.dtype))
)
if array.dtype == np.bool_:
array._handle_error(
array.nplike[
"awkward_reduce_prod_bool",
array.dtype.type,
array.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
else:
array._handle_error(
array.nplike[
"awkward_reduce_prod",
result.dtype.type,
array.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
return ak._v2.contents.NumpyArray(result)


class Any(Reducer):
name = "any"
preferred_dtype = np.bool_

@classmethod
def apply(cls, array, parents, outlength):
result = ak._v2.contents.NumpyArray(
array.nplike.empty(outlength, dtype=cls.preferred_dtype)
)
array._handle_error(
array.nplike[
"awkward_reduce_sum_bool",
result.dtype.type,
array._data.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
return ak._v2.contents.NumpyArray(result)


class All(Reducer):
name = "all"
preferred_dtype = np.bool_

@classmethod
def apply(cls, array, parents, outlength):
result = ak._v2.contents.NumpyArray(
array.nplike.empty(outlength, dtype=cls.preferred_dtype)
)
array._handle_error(
array.nplike[
"awkward_reduce_prod_bool",
result.dtype.type,
array._data.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
return ak._v2.contents.NumpyArray(result)


class Min(Reducer):
name = "min"
preferred_dtype = np.float64

@classmethod
def apply(cls, array, parents, outlength):
dtype = array.dtype
result = ak._v2.contents.NumpyArray(array.nplike.empty(outlength, dtype))

if array.dtype == np.bool_:
array._handle_error(
array.nplike[
"awkward_reduce_prod_bool",
result.dtype.type,
array.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
else:
array._handle_error(
array.nplike[
"awkward_reduce_min",
result.dtype.type,
array.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
np.inf, # FIXME: set initial
)
)
return ak._v2.contents.NumpyArray(result)


class Max(Reducer):
name = "max"
preferred_dtype = np.float64

@classmethod
def apply(cls, array, parents, outlength):
dtype = array.dtype
result = ak._v2.contents.NumpyArray(array.nplike.empty(outlength, dtype))

if array.dtype == np.bool_:
array._handle_error(
array.nplike[
"awkward_reduce_sum_bool",
result.dtype.type,
array.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
)
)
else:
array._handle_error(
array.nplike[
"awkward_reduce_max",
result.dtype.type,
array.dtype.type,
parents.dtype.type,
](
result._data,
array._data,
parents.to(array.nplike),
len(parents),
outlength,
-np.inf, # FIXME: set initial
)
)
return ak._v2.contents.NumpyArray(result)
Loading