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.

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

feat: support Series.dt.normalize #483

Merged
merged 11 commits into from
Mar 21, 2024
Merged
7 changes: 7 additions & 0 deletions bigframes/core/compile/scalar_op_compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -632,6 +632,13 @@ def year_op_impl(x: ibis_types.Value):
return typing.cast(ibis_types.TimestampValue, x).year().cast(ibis_dtypes.int64)


@scalar_op_compiler.register_unary_op(ops.normalize_op)
def normalize_op_impl(x: ibis_types.Value):
result_type = x.type()
result = x.truncate("D")
return result.cast(result_type)


# Parameterized ops
@scalar_op_compiler.register_unary_op(ops.StructFieldOp, pass_op=True)
def struct_field_op_impl(x: ibis_types.Value, op: ops.StructFieldOp):
Expand Down
1 change: 1 addition & 0 deletions bigframes/operations/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,7 @@ def create_ternary_op(
second_op = create_unary_op(name="second", type_rule=op_typing.INTEGER)
time_op = create_unary_op(name="time", type_rule=op_typing.INTEGER)
year_op = create_unary_op(name="year", type_rule=op_typing.INTEGER)
normalize_op = create_unary_op(name="normalize")
## Trigonometry Ops
sin_op = create_unary_op(name="sin", type_rule=op_typing.REAL_NUMERIC)
cos_op = create_unary_op(name="cos", type_rule=op_typing.REAL_NUMERIC)
Expand Down
3 changes: 3 additions & 0 deletions bigframes/operations/datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,3 +94,6 @@ def unit(self) -> str:

def strftime(self, date_format: str) -> series.Series:
return self._apply_unary_op(ops.StrftimeOp(date_format=date_format))

def normalize(self) -> series.Series:
return self._apply_unary_op(ops.normalize_op)
17 changes: 17 additions & 0 deletions tests/system/small/operations/test_datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -266,3 +266,20 @@ def test_dt_strftime_time():
bf_result, expected_result, check_index_type=False, check_dtype=False
)
assert bf_result.dtype == "string[pyarrow]"


@pytest.mark.parametrize(
("col_name",),
DATETIME_COL_NAMES,
)
@skip_legacy_pandas
def test_dt_normalize(scalars_dfs, col_name):
scalars_df, scalars_pandas_df = scalars_dfs
bf_series: bigframes.series.Series = scalars_df[col_name]
bf_result = bf_series.dt.normalize().to_pandas()
Copy link
Contributor

Choose a reason for hiding this comment

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

we don't define bf_series, but just call scalars_df[col_name].dt.normalize().to_pandas(). Would that work also?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

True, bf_series is unnecessary. Fixed, thanks!

pd_result = scalars_pandas_df[col_name].dt.normalize()

assert_series_equal(
pd_result.astype(bf_series.dtype), # normalize preserves type
bf_result,
)
Original file line number Diff line number Diff line change
Expand Up @@ -36,3 +36,34 @@ def strftime(self, date_format: str):
bigframes.series.Series of formatted strings.
"""
raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)

def normalize(self):
"""
Convert times to midnight.

The time component of the date-time is converted to midnight i.e.
00:00:00. This is useful in cases when the time does not matter.
The return dtype will match the source series.

This method is available on Series with datetime values under the
.dt accessor.

**Examples:**

>>> import pandas as pd
>>> import bigframes.pandas as bpd
>>> s = bpd.Series(pd.date_range(
... start='2014-08-01 10:00',
... freq='h',
... periods=3,
... tz='Asia/Calcutta')) # note timezones will be converted to UTC here
>>> s.dt.normalize()
0 2014-08-01 00:00:00+00:00
1 2014-08-01 00:00:00+00:00
2 2014-08-01 00:00:00+00:00
dtype: timestamp[us, tz=UTC][pyarrow]

Returns:
bigframes.series.Series of the same dtype as the data.
"""
raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)