[Python] Snowflake Python UDF that acts as a window function

I would like to use a vectorized (pandas) udf in Snowflake Snowpark as a window function.

from snowflake.snowpark import Window
data.withColumn("result", my_udf('column1', 'column2').over(Window.partitionBy('column3')))


where my_udf is a function that receives a Pandas DataFrame with two columns:

def my_udf(data: pd.DataFrame):
...


I tried

import pandas as pd
from snowflake.snowpark import Session
from snowflake.snowpark.types import PandasDataFrame

class my_func:
def process(self, df: PandasDataFrame[int, int]) -> PandasDataFrame[int]:
return df[["col2"]]
process._sf_vectorized_input = pd.DataFrame


my_udtf = session.udtf.register(
my_func,
output_schema=["result"],
input_names=['"col1"', '"col2"']
)

data.withColumn(
"result",
my_udtf('col1', 'col2').over(Window.partitionBy("col3"))
)


but it fails with "TypeError: over() takes 1 positional argument but 2 were given"

Continue reading...