[GitHub] [beam] ibzib commented on a change in pull request #14940: [BEAM-9547] Add support for `Series.dt` methods

[GitBox <gi...@apache.org>]

ibzib commented on a change in pull request #14940:
URL: https://github.com/apache/beam/pull/14940#discussion_r656594314



##########
File path: sdks/python/apache_beam/dataframe/pandas_doctests_test.py
##########
@@ -545,43 +543,46 @@ def test_string_tests(self):
 
   def test_datetime_tests(self):
     # TODO(BEAM-10721)
-    datetimelike_result = doctests.testmod(
-        pd.core.arrays.datetimelike,
+    indexes_accessors_result = doctests.testmod(
+        pd.core.indexes.accessors,
         use_beam=False,
         skip={
-            'pandas.core.arrays.datetimelike.AttributesMixin._unbox_scalar': [
+            'pandas.core.indexes.accessors.TimedeltaProperties': [
+                # Seems like an upstream bug. The property is 'second'
+                'seconds_series.dt.seconds'
+            ],
+
+            # Test data creation fails for these:

Review comment:
       Why?

##########
File path: sdks/python/apache_beam/dataframe/frames.py
##########
@@ -4042,6 +4047,134 @@ def codes(self):
               make_cat_func(method), name=method,
               base=pd.core.arrays.categorical.CategoricalAccessor))
 
+class _DeferredDatetimeMethods(frame_base.DeferredBase):
+  @property  # type: ignore
+  @frame_base.with_docs_from(pd.core.indexes.accessors.DatetimeProperties)
+  def tz(self):
+    return self._expr.proxy().dt.tz
+
+  @property  # type: ignore
+  @frame_base.with_docs_from(pd.core.indexes.accessors.DatetimeProperties)
+  def freq(self):
+    return self._expr.proxy().dt.freq
+
+  @frame_base.with_docs_from(pd.core.indexes.accessors.DatetimeProperties)
+  def tz_localize(self, *args, ambiguous='infer', **kwargs):
+    """``ambiguous`` cannot be set to ``"infer"`` as its semantics are
+    order-sensitive. Similarly, specifying ``ambiguous`` as an
+    :class:`~numpy.ndarray` is order-sensitive, but you can achieve similar
+    functionality by specifying ``ambiguous`` as a Series."""
+    if isinstance(ambiguous, np.ndarray):
+      raise frame_base.WontImplementError(
+          "tz_localize(ambiguous=ndarray) is not supported because it makes "
+          "this operation sensitive to the order of the data. Please use a "
+          "DeferredSeries instead.",
+          reason="order-sensitive")
+    elif isinstance(ambiguous, frame_base.DeferredFrame):
+      return frame_base.DeferredFrame.wrap(
+          expressions.ComputedExpression(
+              'tz_localize',
+              lambda s,
+              ambiguous: s.dt.tz_localize(*args, ambiguous=ambiguous, **kwargs),
+              [self._expr, ambiguous._expr],
+              requires_partition_by=partitionings.Index(),
+              preserves_partition_by=partitionings.Arbitrary()))
+    elif ambiguous == 'infer':
+      # infer attempts to infer based on the order of the timestamps
+      raise frame_base.WontImplementError(
+          f"tz_localize(ambiguous={ambiguous!r}) is not allowed because it "
+          "makes this operation sensitive to the order of the data.",
+          reason="order-sensitive")
+
+    return frame_base.DeferredFrame.wrap(
+        expressions.ComputedExpression(
+            'tz_localize',
+            lambda s: s.dt.tz_localize(*args, ambiguous=ambiguous, **kwargs),
+            [self._expr],
+            requires_partition_by=partitionings.Arbitrary(),
+            preserves_partition_by=partitionings.Arbitrary()))
+
+
+  to_period = frame_base.not_implemented_method(

Review comment:
       Why was this left unimplemented? It might be helpful to leave a comment for later.

##########
File path: sdks/python/apache_beam/dataframe/frames_test.py
##########
@@ -891,6 +891,74 @@ def test_cat(self):
     self._run_test(lambda s: s.cat.as_unordered(), s)
     self._run_test(lambda s: s.cat.codes, s)
 
+  @parameterized.expand(frames.ELEMENTWISE_DATETIME_PROPERTIES)
+  def test_dt_property(self, prop_name):
+    # Generate a series with a lot of unique timestamps
+    s = pd.Series(
+        pd.date_range('1/1/2000', periods=100, freq='m') +

Review comment:
       Should we add a timezone here too, or does it not matter?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org