PERF: array_strptime (#55898)

* PERF: array_strptime avoid object path

* creso fixup

* Fixup leftover assertion

* object instead of raise

* post-merge fixup

* post-merge fixup

Author: jbrockmendel

pandas/_libs/tslibs/strptime.pyx

@@ -47,14 +47,18 @@ from numpy cimport (
 )
 
 from pandas._libs.missing cimport checknull_with_nat_and_na
-from pandas._libs.tslibs.conversion cimport get_datetime64_nanos
+from pandas._libs.tslibs.conversion cimport (
+    get_datetime64_nanos,
+    parse_pydatetime,
+)
 from pandas._libs.tslibs.dtypes cimport (
     get_supported_reso,
     npy_unit_to_abbrev,
     npy_unit_to_attrname,
 )
 from pandas._libs.tslibs.nattype cimport (
     NPY_NAT,
+    c_NaT as NaT,
     c_nat_strings as nat_strings,
 )
 from pandas._libs.tslibs.np_datetime cimport (
@@ -65,7 +69,6 @@ from pandas._libs.tslibs.np_datetime cimport (
     npy_datetimestruct,
     npy_datetimestruct_to_datetime,
     pydate_to_dt64,
-    pydatetime_to_dt64,
     string_to_dts,
 )
 
@@ -82,6 +85,8 @@ from pandas._libs.util cimport (
 
 from pandas._libs.tslibs.timestamps import Timestamp
 
+from pandas._libs.tslibs.tzconversion cimport tz_localize_to_utc_single
+
 cnp.import_array()
 
 
@@ -314,11 +319,13 @@ def array_strptime(
         Py_ssize_t i, n = len(values)
         npy_datetimestruct dts
         int64_t[::1] iresult
-        object[::1] result_timezone
         object val, tz
+        bint seen_datetime_offset = False
         bint is_raise = errors=="raise"
         bint is_ignore = errors=="ignore"
         bint is_coerce = errors=="coerce"
+        bint is_same_offsets
+        set out_tzoffset_vals = set()
         tzinfo tz_out = None
         bint iso_format = format_is_iso(fmt)
         NPY_DATETIMEUNIT out_bestunit, item_reso
@@ -338,7 +345,6 @@ def array_strptime(
         abbrev = npy_unit_to_abbrev(creso)
     result = np.empty(n, dtype=f"M8[{abbrev}]")
     iresult = result.view("i8")
-    result_timezone = np.empty(n, dtype="object")
 
     dts.us = dts.ps = dts.as = 0
 
@@ -361,23 +367,18 @@ def array_strptime(
                 if infer_reso:
                     creso = state.creso
                 tz_out = state.process_datetime(val, tz_out, utc)
-                if isinstance(val, _Timestamp):
-                    val = (<_Timestamp>val)._as_creso(creso)
-                    iresult[i] = val.tz_localize(None)._value
-                else:
-                    iresult[i] = pydatetime_to_dt64(
-                        val.replace(tzinfo=None), &dts, reso=creso
-                    )
-                result_timezone[i] = val.tzinfo
+                iresult[i] = parse_pydatetime(val, &dts, state.creso)
                 continue
             elif PyDate_Check(val):
+                state.found_other = True
                 item_reso = NPY_DATETIMEUNIT.NPY_FR_s
                 state.update_creso(item_reso)
                 if infer_reso:
                     creso = state.creso
                 iresult[i] = pydate_to_dt64(val, &dts, reso=creso)
                 continue
             elif cnp.is_datetime64_object(val):
+                state.found_other = True
                 item_reso = get_supported_reso(get_datetime64_unit(val))
                 state.update_creso(item_reso)
                 if infer_reso:
@@ -418,13 +419,17 @@ def array_strptime(
                         f"Out of bounds {attrname} timestamp: {val}"
                     ) from err
                 if out_local == 1:
-                    # Store the out_tzoffset in seconds
-                    # since we store the total_seconds of
-                    # dateutil.tz.tzoffset objects
+                    nsecs = out_tzoffset * 60
+                    out_tzoffset_vals.add(nsecs)
+                    seen_datetime_offset = True
                     tz = timezone(timedelta(minutes=out_tzoffset))
-                    result_timezone[i] = tz
-                    out_local = 0
-                    out_tzoffset = 0
+                    value = tz_localize_to_utc_single(
+                        value, tz, ambiguous="raise", nonexistent=None, creso=creso
+                    )
+                else:
+                    tz = None
+                    out_tzoffset_vals.add("naive")
+                    state.found_naive_str = True
                 iresult[i] = value
                 continue
 
@@ -450,14 +455,34 @@ def array_strptime(
             state.update_creso(item_reso)
             if infer_reso:
                 creso = state.creso
+
             try:
                 iresult[i] = npy_datetimestruct_to_datetime(creso, &dts)
             except OverflowError as err:
                 attrname = npy_unit_to_attrname[creso]
                 raise OutOfBoundsDatetime(
                     f"Out of bounds {attrname} timestamp: {val}"
                 ) from err
-            result_timezone[i] = tz
+
+            if tz is not None:
+                ival = iresult[i]
+                iresult[i] = tz_localize_to_utc_single(
+                    ival, tz, ambiguous="raise", nonexistent=None, creso=creso
+                )
+                nsecs = (ival - iresult[i])
+                if creso == NPY_FR_ns:
+                    nsecs = nsecs // 10**9
+                elif creso == NPY_DATETIMEUNIT.NPY_FR_us:
+                    nsecs = nsecs // 10**6
+                elif creso == NPY_DATETIMEUNIT.NPY_FR_ms:
+                    nsecs = nsecs // 10**3
+
+                out_tzoffset_vals.add(nsecs)
+                seen_datetime_offset = True
+            else:
+                state.found_naive_str = True
+                tz = None
+                out_tzoffset_vals.add("naive")
 
         except (ValueError, OutOfBoundsDatetime) as ex:
             ex.args = (
@@ -474,7 +499,37 @@ def array_strptime(
                 continue
             elif is_raise:
                 raise
-            return values, []
+            return values, None
+
+    if seen_datetime_offset and not utc:
+        is_same_offsets = len(out_tzoffset_vals) == 1
+        if not is_same_offsets or (state.found_naive or state.found_other):
+            result2 = _array_strptime_object_fallback(
+                values, fmt=fmt, exact=exact, errors=errors, utc=utc
+            )
+            return result2, None
+        elif tz_out is not None:
+            # GH#55693
+            tz_offset = out_tzoffset_vals.pop()
+            tz_out2 = timezone(timedelta(seconds=tz_offset))
+            if not tz_compare(tz_out, tz_out2):
+                # e.g. test_to_datetime_mixed_offsets_with_utc_false_deprecated
+                result2 = _array_strptime_object_fallback(
+                    values, fmt=fmt, exact=exact, errors=errors, utc=utc
+                )
+                return result2, None
+            # e.g. test_guess_datetime_format_with_parseable_formats
+        else:
+            # e.g. test_to_datetime_iso8601_with_timezone_valid
+            tz_offset = out_tzoffset_vals.pop()
+            tz_out = timezone(timedelta(seconds=tz_offset))
+    elif not utc:
+        if tz_out and (state.found_other or state.found_naive_str):
+            # found_other indicates a tz-naive int, float, dt64, or date
+            result2 = _array_strptime_object_fallback(
+                values, fmt=fmt, exact=exact, errors=errors, utc=utc
+            )
+            return result2, None
 
     if infer_reso:
         if state.creso_ever_changed:
@@ -488,7 +543,6 @@ def array_strptime(
                 utc=utc,
                 creso=state.creso,
             )
-
         elif state.creso == NPY_DATETIMEUNIT.NPY_FR_GENERIC:
             # i.e. we never encountered anything non-NaT, default to "s". This
             # ensures that insert and concat-like operations with NaT
@@ -499,7 +553,7 @@ def array_strptime(
             #  a second pass.
             abbrev = npy_unit_to_abbrev(state.creso)
             result = iresult.base.view(f"M8[{abbrev}]")
-    return result, result_timezone.base
+    return result, tz_out
 
 
 cdef tzinfo _parse_with_format(
@@ -737,6 +791,157 @@ cdef tzinfo _parse_with_format(
     return tz
 
 
+def _array_strptime_object_fallback(
+    ndarray[object] values,
+    str fmt,
+    bint exact=True,
+    errors="raise",
+    bint utc=False,
+):
+
+    cdef:
+        Py_ssize_t i, n = len(values)
+        npy_datetimestruct dts
+        int64_t iresult
+        object val
+        tzinfo tz
+        bint is_raise = errors=="raise"
+        bint is_ignore = errors=="ignore"
+        bint is_coerce = errors=="coerce"
+        bint iso_format = format_is_iso(fmt)
+        NPY_DATETIMEUNIT creso, out_bestunit, item_reso
+        int out_local = 0, out_tzoffset = 0
+        bint string_to_dts_succeeded = 0
+
+    assert is_raise or is_ignore or is_coerce
+
+    item_reso = NPY_DATETIMEUNIT.NPY_FR_GENERIC
+    format_regex, locale_time = _get_format_regex(fmt)
+
+    result = np.empty(n, dtype=object)
+
+    dts.us = dts.ps = dts.as = 0
+
+    for i in range(n):
+        val = values[i]
+        try:
+            if isinstance(val, str):
+                if len(val) == 0 or val in nat_strings:
+                    result[i] = NaT
+                    continue
+            elif checknull_with_nat_and_na(val):
+                result[i] = NaT
+                continue
+            elif PyDateTime_Check(val):
+                result[i] = Timestamp(val)
+                continue
+            elif PyDate_Check(val):
+                result[i] = Timestamp(val)
+                continue
+            elif cnp.is_datetime64_object(val):
+                result[i] = Timestamp(val)
+                continue
+            elif (
+                    (is_integer_object(val) or is_float_object(val))
+                    and (val != val or val == NPY_NAT)
+            ):
+                result[i] = NaT
+                continue
+            else:
+                val = str(val)
+
+            if fmt == "ISO8601":
+                string_to_dts_succeeded = not string_to_dts(
+                    val, &dts, &out_bestunit, &out_local,
+                    &out_tzoffset, False, None, False
+                )
+            elif iso_format:
+                string_to_dts_succeeded = not string_to_dts(
+                    val, &dts, &out_bestunit, &out_local,
+                    &out_tzoffset, False, fmt, exact
+                )
+            if string_to_dts_succeeded:
+                # No error reported by string_to_dts, pick back up
+                # where we left off
+                creso = get_supported_reso(out_bestunit)
+                try:
+                    value = npy_datetimestruct_to_datetime(creso, &dts)
+                except OverflowError as err:
+                    raise OutOfBoundsDatetime(
+                        f"Out of bounds nanosecond timestamp: {val}"
+                    ) from err
+                if out_local == 1:
+                    tz = timezone(timedelta(minutes=out_tzoffset))
+                    value = tz_localize_to_utc_single(
+                        value, tz, ambiguous="raise", nonexistent=None, creso=creso
+                    )
+                else:
+                    tz = None
+                ts = Timestamp._from_value_and_reso(value, creso, tz)
+                result[i] = ts
+                continue
+
+            if parse_today_now(val, &iresult, utc, NPY_FR_ns):
+                result[i] = Timestamp(val)
+                continue
+
+            # Some ISO formats can't be parsed by string_to_dts
+            # For example, 6-digit YYYYMD. So, if there's an error, and a format
+            # was specified, then try the string-matching code below. If the format
+            # specified was 'ISO8601', then we need to error, because
+            # only string_to_dts handles mixed ISO8601 formats.
+            if not string_to_dts_succeeded and fmt == "ISO8601":
+                raise ValueError(f"Time data {val} is not ISO8601 format")
+
+            tz = _parse_with_format(
+                val, fmt, exact, format_regex, locale_time, &dts, &item_reso
+            )
+            try:
+                iresult = npy_datetimestruct_to_datetime(item_reso, &dts)
+            except OverflowError as err:
+                raise OutOfBoundsDatetime(
+                    f"Out of bounds nanosecond timestamp: {val}"
+                ) from err
+            if tz is not None:
+                iresult = tz_localize_to_utc_single(
+                    iresult, tz, ambiguous="raise", nonexistent=None, creso=item_reso
+                )
+            ts = Timestamp._from_value_and_reso(iresult, item_reso, tz)
+            result[i] = ts
+
+        except (ValueError, OutOfBoundsDatetime) as ex:
+            ex.args = (
+                f"{str(ex)}, at position {i}. You might want to try:\n"
+                "    - passing `format` if your strings have a consistent format;\n"
+                "    - passing `format='ISO8601'` if your strings are "
+                "all ISO8601 but not necessarily in exactly the same format;\n"
+                "    - passing `format='mixed'`, and the format will be "
+                "inferred for each element individually. "
+                "You might want to use `dayfirst` alongside this.",
+            )
+            if is_coerce:
+                result[i] = NaT
+                continue
+            elif is_raise:
+                raise
+            return values
+
+    import warnings
+
+    from pandas.util._exceptions import find_stack_level
+    warnings.warn(
+        "In a future version of pandas, parsing datetimes with mixed time "
+        "zones will raise an error unless `utc=True`. Please specify `utc=True` "
+        "to opt in to the new behaviour and silence this warning. "
+        "To create a `Series` with mixed offsets and `object` dtype, "
+        "please use `apply` and `datetime.datetime.strptime`",
+        FutureWarning,
+        stacklevel=find_stack_level(),
+    )
+
+    return result
+
+
 class TimeRE(_TimeRE):
     """
     Handle conversion from format directives to regexes.