duckdb 1.5.0.dev56__cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl

duckdb/experimental/spark/errors/exceptions/__init__.py

@@ -0,0 +1,16 @@
+#  # noqa: D104
+# 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.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#

duckdb/experimental/spark/errors/exceptions/base.py

@@ -0,0 +1,168 @@
+from typing import Optional, cast  # noqa: D100
+
+from ..utils import ErrorClassesReader
+
+
+class PySparkException(Exception):
+    """Base Exception for handling errors generated from PySpark."""
+
+    def __init__(  # noqa: D107
+        self,
+        message: Optional[str] = None,
+        # The error class, decides the message format, must be one of the valid options listed in 'error_classes.py'
+        error_class: Optional[str] = None,
+        # The dictionary listing the arguments specified in the message (or the error_class)
+        message_parameters: Optional[dict[str, str]] = None,
+    ) -> None:
+        # `message` vs `error_class` & `message_parameters` are mutually exclusive.
+        assert (message is not None and (error_class is None and message_parameters is None)) or (
+            message is None and (error_class is not None and message_parameters is not None)
+        )
+
+        self.error_reader = ErrorClassesReader()
+
+        if message is None:
+            self.message = self.error_reader.get_error_message(
+                cast("str", error_class), cast("dict[str, str]", message_parameters)
+            )
+        else:
+            self.message = message
+
+        self.error_class = error_class
+        self.message_parameters = message_parameters
+
+    def getErrorClass(self) -> Optional[str]:
+        """Returns an error class as a string.
+
+        .. versionadded:: 3.4.0
+
+        See Also:
+        --------
+        :meth:`PySparkException.getMessageParameters`
+        :meth:`PySparkException.getSqlState`
+        """
+        return self.error_class
+
+    def getMessageParameters(self) -> Optional[dict[str, str]]:
+        """Returns a message parameters as a dictionary.
+
+        .. versionadded:: 3.4.0
+
+        See Also:
+        --------
+        :meth:`PySparkException.getErrorClass`
+        :meth:`PySparkException.getSqlState`
+        """
+        return self.message_parameters
+
+    def getSqlState(self) -> None:
+        """Returns an SQLSTATE as a string.
+
+        Errors generated in Python have no SQLSTATE, so it always returns None.
+
+        .. versionadded:: 3.4.0
+
+        See Also:
+        --------
+        :meth:`PySparkException.getErrorClass`
+        :meth:`PySparkException.getMessageParameters`
+        """
+        return None
+
+    def __str__(self) -> str:  # noqa: D105
+        if self.getErrorClass() is not None:
+            return f"[{self.getErrorClass()}] {self.message}"
+        else:
+            return self.message
+
+
+class AnalysisException(PySparkException):
+    """Failed to analyze a SQL query plan."""
+
+
+class SessionNotSameException(PySparkException):
+    """Performed the same operation on different SparkSession."""
+
+
+class TempTableAlreadyExistsException(AnalysisException):
+    """Failed to create temp view since it is already exists."""
+
+
+class ParseException(AnalysisException):
+    """Failed to parse a SQL command."""
+
+
+class IllegalArgumentException(PySparkException):
+    """Passed an illegal or inappropriate argument."""
+
+
+class ArithmeticException(PySparkException):
+    """Arithmetic exception thrown from Spark with an error class."""
+
+
+class UnsupportedOperationException(PySparkException):
+    """Unsupported operation exception thrown from Spark with an error class."""
+
+
+class ArrayIndexOutOfBoundsException(PySparkException):
+    """Array index out of bounds exception thrown from Spark with an error class."""
+
+
+class DateTimeException(PySparkException):
+    """Datetime exception thrown from Spark with an error class."""
+
+
+class NumberFormatException(IllegalArgumentException):
+    """Number format exception thrown from Spark with an error class."""
+
+
+class StreamingQueryException(PySparkException):
+    """Exception that stopped a :class:`StreamingQuery`."""
+
+
+class QueryExecutionException(PySparkException):
+    """Failed to execute a query."""
+
+
+class PythonException(PySparkException):
+    """Exceptions thrown from Python workers."""
+
+
+class SparkRuntimeException(PySparkException):
+    """Runtime exception thrown from Spark with an error class."""
+
+
+class SparkUpgradeException(PySparkException):
+    """Exception thrown because of Spark upgrade."""
+
+
+class UnknownException(PySparkException):
+    """None of the above exceptions."""
+
+
+class PySparkValueError(PySparkException, ValueError):
+    """Wrapper class for ValueError to support error classes."""
+
+
+class PySparkIndexError(PySparkException, IndexError):
+    """Wrapper class for IndexError to support error classes."""
+
+
+class PySparkTypeError(PySparkException, TypeError):
+    """Wrapper class for TypeError to support error classes."""
+
+
+class PySparkAttributeError(PySparkException, AttributeError):
+    """Wrapper class for AttributeError to support error classes."""
+
+
+class PySparkRuntimeError(PySparkException, RuntimeError):
+    """Wrapper class for RuntimeError to support error classes."""
+
+
+class PySparkAssertionError(PySparkException, AssertionError):
+    """Wrapper class for AssertionError to support error classes."""
+
+
+class PySparkNotImplementedError(PySparkException, NotImplementedError):
+    """Wrapper class for NotImplementedError to support error classes."""

duckdb/experimental/spark/errors/utils.py

@@ -0,0 +1,111 @@
+#  # 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.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import re
+
+from .error_classes import ERROR_CLASSES_MAP
+
+
+class ErrorClassesReader:
+    """A reader to load error information from error_classes.py."""
+
+    def __init__(self) -> None:  # noqa: D107
+        self.error_info_map = ERROR_CLASSES_MAP
+
+    def get_error_message(self, error_class: str, message_parameters: dict[str, str]) -> str:
+        """Returns the completed error message by applying message parameters to the message template."""
+        message_template = self.get_message_template(error_class)
+        # Verify message parameters.
+        message_parameters_from_template = re.findall("<([a-zA-Z0-9_-]+)>", message_template)
+        assert set(message_parameters_from_template) == set(message_parameters), (
+            f"Undefined error message parameter for error class: {error_class}. Parameters: {message_parameters}"
+        )
+        table = str.maketrans("<>", "{}")
+
+        return message_template.translate(table).format(**message_parameters)
+
+    def get_message_template(self, error_class: str) -> str:
+        """Returns the message template for corresponding error class from error_classes.py.
+
+        For example,
+        when given `error_class` is "EXAMPLE_ERROR_CLASS",
+        and corresponding error class in error_classes.py looks like the below:
+
+        .. code-block:: python
+
+            "EXAMPLE_ERROR_CLASS" : {
+              "message" : [
+                "Problem <A> because of <B>."
+              ]
+            }
+
+        In this case, this function returns:
+        "Problem <A> because of <B>."
+
+        For sub error class, when given `error_class` is "EXAMPLE_ERROR_CLASS.SUB_ERROR_CLASS",
+        and corresponding error class in error_classes.py looks like the below:
+
+        .. code-block:: python
+
+            "EXAMPLE_ERROR_CLASS" : {
+              "message" : [
+                "Problem <A> because of <B>."
+              ],
+              "sub_class" : {
+                "SUB_ERROR_CLASS" : {
+                  "message" : [
+                    "Do <C> to fix the problem."
+                  ]
+                }
+              }
+            }
+
+        In this case, this function returns:
+        "Problem <A> because <B>. Do <C> to fix the problem."
+        """
+        error_classes = error_class.split(".")
+        len_error_classes = len(error_classes)
+        assert len_error_classes in (1, 2)
+
+        # Generate message template for main error class.
+        main_error_class = error_classes[0]
+        if main_error_class in self.error_info_map:
+            main_error_class_info_map = self.error_info_map[main_error_class]
+        else:
+            msg = f"Cannot find main error class '{main_error_class}'"
+            raise ValueError(msg)
+
+        main_message_template = "\n".join(main_error_class_info_map["message"])
+
+        has_sub_class = len_error_classes == 2
+
+        if not has_sub_class:
+            message_template = main_message_template
+        else:
+            # Generate message template for sub error class if exists.
+            sub_error_class = error_classes[1]
+            main_error_class_subclass_info_map = main_error_class_info_map["sub_class"]
+            if sub_error_class in main_error_class_subclass_info_map:
+                sub_error_class_info_map = main_error_class_subclass_info_map[sub_error_class]
+            else:
+                msg = f"Cannot find sub error class '{sub_error_class}'"
+                raise ValueError(msg)
+
+            sub_message_template = "\n".join(sub_error_class_info_map["message"])
+            message_template = main_message_template + " " + sub_message_template
+
+        return message_template

duckdb/experimental/spark/exception.py

@@ -0,0 +1,18 @@
+# ruff: noqa: D100
+from typing import Optional
+
+
+class ContributionsAcceptedError(NotImplementedError):
+    """This method is not planned to be implemented, if you would like to implement this method
+    or show your interest in this method to other members of the community,
+    feel free to open up a PR or a Discussion over on https://github.com/duckdb/duckdb.
+    """  # noqa: D205
+
+    def __init__(self, message: Optional[str] = None) -> None:  # noqa: D107
+        doc = self.__class__.__doc__
+        if message:
+            doc = message + "\n" + doc
+        super().__init__(doc)
+
+
+__all__ = ["ContributionsAcceptedError"]

duckdb/experimental/spark/sql/__init__.py

@@ -0,0 +1,7 @@
+from .catalog import Catalog  # noqa: D104
+from .conf import RuntimeConfig
+from .dataframe import DataFrame
+from .readwriter import DataFrameWriter
+from .session import SparkSession
+
+__all__ = ["Catalog", "DataFrame", "DataFrameWriter", "RuntimeConfig", "SparkSession"]

duckdb/experimental/spark/sql/_typing.py

@@ -0,0 +1,86 @@
+#
+# 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.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import (
+    Any,
+    Callable,
+    Optional,
+    TypeVar,
+    Union,
+)
+
+try:
+    from typing import Literal, Protocol
+except ImportError:
+    from typing_extensions import Literal, Protocol
+
+import datetime
+import decimal
+
+from .._typing import PrimitiveType
+from . import types
+from .column import Column
+
+ColumnOrName = Union[Column, str]
+ColumnOrName_ = TypeVar("ColumnOrName_", bound=ColumnOrName)
+DecimalLiteral = decimal.Decimal
+DateTimeLiteral = Union[datetime.datetime, datetime.date]
+LiteralType = PrimitiveType
+AtomicDataTypeOrString = Union[types.AtomicType, str]
+DataTypeOrString = Union[types.DataType, str]
+OptionalPrimitiveType = Optional[PrimitiveType]
+
+AtomicValue = TypeVar(
+    "AtomicValue",
+    datetime.datetime,
+    datetime.date,
+    decimal.Decimal,
+    bool,
+    str,
+    int,
+    float,
+)
+
+RowLike = TypeVar("RowLike", list[Any], tuple[Any, ...], types.Row)
+
+SQLBatchedUDFType = Literal[100]
+
+
+class SupportsOpen(Protocol):
+    def open(self, partition_id: int, epoch_id: int) -> bool: ...
+
+
+class SupportsProcess(Protocol):
+    def process(self, row: types.Row) -> None: ...
+
+
+class SupportsClose(Protocol):
+    def close(self, error: Exception) -> None: ...
+
+
+class UserDefinedFunctionLike(Protocol):
+    func: Callable[..., Any]
+    evalType: int
+    deterministic: bool
+
+    @property
+    def returnType(self) -> types.DataType: ...
+
+    def __call__(self, *args: ColumnOrName) -> Column: ...
+
+    def asNondeterministic(self) -> "UserDefinedFunctionLike": ...

duckdb/experimental/spark/sql/catalog.py

@@ -0,0 +1,79 @@
+from typing import NamedTuple, Optional, Union  # noqa: D100
+
+from .session import SparkSession
+
+
+class Database(NamedTuple):  # noqa: D101
+    name: str
+    description: Optional[str]
+    locationUri: str
+
+
+class Table(NamedTuple):  # noqa: D101
+    name: str
+    database: Optional[str]
+    description: Optional[str]
+    tableType: str
+    isTemporary: bool
+
+
+class Column(NamedTuple):  # noqa: D101
+    name: str
+    description: Optional[str]
+    dataType: str
+    nullable: bool
+    isPartition: bool
+    isBucket: bool
+
+
+class Function(NamedTuple):  # noqa: D101
+    name: str
+    description: Optional[str]
+    className: str
+    isTemporary: bool
+
+
+class Catalog:  # noqa: D101
+    def __init__(self, session: SparkSession) -> None:  # noqa: D107
+        self._session = session
+
+    def listDatabases(self) -> list[Database]:  # noqa: D102
+        res = self._session.conn.sql("select database_name from duckdb_databases()").fetchall()
+
+        def transform_to_database(x: list[str]) -> Database:
+            return Database(name=x[0], description=None, locationUri="")
+
+        databases = [transform_to_database(x) for x in res]
+        return databases
+
+    def listTables(self) -> list[Table]:  # noqa: D102
+        res = self._session.conn.sql("select table_name, database_name, sql, temporary from duckdb_tables()").fetchall()
+
+        def transform_to_table(x: list[str]) -> Table:
+            return Table(name=x[0], database=x[1], description=x[2], tableType="", isTemporary=x[3])
+
+        tables = [transform_to_table(x) for x in res]
+        return tables
+
+    def listColumns(self, tableName: str, dbName: Optional[str] = None) -> list[Column]:  # noqa: D102
+        query = f"""
+			select column_name, data_type, is_nullable from duckdb_columns() where table_name = '{tableName}'
+		"""
+        if dbName:
+            query += f" and database_name = '{dbName}'"
+        res = self._session.conn.sql(query).fetchall()
+
+        def transform_to_column(x: list[Union[str, bool]]) -> Column:
+            return Column(name=x[0], description=None, dataType=x[1], nullable=x[2], isPartition=False, isBucket=False)
+
+        columns = [transform_to_column(x) for x in res]
+        return columns
+
+    def listFunctions(self, dbName: Optional[str] = None) -> list[Function]:  # noqa: D102
+        raise NotImplementedError
+
+    def setCurrentDatabase(self, dbName: str) -> None:  # noqa: D102
+        raise NotImplementedError
+
+
+__all__ = ["Catalog", "Column", "Database", "Function", "Table"]