feat: Allow tuples, strings and lists for fillna

The docstring refers to the word list, but the API will allow str, list and tuple.
zero323 · Sep 2, 2020 · 466f8a1 · 466f8a1
1 parent d0c15b9
commit 466f8a1
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git 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
@@ -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: ...