duckdb 1.4.1.dev141__cp310-cp310-macosx_10_9_universal2.whl → 1.5.0.dev44__cp310-cp310-macosx_10_9_universal2.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,27 +15,26 @@
 # limitations under the License.
 #
 
-from typing import TYPE_CHECKING, Callable, Union, overload
-
 from ..exception import ContributionsAcceptedError
+from typing import Callable, TYPE_CHECKING, overload, Dict, Union, List
+
 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 ColumnOrName
-    from .session import SparkSession
+    from ._typing import LiteralType
 
 __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 = self._df.relation.apply(
+    jdf = getattr(self._df.relation, "apply")(
         function_name=name,  # aggregate function
         function_aggr=expressions,  # inputs to aggregate
         group_expr=group_by,  # groups
@@ -43,7 +42,6 @@ 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__
@@ -54,49 +52,49 @@ def df_varargs_api(f: Callable[..., DataFrame]) -> Callable[..., DataFrame]:
     return _api
 
 
-class Grouping:  # noqa: D101
-    def __init__(self, *cols: "ColumnOrName", **kwargs) -> None:  # noqa: D107
+class Grouping:
+    def __init__(self, *cols: "ColumnOrName", **kwargs):
         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:  # noqa: D102
+    def get_columns(self) -> str:
         columns = ",".join([str(x) for x in self._cols])
         return columns
 
-    def __str__(self) -> str:  # noqa: D105
+    def __str__(self):
         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) -> None:  # noqa: D107
+    def __init__(self, grouping: Grouping, df: DataFrame):
         self._grouping = grouping
         self._df = df
         self.session: SparkSession = df.session
 
-    def __repr__(self) -> str:  # noqa: D105
+    def __repr__(self) -> str:
         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|
@@ -117,7 +115,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:
@@ -141,12 +139,11 @@ 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|
@@ -159,7 +156,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)|
         +-----+--------+
@@ -169,7 +166,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)|
         +--------+-----------+
@@ -180,19 +177,18 @@ 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|
@@ -232,12 +228,11 @@ 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|
@@ -277,12 +272,11 @@ 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|
@@ -314,12 +308,14 @@ class GroupedData:
         """
 
     @overload
-    def agg(self, *exprs: Column) -> DataFrame: ...
+    def agg(self, *exprs: Column) -> DataFrame:
+        ...
 
     @overload
-    def agg(self, __exprs: dict[str, str]) -> DataFrame: ...  # noqa: PYI063
+    def agg(self, __exprs: Dict[str, str]) -> DataFrame:
+        ...
 
-    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:
@@ -351,18 +347,17 @@ 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|
@@ -398,9 +393,10 @@ 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)|
@@ -421,4 +417,4 @@ class GroupedData:
             rel = self._df.relation.select(*expressions, groups=group_by)
         return DataFrame(rel, self.session)
 
-    # TODO: add 'pivot'  # noqa: TD002, TD003
+    # TODO: add 'pivot'

duckdb/experimental/spark/sql/readwriter.py

@@ -1,9 +1,11 @@
-from typing import TYPE_CHECKING, Optional, Union, cast  # noqa: D100
+from typing import TYPE_CHECKING, List, Optional, Union, cast
 
-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]
 
@@ -12,19 +14,19 @@ if TYPE_CHECKING:
     from duckdb.experimental.spark.sql.session import SparkSession
 
 
-class DataFrameWriter:  # noqa: D101
-    def __init__(self, dataframe: "DataFrame") -> None:  # noqa: D107
+class DataFrameWriter:
+    def __init__(self, dataframe: "DataFrame"):
         self.dataframe = dataframe
 
-    def saveAsTable(self, table_name: str) -> None:  # noqa: D102
+    def saveAsTable(self, table_name: str) -> None:
         relation = self.dataframe.relation
         relation.create(table_name)
 
-    def parquet(  # noqa: D102
+    def parquet(
         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
@@ -35,7 +37,7 @@ class DataFrameWriter:  # noqa: D101
 
         relation.write_parquet(path, compression=compression)
 
-    def csv(  # noqa: D102
+    def csv(
         self,
         path: str,
         mode: Optional[str] = None,
@@ -55,7 +57,7 @@ class DataFrameWriter:  # noqa: D101
         encoding: Optional[str] = None,
         emptyValue: Optional[str] = None,
         lineSep: Optional[str] = None,
-    ) -> None:
+    ):
         if mode not in (None, "overwrite"):
             raise NotImplementedError
         if escapeQuotes:
@@ -86,13 +88,13 @@ class DataFrameWriter:  # noqa: D101
         )
 
 
-class DataFrameReader:  # noqa: D101
-    def __init__(self, session: "SparkSession") -> None:  # noqa: D107
+class DataFrameReader:
+    def __init__(self, session: "SparkSession"):
         self.session = session
 
-    def load(  # noqa: D102
+    def load(
         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,
@@ -100,7 +102,7 @@ class DataFrameReader:  # noqa: D101
         from duckdb.experimental.spark.sql.dataframe import DataFrame
 
         if not isinstance(path, str):
-            raise TypeError
+            raise ImportError
         if options:
             raise ContributionsAcceptedError
 
@@ -121,15 +123,15 @@ class DataFrameReader:  # noqa: D101
         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(  # noqa: D102
+    def csv(
         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,
@@ -223,7 +225,7 @@ class DataFrameReader:  # noqa: D101
         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(
@@ -245,15 +247,13 @@ class DataFrameReader:  # noqa: D101
             df = df.toDF(*names)
         return df
 
-    def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame":  # noqa: D102
+    def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame":
         input = list(paths)
         if len(input) != 1:
-            msg = "Only single paths are supported for now"
-            raise NotImplementedError(msg)
+            raise NotImplementedError("Only single paths are supported for now")
         option_amount = len(options.keys())
         if option_amount != 0:
-            msg = "Options are not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError("Options are not supported")
         path = input[0]
         rel = self.session.conn.read_parquet(path)
         from ..sql.dataframe import DataFrame
@@ -263,7 +263,7 @@ class DataFrameReader:  # noqa: D101
 
     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,7 +289,8 @@ class DataFrameReader:  # noqa: D101
         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``.
@@ -320,16 +321,16 @@ class DataFrameReader:  # noqa: D101
 
             .. # 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()
@@ -339,89 +340,102 @@ class DataFrameReader:  # noqa: D101
         |100|Hyukjin Kwon|
         +---+------------+
         """
+
         if schema is not None:
-            msg = "The 'schema' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError("The 'schema' option is not supported")
         if primitivesAsString is not None:
-            msg = "The 'primitivesAsString' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'primitivesAsString' option is not supported"
+            )
         if prefersDecimal is not None:
-            msg = "The 'prefersDecimal' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'prefersDecimal' option is not supported"
+            )
         if allowComments is not None:
-            msg = "The 'allowComments' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'allowComments' option is not supported"
+            )
         if allowUnquotedFieldNames is not None:
-            msg = "The 'allowUnquotedFieldNames' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'allowUnquotedFieldNames' option is not supported"
+            )
         if allowSingleQuotes is not None:
-            msg = "The 'allowSingleQuotes' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'allowSingleQuotes' option is not supported"
+            )
         if allowNumericLeadingZero is not None:
-            msg = "The 'allowNumericLeadingZero' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'allowNumericLeadingZero' option is not supported"
+            )
         if allowBackslashEscapingAnyCharacter is not None:
-            msg = "The 'allowBackslashEscapingAnyCharacter' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'allowBackslashEscapingAnyCharacter' option is not supported"
+            )
         if mode is not None:
-            msg = "The 'mode' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError("The 'mode' option is not supported")
         if columnNameOfCorruptRecord is not None:
-            msg = "The 'columnNameOfCorruptRecord' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'columnNameOfCorruptRecord' option is not supported"
+            )
         if dateFormat is not None:
-            msg = "The 'dateFormat' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError("The 'dateFormat' option is not supported")
         if timestampFormat is not None:
-            msg = "The 'timestampFormat' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'timestampFormat' option is not supported"
+            )
         if multiLine is not None:
-            msg = "The 'multiLine' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError("The 'multiLine' option is not supported")
         if allowUnquotedControlChars is not None:
-            msg = "The 'allowUnquotedControlChars' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'allowUnquotedControlChars' option is not supported"
+            )
         if lineSep is not None:
-            msg = "The 'lineSep' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError("The 'lineSep' option is not supported")
         if samplingRatio is not None:
-            msg = "The 'samplingRatio' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'samplingRatio' option is not supported"
+            )
         if dropFieldIfAllNull is not None:
-            msg = "The 'dropFieldIfAllNull' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'dropFieldIfAllNull' option is not supported"
+            )
         if encoding is not None:
-            msg = "The 'encoding' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError("The 'encoding' option is not supported")
         if locale is not None:
-            msg = "The 'locale' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError("The 'locale' option is not supported")
         if pathGlobFilter is not None:
-            msg = "The 'pathGlobFilter' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'pathGlobFilter' option is not supported"
+            )
         if recursiveFileLookup is not None:
-            msg = "The 'recursiveFileLookup' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'recursiveFileLookup' option is not supported"
+            )
         if modifiedBefore is not None:
-            msg = "The 'modifiedBefore' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'modifiedBefore' option is not supported"
+            )
         if modifiedAfter is not None:
-            msg = "The 'modifiedAfter' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'modifiedAfter' option is not supported"
+            )
         if allowNonNumericNumbers is not None:
-            msg = "The 'allowNonNumericNumbers' option is not supported"
-            raise ContributionsAcceptedError(msg)
+            raise ContributionsAcceptedError(
+                "The 'allowNonNumericNumbers' option is not supported"
+            )
 
         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",
@@ -432,4 +446,4 @@ class DataFrameReader:  # noqa: D101
             )
 
 
-__all__ = ["DataFrameReader", "DataFrameWriter"]
+__all__ = ["DataFrameWriter", "DataFrameReader"]