From 466f8a106c55e24d7411ba0d323b30a57b132af5 Mon Sep 17 00:00:00 2001
From: Oliver Willekens <oliver.willekens@dataminded.be>
Date: Fri, 28 Aug 2020 21:37:00 +0200
Subject: [PATCH] feat: Allow tuples, strings and lists for fillna

The docstring refers to the word list, but the API will allow str, list
and tuple.
---
 test-data/unit/sql-dataframe.test       | 12 ++++++++++++
 third_party/3/pyspark/sql/dataframe.pyi |  4 +++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/test-data/unit/sql-dataframe.test b/test-data/unit/sql-dataframe.test
index 6ec0cb82..3d2267ee 100644
--- a/test-data/unit/sql-dataframe.test
+++ b/test-data/unit/sql-dataframe.test
@@ -108,3 +108,15 @@ df.drop(col("id"), col("foo"))  # E: No overload variant of "drop" of "DataFrame
                                 # N:     <1 more non-matching overload not shown>
 
 [out]
+
+[case fillNullValues]
+from pyspark.sql import SparkSession
+
+spark = SparkSession.builder.getOrCreate()
+df = spark.createDataFrame([(1,2)], schema=("id1", "id2"))
+
+df.fillna(value=1, subset="id1")
+df.fillna(value=1, subset=("id1", "id2"))
+df.fillna(value=1, subset=["id1"])
+
+[out]
diff --git a/third_party/3/pyspark/sql/dataframe.pyi b/third_party/3/pyspark/sql/dataframe.pyi
index 25013076..5977ab73 100644
--- a/third_party/3/pyspark/sql/dataframe.pyi
+++ b/third_party/3/pyspark/sql/dataframe.pyi
@@ -198,7 +198,9 @@ class DataFrame(PandasMapOpsMixin, PandasConversionMixin):
     ) -> DataFrame: ...
     @overload
     def fillna(
-        self, value: LiteralType, subset: Optional[List[str]] = ...
+        self,
+        value: LiteralType,
+        subset: Optional[Union[str, Tuple[str, ...], List[str]]] = ...,
     ) -> DataFrame: ...
     @overload
     def fillna(self, value: Dict[str, LiteralType]) -> DataFrame: ...