Skip to content

REF: handle 2D in tslibs.vectorized #46886

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

Merged
merged 4 commits into from
May 6, 2022
Merged

REF: handle 2D in tslibs.vectorized #46886

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
552d420
REF: handle 2D in tslibs.vectorized
jbrockmendel Apr 27, 2022
bab9d7e
Merge branch 'main' into ref-tzconversion-7
jbrockmendel May 2, 2022
52c5e2e
Merge branch 'main' into ref-tzconversion-7
jbrockmendel May 4, 2022
403f3b0
Use flatiter
jbrockmendel May 4, 2022
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
1 change: 1 addition & 0 deletions pandas/_libs/tslibs/timedeltas.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,7 @@ def array_to_timedelta64(
raise ValueError(
"unit must not be specified if the input contains a str"
)
cnp.PyArray_ITER_NEXT(it)

# Usually, we have all strings. If so, we hit the fast path.
# If this path fails, we try conversion a different way, and
Expand Down
12 changes: 6 additions & 6 deletions pandas/_libs/tslibs/vectorized.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,24 +11,24 @@ from pandas._libs.tslibs.offsets import BaseOffset
from pandas._typing import npt

def dt64arr_to_periodarr(
stamps: npt.NDArray[np.int64], # const int64_t[:]
stamps: npt.NDArray[np.int64],
freq: int,
tz: tzinfo | None,
) -> npt.NDArray[np.int64]: ... # np.ndarray[np.int64, ndim=1]
) -> npt.NDArray[np.int64]: ...
def is_date_array_normalized(
stamps: npt.NDArray[np.int64], # const int64_t[:]
stamps: npt.NDArray[np.int64],
tz: tzinfo | None = ...,
) -> bool: ...
def normalize_i8_timestamps(
stamps: npt.NDArray[np.int64], # const int64_t[:]
stamps: npt.NDArray[np.int64],
tz: tzinfo | None,
) -> npt.NDArray[np.int64]: ...
def get_resolution(
stamps: npt.NDArray[np.int64], # const int64_t[:]
stamps: npt.NDArray[np.int64],
tz: tzinfo | None = ...,
) -> Resolution: ...
def ints_to_pydatetime(
arr: npt.NDArray[np.int64], # const int64_t[:}]
arr: npt.NDArray[np.int64],
tz: tzinfo | None = ...,
freq: BaseOffset | None = ...,
fold: bool = ...,
Expand Down
160 changes: 104 additions & 56 deletions pandas/_libs/tslibs/vectorized.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,8 @@ from .tzconversion cimport Localizer

@cython.boundscheck(False)
@cython.wraparound(False)
def tz_convert_from_utc(const int64_t[:] stamps, tzinfo tz):
def tz_convert_from_utc(ndarray stamps, tzinfo tz):
# stamps is int64_t, arbitrary ndim
"""
Convert the values (in i8) from UTC to tz

Expand All @@ -54,27 +55,33 @@ def tz_convert_from_utc(const int64_t[:] stamps, tzinfo tz):
cdef:
Localizer info = Localizer(tz)
int64_t utc_val, local_val
Py_ssize_t pos, i, n = stamps.shape[0]
Py_ssize_t pos, i, n = stamps.size

int64_t[::1] result
ndarray result
cnp.broadcast mi

if tz is None or is_utc(tz) or stamps.size == 0:
# Much faster than going through the "standard" pattern below
return stamps.base.copy()
return stamps.copy()

result = np.empty(n, dtype=np.int64)
result = cnp.PyArray_EMPTY(stamps.ndim, stamps.shape, cnp.NPY_INT64, 0)
mi = cnp.PyArray_MultiIterNew2(result, stamps)

for i in range(n):
utc_val = stamps[i]
# Analogous to: utc_val = stamps[i]
utc_val = (<int64_t*>cnp.PyArray_MultiIter_DATA(mi, 1))[0]

if utc_val == NPY_NAT:
result[i] = NPY_NAT
continue
local_val = NPY_NAT
else:
local_val = info.utc_val_to_local_val(utc_val, &pos)

local_val = info.utc_val_to_local_val(utc_val, &pos)
# Analogous to: result[i] = local_val
(<int64_t*>cnp.PyArray_MultiIter_DATA(mi, 0))[0] = local_val

result[i] = local_val
cnp.PyArray_MultiIter_NEXT(mi)

return result.base
return result


# -------------------------------------------------------------------------
Expand All @@ -83,12 +90,13 @@ def tz_convert_from_utc(const int64_t[:] stamps, tzinfo tz):
@cython.wraparound(False)
@cython.boundscheck(False)
def ints_to_pydatetime(
const int64_t[:] stamps,
ndarray stamps,
tzinfo tz=None,
BaseOffset freq=None,
bint fold=False,
str box="datetime"
) -> np.ndarray:
# stamps is int64, arbitrary ndim
"""
Convert an i8 repr to an ndarray of datetimes, date, time or Timestamp.

Expand Down Expand Up @@ -119,13 +127,21 @@ def ints_to_pydatetime(
cdef:
Localizer info = Localizer(tz)
int64_t utc_val, local_val
Py_ssize_t i, n = stamps.shape[0]
Py_ssize_t i, n = stamps.size
Py_ssize_t pos = -1 # unused, avoid not-initialized warning

npy_datetimestruct dts
tzinfo new_tz
ndarray[object] result = np.empty(n, dtype=object)
bint use_date = False, use_time = False, use_ts = False, use_pydt = False
object res_val

# Note that `result` (and thus `result_flat`) is C-order and
# `it` iterates C-order as well, so the iteration matches
# See discussion at
# github.com/pandas-dev/pandas/pull/46886#discussion_r860261305
ndarray result = cnp.PyArray_EMPTY(stamps.ndim, stamps.shape, cnp.NPY_OBJECT, 0)
object[::1] res_flat = result.ravel() # should NOT be a copy
cnp.flatiter it = cnp.PyArray_IterNew(stamps)

if box == "date":
assert (tz is None), "tz should be None when converting to date"
Expand All @@ -142,31 +158,44 @@ def ints_to_pydatetime(
)

for i in range(n):
utc_val = stamps[i]
# Analogous to: utc_val = stamps[i]
utc_val = (<int64_t*>cnp.PyArray_ITER_DATA(it))[0]

new_tz = tz

if utc_val == NPY_NAT:
result[i] = <object>NaT
continue
res_val = <object>NaT

local_val = info.utc_val_to_local_val(utc_val, &pos)
if info.use_pytz:
# find right representation of dst etc in pytz timezone
new_tz = tz._tzinfos[tz._transition_info[pos]]

dt64_to_dtstruct(local_val, &dts)

if use_ts:
result[i] = create_timestamp_from_ts(utc_val, dts, new_tz, freq, fold)
elif use_pydt:
result[i] = datetime(
dts.year, dts.month, dts.day, dts.hour, dts.min, dts.sec, dts.us,
new_tz, fold=fold,
)
elif use_date:
result[i] = date(dts.year, dts.month, dts.day)
else:
result[i] = time(dts.hour, dts.min, dts.sec, dts.us, new_tz, fold=fold)

local_val = info.utc_val_to_local_val(utc_val, &pos)
if info.use_pytz:
# find right representation of dst etc in pytz timezone
new_tz = tz._tzinfos[tz._transition_info[pos]]

dt64_to_dtstruct(local_val, &dts)

if use_ts:
res_val = create_timestamp_from_ts(utc_val, dts, new_tz, freq, fold)
elif use_pydt:
res_val = datetime(
dts.year, dts.month, dts.day, dts.hour, dts.min, dts.sec, dts.us,
new_tz, fold=fold,
)
elif use_date:
res_val = date(dts.year, dts.month, dts.day)
else:
res_val = time(dts.hour, dts.min, dts.sec, dts.us, new_tz, fold=fold)

# Note: we can index result directly instead of using PyArray_MultiIter_DATA
# like we do for the other functions because result is known C-contiguous
# and is the first argument to PyArray_MultiIterNew2. The usual pattern
# does not seem to work with object dtype.
# See discussion at
# github.com/pandas-dev/pandas/pull/46886#discussion_r860261305
res_flat[i] = res_val
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@seberg is there a more idiomatic way to do this __setitem__?

Copy link
Contributor

@seberg seberg Apr 28, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While I want to replace it (it is not quite ideal, NumPy uses a different functionality internally now). This looks like it should be fine to use PyArray_SETITEM(arr, pointer, res_val) as the "idiomatic" way.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. To clarify though, I was hoping for something that didn't rely on setting into res_flat, as that requires knowledge about the inner workings of the cnp.broadcast object that i only kinda-get, and that only because you and bashtage explained it in the comments to one of his PRs.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are not trying to avoid Cython memory-views due to the buffer export overhead for smaller arrays, right?

In that case, I would suggest using (untested):

object res_flat[::1] = PyArray_BYTES(arr)  # if that works, or the ravel()

That way, cython can deal with the correct assignment, and you know that object is right here. (You can avoid the memoryview, but I think you may have to Py_INCREF manually then.)

There could be faster/better ways to iterate the other array as well, but I suppose since this works with Python objects, it won't matter much. However, I think you can still replace the MultiIter with a single array iterator in any case (just make sure that one iterates in C order). There is no reason to iteraterate res, since that is effectively unused.

It would be really nice to have a "ufunc-like" pattern for this type of thing, but I would need more dedicated time to figure that out.
(The point being, that you would define a single "strided loop" for 1-D arrays, and than add machinery/code generation to enable to for the N-D case)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are not trying to avoid Cython memory-views due to the buffer export overhead for smaller arrays, right?

Never crossed my mind.

The question may be clarified by describing the failure mode I am looking to avoid:

In 6 months another contributor goes to make the analogous change to another function, uses this pattern as a template, but instead of mi = cnp.PyArray_MultiIterNew2(result, stamps) reverses the arguments as mi = cnp.PyArray_MultiIterNew2(stamps, result), so setting res_flat[i] = res_val is no longer necessarily correct.

i.e. using res_flat at all is implicitly relying on knowledge of MultiIter internals that AFAIK is not documented. It just seems like a code smell. Is that clearer?

but I think you may have to Py_INCREF manually then

Definitely don't want to be in that business.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are only relying on the C-order iteration I think. Which I believe MultiIter guarantees (but I am not sure). In any case, you do not need MultiIter, you only need PyArray_IterNew with a comment that it iterates in C-order and is thus compatible with the res_flat.
Then you can use object res_flat[::1] = res.ravel() for assignment, relying on the fact that the typed memoryview is typed correctly.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK, I think I was a bit too fuzzy about what I meant, so here is the working diff in the details.
So long the result is allocated as C-order, the same pattern could also be used elsewhere for small speedups ("iterating" one instead of two arrays), but I think here it may be particularly worthwhile, because I am not certain that cython can optimize the item assignment otherwise.

It would be cool if cython had a way to use a more "ufunc-like" pattern for these iterations, but I guess that that would be a pretty big new thing...

diff --git a/pandas/_libs/tslibs/vectorized.pyx b/pandas/_libs/tslibs/vectorized.pyx
index 3f0f06d874..6764127552 100644
--- a/pandas/_libs/tslibs/vectorized.pyx
+++ b/pandas/_libs/tslibs/vectorized.pyx
@@ -135,9 +135,11 @@ def ints_to_pydatetime(
         bint use_date = False, use_time = False, use_ts = False, use_pydt = False
         object res_val
 
+        # Note that `result` (and thus `result_flat`) is C-order and
+        # `stamps_iter` iterates  C-order as well, so the iteration matches:
         ndarray result = cnp.PyArray_EMPTY(stamps.ndim, stamps.shape, cnp.NPY_OBJECT, 0)
-        ndarray res_flat = result.ravel()  # should NOT be a copy
-        cnp.broadcast mi = cnp.PyArray_MultiIterNew2(result, stamps)
+        object[::1] res_flat = result.ravel()  # should NOT be a copy
+        cnp.flatiter stamps_iter = cnp.PyArray_IterNew(stamps)
 
     if box == "date":
         assert (tz is None), "tz should be None when converting to date"
@@ -155,7 +157,7 @@ def ints_to_pydatetime(
 
     for i in range(n):
         # Analogous to: utc_val = stamps[i]
-        utc_val = (<int64_t*>cnp.PyArray_MultiIter_DATA(mi, 1))[0]
+        utc_val = (<int64_t*>cnp.PyArray_ITER_DATA(stamps_iter))[0]
 
         new_tz = tz
 
@@ -183,15 +185,11 @@ def ints_to_pydatetime(
             else:
                 res_val = time(dts.hour, dts.min, dts.sec, dts.us, new_tz, fold=fold)
 
-        # Note: we can index result directly instead of using PyArray_MultiIter_DATA
-        #  like we do for the other functions because result is known C-contiguous
-        #  and is the first argument to PyArray_MultiIterNew2.  The usual pattern
-        #  does not seem to work with object dtype.
-        #  See discussion at
-        #  github.com/pandas-dev/pandas/pull/46886#discussion_r860261305
+        # Note: We use a typed memory-view here to use cython for handling
+        #       need for object ownership handling (reference counting).
         res_flat[i] = res_val
 
-        cnp.PyArray_MultiIter_NEXT(mi)
+        cnp.PyArray_ITER_NEXT(stamps_iter)
 
     return result

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you only need PyArray_IterNew with a comment that it iterates in C-order

huh i guess i assumed it iterated in a cheapest-available order, so would be F-order if the array is F-contiguous. Good to know, thanks!


cnp.PyArray_ITER_NEXT(it)

return result

Expand All @@ -190,27 +219,33 @@ cdef inline c_Resolution _reso_stamp(npy_datetimestruct *dts):

@cython.wraparound(False)
@cython.boundscheck(False)
def get_resolution(const int64_t[:] stamps, tzinfo tz=None) -> Resolution:
def get_resolution(ndarray stamps, tzinfo tz=None) -> Resolution:
# stamps is int64_t, any ndim
cdef:
Localizer info = Localizer(tz)
int64_t utc_val, local_val
Py_ssize_t i, n = stamps.shape[0]
Py_ssize_t i, n = stamps.size
Py_ssize_t pos = -1 # unused, avoid not-initialized warning
cnp.flatiter it = cnp.PyArray_IterNew(stamps)

npy_datetimestruct dts
c_Resolution reso = c_Resolution.RESO_DAY, curr_reso

for i in range(n):
utc_val = stamps[i]
# Analogous to: utc_val = stamps[i]
utc_val = cnp.PyArray_GETITEM(stamps, cnp.PyArray_ITER_DATA(it))

if utc_val == NPY_NAT:
continue
pass
else:
local_val = info.utc_val_to_local_val(utc_val, &pos)

local_val = info.utc_val_to_local_val(utc_val, &pos)
dt64_to_dtstruct(local_val, &dts)
curr_reso = _reso_stamp(&dts)
if curr_reso < reso:
reso = curr_reso

dt64_to_dtstruct(local_val, &dts)
curr_reso = _reso_stamp(&dts)
if curr_reso < reso:
reso = curr_reso
cnp.PyArray_ITER_NEXT(it)

return Resolution(reso)

Expand All @@ -221,7 +256,8 @@ def get_resolution(const int64_t[:] stamps, tzinfo tz=None) -> Resolution:
@cython.cdivision(False)
@cython.wraparound(False)
@cython.boundscheck(False)
cpdef ndarray[int64_t] normalize_i8_timestamps(const int64_t[:] stamps, tzinfo tz):
cpdef ndarray normalize_i8_timestamps(ndarray stamps, tzinfo tz):
# stamps is int64_t, arbitrary ndim
"""
Normalize each of the (nanosecond) timezone aware timestamps in the given
array by rounding down to the beginning of the day (i.e. midnight).
Expand All @@ -238,28 +274,35 @@ cpdef ndarray[int64_t] normalize_i8_timestamps(const int64_t[:] stamps, tzinfo t
"""
cdef:
Localizer info = Localizer(tz)
int64_t utc_val, local_val
Py_ssize_t i, n = stamps.shape[0]
int64_t utc_val, local_val, res_val
Py_ssize_t i, n = stamps.size
Py_ssize_t pos = -1 # unused, avoid not-initialized warning

int64_t[::1] result = np.empty(n, dtype=np.int64)
ndarray result = cnp.PyArray_EMPTY(stamps.ndim, stamps.shape, cnp.NPY_INT64, 0)
cnp.broadcast mi = cnp.PyArray_MultiIterNew2(result, stamps)

for i in range(n):
utc_val = stamps[i]
# Analogous to: utc_val = stamps[i]
utc_val = (<int64_t*>cnp.PyArray_MultiIter_DATA(mi, 1))[0]

if utc_val == NPY_NAT:
result[i] = NPY_NAT
continue
res_val = NPY_NAT
else:
local_val = info.utc_val_to_local_val(utc_val, &pos)
res_val = local_val - (local_val % DAY_NANOS)

local_val = info.utc_val_to_local_val(utc_val, &pos)
# Analogous to: result[i] = res_val
(<int64_t*>cnp.PyArray_MultiIter_DATA(mi, 0))[0] = res_val

result[i] = local_val - (local_val % DAY_NANOS)
cnp.PyArray_MultiIter_NEXT(mi)

return result.base # `.base` to access underlying ndarray
return result


@cython.wraparound(False)
@cython.boundscheck(False)
def is_date_array_normalized(const int64_t[:] stamps, tzinfo tz=None) -> bool:
def is_date_array_normalized(ndarray stamps, tzinfo tz=None) -> bool:
# stamps is int64_t, arbitrary ndim
"""
Check if all of the given (nanosecond) timestamps are normalized to
midnight, i.e. hour == minute == second == 0. If the optional timezone
Expand All @@ -277,16 +320,21 @@ def is_date_array_normalized(const int64_t[:] stamps, tzinfo tz=None) -> bool:
cdef:
Localizer info = Localizer(tz)
int64_t utc_val, local_val
Py_ssize_t i, n = stamps.shape[0]
Py_ssize_t i, n = stamps.size
Py_ssize_t pos = -1 # unused, avoid not-initialized warning
cnp.flatiter it = cnp.PyArray_IterNew(stamps)

for i in range(n):
utc_val = stamps[i]
# Analogous to: utc_val = stamps[i]
utc_val = cnp.PyArray_GETITEM(stamps, cnp.PyArray_ITER_DATA(it))

local_val = info.utc_val_to_local_val(utc_val, &pos)

if local_val % DAY_NANOS != 0:
return False

cnp.PyArray_ITER_NEXT(it)

return True


Expand Down