Make the experimental Spark API work as advertised against PySpark 4.2.0 and 3.5.9

duckdb/experimental/spark/_globals.py

@@ -33,7 +33,7 @@ def foo(arg=pyducdkb.spark._NoValue):
 
 __ALL__ = ["_NoValue"]
 
-from typing_extensions import Self
+from typing import Self
 
 # Disallow reloading this module so as to preserve the identities of the
 # classes defined here.

duckdb/experimental/spark/_typing.py

@@ -17,11 +17,7 @@
 # under the License.
 
 from collections.abc import Callable, Iterable, Sized
-from typing import Literal, TypeVar
-
-from typing_extensions import Protocol, Self
-
-from numpy import float32, float64, int32, int64, ndarray
+from typing import Literal, Protocol, Self, TypeVar
 
 F = TypeVar("F", bound=Callable)
 T_co = TypeVar("T_co", covariant=True)
@@ -43,5 +39,3 @@ class SizedIterable(Protocol, Sized, Iterable[T_co]): ...
 
 
 S = TypeVar("S", bound=SupportsOrdering)
-
-NumberOrArray = TypeVar("NumberOrArray", float, int, complex, int32, int64, float32, float64, ndarray)

duckdb/experimental/spark/errors/error_classes.py

@@ -18,6 +18,12 @@
 
 ERROR_CLASSES_JSON = """
 {
+  "AMBIGUOUS_REFERENCE" : {
+    "message" : [
+      "Reference <name> is ambiguous, could be: <referenceNames>."
+    ],
+    "sqlState" : "42704"
+  },
   "APPLICATION_NAME_NOT_SET" : {
     "message" : [
       "An application name must be set in your configuration."
@@ -106,22 +112,26 @@
   "CANNOT_INFER_EMPTY_SCHEMA": {
     "message": [
       "Can not infer schema from empty dataset."
-    ]
+    ],
+    "sqlState": "42K09"
   },
   "CANNOT_INFER_SCHEMA_FOR_TYPE": {
     "message": [
       "Can not infer schema for type: `<data_type>`."
-    ]
+    ],
+    "sqlState": "42K09"
   },
   "CANNOT_INFER_TYPE_FOR_FIELD": {
     "message": [
       "Unable to infer the type of the field `<field_name>`."
-    ]
+    ],
+    "sqlState": "42K09"
   },
   "CANNOT_MERGE_TYPE": {
     "message": [
       "Can not merge type `<data_type1>` and `<data_type2>`."
-    ]
+    ],
+    "sqlState": "42K09"
   },
   "CANNOT_OPEN_SOCKET": {
     "message": [
@@ -148,6 +158,12 @@
       "returnType can not be specified when `<arg_name>` is a user-defined function, but got <return_type>."
     ]
   },
+  "COLUMN_ALREADY_EXISTS" : {
+    "message" : [
+      "The column <columnName> already exists. Choose another name or rename the existing column."
+    ],
+    "sqlState" : "42711"
+  },
   "COLUMN_IN_LIST": {
     "message": [
       "`<func_name>` does not allow a Column in a list."
@@ -230,11 +246,23 @@
       "Duplicated field names in Arrow Struct are not allowed, got <field_names>"
     ]
   },
+  "FIELD_NOT_FOUND" : {
+    "message" : [
+      "No such struct field <fieldName> in <fields>."
+    ],
+    "sqlState" : "42704"
+  },
   "HIGHER_ORDER_FUNCTION_SHOULD_RETURN_COLUMN" : {
     "message" : [
       "Function `<func_name>` should return Column, got <return_type>."
     ]
   },
+  "INCOMPATIBLE_COLUMN_TYPE" : {
+    "message" : [
+      "<operator> can only be performed on tables with compatible column types. The <columnOrdinalNumber> column of the <tableOrdinalNumber> table is <dataType1> type which is not compatible with <dataType2> at the same column of the first table.<hint>."
+    ],
+    "sqlState" : "42825"
+  },
   "INCORRECT_CONF_FOR_PROFILE" : {
     "message" : [
       "`spark.python.profile` or `spark.python.profile.memory` configuration",
@@ -286,6 +314,12 @@
       "Pandas UDF should return StructType for <eval_type>, got <return_type>."
     ]
   },
+  "INVALID_SAVE_MODE" : {
+    "message" : [
+      "The specified save mode <mode> is invalid. Valid save modes include \\"append\\", \\"overwrite\\", \\"ignore\\", \\"error\\", \\"errorifexists\\", and \\"default\\"."
+    ],
+    "sqlState" : "42000"
+  },
   "INVALID_TIMEOUT_TIMESTAMP" : {
     "message" : [
       "Timeout timestamp (<timestamp>) cannot be earlier than the current watermark (<watermark>)."
@@ -371,6 +405,12 @@
       "A master URL must be set in your configuration."
     ]
   },
+  "MISSING_GROUP_BY" : {
+    "message" : [
+      "The query does not include a GROUP BY clause. Add GROUP BY or turn it into the window functions using OVER clauses."
+    ],
+    "sqlState" : "42803"
+  },
   "MISSING_LIBRARY_FOR_PROFILER" : {
     "message" : [
       "Install the 'memory_profiler' library in the cluster to enable memory profiling."
@@ -626,6 +666,12 @@
       "No active Spark session found. Please create a new Spark session before running the code."
     ]
   },
+  "NUM_COLUMNS_MISMATCH" : {
+    "message" : [
+      "<operator> can only be performed on inputs with the same number of columns, but the first input has <firstNumColumns> columns and the <invalidOrdinalNum> input has <invalidNumColumns> columns."
+    ],
+    "sqlState" : "42826"
+  },
   "ONLY_ALLOWED_FOR_SINGLE_COLUMN" : {
     "message" : [
       "Argument `<arg_name>` can only be provided for a single column."
@@ -636,6 +682,18 @@
       "Only a single trigger is allowed."
     ]
   },
+  "PARSE_SYNTAX_ERROR" : {
+    "message" : [
+      "Syntax error at or near <error><hint>."
+    ],
+    "sqlState" : "42601"
+  },
+  "PATH_ALREADY_EXISTS" : {
+    "message" : [
+      "Path <outputPath> already exists. Set mode as \\"overwrite\\" to overwrite the existing path."
+    ],
+    "sqlState" : "42K04"
+  },
   "PIPE_FUNCTION_EXITED" : {
     "message" : [
       "Pipe function `<func_name>` exited with error code <error_code>."
@@ -797,11 +855,41 @@
       "Unknown response: <response>."
     ]
   },
+  "UNRESOLVED_COLUMN" : {
+    "message" : [
+      "A column, variable, or function parameter with name <objectName> cannot be resolved."
+    ],
+    "sub_class" : {
+      "WITHOUT_SUGGESTION" : {
+        "message" : [
+          ""
+        ]
+      },
+      "WITH_SUGGESTION" : {
+        "message" : [
+          "Did you mean one of the following? [<proposal>]."
+        ]
+      }
+    },
+    "sqlState" : "42703"
+  },
+  "UNRESOLVED_COLUMN_AMONG_FIELD_NAMES" : {
+    "message" : [
+      "Cannot resolve column name \\"<colName>\\" among (<fieldNames>)."
+    ],
+    "sqlState" : "42703"
+  },
   "UNSUPPORTED_DATA_TYPE" : {
     "message" : [
       "Unsupported DataType `<data_type>`."
     ]
   },
+  "UNSUPPORTED_DATATYPE" : {
+    "message" : [
+      "Unsupported data type <typeName>."
+    ],
+    "sqlState" : "0A000"
+  },
   "UNSUPPORTED_DATA_TYPE_FOR_ARROW" : {
     "message" : [
       "Single data type <data_type> is not supported with Arrow."
@@ -817,6 +905,25 @@
       "<data_type> is only supported with pyarrow 2.0.0 and above."
     ]
   },
+  "UNSUPPORTED_DATA_TYPE_FOR_DATASOURCE" : {
+    "message" : [
+      "The <format> datasource doesn't support the column <columnName> of the type <columnType>."
+    ],
+    "sqlState" : "0A000"
+  },
+  "UNSUPPORTED_FEATURE" : {
+    "message" : [
+      "The feature is not supported:"
+    ],
+    "sub_class" : {
+      "SET_OPERATION_ON_MAP_TYPE" : {
+        "message" : [
+          "Cannot have MAP type columns in DataFrame which calls set operations (INTERSECT, EXCEPT, etc.), but the type of column <colName> is <dataType>."
+        ]
+      }
+    },
+    "sqlState" : "0A000"
+  },
   "UNSUPPORTED_JOIN_TYPE" : {
     "message" : [
       "Unsupported join type: <join_type>. Supported join types include: \\"inner\\", \\"outer\\", \\"full\\", \\"fullouter\\", \\"full_outer\\", \\"leftouter\\", \\"left\\", \\"left_outer\\", \\"rightouter\\", \\"right\\", \\"right_outer\\", \\"leftsemi\\", \\"left_semi\\", \\"semi\\", \\"leftanti\\", \\"left_anti\\", \\"anti\\", \\"cross\\"."
@@ -911,6 +1018,11 @@
     "message": [
       "Index must be non-zero."
     ]
+  },
+  "_LEGACY_ERROR_TEMP_1323" : {
+    "message" : [
+      "\\"<colName>\\" is not a numeric column. Aggregation function can only be applied on a numeric column."
+    ]
   }
 }
 """

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

@@ -31,6 +31,18 @@ def __init__(  # noqa: D107
         self.error_class = error_class
         self.message_parameters = message_parameters
 
+    def getCondition(self) -> str | None:
+        """Returns an error condition, the name Spark 4 uses for an error class.
+
+        .. versionadded:: 4.0.0
+
+        See Also:
+        --------
+        :meth:`PySparkException.getMessageParameters`
+        :meth:`PySparkException.getSqlState`
+        """
+        return self.error_class
+
     def getErrorClass(self) -> str | None:
         """Returns an error class as a string.
 
@@ -55,10 +67,8 @@ def getMessageParameters(self) -> dict[str, str] | None:
         """
         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.
+    def getSqlState(self) -> str | None:
+        """Returns an SQLSTATE as a string, or None when the error class has none.
 
         .. versionadded:: 3.4.0
 
@@ -67,7 +77,7 @@ def getSqlState(self) -> None:
         :meth:`PySparkException.getErrorClass`
         :meth:`PySparkException.getMessageParameters`
         """
-        return None
+        return self.error_reader.get_sqlstate(self.error_class)
 
     def __str__(self) -> str:  # noqa: D105
         if self.getErrorClass() is not None:

duckdb/experimental/spark/errors/utils.py

@@ -38,6 +38,12 @@ def get_error_message(self, error_class: str, message_parameters: dict[str, str]
 
         return message_template.translate(table).format(**message_parameters)
 
+    def get_sqlstate(self, error_class: str | None) -> str | None:
+        """Returns the SQLSTATE of an error class, or None when it has none."""
+        if error_class is None:
+            return None
+        return self.error_info_map.get(error_class.split(".")[0], {}).get("sqlState")
+
     def get_message_template(self, error_class: str) -> str:
         """Returns the message template for corresponding error class from error_classes.py.
 

duckdb/experimental/spark/sql/__init__.py

@@ -3,5 +3,6 @@
 from .dataframe import DataFrame
 from .readwriter import DataFrameWriter
 from .session import SparkSession
+from .types import Row
 
-__all__ = ["Catalog", "DataFrame", "DataFrameWriter", "RuntimeConfig", "SparkSession"]
+__all__ = ["Catalog", "DataFrame", "DataFrameWriter", "Row", "RuntimeConfig", "SparkSession"]

duckdb/experimental/spark/sql/_typing.py

@@ -16,21 +16,10 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from collections.abc import Callable
-from typing import (
-    Any,
-    TypeVar,
-)
-
-try:
-    from typing import Literal, Protocol
-except ImportError:
-    from typing import Literal
-
-    from typing_extensions import Protocol
-
 import datetime
 import decimal
+from collections.abc import Callable
+from typing import Any, Literal, Protocol, TypeVar
 
 from .._typing import PrimitiveType
 from . import types