From c5809e18df138d4f5d380d114b7174a7d0c1bad1 Mon Sep 17 00:00:00 2001 From: Dangercrow Date: Sat, 11 May 2019 18:27:19 +0100 Subject: [PATCH 1/2] Fix return type for DataFrame.groupBy / cube / rollup --- third_party/3/pyspark/sql/dataframe.pyi | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/third_party/3/pyspark/sql/dataframe.pyi b/third_party/3/pyspark/sql/dataframe.pyi index 03d96628..557c86ad 100644 --- a/third_party/3/pyspark/sql/dataframe.pyi +++ b/third_party/3/pyspark/sql/dataframe.pyi @@ -10,6 +10,7 @@ from py4j.java_gateway import JavaObject # type: ignore from pyspark.sql._typing import ColumnOrName, Literal, LiteralType from pyspark.sql.types import * from pyspark.sql.context import SQLContext +from pyspark.sql.group import GroupedData from pyspark.sql.readwriter import DataFrameWriter from pyspark.sql.streaming import DataStreamWriter from pyspark.sql.column import Column @@ -91,17 +92,17 @@ class DataFrame: def selectExpr(self, *expr: List[str]) -> 'DataFrame': ... def filter(self, condition: ColumnOrName) -> 'DataFrame': ... @overload - def groupBy(self, *cols: ColumnOrName) -> 'DataFrame': ... + def groupBy(self, *cols: ColumnOrName) -> GroupedData': ... @overload - def groupBy(self, __cols: List[ColumnOrName]) -> 'DataFrame': ... + def groupBy(self, __cols: List[ColumnOrName]) -> GroupedData: ... @overload - def rollup(self, *cols: ColumnOrName) -> 'DataFrame': ... + def rollup(self, *cols: ColumnOrName) -> GroupedData: ... @overload - def rollup(self, __cols: List[ColumnOrName]) -> 'DataFrame': ... + def rollup(self, __cols: List[ColumnOrName]) -> GroupedData: ... @overload - def cube(self, *cols: ColumnOrName) -> 'DataFrame': ... + def cube(self, *cols: ColumnOrName) -> GroupedData: ... @overload - def cube(self, __cols: List[ColumnOrName]) -> 'DataFrame': ... + def cube(self, __cols: List[ColumnOrName]) -> GroupedData: ... def agg(self, *exprs: Union[Column, Dict[str, str]]) -> 'DataFrame': ... def union(self, other: 'DataFrame') -> 'DataFrame': ... def unionAll(self, other: 'DataFrame') -> 'DataFrame': ... From 804bd41ee8c99ffb96a9755024fe61dea6621ed5 Mon Sep 17 00:00:00 2001 From: dangercrow Date: Mon, 20 May 2019 10:10:33 +0100 Subject: [PATCH 2/2] Remove dangling quote I should really install a proper IDE on my machine. --- third_party/3/pyspark/sql/dataframe.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/3/pyspark/sql/dataframe.pyi b/third_party/3/pyspark/sql/dataframe.pyi index 557c86ad..9a025c17 100644 --- a/third_party/3/pyspark/sql/dataframe.pyi +++ b/third_party/3/pyspark/sql/dataframe.pyi @@ -92,7 +92,7 @@ class DataFrame: def selectExpr(self, *expr: List[str]) -> 'DataFrame': ... def filter(self, condition: ColumnOrName) -> 'DataFrame': ... @overload - def groupBy(self, *cols: ColumnOrName) -> GroupedData': ... + def groupBy(self, *cols: ColumnOrName) -> GroupedData: ... @overload def groupBy(self, __cols: List[ColumnOrName]) -> GroupedData: ... @overload