You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2024/07/18 00:49:07 UTC
(spark) branch master updated: [SPARK-48924][PS] Add a pandas-like `make_interval` helper function
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 171c7947d1af [SPARK-48924][PS] Add a pandas-like `make_interval` helper function
171c7947d1af is described below
commit 171c7947d1afaf6a423ab2532c98b7aabdacfffc
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Thu Jul 18 09:48:55 2024 +0900
[SPARK-48924][PS] Add a pandas-like `make_interval` helper function
### What changes were proposed in this pull request?
Add a pandas-like `make_interval` helper function
### Why are the changes needed?
factor it out as a helper function to be reusable
### Does this PR introduce _any_ user-facing change?
No, internal change only
### How was this patch tested?
CI
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #47385 from zhengruifeng/ps_simplify_make_interval.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
python/pyspark/pandas/resample.py | 35 ++++++++++----------------------
python/pyspark/pandas/spark/functions.py | 16 ++++++++++++++-
2 files changed, 26 insertions(+), 25 deletions(-)
diff --git a/python/pyspark/pandas/resample.py b/python/pyspark/pandas/resample.py
index 5557ca2af773..02db8fd91b9b 100644
--- a/python/pyspark/pandas/resample.py
+++ b/python/pyspark/pandas/resample.py
@@ -25,7 +25,6 @@ from typing import (
Generic,
List,
Optional,
- Union,
)
import numpy as np
@@ -42,6 +41,7 @@ from pyspark.sql.types import (
from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm.
from pyspark.pandas._typing import FrameLike
from pyspark.pandas.frame import DataFrame
+from pyspark.pandas.spark import functions as SF
from pyspark.pandas.internal import (
InternalField,
InternalFrame,
@@ -130,19 +130,6 @@ class Resampler(Generic[FrameLike], metaclass=ABCMeta):
def _agg_columns_scols(self) -> List[Column]:
return [s.spark.column for s in self._agg_columns]
- def get_make_interval( # type: ignore[return]
- self, unit: str, col: Union[Column, int, float]
- ) -> Column:
- col = col if not isinstance(col, (int, float)) else F.lit(col)
- if unit == "MONTH":
- return F.make_interval(months=col)
- if unit == "HOUR":
- return F.make_interval(hours=col)
- if unit == "MINUTE":
- return F.make_interval(mins=col)
- if unit == "SECOND":
- return F.make_interval(secs=col)
-
def _bin_timestamp(self, origin: pd.Timestamp, ts_scol: Column) -> Column:
key_type = self._resamplekey_type
origin_scol = F.lit(origin)
@@ -203,18 +190,18 @@ class Resampler(Generic[FrameLike], metaclass=ABCMeta):
truncated_ts_scol = F.date_trunc("MONTH", ts_scol)
edge_label = truncated_ts_scol
if left_closed and right_labeled:
- edge_label += self.get_make_interval("MONTH", n)
+ edge_label += SF.make_interval("MONTH", n)
elif right_closed and left_labeled:
- edge_label -= self.get_make_interval("MONTH", n)
+ edge_label -= SF.make_interval("MONTH", n)
if left_labeled:
non_edge_label = F.when(
mod == 0,
- truncated_ts_scol - self.get_make_interval("MONTH", n),
- ).otherwise(truncated_ts_scol - self.get_make_interval("MONTH", mod))
+ truncated_ts_scol - SF.make_interval("MONTH", n),
+ ).otherwise(truncated_ts_scol - SF.make_interval("MONTH", mod))
else:
non_edge_label = F.when(mod == 0, truncated_ts_scol).otherwise(
- truncated_ts_scol - self.get_make_interval("MONTH", mod - n)
+ truncated_ts_scol - SF.make_interval("MONTH", mod - n)
)
ret = F.to_timestamp(
@@ -292,19 +279,19 @@ class Resampler(Generic[FrameLike], metaclass=ABCMeta):
edge_label = truncated_ts_scol
if left_closed and right_labeled:
- edge_label += self.get_make_interval(unit_str, n)
+ edge_label += SF.make_interval(unit_str, n)
elif right_closed and left_labeled:
- edge_label -= self.get_make_interval(unit_str, n)
+ edge_label -= SF.make_interval(unit_str, n)
if left_labeled:
non_edge_label = F.when(mod == 0, truncated_ts_scol).otherwise(
- truncated_ts_scol - self.get_make_interval(unit_str, mod)
+ truncated_ts_scol - SF.make_interval(unit_str, mod)
)
else:
non_edge_label = F.when(
mod == 0,
- truncated_ts_scol + self.get_make_interval(unit_str, n),
- ).otherwise(truncated_ts_scol - self.get_make_interval(unit_str, mod - n))
+ truncated_ts_scol + SF.make_interval(unit_str, n),
+ ).otherwise(truncated_ts_scol - SF.make_interval(unit_str, mod - n))
ret = F.when(edge_cond, edge_label).otherwise(non_edge_label)
diff --git a/python/pyspark/pandas/spark/functions.py b/python/pyspark/pandas/spark/functions.py
index db1cc423078a..8abeff655ea5 100644
--- a/python/pyspark/pandas/spark/functions.py
+++ b/python/pyspark/pandas/spark/functions.py
@@ -17,8 +17,9 @@
"""
Additional Spark functions used in pandas-on-Spark.
"""
-from pyspark.sql.column import Column
+from pyspark.sql import Column, functions as F
from pyspark.sql.utils import is_remote
+from typing import Union
def product(col: Column, dropna: bool) -> Column:
@@ -171,3 +172,16 @@ def null_index(col: Column) -> Column:
sc = SparkContext._active_spark_context
return Column(sc._jvm.PythonSQLUtils.nullIndex(col._jc))
+
+
+def make_interval(unit: str, e: Union[Column, int, float]) -> Column:
+ unit_mapping = {
+ "YEAR": "years",
+ "MONTH": "months",
+ "WEEK": "weeks",
+ "DAY": "days",
+ "HOUR": "hours",
+ "MINUTE": "mins",
+ "SECOND": "secs",
+ }
+ return F.make_interval(**{unit_mapping[unit]: F.lit(e)})
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org