Skip to content

Commit 8a555ef

Browse files
sfc-gh-helmeleegysfc-gh-helmeleegy
authored
SNOW-2476996: Add support for rolling.min/max/count/sum/mean/std/var/sem/corr in faster pandas (#3972)
1 parent b1c9bff commit 8a555ef

File tree

3 files changed

+414
-0
lines changed

3 files changed

+414
-0
lines changed

‎CHANGELOG.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,15 @@
8383
- `dt.month_name`
8484
- `dt.day_name`
8585
- `dt.strftime`
86+
- `rolling.min`
87+
- `rolling.max`
88+
- `rolling.count`
89+
- `rolling.sum`
90+
- `rolling.mean`
91+
- `rolling.std`
92+
- `rolling.var`
93+
- `rolling.sem`
94+
- `rolling.corr`
8695
- Make faster pandas disabled by default (opt-in instead of opt-out).
8796

8897
## 1.42.0 (2025-10-28)

‎src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py‎

Lines changed: 316 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16893,6 +16893,35 @@ def rolling_count(
1689316893
numeric_only: bool = False,
1689416894
*args: Any,
1689516895
**kwargs: Any,
16896+
) -> "SnowflakeQueryCompiler":
16897+
"""
16898+
Wrapper around _rolling_count_internal to be supported in faster pandas.
16899+
"""
16900+
relaxed_query_compiler = None
16901+
if self._relaxed_query_compiler is not None:
16902+
relaxed_query_compiler = (
16903+
self._relaxed_query_compiler._rolling_count_internal(
16904+
fold_axis=fold_axis,
16905+
rolling_kwargs=rolling_kwargs,
16906+
numeric_only=numeric_only,
16907+
**kwargs,
16908+
)
16909+
)
16910+
qc = self._rolling_count_internal(
16911+
fold_axis=fold_axis,
16912+
rolling_kwargs=rolling_kwargs,
16913+
numeric_only=numeric_only,
16914+
**kwargs,
16915+
)
16916+
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
16917+
16918+
def _rolling_count_internal(
16919+
self,
16920+
fold_axis: Union[int, str],
16921+
rolling_kwargs: dict,
16922+
numeric_only: bool = False,
16923+
*args: Any,
16924+
**kwargs: Any,
1689616925
) -> "SnowflakeQueryCompiler":
1689716926
return self._window_agg(
1689816927
window_func=WindowFunction.ROLLING,
@@ -16910,6 +16939,39 @@ def rolling_sum(
1691016939
engine_kwargs: Optional[dict[str, bool]] = None,
1691116940
*args: Any,
1691216941
**kwargs: Any,
16942+
) -> "SnowflakeQueryCompiler":
16943+
"""
16944+
Wrapper around _rolling_sum_internal to be supported in faster pandas.
16945+
"""
16946+
relaxed_query_compiler = None
16947+
if self._relaxed_query_compiler is not None:
16948+
relaxed_query_compiler = self._relaxed_query_compiler._rolling_sum_internal(
16949+
fold_axis=fold_axis,
16950+
rolling_kwargs=rolling_kwargs,
16951+
numeric_only=numeric_only,
16952+
engine=engine,
16953+
engine_kwargs=engine_kwargs,
16954+
**kwargs,
16955+
)
16956+
qc = self._rolling_sum_internal(
16957+
fold_axis=fold_axis,
16958+
rolling_kwargs=rolling_kwargs,
16959+
numeric_only=numeric_only,
16960+
engine=engine,
16961+
engine_kwargs=engine_kwargs,
16962+
**kwargs,
16963+
)
16964+
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
16965+
16966+
def _rolling_sum_internal(
16967+
self,
16968+
fold_axis: Union[int, str],
16969+
rolling_kwargs: dict,
16970+
numeric_only: bool = False,
16971+
engine: Optional[Literal["cython", "numba"]] = None,
16972+
engine_kwargs: Optional[dict[str, bool]] = None,
16973+
*args: Any,
16974+
**kwargs: Any,
1691316975
) -> "SnowflakeQueryCompiler":
1691416976
WarningMessage.warning_if_engine_args_is_set(
1691516977
"rolling_sum", engine, engine_kwargs
@@ -16930,6 +16992,41 @@ def rolling_mean(
1693016992
engine_kwargs: Optional[dict[str, bool]] = None,
1693116993
*args: Any,
1693216994
**kwargs: Any,
16995+
) -> "SnowflakeQueryCompiler":
16996+
"""
16997+
Wrapper around _rolling_mean_internal to be supported in faster pandas.
16998+
"""
16999+
relaxed_query_compiler = None
17000+
if self._relaxed_query_compiler is not None:
17001+
relaxed_query_compiler = (
17002+
self._relaxed_query_compiler._rolling_mean_internal(
17003+
fold_axis=fold_axis,
17004+
rolling_kwargs=rolling_kwargs,
17005+
numeric_only=numeric_only,
17006+
engine=engine,
17007+
engine_kwargs=engine_kwargs,
17008+
**kwargs,
17009+
)
17010+
)
17011+
qc = self._rolling_mean_internal(
17012+
fold_axis=fold_axis,
17013+
rolling_kwargs=rolling_kwargs,
17014+
numeric_only=numeric_only,
17015+
engine=engine,
17016+
engine_kwargs=engine_kwargs,
17017+
**kwargs,
17018+
)
17019+
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
17020+
17021+
def _rolling_mean_internal(
17022+
self,
17023+
fold_axis: Union[int, str],
17024+
rolling_kwargs: dict,
17025+
numeric_only: bool = False,
17026+
engine: Optional[Literal["cython", "numba"]] = None,
17027+
engine_kwargs: Optional[dict[str, bool]] = None,
17028+
*args: Any,
17029+
**kwargs: Any,
1693317030
) -> "SnowflakeQueryCompiler":
1693417031
WarningMessage.warning_if_engine_args_is_set(
1693517032
"rolling_mean", engine, engine_kwargs
@@ -16962,6 +17059,42 @@ def rolling_var(
1696217059
engine_kwargs: Optional[dict[str, bool]] = None,
1696317060
*args: Any,
1696417061
**kwargs: Any,
17062+
) -> "SnowflakeQueryCompiler":
17063+
"""
17064+
Wrapper around _rolling_var_internal to be supported in faster pandas.
17065+
"""
17066+
relaxed_query_compiler = None
17067+
if self._relaxed_query_compiler is not None:
17068+
relaxed_query_compiler = self._relaxed_query_compiler._rolling_var_internal(
17069+
fold_axis=fold_axis,
17070+
rolling_kwargs=rolling_kwargs,
17071+
ddof=ddof,
17072+
numeric_only=numeric_only,
17073+
engine=engine,
17074+
engine_kwargs=engine_kwargs,
17075+
**kwargs,
17076+
)
17077+
qc = self._rolling_var_internal(
17078+
fold_axis=fold_axis,
17079+
rolling_kwargs=rolling_kwargs,
17080+
ddof=ddof,
17081+
numeric_only=numeric_only,
17082+
engine=engine,
17083+
engine_kwargs=engine_kwargs,
17084+
**kwargs,
17085+
)
17086+
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
17087+
17088+
def _rolling_var_internal(
17089+
self,
17090+
fold_axis: Union[int, str],
17091+
rolling_kwargs: dict,
17092+
ddof: int = 1,
17093+
numeric_only: bool = False,
17094+
engine: Optional[Literal["cython", "numba"]] = None,
17095+
engine_kwargs: Optional[dict[str, bool]] = None,
17096+
*args: Any,
17097+
**kwargs: Any,
1696517098
) -> "SnowflakeQueryCompiler":
1696617099
WarningMessage.warning_if_engine_args_is_set(
1696717100
"rolling_var", engine, engine_kwargs
@@ -16983,6 +17116,42 @@ def rolling_std(
1698317116
engine_kwargs: Optional[dict[str, bool]] = None,
1698417117
*args: Any,
1698517118
**kwargs: Any,
17119+
) -> "SnowflakeQueryCompiler":
17120+
"""
17121+
Wrapper around _rolling_std_internal to be supported in faster pandas.
17122+
"""
17123+
relaxed_query_compiler = None
17124+
if self._relaxed_query_compiler is not None:
17125+
relaxed_query_compiler = self._relaxed_query_compiler._rolling_std_internal(
17126+
fold_axis=fold_axis,
17127+
rolling_kwargs=rolling_kwargs,
17128+
ddof=ddof,
17129+
numeric_only=numeric_only,
17130+
engine=engine,
17131+
engine_kwargs=engine_kwargs,
17132+
**kwargs,
17133+
)
17134+
qc = self._rolling_std_internal(
17135+
fold_axis=fold_axis,
17136+
rolling_kwargs=rolling_kwargs,
17137+
ddof=ddof,
17138+
numeric_only=numeric_only,
17139+
engine=engine,
17140+
engine_kwargs=engine_kwargs,
17141+
**kwargs,
17142+
)
17143+
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
17144+
17145+
def _rolling_std_internal(
17146+
self,
17147+
fold_axis: Union[int, str],
17148+
rolling_kwargs: dict,
17149+
ddof: int = 1,
17150+
numeric_only: bool = False,
17151+
engine: Optional[Literal["cython", "numba"]] = None,
17152+
engine_kwargs: Optional[dict[str, bool]] = None,
17153+
*args: Any,
17154+
**kwargs: Any,
1698617155
) -> "SnowflakeQueryCompiler":
1698717156
WarningMessage.warning_if_engine_args_is_set(
1698817157
"rolling_var", engine, engine_kwargs
@@ -17003,6 +17172,39 @@ def rolling_min(
1700317172
engine_kwargs: Optional[dict[str, bool]] = None,
1700417173
*args: Any,
1700517174
**kwargs: Any,
17175+
) -> "SnowflakeQueryCompiler":
17176+
"""
17177+
Wrapper around _rolling_min_internal to be supported in faster pandas.
17178+
"""
17179+
relaxed_query_compiler = None
17180+
if self._relaxed_query_compiler is not None:
17181+
relaxed_query_compiler = self._relaxed_query_compiler._rolling_min_internal(
17182+
fold_axis=fold_axis,
17183+
rolling_kwargs=rolling_kwargs,
17184+
numeric_only=numeric_only,
17185+
engine=engine,
17186+
engine_kwargs=engine_kwargs,
17187+
**kwargs,
17188+
)
17189+
qc = self._rolling_min_internal(
17190+
fold_axis=fold_axis,
17191+
rolling_kwargs=rolling_kwargs,
17192+
numeric_only=numeric_only,
17193+
engine=engine,
17194+
engine_kwargs=engine_kwargs,
17195+
**kwargs,
17196+
)
17197+
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
17198+
17199+
def _rolling_min_internal(
17200+
self,
17201+
fold_axis: Union[int, str],
17202+
rolling_kwargs: dict,
17203+
numeric_only: bool = False,
17204+
engine: Optional[Literal["cython", "numba"]] = None,
17205+
engine_kwargs: Optional[dict[str, bool]] = None,
17206+
*args: Any,
17207+
**kwargs: Any,
1700617208
) -> "SnowflakeQueryCompiler":
1700717209
WarningMessage.warning_if_engine_args_is_set(
1700817210
"rolling_min", engine, engine_kwargs
@@ -17023,6 +17225,39 @@ def rolling_max(
1702317225
engine_kwargs: Optional[dict[str, bool]] = None,
1702417226
*args: Any,
1702517227
**kwargs: Any,
17228+
) -> "SnowflakeQueryCompiler":
17229+
"""
17230+
Wrapper around _rolling_max_internal to be supported in faster pandas.
17231+
"""
17232+
relaxed_query_compiler = None
17233+
if self._relaxed_query_compiler is not None:
17234+
relaxed_query_compiler = self._relaxed_query_compiler._rolling_max_internal(
17235+
fold_axis=fold_axis,
17236+
rolling_kwargs=rolling_kwargs,
17237+
numeric_only=numeric_only,
17238+
engine=engine,
17239+
engine_kwargs=engine_kwargs,
17240+
**kwargs,
17241+
)
17242+
qc = self._rolling_max_internal(
17243+
fold_axis=fold_axis,
17244+
rolling_kwargs=rolling_kwargs,
17245+
numeric_only=numeric_only,
17246+
engine=engine,
17247+
engine_kwargs=engine_kwargs,
17248+
**kwargs,
17249+
)
17250+
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
17251+
17252+
def _rolling_max_internal(
17253+
self,
17254+
fold_axis: Union[int, str],
17255+
rolling_kwargs: dict,
17256+
numeric_only: bool = False,
17257+
engine: Optional[Literal["cython", "numba"]] = None,
17258+
engine_kwargs: Optional[dict[str, bool]] = None,
17259+
*args: Any,
17260+
**kwargs: Any,
1702617261
) -> "SnowflakeQueryCompiler":
1702717262
WarningMessage.warning_if_engine_args_is_set(
1702817263
"rolling_max", engine, engine_kwargs
@@ -17043,6 +17278,57 @@ def rolling_corr(
1704317278
ddof: int = 1,
1704417279
numeric_only: bool = False,
1704517280
**kwargs: Any,
17281+
) -> "SnowflakeQueryCompiler":
17282+
"""
17283+
Wrapper around _rolling_corr_internal to be supported in faster pandas.
17284+
"""
17285+
relaxed_query_compiler = None
17286+
if self._relaxed_query_compiler is not None and (
17287+
not isinstance(other, (Series, DataFrame))
17288+
or other._query_compiler._relaxed_query_compiler is not None
17289+
):
17290+
if isinstance(other, (Series, DataFrame)):
17291+
if isinstance(other, Series):
17292+
new_other = Series(
17293+
query_compiler=other._query_compiler._relaxed_query_compiler
17294+
)
17295+
else: # DataFrame
17296+
new_other = DataFrame(
17297+
query_compiler=other._query_compiler._relaxed_query_compiler
17298+
)
17299+
else:
17300+
new_other = other
17301+
relaxed_query_compiler = (
17302+
self._relaxed_query_compiler._rolling_corr_internal(
17303+
fold_axis=fold_axis,
17304+
rolling_kwargs=rolling_kwargs,
17305+
other=new_other,
17306+
pairwise=pairwise,
17307+
ddof=ddof,
17308+
numeric_only=numeric_only,
17309+
**kwargs,
17310+
)
17311+
)
17312+
qc = self._rolling_corr_internal(
17313+
fold_axis=fold_axis,
17314+
rolling_kwargs=rolling_kwargs,
17315+
other=other,
17316+
pairwise=pairwise,
17317+
ddof=ddof,
17318+
numeric_only=numeric_only,
17319+
**kwargs,
17320+
)
17321+
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
17322+
17323+
def _rolling_corr_internal(
17324+
self,
17325+
fold_axis: Union[int, str],
17326+
rolling_kwargs: dict,
17327+
other: Optional[SnowparkDataFrame] = None,
17328+
pairwise: Optional[bool] = None,
17329+
ddof: int = 1,
17330+
numeric_only: bool = False,
17331+
**kwargs: Any,
1704617332
) -> "SnowflakeQueryCompiler":
1704717333
if other is None:
1704817334
ErrorMessage.parameter_not_implemented_error(
@@ -17133,6 +17419,36 @@ def rolling_sem(
1713317419
numeric_only: bool = False,
1713417420
*args: Any,
1713517421
**kwargs: Any,
17422+
) -> "SnowflakeQueryCompiler":
17423+
"""
17424+
Wrapper around _rolling_sem_internal to be supported in faster pandas.
17425+
"""
17426+
relaxed_query_compiler = None
17427+
if self._relaxed_query_compiler is not None:
17428+
relaxed_query_compiler = self._relaxed_query_compiler._rolling_sem_internal(
17429+
fold_axis=fold_axis,
17430+
rolling_kwargs=rolling_kwargs,
17431+
ddof=ddof,
17432+
numeric_only=numeric_only,
17433+
**kwargs,
17434+
)
17435+
qc = self._rolling_sem_internal(
17436+
fold_axis=fold_axis,
17437+
rolling_kwargs=rolling_kwargs,
17438+
ddof=ddof,
17439+
numeric_only=numeric_only,
17440+
**kwargs,
17441+
)
17442+
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
17443+
17444+
def _rolling_sem_internal(
17445+
self,
17446+
fold_axis: Union[int, str],
17447+
rolling_kwargs: dict,
17448+
ddof: int = 1,
17449+
numeric_only: bool = False,
17450+
*args: Any,
17451+
**kwargs: Any,
1713617452
) -> "SnowflakeQueryCompiler":
1713717453
return self._window_agg(
1713817454
window_func=WindowFunction.ROLLING,

0 commit comments

Comments
 (0)