duckdb 1.5.0.dev44__cp39-cp39-macosx_11_0_arm64.whl → 1.5.0.dev94__cp39-cp39-macosx_11_0_arm64.whl

duckdb/experimental/spark/sql/group.py

@@ -1,4 +1,4 @@
-#
+#  # noqa: D100
 # Licensed to the Apache Software Foundation (ASF) under one or more
 # contributor license agreements.  See the NOTICE file distributed with
 # this work for additional information regarding copyright ownership.
@@ -15,26 +15,27 @@
 # limitations under the License.
 #
 
-from ..exception import ContributionsAcceptedError
-from typing import Callable, TYPE_CHECKING, overload, Dict, Union, List
+from typing import TYPE_CHECKING, Callable, Union, overload
 
+from ..exception import ContributionsAcceptedError
 from .column import Column
-from .session import SparkSession
 from .dataframe import DataFrame
 from .functions import _to_column_expr
-from ._typing import ColumnOrName
 from .types import NumericType
 
+# Only import symbols needed for type checking if something is type checking
 if TYPE_CHECKING:
-    from ._typing import LiteralType
+    from ._typing import ColumnOrName
+    from .session import SparkSession
 
 __all__ = ["GroupedData", "Grouping"]
 
+
 def _api_internal(self: "GroupedData", name: str, *cols: str) -> DataFrame:
     expressions = ",".join(list(cols))
     group_by = str(self._grouping) if self._grouping else ""
     projections = self._grouping.get_columns()
-    jdf = getattr(self._df.relation, "apply")(
+    jdf = self._df.relation.apply(
         function_name=name,  # aggregate function
         function_aggr=expressions,  # inputs to aggregate
         group_expr=group_by,  # groups
@@ -42,6 +43,7 @@ def _api_internal(self: "GroupedData", name: str, *cols: str) -> DataFrame:
     )
     return DataFrame(jdf, self.session)
 
+
 def df_varargs_api(f: Callable[..., DataFrame]) -> Callable[..., DataFrame]:
     def _api(self: "GroupedData", *cols: str) -> DataFrame:
         name = f.__name__
@@ -52,49 +54,49 @@ def df_varargs_api(f: Callable[..., DataFrame]) -> Callable[..., DataFrame]:
     return _api
 
 
-class Grouping:
-    def __init__(self, *cols: "ColumnOrName", **kwargs):
+class Grouping:  # noqa: D101
+    def __init__(self, *cols: "ColumnOrName", **kwargs) -> None:  # noqa: D107
         self._type = ""
         self._cols = [_to_column_expr(x) for x in cols]
-        if 'special' in kwargs:
-            special = kwargs['special']
+        if "special" in kwargs:
+            special = kwargs["special"]
             accepted_special = ["cube", "rollup"]
             assert special in accepted_special
             self._type = special
 
-    def get_columns(self) -> str:
+    def get_columns(self) -> str:  # noqa: D102
         columns = ",".join([str(x) for x in self._cols])
         return columns
 
-    def __str__(self):
+    def __str__(self) -> str:  # noqa: D105
         columns = self.get_columns()
         if self._type:
-            return self._type + '(' + columns + ')'
+            return self._type + "(" + columns + ")"
         return columns
 
 
 class GroupedData:
-    """
-    A set of methods for aggregations on a :class:`DataFrame`,
+    """A set of methods for aggregations on a :class:`DataFrame`,
     created by :func:`DataFrame.groupBy`.
 
-    """
+    """  # noqa: D205
 
-    def __init__(self, grouping: Grouping, df: DataFrame):
+    def __init__(self, grouping: Grouping, df: DataFrame) -> None:  # noqa: D107
         self._grouping = grouping
         self._df = df
         self.session: SparkSession = df.session
 
-    def __repr__(self) -> str:
+    def __repr__(self) -> str:  # noqa: D105
         return str(self._df)
 
     def count(self) -> DataFrame:
         """Counts the number of records for each group.
 
-        Examples
+        Examples:
         --------
         >>> df = spark.createDataFrame(
-        ...      [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")], ["age", "name"])
+        ...     [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")], ["age", "name"]
+        ... )
         >>> df.show()
         +---+-----+
         |age| name|
@@ -115,7 +117,7 @@ class GroupedData:
         |  Bob|    2|
         +-----+-----+
         """
-        return _api_internal(self, "count").withColumnRenamed('count_star()', 'count')
+        return _api_internal(self, "count").withColumnRenamed("count_star()", "count")
 
     @df_varargs_api
     def mean(self, *cols: str) -> DataFrame:
@@ -139,11 +141,12 @@ class GroupedData:
         cols : str
             column names. Non-numeric columns are ignored.
 
-        Examples
+        Examples:
         --------
-        >>> df = spark.createDataFrame([
-        ...     (2, "Alice", 80), (3, "Alice", 100),
-        ...     (5, "Bob", 120), (10, "Bob", 140)], ["age", "name", "height"])
+        >>> df = spark.createDataFrame(
+        ...     [(2, "Alice", 80), (3, "Alice", 100), (5, "Bob", 120), (10, "Bob", 140)],
+        ...     ["age", "name", "height"],
+        ... )
         >>> df.show()
         +---+-----+------+
         |age| name|height|
@@ -156,7 +159,7 @@ class GroupedData:
 
         Group-by name, and calculate the mean of the age in each group.
 
-        >>> df.groupBy("name").avg('age').sort("name").show()
+        >>> df.groupBy("name").avg("age").sort("name").show()
         +-----+--------+
         | name|avg(age)|
         +-----+--------+
@@ -166,7 +169,7 @@ class GroupedData:
 
         Calculate the mean of the age and height in all data.
 
-        >>> df.groupBy().avg('age', 'height').show()
+        >>> df.groupBy().avg("age", "height").show()
         +--------+-----------+
         |avg(age)|avg(height)|
         +--------+-----------+
@@ -177,18 +180,19 @@ class GroupedData:
         if len(columns) == 0:
             schema = self._df.schema
             # Take only the numeric types of the relation
-            columns: List[str] = [x.name for x in schema.fields if isinstance(x.dataType, NumericType)]
+            columns: list[str] = [x.name for x in schema.fields if isinstance(x.dataType, NumericType)]
         return _api_internal(self, "avg", *columns)
 
     @df_varargs_api
     def max(self, *cols: str) -> DataFrame:
         """Computes the max value for each numeric columns for each group.
 
-        Examples
+        Examples:
         --------
-        >>> df = spark.createDataFrame([
-        ...     (2, "Alice", 80), (3, "Alice", 100),
-        ...     (5, "Bob", 120), (10, "Bob", 140)], ["age", "name", "height"])
+        >>> df = spark.createDataFrame(
+        ...     [(2, "Alice", 80), (3, "Alice", 100), (5, "Bob", 120), (10, "Bob", 140)],
+        ...     ["age", "name", "height"],
+        ... )
         >>> df.show()
         +---+-----+------+
         |age| name|height|
@@ -228,11 +232,12 @@ class GroupedData:
         cols : str
             column names. Non-numeric columns are ignored.
 
-        Examples
+        Examples:
         --------
-        >>> df = spark.createDataFrame([
-        ...     (2, "Alice", 80), (3, "Alice", 100),
-        ...     (5, "Bob", 120), (10, "Bob", 140)], ["age", "name", "height"])
+        >>> df = spark.createDataFrame(
+        ...     [(2, "Alice", 80), (3, "Alice", 100), (5, "Bob", 120), (10, "Bob", 140)],
+        ...     ["age", "name", "height"],
+        ... )
         >>> df.show()
         +---+-----+------+
         |age| name|height|
@@ -272,11 +277,12 @@ class GroupedData:
         cols : str
             column names. Non-numeric columns are ignored.
 
-        Examples
+        Examples:
         --------
-        >>> df = spark.createDataFrame([
-        ...     (2, "Alice", 80), (3, "Alice", 100),
-        ...     (5, "Bob", 120), (10, "Bob", 140)], ["age", "name", "height"])
+        >>> df = spark.createDataFrame(
+        ...     [(2, "Alice", 80), (3, "Alice", 100), (5, "Bob", 120), (10, "Bob", 140)],
+        ...     ["age", "name", "height"],
+        ... )
         >>> df.show()
         +---+-----+------+
         |age| name|height|
@@ -308,14 +314,12 @@ class GroupedData:
         """
 
     @overload
-    def agg(self, *exprs: Column) -> DataFrame:
-        ...
+    def agg(self, *exprs: Column) -> DataFrame: ...
 
     @overload
-    def agg(self, __exprs: Dict[str, str]) -> DataFrame:
-        ...
+    def agg(self, __exprs: dict[str, str]) -> DataFrame: ...  # noqa: PYI063
 
-    def agg(self, *exprs: Union[Column, Dict[str, str]]) -> DataFrame:
+    def agg(self, *exprs: Union[Column, dict[str, str]]) -> DataFrame:
         """Compute aggregates and returns the result as a :class:`DataFrame`.
 
         The available aggregate functions can be:
@@ -347,17 +351,18 @@ class GroupedData:
             a dict mapping from column name (string) to aggregate functions (string),
             or a list of :class:`Column`.
 
-        Notes
+        Notes:
         -----
         Built-in aggregation functions and group aggregate pandas UDFs cannot be mixed
         in a single call to this function.
 
-        Examples
+        Examples:
         --------
         >>> from pyspark.sql import functions as F
         >>> from pyspark.sql.functions import pandas_udf, PandasUDFType
         >>> df = spark.createDataFrame(
-        ...      [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")], ["age", "name"])
+        ...     [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")], ["age", "name"]
+        ... )
         >>> df.show()
         +---+-----+
         |age| name|
@@ -393,10 +398,9 @@ class GroupedData:
 
         Same as above but uses pandas UDF.
 
-        >>> @pandas_udf('int', PandasUDFType.GROUPED_AGG)  # doctest: +SKIP
+        >>> @pandas_udf("int", PandasUDFType.GROUPED_AGG)  # doctest: +SKIP
         ... def min_udf(v):
         ...     return v.min()
-        ...
         >>> df.groupBy(df.name).agg(min_udf(df.age)).sort("name").show()  # doctest: +SKIP
         +-----+------------+
         | name|min_udf(age)|
@@ -417,4 +421,4 @@ class GroupedData:
             rel = self._df.relation.select(*expressions, groups=group_by)
         return DataFrame(rel, self.session)
 
-    # TODO: add 'pivot'
+    # TODO: add 'pivot'  # noqa: TD002, TD003

duckdb/experimental/spark/sql/readwriter.py

@@ -1,11 +1,9 @@
-from typing import TYPE_CHECKING, List, Optional, Union, cast
+from typing import TYPE_CHECKING, Optional, Union, cast  # noqa: D100
 
+from ..errors import PySparkNotImplementedError, PySparkTypeError
 from ..exception import ContributionsAcceptedError
 from .types import StructType
 
-
-from ..errors import PySparkNotImplementedError, PySparkTypeError
-
 PrimitiveType = Union[bool, float, int, str]
 OptionalPrimitiveType = Optional[PrimitiveType]
 
@@ -14,19 +12,19 @@ if TYPE_CHECKING:
     from duckdb.experimental.spark.sql.session import SparkSession
 
 
-class DataFrameWriter:
-    def __init__(self, dataframe: "DataFrame"):
+class DataFrameWriter:  # noqa: D101
+    def __init__(self, dataframe: "DataFrame") -> None:  # noqa: D107
         self.dataframe = dataframe
 
-    def saveAsTable(self, table_name: str) -> None:
+    def saveAsTable(self, table_name: str) -> None:  # noqa: D102
         relation = self.dataframe.relation
         relation.create(table_name)
 
-    def parquet(
+    def parquet(  # noqa: D102
         self,
         path: str,
         mode: Optional[str] = None,
-        partitionBy: Union[str, List[str], None] = None,
+        partitionBy: Union[str, list[str], None] = None,
         compression: Optional[str] = None,
     ) -> None:
         relation = self.dataframe.relation
@@ -37,7 +35,7 @@ class DataFrameWriter:
 
         relation.write_parquet(path, compression=compression)
 
-    def csv(
+    def csv(  # noqa: D102
         self,
         path: str,
         mode: Optional[str] = None,
@@ -57,7 +55,7 @@ class DataFrameWriter:
         encoding: Optional[str] = None,
         emptyValue: Optional[str] = None,
         lineSep: Optional[str] = None,
-    ):
+    ) -> None:
         if mode not in (None, "overwrite"):
             raise NotImplementedError
         if escapeQuotes:
@@ -88,13 +86,13 @@ class DataFrameWriter:
         )
 
 
-class DataFrameReader:
-    def __init__(self, session: "SparkSession"):
+class DataFrameReader:  # noqa: D101
+    def __init__(self, session: "SparkSession") -> None:  # noqa: D107
         self.session = session
 
-    def load(
+    def load(  # noqa: D102
         self,
-        path: Optional[Union[str, List[str]]] = None,
+        path: Optional[Union[str, list[str]]] = None,
         format: Optional[str] = None,
         schema: Optional[Union[StructType, str]] = None,
         **options: OptionalPrimitiveType,
@@ -102,7 +100,7 @@ class DataFrameReader:
         from duckdb.experimental.spark.sql.dataframe import DataFrame
 
         if not isinstance(path, str):
-            raise ImportError
+            raise TypeError
         if options:
             raise ContributionsAcceptedError
 
@@ -123,15 +121,15 @@ class DataFrameReader:
         if schema:
             if not isinstance(schema, StructType):
                 raise ContributionsAcceptedError
-            schema = cast(StructType, schema)
+            schema = cast("StructType", schema)
             types, names = schema.extract_types_and_names()
             df = df._cast_types(types)
             df = df.toDF(names)
         raise NotImplementedError
 
-    def csv(
+    def csv(  # noqa: D102
         self,
-        path: Union[str, List[str]],
+        path: Union[str, list[str]],
         schema: Optional[Union[StructType, str]] = None,
         sep: Optional[str] = None,
         encoding: Optional[str] = None,
@@ -225,7 +223,7 @@ class DataFrameReader:
         dtype = None
         names = None
         if schema:
-            schema = cast(StructType, schema)
+            schema = cast("StructType", schema)
             dtype, names = schema.extract_types_and_names()
 
         rel = self.session.conn.read_csv(
@@ -247,13 +245,15 @@ class DataFrameReader:
             df = df.toDF(*names)
         return df
 
-    def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame":
+    def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame":  # noqa: D102
         input = list(paths)
         if len(input) != 1:
-            raise NotImplementedError("Only single paths are supported for now")
+            msg = "Only single paths are supported for now"
+            raise NotImplementedError(msg)
         option_amount = len(options.keys())
         if option_amount != 0:
-            raise ContributionsAcceptedError("Options are not supported")
+            msg = "Options are not supported"
+            raise ContributionsAcceptedError(msg)
         path = input[0]
         rel = self.session.conn.read_parquet(path)
         from ..sql.dataframe import DataFrame
@@ -263,7 +263,7 @@ class DataFrameReader:
 
     def json(
         self,
-        path: Union[str, List[str]],
+        path: Union[str, list[str]],
         schema: Optional[Union[StructType, str]] = None,
         primitivesAsString: Optional[Union[bool, str]] = None,
         prefersDecimal: Optional[Union[bool, str]] = None,
@@ -289,8 +289,7 @@ class DataFrameReader:
         modifiedAfter: Optional[Union[bool, str]] = None,
         allowNonNumericNumbers: Optional[Union[bool, str]] = None,
     ) -> "DataFrame":
-        """
-        Loads JSON files and returns the results as a :class:`DataFrame`.
+        """Loads JSON files and returns the results as a :class:`DataFrame`.
 
         `JSON Lines <http://jsonlines.org/>`_ (newline-delimited JSON) is supported by default.
         For JSON (one record per file), set the ``multiLine`` parameter to ``true``.
@@ -321,16 +320,16 @@ class DataFrameReader:
 
             .. # noqa
 
-        Examples
+        Examples:
         --------
         Write a DataFrame into a JSON file and read it back.
 
         >>> import tempfile
         >>> with tempfile.TemporaryDirectory() as d:
         ...     # Write a DataFrame into a JSON file
-        ...     spark.createDataFrame(
-        ...         [{"age": 100, "name": "Hyukjin Kwon"}]
-        ...     ).write.mode("overwrite").format("json").save(d)
+        ...     spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}]).write.mode(
+        ...         "overwrite"
+        ...     ).format("json").save(d)
         ...
         ...     # Read the JSON file as a DataFrame.
         ...     spark.read.json(d).show()
@@ -340,102 +339,89 @@ class DataFrameReader:
         |100|Hyukjin Kwon|
         +---+------------+
         """
-
         if schema is not None:
-            raise ContributionsAcceptedError("The 'schema' option is not supported")
+            msg = "The 'schema' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if primitivesAsString is not None:
-            raise ContributionsAcceptedError(
-                "The 'primitivesAsString' option is not supported"
-            )
+            msg = "The 'primitivesAsString' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if prefersDecimal is not None:
-            raise ContributionsAcceptedError(
-                "The 'prefersDecimal' option is not supported"
-            )
+            msg = "The 'prefersDecimal' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if allowComments is not None:
-            raise ContributionsAcceptedError(
-                "The 'allowComments' option is not supported"
-            )
+            msg = "The 'allowComments' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if allowUnquotedFieldNames is not None:
-            raise ContributionsAcceptedError(
-                "The 'allowUnquotedFieldNames' option is not supported"
-            )
+            msg = "The 'allowUnquotedFieldNames' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if allowSingleQuotes is not None:
-            raise ContributionsAcceptedError(
-                "The 'allowSingleQuotes' option is not supported"
-            )
+            msg = "The 'allowSingleQuotes' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if allowNumericLeadingZero is not None:
-            raise ContributionsAcceptedError(
-                "The 'allowNumericLeadingZero' option is not supported"
-            )
+            msg = "The 'allowNumericLeadingZero' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if allowBackslashEscapingAnyCharacter is not None:
-            raise ContributionsAcceptedError(
-                "The 'allowBackslashEscapingAnyCharacter' option is not supported"
-            )
+            msg = "The 'allowBackslashEscapingAnyCharacter' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if mode is not None:
-            raise ContributionsAcceptedError("The 'mode' option is not supported")
+            msg = "The 'mode' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if columnNameOfCorruptRecord is not None:
-            raise ContributionsAcceptedError(
-                "The 'columnNameOfCorruptRecord' option is not supported"
-            )
+            msg = "The 'columnNameOfCorruptRecord' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if dateFormat is not None:
-            raise ContributionsAcceptedError("The 'dateFormat' option is not supported")
+            msg = "The 'dateFormat' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if timestampFormat is not None:
-            raise ContributionsAcceptedError(
-                "The 'timestampFormat' option is not supported"
-            )
+            msg = "The 'timestampFormat' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if multiLine is not None:
-            raise ContributionsAcceptedError("The 'multiLine' option is not supported")
+            msg = "The 'multiLine' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if allowUnquotedControlChars is not None:
-            raise ContributionsAcceptedError(
-                "The 'allowUnquotedControlChars' option is not supported"
-            )
+            msg = "The 'allowUnquotedControlChars' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if lineSep is not None:
-            raise ContributionsAcceptedError("The 'lineSep' option is not supported")
+            msg = "The 'lineSep' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if samplingRatio is not None:
-            raise ContributionsAcceptedError(
-                "The 'samplingRatio' option is not supported"
-            )
+            msg = "The 'samplingRatio' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if dropFieldIfAllNull is not None:
-            raise ContributionsAcceptedError(
-                "The 'dropFieldIfAllNull' option is not supported"
-            )
+            msg = "The 'dropFieldIfAllNull' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if encoding is not None:
-            raise ContributionsAcceptedError("The 'encoding' option is not supported")
+            msg = "The 'encoding' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if locale is not None:
-            raise ContributionsAcceptedError("The 'locale' option is not supported")
+            msg = "The 'locale' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if pathGlobFilter is not None:
-            raise ContributionsAcceptedError(
-                "The 'pathGlobFilter' option is not supported"
-            )
+            msg = "The 'pathGlobFilter' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if recursiveFileLookup is not None:
-            raise ContributionsAcceptedError(
-                "The 'recursiveFileLookup' option is not supported"
-            )
+            msg = "The 'recursiveFileLookup' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if modifiedBefore is not None:
-            raise ContributionsAcceptedError(
-                "The 'modifiedBefore' option is not supported"
-            )
+            msg = "The 'modifiedBefore' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if modifiedAfter is not None:
-            raise ContributionsAcceptedError(
-                "The 'modifiedAfter' option is not supported"
-            )
+            msg = "The 'modifiedAfter' option is not supported"
+            raise ContributionsAcceptedError(msg)
         if allowNonNumericNumbers is not None:
-            raise ContributionsAcceptedError(
-                "The 'allowNonNumericNumbers' option is not supported"
-            )
+            msg = "The 'allowNonNumericNumbers' option is not supported"
+            raise ContributionsAcceptedError(msg)
 
         if isinstance(path, str):
             path = [path]
-        if  isinstance(path, list):
+        if isinstance(path, list):
             if len(path) == 1:
                 rel = self.session.conn.read_json(path[0])
                 from .dataframe import DataFrame
 
                 df = DataFrame(rel, self.session)
                 return df
-            raise PySparkNotImplementedError(
-                message="Only a single path is supported for now"
-            )
+            raise PySparkNotImplementedError(message="Only a single path is supported for now")
         else:
             raise PySparkTypeError(
                 error_class="NOT_STR_OR_LIST_OF_RDD",
@@ -446,4 +432,4 @@ class DataFrameReader:
             )
 
 
-__all__ = ["DataFrameWriter", "DataFrameReader"]
+__all__ = ["DataFrameReader", "DataFrameWriter"]