From 39ce9c70f960e85d3cea828d3ef9461b24d8c0ce Mon Sep 17 00:00:00 2001 From: Tristan Nixon Date: Tue, 2 Jan 2024 23:03:18 -0800 Subject: [PATCH] can't import from private pyspark typing modules workaround is to copy definitions locally --- python/tempo/tsdf.py | 3 +-- python/tempo/typing.py | 20 ++++++++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) create mode 100644 python/tempo/typing.py diff --git a/python/tempo/tsdf.py b/python/tempo/tsdf.py index 02358e6c..0d7b549c 100644 --- a/python/tempo/tsdf.py +++ b/python/tempo/tsdf.py @@ -16,8 +16,6 @@ from pyspark.sql.column import Column from pyspark.sql.dataframe import DataFrame from pyspark.sql.types import DataType, StructType -from pyspark.sql._typing import ColumnOrName -from pyspark.sql.pandas._typing import PandasMapIterFunction, PandasGroupedMapFunction from pyspark.sql.window import Window, WindowSpec import tempo.interpol as t_interpolation @@ -26,6 +24,7 @@ import tempo.utils as t_utils from tempo.intervals import IntervalsDF from tempo.tsschema import CompositeTSIndex, TSIndex, TSSchema, WindowBuilder +from tempo.typing import ColumnOrName, PandasMapIterFunction, PandasGroupedMapFunction logger = logging.getLogger(__name__) diff --git a/python/tempo/typing.py b/python/tempo/typing.py new file mode 100644 index 00000000..a29c3d9f --- /dev/null +++ b/python/tempo/typing.py @@ -0,0 +1,20 @@ +from typing import Union, Callable, Iterable, Any + +from pyspark.sql import Column + +from pandas.core.frame import DataFrame as PandasDataFrame + +from pyspark.sql.pandas._typing import PandasMapIterFunction, PandasGroupedMapFunction + +# These definitions were copied from private pypark modules: +# - pyspark.sql._typing +# - pyspark.sql.pandas._typing + +ColumnOrName = Union[Column, str] + +PandasMapIterFunction = Callable[[Iterable[PandasDataFrame]], Iterable[PandasDataFrame]] + +PandasGroupedMapFunction = Union[ + Callable[[PandasDataFrame], PandasDataFrame], + Callable[[Any, PandasDataFrame], PandasDataFrame], +]