Пользовательская функция Snowflake Python, действующая как оконная функция ⇐ Python

[Anonymous]

Я хотел бы использовать векторизованный (pandas) udf в Snowflake Snowpark в качестве оконной функции.

Код: Выделить всё

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

где my_udf — это функция, которая получает кадр данных Pandas с двумя столбцами:

Код: Выделить всё

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

Я пробовал

Код: Выделить всё

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"))
)

но это не удается: «TypeError: over() принимает 1 позиционный аргумент, но были заданы 2»

Подробнее здесь: https://stackoverflow.com/questions/792 ... w-function