diff --git a/duckdb/experimental/spark/_globals.py b/duckdb/experimental/spark/_globals.py index 23a6f1713..3dd7232f1 100644 --- a/duckdb/experimental/spark/_globals.py +++ b/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. diff --git a/duckdb/experimental/spark/_typing.py b/duckdb/experimental/spark/_typing.py index 19663281d..a2aac22c7 100644 --- a/duckdb/experimental/spark/_typing.py +++ b/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) diff --git a/duckdb/experimental/spark/errors/error_classes.py b/duckdb/experimental/spark/errors/error_classes.py index c43a5f189..8b56e3ad3 100644 --- a/duckdb/experimental/spark/errors/error_classes.py +++ b/duckdb/experimental/spark/errors/error_classes.py @@ -18,6 +18,12 @@ ERROR_CLASSES_JSON = """ { + "AMBIGUOUS_REFERENCE" : { + "message" : [ + "Reference is ambiguous, could be: ." + ], + "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: ``." - ] + ], + "sqlState": "42K09" }, "CANNOT_INFER_TYPE_FOR_FIELD": { "message": [ "Unable to infer the type of the field ``." - ] + ], + "sqlState": "42K09" }, "CANNOT_MERGE_TYPE": { "message": [ "Can not merge type `` and ``." - ] + ], + "sqlState": "42K09" }, "CANNOT_OPEN_SOCKET": { "message": [ @@ -148,6 +158,12 @@ "returnType can not be specified when `` is a user-defined function, but got ." ] }, + "COLUMN_ALREADY_EXISTS" : { + "message" : [ + "The column already exists. Choose another name or rename the existing column." + ], + "sqlState" : "42711" + }, "COLUMN_IN_LIST": { "message": [ "`` does not allow a Column in a list." @@ -230,11 +246,23 @@ "Duplicated field names in Arrow Struct are not allowed, got " ] }, + "FIELD_NOT_FOUND" : { + "message" : [ + "No such struct field in ." + ], + "sqlState" : "42704" + }, "HIGHER_ORDER_FUNCTION_SHOULD_RETURN_COLUMN" : { "message" : [ "Function `` should return Column, got ." ] }, + "INCOMPATIBLE_COLUMN_TYPE" : { + "message" : [ + " can only be performed on tables with compatible column types. The column of the table is type which is not compatible with at the same column of the first table.." + ], + "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 , got ." ] }, + "INVALID_SAVE_MODE" : { + "message" : [ + "The specified save mode is invalid. Valid save modes include \\"append\\", \\"overwrite\\", \\"ignore\\", \\"error\\", \\"errorifexists\\", and \\"default\\"." + ], + "sqlState" : "42000" + }, "INVALID_TIMEOUT_TIMESTAMP" : { "message" : [ "Timeout timestamp () cannot be earlier than the current 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" : [ + " can only be performed on inputs with the same number of columns, but the first input has columns and the input has columns." + ], + "sqlState" : "42826" + }, "ONLY_ALLOWED_FOR_SINGLE_COLUMN" : { "message" : [ "Argument `` 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 ." + ], + "sqlState" : "42601" + }, + "PATH_ALREADY_EXISTS" : { + "message" : [ + "Path already exists. Set mode as \\"overwrite\\" to overwrite the existing path." + ], + "sqlState" : "42K04" + }, "PIPE_FUNCTION_EXITED" : { "message" : [ "Pipe function `` exited with error code ." @@ -797,11 +855,41 @@ "Unknown response: ." ] }, + "UNRESOLVED_COLUMN" : { + "message" : [ + "A column, variable, or function parameter with name cannot be resolved." + ], + "sub_class" : { + "WITHOUT_SUGGESTION" : { + "message" : [ + "" + ] + }, + "WITH_SUGGESTION" : { + "message" : [ + "Did you mean one of the following? []." + ] + } + }, + "sqlState" : "42703" + }, + "UNRESOLVED_COLUMN_AMONG_FIELD_NAMES" : { + "message" : [ + "Cannot resolve column name \\"\\" among ()." + ], + "sqlState" : "42703" + }, "UNSUPPORTED_DATA_TYPE" : { "message" : [ "Unsupported DataType ``." ] }, + "UNSUPPORTED_DATATYPE" : { + "message" : [ + "Unsupported data type ." + ], + "sqlState" : "0A000" + }, "UNSUPPORTED_DATA_TYPE_FOR_ARROW" : { "message" : [ "Single data type is not supported with Arrow." @@ -817,6 +905,25 @@ " is only supported with pyarrow 2.0.0 and above." ] }, + "UNSUPPORTED_DATA_TYPE_FOR_DATASOURCE" : { + "message" : [ + "The datasource doesn't support the column of the type ." + ], + "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 is ." + ] + } + }, + "sqlState" : "0A000" + }, "UNSUPPORTED_JOIN_TYPE" : { "message" : [ "Unsupported 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" : [ + "\\"\\" is not a numeric column. Aggregation function can only be applied on a numeric column." + ] } } """ diff --git a/duckdb/experimental/spark/errors/exceptions/base.py b/duckdb/experimental/spark/errors/exceptions/base.py index 9a60512f1..cba62a59a 100644 --- a/duckdb/experimental/spark/errors/exceptions/base.py +++ b/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: diff --git a/duckdb/experimental/spark/errors/utils.py b/duckdb/experimental/spark/errors/utils.py index f2962fc81..68f9c8ed5 100644 --- a/duckdb/experimental/spark/errors/utils.py +++ b/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. diff --git a/duckdb/experimental/spark/sql/__init__.py b/duckdb/experimental/spark/sql/__init__.py index 418273f0a..9130e13b4 100644 --- a/duckdb/experimental/spark/sql/__init__.py +++ b/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"] diff --git a/duckdb/experimental/spark/sql/_typing.py b/duckdb/experimental/spark/sql/_typing.py index cf0b15e1a..d60e14521 100644 --- a/duckdb/experimental/spark/sql/_typing.py +++ b/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 diff --git a/duckdb/experimental/spark/sql/column.py b/duckdb/experimental/spark/sql/column.py index e013a56de..fd66ee5f7 100644 --- a/duckdb/experimental/spark/sql/column.py +++ b/duckdb/experimental/spark/sql/column.py @@ -1,44 +1,557 @@ +import contextlib +import copy +import datetime +import decimal +import inspect +import itertools +import math +import operator +import re +import string +import struct from collections.abc import ( Callable, + Collection, Iterable, + Iterator, ) from typing import TYPE_CHECKING, Any, Union, cast +from ..errors import AnalysisException, ParseException, PySparkTypeError, PySparkValueError, SparkRuntimeException from ..exception import ContributionsAcceptedError -from .types import DataType +from .types import DataType, TimestampType, _instant_to_local_time, _parse_datatype_string if TYPE_CHECKING: from ._typing import DateTimeLiteral, DecimalLiteral, LiteralType -from duckdb import ColumnExpression, ConstantExpression, Expression, FunctionExpression -from duckdb.sqltypes import DuckDBPyType +import duckdb +from duckdb import ( + CaseExpression, + ColumnExpression, + ConstantExpression, + Expression, + FunctionExpression, + LambdaExpression, + SQLExpression, +) +from duckdb.sqltypes import BLOB, DOUBLE, TIMESTAMP_TZ, VARCHAR, DuckDBPyType __all__ = ["Column"] +_INTEGRAL_TYPE_IDS = frozenset( + {"tinyint", "smallint", "integer", "bigint", "hugeint", "utinyint", "usmallint", "uinteger", "ubigint", "uhugeint"} +) + + +# The Spark types PySpark gives NumPy values in a literal, by dtype name. +_NUMPY_LITERAL_TYPES = { + "bool": "BOOLEAN", + "int8": "TINYINT", + "int16": "SMALLINT", + "int32": "INTEGER", + "int64": "BIGINT", + "float32": "FLOAT", + "float64": "DOUBLE", +} + + +class _Ref: + """A column reference inside a Spark column name, which leaves out a qualifier; see _unqualified.""" + + __slots__ = ("parts",) + + def __init__(self, parts: tuple[str, ...]) -> None: + self.parts = parts + + def render(self, columns: Collection[str] | None, references: Collection[str] = ()) -> str: + return ".".join(self.parts if columns is None else _unqualified(self.parts, columns, references)) + + +def _unqualified(parts: tuple[str, ...], columns: Collection[str], references: Collection[str]) -> tuple[str, ...]: + """The parts of a column reference without a leading part that names a relation, as Spark names the column. + + ``columns`` holds the DataFrame's column names and ``references`` the renderings of the qualified references to + them, both folded to lower case. As in Spark, a leading part is a relation where it is no column, or where it + and the next part are a qualified reference, which Spark resolves before a struct field of a column. + """ + if len(parts) > 1 and ( + parts[0].casefold() not in columns or str(ColumnExpression(*parts[:2])).casefold() in references + ): + return parts[1:] + return parts + + +Name = tuple[str | _Ref, ...] + + +def _invalid_attribute_name(name: str) -> AnalysisException: + return AnalysisException( + f"[INVALID_ATTRIBUTE_NAME_SYNTAX] Syntax error in the attribute name: {name}. Check that backticks appear in " + "pairs, a quoted string is a complete name part and use a backtick only inside quoted name parts." + ) + + +def _name_parts(name: str) -> tuple[str, ...]: + """Split a Spark attribute name such as ``a.b`` or ```a.b`.c`` into its parts, with Spark's syntax rules. + + A quoted part must be a whole part, a backtick in it is doubled, and no part of an unquoted name is empty. + """ + parts: list[str] = [] + part: list[str] = [] + quoted = False + i = 0 + while i < len(name): + c = name[i] + if quoted: + if c != "`": + part.append(c) + elif name[i + 1 : i + 2] == "`": + part.append("`") + i += 1 + elif name[i + 1 : i + 2] not in ("", "."): + raise _invalid_attribute_name(name) + else: + quoted = False + elif c == "`": + if part: + raise _invalid_attribute_name(name) + quoted = True + elif c == ".": + if i == 0 or name[i - 1] == "." or i == len(name) - 1: + raise _invalid_attribute_name(name) + parts.append("".join(part)) + part = [] + else: + part.append(c) + i += 1 + if quoted: + raise _invalid_attribute_name(name) + parts.append("".join(part)) + return tuple(parts) + + +def _ref_name(name: str) -> Name | None: + return None if name == "*" else (_Ref(_name_parts(name)),) + + +def _bare_column(name: str) -> Expression: + """A reference to the column ``name`` without a qualifier, whatever characters the name holds.""" + # A single argument is read as a DuckDB qualified name, so quoting keeps a dot or a double quote in the name. + return ColumnExpression('"' + name.replace('"', '""') + '"') + + +def _struct_extract(expr: Expression, name: str) -> Expression: + return FunctionExpression("struct_extract", expr, ConstantExpression(name)) + + +def _qualified_star(parts: tuple[str, ...]) -> Expression: + """The star of the relation or struct column that ``parts``, the parts of a name before ``.*``, name.""" + return SQLExpression(".".join('"' + part.replace('"', '""') + '"' for part in parts) + ".*") + + +def _column_reference(name: str) -> Expression: + """The column, qualified column or struct field a Spark attribute name such as ```a.b`.c`` refers to. + + A name such as ``q.*`` binds only where ``select`` expands it. + """ + if name == "*": + return ColumnExpression(name) + parts = _name_parts(name) + return ColumnExpression(*parts) if len(parts) > 1 else _bare_column(parts[0]) + + +def _shortest_float32(value: float) -> str: + for precision in range(9): + text = f"{value:.{precision}e}" + with contextlib.suppress(OverflowError): + if struct.unpack("f", struct.pack("f", float(text)))[0] == value: + return text + return repr(value) + + +def _java_float_string(value: float, single_precision: bool) -> str: + """Formats a float the way Java's Float.toString and Double.toString do.""" + if math.isnan(value): + return "NaN" + if math.isinf(value): + return "Infinity" if value > 0 else "-Infinity" + if value == 0: + return "-0.0" if math.copysign(1.0, value) < 0 else "0.0" + text = _shortest_float32(value) if single_precision else repr(value) + if len(decimal.Decimal(text).normalize().as_tuple().digits) == 1: + # Java prints at least two significant digits and picks the closest such decimal. + text = f"{value:.1e}" + sign, digit_tuple, exponent = decimal.Decimal(text).as_tuple() + scientific_exponent = len(digit_tuple) + cast("int", exponent) - 1 + digits = "".join(map(str, digit_tuple)).rstrip("0") + if -3 <= scientific_exponent < 7: + if scientific_exponent >= 0: + integer = digits[: scientific_exponent + 1].ljust(scientific_exponent + 1, "0") + fraction = digits[scientific_exponent + 1 :] + else: + integer = "0" + fraction = "0" * (-scientific_exponent - 1) + digits + body = f"{integer}.{fraction or '0'}" + else: + body = f"{digits[0]}.{digits[1:] or '0'}E{scientific_exponent}" + return f"-{body}" if sign else body + + +def _is_numpy(value: object) -> bool: + return type(value).__module__ == "numpy" + + +def _python_value(value: Any) -> Any: # noqa: ANN401 + """The Python value of a NumPy scalar, which is how Spark reads one outside ``lit``; other values unchanged.""" + return value.item() if _is_numpy(value) and getattr(value, "ndim", None) == 0 else value + + +def _literal_name(value: Any) -> str: # noqa: ANN401 + value = _python_value(value) + if value is None: + return "NULL" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, float): + return _java_float_string(value, single_precision=False) + if isinstance(value, decimal.Decimal): + return format(value, "f") + if isinstance(value, datetime.datetime): + # Spark names a time with a zone after the local time of its instant, as collect() returns it. + local = value if value.tzinfo is None else _instant_to_local_time(TimestampType().toInternal(value)) + text = local.isoformat(" ") + return f"TIMESTAMP '{text.rstrip('0') if '.' in text else text}'" + if isinstance(value, datetime.date): + return f"DATE '{value.isoformat()}'" + if isinstance(value, (bytes, bytearray)): + return f"X'{value.hex().upper()}'" + if isinstance(value, (list, tuple)): + return f"array({', '.join(_literal_name(v) for v in value)})" + return str(value) + + +def _sql_name(*pieces: Any) -> Name: # noqa: ANN401 + """Build a Spark column name from text, Columns, column references and literal values.""" + out: list[str | _Ref] = [] + for piece in pieces: + if isinstance(piece, Column): + # Unlike its rendering, an expression's own name leaves out the relation a column reference is bound to. + out.extend(piece._name if piece._name is not None else (piece.expr.get_name(),)) + elif isinstance(piece, (str, _Ref)): + out.append(piece) + else: + out.append(_literal_name(piece)) + return tuple(out) + + +def _format_name(template: str, *args: Any) -> Name: # noqa: ANN401 + """Build a Spark column name from a template such as ``"({0} + {1})"``.""" + pieces: list[Any] = [] + for text, field, _, _ in string.Formatter().parse(template): + pieces.append(text) + if field is not None: + pieces.append(args[int(field)]) + return _sql_name(*pieces) + + +def _comma_separated(values: Iterable[Any]) -> list[Any]: + pieces: list[Any] = [] + for i, value in enumerate(values): + if i: + pieces.append(", ") + pieces.append(value) + return pieces + + +def _call_name(function: str, *args: Any) -> Name: # noqa: ANN401 + return _sql_name(function, "(", *_comma_separated(args), ")") + + +def _render(name: Name, columns: Collection[str] | None = None, references: Collection[str] = ()) -> str: + return "".join(t if isinstance(t, str) else t.render(columns, references) for t in name) + + +_sql_type_names = { + "boolean": "BOOLEAN", + "tinyint": "TINYINT", + "smallint": "SMALLINT", + "integer": "INT", + "bigint": "BIGINT", + "float": "FLOAT", + "double": "DOUBLE", + "varchar": "STRING", + "blob": "BINARY", + "date": "DATE", + "timestamp": "TIMESTAMP_NTZ", + "timestamp with time zone": "TIMESTAMP", + "null": "VOID", +} + + +def _sql_type_name(dtype: DuckDBPyType) -> str: + """The Spark SQL name of the atomic type a DuckDB type stands for, as used in names of casts.""" + if dtype.id == "decimal": + (_, precision), (_, scale) = dtype.children + return f"DECIMAL({precision},{scale})" + return _sql_type_names.get(dtype.id, str(dtype).upper()) + + +def _quoted_identifier(name: str) -> str: + """``name`` as Spark 4 writes an identifier inside a column name: in backticks unless it is a plain word.""" + return name if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name) else "`" + name.replace("`", "``") + "`" + + +# The public methods of pyspark.sql.Column in PySpark 3.5.9 and 4.2.0, so that a missing one raises instead of +# becoming a struct field access. +_SPARK_COLUMN_METHODS = frozenset( + { + "alias", + "asc", + "asc_nulls_first", + "asc_nulls_last", + "astype", + "between", + "bitwiseAND", + "bitwiseOR", + "bitwiseXOR", + "cast", + "contains", + "desc", + "desc_nulls_first", + "desc_nulls_last", + "dropFields", + "endswith", + "eqNullSafe", + "getField", + "getItem", + "ilike", + "isNaN", + "isNotNull", + "isNull", + "isin", + "like", + "name", + "otherwise", + "outer", + "over", + "rlike", + "startswith", + "substr", + "transform", + "try_cast", + "when", + "withField", + } +) + + +class _MissingColumnMethodError(ContributionsAcceptedError, AttributeError): + """A method of Spark's Column that the module lacks; an AttributeError so ``hasattr`` keeps working.""" + + +def _constant(value: Any) -> Expression: # noqa: ANN401 + value = _python_value(value) + # ConstantExpression converts a float NaN to NULL, because Python value conversion treats NaN as missing data, + # also inside a list, so lists are built from their elements. + if isinstance(value, float) and math.isnan(value): + return ConstantExpression("NaN").cast(DOUBLE) + if isinstance(value, (list, tuple)): + return FunctionExpression("list_value", *(_constant(v) for v in value)) + if isinstance(value, dict): + msg = f"[UNSUPPORTED_FEATURE.LITERAL_TYPE] The feature is not supported: Literal for {value!r} of class dict." + raise SparkRuntimeException(message=msg) + if isinstance(value, datetime.datetime) and value.tzinfo is None: + return ConstantExpression(value).cast(TIMESTAMP_TZ) + return ConstantExpression(value) + + +def _numpy_literal(value: Any) -> tuple[Expression, Name, str | None]: # noqa: ANN401 + """The expression, Spark's name and alias of ``lit`` of a NumPy value, typed by its dtype as Spark does.""" + dtype = value.dtype.name + sql_type = "VARCHAR" if value.dtype.kind == "U" else _NUMPY_LITERAL_TYPES.get(dtype) + if value.ndim > 1: + msg = ( + f"[UNSUPPORTED_FEATURE.LITERAL_TYPE] The feature is not supported: Literal for {value!r} of class ndarray." + ) + raise SparkRuntimeException(message=msg) + if sql_type is None: + raise PySparkTypeError(error_class="UNSUPPORTED_NUMPY_ARRAY_SCALAR", message_parameters={"dtype": dtype}) + if value.ndim == 0: + duckdb_type = DuckDBPyType(sql_type) + expr = _constant(value).cast(duckdb_type) + if sql_type in ("BOOLEAN", "VARCHAR"): + return expr, (_literal_name(value),), None + # Spark casts the Python value of a NumPy number to its type and aliases the cast to the number's str(). + return expr, _sql_name("CAST(", value, " AS ", _sql_type_name(duckdb_type), ")"), str(value) + elements = value.tolist() + return _constant(elements).cast(DuckDBPyType(f"{sql_type}[]")), (_literal_name(elements),), None + + def _get_expr(x: Union["Column", str]) -> Expression: - return x.expr if isinstance(x, Column) else ConstantExpression(x) + return x.expr if isinstance(x, Column) else _constant(x) -def _func_op(name: str, doc: str = "") -> Callable[["Column"], "Column"]: - def _(self: "Column") -> "Column": - njc = getattr(self.expr, name)() - return Column(njc) +def _lambda(body: Callable[..., Expression], *captured: Expression) -> LambdaExpression: + """A lambda over ``body`` with parameter names that occur in none of ``captured``, the outer expressions it uses. - _.__doc__ = doc - return _ + So a parameter cannot shadow a column that ``captured`` refers to. + """ + taken = " ".join(map(str, captured)).casefold() + free = (name for name in (f"x{i}" for i in itertools.count()) if name not in taken) + names = tuple(itertools.islice(free, len(inspect.signature(body).parameters))) + return LambdaExpression(names[0] if len(names) == 1 else names, body(*map(ColumnExpression, names))) -def _unary_op( - name: str, - doc: str = "unary operator", -) -> Callable[["Column"], "Column"]: - """Create a method for given unary operator.""" +def _evaluated_once(expr: Expression, body: Callable[[Expression], Expression]) -> Expression: + """``body`` applied to ``expr``, which DuckDB evaluates once per row however often ``body`` uses it.""" + values = FunctionExpression("list_transform", FunctionExpression("list_value", expr), _lambda(body)) + return FunctionExpression("list_extract", values, ConstantExpression(1)) + + +def _truncate(value: Expression, target: DuckDBPyType) -> Expression: + """Cast to an integral type truncating toward zero like Spark, where DuckDB rounds to the nearest integer. + + A string must be digits with an optional fraction, as Spark without ANSI mode reads it, and otherwise raises as + Spark in ANSI mode does. + """ + source_type = FunctionExpression("typeof", value) + is_string = source_type == ConstantExpression("VARCHAR") + text = value.cast(VARCHAR) + is_number = FunctionExpression( + "regexp_full_match", text, ConstantExpression(r"\s*[+-]?([0-9]+\.?[0-9]*|\.[0-9]+)\s*") + ) + malformed = FunctionExpression( + "concat", + ConstantExpression("[CAST_INVALID_INPUT] The value '"), + text, + ConstantExpression( + f'\' of the type "STRING" cannot be cast to "{_sql_type_name(target)}" because it is malformed.' + ), + ) + as_double = value.cast(DOUBLE) + truncated = FunctionExpression("trunc", as_double) + # A decimal or a string can have a fraction that its nearest double drops, so its whole part is cut from its text; + # the leading 0 keeps a bare sign or an empty whole part a valid integer. + may_drop_fraction = ( + is_string & FunctionExpression("contains", text, ConstantExpression(".")) + ) | FunctionExpression("starts_with", source_type, ConstantExpression("DECIMAL")) + whole_part = FunctionExpression( + "regexp_replace", text, ConstantExpression(r"^\s*([+-]?)([0-9]*)(\.[0-9]*)?\s*$"), ConstantExpression(r"\10\2") + ) + return ( + CaseExpression(is_string & ~is_number, FunctionExpression("error", malformed)) + .when(as_double != truncated, truncated.cast(target)) + .when(may_drop_fraction, whole_part.cast(target)) + .otherwise(value.cast(target)) + ) + + +def _to_binary(value: Expression) -> Expression: + """Cast to binary like Spark, which takes a string's UTF-8 bytes where DuckDB reads escape sequences.""" + is_string = FunctionExpression("typeof", value) == ConstantExpression("VARCHAR") + return CaseExpression(is_string, FunctionExpression("encode", value.cast(VARCHAR))).otherwise(value.cast(BLOB)) + + +def _children(dtype: DuckDBPyType) -> list[DuckDBPyType]: + if dtype.id == "map": + return [dtype.key, dtype.value] + if dtype.id in ("list", "array"): + return [dtype.children[0][1]] + return [t for _, t in dtype.children] if dtype.id == "struct" else [] + + +def _cast_differs(dtype: DuckDBPyType) -> bool: + """Whether a DuckDB cast to ``dtype`` can differ from Spark's: an integral, binary or struct type, or one inside.""" + return dtype.id in _INTEGRAL_TYPE_IDS or dtype.id in ("blob", "struct") or any(map(_cast_differs, _children(dtype))) + + +def _convert(value: Expression, target: DuckDBPyType) -> Expression: + if not _cast_differs(target): + return value.cast(target) + if target.id in _INTEGRAL_TYPE_IDS: + return _truncate(value, target) + if target.id == "blob": + return _to_binary(value) + if target.id in ("list", "array"): + (element,) = _children(target) + return FunctionExpression("list_transform", value, _lambda(lambda e: _convert(e, element))).cast(target) + if target.id == "map": + key, item = _children(target) + entry = _lambda( + lambda kv: FunctionExpression( + "row", + _convert(FunctionExpression("struct_extract_at", kv, ConstantExpression(1)), key), + _convert(FunctionExpression("struct_extract_at", kv, ConstantExpression(2)), item), + ) + ) + entries = FunctionExpression("list_transform", FunctionExpression("map_entries", value), entry) + return FunctionExpression("map_from_entries", entries) + fields = _children(target) + # Spark maps struct fields by position, where DuckDB matches them by name. + converted = [ + _convert(FunctionExpression("struct_extract_at", value, ConstantExpression(i)), t) + for i, t in enumerate(fields, 1) + ] + # The cast to a tuple binds only for a struct of exactly this many fields, and typeof keeps it from running. + same_size = FunctionExpression("typeof", value.cast(DuckDBPyType(f"TUPLE({', '.join(['VARCHAR'] * len(fields))})"))) + is_null = value.isnull() | same_size.isnull() + return ( + CaseExpression(is_null, ConstantExpression(None)).otherwise(FunctionExpression("row", *converted)).cast(target) + ) + + +def _reference(column: "Column") -> _Ref | None: + """The column reference ``column`` is, or None for any other expression, such as a cast named after one.""" + name = column._name + is_reference = column._output is None and name is not None and len(name) == 1 and isinstance(name[0], _Ref) + return name[0] if is_reference else None + + +def _spark_cast(column: "Column", target: DuckDBPyType) -> Expression: + """Cast like Spark, which truncates to integral types, takes UTF-8 bytes and maps struct fields by position.""" + if not _cast_differs(target): + return column.expr.cast(target) + # A column reference costs nothing to repeat, and evaluating anything else once makes the cast about 3 times slower. + if _reference(column) is not None: + return _convert(column.expr, target) + return _evaluated_once(column.expr, lambda value: _convert(value, target)) + + +def _duckdb_type(data_type: DataType | str, cast: bool) -> DuckDBPyType: + """The DuckDB type of ``data_type``, read from text as a Spark type, or as a DuckDB type if Spark rejects it. + + ``cast`` reads the target of a cast, where Spark allows char and varchar. + """ + if not isinstance(data_type, str): + return data_type.duckdb_type + try: + return _parse_datatype_string(data_type, cast=cast).duckdb_type + except ParseException as spark_error: + try: + return DuckDBPyType(data_type) + except duckdb.Error: + raise spark_error from None + +def _derived(expr: Expression, name: Name | None, *sources: object, output: str | _Ref | None = None) -> "Column": + """Wrap ``expr`` in a Column that is an aggregate if any of the Columns it was built from is one. + + It is a constant if it was built only from constant Columns and Python values. + """ + column = Column(expr, name, output) + column._aggregate_inputs = tuple(x for s in sources if isinstance(s, Column) for x in s._aggregate_inputs) + if any(isinstance(s, Column) and s._is_aggregate for s in sources): + column._is_aggregate = True + elif all(not isinstance(s, Column) or s._is_constant for s in sources): + column._is_constant = True + return column + + +def _func_op(name: str, sql: str, doc: str = "") -> Callable[["Column"], "Column"]: def _(self: "Column") -> "Column": - # Call the function identified by 'name' on the internal Expression object - expr = getattr(self.expr, name)() - return Column(expr) + njc = getattr(self.expr, name)() + return _derived(njc, _format_name(sql, self), self) _.__doc__ = doc return _ @@ -46,9 +559,13 @@ def _(self: "Column") -> "Column": def _bin_op( name: str, + sql: str, doc: str = "binary operator", ) -> Callable[["Column", Union["Column", "LiteralType", "DecimalLiteral", "DateTimeLiteral"]], "Column"]: - """Create a method for given binary operator.""" + """Create a method for given binary operator. + + ``sql`` is Spark's name for the result, with ``{0}`` for this column and ``{1}`` for the other operand. + """ def _( self: "Column", @@ -56,14 +573,41 @@ def _( ) -> "Column": jc = _get_expr(other) njc = getattr(self.expr, name)(jc) - return Column(njc) + return _derived(njc, _format_name(sql, self, other), self, other) _.__doc__ = doc return _ +def _division( + divide: Callable[[Expression, Expression], Expression], sql: str, reflected: bool +) -> Callable[["Column", Any], "Column"]: + """Divide or take a remainder like Spark without ANSI mode, which gives NULL for a zero divisor. + + DuckDB gives infinity for ``/`` by zero, and for ``%`` by zero raises on integers and gives NaN on doubles. + """ + + def _(self: "Column", other: Any) -> "Column": # noqa: ANN401 + left, right = (_get_expr(other), self.expr) if reflected else (self.expr, _get_expr(other)) + result = divide(left, FunctionExpression("nullif", right, ConstantExpression(0))) + return _derived(result, _format_name(sql, self, other), self, other) + + return _ + + +def _like(name: str, sql: str) -> Callable[["Column", Any], "Column"]: + """Match a LIKE pattern with Spark's default escape character, a backslash.""" + + def _(self: "Column", other: Any) -> "Column": # noqa: ANN401 + func = FunctionExpression(name, self.expr, _get_expr(other), ConstantExpression("\\")) + return _derived(func, _format_name(sql, self, other), self, other) + + return _ + + def _bin_func( name: str, + sql: str, doc: str = "binary function", ) -> Callable[["Column", Union["Column", "LiteralType", "DecimalLiteral", "DateTimeLiteral"]], "Column"]: """Create a function expression for the given binary function.""" @@ -72,9 +616,8 @@ def _( self: "Column", other: Union["Column", "LiteralType", "DecimalLiteral", "DateTimeLiteral"], ) -> "Column": - other = _get_expr(other) - func = FunctionExpression(name, self.expr, other) - return Column(func) + func = FunctionExpression(name, self.expr, _get_expr(other)) + return _derived(func, _format_name(sql, self, other), self, other) _.__doc__ = doc return _ @@ -97,48 +640,106 @@ class Column: .. versionadded:: 1.3.0 """ - def __init__(self, expr: Expression) -> None: # noqa: D107 + # Set on Columns that contain an aggregate function, so select() knows to aggregate. + _is_aggregate = False + # Set on Columns that reference no column, which Spark accepts next to aggregates without a GROUP BY. + _is_constant = False + # The arguments of the aggregate functions in the Column, whose types decide the type Spark gives the result. + _aggregate_inputs: tuple[Expression, ...] = () + # Set on aliases and struct field accesses, which Spark names struct fields after; see _struct_field. + _field: str | None = None + + def __init__(self, expr: Expression, name: Name | None = None, output: str | _Ref | None = None) -> None: + """Wrap an expression. + + ``name`` is Spark's name for the expression, used for it inside other names and, unless ``output`` is + set, as its column name; without it, DuckDB names the column. ``output`` is the column name of an + aliased expression and of a cast of one, or the column reference under a cast. + """ self.expr = expr + self._name = name + self._output = output + + def __bool__(self) -> bool: # noqa: D105 + raise PySparkValueError(error_class="CANNOT_CONVERT_COLUMN_INTO_BOOL", message_parameters={}) + + def __repr__(self) -> str: # noqa: D105 + return f"Column<'{_render(self._name) if self._name is not None else self.expr.get_name()}'>" + + def _with_name(self, name: Name | None) -> "Column": + named = copy.copy(self) + named._name = name + named._output = None + named._field = None + return named + + def _struct_field(self) -> str | None: + """The name Spark gives this Column as a field of a struct, or None where Spark numbers the field.""" + reference = _reference(self) + return reference.parts[-1] if reference is not None else self._field + + def _projection_alias(self, columns: Iterable[str], references: Collection[str] = ()) -> str | None: + """The alias that gives this expression Spark's column name in a projection over ``columns``. + + ``references`` holds the renderings of the qualified references to the columns, folded to lower case. None + when DuckDB already gives it Spark's name, or when Spark's name is not known. + """ + stored = {c.casefold(): c for c in columns} + if isinstance(self._output, _Ref): + parts = _unqualified(self._output.parts, stored.keys(), references) + if len(parts) == 1: + return parts[0] + elif self._output is not None: + return self._output + if self._name is None: + return None + reference = _reference(self) + if reference is not None: + # Aliasing a reference that does not resolve would turn DuckDB's "not found" error into a confusing + # one about a select list alias referencing itself. + name = reference.parts[-1] + return name if stored.get(name.casefold(), name) != name else None + return _render(self._name, stored.keys(), references) # arithmetic operators def __neg__(self) -> "Column": # noqa: D105 - return Column(-self.expr) + return _derived(-self.expr, _format_name("negative({0})", self), self) # `and`, `or`, `not` cannot be overloaded in Python, # so use bitwise operators as boolean operators - __and__ = _bin_op("__and__") - __or__ = _bin_op("__or__") - __invert__ = _func_op("__invert__") - __rand__ = _bin_op("__rand__") - __ror__ = _bin_op("__ror__") + __and__ = _bin_op("__and__", "({0} AND {1})") + __or__ = _bin_op("__or__", "({0} OR {1})") + __invert__ = _func_op("__invert__", "(NOT {0})") + __rand__ = _bin_op("__rand__", "({0} AND {1})") + __ror__ = _bin_op("__ror__", "({0} OR {1})") - __add__ = _bin_op("__add__") + __add__ = _bin_op("__add__", "({0} + {1})") - __sub__ = _bin_op("__sub__") + __sub__ = _bin_op("__sub__", "({0} - {1})") - __mul__ = _bin_op("__mul__") + __mul__ = _bin_op("__mul__", "({0} * {1})") - __div__ = _bin_op("__div__") + __div__ = _division(operator.truediv, "({0} / {1})", reflected=False) - __truediv__ = _bin_op("__truediv__") + __truediv__ = _division(operator.truediv, "({0} / {1})", reflected=False) - __mod__ = _bin_op("__mod__") + __mod__ = _division(operator.mod, "({0} % {1})", reflected=False) - __pow__ = _bin_op("__pow__") + __pow__ = _bin_op("__pow__", "POWER({0}, {1})") - __radd__ = _bin_op("__radd__") + __radd__ = _bin_op("__radd__", "({0} + {1})") - __rsub__ = _bin_op("__rsub__") + __rsub__ = _bin_op("__rsub__", "({1} - {0})") - __rmul__ = _bin_op("__rmul__") + __rmul__ = _bin_op("__rmul__", "({0} * {1})") - __rdiv__ = _bin_op("__rdiv__") + __rdiv__ = _division(operator.truediv, "({1} / {0})", reflected=True) - __rtruediv__ = _bin_op("__rtruediv__") + __rtruediv__ = _division(operator.truediv, "({1} / {0})", reflected=True) - __rmod__ = _bin_op("__rmod__") + __rmod__ = _division(operator.mod, "({1} % {0})", reflected=True) - __rpow__ = _bin_op("__rpow__") + __rpow__ = _bin_op("__rpow__", "POWER({1}, {0})") def __getitem__(self, k: Any) -> "Column": # noqa: ANN401 """An expression that gets an item at position ``ordinal`` out of a list, @@ -175,10 +776,14 @@ def __getitem__(self, k: Any) -> "Column": # noqa: ANN401 # if k.step is not None: # raise ValueError("Using a slice with a step value is not supported") # return self.substr(k.start, k.stop) + elif isinstance(k, str): + field = _derived(_struct_extract(self.expr, k), _sql_name(self, ".", k), self) + field._field = k + return field else: # TODO: this is super hacky # noqa: TD002, TD003 expr_str = str(self.expr) + "." + str(k) - return Column(ColumnExpression(expr_str)) + return _derived(ColumnExpression(expr_str), _sql_name(self, "[", k, "]"), self) def __getattr__(self, item: Any) -> "Column": # noqa: ANN401 """An expression that gets an item at position ``ordinal`` out of a list, @@ -207,10 +812,29 @@ def __getattr__(self, item: Any) -> "Column": # noqa: ANN401 if item.startswith("__"): msg = "Can not access __ (dunder) method" raise AttributeError(msg) + if item in _SPARK_COLUMN_METHODS: + msg = f"Column.{item} is not implemented" + raise _MissingColumnMethodError(msg) return self[item] + def __iter__(self) -> Iterator["Column"]: # noqa: D105 + # Without it, iteration falls back to __getitem__ with 0, 1, ..., which never ends. + raise PySparkTypeError(error_class="NOT_ITERABLE", message_parameters={"objectName": "Column"}) + + def __contains__(self, item: object) -> bool: # noqa: D105 + raise PySparkValueError(error_class="CANNOT_APPLY_IN_FOR_COLUMN", message_parameters={}) + def alias(self, alias: str) -> "Column": # noqa: D102 - return Column(self.expr.alias(alias)) + name = _sql_name(self, " AS ", _quoted_identifier(alias)) + aliased = _derived(self.expr.alias(alias), name, self, output=alias) + aliased._field = alias + return aliased + + def _case_name(self, *pieces: Any) -> Name | None: # noqa: ANN401 + """Extend the name of a CASE expression, which ends in `` END``, with another branch.""" + if self._name is None: + return None + return self._name[:-1] + _sql_name(*pieces, " END") def when(self, condition: "Column", value: Union["Column", str]) -> "Column": # noqa: D102 if not isinstance(condition, Column): @@ -218,27 +842,33 @@ def when(self, condition: "Column", value: Union["Column", str]) -> "Column": # raise TypeError(msg) v = _get_expr(value) expr = self.expr.when(condition.expr, v) - return Column(expr) + return _derived(expr, self._case_name(" WHEN ", condition, " THEN ", value), self, condition, value) def otherwise(self, value: Union["Column", str]) -> "Column": # noqa: D102 v = _get_expr(value) expr = self.expr.otherwise(v) - return Column(expr) + return _derived(expr, self._case_name(" ELSE ", value), self, value) def cast(self, dataType: DataType | str) -> "Column": # noqa: D102 - internal_type = DuckDBPyType(dataType) if isinstance(dataType, str) else dataType.duckdb_type - return Column(self.expr.cast(internal_type)) + internal_type = _duckdb_type(dataType, cast=True) + # Spark keeps the name of an aliased expression or a top level column under a cast, not of a struct field. + output = self._output if self._output is not None else _reference(self) + if internal_type.id in ("list", "array", "map", "struct"): + # Spark's name of a cast to a complex type leaves the cast out, as Hive SQL has no such casts. + name = _sql_name(self) + else: + name = _sql_name("CAST(", self, " AS ", _sql_type_name(internal_type), ")") + return _derived(_spark_cast(self, internal_type), name, self, output=output) def isin(self, *cols: Iterable[Union["Column", str]] | Union["Column", str]) -> "Column": # noqa: D102 if len(cols) == 1 and isinstance(cols[0], (list, set)): # Only one argument supplied, it's a list cols = cast("tuple", cols[0]) - cols = cast( - "tuple", - [_get_expr(c) for c in cols], - ) - return Column(self.expr.isin(*cols)) + name = _sql_name("(", self, " IN (", *_comma_separated(cols), "))") + if not cols: + return _derived(ConstantExpression(False), name, self) + return _derived(self.expr.isin(*[_get_expr(c) for c in cols]), name, self, *cols) # logistic operators def __eq__( # type: ignore[override] @@ -246,31 +876,32 @@ def __eq__( # type: ignore[override] other: Union["Column", "LiteralType", "DecimalLiteral", "DateTimeLiteral"], ) -> "Column": """Binary function.""" - return Column(self.expr == (_get_expr(other))) + return _derived(self.expr == (_get_expr(other)), _format_name("({0} = {1})", self, other), self, other) def __ne__( # type: ignore[override] self, other: object, ) -> "Column": """Binary function.""" - return Column(self.expr != (_get_expr(other))) + name = _format_name("(NOT ({0} = {1}))", self, other) + return _derived(self.expr != (_get_expr(other)), name, self, other) - __lt__ = _bin_op("__lt__") + __lt__ = _bin_op("__lt__", "({0} < {1})") - __le__ = _bin_op("__le__") + __le__ = _bin_op("__le__", "({0} <= {1})") - __ge__ = _bin_op("__ge__") + __ge__ = _bin_op("__ge__", "({0} >= {1})") - __gt__ = _bin_op("__gt__") + __gt__ = _bin_op("__gt__", "({0} > {1})") # String interrogation methods - contains = _bin_func("contains") - rlike = _bin_func("regexp_matches") - like = _bin_func("~~") - ilike = _bin_func("~~*") - startswith = _bin_func("starts_with") - endswith = _bin_func("suffix") + contains = _bin_func("contains", "contains({0}, {1})") + rlike = _bin_func("regexp_matches", "RLIKE({0}, {1})") + like = _like("like_escape", "{0} LIKE {1}") + ilike = _like("ilike_escape", "ilike({0}, {1})") + startswith = _bin_func("starts_with", "startswith({0}, {1})") + endswith = _bin_func("suffix", "endswith({0}, {1})") # order _asc_doc = """ @@ -340,10 +971,26 @@ def __ne__( # type: ignore[override] [Row(name='Tom'), Row(name='Alice'), Row(name=None)] """ - asc = _unary_op("asc", _asc_doc) - desc = _unary_op("desc", _desc_doc) - nulls_first = _unary_op("nulls_first") - nulls_last = _unary_op("nulls_last") + def asc(self) -> "Column": # noqa: D102 + return _derived(self.expr.asc(), _sql_name(self, " ASC", " NULLS FIRST"), self) + + def desc(self) -> "Column": # noqa: D102 + return _derived(self.expr.desc(), _sql_name(self, " DESC", " NULLS LAST"), self) + + asc.__doc__ = _asc_doc + desc.__doc__ = _desc_doc + + def _null_ordering_name(self, ordering: str) -> Name | None: + if self._name is None: + return None + sorted_by = self._name[:-1] if self._name[-1] in (" NULLS FIRST", " NULLS LAST") else self._name + return (*sorted_by, ordering) + + def nulls_first(self) -> "Column": # noqa: D102 + return _derived(self.expr.nulls_first(), self._null_ordering_name(" NULLS FIRST"), self) + + def nulls_last(self) -> "Column": # noqa: D102 + return _derived(self.expr.nulls_last(), self._null_ordering_name(" NULLS LAST"), self) def asc_nulls_first(self) -> "Column": # noqa: D102 return self.asc().nulls_first() @@ -358,7 +1005,26 @@ def desc_nulls_last(self) -> "Column": # noqa: D102 return self.desc().nulls_last() def isNull(self) -> "Column": # noqa: D102 - return Column(self.expr.isnull()) + return _derived(self.expr.isnull(), _format_name("({0} IS NULL)", self), self) def isNotNull(self) -> "Column": # noqa: D102 - return Column(self.expr.isnotnull()) + return _derived(self.expr.isnotnull(), _format_name("({0} IS NOT NULL)", self), self) + + +def _column_name(name: str) -> str: + """``name`` as the name of a column, which DuckDB cannot leave empty: it reads an empty alias as none.""" + if name == "": + msg = "a column with an empty name is not supported" + raise ContributionsAcceptedError(msg) + return name + + +def _named_expr(column: Column, columns: Iterable[str], references: Collection[str] = ()) -> Expression: + """The expression of ``column`` aliased to the name Spark gives it in a projection over ``columns``. + + ``references`` holds the renderings of the qualified references to the columns, folded to lower case. + """ + if isinstance(column._output, str): + _column_name(column._output) + alias = column._projection_alias(columns, references) + return column.expr if alias is None else column.expr.alias(alias) diff --git a/duckdb/experimental/spark/sql/dataframe.py b/duckdb/experimental/spark/sql/dataframe.py index ae0935ecd..bcd3b0d04 100644 --- a/duckdb/experimental/spark/sql/dataframe.py +++ b/duckdb/experimental/spark/sql/dataframe.py @@ -1,23 +1,102 @@ +import copy +import datetime +import re import uuid -from collections.abc import Callable +from collections import Counter +from collections.abc import Callable, Iterable from functools import reduce from keyword import iskeyword from typing import ( TYPE_CHECKING, Any, + NamedTuple, + NoReturn, + TypeVar, Union, cast, overload, ) import duckdb -from duckdb import ColumnExpression, Expression, StarExpression - -from ..errors import PySparkIndexError, PySparkTypeError, PySparkValueError -from .column import Column -from .readwriter import DataFrameWriter +from duckdb import ( + CaseExpression, + ColumnExpression, + ConstantExpression, + Expression, + FunctionExpression, + LambdaExpression, + SQLExpression, + StarExpression, +) +from duckdb.sqltypes import DuckDBPyType + +from ..errors import ( + AnalysisException, + PySparkIndexError, + PySparkNotImplementedError, + PySparkTypeError, + PySparkValueError, +) +from ..exception import ContributionsAcceptedError +from .column import ( + Column, + _bare_column, + _column_name, + _column_reference, + _convert, + _java_float_string, + _name_parts, + _named_expr, + _qualified_star, + _quoted_identifier, + _Ref, + _ref_name, + _reference, + _struct_extract, +) +from .readwriter import DataFrameWriter, _spark_sql_type from .type_utils import duckdb_to_spark_schema -from .types import Row, StructType +from .types import ( + _INFINITE_MICROSECONDS, + ArrayType, + AtomicType, + BinaryType, + BooleanType, + DataType, + DateType, + DayTimeIntervalType, + DecimalType, + DoubleType, + FloatType, + HugeIntegerType, + LongType, + MapType, + NullType, + NumericType, + Row, + StringType, + StructField, + StructType, + TimeNSType, + TimeNTZType, + TimestampMillisecondNTZType, + TimestampNanosecondNTZType, + TimestampNanosecondType, + TimestampNTZType, + TimestampSecondNTZType, + TimestampType, + TimeType, + UnsignedByteType, + UnsignedHugeIntegerType, + UnsignedIntegerType, + UnsignedLongType, + UnsignedShortType, + _apply, + _contains_type, + _duckdb_to_python, + _escape_meta_characters, + _replace_types, +) if TYPE_CHECKING: import pyarrow as pa @@ -29,20 +108,590 @@ from duckdb.experimental.spark.sql import functions as spark_sql_functions +# Spark caps the row count of show() one below its largest array length. +_SHOW_MAX_ROWS = 2_147_483_631 +_FULL_WIDTH = re.compile( + "[\u1100-\u115f\u2e80-\ua4cf\uac00-\ud7a3\uf900-\ufaff\ufe10-\ufe19\ufe30-\ufe6f\uff00-\uff60\uffe0-\uffe6]" +) + + +def _string_half_width(text: str) -> int: + # Spark counts UTF-16 code units and one more for each full width character. + return len(text.encode("utf-16-le")) // 2 + len(_FULL_WIDTH.findall(text)) + + +def _pad(text: str, width: int, align_right: bool) -> str: + fill = " " * (width - _string_half_width(text)) + return fill + text if align_right else text + fill + + +def _truncate_cell(text: str, truncate: int) -> str: + if 0 < truncate < len(text): + return text[:truncate] if truncate < 4 else text[: truncate - 3] + "..." + return text + + +def _fraction_string(microsecond: int) -> str: + return f".{microsecond:06d}".rstrip("0") if microsecond else "" + + +def _day_time_interval_string(value: datetime.timedelta) -> str: + micros = (value.days * 86_400 + value.seconds) * 1_000_000 + value.microseconds + days, rest = divmod(abs(micros), 86_400_000_000) + hours, rest = divmod(rest, 3_600_000_000) + minutes, rest = divmod(rest, 60_000_000) + seconds, microsecond = divmod(rest, 1_000_000) + sign = "-" if micros < 0 else "" + return ( + f"INTERVAL '{sign}{days} {hours:02d}:{minutes:02d}:{seconds:02d}{_fraction_string(microsecond)}' DAY TO SECOND" + ) + + +def _to_pretty_string(value: Any, data_type: DataType) -> str: # noqa: ANN401 + """Renders a value as Spark's show() prints it.""" + if value is None: + return "NULL" + if isinstance(data_type, BooleanType): + return "true" if value else "false" + if isinstance(data_type, (FloatType, DoubleType)): + return _java_float_string(value, isinstance(data_type, FloatType)) + if isinstance(data_type, DecimalType): + return f"{value:f}" + if isinstance(data_type, BinaryType): + return "[" + " ".join(f"{byte:02X}" for byte in value) + "]" + if isinstance(data_type, DayTimeIntervalType): + return _day_time_interval_string(value) + if isinstance(data_type, ArrayType): + return "[" + ", ".join(_to_pretty_string(v, data_type.elementType) for v in value) + "]" + if isinstance(data_type, MapType): + # DuckDB returns maps with unhashable keys as a dict of a key list and a value list. + nested_keys = isinstance(data_type.keyType, (ArrayType, MapType, StructType)) + entries = zip(value["key"], value["value"], strict=True) if nested_keys else value.items() + rendered = ( + f"{_to_pretty_string(k, data_type.keyType)} -> {_to_pretty_string(v, data_type.valueType)}" + for k, v in entries + ) + return "{" + ", ".join(rendered) + "}" + if isinstance(data_type, StructType): + values = value.values() if isinstance(value, dict) else value + return ( + "{" + + ", ".join(_to_pretty_string(v, f.dataType) for v, f in zip(values, data_type.fields, strict=True)) + + "}" + ) + if isinstance(value, datetime.datetime): + return value.replace(tzinfo=None).isoformat(" ", "seconds") + _fraction_string(value.microsecond) + if isinstance(value, datetime.time): + return value.replace(tzinfo=None).isoformat("seconds") + _fraction_string(value.microsecond) + if isinstance(value, datetime.date): + return value.isoformat() + return str(value) + + +def _table_lines(header: list[str], cells: list[list[str]], align_right: bool) -> list[str]: + widths = [max([3, *(_string_half_width(row[i]) for row in [header, *cells])]) for i in range(len(header))] + separator = "+" + "+".join("-" * width for width in widths) + "+" + + def line(row: list[str]) -> str: + return "|" + "|".join(_pad(cell, width, align_right) for cell, width in zip(row, widths, strict=True)) + "|" + + return [separator, line(header), separator, *(line(row) for row in cells), separator] + + +def _vertical_lines(header: list[str], cells: list[list[str]]) -> list[str]: + name_width = max([3, *(_string_half_width(name) for name in header)]) + data_width = max([3, *(_string_half_width(cell) for row in cells for cell in row)]) + lines = [] + for i, row in enumerate(cells): + lines.append(f"-RECORD {i}".ljust(name_width + data_width + 5, "-")) + lines.extend( + f" {_pad(name, name_width, False)} | {_pad(cell, data_width, False)} " + for name, cell in zip(header, row, strict=True) + ) + return lines + + +def _raise_no_columns() -> NoReturn: + raise PySparkNotImplementedError( + message="This DataFrame has no columns. DuckDB has no relation without columns, so only schema, " + "columns, dtypes and printSchema are available." + ) + + +class _NoColumnsRelation: + """Stands in for a relation without columns, such as the result of a statement that returns none.""" + + def __init__(self) -> None: + self.columns: list[str] = [] + self.types: list[DuckDBPyType] = [] + + def __getattr__(self, name: str) -> NoReturn: + _raise_no_columns() + + def __contains__(self, item: str) -> NoReturn: + _raise_no_columns() + + def __str__(self) -> str: + return "DataFrame[]" + + +_T = TypeVar("_T") + +_JOIN_RELATIONS = frozenset({"JOIN_RELATION", "CROSS_PRODUCT_RELATION"}) +# DuckDB relations that keep their child's bindings (Relation::InheritsColumnBindings in the engine); a subquery +# keeps them when its child does. +_INHERITS_BINDINGS = frozenset({"FILTER_RELATION", "ORDER_RELATION", "LIMIT_RELATION", "DISTINCT_RELATION"}) +# DuckDB relations whose columns are those of their (first) child. +_KEEPS_COLUMNS = _INHERITS_BINDINGS | {"SUBQUERY_RELATION", "SET_OPERATION_RELATION"} + +# Spark's join types, after it lowers the case and drops underscores, and the DuckDB join each one is. +_JOIN_TYPES = { + "inner": "inner", + "cross": "inner", + "outer": "outer", + "full": "outer", + "fullouter": "outer", + "left": "left", + "leftouter": "left", + "right": "right", + "rightouter": "right", + "semi": "semi", + "leftsemi": "semi", + "anti": "anti", + "leftanti": "anti", +} + + +def _name_key(name: str) -> str: + # Spark compares names ignoring case letter by letter, so unlike casefold "straße" and "STRASSE" differ. + return name.lower() + + +def _join_type(how: str) -> str: + join_type = _JOIN_TYPES.get(how.lower().replace("_", "")) + if join_type is None: + raise AnalysisException(error_class="UNSUPPORTED_JOIN_TYPE", message_parameters={"join_type": f"'{how}'"}) + return join_type + + +def _position(index: int) -> Expression: + """A reference to the column at ``index`` of the relation a projection reads, which tells apart equal names.""" + return SQLExpression(f"#{index + 1}") + + +def _ambiguous(name: str) -> AnalysisException: + return AnalysisException( + error_class="AMBIGUOUS_REFERENCE", + message_parameters={"name": f"`{name}`", "referenceNames": f"[`{name}`, `{name}`]"}, + ) + + +def _unresolved(name: str, columns: list[str]) -> AnalysisException: + return AnalysisException( + error_class="UNRESOLVED_COLUMN.WITH_SUGGESTION", + message_parameters={"objectName": f"`{name}`", "proposal": ", ".join(f"`{c}`" for c in columns)}, + ) + + +def _distinct_names(names: list[str]) -> list[str]: + """``names`` with each repeat of an earlier name, ignoring case, renamed to the first ``_`` none has. + + DuckDB's binder renames repeats in a subquery to ``_`` too, but only avoids the names before it, so + a renamed column can take the name of a later column. + """ + taken = {_name_key(name) for name in names} + seen: set[str] = set() + result = [] + for name in names: + key = _name_key(name) + if key in seen: + suffix = 1 + while _name_key(f"{name}_{suffix}") in taken: + suffix += 1 + name = f"{name}_{suffix}" + taken.add(_name_key(name)) + seen.add(key) + result.append(name) + return result + + +class _BoundRef(_Ref): + """A column reference a DataFrame handed out, such as ``df.a``, with the rendering of its qualified expression.""" + + __slots__ = ("reference",) + + def __init__(self, parts: tuple[str, ...], reference: str) -> None: + super().__init__(parts) + self.reference = reference + + +def _join_input_names(names: list[str]) -> list[str]: + """``names`` with each repeat of an earlier name, ignoring case, renamed to a name of its own. + + No column of another DataFrame has such a name, which matters because a join shows the DuckDB names of all its + inputs to a bare column reference. + """ + seen: set[str] = set() + result = [] + for name in names: + result.append(f"{name}_{uuid.uuid4().hex}" if _name_key(name) in seen else name) + seen.add(_name_key(name)) + return result + + +def _check_references( + cols: Iterable[Column], + column_expr: Callable[[str], Expression], + references: Callable[[], list[set[str]]], + bindings: frozenset[str], + check_hidden: Callable[[str], None], +) -> None: + """Raise where a Column names a column by a name that more than one column in a scope has, or that none has. + + ``column_expr`` resolves a name in the scope and raises for an ambiguous or unresolved one; ``references`` gives + the qualified references that name each column of the scope, whose relation names are ``bindings``; + ``check_hidden`` raises for a name DuckDB gave a repeated column, which no Spark name refers to. As in Spark, a + leading part that names a relation is a qualifier before it is a column. + """ + counts: Counter[str] | None = None + for col in cols: + for ref in (col._name if isinstance(col, Column) else None) or (): + if not isinstance(ref, _Ref) or ref.parts[-1] == "*": + continue + if isinstance(ref, _BoundRef): + qualified, name = ref.reference, ref.parts[0] + elif len(ref.parts) > 1 and ref.parts[0].casefold() in bindings: + qualified, name = str(ColumnExpression(*ref.parts[:2])).casefold(), ref.parts[1] + else: + column_expr(ref.parts[0]) + continue + if counts is None: + counts = Counter(reference for names in references() for reference in names) + if counts[qualified] > 1: + raise _ambiguous(name) + if counts[qualified] == 0: + check_hidden(name) + + +class _JoinInputs(NamedTuple): + """The inputs of a join; the relation names they expose stay visible in the scope of the joined relation.""" + + left: "DataFrame" + right: "DataFrame" + using: list[str] + how: str + + def binding_names(self) -> frozenset[str]: + return self.left._binding_names() | self.right._binding_names() + + def binding_width(self) -> int: + """The number of columns the inputs add to the FROM clause; the binder hides a semi or anti join's right.""" + width = self.left._binding_width() + return width if self.how in ("semi", "anti") else width + self.right._binding_width() + + def column_expr(self, name: str) -> Expression | None: + """The expression of the joined relation's column ``name``, or None where neither input has the column.""" + if _name_key(name) in (_name_key(key) for key in self.using): + if self.how == "outer": + return _bare_column(name) + return (self.right if self.how == "right" else self.left)._join_input_key_expr(name) + outputs = (self.left,) if self.how in ("semi", "anti") else (self.left, self.right) + sides = [side for side in outputs if side._has_column(name)] + if len(sides) > 1: + raise _ambiguous(name) + if not sides: + self.check_hidden(name) + return None + return sides[0]._join_input_column_expr(name) + + def check_hidden(self, name: str) -> None: + """Raise where ``name`` is a name DuckDB gave a repeated column of an input, which no Spark name refers to.""" + if _name_key(name) in self.left._hidden_names() | self.right._hidden_names(): + raise _unresolved(name, self.column_names()) + + def column_exprs(self, offset: int = 0) -> list[Expression]: + """One expression per column of the joined relation, in its order, which tells apart equal names. + + ``offset`` is the position of the join's first input column among the columns of the FROM clause it is in. + """ + left = self.left._join_input_column_exprs(offset) + right = self.right._join_input_column_exprs(offset + self.left._binding_width()) + return self._output_columns(left, right, self.column_expr) + + def relation_columns(self, binding: str, offset: int) -> list[tuple[Expression, str]] | None: + """The columns and names of the input relation called ``binding``, if there is one; see column_exprs.""" + columns = self.left._join_input_relation_columns(binding, offset) + if columns is None and self.how not in ("semi", "anti"): + columns = self.right._join_input_relation_columns(binding, offset + self.left._binding_width()) + return columns + + def column_names(self) -> list[str]: + """Spark's names of the joined relation's columns; a key of a right join has the right input's name.""" + left, right = self.left.columns, self.right.columns + if self.how != "right": + return self._output_columns(left, right, lambda name: name) + right_names = {_name_key(name): name for name in reversed(right)} + return self._output_columns(left, right, lambda name: right_names[_name_key(name)]) + + def column_references(self) -> list[set[str]]: + """For each column of the joined relation, the renderings of the qualified references that name it.""" + left, right = self.left._join_input_references(), self.right._join_input_references() + return self._output_columns(left, right, lambda name: {str(self.column_expr(name)).casefold()}) + + def _output_columns(self, left: list[_T], right: list[_T], join_key: Callable[[str], _T]) -> list[_T]: + """The items of the joined relation's columns, given one per input column and ``join_key`` for each key.""" + names = self.left.columns + left_keys = self._key_positions(names) + items = [join_key(names[i]) if i in left_keys else item for i, item in enumerate(left)] + if self.how in ("semi", "anti"): + return items + right_keys = self._key_positions(self.right.columns) + return items + [item for i, item in enumerate(right) if i not in right_keys] + + def _key_positions(self, names: list[str]) -> set[int]: + """The positions of the join keys among ``names``: a key is the first column of its name, as in Spark.""" + first = _first_positions(names) + return {first[_name_key(key)] for key in self.using if _name_key(key) in first} + + +def _first_positions(names: list[str]) -> dict[str, int]: + """The position of the first of ``names`` with each name.""" + first: dict[str, int] = {} + for i, name in enumerate(names): + first.setdefault(_name_key(name), i) + return first + + +def _aggregate( + df: "DataFrame", groups: list[Expression], aggregates: list[Column], group_by: str = "" +) -> duckdb.DuckDBPyRelation: + """Aggregate the relation of ``df`` with Spark's result types.""" + if not group_by and not all(isinstance(c, Column) and (c._is_aggregate or c._is_constant) for c in aggregates): + raise AnalysisException(error_class="MISSING_GROUP_BY", message_parameters={}) + relation = df.relation + projections = groups + [_named_expr(c, df.columns, df._qualified_references()) for c in aggregates] + # Without groups the relational API aggregates only when an aggregate function is present, and Spark always does. + extra = [] if group_by or any(c._is_aggregate for c in aggregates) else [FunctionExpression("count_star")] + rel = relation.aggregate(projections + extra, group_by) + return _narrow_aggregates(rel, relation, len(groups), [c._aggregate_inputs for c in aggregates]) + + +# An aggregate over one of these can be a HUGEINT without being a widened sum of a Spark integer. +_HUGEINT_SOURCES = ( + HugeIntegerType, + UnsignedHugeIntegerType, + UnsignedLongType, + UnsignedIntegerType, + UnsignedShortType, + UnsignedByteType, +) + + +def _narrow_aggregates( + relation: duckdb.DuckDBPyRelation, + source: duckdb.DuckDBPyRelation, + group_count: int, + inputs: list[tuple[Expression, ...]], +) -> duckdb.DuckDBPyRelation: + """Give aggregate results Spark's bigint where DuckDB widens integer sums to HUGEINT, which Spark lacks. + + ``inputs`` holds the arguments of the aggregate functions of each result, as expressions over ``source``. A + result keeps its HUGEINT when one of them has a DuckDB only integer type. Columns of ``relation`` past the + groups and those results are left out. + """ + schema = duckdb_to_spark_schema(relation.columns, relation.types) + fields = schema.fields[: group_count + len(inputs)] + columns = [] + narrowed = len(fields) < len(schema.fields) + for i, field in enumerate(fields): + column = _position(i) + if ( + i >= group_count + and _contains_type(field.dataType, (HugeIntegerType,)) + and not _reads_hugeint_source(source, inputs[i - group_count]) + ): + narrow = _replace_types(field.dataType, {HugeIntegerType: LongType()}) + column = column.cast(narrow.duckdb_type) + narrowed = True + columns.append(column.alias(field.name)) + return relation.select(*columns) if narrowed else relation + + +# The timestamp types with a time zone, which DuckDB returns to Python only with a time zone package. +_INSTANT_TYPES = (TimestampType, TimestampNanosecondType) + + +def _epoch_microseconds(expr: Expression, data_type: DataType) -> Expression: + """``expr`` with each timestamp with a time zone in it, at any depth, as its microseconds since the epoch. + + An infinite timestamp becomes the largest or the smallest BIGINT, which no finite one reaches. + """ + if not _contains_type(data_type, _INSTANT_TYPES): + return expr + if isinstance(data_type, _INSTANT_TYPES): + # A cast to microseconds rounds, where DuckDB's own conversion of nanoseconds truncates. + if isinstance(data_type, TimestampNanosecondType): + microseconds = FunctionExpression("epoch_ns", expr) // ConstantExpression(1000) + else: + microseconds = FunctionExpression("epoch_us", expr) + return ( + CaseExpression(expr == ConstantExpression("infinity"), ConstantExpression(_INFINITE_MICROSECONDS)) + .when(expr == ConstantExpression("-infinity"), ConstantExpression(-_INFINITE_MICROSECONDS)) + .otherwise(microseconds) + ) + if isinstance(data_type, ArrayType): + element = _epoch_microseconds(ColumnExpression("x"), data_type.elementType) + return FunctionExpression("list_transform", expr, LambdaExpression("x", element)) + if isinstance(data_type, MapType): + entry = ColumnExpression("e") + key = _epoch_microseconds(_struct_extract(entry, "key"), data_type.keyType).alias("key") + value = _epoch_microseconds(_struct_extract(entry, "value"), data_type.valueType).alias("value") + entries = FunctionExpression( + "list_transform", + FunctionExpression("map_entries", expr), + LambdaExpression("e", FunctionExpression("struct_pack", key, value)), + ) + return FunctionExpression("map_from_entries", entries) + fields = [ + _epoch_microseconds(_struct_extract(expr, f.name), f.dataType).alias(f.name) + for f in cast("StructType", data_type).fields + ] + return CaseExpression(expr.isnull(), ConstantExpression(None)).otherwise(FunctionExpression("struct_pack", *fields)) + + +def _reads_hugeint_source(source: duckdb.DuckDBPyRelation, inputs: tuple[Expression, ...]) -> bool: + if not inputs: + return False + types = duckdb_to_spark_schema([f"c{i}" for i in range(len(inputs))], source.select(*inputs).types) + return any(_contains_type(f.dataType, _HUGEINT_SOURCES) for f in types.fields) + class DataFrame: # noqa: D101 - def __init__(self, relation: duckdb.DuckDBPyRelation, session: "SparkSession") -> None: # noqa: D107 - self.relation = relation + def __init__( # noqa: D107 + self, + relation: duckdb.DuckDBPyRelation | None, + session: "SparkSession", + names: list[str] | None = None, + join_inputs: _JoinInputs | None = None, + ) -> None: + self.relation = relation if relation is not None else _NoColumnsRelation() self.session = session - self._schema = None - if self.relation is not None: - self._schema = duckdb_to_spark_schema(self.relation.columns, self.relation.types) + self._join_scope: tuple[duckdb.DuckDBPyRelation, _JoinInputs] | None = None + if join_inputs is not None: + self._join_scope = (self.relation, join_inputs) + # For each column, the references that named it in the DataFrame this one filters, sorts or renames. + self._inherited_references: list[set[str]] | None = None + names = self.relation.columns if names is None else names + if relation is not None and join_inputs is None and relation.type not in _JOIN_RELATIONS: + # The binder renames repeated names of a relation it reads as a subquery, which Spark keeps, and can + # rename one to the name of another column, so repeats get DuckDB names that no Spark name refers to. + distinct = _distinct_names(names) + if list(map(_name_key, relation.columns)) != list(map(_name_key, distinct)): + self.relation = relation.select(*(_position(i).alias(name) for i, name in enumerate(distinct))) + self._schema = duckdb_to_spark_schema(names, self.relation.types) + self._name_counts = Counter(_name_key(f.name) for f in self._schema.fields) + + def show(self, n: int = 20, truncate: bool | int = True, vertical: bool = False) -> None: + """Prints the first ``n`` rows of the DataFrame to the console. + + Parameters + ---------- + n : int, optional + Number of rows to show. + truncate : bool or int, optional + If set to ``True``, truncate strings longer than 20 chars by default. + If set to a number greater than one, truncates long strings to length ``truncate`` + and align cells right. + vertical : bool, optional + If set to ``True``, print output rows vertically (one line per column value). + + Examples: + -------- + >>> df = spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + >>> df.show(2, vertical=True) + -RECORD 0----- + age | 14 + name | Tom + -RECORD 1----- + age | 23 + name | Alice + only showing top 2 rows + """ + print(self._show_string(n, truncate, vertical)) - def show(self, **kwargs) -> None: # noqa: D102 - self.relation.show() + def _show_string(self, n: int = 20, truncate: bool | int = True, vertical: bool = False) -> str: + if not isinstance(n, int) or isinstance(n, bool): + raise PySparkTypeError( + error_class="NOT_INT", + message_parameters={"arg_name": "n", "arg_type": type(n).__name__}, + ) + if not isinstance(vertical, bool): + raise PySparkTypeError( + error_class="NOT_BOOL", + message_parameters={"arg_name": "vertical", "arg_type": type(vertical).__name__}, + ) + if isinstance(truncate, bool) and truncate: + int_truncate = 20 + else: + try: + int_truncate = int(truncate) + except ValueError: + raise PySparkTypeError( + error_class="NOT_BOOL", + message_parameters={"arg_name": "truncate", "arg_type": type(truncate).__name__}, + ) from None + + num_rows = min(max(n, 0), _SHOW_MAX_ROWS) + local = self._in_session_time_zone() + rows = local._fetchall(num_rows + 1) + fields = local.schema.fields + header = [_escape_meta_characters(field.name) for field in fields] + cells = [ + [ + _truncate_cell(_escape_meta_characters(_to_pretty_string(value, field.dataType)), int_truncate) + for value, field in zip(row, fields, strict=True) + ] + for row in rows[:num_rows] + ] + lines = _vertical_lines(header, cells) if vertical else _table_lines(header, cells, int_truncate > 0) + if vertical and not cells: + lines.append("(0 rows)") + elif len(rows) > num_rows: + lines.append(f"only showing top {num_rows} {'row' if num_rows == 1 else 'rows'}") + else: + lines.append("") + return "\n".join(lines) def toPandas(self) -> "PandasDataFrame": # noqa: D102 - return self.relation.df() + # relation.df() reports a missing pandas as InvalidInputException, or names numpy when both are missing. + import pandas # noqa: F401 + + result = self._in_session_time_zone().relation.df(date_as_object=True) + for i, field in enumerate(self.schema.fields): + if isinstance(field.dataType, DateType): + # DuckDB gives a missing date as NaT, and a column of missing dates the datetime64 type. + column = result.iloc[:, i] + result.isetitem(i, column.astype(object).where(column.notna(), None)) + # DuckDB renames repeated column names, which Spark keeps. + result.columns = self.columns + return result + + def _in_session_time_zone(self) -> "DataFrame": + """This DataFrame with each timestamp as its wall clock time in the session time zone, as Spark shows it.""" + local = _replace_types( + self.schema, {TimestampType: TimestampNTZType(), TimestampNanosecondType: TimestampNanosecondNTZType()} + ) + return self._with_types(cast("StructType", local)) + + def _with_types(self, schema: StructType) -> "DataFrame": + """This DataFrame with each column cast to the type of the field at its position in ``schema``.""" + if schema == self.schema: + return self + columns = [ + expr if new == old else expr.cast(new.dataType.duckdb_type).alias(new.name) + for expr, old, new in zip(self._column_exprs(), self.schema.fields, schema.fields, strict=True) + ] + return self._derive(self.relation.select(*columns)) def toArrow(self) -> "pa.Table": """Returns the contents of this :class:`DataFrame` as PyArrow ``pyarrow.Table``. @@ -68,7 +717,7 @@ def toArrow(self) -> "pa.Table": age: [[2,5]] name: [["Alice","Bob"]] """ - return self.relation.to_arrow_table() + return self.relation.to_arrow_table().rename_columns(self.columns) def createOrReplaceTempView(self, name: str) -> None: """Creates or replaces a local temporary view with this :class:`DataFrame`. @@ -99,23 +748,26 @@ def createOrReplaceTempView(self, name: str) -> None: True """ - self.relation.create_view(name, True) + relation = self.relation + names = _distinct_names(self.columns) + if list(map(_name_key, relation.columns)) != list(map(_name_key, names)): + relation = self._relation_named(names) + relation.create_view(name, True) def createGlobalTempView(self, name: str) -> None: # noqa: D102 raise NotImplementedError def withColumnRenamed(self, columnName: str, newName: str) -> "DataFrame": # noqa: D102 - if columnName not in self.relation: + if not self._has_column(columnName): msg = f"DataFrame does not contain a column named {columnName}" raise ValueError(msg) cols = [] - for x in self.relation.columns: - col = ColumnExpression(x) - if x.casefold() == columnName.casefold(): - col = col.alias(newName) + for x, col in zip(self.columns, self._column_exprs(), strict=True): + if _name_key(x) == _name_key(columnName): + col = col.alias(_column_name(newName)) cols.append(col) rel = self.relation.select(*cols) - return DataFrame(rel, self.session) + return self._derive(rel) def withColumn(self, columnName: str, col: Column) -> "DataFrame": # noqa: D102 if not isinstance(col, Column): @@ -123,19 +775,20 @@ def withColumn(self, columnName: str, col: Column) -> "DataFrame": # noqa: D102 error_class="NOT_COLUMN", message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, ) - if columnName in self.relation: + self._check_references(col) + col = self._through_inherited_reference(col) + if self._has_column(columnName): # We want to replace the existing column with this new expression cols = [] - for x in self.relation.columns: - if x.casefold() == columnName.casefold(): - cols.append(col.expr.alias(columnName)) + for x, existing in zip(self.columns, self._column_exprs(), strict=True): + if _name_key(x) == _name_key(columnName): + cols.append(col.expr.alias(_column_name(columnName))) else: - cols.append(ColumnExpression(x)) + cols.append(existing) else: - cols = [ColumnExpression(x) for x in self.relation.columns] - cols.append(col.expr.alias(columnName)) + cols = [*self._column_exprs(), col.expr.alias(_column_name(columnName))] rel = self.relation.select(*cols) - return DataFrame(rel, self.session) + return self._derive(rel) def withColumns(self, *colsMap: dict[str, Column]) -> "DataFrame": """Returns a new :class:`DataFrame` by adding multiple columns or replacing the @@ -184,32 +837,27 @@ def withColumns(self, *colsMap: dict[str, Column]) -> "DataFrame": }, ) - column_names = list(colsMap.keys()) - columns = list(colsMap.values()) - - # Compute this only once - column_names_for_comparison = [x.casefold() for x in column_names] - + self._check_references(*colsMap.values()) + colsMap = {name: self._through_inherited_reference(col) for name, col in colsMap.items()} + replacements: dict[str, tuple[str, Column]] = {} + for name, col in colsMap.items(): + if _name_key(name) in replacements: + raise AnalysisException( + error_class="COLUMN_ALREADY_EXISTS", message_parameters={"columnName": f"`{name.lower()}`"} + ) + replacements[_name_key(name)] = (name, col) cols = [] - for x in self.relation.columns: - if x.casefold() in column_names_for_comparison: - idx = column_names_for_comparison.index(x) - # We extract the column name from the originally passed - # in ones, as the casing might be different than the one - # in the relation - col_name = column_names.pop(idx) - col = columns.pop(idx) + for x, existing in zip(self.columns, self._column_exprs(), strict=True): + if _name_key(x) in replacements: + col_name, col = replacements[_name_key(x)] cols.append(col.expr.alias(col_name)) else: - cols.append(ColumnExpression(x)) + cols.append(existing) - # In case anything is remaining, these are new columns - # that we need to add to the DataFrame - for col_name, col in zip(column_names, columns, strict=False): - cols.append(col.expr.alias(col_name)) + cols.extend(col.expr.alias(_column_name(name)) for name, col in colsMap.items() if not self._has_column(name)) rel = self.relation.select(*cols) - return DataFrame(rel, self.session) + return self._derive(rel) def withColumnsRenamed(self, colsMap: dict[str, str]) -> "DataFrame": """Returns a new :class:`DataFrame` by renaming multiple columns. @@ -255,30 +903,20 @@ def withColumnsRenamed(self, colsMap: dict[str, str]) -> "DataFrame": message_parameters={"arg_name": "colsMap", "arg_type": type(colsMap).__name__}, ) - unknown_columns = set(colsMap.keys()) - set(self.relation.columns) + unknown_columns = [name for name in colsMap if not self._has_column(name)] if unknown_columns: msg = f"DataFrame does not contain column(s): {', '.join(unknown_columns)}" raise ValueError(msg) - # Compute this only once - old_column_names = list(colsMap.keys()) - old_column_names_for_comparison = [x.casefold() for x in old_column_names] - + new_names = {_name_key(old): new for old, new in colsMap.items()} cols = [] - for x in self.relation.columns: - col = ColumnExpression(x) - if x.casefold() in old_column_names_for_comparison: - idx = old_column_names.index(x) - # We extract the column name from the originally passed - # in ones, as the casing might be different than the one - # in the relation - col_name = old_column_names.pop(idx) - new_col_name = colsMap[col_name] - col = col.alias(new_col_name) + for x, col in zip(self.columns, self._column_exprs(), strict=True): + if _name_key(x) in new_names: + col = col.alias(_column_name(new_names[_name_key(x)])) cols.append(col) rel = self.relation.select(*cols) - return DataFrame(rel, self.session) + return self._derive(rel) def transform(self, func: Callable[..., "DataFrame"], *args: Any, **kwargs: Any) -> "DataFrame": # noqa: ANN401 """Returns a new :class:`DataFrame`. Concise syntax for chaining custom transformations. @@ -464,9 +1102,10 @@ def sort(self, *cols: str | Column | list[str | Column], **kwargs: Any) -> "Data message_parameters={"arg_name": "ascending", "arg_type": type(ascending).__name__}, ) + self._check_references(*columns) columns = [spark_sql_functions._to_column_expr(c) for c in columns] rel = self.relation.sort(*columns) - return DataFrame(rel, self.session) + return self._derive(rel) orderBy = sort @@ -532,6 +1171,7 @@ def filter(self, condition: "ColumnOrName") -> "DataFrame": +---+-----+ """ if isinstance(condition, Column): + self._check_references(condition) cond = condition.expr elif isinstance(condition, str): cond = condition @@ -541,7 +1181,7 @@ def filter(self, condition: "ColumnOrName") -> "DataFrame": message_parameters={"arg_name": "condition", "arg_type": type(condition).__name__}, ) rel = self.relation.filter(cond) - return DataFrame(rel, self.session) + return self._derive(rel) where = filter @@ -549,12 +1189,44 @@ def select(self, *cols) -> "DataFrame": # noqa: D102 cols = list(cols) if len(cols) == 1: cols = cols[0] - if isinstance(cols, list): - projections = [x.expr if isinstance(x, Column) else ColumnExpression(x) for x in cols] - else: - projections = [cols.expr if isinstance(cols, Column) else ColumnExpression(cols)] + if not isinstance(cols, list): + cols = [cols] + cols = [column for x in cols for column in self._selected_columns(x)] + self._check_references(*cols) + cols = [self._through_inherited_reference(x) for x in cols] + if any(x._is_aggregate for x in cols): + return DataFrame(_aggregate(self, [], cols), self.session) + columns, references = self.columns, self._qualified_references() + projections = [_named_expr(x, columns, references) for x in cols] rel = self.relation.select(*projections) - return DataFrame(rel, self.session) + return self._derive(rel) + + def _selected_columns(self, col: "ColumnOrName") -> list[Column]: + """The Columns ``select`` gives for ``col``; a star stands for every column, under its Spark name.""" + if (isinstance(col, str) and col == "*") or (isinstance(col, Column) and str(col.expr) == "*"): + return [Column(expr.alias(name)) for expr, name in zip(self._column_exprs(), self.columns, strict=True)] + if isinstance(col, Column): + reference = _reference(col) + parts = () if reference is None else reference.parts + else: + parts = _name_parts(col) + if len(parts) < 2 or parts[-1] != "*": + return [col] if isinstance(col, Column) else [Column(self._scope_column_expr(col), _ref_name(col))] + columns = self._relation_columns(parts[0].casefold()) if len(parts) == 2 else None + if columns is not None: + return [Column(expr.alias(name)) for expr, name in columns] + # A struct column, whose fields DuckDB names as Spark does; a repeated name raises, unless it is a qualifier. + if parts[0].casefold() not in self._check_bindings(): + self._name_column_expr(parts[0]) + return [Column(_qualified_star(parts[:-1]))] + + def _relation_columns(self, binding: str) -> list[tuple[Expression, str]] | None: + """The columns and names of the relation called ``binding`` in this DataFrame's scope, if there is one.""" + if self._join_inputs is not None: + return self._join_inputs.relation_columns(binding, 0) + if binding != self.relation.alias.casefold(): + return None + return list(zip(self._column_exprs(), self.columns, strict=True)) @property def columns(self) -> list[str]: @@ -689,6 +1361,7 @@ def join( """ if on is not None and not isinstance(on, list): on = [on] # type: ignore[assignment] + conditions = [x for x in on or [] if isinstance(x, Column)] if on is not None and not all(isinstance(x, str) for x in on): assert isinstance(on, list) # Get (or create) the Expressions from the list of Columns @@ -698,37 +1371,248 @@ def join( assert isinstance(on[0], Expression), "on should be Column or list of Column" on = reduce(lambda x, y: x.__and__(y), cast("list[Expression]", on)) + left, right = self._disjoint_join_inputs(other) + scope = _JoinInputs(left, right, [], "inner") + _check_references( + conditions, scope.column_expr, scope.column_references, scope.binding_names(), scope.check_hidden + ) + using: list[str] = [] if on is None and how is None: - result = self.relation.join(other.relation) + how = "inner" + result = left.relation.cross(right.relation) else: if how is None: how = "inner" if on is None: on = "true" elif isinstance(on, list) and all(isinstance(x, str) for x in on): - # Passed directly through as a list of strings - on = on + using = on else: on = str(on) assert isinstance(how, str), "how should be a string" - def map_to_recognized_jointype(how: str) -> str: - known_aliases = { - "inner": [], - "outer": ["full", "fullouter", "full_outer"], - "left": ["leftouter", "left_outer"], - "right": ["rightouter", "right_outer"], - "anti": ["leftanti", "left_anti"], - "semi": ["leftsemi", "left_semi"], - } - for type, aliases in known_aliases.items(): - if how == type or how in aliases: - return type - return how - - how = map_to_recognized_jointype(how) - result = self.relation.join(other.relation, on, how) - return DataFrame(result, self.session) + how = _join_type(how) + result = left.relation.join(right.relation, on, how) + return left._joined(result, right, using, how) + + @property + def _join_inputs(self) -> _JoinInputs | None: + """The inputs of the join whose scope this DataFrame's relation shares, if any.""" + if self._join_scope is None or self._join_scope[0] is not self.relation: + return None + return self._join_scope[1] + + def _derive(self, relation: duckdb.DuckDBPyRelation | None) -> "DataFrame": + result = self._derive_scope(relation) + if relation is not None and (relation.type in _INHERITS_BINDINGS or relation.type == "SUBQUERY_RELATION"): + result._inherited_references = self._column_references() + return result + + def _derive_scope(self, relation: duckdb.DuckDBPyRelation | None) -> "DataFrame": + names = self.columns if relation is not None and relation.type in _KEEPS_COLUMNS else None + if self._join_inputs is None or relation is None or self.relation.type == "CROSS_PRODUCT_RELATION": + return DataFrame(relation, self.session, names) + if relation.type not in _INHERITS_BINDINGS and ( + relation.type != "SUBQUERY_RELATION" or self.relation.type not in _INHERITS_BINDINGS + ): + return DataFrame(relation, self.session, names) + if relation.type == "FILTER_RELATION": + return DataFrame(relation, self.session, names, self._join_inputs) + # The engine binds a filter, projection or aggregate over this relation inside the join's own query, where + # it would land below the ORDER BY, LIMIT or DISTINCT, or lose the subquery's alias; a projection prevents it. + if relation.type == "ORDER_RELATION" and self.relation.type == "FILTER_RELATION": + # The sort reads the filter as a subquery, which renames repeated column names. + return DataFrame(relation.select(*self._column_exprs_by_position()), self.session) + return DataFrame(relation.select(StarExpression()), self.session, self.columns) + + def _column_exprs_by_position(self) -> list[Expression]: + return [_position(i).alias(name) for i, name in enumerate(self.columns)] + + def _has_column(self, name: str) -> bool: + return _name_key(name) in self._name_counts + + def _exposes_join_inputs(self) -> bool: + return self._join_inputs is not None and self.relation.type in _JOIN_RELATIONS + + def _binding_names(self) -> frozenset[str]: + if self._exposes_join_inputs(): + return self._join_inputs.binding_names() + return frozenset({self.relation.alias.casefold()}) + + def _disjoint_join_inputs(self, other: "DataFrame") -> tuple["DataFrame", "DataFrame"]: + left, right = self._as_join_input(), other._as_join_input() + # DataFrames derived from one another share their relation's alias, so a join can see one name twice; + # renaming both inputs makes a reference through that name fail instead of silently picking one side. + if left._binding_names().isdisjoint(right._binding_names()): + return left, right + return left._with_fresh_alias(), right._with_fresh_alias() + + def _as_join_input(self) -> "DataFrame": + """This DataFrame as an input of a join, with its repeated names under DuckDB names of their own. + + A join reads an input that is not itself a join as a subquery, where DuckDB's binder would rename repeated + names, and shows the DuckDB names of all its inputs to a bare column reference. The inputs of a join that is + an input were renamed when that join was made. + """ + if self._exposes_join_inputs() or all(count == 1 for count in self._name_counts.values()): + return self + renamed = copy.copy(self) + renamed.relation = self._relation_named(_join_input_names(self.columns)) + renamed._join_scope = None + renamed._inherited_references = self._column_references() + return renamed + + def _relation_named(self, duckdb_names: list[str]) -> duckdb.DuckDBPyRelation: + """This DataFrame's columns under ``duckdb_names``; over a join it reads the join as one relation.""" + return self.relation.select( + *(expr.alias(name) for expr, name in zip(self._column_exprs(), duckdb_names, strict=True)) + ) + + def _with_fresh_alias(self) -> "DataFrame": + if any(count > 1 for count in self._name_counts.values()): + msg = "Joining DataFrames that share an input is not supported when one of them has duplicate column names" + raise ContributionsAcceptedError(msg) + return self.alias(_fresh_alias()) + + def _joined(self, relation: duckdb.DuckDBPyRelation, right: "DataFrame", using: list[str], how: str) -> "DataFrame": + inputs = _JoinInputs(self, right, using, how) + return DataFrame(relation, self.session, inputs.column_names(), inputs) + + def _column_expr(self, name: str) -> Expression: + if self._join_inputs is not None: + expr = self._join_inputs.column_expr(name) + # DuckDB reads a name neither input has as a qualified name or a struct field, as Spark does. + return ColumnExpression(name) if expr is None else expr + self._check_name(name) + return ColumnExpression(self.relation.alias, name) + + def _scope_column_expr(self, name: str) -> Expression: + parts = _name_parts(name) + if len(parts) > 1 or parts[0] == "*": + return _column_reference(name) + return self._name_column_expr(parts[0]) + + def _name_column_expr(self, name: str) -> Expression: + """The column called ``name``, a name without quoting or a qualifier, in this DataFrame's scope.""" + if self._join_inputs is not None: + expr = self._join_inputs.column_expr(name) + return _bare_column(name) if expr is None else expr + self._check_name(name) + return _bare_column(name) + + def _check_name(self, name: str) -> None: + count = self._name_counts[_name_key(name)] + if count > 1: + raise _ambiguous(name) + if count == 0: + self._check_hidden(name) + + def _check_hidden(self, name: str) -> None: + if self._join_inputs is not None: + self._join_inputs.check_hidden(name) + elif _name_key(name) in self._hidden_names(): + raise _unresolved(name, self.columns) + + def _hidden_names(self) -> frozenset[str]: + """The keys of the DuckDB names of this DataFrame's columns that are none of its Spark names.""" + if self._join_inputs is not None: + return self._join_inputs.left._hidden_names() | self._join_inputs.right._hidden_names() + return frozenset(map(_name_key, self.relation.columns)) - self._name_counts.keys() + + def _check_bindings(self) -> frozenset[str]: + return self._join_inputs.binding_names() if self._join_inputs is not None else self._binding_names() + + def _check_references(self, *cols: Column) -> None: + _check_references( + cols, self._name_column_expr, self._column_references, self._check_bindings(), self._check_hidden + ) + + def _column_exprs(self) -> list[Expression]: + if self._join_inputs is not None: + return self._join_inputs.column_exprs() + return self._column_exprs_by_position() + + def _binding_width(self) -> int: + if self._exposes_join_inputs(): + return self._join_inputs.binding_width() + return len(self.relation.columns) + + def _join_input_column_expr(self, name: str) -> Expression: + if not self._exposes_join_inputs(): + self._check_name(name) + return self._join_input_key_expr(name) + + def _join_input_key_expr(self, name: str) -> Expression: + """The first column named ``name`` of this join input, which a join key on the name refers to.""" + if self._exposes_join_inputs(): + expr = self._join_inputs.column_expr(name) + return _bare_column(name) if expr is None else expr + return ColumnExpression(self.relation.alias, name) + + def _join_input_column_exprs(self, offset: int) -> list[Expression]: + if self._exposes_join_inputs(): + return self._join_inputs.column_exprs(offset) + # The binder reads a qualified name as the first column of the name, so a repeated name needs its position. + return [ + _position(offset + i).alias(name) + if self._name_counts[_name_key(name)] > 1 + else ColumnExpression(self.relation.alias, name) + for i, name in enumerate(self.columns) + ] + + def _join_input_relation_columns(self, binding: str, offset: int) -> list[tuple[Expression, str]] | None: + if self._exposes_join_inputs(): + return self._join_inputs.relation_columns(binding, offset) + if binding != self.relation.alias.casefold(): + return None + return list(zip(self._join_input_column_exprs(offset), self.columns, strict=True)) + + def _join_input_references(self) -> list[set[str]]: + if self._exposes_join_inputs(): + return self._join_inputs.column_references() + return [{str(ColumnExpression(self.relation.alias, name)).casefold()} for name in self.columns] + + def _set_operand(self) -> duckdb.DuckDBPyRelation: + if isinstance(self.relation, _NoColumnsRelation): + _raise_no_columns() + return self.relation + + def _checked_set_operand(self, other: "DataFrame", operator: str) -> duckdb.DuckDBPyRelation: + """``other``'s relation as the second input of a set operation with this DataFrame, of the same width.""" + self._set_operand() + relation = other._set_operand() + if len(self.columns) != len(other.columns): + raise _num_columns_mismatch(operator, len(self.columns), len(other.columns)) + return relation + + def _check_compatible_types( + self, other: "DataFrame", operator: str, by_name: bool = False, allow_missing: bool = False + ) -> None: + for i, (mine, theirs) in enumerate(zip(self.schema.fields, other.schema.fields, strict=True)): + if not _compatible(mine.dataType, theirs.dataType, by_name, allow_missing): + raise _incompatible_column_type(operator, i, mine.dataType, theirs.dataType) + + def _positional_operand(self, other: "DataFrame", operator: str) -> duckdb.DuckDBPyRelation: + """``other``'s relation as the second input of a set operation, which Spark resolves by position. + + DuckDB matches struct fields by name. Spark's union matches them by position and keeps the first input's + names; its other set operations reject struct fields of other names. + """ + relation = self._checked_set_operand(other, operator) + self._check_compatible_types(other, operator) + columns = other._column_exprs() + converted = False + for i, (mine, theirs) in enumerate(zip(self.schema.fields, other.schema.fields, strict=True)): + if not _renames_struct_fields(mine.dataType, theirs.dataType): + continue + if operator != "UNION" or not _same_shape(mine.dataType, theirs.dataType): + raise _incompatible_column_type(operator, i, mine.dataType, theirs.dataType) + columns[i] = _convert(columns[i], mine.dataType.duckdb_type).alias(theirs.name) + converted = True + if operator != "UNION": + for mine, theirs in zip(self.schema.fields, other.schema.fields, strict=True): + _check_comparable(mine.name, mine.dataType, theirs.dataType) + return relation.select(*columns) if converted else relation def crossJoin(self, other: "DataFrame") -> "DataFrame": """Returns the cartesian product with another :class:`DataFrame`. @@ -765,7 +1649,8 @@ def crossJoin(self, other: "DataFrame") -> "DataFrame": | 16| Bob| 85| +---+-----+------+ """ - return DataFrame(self.relation.cross(other.relation), self.session) + left, right = self._disjoint_join_inputs(other) + return left._joined(left.relation.cross(right.relation), right, [], "inner") def alias(self, alias: str) -> "DataFrame": """Returns a new :class:`DataFrame` with an alias set. @@ -799,27 +1684,75 @@ def alias(self, alias: str) -> "DataFrame": +-----+-----+---+ """ assert isinstance(alias, str), "alias should be a string" - return DataFrame(self.relation.set_alias(alias), self.session) + return self._derive(self.relation.set_alias(alias)) def drop(self, *cols: "ColumnOrName") -> "DataFrame": # type: ignore[misc] # noqa: D102 - exclude = [] + names = [_name_key(name) for name in self.columns] + references = self._column_references() + dropped: set[int] = set() for col in cols: if isinstance(col, str): - exclude.append(col) + dropped.update(i for i, name in enumerate(names) if name == _name_key(col)) elif isinstance(col, Column): - exclude.append(col.expr.get_name()) + dropped.update(self._dropped_positions(col, names, references)) else: raise PySparkTypeError( error_class="NOT_COLUMN_OR_STR", message_parameters={"arg_name": "col", "arg_type": type(col).__name__}, ) - # Filter out the columns that don't exist in the relation - exclude = [x for x in exclude if x in self.relation.columns] - expr = StarExpression(exclude=exclude) - return DataFrame(self.relation.select(expr), self.session) + kept = [expr for i, expr in enumerate(self._column_exprs()) if i not in dropped] + return self._derive(self.relation.select(*kept) if kept else None) + + def _column_references(self) -> list[set[str]]: + """For each column, the renderings of the qualified references such as ``df.a`` that name it.""" + if self._join_inputs is not None: + own = self._join_inputs.column_references() + else: + own = [{str(ColumnExpression(self.relation.alias, name)).casefold()} for name in self.columns] + if self._inherited_references is None: + return own + return [a | b for a, b in zip(own, self._inherited_references, strict=True)] + + def _qualified_references(self) -> frozenset[str]: + """The renderings of the qualified references that name a column, folded to lower case.""" + return frozenset(reference for references in self._column_references() for reference in references) + + def _through_inherited_reference(self, col: Column) -> Column: + """``col``, or where it is a qualified reference this DataFrame's own scope no longer sees, its column. + + An alias of such a reference is the column under the alias. + """ + if self._join_inputs is not None or self._inherited_references is None or not isinstance(col, Column): + return col + if _is_qualified_reference(col): + rendered, name = str(col.expr).casefold(), None + else: + ref = _aliased_reference(col) + if ref is None: + return col + rendered = ref.reference if isinstance(ref, _BoundRef) else str(ColumnExpression(*ref.parts)).casefold() + name = col._output + positions = [i for i, references in enumerate(self._inherited_references) if rendered in references] + if len(positions) != 1: + return col + return Column(_position(positions[0]).alias(name or self.columns[positions[0]]), col._name, col._output) + + @staticmethod + def _dropped_positions(col: Column, names: list[str], references: list[set[str]]) -> list[int]: + """The positions a Column given to drop() refers to; anything but a column reference refers to none.""" + if not _is_reference(col): + return [] + if _is_qualified_reference(col): + rendered = str(col.expr).casefold() + return [i for i, reference in enumerate(references) if rendered in reference] + parts = col._name[0].parts + positions = [i for i, name in enumerate(names) if name == _name_key(parts[0])] + if len(positions) > 1: + raise _ambiguous(parts[0]) + return positions def __repr__(self) -> str: # noqa: D105 - return str(self.relation) + return "DataFrame[{}]".format(", ".join(f"{name}: {type_name}" for name, type_name in self.dtypes)) def limit(self, num: int) -> "DataFrame": """Limits the result count to the number specified. @@ -851,11 +1784,13 @@ def limit(self, num: int) -> "DataFrame": +---+----+ """ rel = self.relation.limit(num) - return DataFrame(rel, self.session) + return self._derive(rel) def __contains__(self, item: str) -> bool: """Check if the :class:`DataFrame` contains a column by the name of `item`.""" - return item in self.relation + if isinstance(self.relation, _NoColumnsRelation): + _raise_no_columns() + return self._has_column(item) @property def schema(self) -> StructType: @@ -890,13 +1825,13 @@ def __getitem__(self, item: int | str | Column | list | tuple) -> Union[Column, [Row(age=5, name='Bob')] """ if isinstance(item, str): - return Column(duckdb.ColumnExpression(self.relation.alias, item)) + return self._bound_column(item) elif isinstance(item, Column): return self.filter(item) elif isinstance(item, (list, tuple)): return self.select(*item) elif isinstance(item, int): - return spark_sql_functions.col(self._schema[item].name) + return self[self.columns[item]] else: msg = f"Unexpected item type: {type(item)}" raise TypeError(msg) @@ -909,10 +1844,15 @@ def __getattr__(self, name: str) -> Column: >>> df.select(df.age).collect() [Row(age=2), Row(age=5)] """ - if name not in self.relation.columns: + schema = self.__dict__.get("_schema") + if schema is None or name not in schema.names: msg = f"'{self.__class__.__name__}' object has no attribute '{name}'" raise AttributeError(msg) - return Column(duckdb.ColumnExpression(self.relation.alias, name)) + return self._bound_column(name) + + def _bound_column(self, name: str) -> Column: + expr = self._column_expr(name) + return Column(expr, (_BoundRef((name,), str(expr).casefold()),)) @overload def groupBy(self, *cols: "ColumnOrName") -> "GroupedData": ... @@ -1002,21 +1942,19 @@ def printSchema(self, level: int | None = None) -> None: Parameters ---------- level : int, optional - How many levels to print for nested schemas. Prints all levels by default. + How many levels to print for nested schemas. Prints all levels when not set or not positive. Examples: -------- >>> df.printSchema() root - |-- age: bigint (nullable = true) + |-- age: long (nullable = true) |-- name: string (nullable = true) """ - if level is not None and level < 0: - raise PySparkValueError( - error_class="NEGATIVE_VALUE", - message_parameters={"arg_name": "level", "arg_value": str(level)}, - ) - print(self.schema.treeString(level)) + if level: + print(self.schema.treeString(level)) + else: + print(self.schema.treeString()) def union(self, other: "DataFrame") -> "DataFrame": """Return a new :class:`DataFrame` containing union of rows in this and another @@ -1061,7 +1999,7 @@ def union(self, other: "DataFrame") -> "DataFrame": | 1| 2| 3| +----+----+----+ """ # noqa: D205 - return DataFrame(self.relation.union(other.relation), self.session) + return self._derive(self.relation.union(self._positional_operand(other, "UNION"))) unionAll = union @@ -1122,28 +2060,40 @@ def unionByName(self, other: "DataFrame", allowMissingColumns: bool = False) -> """ if allowMissingColumns: df1 = self.select( - *self.relation.columns, - *[ - spark_sql_functions.lit(None).alias(c) - for c in other.relation.columns - if c not in self.relation.columns - ], + *[self[c] for c in self.columns], + *[spark_sql_functions.lit(None).alias(c) for c in other.columns if not self._has_column(c)], ) - df2 = other.select( - *[ - spark_sql_functions.col(c) - if c in other.relation.columns - else spark_sql_functions.lit(None).alias(c) - for c in df1.relation.columns - ] + *[other[c] if other._has_column(c) else spark_sql_functions.lit(None).alias(c) for c in df1.columns] ) - - return df1.unionByName(df2, allowMissingColumns=False) - else: - other = other.select(*self.relation.columns) - - return DataFrame(self.relation.union(other.relation), self.session) + df1._checked_set_operand(df2, "UNION") + df2 = df2._with_map_fields_of(df1) + df1._check_compatible_types(df2, "UNION", by_name=True, allow_missing=True) + # DuckDB's union fills the struct fields one input lacks, as Spark does with missing columns allowed. + return df1._derive(df1.relation.union(df2._set_operand())) + # Stored names are column names, which select would read as Spark attribute names. + reordered = other.select(*[other[c] for c in self.columns]) + self._checked_set_operand(other, "UNION") + for mine, theirs in zip(self.schema.fields, reordered.schema.fields, strict=True): + _check_struct_fields(mine.dataType, theirs.dataType) + reordered = reordered._with_map_fields_of(self) + self._check_compatible_types(reordered, "UNION", by_name=True) + return self._derive(self.relation.union(reordered._set_operand())) + + def _with_map_fields_of(self, first: "DataFrame") -> "DataFrame": + """This DataFrame as the second input of ``first.unionByName``, each map in it with ``first``'s field names. + + DuckDB's union pairs struct fields by name at any depth, where Spark's unionByName pairs those inside a map by + position, as its union does. + """ + columns = self._column_exprs() + converted = False + for i, (mine, theirs) in enumerate(zip(first.schema.fields, self.schema.fields, strict=True)): + target = _maps_named_as(mine.dataType, theirs.dataType) + if target != theirs.dataType: + columns[i] = _convert(columns[i], target.duckdb_type).alias(theirs.name) + converted = True + return self._derive(self.relation.select(*columns)) if converted else self def intersect(self, other: "DataFrame") -> "DataFrame": """Return a new :class:`DataFrame` containing rows only in @@ -1182,7 +2132,7 @@ def intersect(self, other: "DataFrame") -> "DataFrame": | a| 1| +---+---+ """ # noqa: D205 - return self.intersectAll(other).drop_duplicates() + return self._derive(self.relation.intersect(self._positional_operand(other, "INTERSECT"))).drop_duplicates() def intersectAll(self, other: "DataFrame") -> "DataFrame": """Return a new :class:`DataFrame` containing rows in both this :class:`DataFrame` @@ -1219,7 +2169,7 @@ def intersectAll(self, other: "DataFrame") -> "DataFrame": | b| 3| +---+---+ """ # noqa: D205 - return DataFrame(self.relation.intersect(other.relation), self.session) + return self._derive(self.relation.intersect(self._positional_operand(other, "INTERSECT ALL"))) def exceptAll(self, other: "DataFrame") -> "DataFrame": """Return a new :class:`DataFrame` containing rows in this :class:`DataFrame` but @@ -1259,7 +2209,7 @@ def exceptAll(self, other: "DataFrame") -> "DataFrame": +---+---+ """ # noqa: D205 - return DataFrame(self.relation.except_(other.relation), self.session) + return self._derive(self.relation.except_(self._positional_operand(other, "EXCEPT ALL"))) def dropDuplicates(self, subset: list[str] | None = None) -> "DataFrame": """Return a new :class:`DataFrame` with duplicate rows removed, @@ -1313,14 +2263,28 @@ def dropDuplicates(self, subset: list[str] | None = None) -> "DataFrame": |Alice| 5| 80| +-----+---+------+ """ # noqa: D205 - if subset: - rn_col = f"tmp_col_{uuid.uuid1().hex}" - subset_str = ", ".join([f'"{c}"' for c in subset]) - window_spec = f"OVER(PARTITION BY {subset_str}) AS {rn_col}" - df = DataFrame(self.relation.row_number(window_spec, "*"), self.session) - return df.filter(f"{rn_col} = 1").drop(rn_col) - - return self.distinct() + if subset is None: + return self.distinct() + for column in subset: + if not self._has_column(column): + raise AnalysisException( + error_class="UNRESOLVED_COLUMN_AMONG_FIELD_NAMES", + message_parameters={"colName": column, "fieldNames": ", ".join(self.columns)}, + ) + rn_col = f"tmp_col_{uuid.uuid1().hex}" + # Spark keeps one row of each group of equal subset values, so an empty subset keeps one row, and compares + # every column of a name in the subset, so the keys are positions. + keys = {_name_key(column) for column in subset} + for field in self.schema.fields: + if _name_key(field.name) in keys: + _check_comparable(field.name, field.dataType) + positions = ", ".join(f"#{i + 1}" for i, name in enumerate(self.columns) if _name_key(name) in keys) + partition = f"PARTITION BY {positions}" if positions else "" + window_spec = f"OVER({partition}) AS {rn_col}" + # A window over a join binds against its inputs, where the positions are those of the join's own columns. + relation = self.relation.select(*self._column_exprs()) if self._join_inputs is not None else self.relation + df = DataFrame(relation.row_number(window_spec, "*"), self.session, [*self.columns, rn_col]) + return df.filter(f"{rn_col} = 1").drop(rn_col) drop_duplicates = dropDuplicates @@ -1341,8 +2305,10 @@ def distinct(self) -> "DataFrame": >>> df.distinct().count() 2 """ + for field in self.schema.fields: + _check_comparable(field.name, field.dataType) distinct_rel = self.relation.distinct() - return DataFrame(distinct_rel, self.session) + return self._derive(distinct_rel) def count(self) -> int: """Returns the number of rows in this :class:`DataFrame`. @@ -1364,33 +2330,41 @@ def count(self) -> int: count_rel = self.relation.count("*") return int(count_rel.fetchall()[0][0]) - def _cast_types(self, *types) -> "DataFrame": - existing_columns = self.relation.columns - types_count = len(types) - assert types_count == len(existing_columns) - - cast_expressions = [ - f"{existing}::{target_type} as {existing}" - for existing, target_type in zip(existing_columns, types, strict=False) + def _cast_types(self, *types: str) -> "DataFrame": + assert len(types) == len(self.columns) + columns = [ + expr.cast(DuckDBPyType(target)).alias(name) + for expr, target, name in zip(self._column_exprs(), types, self.columns, strict=True) ] - cast_expressions = ", ".join(cast_expressions) - new_rel = self.relation.project(cast_expressions) - return DataFrame(new_rel, self.session) + return self._derive(self.relation.select(*columns)) def toDF(self, *cols) -> "DataFrame": # noqa: D102 - existing_columns = self.relation.columns - column_count = len(cols) - if column_count != len(existing_columns): + if len(cols) != len(self.columns): raise PySparkValueError(message="Provided column names and number of columns in the DataFrame don't match") - - existing_columns = [ColumnExpression(x) for x in existing_columns] - projections = [existing.alias(new) for existing, new in zip(existing_columns, cols, strict=False)] - new_rel = self.relation.project(*projections) - return DataFrame(new_rel, self.session) + names = [_column_name(new) for new in cols] + projections = [existing.alias(new) for existing, new in zip(self._column_exprs(), names, strict=True)] + return self._derive(self.relation.select(*projections)) + + def _fetchall(self, limit: int | None = None) -> list[tuple]: + """The rows, with each timestamp with a time zone as its microseconds since the epoch.""" + relation = self.relation + if _contains_type(self.schema, _INSTANT_TYPES): + relation = relation.select( + *( + _epoch_microseconds(expr, field.dataType).alias(field.name) + for expr, field in zip(self._column_exprs(), self.schema.fields, strict=True) + ) + ) + if limit is not None: + relation = relation.limit(limit) + return relation.fetchall() def collect(self) -> list[Row]: # noqa: D102 - columns = self.relation.columns - result = self.relation.fetchall() + columns = self.columns + result = self._fetchall() + converters = [_duckdb_to_python(field.dataType) for field in self.schema.fields] + if any(converters): + result = [[_apply(convert, v) for convert, v in zip(converters, row, strict=True)] for row in result] def construct_row(values: list, names: list[str]) -> Row: row = tuple.__new__(Row, list(values)) @@ -1427,8 +2401,210 @@ def cache(self) -> "DataFrame": == Physical Plan == InMemoryTableScan ... """ - cached_relation = self.relation.materialize() - return DataFrame(cached_relation, self.session) + # A materialized relation drops the join scope and renames duplicate columns, so such frames stay lazy. + if self._join_inputs is None and all(count == 1 for count in self._name_counts.values()) and self.columns: + self.relation = self.relation.materialize() + return self + + +def _is_reference(col: Column) -> bool: + return _reference(col) is not None + + +def _aliased_reference(col: Column) -> _Ref | None: + """The qualified column reference ``col`` is an alias of, such as ``df.a.alias("b")``, or None.""" + name = col._name + if name is None or len(name) != 3 or not isinstance(col._output, str) or not isinstance(name[0], _Ref): + return None + if name[1:] != (" AS ", _quoted_identifier(col._output)): + return None + ref = name[0] + return ref if isinstance(ref, _BoundRef) or len(ref.parts) > 1 else None + + +def _is_qualified_reference(col: Column) -> bool: + """Whether ``col`` is a column reference bound to a relation, such as ``df.a``, rather than a bare name.""" + if not _is_reference(col): + return False + parts = col._name[0].parts + return len(parts) > 1 or str(col.expr).casefold() != str(_bare_column(parts[0])).casefold() + + +def _ordinal(index: int) -> str: + """Spark's ordinal word in error messages for the position ``index``, counted from zero.""" + return ("first", "second", "third")[index] if index < 3 else f"{index + 1}th" + + +def _num_columns_mismatch(operator: str, first: int, second: int) -> AnalysisException: + return AnalysisException( + error_class="NUM_COLUMNS_MISMATCH", + message_parameters={ + "operator": operator, + "firstNumColumns": str(first), + "invalidOrdinalNum": "second", + "invalidNumColumns": str(second), + }, + ) + + +def _incompatible_column_type(operator: str, index: int, first: DataType, second: DataType) -> AnalysisException: + return AnalysisException( + error_class="INCOMPATIBLE_COLUMN_TYPE", + message_parameters={ + "operator": operator, + "columnOrdinalNumber": _ordinal(index), + "tableOrdinalNumber": "second", + "dataType1": f'"{_spark_sql_type(second)}"', + "dataType2": f'"{_spark_sql_type(first)}"', + "hint": "", + }, + ) + + +def _type_family(data_type: AtomicType) -> type[DataType]: + """The family of an atomic type; Spark's set operations widen a type to another of its family.""" + if isinstance(data_type, (TimeType, TimeNTZType, TimeNSType)): + return TimeType + if isinstance(data_type, NumericType): + return NumericType + if isinstance( + data_type, + ( + DateType, + TimestampType, + TimestampNanosecondType, + TimestampNTZType, + TimestampSecondNTZType, + TimestampMillisecondNTZType, + TimestampNanosecondNTZType, + ), + ): + return DateType + return type(data_type) + + +def _compatible(left: DataType, right: DataType, by_name: bool, allow_missing: bool) -> bool: + """Whether a set operation accepts ``right`` at the place of ``left``, as both Spark versions do, at any depth. + + Atomic types are compatible within a family, and a string with any atomic type but a boolean, which Spark 3.5.9 + widens to a string, and Spark 4.2.0 a binary to a binary, as DuckDB does. Struct fields pair by position, or by + name with ``by_name``, where ``allow_missing`` accepts a field that only one of the structs has; inside a map they + pair by position, and with ``by_name`` they need the same names. + """ + if left == right or isinstance(left, NullType) or isinstance(right, NullType): + return True + if isinstance(left, StructType) and isinstance(right, StructType): + if by_name: + fields = {_name_key(f.name): f for f in right.fields} + pairs = [(f, fields[_name_key(f.name)]) for f in left.fields if _name_key(f.name) in fields] + if not allow_missing and not len(pairs) == len(left.fields) == len(right.fields): + return False + elif len(left.fields) == len(right.fields): + pairs = list(zip(left.fields, right.fields, strict=True)) + else: + return False + return all(_compatible(a.dataType, b.dataType, by_name, allow_missing) for a, b in pairs) + if isinstance(left, ArrayType) and isinstance(right, ArrayType): + return _compatible(left.elementType, right.elementType, by_name, allow_missing) + if isinstance(left, MapType) and isinstance(right, MapType): + return all( + _compatible(a, b, by_name=False, allow_missing=False) and not (by_name and _renames_struct_fields(a, b)) + for a, b in ((left.keyType, right.keyType), (left.valueType, right.valueType)) + ) + if not isinstance(left, AtomicType) or not isinstance(right, AtomicType): + return False + families = {_type_family(left), _type_family(right)} + return len(families) == 1 or (StringType in families and BooleanType not in families) + + +def _check_struct_fields(left: DataType, right: DataType) -> None: + """Raise where a struct in ``left`` has a field the struct at its place in ``right`` lacks, as unionByName does. + + Spark pairs struct fields by name at any depth outside a map, and by position inside one. + """ + if isinstance(left, StructType) and isinstance(right, StructType): + fields = {_name_key(f.name): f for f in right.fields} + for field in left.fields: + if _name_key(field.name) not in fields: + raise AnalysisException( + error_class="FIELD_NOT_FOUND", + message_parameters={ + "fieldName": f"`{field.name}`", + "fields": ", ".join(f"`{f.name}`" for f in right.fields), + }, + ) + _check_struct_fields(field.dataType, fields[_name_key(field.name)].dataType) + elif isinstance(left, ArrayType) and isinstance(right, ArrayType): + _check_struct_fields(left.elementType, right.elementType) + + +def _maps_named_as(left: DataType, right: DataType) -> DataType: + """``right`` with each map in it given the type of the map at its place in ``left``, where they have one shape. + + Fields pair by name outside a map, at any depth, so only maps that differ in struct field names change. + """ + if isinstance(left, StructType) and isinstance(right, StructType): + fields = {_name_key(f.name): f for f in left.fields} + return StructType( + [ + StructField(f.name, _maps_named_as(fields[_name_key(f.name)].dataType, f.dataType), f.nullable) + if _name_key(f.name) in fields + else f + for f in right.fields + ] + ) + if isinstance(left, ArrayType) and isinstance(right, ArrayType): + return ArrayType(_maps_named_as(left.elementType, right.elementType), right.containsNull) + if isinstance(left, MapType) and isinstance(right, MapType) and _same_shape(left, right): + return left + return right + + +def _check_comparable(name: str, *data_types: DataType) -> None: + """Raise as Spark does where a column that a set operation or a deduplication compares has a map at any depth.""" + for data_type in data_types: + if _contains_type(data_type, (MapType,)): + raise AnalysisException( + error_class="UNSUPPORTED_FEATURE.SET_OPERATION_ON_MAP_TYPE", + message_parameters={"colName": f"`{name}`", "dataType": f'"{_spark_sql_type(data_type)}"'}, + ) + + +def _renames_struct_fields(left: DataType, right: DataType) -> bool: + """Whether a struct in ``right`` has other field names than the struct at its place in ``left``.""" + if isinstance(left, StructType) and isinstance(right, StructType): + return list(map(_name_key, left.names)) != list(map(_name_key, right.names)) or any( + _renames_struct_fields(a.dataType, b.dataType) for a, b in zip(left.fields, right.fields, strict=True) + ) + if isinstance(left, ArrayType) and isinstance(right, ArrayType): + return _renames_struct_fields(left.elementType, right.elementType) + if isinstance(left, MapType) and isinstance(right, MapType): + return _renames_struct_fields(left.keyType, right.keyType) or _renames_struct_fields( + left.valueType, right.valueType + ) + return False + + +def _same_shape(left: DataType, right: DataType) -> bool: + """Whether ``left`` and ``right`` are one type but for the names of struct fields.""" + if isinstance(left, StructType) and isinstance(right, StructType): + return len(left.fields) == len(right.fields) and all( + _same_shape(a.dataType, b.dataType) for a, b in zip(left.fields, right.fields, strict=True) + ) + if isinstance(left, ArrayType) and isinstance(right, ArrayType): + return _same_shape(left.elementType, right.elementType) + if isinstance(left, MapType) and isinstance(right, MapType): + return _same_shape(left.keyType, right.keyType) and _same_shape(left.valueType, right.valueType) + return left == right + + +def _fresh_alias() -> str: + return f"relation_{uuid.uuid4().hex}" + + +def _source_frame(relation: duckdb.DuckDBPyRelation, session: "SparkSession") -> DataFrame: + """A DataFrame reading a source under a name of its own, so a reference to it from another DataFrame fails.""" + return DataFrame(relation.set_alias(_fresh_alias()), session) __all__ = ["DataFrame"] diff --git a/duckdb/experimental/spark/sql/functions.py b/duckdb/experimental/spark/sql/functions.py index 71ff8c59d..a2fcc70aa 100644 --- a/duckdb/experimental/spark/sql/functions.py +++ b/duckdb/experimental/spark/sql/functions.py @@ -1,7 +1,13 @@ +import builtins +import functools +import inspect +import re import warnings from collections.abc import Callable +from decimal import Decimal from typing import TYPE_CHECKING, Any, Optional, Union, overload +import duckdb from duckdb import ( CaseExpression, CoalesceOperator, @@ -9,18 +15,40 @@ ConstantExpression, Expression, FunctionExpression, - LambdaExpression, SQLExpression, ) if TYPE_CHECKING: from .dataframe import DataFrame -from ..errors import PySparkTypeError +from duckdb.sqltypes import BIGINT, DOUBLE, INTEGER, VARCHAR, DuckDBPyType + +from ..errors import AnalysisException, IllegalArgumentException, PySparkTypeError, PySparkValueError from ..exception import ContributionsAcceptedError from . import types as _types from ._typing import ColumnOrName -from .column import Column, _get_expr +from .column import ( + _INTEGRAL_TYPE_IDS, + Column, + Name, + _call_name, + _column_reference, + _constant, + _evaluated_once, + _format_name, + _get_expr, + _is_numpy, + _lambda, + _literal_name, + _name_parts, + _numpy_literal, + _Ref, + _ref_name, + _spark_cast, + _sql_name, + _to_binary, + _truncate, +) def _invoke_function_over_columns(name: str, *cols: "ColumnOrName") -> Column: @@ -31,27 +59,20 @@ def _invoke_function_over_columns(name: str, *cols: "ColumnOrName") -> Column: return _invoke_function(name, *cols) -def _nan_constant() -> Expression: - """Create a NaN constant expression. - - Note: ConstantExpression(float("nan")) returns NULL instead of NaN because - TransformPythonValue() in the C++ layer has nan_as_null=true by default. - This is intentional for data import scenarios (CSV, Pandas, etc.) where NaN - represents missing data. +def _aggregate_column(expr: Expression, *inputs: Expression) -> Column: + column = Column(expr) + column._is_aggregate = True + column._aggregate_inputs = inputs + return column - For mathematical functions that need to return NaN (not NULL) for out-of-range - inputs per PySpark/IEEE 754 semantics, we use SQLExpression as a workaround. - Returns: - ------- - Expression - An expression that evaluates to NaN (not NULL) - """ - return SQLExpression("'NaN'::DOUBLE") +def _invoke_aggregate_over_columns(name: str, *cols: "ColumnOrName") -> Column: + inputs = [_to_column_expr(c) for c in cols] + return _aggregate_column(_invoke_function(name, *inputs).expr, *inputs) def col(column: str) -> Column: # noqa: D103 - return Column(ColumnExpression(column)) + return Column(_column_reference(column), _ref_name(column)) def upper(col: "ColumnOrName") -> Column: @@ -116,15 +137,28 @@ def when(condition: "Column", value: Column | str) -> Column: # noqa: D103 raise TypeError(msg) v = _get_expr(value) expr = CaseExpression(condition.expr, v) - return Column(expr) + return Column(expr, _sql_name("CASE WHEN ", condition, " THEN ", value, " END")) -def _inner_expr_or_val(val: Column | str) -> Column | str: - return val.expr if isinstance(val, Column) else val +def struct(*cols: Union["ColumnOrName", list["ColumnOrName"], set["ColumnOrName"]]) -> Column: # noqa: D103 + if len(cols) == 1 and isinstance(cols[0], (list, set)): + cols = cols[0] + columns = [_as_column(c) for c in cols] + names = [name if (name := c._struct_field()) is not None else f"col{i}" for i, c in enumerate(columns, 1)] + return Column(_struct_pack(names, [c.expr for c in columns])) + +def _struct_pack(names: list[str], values: list[Expression]) -> Expression: + """A struct of ``values`` with the field ``names``, which Spark allows to be empty. -def struct(*cols: Column) -> Column: # noqa: D103 - return Column(FunctionExpression("struct_pack", *[_inner_expr_or_val(x) for x in cols])) + DuckDB's unnamed struct has empty field names, but DuckDB has no struct with both empty and other field names. + """ + if names and not any(names): + return FunctionExpression("row", *values) + if not all(names): + msg = f"a struct with both empty and other field names is not supported: {names}" + raise ContributionsAcceptedError(msg) + return FunctionExpression("struct_pack", *(value.alias(name) for name, value in zip(names, values, strict=True))) def array(*cols: Union["ColumnOrName", list["ColumnOrName"] | tuple["ColumnOrName", ...]]) -> Column: @@ -164,18 +198,87 @@ def array(*cols: Union["ColumnOrName", list["ColumnOrName"] | tuple["ColumnOrNam def lit(col: Any) -> Column: # noqa: D103, ANN401 - return col if isinstance(col, Column) else Column(ConstantExpression(col)) + if isinstance(col, Column): + return col + expr, name, alias = _numpy_literal(col) if _is_numpy(col) else (_constant(col), (_literal_name(col),), None) + literal = Column(expr, name) + literal._is_constant = True + return literal if alias is None else literal.alias(alias) def _invoke_function(function: str, *arguments) -> Column: return Column(FunctionExpression(function, *arguments)) +def _result_type(column: Column, duckdb_type: DuckDBPyType) -> Column: + """Give a result the type Spark gives it, where DuckDB's function returns a wider or different type.""" + return Column(column.expr.cast(duckdb_type)) + + +def _rounded_to_integral(function: str, col: "ColumnOrName") -> Column: + """``function``, ceil or floor, as a BIGINT, which is the input itself for an integral input. + + DuckDB has no integral overload of either, so it would round a BIGINT through a DOUBLE. + """ + expr = _to_column_expr(col) + names = sorted(name.upper() for name in _INTEGRAL_TYPE_IDS) + is_integral = FunctionExpression("typeof", expr).isin(*map(ConstantExpression, names)) + return Column( + CaseExpression(is_integral, expr.cast(BIGINT)).otherwise(FunctionExpression(function, expr).cast(BIGINT)) + ) + + +def _null_if_null(result: Expression, *inputs: Expression) -> Column: + """``result``, or NULL where any of ``inputs`` is NULL, as Spark gives for a function's null intolerant inputs.""" + any_null = functools.reduce(lambda a, b: a | b, [i.isnull() for i in inputs]) + return Column(CaseExpression(any_null, ConstantExpression(None)).otherwise(result)) + + +def _is_null_literal(col: object) -> bool: + """Whether ``col`` is ``lit(None)``, which has Spark's NULL type and DuckDB's untyped NULL.""" + return isinstance(col, Column) and str(col.expr) == "NULL" + + +def _spark_error(message: str) -> Expression: + """Raise ``message`` when a row reaches this expression, for inputs Spark rejects at run time.""" + return FunctionExpression("error", ConstantExpression(message)) + + +def _distinct_in_order(arr: Expression) -> Expression: + """Remove repeats keeping first occurrences in order, as Spark does; DuckDB's list_distinct reorders, drops NULL. + + The list is evaluated once, so a nondeterministic element has one value in the list it is looked up in. + """ + + def distinct(values: Expression) -> Expression: + first = _lambda(lambda x, i: FunctionExpression("list_position", values, x) == i, values) + return FunctionExpression("list_filter", values, first) + + return _evaluated_once(arr, distinct) + + +def _in_common_type( + arr1: Expression, arr2: Expression, body: Callable[[Expression, Expression], Expression] +) -> Expression: + """``body`` of two lists with their elements converted to one type, the one DuckDB gives a list of both lists.""" + return _evaluated_once( + FunctionExpression("list_value", arr1, arr2), + lambda pair: body( + FunctionExpression("list_extract", pair, ConstantExpression(1)), + FunctionExpression("list_extract", pair, ConstantExpression(2)), + ), + ) + + +def _lit_or_column(x: Union["ColumnOrName", float]) -> "ColumnOrName": + return lit(x) if isinstance(x, (int, float)) else x + + def _to_column_expr(col: ColumnOrName) -> Expression: if isinstance(col, Column): return col.expr elif isinstance(col, str): - return ColumnExpression(col) + return _column_reference(col) else: raise PySparkTypeError( error_class="NOT_COLUMN_OR_STR", @@ -183,7 +286,11 @@ def _to_column_expr(col: ColumnOrName) -> Expression: ) -def regexp_replace(str: "ColumnOrName", pattern: str, replacement: str) -> Column: +def _as_column(col: ColumnOrName) -> Column: + return col if isinstance(col, Column) else Column(_to_column_expr(col), _ref_name(col)) + + +def regexp_replace(str: "ColumnOrName", pattern: str | Column, replacement: str | Column) -> Column: r"""Replace all substrings of the specified string value that match regexp with rep. .. versionadded:: 1.5.0 @@ -194,15 +301,45 @@ def regexp_replace(str: "ColumnOrName", pattern: str, replacement: str) -> Colum >>> df.select(regexp_replace("str", r"(\d+)", "--").alias("d")).collect() [Row(d='-----')] """ + literal = replacement if isinstance(replacement, builtins.str) else _string_literal(replacement) + if literal is None and not _is_null_literal(replacement): + msg = "Only a literal replacement, such as lit('$1'), is supported in regexp_replace; a column is not." + raise ContributionsAcceptedError(msg) return _invoke_function( "regexp_replace", _to_column_expr(str), - ConstantExpression(pattern), - ConstantExpression(replacement), + _get_expr(pattern), + ConstantExpression(None if literal is None else _java_replacement(literal)), ConstantExpression("g"), ) +_JAVA_REPLACEMENT_TOKEN = re.compile(r"(?s)\\(?P.?)|\$(?P[0-9]*)|(?P.)") + + +def _java_replacement(replacement: str) -> str: + r"""Translate Java's replacement syntax, where $1 is a group and a backslash escapes, to RE2's \1. + + Java reads a group number of two or more digits by the number of groups in the pattern, and RE2 refers to groups + 0 to 9 only, so such a reference raises. + """ + translated = [] + for token in _JAVA_REPLACEMENT_TOKEN.finditer(replacement): + escaped, digits = token["escaped"], token["group"] + if escaped is not None: + if not escaped: + raise PySparkValueError(message=f"character to be escaped is missing in replacement {replacement!r}") + translated.append("\\\\" if escaped == "\\" else escaped) + elif digits is not None: + if len(digits) != 1: + msg = f"only $0 to $9 group references are supported in replacement {replacement!r}" + raise ContributionsAcceptedError(msg) + translated.append(f"\\{digits}") + else: + translated.append(token["text"]) + return "".join(translated) + + def slice(x: "ColumnOrName", start: Union["ColumnOrName", int], length: Union["ColumnOrName", int]) -> Column: """Collection function: returns an array containing all the elements in `x` from index `start` (array indices start at 1, or from the end if `start` is negative) with the specified `length`. @@ -235,9 +372,29 @@ def slice(x: "ColumnOrName", start: Union["ColumnOrName", int], length: Union["C start = ConstantExpression(start) if isinstance(start, int) else _to_column_expr(start) length = ConstantExpression(length) if isinstance(length, int) else _to_column_expr(length) - end = start + length - - return _invoke_function("list_slice", _to_column_expr(x), start, end) + arr = _to_column_expr(x) + begin = CaseExpression(start < 0, FunctionExpression("len", arr) + start + 1).otherwise(start) + return _null_if_null( + CaseExpression( + start == 0, + _spark_error( + "[INVALID_PARAMETER_VALUE.START] The value of parameter(s) `start` in `slice` is invalid: " + "Expects a positive or a negative value for `start`, but got 0." + ), + ) + .when( + length < 0, + _spark_error( + "[INVALID_PARAMETER_VALUE.LENGTH] The value of parameter(s) `length` in `slice` is invalid: " + "Expects `length` greater than or equal to 0." + ), + ) + .when(begin < 1, FunctionExpression("list_slice", arr, 1, 0)) + .otherwise(FunctionExpression("list_slice", arr, begin, begin + length - 1)), + arr, + start, + length, + ) def asc_nulls_first(col: "ColumnOrName") -> Column: @@ -407,11 +564,12 @@ def left(str: "ColumnOrName", len: "ColumnOrName") -> Column: >>> df.select(left(df.a, df.b).alias("r")).collect() [Row(r='Spa')] """ # noqa: D205 - len = _to_column_expr(len) + # Spark casts the length to an integer, which truncates it, where DuckDB would round it. + str, len = _to_column_expr(str), _truncate(_to_column_expr(len), INTEGER) return Column( - CaseExpression(len <= ConstantExpression(0), ConstantExpression("")).otherwise( - FunctionExpression("array_slice", _to_column_expr(str), ConstantExpression(0), len) - ) + CaseExpression(str.isnull(), ConstantExpression(None)) + .when(len <= ConstantExpression(0), ConstantExpression("")) + .otherwise(FunctionExpression("array_slice", str, ConstantExpression(0), len)) ) @@ -442,11 +600,11 @@ def right(str: "ColumnOrName", len: "ColumnOrName") -> Column: >>> df.select(right(df.a, df.b).alias("r")).collect() [Row(r='SQL')] """ # noqa: D205 - len = _to_column_expr(len) + str, len = _to_column_expr(str), _truncate(_to_column_expr(len), INTEGER) return Column( - CaseExpression(len <= ConstantExpression(0), ConstantExpression("")).otherwise( - FunctionExpression("array_slice", _to_column_expr(str), -len, ConstantExpression(-1)) - ) + CaseExpression(str.isnull(), ConstantExpression(None)) + .when(len <= ConstantExpression(0), ConstantExpression("")) + .otherwise(FunctionExpression("array_slice", str, -len, ConstantExpression(-1))) ) @@ -492,17 +650,14 @@ def levenshtein(left: "ColumnOrName", right: "ColumnOrName", threshold: int | No >>> df0.select(levenshtein("l", "r", 2).alias("d")).collect() [Row(d=-1)] """ - distance = _invoke_function_over_columns("levenshtein", left, right) + distance = _result_type(_invoke_function_over_columns("levenshtein", left, right), INTEGER) if threshold is None: return distance - else: - distance = _to_column_expr(distance) - return Column( - CaseExpression(distance <= ConstantExpression(threshold), distance).otherwise(ConstantExpression(-1)) - ) + distance = _to_column_expr(distance) + return Column(CaseExpression(distance > ConstantExpression(threshold), ConstantExpression(-1)).otherwise(distance)) -def lpad(col: "ColumnOrName", len: int, pad: str) -> Column: +def lpad(col: "ColumnOrName", len: Column | int, pad: Column | str) -> Column: """Left-pad the string column to width `len` with `pad`. .. versionadded:: 1.5.0 @@ -535,10 +690,10 @@ def lpad(col: "ColumnOrName", len: int, pad: str) -> Column: >>> df.select(lpad(df.s, 6, "#").alias("s")).collect() [Row(s='##abcd')] """ - return _invoke_function("lpad", _to_column_expr(col), ConstantExpression(len), ConstantExpression(pad)) + return _pad("lpad", col, len, pad) -def rpad(col: "ColumnOrName", len: int, pad: str) -> Column: +def rpad(col: "ColumnOrName", len: Column | int, pad: Column | str) -> Column: """Right-pad the string column to width `len` with `pad`. .. versionadded:: 1.5.0 @@ -571,7 +726,17 @@ def rpad(col: "ColumnOrName", len: int, pad: str) -> Column: >>> df.select(rpad(df.s, 6, "#").alias("s")).collect() [Row(s='abcd##')] """ - return _invoke_function("rpad", _to_column_expr(col), ConstantExpression(len), ConstantExpression(pad)) + return _pad("rpad", col, len, pad) + + +def _pad(function: str, col: "ColumnOrName", length: Column | int, pad: Column | str) -> Column: + string, width, padding = _to_column_expr(col), _get_expr(length), _get_expr(pad) + # Spark truncates to the length with an empty pad, where DuckDB raises when the string is shorter; greatest would + # turn a NULL width into 0, and left keeps all but the last characters for a negative one. + non_negative = CaseExpression(width < ConstantExpression(0), ConstantExpression(0)).otherwise(width) + truncated = FunctionExpression("left", string, non_negative) + padded = FunctionExpression(function, string, width, padding) + return Column(CaseExpression(padding == ConstantExpression(""), truncated).otherwise(padded)) def ascii(col: "ColumnOrName") -> Column: @@ -639,7 +804,7 @@ def asin(col: "ColumnOrName") -> Column: col = _to_column_expr(col) # asin domain is [-1, 1]; return NaN for out-of-range values per PySpark semantics return Column( - CaseExpression((col < -1.0) | (col > 1.0), _nan_constant()).otherwise(FunctionExpression("asin", col)) + CaseExpression((col < -1.0) | (col > 1.0), _constant(float("nan"))).otherwise(FunctionExpression("asin", col)) ) @@ -732,7 +897,13 @@ def array_agg(col: "ColumnOrName") -> Column: >>> df.agg(array_agg("c").alias("r")).collect() [Row(r=[1, 1, 2])] """ - return _invoke_function_over_columns("list", col) + x = _to_column_expr(col) + values = FunctionExpression( + "list_filter", + FunctionExpression("list", x), + _lambda(lambda v: v.isnotnull()), + ) + return _aggregate_column(CoalesceOperator(values, FunctionExpression("list_value")), x) def collect_list(col: "ColumnOrName") -> Column: @@ -798,7 +969,9 @@ def array_append(col: "ColumnOrName", value: Column | str) -> Column: >>> df.select(array_append(df.c1, "x")).collect() [Row(array_append(c1, x)=['b', 'a', 'c', 'x'])] """ # noqa: D205 - return _invoke_function("list_append", _to_column_expr(col), _get_expr(value)) + arr = _to_column_expr(col) + appended = FunctionExpression("list_append", arr, _get_expr(value)) + return _null_if_null(appended, arr) def array_insert(arr: "ColumnOrName", pos: Union["ColumnOrName", int], value: Column | str) -> Column: @@ -838,59 +1011,38 @@ def array_insert(arr: "ColumnOrName", pos: Union["ColumnOrName", int], value: Co >>> df.select(array_insert(df.data, 5, "hello").alias("data")).collect() [Row(data=['a', 'b', 'c', None, 'hello']), Row(data=['c', 'b', 'a', None, 'hello'])] """ # noqa: D205 - pos = _get_expr(pos) + pos = ConstantExpression(pos) if isinstance(pos, int) else _to_column_expr(pos) arr = _to_column_expr(arr) - # Depending on if the position is positive or not, we need to interpret it differently. - # This is because negative numbers are relative to the end of the NEW list. - # For example, if it's -2, it's expected that the inserted value is at the second position - # in the NEW list. - pos_is_positive = pos > 0 - - list_length_plus_1 = FunctionExpression("add", FunctionExpression("len", arr), 1) - - # If the position is above the list size plus 1, first extend the list with the - # relevant number of nulls to get the list to the size of (pos - 1). - list_ = CaseExpression( - pos > list_length_plus_1, - FunctionExpression("list_resize", arr, FunctionExpression("subtract", pos, 1)), - ).otherwise( - CaseExpression( - (pos < 0) & (FunctionExpression("abs", pos) > list_length_plus_1), - FunctionExpression( - "list_concat", - FunctionExpression( - "list_resize", - FunctionExpression("list_value", None), - FunctionExpression("subtract", FunctionExpression("abs", pos), list_length_plus_1), - ), - arr, - ), - ).otherwise(arr) + value = FunctionExpression("list_value", _get_expr(value)) + length = FunctionExpression("len", arr) + # Position in the result; a negative position counts from the end, -1 appending. + index = CaseExpression(pos > 0, pos).otherwise(length + 2 + pos) + inserted = FunctionExpression( + "list_concat", + FunctionExpression("list_concat", FunctionExpression("list_slice", arr, 1, index - 1), value), + FunctionExpression("list_slice", arr, index, length), ) - - # We slice the array into two parts, insert the value in between and concatenate the two parts - # together again. - return _invoke_function( + nulls_after = FunctionExpression("list_concat", FunctionExpression("list_resize", arr, pos - 1), value) + nulls_before = FunctionExpression( "list_concat", FunctionExpression( - "list_concat", - # First part of the list - FunctionExpression( - "list_slice", - list_, - 1, - CaseExpression(pos_is_positive, FunctionExpression("subtract", pos, 1)).otherwise(pos), - ), - # Here we insert the value at the specified position - FunctionExpression("list_value", _get_expr(value)), - ), - # The remainder of the list - FunctionExpression( - "list_slice", - list_, - CaseExpression(pos_is_positive, pos).otherwise(FunctionExpression("add", pos, 1)), - -1, + "list_concat", value, FunctionExpression("list_resize", FunctionExpression("list_value", None), 1 - index) ), + arr, + ) + return _null_if_null( + CaseExpression( + pos == 0, + _spark_error( + "[INVALID_INDEX_OF_ZERO] The index 0 is invalid. " + "An index shall be either < 0 or > 0 (the first element has index 1)." + ), + ) + .when(index > length + 1, nulls_after) + .when(index < 1, nulls_before) + .otherwise(inserted), + arr, + pos, ) @@ -918,8 +1070,16 @@ def array_contains(col: "ColumnOrName", value: Column | str) -> Column: >>> df.select(array_contains(df.data, lit("a"))).collect() [Row(array_contains(data, a)=True), Row(array_contains(data, a)=False)] """ # noqa: D205 - value = _get_expr(value) - return _invoke_function("array_contains", _to_column_expr(col), value) + if value is None or _is_null_literal(value): + msg = ( + "[DATATYPE_MISMATCH.NULL_TYPE] Cannot resolve array_contains due to data type mismatch: " + "Null typed values cannot be used as arguments of `array_contains`." + ) + raise AnalysisException(msg) + arr = _to_column_expr(col) + found = FunctionExpression("list_contains", arr, _get_expr(value)) + # Spark answers NULL, not false, when the value is missing from an array that holds a NULL. + return Column(CaseExpression(_list_contains_null(arr) & ~found, ConstantExpression(None)).otherwise(found)) def array_distinct(col: "ColumnOrName") -> Column: @@ -946,7 +1106,7 @@ def array_distinct(col: "ColumnOrName") -> Column: >>> df.select(array_distinct(df.data)).collect() [Row(array_distinct(data)=[1, 2, 3]), Row(array_distinct(data)=[4, 5])] """ - return _invoke_function_over_columns("array_distinct", col) + return Column(_distinct_in_order(_to_column_expr(col))) def array_intersect(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: @@ -977,7 +1137,16 @@ def array_intersect(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: >>> df.select(array_intersect(df.c1, df.c2)).collect() [Row(array_intersect(c1, c2)=['a', 'c'])] """ # noqa: D205 - return _invoke_function_over_columns("array_intersect", col1, col2) + arr1, arr2 = _to_column_expr(col1), _to_column_expr(col2) + + def in_both(left: Expression, right: Expression) -> Expression: + return FunctionExpression( + "list_filter", + _distinct_in_order(left), + _lambda(lambda x: FunctionExpression("list_position", right, x).isnotnull(), right), + ) + + return _null_if_null(_in_common_type(arr1, arr2, in_both), arr1, arr2) def array_union(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: @@ -1008,7 +1177,7 @@ def array_union(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: >>> df.select(array_union(df.c1, df.c2)).collect() [Row(array_union(c1, c2)=['b', 'a', 'c', 'd', 'f'])] """ # noqa: D205 - return _invoke_function_over_columns("array_distinct", _invoke_function_over_columns("array_concat", col1, col2)) + return Column(_distinct_in_order(FunctionExpression("||", _to_column_expr(col1), _to_column_expr(col2)))) def array_max(col: "ColumnOrName") -> Column: @@ -1035,9 +1204,7 @@ def array_max(col: "ColumnOrName") -> Column: >>> df.select(array_max(df.data).alias("max")).collect() [Row(max=3), Row(max=10)] """ - return _invoke_function( - "array_extract", _to_column_expr(_invoke_function_over_columns("array_sort", col)), _get_expr(-1) - ) + return _invoke_function_over_columns("list_max", col) def array_min(col: "ColumnOrName") -> Column: @@ -1064,9 +1231,7 @@ def array_min(col: "ColumnOrName") -> Column: >>> df.select(array_min(df.data).alias("min")).collect() [Row(min=1), Row(min=-1)] """ - return _invoke_function( - "array_extract", _to_column_expr(_invoke_function_over_columns("array_sort", col)), _get_expr(1) - ) + return _invoke_function_over_columns("list_min", col) def avg(col: "ColumnOrName") -> Column: @@ -1097,7 +1262,7 @@ def avg(col: "ColumnOrName") -> Column: | 4.5| +-------+ """ - return _invoke_function_over_columns("avg", col) + return _invoke_aggregate_over_columns("avg", col) def sum(col: "ColumnOrName") -> Column: @@ -1128,7 +1293,7 @@ def sum(col: "ColumnOrName") -> Column: | 45| +-------+ """ - return _invoke_function_over_columns("sum", col) + return _invoke_aggregate_over_columns("sum", col) def max(col: "ColumnOrName") -> Column: @@ -1159,7 +1324,7 @@ def max(col: "ColumnOrName") -> Column: | 9| +-------+ """ - return _invoke_function_over_columns("max", col) + return _invoke_aggregate_over_columns("max", col) def mean(col: "ColumnOrName") -> Column: @@ -1191,7 +1356,7 @@ def mean(col: "ColumnOrName") -> Column: | 4.5| +-------+ """ # noqa: D205 - return _invoke_function_over_columns("mean", col) + return _invoke_aggregate_over_columns("mean", col) def median(col: "ColumnOrName") -> Column: @@ -1234,7 +1399,7 @@ def median(col: "ColumnOrName") -> Column: |dotNET| 10000.0| +------+----------------+ """ - return _invoke_function_over_columns("median", col) + return _invoke_aggregate_over_columns("median", col) def mode(col: "ColumnOrName") -> Column: @@ -1277,7 +1442,7 @@ def mode(col: "ColumnOrName") -> Column: |dotNET| 2012| +------+----------+ """ - return _invoke_function_over_columns("mode", col) + return _invoke_aggregate_over_columns("mode", col) def min(col: "ColumnOrName") -> Column: @@ -1308,7 +1473,7 @@ def min(col: "ColumnOrName") -> Column: | 0| +-------+ """ - return _invoke_function_over_columns("min", col) + return _invoke_aggregate_over_columns("min", col) def any_value(col: "ColumnOrName") -> Column: @@ -1338,7 +1503,7 @@ def any_value(col: "ColumnOrName") -> Column: >>> df.select(any_value("c1", True), any_value("c2", True)).collect() [Row(any_value(c1)='a', any_value(c2)=1)] """ - return _invoke_function_over_columns("any_value", col) + return _invoke_aggregate_over_columns("first", col) def count(col: "ColumnOrName") -> Column: @@ -1371,7 +1536,9 @@ def count(col: "ColumnOrName") -> Column: | 4| 3| +--------+----------------+ """ - return _invoke_function_over_columns("count", col) + if _is_star(col): + return _aggregate_column(FunctionExpression("count_star")) + return _invoke_aggregate_over_columns("count", col) def approx_count_distinct(col: "ColumnOrName", rsd: float | None = None) -> Column: @@ -1408,7 +1575,7 @@ def approx_count_distinct(col: "ColumnOrName", rsd: float | None = None) -> Colu if rsd is not None: msg = "rsd is not supported by DuckDB" raise ValueError(msg) - return _invoke_function_over_columns("approx_count_distinct", col) + return _invoke_aggregate_over_columns("approx_count_distinct", col) def approxCountDistinct(col: "ColumnOrName", rsd: float | None = None) -> Column: @@ -1602,7 +1769,7 @@ def ceil(col: "ColumnOrName") -> Column: | 0| +----------+ """ - return _invoke_function_over_columns("ceil", col) + return _rounded_to_integral("ceil", col) def ceiling(col: "ColumnOrName") -> Column: # noqa: D103 @@ -1637,7 +1804,7 @@ def floor(col: "ColumnOrName") -> Column: | 2| +----------+ """ - return _invoke_function_over_columns("floor", col) + return _rounded_to_integral("floor", col) def abs(col: "ColumnOrName") -> Column: @@ -1700,7 +1867,7 @@ def isnan(col: "ColumnOrName") -> Column: |NaN|2.0| true|false| +---+---+-----+-----+ """ - return _invoke_function_over_columns("isnan", col) + return Column(CoalesceOperator(FunctionExpression("isnan", _to_column_expr(col)), ConstantExpression(False))) def isnull(col: "ColumnOrName") -> Column: @@ -1857,9 +2024,7 @@ def array_compact(col: "ColumnOrName") -> Column: >>> df.select(array_compact(df.data)).collect() [Row(array_compact(data)=[1, 2, 3]), Row(array_compact(data)=[4, 5, 4])] """ - return _invoke_function( - "list_filter", _to_column_expr(col), LambdaExpression("x", ColumnExpression("x").isnotnull()) - ) + return _invoke_function("list_filter", _to_column_expr(col), _lambda(lambda x: x.isnotnull())) def array_remove(col: "ColumnOrName", element: Any) -> Column: # noqa: ANN401 @@ -1888,9 +2053,9 @@ def array_remove(col: "ColumnOrName", element: Any) -> Column: # noqa: ANN401 >>> df.select(array_remove(df.data, 1)).collect() [Row(array_remove(data, 1)=[2, 3]), Row(array_remove(data, 1)=[])] """ - return _invoke_function( - "list_filter", _to_column_expr(col), LambdaExpression("x", ColumnExpression("x") != ConstantExpression(element)) - ) + arr, removed = _to_column_expr(col), _get_expr(element) + kept = _lambda(lambda x: CoalesceOperator(x != removed, ConstantExpression(True)), removed) + return _null_if_null(FunctionExpression("list_filter", arr, kept), arr, removed) def last_day(date: "ColumnOrName") -> Column: @@ -2003,8 +2168,14 @@ def char(col: "ColumnOrName") -> Column: | A| +--------+ """ # noqa: D205 - col = _to_column_expr(col) - return Column(FunctionExpression("chr", CaseExpression(col > 256, col % 256).otherwise(col))) + return Column( + _evaluated_once( + _spark_cast(_as_column(col), BIGINT), + lambda n: CaseExpression(n < 0, ConstantExpression("")).otherwise( + FunctionExpression("chr", (n % 256).cast(INTEGER)) + ), + ) + ) def corr(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: @@ -2036,7 +2207,11 @@ def corr(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: >>> df.agg(corr("a", "b").alias("c")).collect() [Row(c=1.0)] """ # noqa: D205 - return _invoke_function_over_columns("corr", col1, col2) + x, y = _to_column_expr(col1), _to_column_expr(col2) + zero_variance = (FunctionExpression("regr_sxx", x, y) == 0) | (FunctionExpression("regr_syy", x, y) == 0) + return _aggregate_column( + CaseExpression(zero_variance, ConstantExpression(None)).otherwise(FunctionExpression("corr", x, y)), x, y + ) def cot(col: "ColumnOrName") -> Column: @@ -2109,7 +2284,7 @@ def negative(col: "ColumnOrName") -> Column: | -2| +------------+ """ # noqa: D205, D415 - return abs(col) * -1 + return Column(-_to_column_expr(col)) def pi() -> Column: @@ -2185,7 +2360,7 @@ def pow(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName", float]) >>> df.select(pow(lit(3), lit(2))).first() Row(POWER(3, 2)=9.0) """ - return _invoke_function_over_columns("pow", col1, col2) + return _invoke_function_over_columns("pow", _lit_or_column(col1), _lit_or_column(col2)) def printf(format: "ColumnOrName", *cols: "ColumnOrName") -> Column: @@ -2253,7 +2428,7 @@ def product(col: "ColumnOrName") -> Column: | 2| 80.0| +----+-------+ """ - return _invoke_function_over_columns("product", col) + return _invoke_aggregate_over_columns("product", col) def rand(seed: int | None = None) -> Column: @@ -2377,7 +2552,9 @@ def regexp_count(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: >>> df.select(regexp_count("str", col("regexp")).alias("d")).collect() [Row(d=3)] """ # noqa: D205 - return _invoke_function_over_columns("len", _invoke_function_over_columns("regexp_extract_all", str, regexp)) + return _result_type( + _invoke_function_over_columns("len", _invoke_function_over_columns("regexp_extract_all", str, regexp)), INTEGER + ) def regexp_extract(str: "ColumnOrName", pattern: str, idx: int) -> Column: @@ -2454,9 +2631,12 @@ def regexp_extract_all(str: "ColumnOrName", regexp: "ColumnOrName", idx: int | C """ # noqa: D205 if idx is None: idx = 1 - return _invoke_function( + matches = FunctionExpression( "regexp_extract_all", _to_column_expr(str), _to_column_expr(regexp), ConstantExpression(idx) ) + # A group that takes no part in a match is NULL in DuckDB and empty in Spark. + no_null = _lambda(lambda x: CoalesceOperator(x, ConstantExpression(""))) + return _invoke_function("list_transform", matches, no_null) def regexp_like(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: @@ -2615,10 +2795,66 @@ def sequence(start: "ColumnOrName", stop: "ColumnOrName", step: Optional["Column >>> df2.select(sequence("C1", "C2", "C3").alias("r")).collect() [Row(r=[4, 2, 0, -2, -4])] """ # noqa: D205 + start_expr, stop_expr = _to_column_expr(start), _to_column_expr(stop) + # Spark's elements have the common type of start and stop, where DuckDB generates bigints or timestamps. + elements = FunctionExpression("list_value", start_expr, stop_expr) if step is None: - return _invoke_function_over_columns("generate_series", start, stop) - else: - return _invoke_function_over_columns("generate_series", start, stop, step) + # Spark steps by one, or by a day for dates, which DuckDB adds to a date as an integer. Adding an integer zero + # to stop keeps the distance of a tinyint or smallint sequence from overflowing, and a date a date. + distance = (stop_expr + ConstantExpression(0).cast(INTEGER)) - start_expr + offsets = FunctionExpression("generate_series", ConstantExpression(0), FunctionExpression("abs", distance)) + steps = _lambda( + lambda k: start_expr + (FunctionExpression("sign", distance) * k).cast(INTEGER), start_expr, stop_expr + ) + series = FunctionExpression("list_transform", offsets, steps) + return Column(FunctionExpression("cast_to_type", series, elements)) + step_expr = _to_column_expr(step) + # Subtracting the step from itself gives a zero of its own type, which also works for interval steps. + zero = step_expr - step_expr + interval_step = FunctionExpression("typeof", zero) == ConstantExpression("INTERVAL") + # Spark widens integral elements to the step's type too, which the element type of start and stop leaves out. + widening_step = ~interval_step & ( + FunctionExpression("typeof", FunctionExpression("list_value", start_expr + zero, stop_expr + zero)) + != FunctionExpression("typeof", elements) + ) + first_step = start_expr + step_expr + # Only a step of whole days keeps the elements of a date sequence dates. + partial_day_step = interval_step & ( + FunctionExpression("cast_to_type", first_step, FunctionExpression("list_extract", elements, 1)) != first_step + ) + # DuckDB adds a step of months to each element in turn, where Spark adds multiples of it to start. + step_interval = SQLExpression(f"TRY_CAST({step_expr} AS INTERVAL)") + month_step = interval_step & ( + (FunctionExpression("date_part", ConstantExpression("year"), step_interval) != ConstantExpression(0)) + | (FunctionExpression("date_part", ConstantExpression("month"), step_interval) != ConstantExpression(0)) + ) + towards_stop = ((step_expr > zero) & (start_expr < stop_expr)) | ((step_expr < zero) & (start_expr > stop_expr)) + message = FunctionExpression( + "concat", + ConstantExpression("Illegal sequence boundaries: "), + start_expr, + ConstantExpression(" to "), + stop_expr, + ConstantExpression(" by "), + step_expr, + ) + series = FunctionExpression("generate_series", start_expr, stop_expr, step_expr) + return Column( + CaseExpression( + widening_step, + _spark_error("sequence with a step of a wider integral type than its start and stop is not supported"), + ) + .when((start_expr != stop_expr) & ~towards_stop, FunctionExpression("error", message)) + .when((step_expr == zero) & (start_expr == stop_expr), FunctionExpression("list_value", start_expr)) + .when( + partial_day_step, + _spark_error( + "sequence step must be an interval day to second of day granularity if start and end values are dates" + ), + ) + .when(month_step, _spark_error("sequence with a step of months is not supported")) + .otherwise(FunctionExpression("cast_to_type", series, elements)) + ) def sign(col: "ColumnOrName") -> Column: @@ -2649,7 +2885,10 @@ def sign(col: "ColumnOrName") -> Column: | -1.0| 1.0| +--------+-------+ """ - return _invoke_function_over_columns("sign", col) + x = _to_column_expr(col).cast(DOUBLE) + # DuckDB's sign of NaN is 0, Spark's is NaN. + signum = FunctionExpression("sign", x).cast(DOUBLE) + return Column(CaseExpression(FunctionExpression("isnan", x), x).otherwise(signum)) def signum(col: "ColumnOrName") -> Column: @@ -2735,7 +2974,20 @@ def skewness(col: "ColumnOrName") -> Column: >>> df.select(skewness(df.c)).first() Row(skewness(c)=0.70710...) """ - return _invoke_function_over_columns("skewness", col) + x = _to_column_expr(col) + # DuckDB's moments raise on a NaN or an infinity, which make Spark's NaN, so they see only finite values. + finite = FunctionExpression("isfinite", x.cast(DOUBLE)) + finite_x = CaseExpression(finite, x) + n = FunctionExpression("count", finite_x).cast(DOUBLE) + # DuckDB's skewness is the sample estimate, undefined below 3 values; Spark's is the population skewness. + population = FunctionExpression("skewness", finite_x) * (n - 2) / FunctionExpression("sqrt", n * (n - 1)) + two_distinct = (n == 2) & (FunctionExpression("var_pop", finite_x) > 0) + skew = ( + CaseExpression(FunctionExpression("bool_or", ~finite), _constant(float("nan"))) + .when(two_distinct, ConstantExpression(0.0)) + .otherwise(population) + ) + return _aggregate_column(skew, x) def encode(col: "ColumnOrName", charset: str) -> Column: @@ -2798,7 +3050,9 @@ def find_in_set(str: "ColumnOrName", str_array: "ColumnOrName") -> Column: str_array = _to_column_expr(str_array) str = _to_column_expr(str) return Column( - CaseExpression(FunctionExpression("contains", str, ConstantExpression(",")), 0).otherwise( + CaseExpression(str.isnull() | str_array.isnull(), ConstantExpression(None)) + .when(FunctionExpression("contains", str, ConstantExpression(",")), ConstantExpression(0)) + .otherwise( CoalesceOperator( FunctionExpression( "list_position", FunctionExpression("string_split", str_array, ConstantExpression(",")), str @@ -2862,7 +3116,7 @@ def first(col: "ColumnOrName", ignorenulls: bool = False) -> Column: """ if ignorenulls: raise ContributionsAcceptedError() - return _invoke_function_over_columns("first", col) + return _invoke_aggregate_over_columns("first", col) def last(col: "ColumnOrName", ignorenulls: bool = False) -> Column: @@ -2917,7 +3171,7 @@ def last(col: "ColumnOrName", ignorenulls: bool = False) -> Column: """ if ignorenulls: raise ContributionsAcceptedError() - return _invoke_function_over_columns("last", col) + return _invoke_aggregate_over_columns("last", col) def greatest(*cols: "ColumnOrName") -> Column: @@ -3250,7 +3504,7 @@ def length(col: "ColumnOrName") -> Column: >>> spark.createDataFrame([("ABC ",)], ["a"]).select(length("a").alias("length")).collect() [Row(length=4)] """ # noqa: D205 - return _invoke_function_over_columns("length", col) + return _result_type(_invoke_function_over_columns("length", col), INTEGER) def coalesce(*cols: "ColumnOrName") -> Column: @@ -3484,7 +3738,7 @@ def sha2(col: "ColumnOrName", numBits: int) -> Column: msg = "numBits should be one of {224, 256, 384, 512, 0}" raise ValueError(msg) - if numBits == 256: + if numBits in {0, 256}: return _invoke_function_over_columns("sha256", col) msg = "SHA-224, SHA-384, and SHA-512 are not supported yet." @@ -3652,6 +3906,25 @@ def asc(col: "ColumnOrName") -> Column: return Column(_to_column_expr(col).asc()) +_DATE_TRUNC_UNITS = { + "year": "year", + "yyyy": "year", + "yy": "year", + "quarter": "quarter", + "month": "month", + "mon": "month", + "mm": "month", + "week": "week", + "day": "day", + "dd": "day", + "hour": "hour", + "minute": "minute", + "second": "second", + "millisecond": "millisecond", + "microsecond": "microsecond", +} + + def date_trunc(format: str, timestamp: "ColumnOrName") -> Column: """Returns timestamp truncated to the unit specified by the format. @@ -3672,16 +3945,27 @@ def date_trunc(format: str, timestamp: "ColumnOrName") -> Column: >>> df.select(date_trunc("mon", df.t).alias("month")).collect() [Row(month=datetime.datetime(1997, 2, 1, 0, 0))] """ - format = format.lower() - if format in ["yyyy", "yy"]: - format = "year" - elif format in ["mon", "mm"]: - format = "month" - elif format == "dd": - format = "day" - elif format == "hh": - format = "hour" - return _invoke_function_over_columns("date_trunc", lit(format), timestamp) + unit = _DATE_TRUNC_UNITS.get(format.lower()) + if unit is None: + return lit(None).cast(_types.TimestampType()) + as_timestamp = Column(_to_column_expr(timestamp)).cast(_types.TimestampType()) + return _invoke_function_over_columns("date_trunc", lit(unit), as_timestamp) + + +_DATE_PART_FIELDS = { + **dict.fromkeys(["YEAR", "Y", "YEARS", "YR", "YRS"], "year"), + "YEAROFWEEK": "isoyear", + **dict.fromkeys(["QUARTER", "QTR"], "quarter"), + **dict.fromkeys(["MONTH", "MON", "MONS", "MONTHS"], "month"), + **dict.fromkeys(["WEEK", "W", "WEEKS"], "week"), + **dict.fromkeys(["DAY", "D", "DAYS"], "day"), + **dict.fromkeys(["DAYOFWEEK", "DOW"], "dayofweek"), + **dict.fromkeys(["DAYOFWEEK_ISO", "DOW_ISO"], "isodow"), + "DOY": "doy", + **dict.fromkeys(["HOUR", "H", "HOURS", "HR", "HRS"], "hour"), + **dict.fromkeys(["MINUTE", "M", "MIN", "MINS", "MINUTES"], "minute"), + **dict.fromkeys(["SECOND", "S", "SEC", "SECONDS", "SECS"], "second"), +} def date_part(field: "ColumnOrName", source: "ColumnOrName") -> Column: @@ -3716,7 +4000,19 @@ def date_part(field: "ColumnOrName", source: "ColumnOrName") -> Column: ... ).collect() [Row(year=2015, month=4, week=15, day=8, minute=8, second=Decimal('15.000000'))] """ - return _invoke_function_over_columns("date_part", field, source) + name = _string_literal(field) + part = None if name is None else _DATE_PART_FIELDS.get(name.upper()) + if part is None: + msg = ( + f"[INVALID_EXTRACT_FIELD] Cannot extract {field if name is None else name!r}; " + "the field must be a literal such as lit('YEAR')." + ) + raise AnalysisException(msg) + if part == "second": + microseconds = _invoke_function_over_columns("date_part", lit("microseconds"), source) + return (microseconds.cast(_types.DecimalType(8, 0)) * lit(Decimal("0.000001"))).cast(_types.DecimalType(8, 6)) + extracted = _invoke_function_over_columns("date_part", lit(part), source) + return _result_type(extracted + lit(1) if part == "dayofweek" else extracted, INTEGER) def extract(field: "ColumnOrName", source: "ColumnOrName") -> Column: @@ -3827,7 +4123,7 @@ def date_diff(end: "ColumnOrName", start: "ColumnOrName") -> Column: +----------+----------+-----------------+ >>> df.select( - ... "*", sf.date_diff(sf.col("d1").cast("DATE"), sf.col("d2").cast("DATE")) + ... "*", sf.date_diff(sf.col("d2").cast("DATE"), sf.col("d1").cast("DATE")) ... ).show() +----------+----------+-----------------+ | d1| d2|date_diff(d2, d1)| @@ -3835,7 +4131,7 @@ def date_diff(end: "ColumnOrName", start: "ColumnOrName") -> Column: |2015-04-08|2015-05-10| 32| +----------+----------+-----------------+ """ - return _invoke_function_over_columns("date_diff", lit("day"), end, start) + return _result_type(_invoke_function_over_columns("date_diff", lit("day"), start, end), INTEGER) def year(col: "ColumnOrName") -> Column: @@ -3862,7 +4158,7 @@ def year(col: "ColumnOrName") -> Column: >>> df.select(year("dt").alias("year")).collect() [Row(year=2015)] """ - return _invoke_function_over_columns("year", col) + return _result_type(_invoke_function_over_columns("year", col), INTEGER) def quarter(col: "ColumnOrName") -> Column: @@ -3889,7 +4185,7 @@ def quarter(col: "ColumnOrName") -> Column: >>> df.select(quarter("dt").alias("quarter")).collect() [Row(quarter=2)] """ - return _invoke_function_over_columns("quarter", col) + return _result_type(_invoke_function_over_columns("quarter", col), INTEGER) def month(col: "ColumnOrName") -> Column: @@ -3916,7 +4212,7 @@ def month(col: "ColumnOrName") -> Column: >>> df.select(month("dt").alias("month")).collect() [Row(month=4)] """ - return _invoke_function_over_columns("month", col) + return _result_type(_invoke_function_over_columns("month", col), INTEGER) def dayofweek(col: "ColumnOrName") -> Column: @@ -3944,7 +4240,7 @@ def dayofweek(col: "ColumnOrName") -> Column: >>> df.select(dayofweek("dt").alias("day")).collect() [Row(day=4)] """ # noqa: D205 - return _invoke_function_over_columns("dayofweek", col) + lit(1) + return _result_type(_invoke_function_over_columns("dayofweek", col) + lit(1), INTEGER) def day(col: "ColumnOrName") -> Column: @@ -3968,7 +4264,7 @@ def day(col: "ColumnOrName") -> Column: >>> df.select(day("dt").alias("day")).collect() [Row(day=8)] """ - return _invoke_function_over_columns("day", col) + return _result_type(_invoke_function_over_columns("day", col), INTEGER) def dayofmonth(col: "ColumnOrName") -> Column: @@ -4022,7 +4318,7 @@ def dayofyear(col: "ColumnOrName") -> Column: >>> df.select(dayofyear("dt").alias("day")).collect() [Row(day=98)] """ - return _invoke_function_over_columns("dayofyear", col) + return _result_type(_invoke_function_over_columns("dayofyear", col), INTEGER) def hour(col: "ColumnOrName") -> Column: @@ -4050,7 +4346,7 @@ def hour(col: "ColumnOrName") -> Column: >>> df.select(hour("ts").alias("hour")).collect() [Row(hour=13)] """ - return _invoke_function_over_columns("hour", col) + return _result_type(_invoke_function_over_columns("hour", col), INTEGER) def minute(col: "ColumnOrName") -> Column: @@ -4078,7 +4374,7 @@ def minute(col: "ColumnOrName") -> Column: >>> df.select(minute("ts").alias("minute")).collect() [Row(minute=8)] """ - return _invoke_function_over_columns("minute", col) + return _result_type(_invoke_function_over_columns("minute", col), INTEGER) def second(col: "ColumnOrName") -> Column: @@ -4106,7 +4402,7 @@ def second(col: "ColumnOrName") -> Column: >>> df.select(second("ts").alias("second")).collect() [Row(second=15)] """ - return _invoke_function_over_columns("second", col) + return _result_type(_invoke_function_over_columns("second", col), INTEGER) def weekofyear(col: "ColumnOrName") -> Column: @@ -4135,7 +4431,7 @@ def weekofyear(col: "ColumnOrName") -> Column: >>> df.select(weekofyear(df.dt).alias("week")).collect() [Row(week=15)] """ # noqa: D205 - return _invoke_function_over_columns("weekofyear", col) + return _result_type(_invoke_function_over_columns("weekofyear", col), INTEGER) def cos(col: "ColumnOrName") -> Column: @@ -4198,7 +4494,7 @@ def acos(col: "ColumnOrName") -> Column: col = _to_column_expr(col) # acos domain is [-1, 1]; return NaN for out-of-range values per PySpark semantics return Column( - CaseExpression((col < -1.0) | (col > 1.0), _nan_constant()).otherwise(FunctionExpression("acos", col)) + CaseExpression((col < -1.0) | (col > 1.0), _constant(float("nan"))).otherwise(FunctionExpression("acos", col)) ) @@ -4265,9 +4561,18 @@ def call_function(funcName: str, *cols: "ColumnOrName") -> Column: | 102.0| +------------------------------------+ """ + if funcName.lower() in _aggregate_function_names(): + return _invoke_aggregate_over_columns(funcName, *cols) return _invoke_function_over_columns(funcName, *cols) +@functools.cache +def _aggregate_function_names() -> frozenset[str]: + with duckdb.connect() as connection: + query = "SELECT DISTINCT function_name FROM duckdb_functions() WHERE function_type = 'aggregate'" + return frozenset(name for (name,) in connection.sql(query).fetchall()) + + def covar_pop(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: """Returns a new :class:`~pyspark.sql.Column` for the population covariance of ``col1`` and ``col2``. @@ -4297,7 +4602,7 @@ def covar_pop(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: >>> df.agg(covar_pop("a", "b").alias("c")).collect() [Row(c=0.0)] """ # noqa: D205 - return _invoke_function_over_columns("covar_pop", col1, col2) + return _invoke_aggregate_over_columns("covar_pop", col1, col2) def covar_samp(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: @@ -4329,7 +4634,7 @@ def covar_samp(col1: "ColumnOrName", col2: "ColumnOrName") -> Column: >>> df.agg(covar_samp("a", "b").alias("c")).collect() [Row(c=0.0)] """ # noqa: D205 - return _invoke_function_over_columns("covar_samp", col1, col2) + return _invoke_aggregate_over_columns("covar_samp", col1, col2) def exp(col: "ColumnOrName") -> Column: @@ -4387,7 +4692,15 @@ def factorial(col: "ColumnOrName") -> Column: >>> df.select(factorial(df.n).alias("f")).collect() [Row(f=120)] """ - return _invoke_function_over_columns("factorial", col) + return Column( + _evaluated_once( + _spark_cast(_as_column(col), INTEGER), + lambda n: CaseExpression( + n.between(ConstantExpression(0), ConstantExpression(20)), + FunctionExpression("factorial", n).cast(BIGINT), + ), + ) + ) def log2(col: "ColumnOrName") -> Column: @@ -4418,7 +4731,7 @@ def log2(col: "ColumnOrName") -> Column: | 2.0| +----+ """ - return _invoke_function_over_columns("log2", col) + return _log("log2", col) def ln(col: "ColumnOrName") -> Column: @@ -4446,7 +4759,13 @@ def ln(col: "ColumnOrName") -> Column: |1.3862943611198906| +------------------+ """ - return _invoke_function_over_columns("ln", col) + return _log("ln", col) + + +def _log(name: str, col: "ColumnOrName") -> Column: + """Spark's logarithms are NULL for values at or below zero, where DuckDB gives -inf or NaN.""" + x = _to_column_expr(col) + return Column(CaseExpression(x <= 0, ConstantExpression(None)).otherwise(FunctionExpression(name, x))) def degrees(col: "ColumnOrName") -> Column: @@ -4565,13 +4884,7 @@ def atan2(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName", float] >>> df.select(atan2(lit(1), lit(2))).first() Row(ATAN2(1, 2)=0.46364...) """ - - def lit_or_column(x: Union["ColumnOrName", float]) -> Column: - if isinstance(x, (int, float)): - return lit(x) - return x - - return _invoke_function_over_columns("atan2", lit_or_column(col1), lit_or_column(col2)) + return _invoke_function_over_columns("atan2", _lit_or_column(col1), _lit_or_column(col2)) def tan(col: "ColumnOrName") -> Column: @@ -4727,10 +5040,9 @@ def get(col: "ColumnOrName", index: Union["ColumnOrName", int]) -> Column: +----------------------+ """ # noqa: D205 index = ConstantExpression(index) if isinstance(index, int) else _to_column_expr(index) - # Spark uses 0-indexing, DuckDB 1-indexing - index = index + 1 - - return _invoke_function("list_extract", _to_column_expr(col), index) + # Spark counts from 0 and gives NULL below it, where DuckDB counts from 1 and counts a negative index from the end. + element = FunctionExpression("list_extract", _to_column_expr(col), index + 1) + return Column(CaseExpression(index < ConstantExpression(0), ConstantExpression(None)).otherwise(element)) def initcap(col: "ColumnOrName") -> Column: @@ -4762,16 +5074,12 @@ def initcap(col: "ColumnOrName") -> Column: FunctionExpression( "list_transform", FunctionExpression("string_split", _to_column_expr(col), ConstantExpression(" ")), - LambdaExpression( - "x", - FunctionExpression( + _lambda( + lambda x: FunctionExpression( "concat", - FunctionExpression( - "upper", - FunctionExpression("array_extract", ColumnExpression("x"), 1), - ), - FunctionExpression("array_slice", ColumnExpression("x"), 2, -1), - ), + FunctionExpression("upper", FunctionExpression("array_extract", x, 1)), + FunctionExpression("lower", FunctionExpression("array_slice", x, 2, -1)), + ) ), ), ConstantExpression(" "), @@ -4804,7 +5112,7 @@ def octet_length(col: "ColumnOrName") -> Column: ... .select(octet_length('cat')).collect() [Row(octet_length(cat)=3), Row(octet_length(cat)=4)] """ - return _invoke_function("octet_length", _to_column_expr(col).cast("blob")) + return _result_type(_invoke_function("octet_length", _to_binary(_to_column_expr(col))), INTEGER) def hex(col: "ColumnOrName") -> Column: @@ -4895,9 +5203,7 @@ def base64(col: "ColumnOrName") -> Column: |UGFuZGFzIEFQSQ==| +----------------+ """ - if isinstance(col, str): - col = Column(ColumnExpression(col)) - return _invoke_function_over_columns("base64", col.cast("BLOB")) + return _invoke_function("base64", _to_binary(_to_column_expr(col))) def unbase64(col: "ColumnOrName") -> Column: @@ -5005,13 +5311,13 @@ def array_join(col: "ColumnOrName", delimiter: str, null_replacement: str | None col = FunctionExpression( "list_transform", col, - LambdaExpression( - "x", - CaseExpression(ColumnExpression("x").isnull(), ConstantExpression(null_replacement)).otherwise( - ColumnExpression("x") - ), + _lambda( + lambda x: CaseExpression(x.isnull(), ConstantExpression(null_replacement)).otherwise(x.cast(VARCHAR)) ), ) + else: + # array_to_string skips NULL elements too, but gives NULL where no element remains. + col = FunctionExpression("list_filter", col, _lambda(lambda x: x.isnotnull())) return _invoke_function("array_to_string", col, ConstantExpression(delimiter)) @@ -5047,11 +5353,9 @@ def array_position(col: "ColumnOrName", value: Any) -> Column: # noqa: ANN401 >>> df.select(array_position(df.data, "a")).collect() [Row(array_position(data, a)=3), Row(array_position(data, a)=0)] """ # noqa: D205 - return Column( - CoalesceOperator( - _to_column_expr(_invoke_function_over_columns("list_position", col, lit(value))), ConstantExpression(0) - ) - ) + arr, element = _to_column_expr(col), _get_expr(value) + position = CoalesceOperator(FunctionExpression("list_position", arr, element), ConstantExpression(0)) + return _null_if_null(position.cast(BIGINT), arr, element) def array_prepend(col: "ColumnOrName", value: Any) -> Column: # noqa: ANN401 @@ -5079,7 +5383,9 @@ def array_prepend(col: "ColumnOrName", value: Any) -> Column: # noqa: ANN401 >>> df.select(array_prepend(df.data, 1)).collect() [Row(array_prepend(data, 1)=[1, 2, 3, 4]), Row(array_prepend(data, 1)=[1])] """ # noqa: D205 - return _invoke_function_over_columns("list_prepend", lit(value), col) + arr = _to_column_expr(col) + prepended = FunctionExpression("list_prepend", _get_expr(value), arr) + return _null_if_null(prepended, arr) def array_repeat(col: "ColumnOrName", count: Union["ColumnOrName", int]) -> Column: @@ -5108,9 +5414,11 @@ def array_repeat(col: "ColumnOrName", count: Union["ColumnOrName", int]) -> Colu >>> df.select(array_repeat(df.data, 3).alias("r")).collect() [Row(r=['ab', 'ab', 'ab'])] """ - count = lit(count) if isinstance(count, int) else count - - return _invoke_function_over_columns("list_resize", _invoke_function_over_columns("list_value", col), count, col) + element = _to_column_expr(col) + count = ConstantExpression(count) if isinstance(count, int) else _to_column_expr(count) + size = FunctionExpression("greatest", count, ConstantExpression(0)) + repeated = FunctionExpression("list_resize", FunctionExpression("list_value", element), size, element) + return _null_if_null(repeated, count) def array_size(col: "ColumnOrName") -> Column: @@ -5134,7 +5442,7 @@ def array_size(col: "ColumnOrName") -> Column: >>> df.select(array_size(df.data).alias("r")).collect() [Row(r=3), Row(r=None)] """ - return _invoke_function_over_columns("len", col) + return _result_type(_invoke_function_over_columns("len", col), INTEGER) def array_sort(col: "ColumnOrName", comparator: Callable[[Column, Column], Column] | None = None) -> Column: @@ -5276,7 +5584,7 @@ def split(str: "ColumnOrName", pattern: str, limit: int = -1) -> Column: >>> df.select(split(df.s, "[ABC]", -1).alias("s")).collect() [Row(s=['one', 'two', 'three', ''])] """ - if limit > 0: + if isinstance(limit, Column) or limit > 0: # Unclear how to implement this in DuckDB as we'd need to map back from the split array # to the original array which is tricky with regular expressions. msg = "limit is not yet supported" @@ -5322,10 +5630,22 @@ def split_part(src: "ColumnOrName", delimiter: "ColumnOrName", partNum: "ColumnO partNum = _to_column_expr(partNum) part = FunctionExpression("split_part", src, delimiter, partNum) - return Column( - CaseExpression(src.isnull() | delimiter.isnull() | partNum.isnull(), ConstantExpression(None)).otherwise( - CaseExpression(delimiter == ConstantExpression(""), ConstantExpression("")).otherwise(part) + unsplit = CaseExpression(partNum.isin(ConstantExpression(1), ConstantExpression(-1)), src).otherwise( + ConstantExpression("") + ) + return _null_if_null( + CaseExpression( + partNum == ConstantExpression(0), + _spark_error( + "[INVALID_INDEX_OF_ZERO] The index 0 is invalid. " + "An index shall be either < 0 or > 0 (the first element has index 1)." + ), ) + .when(delimiter == ConstantExpression(""), unsplit) + .otherwise(part), + src, + delimiter, + partNum, ) @@ -5358,7 +5678,7 @@ def stddev_samp(col: "ColumnOrName") -> Column: |1.8708286933869...| +------------------+ """ # noqa: D205 - return _invoke_function_over_columns("stddev_samp", col) + return _invoke_aggregate_over_columns("stddev_samp", col) def stddev(col: "ColumnOrName") -> Column: @@ -5449,7 +5769,7 @@ def stddev_pop(col: "ColumnOrName") -> Column: |1.707825127659...| +-----------------+ """ # noqa: D205 - return _invoke_function_over_columns("stddev_pop", col) + return _invoke_aggregate_over_columns("stddev_pop", col) def var_pop(col: "ColumnOrName") -> Column: @@ -5476,7 +5796,7 @@ def var_pop(col: "ColumnOrName") -> Column: >>> df.select(var_pop(df.id)).first() Row(var_pop(id)=2.91666...) """ - return _invoke_function_over_columns("var_pop", col) + return _invoke_aggregate_over_columns("var_pop", col) def var_samp(col: "ColumnOrName") -> Column: @@ -5508,7 +5828,7 @@ def var_samp(col: "ColumnOrName") -> Column: | 3.5| +------------+ """ # noqa: D205 - return _invoke_function_over_columns("var_samp", col) + return _invoke_aggregate_over_columns("var_samp", col) def variance(col: "ColumnOrName") -> Column: @@ -5567,7 +5887,7 @@ def weekday(col: "ColumnOrName") -> Column: | 2| +---+ """ - return _invoke_function_over_columns("isodow", col) - 1 + return _result_type(_invoke_function_over_columns("isodow", col) - 1, INTEGER) def zeroifnull(col: "ColumnOrName") -> Column: @@ -5593,12 +5913,197 @@ def zeroifnull(col: "ColumnOrName") -> Column: return coalesce(col, lit(0)) -def _to_date_or_timestamp(col: "ColumnOrName", spark_datatype: _types.DataType, format: str | None = None) -> Column: - if format is not None: - raise ContributionsAcceptedError( - "format is not yet supported as DuckDB and PySpark use a different way of specifying them." - + " As a workaround, you can use F.call_function('strptime', col, format)" +_YEAR_PATTERN_FIELDS = {1: "variable", 3: "variable", 4: "fixed"} +_NUMERIC_PATTERN_FIELDS = { + "M": ("month", "%m"), + "L": ("month", "%m"), + "d": ("day", "%d"), + "H": ("hour", "%H"), + "h": ("hour12", "%I"), + "m": ("minute", "%M"), + "s": ("second", "%S"), +} +_UNSUPPORTED_PATTERN_LETTERS = { + "G": "era", + "D": "day of year", + "K": "hour of am-pm (0-11)", + "k": "clock hour of day (1-24)", + "V": "time zone ID", + "z": "time zone name", + "O": "localized zone offset", + "X": "zone offset", + "x": "zone offset", + "Z": "zone offset", +} + + +def _datetime_pattern_error(pattern: str, reason: str) -> IllegalArgumentException: + return IllegalArgumentException(f"Invalid datetime pattern {pattern!r}: {reason}") + + +def _unsupported_datetime_pattern(pattern: str, reason: str) -> ContributionsAcceptedError: + return ContributionsAcceptedError(f"Datetime pattern {pattern!r} is not supported for parsing: {reason}") + + +def _tokenize_datetime_pattern(pattern: str) -> list[tuple[str, int] | str]: + """Split a Spark datetime pattern into (letter, count) fields and literal strings.""" + tokens: list[tuple[str, int] | str] = [] + i = 0 + while i < len(pattern): + char = pattern[i] + if char == "'": + if pattern.startswith("''", i): + tokens.append("'") + i += 2 + continue + text = [] + i += 1 + while True: + if i == len(pattern): + raise _datetime_pattern_error(pattern, "it ends with an incomplete quoted literal") + if pattern.startswith("''", i): + text.append("'") + i += 2 + elif pattern[i] == "'": + i += 1 + break + else: + text.append(pattern[i]) + i += 1 + tokens.append("".join(text)) + elif char.isascii() and char.isalpha(): + end = i + while end < len(pattern) and pattern[end] == char: + end += 1 + tokens.append((char, end - i)) + i = end + elif char in "#{}": + raise _datetime_pattern_error(pattern, f"{char!r} is reserved") + elif char in "[]": + raise _unsupported_datetime_pattern(pattern, "optional sections") + else: + tokens.append(char) + i += 1 + return tokens + + +def _translate_datetime_pattern_field(pattern: str, letter: str, count: int) -> tuple[str, str, str]: + """Translate one pattern field to (field, DuckDB strptime code, kind). + + The kind is "fixed" or "variable" for digits of a width Spark fixes or not, "fraction" or "text". + """ + if letter in _UNSUPPORTED_PATTERN_LETTERS: + raise _unsupported_datetime_pattern(pattern, f"{letter!r} ({_UNSUPPORTED_PATTERN_LETTERS[letter]})") + if letter == "y": + if count > 6: + raise _datetime_pattern_error(pattern, "too many 'y' letters") + if count == 2: + reason = "'yy', which DuckDB reads as the years 1969 to 2068 and Spark as 2000 to 2099" + raise _unsupported_datetime_pattern(pattern, reason) + if count not in _YEAR_PATTERN_FIELDS: + raise _unsupported_datetime_pattern(pattern, f"{'y' * count!r}, as DuckDB reads at most four year digits") + return "year", "%Y", _YEAR_PATTERN_FIELDS[count] + if letter in "ML" and count in (3, 4): + return "month", "%b" if count == 3 else "%B", "text" + if letter in _NUMERIC_PATTERN_FIELDS: + if count > 2: + raise _datetime_pattern_error(pattern, f"too many {letter!r} letters") + field, code = _NUMERIC_PATTERN_FIELDS[letter] + return field, code, "variable" if count == 1 else "fixed" + if letter == "S": + if count > 9: + raise _datetime_pattern_error(pattern, "too many 'S' letters") + if count > 6: + raise _unsupported_datetime_pattern(pattern, "nanoseconds, which DuckDB rounds where Spark truncates them") + return "fraction", "%g" if count <= 3 else "%f", "fraction" + if letter == "a": + if count > 1: + raise _datetime_pattern_error(pattern, "too many 'a' letters") + return "ampm", "%p", "text" + if letter in "EFQq": + raise _datetime_pattern_error(pattern, f"{letter!r} can only be used for formatting") + if letter in "YWceuw": + raise _datetime_pattern_error(pattern, f"{letter!r} is a week-based field, which Spark does not support") + raise _datetime_pattern_error(pattern, f"{letter!r} is not a pattern letter; quote literal text, as in 'T'") + + +def _translate_datetime_pattern(pattern: str) -> str: + """Translate a Spark datetime pattern for parsing into the DuckDB strptime format of the same fields. + + DuckDB's strptime accepts fewer digits and more spaces than Spark; a pattern for which it would read text that + Spark accepts to another value raises. + """ + duckdb_format, fields = [], [] + previous_kind = "text" + for token in _tokenize_datetime_pattern(pattern): + if isinstance(token, str): + if previous_kind in ("variable", "fraction") and token[:1].isdigit(): + raise _unsupported_datetime_pattern(pattern, "a variable width field followed by a digit") + duckdb_format.append(token.replace("%", "%%")) + previous_kind = "text" + continue + field, code, kind = _translate_datetime_pattern_field(pattern, *token) + # DuckDB reads digits greedily; Spark fails on a variable width field next to another number. + follows_variable = previous_kind in ("variable", "fraction") and kind != "text" + if follows_variable or (previous_kind != "text" and kind == "variable"): + raise _unsupported_datetime_pattern(pattern, "a variable width field next to another numeric field") + if field in fields: + raise _unsupported_datetime_pattern(pattern, f"the field of {token[0]!r} appears more than once") + fields.append(field) + duckdb_format.append(code) + previous_kind = kind + if "hour" in fields and ("hour12" in fields or "ampm" in fields): + raise _unsupported_datetime_pattern(pattern, "'H' combined with 'h' or 'a'") + if "year" not in fields: + raise _unsupported_datetime_pattern(pattern, "no year, where DuckDB reads the year 1900 and Spark 1970") + if "hour12" in fields and "ampm" not in fields: + reason = "'h' without 'a', where DuckDB reads 12 as noon and Spark as midnight" + raise _unsupported_datetime_pattern(pattern, reason) + return "".join(duckdb_format) + + +def _string_literal(col: "ColumnOrName") -> str | None: + """Return the string behind a literal column such as ``lit("yyyy-MM-dd")``, or None for any other column.""" + rendered = str(col.expr) if isinstance(col, Column) else "" + if len(rendered) < 2 or rendered[0] != "'" or rendered[-1] != "'" or "'" in rendered[1:-1].replace("''", ""): + return None + return rendered[1:-1].replace("''", "'") + + +def _datetime_pattern_argument(format: Optional["ColumnOrName"]) -> str | None: + if format is None: + return None + pattern = _string_literal(format) + if pattern is None: + msg = ( + "Only a literal datetime pattern, such as lit('yyyy-MM-dd'), is supported as format; " + "a column, or a str, which names a column, is not." ) + raise ContributionsAcceptedError(msg) + return pattern + + +def _parse_datetime(col: "ColumnOrName", pattern: str, *, raise_on_error: bool) -> Column: + source = _to_column_expr(col) + duckdb_format = ConstantExpression(_translate_datetime_pattern(pattern)) + parsed = FunctionExpression("strptime", source, duckdb_format) + if raise_on_error: + return Column(parsed) + # DuckDB's try_strptime reads the words epoch and infinity as 1900-01-01, where strptime reads what they name. + return Column(CaseExpression(FunctionExpression("try_strptime", source, duckdb_format).isnotnull(), parsed)) + + +def _to_date_or_timestamp( + col: "ColumnOrName", spark_datatype: _types.DataType, format: str | None = None, *, try_cast: bool = False +) -> Column: + if format is not None: + if not isinstance(format, str): + raise PySparkTypeError( + error_class="NOT_STR", message_parameters={"arg_name": "format", "arg_type": type(format).__name__} + ) + return _parse_datetime(col, format, raise_on_error=not try_cast).cast(spark_datatype) + if try_cast: + return Column(SQLExpression(f"TRY_CAST({_to_column_expr(col)} AS {spark_datatype.duckdb_type})")) return Column(_to_column_expr(col)).cast(spark_datatype) @@ -5675,7 +6180,7 @@ def to_timestamp(col: "ColumnOrName", format: str | None = None) -> Column: >>> df.select(to_timestamp(df.t, "yyyy-MM-dd HH:mm:ss").alias("dt")).collect() [Row(dt=datetime.datetime(1997, 2, 28, 10, 30))] """ # noqa: D205 - return _to_date_or_timestamp(col, _types.TimestampNTZType(), format) + return _to_date_or_timestamp(col, _types.TimestampType(), format) def to_timestamp_ltz( @@ -5706,7 +6211,8 @@ def to_timestamp_ltz( ... # doctest: +SKIP [Row(r=datetime.datetime(2016, 12, 31, 0, 0))] """ # noqa: D205 - return _to_date_or_timestamp(timestamp, _types.TimestampNTZType(), format) + pattern = _datetime_pattern_argument(format) + return _to_date_or_timestamp(timestamp, _types.TimestampType(), pattern) def to_timestamp_ntz( @@ -5737,7 +6243,8 @@ def to_timestamp_ntz( ... # doctest: +SKIP [Row(r=datetime.datetime(2016, 4, 8, 0, 0))] """ # noqa: D205 - return _to_date_or_timestamp(timestamp, _types.TimestampNTZType(), format) + pattern = _datetime_pattern_argument(format) + return _to_date_or_timestamp(timestamp, _types.TimestampNTZType(), pattern) def try_to_timestamp(col: "ColumnOrName", format: Optional["ColumnOrName"] = None) -> Column: @@ -5760,10 +6267,8 @@ def try_to_timestamp(col: "ColumnOrName", format: Optional["ColumnOrName"] = Non >>> df.select(try_to_timestamp(df.t, lit("yyyy-MM-dd HH:mm:ss")).alias("dt")).collect() [Row(dt=datetime.datetime(1997, 2, 28, 10, 30))] """ # noqa: D205, D415 - if format is None: - format = lit(["%Y-%m-%d", "%Y-%m-%d %H:%M:%S"]) - - return _invoke_function_over_columns("try_strptime", col, format) + pattern = _datetime_pattern_argument(format) + return _to_date_or_timestamp(col, _types.TimestampType(), pattern, try_cast=True) def substr(str: "ColumnOrName", pos: "ColumnOrName", len: Optional["ColumnOrName"] = None) -> Column: @@ -5817,16 +6322,13 @@ def substr(str: "ColumnOrName", pos: "ColumnOrName", len: Optional["ColumnOrName | k SQL| +------------------------+ """ # noqa: D205 - if len is not None: - return _invoke_function_over_columns("substring", str, pos, len) - else: - return _invoke_function_over_columns("substring", str, pos) + return _spark_substring(_to_column_expr(str), _to_column_expr(pos), None if len is None else _to_column_expr(len)) def _unix_diff(col: "ColumnOrName", part: str) -> Column: - return _invoke_function_over_columns( - "date_diff", lit(part), lit("1970-01-01 00:00:00+00:00").cast("timestamp"), col - ) + # Days count calendar dates; the smaller units count from the epoch instant, which is midnight UTC. + epoch = lit("1970-01-01").cast("date") if part == "days" else lit("1970-01-01 00:00:00+00:00").cast("timestamp") + return _invoke_function_over_columns("date_diff", lit(part), epoch, col) def unix_date(col: "ColumnOrName") -> Column: @@ -5842,7 +6344,7 @@ def unix_date(col: "ColumnOrName") -> Column: [Row(n=1)] >>> spark.conf.unset("spark.sql.session.timeZone") """ - return _unix_diff(col, "days") + return _result_type(_unix_diff(col, "days"), INTEGER) def unix_micros(col: "ColumnOrName") -> Column: @@ -5922,22 +6424,20 @@ def arrays_overlap(a1: "ColumnOrName", a2: "ColumnOrName") -> Column: a1_has_null = _list_contains_null(a1) a2_has_null = _list_contains_null(a2) - return Column( + return _null_if_null( CaseExpression(FunctionExpression("list_has_any", a1, a2), ConstantExpression(True)).otherwise( CaseExpression( (FunctionExpression("len", a1) > 0) & (FunctionExpression("len", a2) > 0) & (a1_has_null | a2_has_null), ConstantExpression(None), ).otherwise(ConstantExpression(False)) - ) + ), + a1, + a2, ) def _list_contains_null(c: ColumnExpression) -> Expression: - return FunctionExpression( - "list_contains", - FunctionExpression("list_transform", c, LambdaExpression("x", ColumnExpression("x").isnull())), - True, - ) + return FunctionExpression("list_count", c) < FunctionExpression("len", c) def arrays_zip(*cols: "ColumnOrName") -> Column: @@ -5981,7 +6481,18 @@ def arrays_zip(*cols: "ColumnOrName") -> Column: | | |-- vals2: long (nullable = true) | | |-- vals3: long (nullable = true) """ # noqa: D205 - return _invoke_function_over_columns("list_zip", *cols) + if not cols: + return Column(FunctionExpression("list_value").cast(duckdb.list_type(duckdb.struct_type({})))) + columns = [_as_column(c) for c in cols] + arrays = [c.expr for c in columns] + names = [name if (name := c._struct_field()) is not None else f"{i}" for i, c in enumerate(columns)] + + def zipped_struct(row: Expression) -> Expression: + fields = [FunctionExpression("struct_extract_at", row, ConstantExpression(i)) for i in range(1, len(names) + 1)] + return _struct_pack(names, fields) + + zipped = FunctionExpression("list_transform", FunctionExpression("list_zip", *arrays), _lambda(zipped_struct)) + return _null_if_null(zipped, *arrays) def substring(str: "ColumnOrName", pos: int, len: int) -> Column: @@ -6020,11 +6531,25 @@ def substring(str: "ColumnOrName", pos: int, len: int) -> Column: >>> df.select(substring(df.s, 1, 2).alias("s")).collect() [Row(s='ab')] """ # noqa: D205 - return _invoke_function( - "substring", - _to_column_expr(str), - ConstantExpression(pos), - ConstantExpression(len), + return _spark_substring(_to_column_expr(str), ConstantExpression(pos), ConstantExpression(len)) + + +def _spark_substring(string: Expression, pos: Expression, length: Expression | None) -> Column: + """Spark reads position 0 as 1, counts a negative position from the end, and gives '' for a negative length.""" + start = ( + CaseExpression(pos > 0, pos - 1) + .when(pos < 0, FunctionExpression("length", string) + pos) + .when(pos == 0, ConstantExpression(0)) + ) + # Not greatest(), which skips a NULL where Spark gives NULL. + begin = CaseExpression(start < 0, ConstantExpression(0)).otherwise(start) + if length is None: + return Column(FunctionExpression("substring", string, begin + 1)) + end = start + length + return Column( + CaseExpression(string.isnotnull() & (begin >= end), ConstantExpression("")).otherwise( + FunctionExpression("substring", string, begin + 1, end - begin) + ) ) @@ -6133,7 +6658,15 @@ def concat(*cols: "ColumnOrName") -> Column: >>> df DataFrame[arr: array] """ # noqa: D205, D415 - return _invoke_function_over_columns("concat", *cols) + if not cols: + return lit("") + # DuckDB's || keeps binary and arrays and propagates NULL like Spark, where its concat makes text and skips NULL; + # Spark reads a NULL literal as a string, where DuckDB would make the whole result untyped. + exprs = [_to_column_expr(c).cast(VARCHAR) if _is_null_literal(c) else _to_column_expr(c) for c in cols] + if len(exprs) == 1: + (value,) = exprs + return Column(FunctionExpression("cast_to_type", value, FunctionExpression("||", value, value))) + return Column(functools.reduce(lambda left, right: FunctionExpression("||", left, right), exprs)) def instr(str: "ColumnOrName", substr: str) -> Column: @@ -6173,7 +6706,7 @@ def instr(str: "ColumnOrName", substr: str) -> Column: >>> df.select(instr(df.s, "b").alias("s")).collect() [Row(s=2)] """ # noqa: D205 - return _invoke_function("instr", _to_column_expr(str), ConstantExpression(substr)) + return _result_type(_invoke_function("instr", _to_column_expr(str), ConstantExpression(substr)), INTEGER) def expr(str: str) -> Column: @@ -6215,3 +6748,156 @@ def broadcast(df: "DataFrame") -> "DataFrame": or optimizations, since broadcasting is not applicable in the DuckDB context. """ # noqa: D205 return df + + +def _is_star(value: Any) -> bool: # noqa: ANN401 + if isinstance(value, Column): + return str(value.expr) == "*" + if isinstance(value, str): + return value == "*" + return isinstance(value, _Ref) and value.parts[-1] == "*" + + +# Spark names a call `name(arg, ...)` after the function's own name, so only the functions whose names differ +# are listed: by the name Spark uses, or by a callable that builds the name from the arguments, including the +# defaults Spark adds and leaving out the arguments Spark leaves out. +_SPARK_NAMES: dict[str, str | Callable[..., Name]] = { + "acos": "ACOS", + "approxCountDistinct": lambda c, *_: _call_name("approx_count_distinct", c), + "approx_count_distinct": lambda c, *_: _call_name("approx_count_distinct", c), + "any_value": lambda c, *_: _call_name("any_value", c), + "array_agg": "collect_list", + "asc": lambda c: _sql_name(c, " ASC", " NULLS FIRST"), + "asc_nulls_first": lambda c: _sql_name(c, " ASC", " NULLS FIRST"), + "asc_nulls_last": lambda c: _sql_name(c, " ASC", " NULLS LAST"), + "asin": "ASIN", + "atan": "ATAN", + "atan2": "ATAN2", + "bround": lambda c, scale=0: _call_name("bround", c, scale), + "call_function": lambda name, *cols: _call_name(name, *cols), + "cbrt": "CBRT", + "ceil": "CEIL", + "cos": "COS", + "cot": "COT", + "count": lambda c: _call_name("count", 1 if _is_star(c) else c), + "curdate": "current_date", + "datepart": lambda field, source: _format_name("datepart({0} FROM {1})", field, source), + "degrees": "DEGREES", + "desc": lambda c: _sql_name(c, " DESC", " NULLS LAST"), + "desc_nulls_first": lambda c: _sql_name(c, " DESC", " NULLS FIRST"), + "desc_nulls_last": lambda c: _sql_name(c, " DESC", " NULLS LAST"), + "e": "E", + "exp": "EXP", + "extract": lambda field, source: _format_name("extract({0} FROM {1})", field, source), + "first": lambda c, *_: _call_name("first", c), + "floor": "FLOOR", + "isnotnull": lambda c: _format_name("({0} IS NOT NULL)", c), + "isnull": lambda c: _format_name("({0} IS NULL)", c), + "last": lambda c, *_: _call_name("last", c), + "ilike": lambda s, pattern, *_: _call_name("ilike", s, pattern), + "like": lambda s, pattern, escape=None: _format_name( + "{0} LIKE {1}" if escape is None else "{0} LIKE {1} ESCAPE '{2}'", s, pattern, escape + ), + "log2": "LOG2", + "mean": "avg", + "pi": "PI", + "positive": lambda c: _format_name("(+ {0})", c), + "pow": "POWER", + "radians": "RADIANS", + "regexp": "REGEXP", + "regexp_extract_all": lambda s, regexp, idx=1: _call_name("regexp_extract_all", s, regexp, idx), + "regexp_like": "REGEXP_LIKE", + "regexp_replace": lambda s, pattern, replacement: _call_name("regexp_replace", s, pattern, replacement, 1), + "round": lambda c, scale=0: _call_name("round", c, scale), + "signum": "SIGNUM", + "sin": "SIN", + "sort_array": lambda c, asc=True: _call_name("sort_array", c, asc), + "split": lambda s, pattern, limit=-1: _call_name("split", s, pattern, limit), + "sqrt": "SQRT", + "substr": lambda s, pos, length=2147483647: _call_name("substr", s, pos, length), + "tan": "TAN", +} + +# These name their results themselves. +_NAMED_IN_BODY = {"col", "lit", "expr", "when"} + + +def _spark_named(function: Callable[..., Any]) -> Callable[..., Any]: + """Give the Columns ``function`` returns the name Spark gives the same call.""" + spark_name = _SPARK_NAMES.get(function.__name__, function.__name__) + signature = inspect.signature(function) + + @functools.wraps(function) + def named(*args: Any, **kwargs: Any) -> Any: # noqa: ANN401 + result = function(*args, **kwargs) + if not isinstance(result, Column): + return result + values: list[Any] = [] + for param_name, value in signature.bind(*args, **kwargs).arguments.items(): + param = signature.parameters[param_name] + if param.kind is param.VAR_POSITIONAL: + given = value[0] if len(value) == 1 and isinstance(value[0], (list, tuple, set)) else value + elif value is None and param.default is None: + continue + else: + given = (value,) + refs = "ColumnOrName" in str(param.annotation) + values.extend(_Ref(_name_parts(v)) if refs and isinstance(v, str) else v for v in given) + name = spark_name(*values) if callable(spark_name) else _call_name(spark_name, *values) + return result._with_name(name) + + return named + + +for _function_name, _function in list(globals().items()): + if ( + inspect.isfunction(_function) + and _function.__module__ == __name__ + and not _function_name.startswith("_") + and _function_name not in _NAMED_IN_BODY + ): + globals()[_function_name] = _spark_named(_function) + + +def _carry_markers(function: Callable[..., Any]) -> Callable[..., Any]: + signature = inspect.signature(function) + + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: # noqa: ANN401 + result = function(*args, **kwargs) + if isinstance(result, Column) and not result._is_aggregate: + arguments: list[Any] = [] + names_a_column = False + for param_name, value in signature.bind(*args, **kwargs).arguments.items(): + param = signature.parameters[param_name] + given = list(value) if param.kind is param.VAR_POSITIONAL else [value] + given += [x for a in given if isinstance(a, (list, tuple, set)) for x in a] + arguments += given + # A str names a column only where the function takes a column. + if "ColumnOrName" in str(param.annotation) and any(isinstance(a, str) for a in given): + names_a_column = True + if any(isinstance(a, Column) and a._is_aggregate for a in arguments): + result._is_aggregate = True + result._aggregate_inputs = tuple( + x for a in arguments if isinstance(a, Column) for x in a._aggregate_inputs + ) + elif not names_a_column and all(a._is_constant for a in arguments if isinstance(a, Column)): + result._is_constant = True + return result + + return wrapper + + +# These read a column name or SQL text from a str, so their result is never a constant. +_READ_COLUMNS_FROM_TEXT = {"col", "expr"} + +# A Column built from an aggregate is itself an aggregate, and one built only from constants is a constant; +# DataFrame.select relies on both. +for _name, _function in list(globals().items()): + if ( + not _name.startswith("_") + and _name not in _READ_COLUMNS_FROM_TEXT + and inspect.isfunction(_function) + and _function.__module__ == __name__ + ): + globals()[_name] = _carry_markers(_function) diff --git a/duckdb/experimental/spark/sql/group.py b/duckdb/experimental/spark/sql/group.py index 5f784453a..127dc6096 100644 --- a/duckdb/experimental/spark/sql/group.py +++ b/duckdb/experimental/spark/sql/group.py @@ -15,13 +15,17 @@ # limitations under the License. # -from collections.abc import Callable +from collections.abc import Callable, Collection from typing import TYPE_CHECKING, overload +from duckdb import Expression + +from ..errors import AnalysisException from ..exception import ContributionsAcceptedError -from .column import Column -from .dataframe import DataFrame -from .functions import _to_column_expr +from . import functions +from .column import Column, _name_parts, _named_expr, _ref_name +from .dataframe import DataFrame, _aggregate, _name_key +from .functions import _to_column_expr, col from .types import NumericType # Only import symbols needed for type checking if something is type checking @@ -32,23 +36,26 @@ __all__ = ["GroupedData", "Grouping"] +def _numeric_columns(df: DataFrame) -> list[str]: + # Spark reads these names as attribute names too, so a name with a backtick raises there as well. + return [x.name for x in df.schema.fields if isinstance(x.dataType, NumericType)] + + 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( - function_name=name, # aggregate function - function_aggr=expressions, # inputs to aggregate - group_expr=group_by, # groups - projected_columns=projections, # projections - ) - return DataFrame(jdf, self.session) + non_numeric = {_name_key(x.name) for x in self._df.schema.fields if not isinstance(x.dataType, NumericType)} + for column in cols: + parts = _name_parts(column) + if len(parts) == 1 and _name_key(parts[0]) in non_numeric: + raise AnalysisException(error_class="_LEGACY_ERROR_TEMP_1323", message_parameters={"colName": column}) + function = getattr(functions, name) + return self._agg([function(Column(self._df._scope_column_expr(c), _ref_name(c))) for c in cols]) def df_varargs_api(f: Callable[..., DataFrame]) -> Callable[..., DataFrame]: def _api(self: "GroupedData", *cols: str) -> DataFrame: name = f.__name__ - return _api_internal(self, name, *cols) + # Without columns Spark aggregates every numeric column, grouping keys included. + return _api_internal(self, name, *(cols or _numeric_columns(self._df))) _api.__name__ = f.__name__ _api.__doc__ = f.__doc__ @@ -58,7 +65,8 @@ def _api(self: "GroupedData", *cols: str) -> DataFrame: 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] + self._columns = [col(x) if isinstance(x, str) else x for x in cols] + self._cols = [_to_column_expr(x) for x in self._columns] if "special" in kwargs: special = kwargs["special"] accepted_special = ["cube", "rollup"] @@ -69,6 +77,13 @@ def get_columns(self) -> str: # noqa: D102 columns = ",".join([str(x) for x in self._cols]) return columns + def named_columns(self, columns: list[str], references: Collection[str]) -> list[Expression]: + """The grouping expressions aliased to the names Spark gives them in a projection over ``columns``. + + ``references`` holds the renderings of the qualified references to the columns, folded to lower case. + """ + return [_named_expr(x, columns, references) for x in self._columns] + def __str__(self) -> str: # noqa: D105 columns = self.get_columns() if self._type: @@ -118,9 +133,8 @@ def count(self) -> DataFrame: | Bob| 2| +-----+-----+ """ - return _api_internal(self, "count").withColumnRenamed("count_star()", "count") + return self._agg([functions.count(functions.lit(1)).alias("count")]) - @df_varargs_api def mean(self, *cols: str) -> DataFrame: """Computes average values for each numeric columns for each group. @@ -131,6 +145,7 @@ def mean(self, *cols: str) -> DataFrame: cols : str column names. Non-numeric columns are ignored. """ + return self.avg(*cols) def avg(self, *cols: str) -> DataFrame: """Computes average values for each numeric columns for each group. @@ -177,12 +192,7 @@ def avg(self, *cols: str) -> DataFrame: | 5.0| 110.0| +--------+-----------+ """ - columns = list(cols) - 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)] - return _api_internal(self, "avg", *columns) + return _api_internal(self, "avg", *(cols or _numeric_columns(self._df))) @df_varargs_api def max(self, *cols: str) -> DataFrame: @@ -413,13 +423,12 @@ def agg(self, *exprs: Column | dict[str, str]) -> DataFrame: assert exprs, "exprs should not be empty" if len(exprs) == 1 and isinstance(exprs[0], dict): raise ContributionsAcceptedError - else: - # Columns - assert all(isinstance(c, Column) for c in exprs), "all exprs should be Column" - expressions = list(self._grouping._cols) - expressions.extend([x.expr for x in exprs]) - group_by = str(self._grouping) - rel = self._df.relation.select(*expressions, groups=group_by) - return DataFrame(rel, self.session) + assert all(isinstance(c, Column) for c in exprs), "all exprs should be Column" + return self._agg(list(exprs)) + + def _agg(self, exprs: list[Column]) -> DataFrame: + self._df._check_references(*self._grouping._columns, *exprs) + groups = self._grouping.named_columns(self._df.columns, self._df._qualified_references()) + return DataFrame(_aggregate(self._df, groups, exprs, str(self._grouping)), self.session) # TODO: add 'pivot' # noqa: TD002, TD003 diff --git a/duckdb/experimental/spark/sql/readwriter.py b/duckdb/experimental/spark/sql/readwriter.py index 230d5d2af..1db3b11dd 100644 --- a/duckdb/experimental/spark/sql/readwriter.py +++ b/duckdb/experimental/spark/sql/readwriter.py @@ -1,24 +1,396 @@ -from typing import TYPE_CHECKING, cast +import pathlib +import urllib.parse +import urllib.request +from collections import Counter +from typing import TYPE_CHECKING, Any, cast -from ..errors import PySparkNotImplementedError, PySparkTypeError +from duckdb import CaseExpression, ColumnExpression, ConstantExpression, FunctionExpression, SQLExpression +from duckdb.sqltypes import DuckDBPyType + +from ..errors import AnalysisException, IllegalArgumentException, PySparkNotImplementedError, PySparkTypeError from ..exception import ContributionsAcceptedError -from .types import StructType +from .column import Column, _quoted_identifier +from .functions import _parse_datetime, _tokenize_datetime_pattern, _translate_datetime_pattern_field +from .type_utils import convert_type +from .types import ( + ArrayType, + DataType, + DateType, + DoubleType, + FloatType, + MapType, + NullType, + StringType, + StructField, + StructType, + TimeNTZType, + TimestampNTZType, + TimestampType, + UUIDType, + _parse_datatype_string, + _replace_types, +) PrimitiveType = bool | float | int | str OptionalPrimitiveType = PrimitiveType | None if TYPE_CHECKING: + import duckdb + from duckdb import Expression from duckdb.experimental.spark.sql.dataframe import DataFrame from duckdb.experimental.spark.sql.session import SparkSession +# The DuckDB strftime code that writes a Spark datetime pattern field, by letter and count, where one is exact. +_STRFTIME_CODES = { + ("y", 2): "%y", + ("y", 4): "%Y", + **{(letter, 2): "%m" for letter in "ML"}, + **{(letter, 3): "%b" for letter in "ML"}, + **{(letter, 4): "%B" for letter in "ML"}, + ("d", 2): "%d", + ("H", 2): "%H", + ("h", 2): "%I", + ("m", 2): "%M", + ("s", 2): "%S", + ("S", 3): "%g", + ("S", 6): "%f", + ("S", 9): "%n", + ("a", 1): "%p", + **{("E", count): "%a" for count in (1, 2, 3)}, + ("E", 4): "%A", +} +# Pattern letters Spark accepts only for formatting, which the translation for parsing rejects. +_FORMATTING_LETTERS = frozenset("EFQq") +# DuckDB's JSON reader detects these types in strings, which Spark's JSON inference leaves strings; a field that is +# only ever NULL is a string to Spark too. +_JSON_STRING_TYPES = { + cls: StringType() for cls in (DateType, TimeNTZType, TimestampNTZType, TimestampType, UUIDType, NullType) +} +_SAVE_MODES = { + "append": "append", + "overwrite": "overwrite", + "ignore": "ignore", + "error": "error", + "errorifexists": "error", + "default": "error", +} +_FORMATS_BY_EXTENSION = {".csv": "csv", ".tsv": "tsv", ".json": "json", ".jsonl": "json", ".ndjson": "json"} +# DuckDB's default type candidates without BIGNUM, which the module has no type for; wider integers read as doubles. +_CSV_TYPE_CANDIDATES = ["BOOLEAN", "BIGINT", "HUGEINT", "DOUBLE", "TIME", "DATE", "TIMESTAMP", "TIMESTAMPTZ", "VARCHAR"] + + +def _parse_schema(schema: StructType | str) -> StructType: + if isinstance(schema, str): + schema = _parse_datatype_string(schema) + if not isinstance(schema, StructType): + raise PySparkTypeError( + error_class="NOT_STR_OR_STRUCT", + message_parameters={"arg_name": "schema", "arg_type": type(schema).__name__}, + ) + _check_unique_names(schema.names) + return schema + + +def _check_unique_names(names: list[str]) -> None: + """Raise where two of ``names`` are one name ignoring case, as Spark does for a schema or a written file.""" + seen: set[str] = set() + for name in names: + if name.lower() in seen: + raise AnalysisException( + error_class="COLUMN_ALREADY_EXISTS", message_parameters={"columnName": f"`{name.lower()}`"} + ) + seen.add(name.lower()) + + +def _flag(name: str, value: bool | str | None, default: bool = False) -> bool: + """A boolean option, which Spark takes as a bool or as the text true or false in any case.""" + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, str) and value.lower() in ("true", "false"): + return value.lower() == "true" + raise IllegalArgumentException(message=f"{name} flag can be true or false.") + + +def _positions_by_name(columns: list[str], schema: StructType, ignore_case: bool) -> list[int | None]: + """For each field of ``schema``, the position of the column of the same name, or None where there is none.""" + positions: dict[str, int] = {} + for i, column in enumerate(columns): + positions.setdefault(column.casefold() if ignore_case else column, i) + return [positions.get(f.name.casefold() if ignore_case else f.name) for f in schema.fields] + + +def _with_schema(df: "DataFrame", schema: StructType, positions: list[int | None]) -> "DataFrame": + """``df`` read as ``schema``: each field from the column at its position, or NULL where it has none.""" + from duckdb.experimental.spark.sql.dataframe import _position + + columns = [ + (ConstantExpression(None) if i is None else _position(i)).cast(f.dataType.duckdb_type).alias(f.name) + for i, f in zip(positions, schema.fields, strict=True) + ] + return df._derive(df.relation.select(*columns)) + + +def _format_of(path: str) -> str: + """The format ``load`` reads a path in without a format: by its extension, otherwise parquet as Spark does.""" + name = path.lower().removesuffix(".gz").removesuffix(".zst") + return _FORMATS_BY_EXTENSION.get(pathlib.PurePath(name).suffix, "parquet") + + +def _local_path(path: str) -> pathlib.Path | None: + """The local file ``path`` names, as DuckDB expands it, or None for a path of another file system.""" + if path.startswith("file:"): + return pathlib.Path(urllib.request.url2pathname(urllib.parse.urlparse(path).path)) + if "://" in path: + return None + return pathlib.Path(path).expanduser() + + +def _path_exists(session: "SparkSession", path: str) -> bool: + local = _local_path(path) + if local is not None: + return local.exists() + # A directory of an object store is the objects under its prefix. + query = "SELECT 1 FROM glob($1) UNION ALL SELECT 1 FROM glob($2) LIMIT 1" + return bool(session.conn.execute(query, [path, path.rstrip("/") + "/*"]).fetchall()) + + +def _duckdb_path(path: str) -> str: + """``path`` as DuckDB opens it: a file URL as the local path it names, since DuckDB reads it as a relative path.""" + local = _local_path(path) + return path if local is None or not path.startswith("file:") else str(local) + + +def _qualified_path(path: str) -> str: + """``path`` as Spark prints it in a message, a local file as a file URI.""" + local = _local_path(path) + return path if local is None else local.absolute().as_uri() + + +def _should_write(session: "SparkSession", path: str, mode: str | None) -> bool: + """Whether a writer in save ``mode`` writes ``path``, raising where Spark does when the path exists.""" + if mode is not None and (not isinstance(mode, str) or mode.lower() not in _SAVE_MODES): + raise AnalysisException(error_class="INVALID_SAVE_MODE", message_parameters={"mode": f'"{mode}"'}) + save_mode = "error" if mode is None else _SAVE_MODES[mode.lower()] + if not _path_exists(session, path): + return True + if save_mode == "error": + local = _local_path(path) + location = path if local is None else f"file:{local.absolute()}" + raise AnalysisException(error_class="PATH_ALREADY_EXISTS", message_parameters={"outputPath": location}) + if save_mode == "append": + msg = "Appending to an existing path is not supported: DuckDB writes a single file" + raise ContributionsAcceptedError(msg) + return save_mode == "overwrite" + + +def _spark_sql_type(data_type: DataType) -> str: + """The SQL text of a type in Spark's error messages.""" + if isinstance(data_type, ArrayType): + return f"ARRAY<{_spark_sql_type(data_type.elementType)}>" + if isinstance(data_type, MapType): + return f"MAP<{_spark_sql_type(data_type.keyType)}, {_spark_sql_type(data_type.valueType)}>" + if isinstance(data_type, StructType): + fields = (f"{_quoted_identifier(f.name)}: {_spark_sql_type(f.dataType)}" for f in data_type.fields) + return f"STRUCT<{', '.join(fields)}>" + return data_type.simpleString().upper() + + +def _unsupported_csv_type(field: StructField) -> AnalysisException: + return AnalysisException( + error_class="UNSUPPORTED_DATA_TYPE_FOR_DATASOURCE", + message_parameters={ + "format": "CSV", + "columnName": f"`{field.name}`", + "columnType": f'"{_spark_sql_type(field.dataType)}"', + }, + ) + + +def _strftime_format(pattern: str | None) -> str | None: + """The DuckDB strftime format that writes values as the Spark datetime ``pattern`` formats them.""" + if pattern is None: + return None + codes = [] + for token in _tokenize_datetime_pattern(pattern): + if isinstance(token, str): + codes.append(token.replace("%", "%%")) + continue + code = _STRFTIME_CODES.get(token) + if code is None: + if token[0] not in _FORMATTING_LETTERS: + _translate_datetime_pattern_field(pattern, *token) + msg = f"Datetime pattern {pattern!r} is not supported for formatting: {token[0] * token[1]!r}" + raise ContributionsAcceptedError(msg) + codes.append(code) + return "".join(codes) + + +def _csv_dialect(sep: str | None, quote: str | None, escape: str | None) -> dict[str, Any]: + """The CSV dialect given, with Spark's separator and quote and DuckDB's escape where none is, never sniffed. + + DuckDB's escape is the quote, so a backslash is text and a quote inside quotes is doubled. + """ + quote = '"' if quote is None else quote + return {"sep": "," if sep is None else sep, "quotechar": quote, "escapechar": quote if escape is None else escape} + + +def _try_cast(value: "Expression", target: DuckDBPyType) -> "Expression": + return SQLExpression(f"TRY_CAST({value} AS {target})") + + +def _csv_value( + text: "Expression", field: StructField, patterns: dict[type[DataType], str], specials: tuple[str, str, str] +) -> "Expression": + """The CSV field ``text`` read as the type of ``field``, NULL where it does not convert, as in Spark. + + A date or timestamp with a Spark datetime pattern in ``patterns`` reads with it, any other type with DuckDB's + cast. ``specials`` are the texts of NaN, infinity and minus infinity. + """ + data_type = field.dataType + kind, target = type(data_type), data_type.duckdb_type + if kind in (ArrayType, MapType, StructType, NullType): + raise _unsupported_csv_type(field) + if kind is StringType: + return text + if kind in patterns: + return _parse_datetime(Column(text), patterns[kind], raise_on_error=False).expr.cast(target) + if kind not in (FloatType, DoubleType): + return _try_cast(text, target) + nan, infinity, minus_infinity = specials + return ( + CaseExpression(text == ConstantExpression(nan), ConstantExpression("NaN").cast(target)) + .when(text == ConstantExpression(infinity), ConstantExpression("Infinity").cast(target)) + .when(text == ConstantExpression(minus_infinity), ConstantExpression("-Infinity").cast(target)) + .otherwise(_try_cast(text, target)) + ) + + +def _csv_header_names(fields: list[str | None], null_value: str | None) -> list[str]: + """Spark's column names for the fields of a CSV header line, made unique by their position.""" + counts = Counter(field.lower() for field in fields if field is not None) + return [ + f"_c{i}" if not field or field == null_value else f"{field}{i}" if counts[field.lower()] > 1 else field + for i, field in enumerate(fields) + ] + + +def _first_csv_line(session: "SparkSession", path: str, options: dict[str, Any]) -> list[str | None] | None: + """The fields of the first line of the CSV files ``path`` names that is not empty, or None.""" + width = len(session.conn.read_csv(path, auto_detect=True, header=False, all_varchar=True, **options).columns) + columns = {f"column{i}": "VARCHAR" for i in range(width)} + # A null string no line contains keeps an empty field apart from a field past the end of the line. + line_options = {**options, "header": False, "na_values": ["\0"]} + rows = session.conn.read_csv(path, auto_detect=False, columns=columns, **line_options).limit(1).fetchall() + if not rows: + return None + fields = list(rows[0]) + while fields and fields[-1] is None: + fields.pop() + return fields + + +def _check_csv_header(fields: list[str | None], schema: StructType, path: str) -> None: + """Raise as Spark does with enforceSchema off where the header does not name the schema's fields in order.""" + header = ["" if field is None else field for field in fields] + names = schema.names + source = f"CSV file: {_qualified_path(path)}" + if len(header) != len(names): + msg = ( + "Number of column in CSV header is not equal to number of fields in the schema:\n" + f" Header length: {len(header)}, schema size: {len(names)}\n{source}" + ) + raise IllegalArgumentException(message=msg) + for column, name in zip(header, names, strict=True): + if column.lower() != name.lower(): + msg = ( + "CSV header does not conform to the schema.\n" + f" Header: {', '.join(header)}\n Schema: {', '.join(names)}\n" + f"Expected: {name} but found: {column}\n{source}" + ) + raise IllegalArgumentException(message=msg) + + +def _spark_json_type(data_type: DataType) -> DataType: + """The type Spark's JSON inference gives a value DuckDB's JSON reader infers ``data_type`` for.""" + if isinstance(data_type, StructType): + fields = sorted(data_type.fields, key=lambda f: f.name) + return StructType([StructField(f.name, _spark_json_type(f.dataType)) for f in fields]) + if isinstance(data_type, ArrayType): + return ArrayType(_spark_json_type(data_type.elementType)) + return _JSON_STRING_TYPES.get(type(data_type), data_type) + + +def _read_json(session: "SparkSession", path: str) -> "duckdb.DuckDBPyRelation | None": + """Read JSON lines with the schema Spark infers from every record: no map type, fields sorted by name. + + None for a file without records. Each record is inferred as one value, so DuckDB raises for keys that differ + only in case, where it would rename one of them. + """ + inferred = session.conn.read_json(path, sample_size=-1, map_inference_threshold=-1, records="false") + (record_type,) = inferred.types + if record_type.id != "struct": + if not inferred.limit(1).fetchall(): + return None + msg = "Reading JSON lines that are not all objects is not supported" + raise ContributionsAcceptedError(msg) + schema = _spark_json_type(convert_type(record_type)) + columns = {f.name: str(f.dataType.duckdb_type) for f in cast("StructType", schema).fields} + return session.conn.read_json(path, columns=columns) + + +def _read_json_with_schema(session: "SparkSession", path: str, schema: StructType) -> "DataFrame": + """Read JSON lines as ``schema`` with DuckDB's JSON conversion, NULL where a value does not convert.""" + from duckdb.experimental.spark.sql.dataframe import _position, _source_frame + + # DuckDB's JSON reader reads a timestamp without an offset as UTC, where its cast reads it in the session zone. + read_types = [_replace_types(f.dataType, {TimestampType: StringType()}) for f in schema.fields] + columns = {f.name: str(t.duckdb_type) for f, t in zip(schema.fields, read_types, strict=True)} + rel = session.conn.read_json(path, columns=columns, records="true", ignore_errors=True) + df = _source_frame(rel, session) + values = [_try_cast(_position(i), f.dataType.duckdb_type).alias(f.name) for i, f in enumerate(schema.fields)] + return df._derive(df.relation.select(*values)) + + +def _read_inferred_csv( + session: "SparkSession", path: str, names: list[str], header: bool, null_value: str, options: dict[str, Any] +) -> "DataFrame": + """The lines of a CSV file read by DuckDB with the types it detects, named ``names``. + + DuckDB detects the types from a sample of the lines, as its own reader does: a later value its reader cannot read + as the detected type raises, and a fraction in a column detected as integers is rounded. + """ + from duckdb.experimental.spark.sql.dataframe import _source_frame + + rel = session.conn.read_csv( + path, + # DuckDB's sniffer fails on a header line of another width than the other lines, so it skips it as a line. + header=False, + skiprows=1 if header else 0, + na_values=["", null_value], + auto_type_candidates=_CSV_TYPE_CANDIDATES, + **options, + ) + fields = [ColumnExpression(name) for name in rel.columns] + values = [] + for i, name in enumerate(names): + data_type = convert_type(rel.types[i]) if i < len(fields) else StringType() + # Spark infers a timestamp as its TimestampType, the text read as a local time. + data_type = TimestampType() if isinstance(data_type, TimestampNTZType) else data_type + value = fields[i] if i < len(fields) else ConstantExpression(None) + values.append(value.cast(data_type.duckdb_type).alias(name)) + df = _source_frame(rel, session) + return df._derive(df.relation.select(*values)) + class DataFrameWriter: # noqa: D101 def __init__(self, dataframe: "DataFrame") -> None: # noqa: D107 self.dataframe = dataframe def saveAsTable(self, table_name: str) -> None: # noqa: D102 - relation = self.dataframe.relation - relation.create(table_name) + _check_unique_names(self.dataframe.columns) + self.dataframe.relation.create(table_name) def parquet( # noqa: D102 self, @@ -28,12 +400,11 @@ def parquet( # noqa: D102 compression: str | None = None, ) -> None: relation = self.dataframe.relation - if mode: - raise NotImplementedError if partitionBy: raise NotImplementedError - - relation.write_parquet(path, compression=compression) + _check_unique_names(self.dataframe.columns) + if _should_write(self.dataframe.session, path, mode): + relation.write_parquet(_duckdb_path(path), compression=compression) def csv( # noqa: D102 self, @@ -56,33 +427,37 @@ def csv( # noqa: D102 emptyValue: str | None = None, lineSep: str | None = None, ) -> None: - if mode not in (None, "overwrite"): - raise NotImplementedError - if escapeQuotes: - raise NotImplementedError - if ignoreLeadingWhiteSpace: - raise NotImplementedError - if ignoreTrailingWhiteSpace: - raise NotImplementedError + if not _flag("escapeQuotes", escapeQuotes, default=True): + msg = "escapeQuotes=False is not supported: DuckDB quotes every value that contains the quote character" + raise ContributionsAcceptedError(msg) + if _flag("ignoreLeadingWhiteSpace", ignoreLeadingWhiteSpace): + msg = "ignoreLeadingWhiteSpace=True is not supported: DuckDB writes strings as they are" + raise ContributionsAcceptedError(msg) + if _flag("ignoreTrailingWhiteSpace", ignoreTrailingWhiteSpace): + msg = "ignoreTrailingWhiteSpace=True is not supported: DuckDB writes strings as they are" + raise ContributionsAcceptedError(msg) if charToEscapeQuoteEscaping: raise NotImplementedError if emptyValue: raise NotImplementedError if lineSep: raise NotImplementedError - relation = self.dataframe.relation - relation.write_csv( - path, - sep=sep, + _check_unique_names(self.dataframe.columns) + if not _should_write(self.dataframe.session, path, mode): + return + for field in self.dataframe.schema.fields: + if isinstance(field.dataType, (ArrayType, MapType, StructType, NullType)): + raise _unsupported_csv_type(field) + self.dataframe.relation.write_csv( + _duckdb_path(path), + **_csv_dialect(sep, quote, escape), na_rep=nullValue, - quotechar=quote, compression=compression, - escapechar=escape, - header=header if isinstance(header, bool) else header == "True", + header=_flag("header", header), encoding=encoding, - quoting=quoteAll, - date_format=dateFormat, - timestamp_format=timestampFormat, + quoting="ALL" if _flag("quoteAll", quoteAll) else None, + date_format=_strftime_format(dateFormat), + timestamp_format=_strftime_format(timestampFormat), ) @@ -97,35 +472,29 @@ def load( # noqa: D102 schema: StructType | str | None = None, **options: OptionalPrimitiveType, ) -> "DataFrame": - from duckdb.experimental.spark.sql.dataframe import DataFrame + from duckdb.experimental.spark.sql.dataframe import DataFrame, _source_frame if not isinstance(path, str): raise TypeError if options: raise ContributionsAcceptedError - rel = None - if format: - format = format.lower() - if format == "csv" or format == "tsv": - rel = self.session.conn.read_csv(path) - elif format == "json": - rel = self.session.conn.read_json(path) - elif format == "parquet": - rel = self.session.conn.read_parquet(path) - else: - raise ContributionsAcceptedError - else: - rel = self.session.conn.sql(f"select * from {path}") - df = DataFrame(rel, self.session) - if schema: - if not isinstance(schema, StructType): - raise ContributionsAcceptedError - schema = cast("StructType", schema) - types, names = schema.extract_types_and_names() - df = df._cast_types(types) - df = df.toDF(names) - return df + format = format.lower() if format else _format_of(path) + if format in ("csv", "tsv"): + return self.csv(path, schema=schema, sep="\t" if format == "tsv" else None) + if format == "json": + if schema: + return _read_json_with_schema(self.session, path, _parse_schema(schema)) + rel = _read_json(self.session, path) + return DataFrame(None, self.session) if rel is None else _source_frame(rel, self.session) + if format != "parquet": + raise ContributionsAcceptedError + df = _source_frame(self.session.conn.read_parquet(path), self.session) + if not schema: + return df + # Spark finds the fields of a parquet schema by name, ignoring case. + schema = _parse_schema(schema) + return _with_schema(df, schema, _positions_by_name(df.relation.columns, schema, ignore_case=True)) def csv( # noqa: D102 self, @@ -166,24 +535,12 @@ def csv( # noqa: D102 ) -> "DataFrame": if not isinstance(path, str): raise NotImplementedError - if schema and not isinstance(schema, StructType): - raise ContributionsAcceptedError if comment: raise ContributionsAcceptedError - if inferSchema: - raise ContributionsAcceptedError - if ignoreLeadingWhiteSpace: + if _flag("ignoreLeadingWhiteSpace", ignoreLeadingWhiteSpace): raise ContributionsAcceptedError - if ignoreTrailingWhiteSpace: + if _flag("ignoreTrailingWhiteSpace", ignoreTrailingWhiteSpace): raise ContributionsAcceptedError - if nanValue: - raise ConnectionAbortedError - if positiveInf: - raise ConnectionAbortedError - if negativeInf: - raise ConnectionAbortedError - if negativeInf: - raise ConnectionAbortedError if maxColumns: raise ContributionsAcceptedError if maxCharsPerColumn: @@ -194,21 +551,17 @@ def csv( # noqa: D102 raise ContributionsAcceptedError if columnNameOfCorruptRecord: raise ContributionsAcceptedError - if multiLine: - raise ContributionsAcceptedError if charToEscapeQuoteEscaping: raise ContributionsAcceptedError if samplingRatio: raise ContributionsAcceptedError - if enforceSchema: - raise ContributionsAcceptedError if emptyValue: raise ContributionsAcceptedError if locale: raise ContributionsAcceptedError if pathGlobFilter: raise ContributionsAcceptedError - if recursiveFileLookup: + if _flag("recursiveFileLookup", recursiveFileLookup): raise ContributionsAcceptedError if modifiedBefore: raise ContributionsAcceptedError @@ -219,31 +572,76 @@ def csv( # noqa: D102 if lineSep: # We have support for custom newline, just needs to be ported to 'read_csv' raise NotImplementedError - - dtype = None - names = None - if schema: - schema = cast("StructType", schema) - dtype, names = schema.extract_types_and_names() - + from ..sql.dataframe import DataFrame, _source_frame + + has_header = _flag("header", header) + enforce = _flag("enforceSchema", enforceSchema, default=True) + infer = _flag("inferSchema", inferSchema) + parsed = _parse_schema(schema) if schema else None + if parsed is None and infer and (dateFormat or timestampFormat): + msg = "Inferring dates and timestamps with dateFormat or timestampFormat is not supported" + raise ContributionsAcceptedError(msg) + null_value = "" if nullValue is None else nullValue + options: dict[str, Any] = { + **_csv_dialect(sep, quote, escape), + "encoding": encoding, + # Spark reads the fields of a line by position, NULL past its end, and drops those past the last column. + "null_padding": True, + "strict_mode": False, + } + if _flag("multiLine", multiLine): + # DuckDB reads quoted line breaks with null padding only in its serial reader. + options["parallel"] = False + fields: list[str | None] = [] + if parsed is None: + first = _first_csv_line(self.session, path, options) + if first is None: + return DataFrame(None, self.session) + fields = first + elif has_header and not enforce: + # Spark checks the header line of each file. + for (file,) in self.session.conn.execute("SELECT file FROM glob($1) ORDER BY file", [path]).fetchall(): + first = _first_csv_line(self.session, file, options) + if first is not None: + _check_csv_header(first, parsed, file) + if parsed is not None: + names = parsed.names + elif has_header: + names = _csv_header_names(fields, nullValue) + else: + names = [f"_c{i}" for i in range(len(fields))] + if parsed is None and infer: + return _read_inferred_csv(self.session, path, names, has_header, null_value, options) + columns = [ColumnExpression(f"column{i}") for i in range(len(names))] rel = self.session.conn.read_csv( path, - header=header if isinstance(header, bool) else header == "True", - sep=sep, - dtype=dtype, - na_values=nullValue, - quotechar=quote, - escapechar=escape, - encoding=encoding, - date_format=dateFormat, - timestamp_format=timestampFormat, + auto_detect=False, + columns={column.get_name(): "VARCHAR" for column in columns}, + header=has_header, + # An unquoted empty field is NULL and a quoted one the empty string, which nullValue then makes NULL. + na_values=[""], + allow_quoted_nulls=False, + **options, ) - from ..sql.dataframe import DataFrame - - df = DataFrame(rel, self.session) - if names: - df = df.toDF(*names) - return df + df = _source_frame(rel, self.session) + texts = df.relation.select( + *(FunctionExpression("nullif", c, ConstantExpression(null_value)).alias(c.get_name()) for c in columns) + ) + if parsed is not None: + schema_fields = parsed.fields + patterns = {kind: p for kind, p in ((DateType, dateFormat), (TimestampType, timestampFormat)) if p} + else: + schema_fields, patterns = [StructField(name, StringType()) for name in names], {} + specials = ( + "NaN" if nanValue is None else nanValue, + "Inf" if positiveInf is None else positiveInf, + "-Inf" if negativeInf is None else negativeInf, + ) + values = [ + _csv_value(column, field, patterns, specials).alias(field.name) + for column, field in zip(columns, schema_fields, strict=True) + ] + return df._derive(texts.select(*values)) def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame": # noqa: D102 input = list(paths) @@ -256,10 +654,9 @@ def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame raise ContributionsAcceptedError(msg) path = input[0] rel = self.session.conn.read_parquet(path) - from ..sql.dataframe import DataFrame + from ..sql.dataframe import _source_frame - df = DataFrame(rel, self.session) - return df + return _source_frame(rel, self.session) def json( self, @@ -416,11 +813,10 @@ def json( path = [path] if isinstance(path, list): if len(path) == 1: - rel = self.session.conn.read_json(path[0]) - from .dataframe import DataFrame + rel = _read_json(self.session, path[0]) + from .dataframe import DataFrame, _source_frame - df = DataFrame(rel, self.session) - return df + return DataFrame(None, self.session) if rel is None else _source_frame(rel, self.session) raise PySparkNotImplementedError(message="Only a single path is supported for now") else: raise PySparkTypeError( diff --git a/duckdb/experimental/spark/sql/session.py b/duckdb/experimental/spark/sql/session.py index c407a9f19..65a22317d 100644 --- a/duckdb/experimental/spark/sql/session.py +++ b/duckdb/experimental/spark/sql/session.py @@ -1,8 +1,12 @@ +import datetime +import math import uuid -from collections.abc import Iterable, Sized -from typing import TYPE_CHECKING, Any, NoReturn, Union +from collections.abc import Iterable +from functools import reduce +from typing import TYPE_CHECKING, Any, NoReturn, Union, cast import duckdb +from duckdb import ColumnExpression if TYPE_CHECKING: from pandas.core.frame import DataFrame as PandasDataFrame @@ -12,13 +16,36 @@ from ..conf import SparkConf from ..context import SparkContext -from ..errors import PySparkTypeError +from ..errors import IllegalArgumentException, PySparkTypeError, PySparkValueError from ..exception import ContributionsAcceptedError +from .column import _is_numpy from .conf import RuntimeConfig -from .dataframe import DataFrame +from .dataframe import DataFrame, _source_frame from .readwriter import DataFrameReader from .streaming import DataStreamReader -from .types import StructType +from .types import ( + ArrayType, + DataType, + DecimalType, + FractionalType, + IntegralType, + MapType, + StructField, + StructType, + TimestampMillisecondNTZType, + TimestampNanosecondNTZType, + TimestampNanosecondType, + TimestampNTZType, + TimestampSecondNTZType, + TimestampType, + _has_nulltype, + _infer_schema, + _make_type_verifier, + _merge_type, + _parse_datatype_string, + _python_to_duckdb, + _replace_types, +) from .udf import UDFRegistration # In spark: @@ -31,16 +58,84 @@ # SparkContext can be compared to our Connection class, and SparkConf to our ClientContext class -# data is a List of rows -# every value in each row needs to be turned into a Value -def _combine_data_and_schema(data: Iterable[Any], schema: StructType) -> list[duckdb.Value]: - from duckdb import Value - - new_data = [] - for row in data: - new_row = [Value(x, dtype.duckdb_type) for x, dtype in zip(row, [y.dataType for y in schema], strict=False)] - new_data.append(new_row) - return new_data +def _sql_value(value: Any, data_type: DataType, parameters: list[Any]) -> str: # noqa: ANN401 + """Returns SQL for ``value`` and appends the parameters it refers to.""" + if value is None: + return "NULL" + nested = isinstance(data_type, (ArrayType, MapType, StructType)) + # duckdb.Value gives DuckDB the type of a nested value, which it would read wrongly from a dict. + parameters.append(duckdb.Value(value, data_type.duckdb_type) if nested else value) + return "?" + + +# Spark reads every pandas datetime column as a timestamp, a naive one as a wall clock time in the session time zone. +_PANDAS_TIMESTAMPS = { + TimestampNTZType: TimestampType(), + TimestampMillisecondNTZType: TimestampType(), + TimestampNanosecondNTZType: TimestampType(), + TimestampSecondNTZType: TimestampType(), + TimestampNanosecondType: TimestampType(), +} + + +def _is_missing(value: Any) -> bool: # noqa: ANN401 + """Whether ``value`` is None or a pandas missing value, which DuckDB reads as NULL.""" + if isinstance(value, float): + return math.isnan(value) + return value is None or type(value).__name__ in ("NAType", "NaTType") + + +def _python_value(value: Any) -> Any: # noqa: ANN401 + """A value of a pandas column as the Python value DuckDB reads, at any depth, which Spark's verifier can check. + + A pandas missing value is None, and NumPy values and pandas timestamps and durations are Python values. + """ + if _is_numpy(value): + value = value.tolist() + if _is_missing(value): + return None + if isinstance(value, datetime.datetime) and hasattr(value, "to_pydatetime"): + return value.to_pydatetime(warn=False) + if isinstance(value, datetime.timedelta) and hasattr(value, "to_pytimedelta"): + return value.to_pytimedelta() + if isinstance(value, dict): + return {_python_value(k): _python_value(v) for k, v in value.items()} + if type(value) in (list, tuple): + return type(value)(map(_python_value, value)) + return value + + +def _verify_pandas_values(data: "PandasDataFrame", schema: StructType) -> None: + """Raise where Spark rejects a value of ``data`` for ``schema``, which it reads through the rows of ``data``. + + A missing value is None in those rows, except in a float column, whose missing values are NaN. Spark keeps a NaN + or another pandas missing value as a value, but DuckDB reads it as NULL, so it counts as None. Spark 4.2.0 takes + a float column for an integral field only where every value is whole or, in a nullable field, missing; it never + takes one for a decimal field. + """ + for (_, series), field in zip(data.items(), schema.fields, strict=False): + verify = _make_type_verifier(field.dataType, field.nullable, name=f"field {field.name}") + if series.dtype.kind == "f": + missing = series.isna() + if isinstance(field.dataType, IntegralType): + rejected = series[(~missing & (series % 1 != 0)) | (missing & (not field.nullable))] + elif isinstance(field.dataType, DecimalType): + rejected = series + elif isinstance(field.dataType, FractionalType): + if missing.any(): + verify(None) + continue + else: + if missing.any(): + verify(None) + rejected = series[~missing] + if len(rejected): + verify(float(rejected.iloc[0])) + elif series.dtype == object: + for value in series: + verify(_python_value(value)) + elif series.isna().any(): + verify(None) class SparkSession: # noqa: D101 @@ -49,77 +144,53 @@ def __init__(self, context: SparkContext) -> None: # noqa: D107 self._context = context self._conf = RuntimeConfig(self.conn) - def _create_dataframe(self, data: Union[Iterable[Any], "PandasDataFrame"]) -> DataFrame: - try: - import pandas - - has_pandas = True - except ImportError: - has_pandas = False - if has_pandas and isinstance(data, pandas.DataFrame): - unique_name = f"pyspark_pandas_df_{uuid.uuid1()}" - self.conn.register(unique_name, data) - return DataFrame(self.conn.sql(f'select * from "{unique_name}"'), self) - - def verify_tuple_integrity(tuples: list[tuple]) -> None: - if len(tuples) <= 1: - return - expected_length = len(tuples[0]) - for i, item in enumerate(tuples[1:]): - actual_length = len(item) - if expected_length == actual_length: - continue - raise PySparkTypeError( - error_class="LENGTH_SHOULD_BE_THE_SAME", - message_parameters={ - "arg1": f"data{i}", - "arg2": f"data{i + 1}", - "arg1_length": str(expected_length), - "arg2_length": str(actual_length), - }, - ) - - if not isinstance(data, list): - data = list(data) - verify_tuple_integrity(data) - - def construct_query(tuples: Iterable) -> str: - def construct_values_list(row: Sized, start_param_idx: int) -> str: - parameter_count = len(row) - parameters = [f"${x + start_param_idx}" for x in range(parameter_count)] - parameters = "(" + ", ".join(parameters) + ")" - return parameters - - row_size = len(tuples[0]) - values_list = [construct_values_list(x, 1 + (i * row_size)) for i, x in enumerate(tuples)] - values_list = ", ".join(values_list) - - query = f""" - select * from (values {values_list}) - """ - return query - - query = construct_query(data) - - def construct_parameters(tuples: Iterable) -> list[list]: - parameters = [] - for row in tuples: - parameters.extend(list(row)) - return parameters - - parameters = construct_parameters(data) - - rel = self.conn.sql(query, params=parameters) + def _create_dataframe(self, data: list[Any], schema: StructType, names: list[str] | None = None) -> DataFrame: + if not schema.fields: + msg = "DuckDB can not create a DataFrame without columns" + raise ContributionsAcceptedError(msg) + column_names = list(schema.names) + for i, name in enumerate(names or []): + column_names[i] = name + convert = _python_to_duckdb(schema) + rows = [convert(row) for row in data] or [tuple(None for _ in schema.fields)] + parameters: list[Any] = [] + values = ", ".join( + "(" + + ", ".join(_sql_value(v, f.dataType, parameters) for v, f in zip(row, schema.fields, strict=True)) + + ")" + for row in rows + ) + rel = self.conn.sql(f"select * from (values {values})", params=parameters) + rel = rel.project( + *[ + ColumnExpression(column).cast(field.dataType.duckdb_type).alias(name) + for column, field, name in zip(rel.columns, schema.fields, column_names, strict=True) + ] + ) + if not data: + rel = rel.limit(0) return DataFrame(rel, self) + def _inferSchemaFromList(self, data: list[Any], names: list[str] | None = None) -> StructType: + if not data: + raise PySparkValueError(error_class="CANNOT_INFER_EMPTY_SCHEMA", message_parameters={}) + schema = reduce(_merge_type, (_infer_schema(row, names) for row in data)) + if _has_nulltype(schema): + raise PySparkValueError(error_class="CANNOT_DETERMINE_TYPE", message_parameters={}) + return schema + def _createDataFrameFromPandas( self, data: "PandasDataFrame", types: list[str] | None, names: list[str] | None ) -> DataFrame: - df = self._create_dataframe(data) + unique_name = f"pyspark_pandas_df_{uuid.uuid1()}" + self.conn.register(unique_name, data) + df = DataFrame(self.conn.sql(f'select * from "{unique_name}"'), self) # Cast to types if types: df = df._cast_types(*types) + else: + df = df._with_types(cast("StructType", _replace_types(df.schema, _PANDAS_TIMESTAMPS))) # Alias to names if names: df = df.toDF(*names) @@ -128,16 +199,12 @@ def _createDataFrameFromPandas( def createDataFrame( # noqa: D102 self, data: Union["PandasDataFrame", Iterable[Any]], - schema: StructType | list[str] | None = None, + schema: DataType | str | list[str] | None = None, samplingRatio: float | None = None, verifySchema: bool = True, ) -> DataFrame: if samplingRatio: raise NotImplementedError - if not verifySchema: - raise NotImplementedError - types = None - names = None if isinstance(data, DataFrame): raise PySparkTypeError( @@ -145,11 +212,15 @@ def createDataFrame( # noqa: D102 message_parameters={"arg_name": "data"}, ) - if schema: - if isinstance(schema, StructType): - types, names = schema.extract_types_and_names() - else: - names = schema + if isinstance(schema, str): + schema = _parse_datatype_string(schema) + elif isinstance(schema, (list, tuple)): + schema = list(schema) + elif schema is not None and not isinstance(schema, DataType): + raise PySparkTypeError( + error_class="NOT_LIST_OR_NONE_OR_STRUCT", + message_parameters={"arg_name": "schema", "arg_type": type(schema).__name__}, + ) try: import pandas @@ -157,36 +228,28 @@ def createDataFrame( # noqa: D102 has_pandas = True except ImportError: has_pandas = False - # Falsey check on pandas dataframe is not defined, so first check if it's not a pandas dataframe - # Then check if 'data' is None or [] if has_pandas and isinstance(data, pandas.DataFrame): - return self._createDataFrameFromPandas(data, types, names) - - # Finally check if a schema was provided - is_empty = False - if not data and names: - # Create NULLs for every type in our dataframe - is_empty = True - data = [tuple(None for _ in names)] - - if schema and isinstance(schema, StructType): - # Transform the data into Values to combine the data+schema - data = _combine_data_and_schema(data, schema) - - df = self._create_dataframe(data) - if is_empty: - rel = df.relation - # Add impossible where clause - rel = rel.filter("1=0") - df = DataFrame(rel, self) + if isinstance(schema, StructType): + if verifySchema: + _verify_pandas_values(data, schema) + types, names = schema.extract_types_and_names() + return self._createDataFrameFromPandas(data, types, names) + if isinstance(schema, DataType): + msg = "A pandas DataFrame with a schema that is not a StructType is not supported" + raise ContributionsAcceptedError(msg) + return self._createDataFrameFromPandas(data, None, schema) - # Cast to types - if types: - df = df._cast_types(*types) - # Alias to names - if names: - df = df.toDF(*names) - return df + if not isinstance(data, list): + data = list(data) + if isinstance(schema, DataType) and verifySchema: + verify = _make_type_verifier(schema, name=None if isinstance(schema, StructType) else "field value") + for obj in data: + verify(obj) + if isinstance(schema, StructType): + return self._create_dataframe(data, schema) + if isinstance(schema, DataType): + return self._create_dataframe([(obj,) for obj in data], StructType([StructField("value", schema)])) + return self._create_dataframe(data, self._inferSchemaFromList(data, names=schema), names=schema) def newSession(self) -> "SparkSession": # noqa: D102 return SparkSession(self._context) @@ -198,14 +261,15 @@ def range( # noqa: D102 step: int = 1, numPartitions: int | None = None, ) -> "DataFrame": - if numPartitions: - raise ContributionsAcceptedError - + if step == 0: + msg = "requirement failed: step (0) cannot be 0" + raise IllegalArgumentException(msg) if end is None: end = start start = 0 - - return DataFrame(self.conn.table_function("range", parameters=[start, end, step]), self) + # numPartitions only splits Spark's rows across partitions; it does not change them. + relation = self.conn.table_function("range", parameters=[start, end, step]) + return _source_frame(relation.select(ColumnExpression("range").alias("id")), self) def sql(self, sqlQuery: str, **kwargs: Any) -> DataFrame: # noqa: D102, ANN401 if kwargs: diff --git a/duckdb/experimental/spark/sql/type_utils.py b/duckdb/experimental/spark/sql/type_utils.py index 65e17662c..3e753d550 100644 --- a/duckdb/experimental/spark/sql/type_utils.py +++ b/duckdb/experimental/spark/sql/type_utils.py @@ -74,7 +74,8 @@ "struct": StructType, "map": MapType, # union - # enum + # Spark has no enum type; a pandas categorical reads as string. + "enum": StringType, # null (???) "float": FloatType, "double": DoubleType, diff --git a/duckdb/experimental/spark/sql/types.py b/duckdb/experimental/spark/sql/types.py index e9609627a..9078df3ca 100644 --- a/duckdb/experimental/spark/sql/types.py +++ b/duckdb/experimental/spark/sql/types.py @@ -3,24 +3,29 @@ import calendar import datetime +import decimal +import json import math import re import time +import uuid from builtins import tuple -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping +from functools import reduce from types import MappingProxyType -from typing import Any, ClassVar, NoReturn, TypeVar, Union, cast, overload - -from typing_extensions import Self +from typing import Any, ClassVar, NoReturn, Self, TypeVar, Union, cast, overload import duckdb from duckdb.sqltypes import DuckDBPyType +from ..errors import ParseException, PySparkNotImplementedError, PySparkTypeError, PySparkValueError from ..exception import ContributionsAcceptedError T = TypeVar("T") U = TypeVar("U") +_JVM_INT_MAX = 2**31 - 1 + __all__ = [ "ArrayType", "BinaryType", @@ -88,10 +93,10 @@ def simpleString(self) -> str: # noqa: D102 return self.typeName() def jsonValue(self) -> str | dict[str, Any]: # noqa: D102 - raise ContributionsAcceptedError + return self.typeName() def json(self) -> str: # noqa: D102 - raise ContributionsAcceptedError + return json.dumps(self.jsonValue(), separators=(",", ":"), sort_keys=True) def needConversion(self) -> bool: """Does this type needs conversion between Python object and internal SQL object. @@ -222,10 +227,6 @@ class TimestampType(AtomicType, metaclass=DataTypeSingleton): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("TIMESTAMPTZ")) - @classmethod - def typeName(cls) -> str: # noqa: D102 - return "timestamptz" - def needConversion(self) -> bool: # noqa: D102 return True @@ -271,7 +272,7 @@ def needConversion(self) -> bool: # noqa: D102 @classmethod def typeName(cls) -> str: # noqa: D102 - return "timestamp" + return "timestamp_ntz" def toInternal(self, dt: datetime.datetime) -> int: # noqa: D102 if dt is not None: @@ -373,6 +374,8 @@ def __init__(self, precision: int = 10, scale: int = 0) -> None: # noqa: D107 def simpleString(self) -> str: # noqa: D102 return f"decimal({int(self.precision):d},{int(self.scale):d})" + jsonValue = simpleString + def __repr__(self) -> str: # noqa: D105 return f"DecimalType({int(self.precision):d},{int(self.scale):d})" @@ -407,7 +410,8 @@ class UnsignedByteType(IntegralType): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("UTINYINT")) - def simpleString(self) -> str: # noqa: D102 + @classmethod + def typeName(cls) -> str: # noqa: D102 return "utinyint" @@ -427,7 +431,8 @@ class UnsignedShortType(IntegralType): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("USMALLINT")) - def simpleString(self) -> str: # noqa: D102 + @classmethod + def typeName(cls) -> str: # noqa: D102 return "usmallint" @@ -438,7 +443,7 @@ def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("INTEGER")) def simpleString(self) -> str: # noqa: D102 - return "integer" + return "int" class UnsignedIntegerType(IntegralType): @@ -447,7 +452,8 @@ class UnsignedIntegerType(IntegralType): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("UINTEGER")) - def simpleString(self) -> str: # noqa: D102 + @classmethod + def typeName(cls) -> str: # noqa: D102 return "uinteger" @@ -475,7 +481,8 @@ class UnsignedLongType(IntegralType): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("UBIGINT")) - def simpleString(self) -> str: # noqa: D102 + @classmethod + def typeName(cls) -> str: # noqa: D102 return "ubigint" @@ -489,7 +496,8 @@ class HugeIntegerType(IntegralType): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("HUGEINT")) - def simpleString(self) -> str: # noqa: D102 + @classmethod + def typeName(cls) -> str: # noqa: D102 return "hugeint" @@ -503,7 +511,8 @@ class UnsignedHugeIntegerType(IntegralType): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("UHUGEINT")) - def simpleString(self) -> str: # noqa: D102 + @classmethod + def typeName(cls) -> str: # noqa: D102 return "uhugeint" @@ -513,7 +522,8 @@ class TimeType(IntegralType): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("TIMETZ")) - def simpleString(self) -> str: # noqa: D102 + @classmethod + def typeName(cls) -> str: # noqa: D102 return "timetz" @@ -523,7 +533,8 @@ class TimeNTZType(IntegralType): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("TIME")) - def simpleString(self) -> str: # noqa: D102 + @classmethod + def typeName(cls) -> str: # noqa: D102 return "time" @@ -533,7 +544,8 @@ class TimeNSType(IntegralType): def __init__(self) -> None: # noqa: D107 super().__init__(DuckDBPyType("TIME_NS")) - def simpleString(self) -> str: # noqa: D102 + @classmethod + def typeName(cls) -> str: # noqa: D102 return "time_ns" @@ -557,7 +569,6 @@ class DayTimeIntervalType(AtomicType): _inverted_fields: Mapping[int, str] = MappingProxyType(dict(zip(_fields.values(), _fields.keys(), strict=False))) def __init__(self, startField: int | None = None, endField: int | None = None) -> None: # noqa: D107 - super().__init__(DuckDBPyType("INTERVAL")) if startField is None and endField is None: # Default matched to scala side. startField = DayTimeIntervalType.DAY @@ -572,6 +583,14 @@ def __init__(self, startField: int | None = None, endField: int | None = None) - self.startField = cast("int", startField) self.endField = cast("int", endField) + @property + def duckdb_type(self) -> DuckDBPyType: # noqa: D102 + # A DuckDB interval has no fields, so a DataFrame could only report it as interval day to second. + if (self.startField, self.endField) != (DayTimeIntervalType.DAY, DayTimeIntervalType.SECOND): + msg = f"The Spark type {self.simpleString()} is not supported" + raise ContributionsAcceptedError(msg) + return DuckDBPyType("INTERVAL") + def _str_repr(self) -> str: fields = DayTimeIntervalType._fields start_field_name = fields[self.startField] @@ -582,6 +601,7 @@ def _str_repr(self) -> str: return f"interval {start_field_name} to {end_field_name}" simpleString = _str_repr + jsonValue = _str_repr def __repr__(self) -> str: # noqa: D105 return f"{type(self).__name__}({int(self.startField):d}, {int(self.endField):d})" @@ -628,6 +648,13 @@ def simpleString(self) -> str: # noqa: D102 def __repr__(self) -> str: # noqa: D105 return f"ArrayType({self.elementType}, {self.containsNull!s})" + def jsonValue(self) -> dict[str, Any]: # noqa: D102 + return { + "type": self.typeName(), + "elementType": self.elementType.jsonValue(), + "containsNull": self.containsNull, + } + def needConversion(self) -> bool: # noqa: D102 return self.elementType.needConversion() @@ -680,6 +707,14 @@ def simpleString(self) -> str: # noqa: D102 def __repr__(self) -> str: # noqa: D105 return f"MapType({self.keyType}, {self.valueType}, {self.valueContainsNull!s})" + def jsonValue(self) -> dict[str, Any]: # noqa: D102 + return { + "type": self.typeName(), + "keyType": self.keyType.jsonValue(), + "valueType": self.valueType.jsonValue(), + "valueContainsNull": self.valueContainsNull, + } + def needConversion(self) -> bool: # noqa: D102 return self.keyType.needConversion() or self.valueType.needConversion() @@ -723,7 +758,6 @@ def __init__( # noqa: D107 nullable: bool = True, metadata: dict[str, Any] | None = None, ) -> None: - super().__init__(dataType.duckdb_type) assert isinstance(dataType, DataType), f"dataType {dataType} should be an instance of {DataType}" assert isinstance(name, str), f"field name {name} should be a string" self.name = name @@ -731,12 +765,24 @@ def __init__( # noqa: D107 self.nullable = nullable self.metadata = metadata or {} + @property + def duckdb_type(self) -> DuckDBPyType: # noqa: D102 + return self.dataType.duckdb_type + def simpleString(self) -> str: # noqa: D102 return f"{self.name}:{self.dataType.simpleString()}" def __repr__(self) -> str: # noqa: D105 return f"StructField('{self.name}', {self.dataType}, {self.nullable!s})" + def jsonValue(self) -> dict[str, Any]: # noqa: D102 + return { + "name": self.name, + "type": self.dataType.jsonValue(), + "nullable": self.nullable, + "metadata": self.metadata, + } + def needConversion(self) -> bool: # noqa: D102 return self.dataType.needConversion() @@ -747,8 +793,7 @@ def fromInternal(self, obj: T) -> T: # noqa: D102 return self.dataType.fromInternal(obj) def typeName(self) -> str: # type: ignore[override] # noqa: D102 - msg = "StructField does not have typeName. Use typeName on its type explicitly instead." - raise TypeError(msg) + raise PySparkTypeError(error_class="INVALID_TYPENAME_CALL", message_parameters={}) class StructType(DataType): @@ -779,9 +824,6 @@ class StructType(DataType): False """ - def _update_internal_duckdb_type(self) -> None: - self.duckdb_type = duckdb.struct_type(dict(zip(self.names, [x.duckdb_type for x in self.fields], strict=False))) - def __init__(self, fields: list[StructField] | None = None) -> None: # noqa: D107 if not fields: self.fields = [] @@ -793,7 +835,18 @@ def __init__(self, fields: list[StructField] | None = None) -> None: # noqa: D1 # Precalculated list of fields that need conversion with fromInternal/toInternal functions self._needConversion = [f.needConversion() for f in self] self._needSerializeAnyField = any(self._needConversion) - super().__init__(duckdb.struct_type(dict(zip(self.names, [x.duckdb_type for x in self.fields], strict=False)))) + + @property + def duckdb_type(self) -> DuckDBPyType: # noqa: D102 + if self.names and not any(self.names): + # The struct DuckDB builds without field names, such as the elements of list_zip. + return DuckDBPyType(f"TUPLE({', '.join(str(f.duckdb_type) for f in self.fields)})") + folded = [name.casefold() for name in self.names] + if len(set(folded)) < len(folded): + # DuckDB matches struct field names case insensitively, so it cannot hold two fields with one name. + msg = f"DuckDB has no struct type with two fields named alike: {self.simpleString()}" + raise PySparkNotImplementedError(message=msg) + return duckdb.struct_type(dict(zip(self.names, [x.duckdb_type for x in self.fields], strict=True))) @overload def add( @@ -860,6 +913,8 @@ def add( if isinstance(field, str) and data_type is None: msg = "Must specify DataType if passing name of struct_field to create." raise ValueError(msg) + elif isinstance(data_type, str): + data_type_f = _parse_type_name(data_type) else: data_type_f = data_type self.fields.append(StructField(field, data_type_f, nullable, metadata)) @@ -867,7 +922,6 @@ def add( # Precalculated list of fields that need conversion with fromInternal/toInternal functions self._needConversion = [f.needConversion() for f in self] self._needSerializeAnyField = any(self._needConversion) - self._update_internal_duckdb_type() return self def __iter__(self) -> Iterator[StructField]: @@ -904,6 +958,9 @@ def simpleString(self) -> str: # noqa: D102 def __repr__(self) -> str: # noqa: D105 return "StructType([{}])".format(", ".join(str(field) for field in self)) + def jsonValue(self) -> dict[str, Any]: # noqa: D102 + return {"type": self.typeName(), "fields": [f.jsonValue() for f in self]} + def __contains__(self, item: str) -> bool: # noqa: D105 return item in self.names @@ -926,13 +983,13 @@ def fieldNames(self) -> list[str]: """ return list(self.names) - def treeString(self, level: int | None = None) -> str: + def treeString(self, maxDepth: int = _JVM_INT_MAX) -> str: """Returns a string representation of the schema in tree format. Parameters ---------- - level : int, optional - Maximum depth to print. If None, prints all levels. + maxDepth : int, optional + Maximum depth to print. Prints all levels when not positive. Returns: ------- @@ -945,57 +1002,10 @@ def treeString(self, level: int | None = None) -> str: >>> print(schema.treeString()) root |-- age: integer (nullable = true) + """ - - def _tree_string(schema: "StructType", depth: int = 0, max_depth: int | None = None) -> list[str]: - """Recursively build tree string lines.""" - lines = [] - if depth == 0: - lines.append("root") - - if max_depth is not None and depth >= max_depth: - return lines - - for field in schema.fields: - indent = " " * depth - prefix = " |-- " - nullable_str = "true" if field.nullable else "false" - - # Handle nested StructType - if isinstance(field.dataType, StructType): - lines.append(f"{indent}{prefix}{field.name}: struct (nullable = {nullable_str})") - # Recursively handle nested struct - don't skip any lines, root only appears at depth 0 - nested_lines = _tree_string(field.dataType, depth + 1, max_depth) - lines.extend(nested_lines) - # Handle ArrayType - elif isinstance(field.dataType, ArrayType): - element_type = field.dataType.elementType - if isinstance(element_type, StructType): - lines.append(f"{indent}{prefix}{field.name}: array (nullable = {nullable_str})") - lines.append( - f"{indent} | |-- element: struct (containsNull = {field.dataType.containsNull})" - ) - nested_lines = _tree_string(element_type, depth + 2, max_depth) - lines.extend(nested_lines) - else: - type_str = element_type.simpleString() - lines.append(f"{indent}{prefix}{field.name}: array<{type_str}> (nullable = {nullable_str})") - # Handle MapType - elif isinstance(field.dataType, MapType): - key_type = field.dataType.keyType.simpleString() - value_type = field.dataType.valueType.simpleString() - lines.append( - f"{indent}{prefix}{field.name}: map<{key_type},{value_type}> (nullable = {nullable_str})" - ) - # Handle simple types - else: - type_str = field.dataType.simpleString() - lines.append(f"{indent}{prefix}{field.name}: {type_str} (nullable = {nullable_str})") - - return lines - - lines = _tree_string(self, 0, level) - return "\n".join(lines) + lines = ["root", *_tree_lines(self, " |", maxDepth if maxDepth > 0 else _JVM_INT_MAX)] + return "\n".join(lines) + "\n" def needConversion(self) -> bool: # noqa: D102 # We need convert Row()/namedtuple into tuple() @@ -1055,6 +1065,46 @@ def fromInternal(self, obj: tuple) -> "Row": # noqa: D102 return _create_row(self.names, values) +def _tree_type_name(data_type: DataType) -> str: + if isinstance(data_type, (DecimalType, DayTimeIntervalType)): + return data_type.simpleString() + return data_type.typeName() + + +def _tree_lines(data_type: DataType, prefix: str, max_depth: int) -> list[str]: + """Returns the lines treeString prints for the children of a nested type.""" + if max_depth <= 0: + return [] + if isinstance(data_type, StructType): + children = [ + (_escape_meta_characters(f.name), f.dataType, f" (nullable = {str(f.nullable).lower()})") + for f in data_type.fields + ] + elif isinstance(data_type, ArrayType): + children = [("element", data_type.elementType, f" (containsNull = {str(data_type.containsNull).lower()})")] + elif isinstance(data_type, MapType): + children = [ + ("key", data_type.keyType, ""), + ("value", data_type.valueType, f" (valueContainsNull = {str(data_type.valueContainsNull).lower()})"), + ] + else: + return [] + lines = [] + for name, child_type, flag in children: + lines.append(f"{prefix}-- {name}: {_tree_type_name(child_type)}{flag}") + lines.extend(_tree_lines(child_type, f"{prefix} |", max_depth - 1)) + return lines + + +_META_CHARACTERS = str.maketrans( + {"\n": "\\n", "\r": "\\r", "\t": "\\t", "\f": "\\f", "\b": "\\b", "\v": "\\v", "\a": "\\a"} +) + + +def _escape_meta_characters(text: str) -> str: + return text.translate(_META_CHARACTERS) + + class UnionType(DataType): def __init__(self) -> None: raise ContributionsAcceptedError @@ -1158,6 +1208,714 @@ def __eq__(self, other: object) -> bool: _INTERVAL_DAYTIME = re.compile(r"interval (day|hour|minute|second)( to (day|hour|minute|second))?") +_DDL_TOKEN = re.compile(r"""\s*(`(?:[^`]|``)*`|'(?:[^'\\]|\\.)*'|"(?:[^"\\]|\\.)*"|\w+|\S)""") + +_DUCKDB_ONLY_ATOMIC_TYPES: list[type[DataType]] = [ + UnsignedByteType, + UnsignedShortType, + UnsignedIntegerType, + UnsignedLongType, + HugeIntegerType, + UnsignedHugeIntegerType, + UUIDType, + BitstringType, + TimeType, + TimeNSType, + TimestampSecondNTZType, + TimestampMillisecondNTZType, + TimestampNanosecondNTZType, + TimestampNanosecondType, +] + +# DuckDB only types are known by the name the module prints for them; Spark's names take precedence. +_DDL_ATOMIC_TYPES: dict[str, type[DataType]] = {t().simpleString(): t for t in _DUCKDB_ONLY_ATOMIC_TYPES} | { + "boolean": BooleanType, + "tinyint": ByteType, + "byte": ByteType, + "smallint": ShortType, + "short": ShortType, + "int": IntegerType, + "integer": IntegerType, + "bigint": LongType, + "long": LongType, + "float": FloatType, + "real": FloatType, + "double": DoubleType, + "string": StringType, + "binary": BinaryType, + "date": DateType, + "timestamp": TimestampType, + "timestamp_ltz": TimestampType, + "timestamp_ntz": TimestampNTZType, + "time": TimeNTZType, + "variant": VariantType, + "void": NullType, +} + +# The types by the name typeName() gives them; Spark's names take precedence. +_TYPE_NAMES: dict[str, type[DataType]] = { + t.typeName(): t for t in [*_DUCKDB_ONLY_ATOMIC_TYPES, TimeNTZType] +} | _all_atomic_types + + +def _parse_type_name(name: str) -> DataType: + """Reads a type from its JSON name, which is what Spark expects of a str given to ``StructType.add``.""" + if name in _TYPE_NAMES: + return _TYPE_NAMES[name]() + if match := _FIXED_DECIMAL.match(name): + return DecimalType(int(match.group(1)), int(match.group(2))) + if _INTERVAL_DAYTIME.match(name): + return _parse_datatype_string(name) + raise PySparkValueError(error_class="CANNOT_PARSE_DATATYPE", message_parameters={"msg": name}) + + +_DDL_DECIMAL_TYPES = frozenset(["decimal", "dec", "numeric"]) +_DDL_UNSUPPORTED_TYPES = frozenset(["char", "character", "varchar"]) +_DDL_DAY_TIME_FIELDS = DayTimeIntervalType._inverted_fields +_DDL_YEAR_MONTH_FIELDS = {"year": 0, "month": 1} + + +def _is_ddl_name(token: str | None) -> bool: + if token is None: + return False + return re.fullmatch(r"\w+", token) is not None or (len(token) > 1 and token[0] == token[-1] == "`") + + +class _DDLParser: + """Parser for Spark's DDL syntax for types and schemas, such as ``a int, b array``.""" + + def __init__(self, text: str, cast: bool = False) -> None: + self.cast = cast + self.tokens: list[str] = [] + pos = 0 + while (match := _DDL_TOKEN.match(text, pos)) is not None: + self.tokens.append(match.group(1)) + pos = match.end() + self.pos = 0 + + def parse(self) -> DataType: + """Parses a list of fields, with or without a colon after each name, or else a single type, as Spark does. + + Text is a list of fields if it is empty or a colon or a name follows its first name, except an interval type, + the one type of two names. The type of a cast is a single type. + """ + second = self.tokens[1].lower() if len(self.tokens) > 1 else None + is_interval = self._peek() == "interval" and second in {*_DDL_DAY_TIME_FIELDS, *_DDL_YEAR_MONTH_FIELDS} + is_fields = not self.tokens or ((second == ":" or _is_ddl_name(second)) and not is_interval) + result = self._fields() if is_fields and not self.cast else self._type() + if self.pos < len(self.tokens): + self._fail() + return result + + def _peek(self) -> str | None: + return self.tokens[self.pos].lower() if self.pos < len(self.tokens) else None + + def _next(self) -> str: + if self.pos == len(self.tokens): + self._fail() + self.pos += 1 + return self.tokens[self.pos - 1] + + def _accept(self, token: str) -> bool: + if self._peek() == token: + self.pos += 1 + return True + return False + + def _expect(self, token: str) -> None: + if not self._accept(token): + self._fail() + + def _fail(self) -> NoReturn: + error = f"'{self.tokens[self.pos]}'" if self.pos < len(self.tokens) else "end of input" + raise ParseException(error_class="PARSE_SYNTAX_ERROR", message_parameters={"error": error, "hint": ""}) + + def _fields(self, closing: str | None = None) -> StructType: + fields = [] + if self._peek() != closing: + fields.append(self._field()) + while self._accept(","): + fields.append(self._field()) + return StructType(fields) + + def _field(self) -> StructField: + name = self._identifier() + self._accept(":") + data_type = self._type() + nullable = True + if self._accept("not"): + self._expect("null") + nullable = False + metadata = {"comment": self._string()} if self._accept("comment") else None + return StructField(name, data_type, nullable, metadata) + + def _string(self) -> str: + token = self._next() + if len(token) < 2 or token[0] not in "'\"" or token[-1] != token[0]: + self.pos -= 1 + self._fail() + return re.sub(r"\\(.)", r"\1", token[1:-1]) + + def _identifier(self) -> str: + token = self._next() + if not _is_ddl_name(token): + self.pos -= 1 + self._fail() + return token[1:-1].replace("``", "`") if token[0] == "`" else token + + def _integer(self) -> int: + token = self._next() + if not token.isdigit(): + self.pos -= 1 + self._fail() + return int(token) + + def _type(self) -> DataType: + name = self._next().lower() + if name == "array": + self._expect("<") + element_type = self._type() + self._expect(">") + return ArrayType(element_type) + if name == "map": + self._expect("<") + key_type = self._type() + self._expect(",") + value_type = self._type() + self._expect(">") + return MapType(key_type, value_type) + if name == "struct": + self._expect("<") + struct = self._fields(">") + self._expect(">") + return struct + if name in _DDL_DECIMAL_TYPES: + return self._decimal() + if name == "interval": + return self._interval() + if name in _DDL_ATOMIC_TYPES: + return _DDL_ATOMIC_TYPES[name]() + if name in _DDL_UNSUPPORTED_TYPES: + self._expect("(") + self._integer() + self._expect(")") + if not self.cast: + msg = f"The Spark type {name} is not supported" + raise ContributionsAcceptedError(msg) + # Spark reads a char or varchar of any length as a string where it casts. + return StringType() + if not re.fullmatch(r"\w+", name): + self.pos -= 1 + self._fail() + raise ParseException(error_class="UNSUPPORTED_DATATYPE", message_parameters={"typeName": f'"{name.upper()}"'}) + + def _decimal(self) -> DecimalType: + if not self._accept("("): + return DecimalType() + start = self.pos + precision = self._integer() + scale = self._integer() if self._accept(",") else 0 + self._expect(")") + if not 0 < precision <= 38 or scale > precision: + self.pos = start + self._fail() + return DecimalType(precision, scale) + + def _interval(self) -> DataType: + if self._peek() not in _DDL_DAY_TIME_FIELDS and self._peek() not in _DDL_YEAR_MONTH_FIELDS: + # Spark 4 reads a bare interval as a calendar interval, a type the module lacks. + msg = "The Spark type interval is not supported" + raise ContributionsAcceptedError(msg) + start = self._next().lower() + fields = _DDL_DAY_TIME_FIELDS if start in _DDL_DAY_TIME_FIELDS else _DDL_YEAR_MONTH_FIELDS + end = start + if self._accept("to"): + end = self._next().lower() + if end not in fields or fields[end] <= fields[start]: + self.pos -= 2 + self._fail() + if fields is _DDL_YEAR_MONTH_FIELDS: + msg = f"The Spark type interval {start} to {end} is not supported" + raise ContributionsAcceptedError(msg) + return DayTimeIntervalType(fields[start], fields[end]) + + +def _parse_datatype_string(s: str, cast: bool = False) -> DataType: + """Parses a type or schema in Spark's DDL syntax. + + A list of fields becomes a :class:`StructType`, with or without a colon between name and type + (``a int, b string`` or ``a: int, b: string``); anything else must be a single type, such as + ``int``, ``array`` or ``struct``. Type names are case insensitive. ``cast`` reads + the type of a cast, where Spark allows char and varchar. + """ + return _DDLParser(s, cast).parse() + + +_type_mappings: dict[type, type[DataType]] = { + type(None): NullType, + bool: BooleanType, + int: LongType, + float: DoubleType, + str: StringType, + bytearray: BinaryType, + bytes: BinaryType, + decimal.Decimal: DecimalType, + datetime.date: DateType, + datetime.datetime: TimestampType, + datetime.time: TimeNTZType, + datetime.timedelta: DayTimeIntervalType, +} + + +def _infer_type(obj: Any) -> DataType: # noqa: ANN401 + if obj is None: + return NullType() + + data_type = _type_mappings.get(type(obj)) + if data_type is DecimalType: + # The precision and scale of each value may differ from row to row. + return DecimalType(38, 18) + if data_type is not None: + return data_type() + + if isinstance(obj, dict): + key_type: DataType = NullType() + value_type: DataType = NullType() + for key, value in obj.items(): + if key is not None: + key_type = _merge_type(key_type, _infer_type(key)) + if value is not None: + value_type = _merge_type(value_type, _infer_type(value)) + return MapType(key_type, value_type) + if isinstance(obj, list): + if obj: + return ArrayType(reduce(_merge_type, (_infer_type(v) for v in obj))) + return ArrayType(NullType()) + try: + return _infer_schema(obj) + except TypeError: + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": type(obj).__name__}, + ) from None + + +def _infer_schema(row: Any, names: list[str] | None = None) -> StructType: # noqa: ANN401 + if isinstance(row, dict): + items = sorted(row.items()) + elif isinstance(row, (tuple, list)): + if hasattr(row, "__fields__"): + items = zip(row.__fields__, tuple(row), strict=False) + elif hasattr(row, "_fields"): + items = zip(row._fields, tuple(row), strict=False) + else: + names = list(names or []) + names.extend(f"_{i}" for i in range(len(names) + 1, len(row) + 1)) + items = zip(names, row, strict=False) + elif hasattr(row, "__dict__"): + items = sorted(row.__dict__.items()) + else: + raise PySparkTypeError( + error_class="CANNOT_INFER_SCHEMA_FOR_TYPE", + message_parameters={"data_type": type(row).__name__}, + ) + + fields = [] + for name, value in items: + try: + fields.append(StructField(name, _infer_type(value), True)) + except TypeError: # noqa: PERF203 + raise PySparkTypeError( + error_class="CANNOT_INFER_TYPE_FOR_FIELD", + message_parameters={"field_name": name}, + ) from None + return StructType(fields) + + +def _has_nulltype(dt: DataType) -> bool: + if isinstance(dt, StructType): + return any(_has_nulltype(f.dataType) for f in dt.fields) + if isinstance(dt, ArrayType): + return _has_nulltype(dt.elementType) + if isinstance(dt, MapType): + return _has_nulltype(dt.keyType) or _has_nulltype(dt.valueType) + return isinstance(dt, NullType) + + +def _merge_type(a: DataType, b: DataType) -> DataType: + if isinstance(a, NullType): + return b + if isinstance(b, NullType): + return a + if isinstance(a, TimestampType) and isinstance(b, TimestampNTZType): + return a + if isinstance(a, TimestampNTZType) and isinstance(b, TimestampType): + return b + if isinstance(a, AtomicType) and isinstance(b, StringType): + return b + if isinstance(a, StringType) and isinstance(b, AtomicType): + return a + if type(a) is not type(b): + raise PySparkTypeError( + error_class="CANNOT_MERGE_TYPE", + message_parameters={"data_type1": type(a).__name__, "data_type2": type(b).__name__}, + ) + + if isinstance(a, StructType): + b_types = {f.name: f.dataType for f in cast("StructType", b).fields} + fields = [StructField(f.name, _merge_type(f.dataType, b_types.get(f.name, NullType()))) for f in a.fields] + names = {f.name for f in fields} + fields.extend(StructField(name, data_type) for name, data_type in b_types.items() if name not in names) + return StructType(fields) + if isinstance(a, ArrayType): + return ArrayType(_merge_type(a.elementType, cast("ArrayType", b).elementType)) + if isinstance(a, MapType): + b = cast("MapType", b) + return MapType(_merge_type(a.keyType, b.keyType), _merge_type(a.valueType, b.valueType)) + return a + + +_acceptable_types: dict[type[DataType], tuple[type, ...]] = { + NullType: (type(None),), + BooleanType: (bool,), + ByteType: (int,), + ShortType: (int,), + IntegerType: (int,), + LongType: (int,), + UnsignedByteType: (int,), + UnsignedShortType: (int,), + UnsignedIntegerType: (int,), + UnsignedLongType: (int,), + HugeIntegerType: (int,), + UnsignedHugeIntegerType: (int,), + FloatType: (float,), + DoubleType: (float,), + DecimalType: (decimal.Decimal,), + StringType: (str,), + BitstringType: (str,), + UUIDType: (uuid.UUID,), + BinaryType: (bytearray, bytes), + DateType: (datetime.date, datetime.datetime), + TimeType: (datetime.time,), + TimeNTZType: (datetime.time,), + TimeNSType: (datetime.time,), + TimestampType: (datetime.datetime,), + TimestampNanosecondType: (datetime.datetime,), + TimestampNTZType: (datetime.datetime,), + TimestampSecondNTZType: (datetime.datetime,), + TimestampMillisecondNTZType: (datetime.datetime,), + TimestampNanosecondNTZType: (datetime.datetime,), + DayTimeIntervalType: (datetime.timedelta,), + ArrayType: (list, tuple), + MapType: (dict,), + StructType: (tuple, list, dict), +} + +_integral_bounds: dict[type[DataType], tuple[int, int]] = { + ByteType: (-128, 127), + ShortType: (-32768, 32767), + IntegerType: (-2147483648, 2147483647), + LongType: (-9223372036854775808, 9223372036854775807), +} + + +def _make_type_verifier(data_type: DataType, nullable: bool = True, name: str | None = None) -> Callable[[Any], None]: + """Returns a function that raises when a Python value does not fit ``data_type``. + + Ported from ``pyspark.sql.types._make_type_verifier`` in PySpark 3.5.9, with its error classes: + a wrong Python type raises ``PySparkTypeError``, a None where the type is not nullable and an + integer out of range raise ``PySparkValueError``. + """ + + def new_name(n: str) -> str: + return f"field {n}" if name is None else f"field {n} in {name}" + + def verify_acceptable_types(obj: Any) -> None: # noqa: ANN401 + # Subclasses are rejected, as in Spark, so a bool is not accepted as an int. + if type(obj) not in _acceptable_types[type(data_type)]: + raise PySparkTypeError( + error_class="CANNOT_ACCEPT_OBJECT_IN_TYPE", + message_parameters={ + "data_type": str(data_type), + "obj_name": str(obj), + "obj_type": type(obj).__name__, + }, + ) + + verify_value: Callable[[Any], None] + if isinstance(data_type, (StringType, VariantType)): + + def verify_value(obj: Any) -> None: # noqa: ANN401 + pass + + elif type(data_type) in _integral_bounds: + lower, upper = _integral_bounds[type(data_type)] + + def verify_value(obj: Any) -> None: # noqa: ANN401 + verify_acceptable_types(obj) + if obj < lower or obj > upper: + raise PySparkValueError( + error_class="VALUE_OUT_OF_BOUND", + message_parameters={ + "arg_name": "obj", + "lower_bound": str(lower), + "upper_bound": str(upper), + "actual": str(obj), + }, + ) + + elif isinstance(data_type, ArrayType): + verify_element = _make_type_verifier(data_type.elementType, data_type.containsNull, f"element in array {name}") + + def verify_value(obj: Any) -> None: # noqa: ANN401 + verify_acceptable_types(obj) + for element in obj: + verify_element(element) + + elif isinstance(data_type, MapType): + verify_key = _make_type_verifier(data_type.keyType, False, f"key of map {name}") + verify_map_value = _make_type_verifier(data_type.valueType, data_type.valueContainsNull, f"value of map {name}") + + def verify_value(obj: Any) -> None: # noqa: ANN401 + verify_acceptable_types(obj) + for key, value in obj.items(): + verify_key(key) + verify_map_value(value) + + elif isinstance(data_type, StructType): + verifiers = [(f.name, _make_type_verifier(f.dataType, f.nullable, new_name(f.name))) for f in data_type.fields] + + def verify_value(obj: Any) -> None: # noqa: ANN401 + if isinstance(obj, dict): + for field_name, verify in verifiers: + verify(obj.get(field_name)) + elif isinstance(obj, (tuple, list)): + if len(obj) != len(verifiers): + raise PySparkValueError( + error_class="LENGTH_SHOULD_BE_THE_SAME", + message_parameters={ + "arg1": "obj", + "arg2": "fields", + "arg1_length": str(len(obj)), + "arg2_length": str(len(verifiers)), + }, + ) + for value, (_, verify) in zip(obj, verifiers, strict=True): + verify(value) + elif hasattr(obj, "__dict__"): + for field_name, verify in verifiers: + verify(obj.__dict__.get(field_name)) + else: + raise PySparkTypeError( + error_class="CANNOT_ACCEPT_OBJECT_IN_TYPE", + message_parameters={ + "data_type": "StructType", + "obj_name": str(obj), + "obj_type": type(obj).__name__, + }, + ) + + else: + verify_value = verify_acceptable_types + + def verify(obj: Any) -> None: # noqa: ANN401 + if obj is not None: + verify_value(obj) + elif not nullable: + raise PySparkValueError(error_class="CANNOT_BE_NONE", message_parameters={"arg_name": "obj"}) + + return verify + + +def _to_string(obj: Any) -> str | None: # noqa: ANN401 + if obj is None or isinstance(obj, str): + return obj + return str(obj).lower() if isinstance(obj, bool) else str(obj) + + +def _apply(convert: Callable[[Any], Any] | None, obj: Any) -> Any: # noqa: ANN401 + return obj if convert is None else convert(obj) + + +def _contains_type(data_type: DataType, classes: tuple[type[DataType], ...]) -> bool: + """Whether ``data_type`` is, or has at any depth, a type of one of ``classes``.""" + if isinstance(data_type, classes): + return True + if isinstance(data_type, ArrayType): + return _contains_type(data_type.elementType, classes) + if isinstance(data_type, MapType): + return _contains_type(data_type.keyType, classes) or _contains_type(data_type.valueType, classes) + if isinstance(data_type, StructType): + return any(_contains_type(f.dataType, classes) for f in data_type.fields) + return False + + +def _replace_types(data_type: DataType, replacements: Mapping[type[DataType], DataType]) -> DataType: + """``data_type`` with each type, at any depth, whose class is in ``replacements`` replaced.""" + if isinstance(data_type, ArrayType): + return ArrayType(_replace_types(data_type.elementType, replacements), data_type.containsNull) + if isinstance(data_type, MapType): + return MapType( + _replace_types(data_type.keyType, replacements), + _replace_types(data_type.valueType, replacements), + data_type.valueContainsNull, + ) + if isinstance(data_type, StructType): + return StructType( + [StructField(f.name, _replace_types(f.dataType, replacements), f.nullable, f.metadata) for f in data_type] + ) + return replacements.get(type(data_type), data_type) + + +def _python_to_duckdb(data_type: DataType) -> Callable[[Any], Any]: + """Returns a function that turns a Python value into the value DuckDB stores for ``data_type``. + + Rows, tuples, dicts and objects become tuples in field order, values in a string field become + strings, and a naive datetime is a local time for a ``TimestampType`` and a wall clock time for + a ``TimestampNTZType``, as in Spark. + """ + if isinstance(data_type, StructType): + names = data_type.names + converters = [_python_to_duckdb(f.dataType) for f in data_type.fields] + + def convert_struct(obj: Any) -> tuple | None: # noqa: ANN401 + if obj is None: + return None + if isinstance(obj, (tuple, list)): + values = obj + elif isinstance(obj, dict): + values = [obj.get(name) for name in names] + elif hasattr(obj, "__dict__"): + values = [obj.__dict__.get(name) for name in names] + else: + raise PySparkTypeError( + error_class="UNSUPPORTED_DATA_TYPE", + message_parameters={"data_type": type(obj).__name__}, + ) + if len(values) != len(converters): + raise PySparkValueError( + error_class="LENGTH_SHOULD_BE_THE_SAME", + message_parameters={ + "arg1": "obj", + "arg2": "fields", + "arg1_length": str(len(values)), + "arg2_length": str(len(converters)), + }, + ) + return tuple(convert(value) for convert, value in zip(converters, values, strict=True)) + + return convert_struct + if isinstance(data_type, ArrayType): + convert_element = _python_to_duckdb(data_type.elementType) + return lambda obj: None if obj is None else [convert_element(v) for v in obj] + if isinstance(data_type, MapType): + convert_key = _python_to_duckdb(data_type.keyType) + convert_value = _python_to_duckdb(data_type.valueType) + return lambda obj: None if obj is None else {convert_key(k): convert_value(v) for k, v in obj.items()} + if isinstance(data_type, StringType): + return _to_string + if isinstance(data_type, NullType): + return lambda obj: None + if isinstance(data_type, TimestampType): + return _local_time_to_instant + if isinstance(data_type, TimestampNTZType): + return lambda obj: obj.replace(tzinfo=None) if isinstance(obj, datetime.datetime) else obj + return lambda obj: obj + + +def _duckdb_to_python(data_type: DataType, aware: bool = False) -> Callable[[Any], Any] | None: + """Returns a function that turns a value DuckDB returns into the value Spark returns. + + Structs become :class:`Row` objects and a ``TimestampType`` value becomes a naive local time. + Returns None where the two already agree. ``aware`` reads instants that come as aware datetimes, as DuckDB gives + them to a Python function, rather than as microseconds since the epoch. + """ + if isinstance(data_type, StructType): + names = data_type.names + converters = [_duckdb_to_python(f.dataType, aware) for f in data_type.fields] + + def convert_struct(obj: dict | tuple | None) -> Row | None: + if obj is None: + return None + values = obj.values() if isinstance(obj, dict) else obj + return _create_row(names, [_apply(convert, v) for convert, v in zip(converters, values, strict=True)]) + + return convert_struct + if isinstance(data_type, ArrayType): + convert_element = _duckdb_to_python(data_type.elementType, aware) + if convert_element is None: + return None + return lambda obj: None if obj is None else [convert_element(v) for v in obj] + if isinstance(data_type, MapType): + convert_key = _duckdb_to_python(data_type.keyType, aware) + convert_value = _duckdb_to_python(data_type.valueType, aware) + if isinstance(data_type.keyType, (ArrayType, MapType, StructType)): + # DuckDB returns a map whose keys Python cannot hash as a dict of a key list and a value list. + return lambda obj: ( + None + if obj is None + else { + _apply(convert_key, k): _apply(convert_value, v) + for k, v in zip(obj["key"], obj["value"], strict=True) + } + ) + if convert_key is None and convert_value is None: + return None + return lambda obj: ( + None if obj is None else {_apply(convert_key, k): _apply(convert_value, v) for k, v in obj.items()} + ) + if isinstance(data_type, TimestampType): + return _aware_to_local_time if aware else _instant_to_local_time + if isinstance(data_type, TimestampNanosecondType) and not aware: + return _instant_to_utc_time + return None + + +def _aware_to_local_time(obj: datetime.datetime | None) -> datetime.datetime | None: + return None if obj is None else _instant_to_local_time(TimestampType().toInternal(obj)) + + +def _local_time_to_instant(obj: Any) -> Any: # noqa: ANN401 + """Reads a naive datetime as a local time the way Spark's ``TimestampType.toInternal`` does.""" + if not isinstance(obj, datetime.datetime) or obj.tzinfo is not None: + return obj + # A fixed offset keeps the wall clock, so a time at either end of the calendar never shifts out of range. + fields = obj.timetuple() + offset = datetime.timedelta(seconds=calendar.timegm(fields) - time.mktime(fields)) + return obj.replace(tzinfo=datetime.timezone(offset)) + + +# The microseconds that stand for an infinite timestamp. +_INFINITE_MICROSECONDS = 2**63 - 1 + + +def _infinity(microseconds: int) -> datetime.datetime | None: + """The end of Python's calendar that DuckDB returns for an infinite timestamp, or None for a finite one.""" + if microseconds == _INFINITE_MICROSECONDS: + return datetime.datetime.max + if microseconds == -_INFINITE_MICROSECONDS: + return datetime.datetime.min + return None + + +def _instant_to_local_time(microseconds: int | None) -> datetime.datetime | None: + """Turns microseconds since the epoch into a naive local time the way Spark's ``TimestampType.fromInternal`` does. + + An instant whose local time is outside Python's range raises, as in Spark. + """ + if microseconds is None: + return None + return _infinity(microseconds) or TimestampType().fromInternal(microseconds) + + +def _instant_to_utc_time(microseconds: int | None) -> datetime.datetime | None: + """Turns microseconds since the epoch into the aware datetime in UTC of that instant.""" + if microseconds is None: + return None + epoch = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) + return _infinity(microseconds) or epoch + datetime.timedelta(microseconds=microseconds) + + def _create_row(fields: Union["Row", list[str]], values: tuple[Any, ...] | list[Any]) -> "Row": row = Row(*values) row.__fields__ = fields diff --git a/duckdb/experimental/spark/sql/udf.py b/duckdb/experimental/spark/sql/udf.py index c22f6be91..3210aab11 100644 --- a/duckdb/experimental/spark/sql/udf.py +++ b/duckdb/experimental/spark/sql/udf.py @@ -1,10 +1,24 @@ # https://sparkbyexamples.com/pyspark/pyspark-udf-user-defined-function/ +import contextlib +import datetime +import inspect +import uuid from typing import TYPE_CHECKING, Any, Optional, TypeVar if TYPE_CHECKING: from collections.abc import Callable -from .types import DataType +import duckdb +from duckdb import Expression, FunctionExpression +from duckdb.sqltypes import DuckDBPyType + +from ..errors import PySparkTypeError +from ..exception import ContributionsAcceptedError +from ._typing import ColumnOrName +from .column import Column, _call_name, _derived, _duckdb_type +from .functions import _as_column +from .type_utils import convert_type +from .types import DataType, StringType, TimestampType, _duckdb_to_python, _local_time_to_instant if TYPE_CHECKING: from .session import SparkSession @@ -13,6 +27,113 @@ UserDefinedFunctionLike = TypeVar("UserDefinedFunctionLike") +def _timestamp(value: Any) -> Any: # noqa: ANN401 + # DuckDB cannot convert a naive datetime to a TIMESTAMP WITH TIME ZONE, and reads a date as midnight in UTC. + if value is None or isinstance(value, datetime.datetime): + return _local_time_to_instant(value) + msg = f"a UDF result of type {type(value).__name__} for a timestamp return type" + raise TypeError(msg) + + +def _result_reader(data_type: DataType | None) -> "Callable[[Any], Any]": + """How a UDF's result goes to DuckDB, which converts it to the return type: a naive datetime as a local time.""" + return _timestamp if isinstance(data_type, TimestampType) else lambda v: v + + +def _spark_type(dtype: DuckDBPyType) -> DataType | None: + """The Spark type a DuckDB type stands for, or None for a DuckDB type that none stands for, such as an enum.""" + try: + return convert_type(dtype) + except (KeyError, ContributionsAcceptedError): + return None + + +def _argument_reader(type_text: str) -> "Callable[[Any], Any] | None": + """How Spark gives a UDF an argument of the DuckDB type ``type_text``: a struct as a Row, also inside others.""" + try: + data_type = _spark_type(DuckDBPyType(type_text)) + except duckdb.Error: + return None + return None if data_type is None else _duckdb_to_python(data_type, aware=True) + + +def _spark_udf(f: "Callable[..., Any]", data_type: DataType | None) -> "Callable[[Any, Any], Any]": + """``f`` called with Spark's values of its arguments. + + DuckDB gives a Python function a struct and a map alike as a dict, so the arguments come with their DuckDB types. + """ + read = _result_reader(data_type) + readers: dict[str, Callable[[Any], Any] | None] = {} + + def call(types: Any, values: Any) -> Any: # noqa: ANN401 + arguments = [] + for type_text, value in zip(types, values, strict=True): + if type_text not in readers: + readers[type_text] = _argument_reader(type_text) + reader = readers[type_text] + arguments.append(value if reader is None or value is None else reader(value)) + return read(f(*arguments)) + + return call + + +def _typed_arguments(arguments: list[Expression]) -> list[Expression]: + """The arguments of a call to a function ``_spark_udf`` made: the arguments' types, and the arguments.""" + types = FunctionExpression("list_value", *(FunctionExpression("typeof", a) for a in arguments)) + return [types.cast(DuckDBPyType("VARCHAR[]")), FunctionExpression("row", *arguments)] + + +_POSITIONAL = frozenset( + [inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.VAR_POSITIONAL] +) + + +def _positional_signature(f: "Callable[..., Any]") -> inspect.Signature: + """The positional parameters of ``f``, without annotations.""" + parameters = inspect.signature(f).parameters.values() + return inspect.Signature([p.replace(annotation=p.empty) for p in parameters if p.kind in _POSITIONAL]) + + +def _sql_identifier(name: str) -> str: + return '"' + name.replace('"', '""') + '"' + + +def _register_sql_name( + connection: duckdb.DuckDBPyConnection, + name: str, + f: "Callable[..., Any]", + data_type: DataType | None, + return_type: DuckDBPyType, +) -> None: + """Make SQL text call ``f`` by ``name``, as DuckDB calls a Python function: with DuckDB's values, a struct a dict. + + SQL text keeps calling a function of DuckDB or of an extension of that name; a Python function or a temporary + macro of that name is replaced, as Spark replaces a temporary function. + """ + connection.execute(f"DROP MACRO IF EXISTS temp.main.{_sql_identifier(name)}") + with contextlib.suppress(duckdb.InvalidInputException): + connection.remove_function(name) + # A failing statement would abort a transaction the session has open, so the catalog tells whether the name is + # taken, not a call that fails to bind. + taken = connection.sql( + "SELECT 1 FROM duckdb_functions() WHERE lower(function_name) = lower(?) " + "AND function_type IN ('scalar', 'aggregate', 'macro', 'window') AND schema_name IN ('main', 'pg_catalog')", + params=[name], + ).fetchall() + if taken: + return + read = _result_reader(data_type) + + def call(*values: Any) -> Any: # noqa: ANN401 + return read(f(*values)) + + # DuckDB reads its parameters from the signature, casting an argument to an annotation, which Spark ignores; + # without a signature it takes any number of arguments. + with contextlib.suppress(TypeError, ValueError): + call.__signature__ = _positional_signature(f) # type: ignore[attr-defined] + connection.create_function(name, call, return_type=return_type, null_handling="special") + + class UDFRegistration: # noqa: D101 def __init__(self, sparkSession: "SparkSession") -> None: # noqa: D107 self.sparkSession = sparkSession @@ -22,8 +143,37 @@ def register( # noqa: D102 name: str, f: "Callable[..., Any] | UserDefinedFunctionLike", returnType: Optional["DataTypeOrString"] = None, - ) -> "UserDefinedFunctionLike": - self.sparkSession.conn.create_function(name, f, return_type=returnType) + ) -> "Callable[..., Column]": + registered = getattr(f, "_spark_udf_function", None) + if registered is not None: + if returnType is not None: + raise PySparkTypeError( + error_class="CANNOT_SPECIFY_RETURN_TYPE_FOR_UDF", + message_parameters={"arg_name": "f", "return_type": str(returnType)}, + ) + f, returnType = registered, f.returnType + if returnType is None: + returnType = StringType() + return_type = returnType if isinstance(returnType, DuckDBPyType) else _duckdb_type(returnType, cast=False) + data_type = returnType if isinstance(returnType, DataType) else _spark_type(return_type) + connection = self.sparkSession.conn + # A function for each registration keeps an earlier callable on its own function, as in Spark, and the + # callable away from a builtin of the same name. + internal_name = f"spark_udf_{uuid.uuid4().hex}" + # Spark calls a Python function on NULL arguments too. + connection.create_function( + internal_name, _spark_udf(f, data_type), return_type=return_type, null_handling="special" + ) + _register_sql_name(connection, name, f, data_type, return_type) + + def call(*cols: ColumnOrName) -> Column: + columns = [_as_column(c) for c in cols] + expr = FunctionExpression(internal_name, *_typed_arguments([c.expr for c in columns])) + return _derived(expr, _call_name(name, *columns), *columns) + + call.returnType = returnType if data_type is None else data_type + call._spark_udf_function = f + return call def registerJavaFunction( # noqa: D102 self, diff --git a/pyproject.toml b/pyproject.toml index 1c89ebcaa..eaa5da1dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ Changelog = "https://github.com/duckdb/duckdb-python/releases" all = [ # users can install duckdb with 'duckdb[all]', which will install this list "ipython", # used in duckdb.query_graph "fsspec", # used in duckdb.filesystem - "numpy", # used in duckdb.experimental.spark and in duckdb.fetchnumpy() + "numpy", # used in duckdb.fetchnumpy() "pandas", # used for pandas dataframes all over the place "pyarrow", # used for pyarrow support "adbc-driver-manager", # for the adbc driver @@ -366,8 +366,10 @@ filterwarnings = [ # Jupyter is throwing DeprecationWarnings "ignore:function ham() is deprecated:DeprecationWarning", # Pyspark is throwing these warnings - "ignore:distutils Version classes are deprecated:DeprecationWarning", - "ignore:is_datetime64tz_dtype is deprecated:DeprecationWarning", + "ignore:distutils Version classes are deprecated:DeprecationWarning:pyspark", + "ignore:is_datetime64tz_dtype is deprecated:DeprecationWarning:pyspark", + "ignore:The distutils package is deprecated:DeprecationWarning:pyspark", + "ignore:The copy keyword is deprecated:DeprecationWarning:pyspark", ] [tool.mypy] diff --git a/src/native/python_conversion.cpp b/src/native/python_conversion.cpp index d6883bbcd..0e7dffb11 100644 --- a/src/native/python_conversion.cpp +++ b/src/native/python_conversion.cpp @@ -121,7 +121,7 @@ bool DictionaryHasMapFormat(const PyDictionary &dict) { } Value TransformDictionaryToStruct(optional_ptr context, const PyDictionary &dict, - const LogicalType &target_type = LogicalType::UNKNOWN) { + const LogicalType &target_type, bool nan_as_null) { auto conversion_error = [&]() { return InvalidInputException("We could not convert the object %s to the desired target type (%s)", dict.ToString(), target_type.ToString()); @@ -136,7 +136,7 @@ Value TransformDictionaryToStruct(optional_ptr context, const PyD // A TUPLE has no member names, so its members take the dict's values in dict order. for (idx_t i = 0; i < dict.len; i++) { auto val = TransformPythonValue(context, dict.values.attr("__getitem__")(i), - StructType::GetChildType(target_type, i)); + StructType::GetChildType(target_type, i), nan_as_null); struct_values.emplace_back(make_pair(StructType::GetChildName(target_type, i), std::move(val))); } return Value::STRUCT(std::move(struct_values)); @@ -145,7 +145,8 @@ Value TransformDictionaryToStruct(optional_ptr context, const PyD auto struct_keys = TransformStructKeys(dict.keys, dict.len); if (!struct_target) { for (idx_t i = 0; i < dict.len; i++) { - auto val = TransformPythonValue(context, dict.values.attr("__getitem__")(i), LogicalType::UNKNOWN); + auto val = + TransformPythonValue(context, dict.values.attr("__getitem__")(i), LogicalType::UNKNOWN, nan_as_null); struct_values.emplace_back(make_pair(std::move(struct_keys[i]), std::move(val))); } return Value::STRUCT(std::move(struct_values)); @@ -172,14 +173,14 @@ Value TransformDictionaryToStruct(optional_ptr context, const PyD } for (idx_t i = 0; i < dict.len; i++) { auto val = TransformPythonValue(context, dict.values.attr("__getitem__")(value_indexes[i]), - StructType::GetChildType(target_type, i)); + StructType::GetChildType(target_type, i), nan_as_null); struct_values.emplace_back(make_pair(StructType::GetChildName(target_type, i), std::move(val))); } return Value::STRUCT(std::move(struct_values)); } Value TransformStructFormatDictionaryToMap(optional_ptr context, const PyDictionary &dict, - const LogicalType &target_type) { + const LogicalType &target_type, bool nan_as_null) { if (dict.len == 0) { return EmptyMapValue(); } @@ -204,8 +205,8 @@ Value TransformStructFormatDictionaryToMap(optional_ptr context, vector elements; for (idx_t i = 0; i < size; i++) { - Value new_key = TransformPythonValue(context, dict.keys.attr("__getitem__")(i), key_target); - Value new_value = TransformPythonValue(context, dict.values.attr("__getitem__")(i), value_target); + Value new_key = TransformPythonValue(context, dict.keys.attr("__getitem__")(i), key_target, nan_as_null); + Value new_value = TransformPythonValue(context, dict.values.attr("__getitem__")(i), value_target, nan_as_null); key_type = ProxiedForceMaxLogicalType(context, key_type, new_key.type()); value_type = ProxiedForceMaxLogicalType(context, value_type, new_value.type()); @@ -229,10 +230,10 @@ Value TransformStructFormatDictionaryToMap(optional_ptr context, } Value TransformDictionaryToMap(optional_ptr context, const PyDictionary &dict, - const LogicalType &target_type = LogicalType::UNKNOWN) { + const LogicalType &target_type, bool nan_as_null) { if (target_type.id() != LogicalTypeId::UNKNOWN && !DictionaryHasMapFormat(dict)) { // dict == { 'k1': v1, 'k2': v2, ..., 'kn': vn } - return TransformStructFormatDictionaryToMap(context, dict, target_type); + return TransformStructFormatDictionaryToMap(context, dict, target_type, nan_as_null); } auto keys = dict.values.attr("__getitem__")(0); @@ -259,8 +260,8 @@ Value TransformDictionaryToMap(optional_ptr context, const PyDict value_target = LogicalType::LIST(MapType::ValueType(target_type)); } - auto key_list = TransformPythonValue(context, keys, key_target); - auto value_list = TransformPythonValue(context, values, value_target); + auto key_list = TransformPythonValue(context, keys, key_target, nan_as_null); + auto value_list = TransformPythonValue(context, values, value_target, nan_as_null); LogicalType key_type = LogicalType::SQLNULL; LogicalType value_type = LogicalType::SQLNULL; @@ -286,8 +287,8 @@ Value TransformDictionaryToMap(optional_ptr context, const PyDict return Value::MAP(ListType::GetChildType(map_type), std::move(elements)); } -Value TransformTupleToStruct(optional_ptr context, nb::handle ele, - const LogicalType &target_type = LogicalType::UNKNOWN) { +Value TransformTupleToStruct(optional_ptr context, nb::handle ele, const LogicalType &target_type, + bool nan_as_null) { auto tuple = nb::cast(ele); auto size = nb::len(tuple); @@ -304,7 +305,7 @@ Value TransformTupleToStruct(optional_ptr context, nb::handle ele auto &type = child_types[i].second; auto &name = StructType::GetChildName(target_type, i); auto element = nb::handle(tuple[i]); - auto converted_value = TransformPythonValue(context, element, type); + auto converted_value = TransformPythonValue(context, element, type, nan_as_null); children.emplace_back(make_pair(name, std::move(converted_value))); } auto result = Value::STRUCT(std::move(children)); @@ -416,7 +417,7 @@ LogicalType SniffPythonIntegerType(nb::handle ele) { return res.type(); } -Value TransformDictionary(optional_ptr context, const PyDictionary &dict) { +Value TransformDictionary(optional_ptr context, const PyDictionary &dict, bool nan_as_null) { //! DICT -> MAP FORMAT // keys() = [key, value] // values() = [ [n keys] ], [ [n values] ] @@ -430,9 +431,9 @@ Value TransformDictionary(optional_ptr context, const PyDictionar } if (DictionaryHasMapFormat(dict)) { - return TransformDictionaryToMap(context, dict); + return TransformDictionaryToMap(context, dict, LogicalType::UNKNOWN, nan_as_null); } - return TransformDictionaryToStruct(context, dict); + return TransformDictionaryToStruct(context, dict, LogicalType::UNKNOWN, nan_as_null); } struct ContiguousBuffer { @@ -607,7 +608,7 @@ struct PythonValueConversion { } static void HandleList(optional_ptr context, Value &result, const LogicalType &target_type, - nb::handle ele, idx_t list_size) { + nb::handle ele, idx_t list_size, bool nan_as_null) { vector values; values.reserve(list_size); @@ -621,7 +622,7 @@ struct PythonValueConversion { } LogicalType element_type = LogicalType::SQLNULL; for (idx_t i = 0; i < list_size; i++) { - Value new_value = TransformPythonValue(context, ele.attr("__getitem__")(i), child_type); + Value new_value = TransformPythonValue(context, ele.attr("__getitem__")(i), child_type, nan_as_null); element_type = ProxiedForceMaxLogicalType(context, element_type, new_value.type()); values.push_back(std::move(new_value)); } @@ -633,12 +634,12 @@ struct PythonValueConversion { } static void HandleTuple(optional_ptr context, Value &result, const LogicalType &target_type, - nb::handle ele, idx_t list_size) { + nb::handle ele, idx_t list_size, bool nan_as_null) { if (target_type.id() == LogicalTypeId::STRUCT || target_type.id() == LogicalTypeId::TUPLE) { - result = TransformTupleToStruct(context, ele, target_type); + result = TransformTupleToStruct(context, ele, target_type, nan_as_null); return; } - HandleList(context, result, target_type, ele, list_size); + HandleList(context, result, target_type, ele, list_size, nan_as_null); } static Value HandleObjectInternal(optional_ptr context, nb::handle ele, PythonObjectType object_type, @@ -661,11 +662,11 @@ struct PythonValueConversion { switch (target_type.id()) { case LogicalTypeId::STRUCT: case LogicalTypeId::TUPLE: - return TransformDictionaryToStruct(context, dict, target_type); + return TransformDictionaryToStruct(context, dict, target_type, nan_as_null); case LogicalTypeId::MAP: - return TransformDictionaryToMap(context, dict, target_type); + return TransformDictionaryToMap(context, dict, target_type, nan_as_null); default: - return TransformDictionary(context, dict); + return TransformDictionary(context, dict, nan_as_null); } } case PythonObjectType::Value: { @@ -678,7 +679,7 @@ struct PythonValueConversion { throw InvalidInputException("The 'type' of a Value should be of type DuckDBPyType, not '%s'", actual_type); } - return TransformPythonValue(context, object, internal_type->Type()); + return TransformPythonValue(context, object, internal_type->Type(), nan_as_null); } default: throw InternalException("Unsupported fallback"); @@ -889,7 +890,7 @@ struct PythonVectorConversion { template static void HandleListFast(optional_ptr context, Vector &result, const idx_t &result_offset, - nb::handle ele, idx_t list_size) { + nb::handle ele, idx_t list_size, bool nan_as_null) { auto &result_type = result.GetType(); if (result_type.id() == LogicalTypeId::ARRAY) { idx_t array_size = ArrayType::GetSize(result_type); @@ -902,7 +903,7 @@ struct PythonVectorConversion { idx_t start_offset = result_offset * array_size; for (idx_t i = 0; i < list_size; i++) { auto child_ele = IS_LIST ? PyList_GetItem(ele.ptr(), i) : PyTuple_GetItem(ele.ptr(), i); - TransformPythonObject(context, child_ele, child_array, start_offset + i); + TransformPythonObject(context, child_ele, child_array, start_offset + i, nan_as_null); } return; } @@ -920,7 +921,7 @@ struct PythonVectorConversion { auto &child_vector = ListVector::GetChildMutable(result); for (idx_t i = 0; i < list_size; i++) { auto child_ele = IS_LIST ? PyList_GetItem(ele.ptr(), i) : PyTuple_GetItem(ele.ptr(), i); - TransformPythonObject(context, child_ele, child_vector, start_offset + i); + TransformPythonObject(context, child_ele, child_vector, start_offset + i, nan_as_null); } ListVector::SetListSize(result, start_offset + list_size); return; @@ -929,20 +930,20 @@ struct PythonVectorConversion { } static void HandleList(optional_ptr context, Vector &result, const idx_t &result_offset, - nb::handle ele, idx_t list_size) { + nb::handle ele, idx_t list_size, bool nan_as_null) { auto &result_type = result.GetType(); if (result_type.id() == LogicalTypeId::ARRAY || result_type.id() == LogicalTypeId::LIST) { - HandleListFast(context, result, result_offset, ele, list_size); + HandleListFast(context, result, result_offset, ele, list_size, nan_as_null); return; } // fallback to value conversion Value result_val; - PythonValueConversion::HandleList(context, result_val, result_type, ele, list_size); + PythonValueConversion::HandleList(context, result_val, result_type, ele, list_size, nan_as_null); FallbackValueConversion(result, result_offset, std::move(result_val)); } static void ConvertTupleToStruct(optional_ptr context, Vector &result, const idx_t &result_offset, - nb::handle ele, idx_t size) { + nb::handle ele, idx_t size, bool nan_as_null) { auto &child_types = StructType::GetChildTypes(result.GetType()); auto child_count = child_types.size(); if (size != child_count) { @@ -954,21 +955,21 @@ struct PythonVectorConversion { auto &struct_children = StructVector::GetEntries(result); for (idx_t i = 0; i < child_count; i++) { auto child_ele = PyTuple_GetItem(ele.ptr(), i); - TransformPythonObject(context, child_ele, struct_children[i], result_offset); + TransformPythonObject(context, child_ele, struct_children[i], result_offset, nan_as_null); } } static void HandleTuple(optional_ptr context, Vector &result, const idx_t &result_offset, - nb::handle ele, idx_t tuple_size) { + nb::handle ele, idx_t tuple_size, bool nan_as_null) { auto &result_type = result.GetType(); switch (result_type.id()) { case LogicalTypeId::STRUCT: case LogicalTypeId::TUPLE: - ConvertTupleToStruct(context, result, result_offset, ele, tuple_size); + ConvertTupleToStruct(context, result, result_offset, ele, tuple_size, nan_as_null); break; case LogicalTypeId::ARRAY: case LogicalTypeId::LIST: - HandleListFast(context, result, result_offset, ele, tuple_size); + HandleListFast(context, result, result_offset, ele, tuple_size, nan_as_null); break; default: throw InternalException("Unsupported type for HandleTuple"); @@ -1051,7 +1052,7 @@ void TransformPythonObjectInternal(optional_ptr context, nb::hand } case PythonObjectType::List: { auto list_size = nb::len(ele); - OP::HandleList(context, result, param, ele, list_size); + OP::HandleList(context, result, param, ele, list_size, nan_as_null); break; } case PythonObjectType::Tuple: { @@ -1063,7 +1064,7 @@ void TransformPythonObjectInternal(optional_ptr context, nb::hand case LogicalTypeId::UNKNOWN: case LogicalTypeId::LIST: case LogicalTypeId::ARRAY: - OP::HandleTuple(context, result, param, ele, list_size); + OP::HandleTuple(context, result, param, ele, list_size, nan_as_null); break; default: throw InvalidInputException("Can't convert tuple to a Value of type %s", conversion_target); diff --git a/src/pyconnection/type_creation.cpp b/src/pyconnection/type_creation.cpp index d517c96b6..65f65db37 100644 --- a/src/pyconnection/type_creation.cpp +++ b/src/pyconnection/type_creation.cpp @@ -55,9 +55,6 @@ static child_list_t GetChildList(const nb::object &container) { std::unique_ptr DuckDBPyConnection::StructType(const nb::object &fields) { child_list_t types = GetChildList(fields); - if (types.empty()) { - throw InvalidInputException("Can not create an empty struct type!"); - } auto struct_type = LogicalType::STRUCT(std::move(types)); return make_uniq(struct_type); } diff --git a/src/pyexpression.cpp b/src/pyexpression.cpp index 998483dc7..385cf9b67 100644 --- a/src/pyexpression.cpp +++ b/src/pyexpression.cpp @@ -296,8 +296,21 @@ static void PopulateExcludeList(qualified_column_set_t &exclude, nb::object list if (expr->GetExpression().GetExpressionType() != ExpressionType::COLUMN_REF) { throw nb::value_error("Only ColumnExpressions are accepted Expression types here"); } - auto &column = expr->GetExpression().Cast(); - exclude.insert(QualifiedColumnName(column.GetColumnName())); + auto &names = expr->GetExpression().Cast().ColumnNames(); + if (names.empty() || names.size() > 4) { + throw nb::value_error("Expected 1 to 4 name parts (catalog.schema.table.column) in an excluded column"); + } + QualifiedColumnName qname(names.back()); + if (names.size() >= 2) { + qname.table = names[names.size() - 2]; + } + if (names.size() >= 3) { + qname.schema = names[names.size() - 3]; + } + if (names.size() == 4) { + qname.catalog = names[0]; + } + exclude.insert(qname); } } @@ -316,14 +329,10 @@ std::unique_ptr DuckDBPyExpression::ColumnExpression(const n return StarExpression(); } - auto qualified_name = QualifiedName::Parse(column_name); - if (!qualified_name.Catalog().empty()) { - column_names.push_back(qualified_name.Catalog()); - } - if (!qualified_name.Schema().empty()) { - column_names.push_back(qualified_name.Schema()); + column_names = QualifiedName::ParseComponents(column_name); + if (column_names.empty()) { + column_names.emplace_back(); } - column_names.push_back(qualified_name.Name()); } else { for (auto part : names) { // nanobind args iteration yields temporary handles; bind by value (cheap handle) column_names.push_back(Identifier(nb::cast(part))); diff --git a/tests/fast/spark/test_spark_aggregates.py b/tests/fast/spark/test_spark_aggregates.py new file mode 100644 index 000000000..c40baa2e6 --- /dev/null +++ b/tests/fast/spark/test_spark_aggregates.py @@ -0,0 +1,371 @@ +import datetime +import math +from decimal import Decimal + +import pytest + +_ = pytest.importorskip("duckdb.experimental.spark") + +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException, error_condition +from spark_namespace.sql import functions as F +from spark_namespace.sql.types import ( + DoubleType, + IntegerType, + LongType, + Row, + ShortType, + StringType, + StructField, + StructType, +) + + +def _ansi(spark): + return USE_ACTUAL_SPARK and spark.conf.get("spark.sql.ansi.enabled") == "true" + + +@pytest.fixture +def numbers(spark): + schema = StructType( + [ + StructField("g", StringType()), + StructField("i", IntegerType()), + StructField("l", LongType()), + StructField("s", ShortType()), + StructField("d", DoubleType()), + ] + ) + rows = [("a", 1, 10, 1, 1.0), ("a", 2, None, 2, 2.0), ("b", 4, 40, 4, 4.0), ("b", None, 50, None, None)] + return spark.createDataFrame(rows, schema) + + +class TestSparkAggregates: + def test_select_global_aggregates(self, numbers): + res = numbers.select( + F.sum("i").alias("sum"), + F.count("i").alias("count"), + F.avg("d").alias("avg"), + (F.max("l") * 2 + F.min("i")).alias("expr"), + F.round(F.avg("l"), 1).alias("rounded"), + F.when(F.count("l") > 2, "many").otherwise("few").alias("when"), + ).collect() + assert res == [Row(sum=7, count=3, avg=pytest.approx(7 / 3), expr=101, rounded=33.3, when="many")] + + def test_select_aggregate_of_empty_input(self, numbers): + res = ( + numbers.filter(F.col("i") > 100) + .select( + F.sum("i").alias("sum"), + F.count("i").alias("count"), + F.max("d").alias("max"), + F.coalesce(F.sum("l"), F.lit(0)).alias("zero"), + F.collect_list("i").alias("list"), + ) + .collect() + ) + assert res == [Row(sum=None, count=0, max=None, zero=0, list=[])] + + def test_global_groupby_agg(self, numbers): + assert numbers.groupBy().agg(F.sum("l").alias("sum")).collect() == [Row(sum=100)] + + def test_constants_next_to_aggregates(self, numbers): + res = numbers.select(F.sum("i"), F.lit(1), (F.lit(2) * 3).alias("six"), F.abs(F.lit(-4)), F.lit("x")) + assert res.columns == ["sum(i)", "1", "six", "abs(-4)", "x"] + assert [tuple(r) for r in res.collect()] == [(7, 1, 6, 4, "x")] + res = numbers.groupBy().agg(F.max("i"), F.lit(None).alias("n")) + assert [tuple(r) for r in res.collect()] == [(4, None)] + + def test_string_arguments_that_are_not_columns_keep_a_constant(self, numbers): + res = numbers.select( + F.sum("i"), F.lpad(F.lit("x"), 3, "-"), F.concat_ws("-", F.lit("a"), F.lit("b")), F.instr(F.lit("ab"), "b") + ) + assert res.columns == ["sum(i)", "lpad(x, 3, -)", "concat_ws(-, a, b)", "instr(ab, b)"] + assert [tuple(r) for r in res.collect()] == [(7, "--x", "a-b", 2)] + + def test_call_function_of_an_aggregate(self, numbers): + res = numbers.groupBy().agg(F.call_function("sum", "l"), F.call_function("max", F.col("i"))) + assert res.columns == ["sum(l)", "max(i)"] + assert res.schema.simpleString() == "struct" + assert res.collect() == [Row(100, 4)] + assert numbers.select(F.call_function("count", "i")).collect() == [Row(3)] + grouped = numbers.groupBy("g").agg(F.call_function("avg", "d").alias("avg")).orderBy("g").collect() + assert grouped == [Row(g="a", avg=1.5), Row(g="b", avg=4.0)] + + @pytest.mark.parametrize( + "make", + [ + lambda df: df.select(F.sum("i"), "i"), + lambda df: df.select(F.sum("i"), F.col("i")), + lambda df: df.groupBy().agg(F.col("i")), + lambda df: df.groupBy().agg(F.sum("i"), F.col("i")), + lambda df: df.select(F.sum("i"), F.col("i") + F.lit(1)), + lambda df: df.select(F.sum("i"), F.abs("i")), + lambda df: df.select(F.sum("i"), F.lpad("g", 3, "-")), + lambda df: df.select(F.sum("i"), F.expr("g")), + lambda df: df.select(F.sum("i"), F.struct({"g"})), + lambda df: df.select(F.sum("i"), F.array({"g"})), + ], + ) + def test_aggregate_mixed_with_plain_column_raises(self, numbers, make): + with pytest.raises(AnalysisException, match="MISSING_GROUP_BY"): + make(numbers).collect() + + def test_sum_of_integers_is_bigint(self, numbers): + df = numbers.groupBy("g").agg(F.sum("i"), F.sum("l"), F.sum("s"), F.sum("d"), F.sum("i") * 2).orderBy("g") + assert [f.dataType for f in df.schema.fields] == [ + StringType(), + LongType(), + LongType(), + LongType(), + DoubleType(), + LongType(), + ] + assert df.collect() == [Row("a", 3, 10, 3, 3.0, 6), Row("b", 4, 90, 4, 4.0, 8)] + global_sum = numbers.select(F.sum("i")) + assert global_sum.schema.fields[0].dataType == LongType() + + def test_grouped_data_sum_is_bigint(self, numbers): + df = numbers.groupBy("g").sum("i", "d").orderBy("g") + assert [f.dataType for f in df.schema.fields] == [StringType(), LongType(), DoubleType()] + assert df.collect() == [Row("a", 3, 3.0), Row("b", 4, 4.0)] + + def test_sum_overflow(self, spark): + df = spark.createDataFrame([(9223372036854775807,), (1,)], StructType([StructField("v", LongType())])) + if USE_ACTUAL_SPARK and not _ansi(spark): + assert df.select(F.sum("v")).collect() == [Row(-9223372036854775808)] + else: + with pytest.raises(Exception, match=r"(?i)overflow|out of range"): + df.select(F.sum("v")).collect() + + def test_sum_of_int_past_int_range(self, spark): + df = spark.createDataFrame([(2147483647,), (1,)], StructType([StructField("v", IntegerType())])) + assert df.select(F.sum("v").alias("s")).collect() == [Row(s=2147483648)] + + def test_aggregate_result_types(self, numbers): + df = numbers.select( + F.avg("i"), F.mean("l"), F.count("i"), F.count("*"), F.min("i"), F.max("l"), F.stddev("i"), F.var_pop("i") + ) + types = [f.dataType for f in df.schema.fields] + assert types == [ + DoubleType(), + DoubleType(), + LongType(), + LongType(), + IntegerType(), + LongType(), + DoubleType(), + DoubleType(), + ] + + def test_count_star_and_literal(self, numbers): + res = ( + numbers.groupBy("g") + .agg(F.count("*").alias("star"), F.count(F.lit(1)).alias("one"), F.count("i").alias("i")) + .orderBy("g") + .collect() + ) + assert res == [Row(g="a", star=2, one=2, i=2), Row(g="b", star=2, one=2, i=1)] + assert numbers.select(F.count("*").alias("n")).collect() == [Row(n=4)] + assert numbers.select(F.count(F.lit(None)).alias("n")).collect() == [Row(n=0)] + + def test_count_star_names(self, numbers): + res = numbers.select(F.count("*"), F.count(F.col("*")), F.count(F.lit(1))) + assert res.columns == ["count(1)", "count(1)", "count(1)"] + assert [tuple(r) for r in res.collect()] == [(4, 4, 4)] + assert numbers.groupBy("g").agg(F.count("*")).columns == ["g", "count(1)"] + + def test_unaliased_struct_names(self, numbers): + res = numbers.select(F.struct("g", "i"), F.struct(["g", numbers.i]), F.struct(F.col("i").alias("z"))) + assert res.columns == ["struct(g, i)", "struct(g, i)", "struct(i AS z)"] + + def test_skewness_population(self, spark): + schema = StructType([StructField("g", StringType()), StructField("x", DoubleType())]) + rows = [ + ("a", 1.0), + ("a", 1.0), + ("a", 2.0), + ("a", 10.0), + ("b", 1.0), + ("b", 3.0), + ("c", 3.0), + ("c", 3.0), + ("d", 5.0), + ] + res = spark.createDataFrame(rows, schema).groupBy("g").agg(F.skewness("x").alias("v")).orderBy("g").collect() + assert res[0].v == pytest.approx(1.1153956690) + assert res[1].v == pytest.approx(0.0) + assert res[2].v is None + assert res[3].v is None + + @pytest.mark.parametrize( + "values", + [[float("nan"), 1.0, 2.0], [float("nan"), 1.0], [float("nan")], [float("inf"), 1.0, 2.0], [-float("inf"), 1.0]], + ) + def test_skewness_of_nan_or_infinity_is_nan(self, spark, values): + df = spark.createDataFrame([(v,) for v in values], StructType([StructField("x", DoubleType())])) + assert math.isnan(df.select(F.skewness("x").alias("v")).collect()[0].v) + + def test_corr_without_variance(self, spark): + schema = StructType([StructField("x", DoubleType()), StructField("y", DoubleType())]) + df = spark.createDataFrame([(1.0, 1.0), (1.0, 2.0)], schema) + assert spark.createDataFrame([(1.0, 1.0)], schema).select(F.corr("x", "y").alias("c")).collect() == [ + Row(c=None) + ] + if _ansi(spark): + with pytest.raises(Exception, match="DIVIDE_BY_ZERO"): + df.select(F.corr("x", "y")).collect() + else: + assert df.select(F.corr("x", "y").alias("c")).collect() == [Row(c=None)] + varied = spark.createDataFrame([(1.0, 2.0), (2.0, 4.5), (4.0, 7.0)], schema) + assert varied.select(F.corr("x", "y").alias("c")).collect()[0].c == pytest.approx(0.9819805061) + + def test_collect_list_skips_nulls(self, spark): + schema = StructType([StructField("g", StringType()), StructField("v", LongType())]) + df = spark.createDataFrame([("a", 1), ("a", None), ("a", 2), ("b", None)], schema) + res = ( + df.groupBy("g") + .agg(F.sort_array(F.collect_list("v")).alias("c"), F.sort_array(F.array_agg("v")).alias("a")) + .orderBy("g") + .collect() + ) + assert res == [Row(g="a", c=[1, 2], a=[1, 2]), Row(g="b", c=[], a=[])] + + def test_any_value_keeps_first_null(self, spark): + schema = StructType([StructField("g", LongType()), StructField("v", LongType())]) + df = spark.createDataFrame([(1, None), (1, 2)], schema) + if USE_ACTUAL_SPARK: + df = df.coalesce(1) + assert df.groupBy("g").agg(F.any_value("v").alias("v")).collect() == [Row(g=1, v=None)] + + def test_missing_group_by_sql_state(self, numbers): + with pytest.raises(AnalysisException) as error: + numbers.select(F.sum("i"), "i").collect() + assert error.value.getSqlState() == "42803" + + +@pytest.fixture +def typed(spark): + schema = "k string, i int, l bigint, d decimal(10,2)" + return spark.createDataFrame([("x", 1, 2, Decimal("1.50")), ("x", 3, 4, Decimal("2.25"))], schema) + + +@pytest.fixture +def duckdb_only(spark): + if USE_ACTUAL_SPARK: + pytest.skip("HUGEINT and UBIGINT are DuckDB only types") + return spark.sql( + f"select 'x' k, 1::int i, {2**127 - 1}::hugeint h, {2**64 - 1}::ubigint u union all " + "select 'x', 2, -1, 0::ubigint" + ) + + +class TestAggregateResultTypes: + def test_sum_and_max_of_spark_types(self, typed): + df = typed.select(F.sum("i"), F.sum("l"), F.max("i"), F.max("l"), F.sum("d"), F.max("d")) + assert df.dtypes[:4] == [("sum(i)", "bigint"), ("sum(l)", "bigint"), ("max(i)", "int"), ("max(l)", "bigint")] + assert df.dtypes[5] == ("max(d)", "decimal(10,2)") + assert df.collect() == [Row(4, 6, 3, 4, Decimal("3.75"), Decimal("2.25"))] + + def test_aggregate_nested_in_an_expression(self, typed): + df = typed.select(F.array(F.sum("i")).alias("a"), F.struct(F.sum("l").alias("s")).alias("t")) + assert df.dtypes == [("a", "array"), ("t", "struct")] + assert df.collect() == [Row(a=[4], t=Row(s=6))] + + def test_aggregates_over_a_join(self, spark, typed): + other = spark.createDataFrame([("x", 10)], "k string, j int") + joined = typed.join(other, "k") + df = joined.groupBy("k").agg(F.sum(typed.i), F.array(F.sum(other.j)).alias("a"), F.max(typed.l)) + assert df.dtypes == [("k", "string"), ("sum(i)", "bigint"), ("a", "array"), ("max(l)", "bigint")] + assert df.collect() == [Row(k="x", **{"sum(i)": 4}, a=[20], **{"max(l)": 4})] + + def test_duckdb_only_inputs_keep_their_types(self, duckdb_only): + df = duckdb_only.select(F.max("h"), F.sum("h"), F.sum("u"), F.max("u"), F.sum("i")) + assert df.dtypes == [ + ("max(h)", "hugeint"), + ("sum(h)", "hugeint"), + ("sum(u)", "hugeint"), + ("max(u)", "ubigint"), + ("sum(i)", "bigint"), + ] + assert df.collect() == [Row(2**127 - 1, 2**127 - 2, 2**64 - 1, 2**64 - 1, 3)] + + def test_duckdb_only_inputs_of_grouped_data(self, duckdb_only): + df = duckdb_only.groupBy("k").sum("i", "h", "u") + assert df.dtypes == [("k", "string"), ("sum(i)", "bigint"), ("sum(h)", "hugeint"), ("sum(u)", "hugeint")] + assert df.collect() == [Row("x", 3, 2**127 - 2, 2**64 - 1)] + assert duckdb_only.groupBy("k").max("h").collect() == [Row("x", 2**127 - 1)] + + def test_duckdb_only_inputs_over_a_join(self, spark, duckdb_only): + other = spark.createDataFrame([("x", 10)], "k string, j int") + df = duckdb_only.join(other, "k").groupBy("k").agg(F.max(duckdb_only.h), F.sum(other.j)) + assert df.dtypes == [("k", "string"), ("max(h)", "hugeint"), ("sum(j)", "bigint")] + assert df.collect() == [Row("x", 2**127 - 1, 20)] + + def test_expression_mixing_a_spark_sum_with_a_duckdb_only_input(self, duckdb_only): + # Every HUGEINT of a result with a DuckDB only input stays, including that of the integer sum beside it. + df = duckdb_only.select(F.struct(F.sum("i").alias("s"), F.max("h").alias("m")).alias("t")) + assert df.dtypes == [("t", "struct")] + assert df.collect() == [Row(t=Row(s=3, m=2**127 - 1))] + + +class TestGroupedDataNumericColumns: + @pytest.mark.parametrize("method", ["max", "min", "sum", "avg", "mean"]) + @pytest.mark.parametrize("column", ["s", "dt", "b", "`s`"]) + def test_non_numeric_column_raises(self, spark, method, column): + df = spark.createDataFrame([(1, "x", datetime.date(2024, 1, 1), True)], "k int, s string, dt date, b boolean") + with pytest.raises(AnalysisException) as error: + getattr(df.groupBy("k"), method)(column).collect() + assert error_condition(error.value) == "_LEGACY_ERROR_TEMP_1323" + assert dict(error.value.getMessageParameters()) == {"colName": column} + + @pytest.mark.parametrize("schema", ["k int, `a``b` int", "k int, `c.d` int"]) + def test_numeric_columns_are_read_as_attribute_names(self, spark, schema): + # Spark reads the names of the numeric columns as attribute names, so these raise there too. + with pytest.raises(Exception): # noqa: B017, PT011 + spark.createDataFrame([(1, 2)], schema).groupBy("k").max().collect() + + def test_numeric_columns_are_aggregated(self, spark): + df = spark.createDataFrame([(1, "x", 2.0)], "k int, s string, d double") + assert df.groupBy("k").max("d", "k").collect() == [Row(k=1, **{"max(d)": 2.0, "max(k)": 1})] + assert df.groupBy("k").max().collect() == [Row(k=1, **{"max(k)": 1, "max(d)": 2.0})] + + +class TestRepeatedResultNames: + @pytest.fixture + def pairs(self, spark): + return spark.createDataFrame([(1, 10, 3), (1, 20, 4)], ["k", "a", "b"]) + + def test_aggregate_named_like_the_key(self, pairs): + res = pairs.groupBy("k").agg(F.sum("a").alias("k")) + assert res.columns == ["k", "k"] + assert res.dtypes == [("k", "bigint"), ("k", "bigint")] + assert [tuple(r) for r in res.collect()] == [(1, 30)] + + def test_aggregates_with_one_name(self, pairs): + res = pairs.select(F.sum("a").alias("x"), F.sum("b").alias("x")) + assert [tuple(r) for r in res.collect()] == [(30, 7)] + + def test_names_that_differ_in_case(self, pairs): + res = pairs.groupBy("k").agg(F.sum("a").alias("K"), F.max("b").alias("A")) + assert res.columns == ["k", "K", "A"] + assert [tuple(r) for r in res.collect()] == [(1, 30, 4)] + + +class TestGlobalAggregateOfConstants: + @pytest.fixture + def two_rows(self, spark): + return spark.createDataFrame([(1, 2), (3, 4)], ["a", "b"]) + + def test_one_row(self, two_rows): + res = two_rows.groupBy().agg(F.lit(1), F.lit("s").alias("s")) + assert res.columns == ["1", "s"] + assert [tuple(r) for r in res.collect()] == [(1, "s")] + + def test_one_row_of_empty_input(self, two_rows): + empty = two_rows.filter(F.col("a") > 100) + assert [tuple(r) for r in empty.groupBy().agg(F.lit(1)).collect()] == [(1,)] + assert [tuple(r) for r in empty.groupBy().agg(F.lit(1), F.sum("a")).collect()] == [(1, None)] + + def test_constants_with_keys(self, two_rows): + assert sorted(tuple(r) for r in two_rows.groupBy("a").agg(F.lit(1)).collect()) == [(1, 1), (3, 1)] diff --git a/tests/fast/spark/test_spark_column.py b/tests/fast/spark/test_spark_column.py index b2656643c..290b2e26e 100644 --- a/tests/fast/spark/test_spark_column.py +++ b/tests/fast/spark/test_spark_column.py @@ -2,37 +2,49 @@ _ = pytest.importorskip("duckdb.experimental.spark") +import datetime import re +from decimal import Decimal from spark_namespace import USE_ACTUAL_SPARK -from spark_namespace.errors import PySparkTypeError +from spark_namespace.errors import ParseException, PySparkTypeError, PySparkValueError, SparkRuntimeException +from spark_namespace.sql import functions as F from spark_namespace.sql.functions import array, col, struct -from spark_namespace.sql.types import Row +from spark_namespace.sql.types import ( + BooleanType, + ByteType, + DoubleType, + FloatType, + IntegerType, + LongType, + Row, + ShortType, + StringType, + StructField, + StructType, +) + +if USE_ACTUAL_SPARK: + import pyspark + + +def _ansi(spark): + return USE_ACTUAL_SPARK and spark.conf.get("spark.sql.ansi.enabled") == "true" class TestSparkColumn: def test_struct_column(self, spark): df = spark.createDataFrame([Row(a=1, b=2, c=3, d=4)]) + df = df.withColumn("struct", struct(df.a, df.b)) + assert "struct" in df.columns + assert df.schema["struct"].dataType.fieldNames() == ["a", "b"] - # TODO: column names should be set explicitly using the Row, rather than letting duckdb # noqa: TD002, TD003 - # assign defaults(col0, col1, etc..) - if USE_ACTUAL_SPARK: - df = df.withColumn("struct", struct(df.a, df.b)) + # The module keeps PySpark 3.5's error table; PySpark 4 renamed this error. + if USE_ACTUAL_SPARK and not pyspark.__version__.startswith("3."): + message = "[NOT_EXPECTED_TYPE] Argument `col` should be Column, got str." else: - df = df.withColumn("struct", struct(df.col0, df.col1)) - assert "struct" in df - new_col = df.schema["struct"] - - if USE_ACTUAL_SPARK: - assert "a" in df.schema["struct"].dataType.fieldNames() - assert "b" in df.schema["struct"].dataType.fieldNames() - else: - assert "col0" in new_col.dataType - assert "col1" in new_col.dataType - - with pytest.raises( - PySparkTypeError, match=re.escape("[NOT_COLUMN] Argument `col` should be a Column, got str.") - ): + message = "[NOT_COLUMN] Argument `col` should be a Column, got str." + with pytest.raises(PySparkTypeError, match=re.escape(message)): df = df.withColumn("struct", "yes") def test_array_column(self, spark): @@ -66,3 +78,450 @@ def test_isnotnull(self, spark): Row(a=1, b=None, r1=True, r2=False), Row(a=None, b=2, r1=False, r2=True), ] + + def test_cast_to_integral_truncates(self, spark): + schema = StructType( + [ + StructField("big", LongType()), + StructField("d", DoubleType()), + StructField("f", FloatType()), + StructField("b", BooleanType()), + StructField("n", LongType()), + ] + ) + df = spark.createDataFrame( + [(9007199254740993, -2.7, 2.7, True, None), (-9007199254740993, 2.5, -0.5, False, None)], schema + ) + res = df.select( + col("big").cast("bigint").alias("big"), + col("big").cast("decimal(20,0)").cast("bigint").alias("big_dec"), + col("d").cast("int").alias("d_int"), + col("d").cast(LongType()).alias("d_long"), + col("d").cast(ShortType()).alias("d_short"), + col("d").cast(ByteType()).alias("d_byte"), + col("f").cast("int").alias("f_int"), + col("d").cast("decimal(5,2)").cast("int").alias("dec_int"), + col("b").cast("int").alias("b_int"), + col("n").cast("int").alias("n_int"), + ) + assert [f.dataType for f in res.schema.fields] == [ + LongType(), + LongType(), + IntegerType(), + LongType(), + ShortType(), + ByteType(), + IntegerType(), + IntegerType(), + IntegerType(), + IntegerType(), + ] + assert res.collect() == [ + Row( + big=9007199254740993, + big_dec=9007199254740993, + d_int=-2, + d_long=-2, + d_short=-2, + d_byte=-2, + f_int=2, + dec_int=-2, + b_int=1, + n_int=None, + ), + Row( + big=-9007199254740993, + big_dec=-9007199254740993, + d_int=2, + d_long=2, + d_short=2, + d_byte=2, + f_int=0, + dec_int=2, + b_int=0, + n_int=None, + ), + ] + + def test_cast_inside_expressions(self, spark): + df = spark.createDataFrame( + [(2.7, -1.5), (3.2, 0.0)], StructType([StructField("a", DoubleType()), StructField("b", DoubleType())]) + ) + res = df.select( + (col("a").cast("int") + col("b").cast("int")).alias("s"), + F.abs(col("b").cast("int")).alias("abs"), + F.concat(col("a").cast("int").cast("string"), F.lit("x")).alias("c"), + F.when(col("a").cast("int") == 2, "two").otherwise("other").alias("w"), + ).collect() + assert res == [Row(s=1, abs=1, c="2x", w="two"), Row(s=3, abs=0, c="3x", w="other")] + assert df.filter(col("a").cast("int") == 3).collect() == [Row(a=3.2, b=0.0)] + + def test_cast_strings_to_integral(self, spark): + schema = StructType([StructField(c, StringType()) for c in "abcd"]) + df = spark.createDataFrame([("7", " 8 ", "-2.7", "1.9")], schema) + assert df.select(col("a").cast("int").alias("a"), col("b").cast("bigint").alias("b")).collect() == [ + Row(a=7, b=8) + ] + if _ansi(spark): + for fractional in ["c", "d"]: + with pytest.raises(Exception, match="CAST_INVALID_INPUT"): + df.select(col(fractional).cast("int")).collect() + else: + res = df.select(col("c").cast("int").alias("c"), col("d").cast("bigint").alias("d")) + assert res.collect() == [Row(c=-2, d=1)] + + @pytest.mark.parametrize("value", ["abc", "1e3", "-", "", " ", "1_0"]) + def test_cast_invalid_string_to_integral(self, spark, value): + df = spark.createDataFrame([(value,)], StructType([StructField("s", StringType())])) + if USE_ACTUAL_SPARK and not _ansi(spark): + assert df.select(col("s").cast("int").alias("v")).collect() == [Row(v=None)] + else: + with pytest.raises(Exception, match=r"(?i)cast|convert"): + df.select(col("s").cast("int")).collect() + + def test_cast_fractional_strings_to_integral(self, spark): + values = [".5", "-.5", "3.", "+2.5", " -7.9 ", "-12.99999999999999999999", "9007199254740993.9"] + df = spark.createDataFrame([(v,) for v in values], StructType([StructField("s", StringType())])) + if _ansi(spark): + for value in values: + with pytest.raises(Exception, match="CAST_INVALID_INPUT"): + df.filter(col("s") == value).select(col("s").cast("int")).collect() + else: + res = df.select(col("s").cast("bigint").alias("l")).collect() + assert res == [Row(l=v) for v in [0, 0, 3, 2, -7, -12, 9007199254740993]] + assert df.filter(col("s") != values[-1]).select(col("s").cast("int").alias("i")).collect() == [ + Row(i=v) for v in [0, 0, 3, 2, -7, -12] + ] + + @pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf"), 1e20]) + def test_cast_out_of_range_double_to_int(self, spark, value): + df = spark.range(1).select(F.lit(value).alias("v")) + if USE_ACTUAL_SPARK and not _ansi(spark): + expected = {"nan": 0, "inf": 2147483647, "-inf": -2147483648}.get(str(value), 2147483647) + assert df.select(col("v").cast("int").alias("v")).collect() == [Row(v=expected)] + else: + with pytest.raises(Exception, match=r"(?i)overflow|out of range"): + df.select(col("v").cast("int")).collect() + + def test_isin_without_values(self, spark): + df = spark.createDataFrame([(1,)], StructType([StructField("a", LongType())])) + assert df.select(col("a").isin().alias("v"), col("a").isin([]).alias("w")).collect() == [Row(v=False, w=False)] + + def test_like_backslash_escape(self, spark): + df = spark.createDataFrame([("a%b",), ("axb",), ("a\\b",)], StructType([StructField("s", StringType())])) + res = df.select( + col("s").like("a\\%b").alias("pct"), + col("s").like("a\\\\b").alias("slash"), + col("s").ilike("A\\%B").alias("ipct"), + ).collect() + assert res == [ + Row(pct=True, slash=False, ipct=True), + Row(pct=False, slash=False, ipct=False), + Row(pct=False, slash=True, ipct=False), + ] + + def test_struct_values_collect_as_rows(self, spark): + schema = StructType([StructField("a", LongType()), StructField("b", LongType())]) + df = spark.createDataFrame([(7, -2)], schema) + res = df.select(struct("a", "b").alias("s"), array(struct(col("a").alias("z"))).alias("l")) + row = res.collect()[0] + assert row.s == Row(a=7, b=-2) + assert row.s.a == 7 + assert row.l == [Row(z=7)] + + @pytest.mark.timeout(60) + def test_column_is_not_iterable(self, spark): + with pytest.raises(PySparkTypeError, match=re.escape("[NOT_ITERABLE] Column is not iterable.")): + iter(col("a")) + with pytest.raises(PySparkTypeError, match="NOT_ITERABLE"): + list(col("a")) + + def test_column_has_no_truth_value(self, spark): + message = "[CANNOT_CONVERT_COLUMN_INTO_BOOL] Cannot convert column into bool: please use '&' for 'and'" + with pytest.raises(PySparkValueError, match=re.escape(message)): + bool(col("a")) + with pytest.raises(PySparkValueError, match="CANNOT_CONVERT_COLUMN_INTO_BOOL"): + _ = col("a") in [col("b")] + with pytest.raises(PySparkValueError, match="CANNOT_APPLY_IN_FOR_COLUMN"): + _ = "x" in col("a") + + def test_remainder(self, spark): + schema = StructType([StructField(c, LongType()) for c in "az"] + [StructField("d", DoubleType())]) + df = spark.createDataFrame([(-7, 3, -7.5)], schema) + res = df.select((col("a") % col("z")).alias("i"), (col("d") % 2).alias("d"), (7 % col("z")).alias("r")) + assert res.columns == ["i", "d", "r"] + assert res.collect() == [Row(i=-1, d=-1.5, r=1)] + assert df.select(col("a") % col("z"), 7 % col("z")).columns == ["(a % z)", "(7 % z)"] + + def test_remainder_by_zero(self, spark): + schema = StructType([StructField(c, LongType()) for c in "az"] + [StructField("d", DoubleType())]) + df = spark.createDataFrame([(7, 0, 7.5)], schema) + for remainder in [col("a") % col("z"), col("d") % col("z"), col("d") % 0.0, 7 % col("z")]: + res = df.select(remainder.alias("r")) + if _ansi(spark): + with pytest.raises(Exception, match="REMAINDER_BY_ZERO"): + res.collect() + else: + assert res.collect() == [Row(r=None)] + + @pytest.mark.parametrize("value", [{"k": 1}, {}, {"key": [1], "value": [2]}, [{"k": 1}]]) + def test_lit_of_a_dict_raises(self, spark, value): + with pytest.raises(SparkRuntimeException, match=re.escape("[UNSUPPORTED_FEATURE.LITERAL_TYPE]")): + spark.range(1).select(F.lit(value).alias("a")).collect() + + def test_struct_field_names(self, spark): + inner = StructType([StructField("a", LongType()), StructField("b", StringType())]) + schema = StructType([StructField("n", LongType()), StructField("s", StringType()), StructField("st", inner)]) + df = spark.createDataFrame([(1, "x", (2, "y"))], schema) + res = df.select( + struct( + col("n") + 1, + "s", + F.lit(2), + col("n").alias("m"), + F.upper("s"), + col("st.a"), + col("st")["b"], + df.n.cast("string"), + ).alias("o") + ) + assert res.schema["o"].dataType.fieldNames() == ["col1", "s", "col3", "m", "col5", "a", "b", "col8"] + assert res.collect() == [Row(o=Row(col1=2, s="x", col3=2, m=1, col5="X", a=2, b="y", col8="1"))] + assert df.select(struct("s", col("n") * 2)).columns == ["struct(s, (n * 2))"] + + def test_item_access_reads_the_field_of_that_name(self, spark): + df = spark.createDataFrame([((1, (2,), 3, 4),)], 's struct<`a.b`:int, a:struct, `c d`:int, `e"f`:int>') + res = df.select(col("s")["a.b"], df.s["a.b"], col("s")["a"]["b"], col("s")["c d"], col("s")['e"f']) + assert res.columns == ["s.a.b", "s.a.b", "s.a.b", "s.c d", 's.e"f'] + assert [tuple(r) for r in res.collect()] == [(1, 1, 2, 3, 4)] + assert df.select(col("s")["A.B"].alias("v")).collect() == [Row(v=1)] + + def test_item_access_on_a_computed_struct(self, spark): + computed = struct(F.lit(1).alias("a"), struct(F.lit(2).alias("c")).alias("b")) + res = spark.range(1).select(computed["a"], computed["b"]["c"]) + assert res.columns == ["struct(1 AS a, struct(2 AS c) AS b).a", "struct(1 AS a, struct(2 AS c) AS b).b.c"] + assert [tuple(r) for r in res.collect()] == [(1, 2)] + + +class TestCastTypeStrings: + @pytest.mark.parametrize( + ("source", "type_string", "simple", "spark_type", "value"), + [ + ("t", "timestamp", "timestamp", "TIMESTAMP", datetime.datetime(2024, 1, 2, 3, 4, 5)), + ("t", "timestamp_ltz", "timestamp", "TIMESTAMP", datetime.datetime(2024, 1, 2, 3, 4, 5)), + ("t", "timestamp_ntz", "timestamp_ntz", "TIMESTAMP_NTZ", datetime.datetime(2024, 1, 2, 3, 4, 5)), + ("t", "DATE", "date", "DATE", datetime.date(2024, 1, 2)), + ("n", "decimal", "decimal(10,0)", "DECIMAL(10,0)", Decimal("13")), + ("n", "decimal(5,2)", "decimal(5,2)", "DECIMAL(5,2)", Decimal("12.70")), + ("n", "float", "float", "FLOAT", 12.699999809265137), + ("i", "byte", "tinyint", "TINYINT", 3), + ("i", "short", "smallint", "SMALLINT", 3), + ("i", "long", "bigint", "BIGINT", 3), + ("i", "string", "string", "STRING", "3"), + ], + ) + def test_atomic_types(self, spark, source, type_string, simple, spark_type, value): + df = spark.createDataFrame([("2024-01-02 03:04:05", "12.7", "3")], "t string, n string, i string") + res = df.select(F.col(source).cast(type_string), F.trim(source).cast(type_string)) + assert res.columns == [source, f"CAST(trim({source}) AS {spark_type})"] + assert res.schema.simpleString() == f"struct<{source}:{simple},CAST(trim({source}) AS {spark_type}):{simple}>" + assert tuple(res.collect()[0]) == (value, value) + + def test_nested_types(self, spark): + schema = "a array, m map, st struct" + df = spark.createDataFrame([([1, 2], {"k": 1}, (1, "x"))], schema) + res = df.select( + F.col("a").cast("array"), + F.col("m").cast("map"), + F.col("st").cast("struct"), + ) + assert res.schema.simpleString() == "struct,m:map,st:struct>" + assert res.collect() == [Row(a=[1, 2], m={"k": 1}, st=Row(p="1", q="x"))] + + def test_duckdb_only_type(self, spark): + df = spark.createDataFrame([(3,)], "i int") + if USE_ACTUAL_SPARK: + with pytest.raises(ParseException): + df.select(F.col("i").cast("hugeint")).collect() + else: + res = df.select(F.col("i").cast("hugeint")) + assert res.schema.simpleString() == "struct" + assert res.collect() == [Row(i=3)] + + @pytest.mark.parametrize("type_string", ["varchar(10)", "char(3)", "CHARACTER(2)"]) + def test_char_and_varchar_cast_to_string(self, spark, type_string): + df = spark.createDataFrame([(" ab ",)], "s string") + res = df.select(F.col("s").cast(type_string), F.trim("s").cast(type_string)) + assert res.columns == ["s", "CAST(trim(s) AS STRING)"] + assert res.schema.simpleString() == "struct" + assert res.collect() == [Row(" ab ", "ab")] + + def test_nested_char_and_varchar_cast_to_string(self, spark): + df = spark.createDataFrame([(["x"], {"k": "v"})], "a array, m map") + res = df.select(F.col("a").cast("array"), F.col("m").cast("map")) + assert res.schema.simpleString() == "struct,m:map>" + assert res.collect() == [Row(a=["x"], m={"k": "v"})] + + @pytest.mark.parametrize( + ("source", "type_string", "value"), + [ + ("s", "VARCHAR", "{}"), + ("s", "json", "{}"), + ("st", "STRUCT(b INT)", Row(b=1)), + ("a", "INTEGER[]", [1]), + ("a", "INTEGER[1]", (1,)), + ], + ) + def test_duckdb_spellings(self, spark, source, type_string, value): + df = spark.createDataFrame([("{}", (1.7,), [1.7])], "s string, st struct, a array") + if USE_ACTUAL_SPARK: + with pytest.raises(ParseException): + df.select(F.col(source).cast(type_string)).collect() + else: + assert df.select(F.col(source).cast(type_string).alias("v")).collect() == [Row(v=value)] + + @pytest.mark.parametrize("type_string", ["inetger", "varchar", "varchar(x)", "array", "a int, b string"]) + def test_invalid_type_string_raises_parse_exception(self, spark, type_string): + df = spark.createDataFrame([("a",)], "s string") + if type_string == "varchar" and not USE_ACTUAL_SPARK: + assert df.select(F.col("s").cast(type_string)).collect() == [Row(s="a")] + return + with pytest.raises(ParseException): + df.select(F.col("s").cast(type_string)).collect() + + def test_bare_interval_raises(self, spark): + df = spark.createDataFrame([("1 day",)], "s string") + if not USE_ACTUAL_SPARK: + from duckdb.experimental.spark.exception import ContributionsAcceptedError as expected + elif pyspark.__version__.startswith("3."): + expected = PySparkValueError + else: + from spark_namespace.errors import PySparkNotImplementedError as expected + with pytest.raises(expected): + df.select(F.col("s").cast("interval")).collect() + + +class TestSparkCastSemantics: + def test_decimal_to_integral_truncates_exactly(self, spark): + df = spark.createDataFrame( + [(Decimal("100000000000.999999"), Decimal("-3.999999999"), Decimal("3.99999999999999999"))], + "a decimal(18,6), b decimal(18,9), c decimal(38,17)", + ) + res = df.select(F.col("a").cast("long"), F.col("b").cast("int"), F.col("c").cast("bigint")) + assert res.collect() == [Row(a=100000000000, b=-3, c=3)] + + def test_decimal_to_small_integral_types(self, spark): + df = spark.createDataFrame( + [(Decimal("1.99999999999999999"), Decimal("-0.99999999999999999999"))], "a decimal(18,17), b decimal(38,20)" + ) + res = df.select(F.col("a").cast("int"), F.col("a").cast("tinyint").alias("t"), F.col("b").cast("short")) + # Spark in ANSI mode casts a decimal to int, smallint or tinyint through its nearest double. + if _ansi(spark): + assert res.collect() == [Row(a=2, t=2, b=-1)] + else: + assert res.collect() == [Row(a=1, t=1, b=0)] + + def test_integral_casts_truncate_inside_nested_types(self, spark): + df = spark.createDataFrame( + [([1.7, -1.7, None], [[1.7, -2.5]], {"a": 1.7, "b": -2.9}, (-1.7, Decimal("2.99")))], + "a array, n array>, m map, s struct", + ) + res = df.select( + F.col("a").cast("array"), + F.col("n").cast("array>"), + F.col("m").cast("map"), + F.col("s").cast("struct"), + ) + assert res.columns == ["a", "n", "m", "s"] + assert res.schema.simpleString() == ( + "struct,n:array>,m:map,s:struct>" + ) + assert res.collect() == [Row(a=[1, -1, None], n=[[1, -2]], m={"a": 1, "b": -2}, s=Row(x=-1, y=2))] + + def test_map_key_cast_truncates(self, spark): + df = spark.createDataFrame([({1.7: "a"},)], "m map") + if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3."): + with pytest.raises(Exception, match="DATATYPE_MISMATCH"): + df.select(F.col("m").cast("map")).collect() + else: + assert df.select(F.col("m").cast("map")).collect() == [Row(m={1: "a"})] + + def test_struct_cast_maps_fields_by_position(self, spark): + df = spark.createDataFrame([((5, "z"),), (None,)], "s struct") + renamed = df.select(F.col("s").cast("struct")) + assert renamed.schema.simpleString() == "struct>" + assert renamed.collect() == [Row(s=Row(p=5, y="z")), Row(s=None)] + swapped = df.select(F.col("s").cast("struct")) + assert swapped.collect() == [Row(s=Row(q="5", p="z")), Row(s=None)] + + @pytest.mark.parametrize("target", ["struct", "struct"]) + def test_struct_cast_with_another_field_count_raises(self, spark, target): + df = spark.createDataFrame([((5, "z"),)], "s struct") + with pytest.raises(Exception, match=r"(?i)cast|size|out of range"): + df.select(F.col("s").cast(target)).collect() + + def test_nested_struct_casts_map_fields_by_position(self, spark): + df = spark.createDataFrame( + [([(5, "z")], {"k": (1, "w")}, ((1.7,), "z"))], + "a array>, m map>, " + "s struct,q:string>", + ) + res = df.select( + F.col("a").cast("array>"), + F.col("m").cast("map>"), + F.col("s").cast("struct,c:string>"), + F.col("s").cast("struct,c:string>").alias("d"), + ) + assert res.collect() == [ + Row( + a=[Row(u=5, v="z")], + m={"k": Row(u=1, v="w")}, + s=Row(a=Row(b=1), c="z"), + d=Row(a=Row(b=1.7), c="z"), + ) + ] + + def test_cast_evaluates_its_input_once(self, spark): + df = spark.createDataFrame([(1,), (2,), (3,)], "v bigint") + if USE_ACTUAL_SPARK: + spark.udf.register("counted", lambda v: v * 1.5, DoubleType()) + else: + calls = [] + + def counted(v): + calls.append(v) + return v * 1.5 + + # A UDF with side effects is evaluated wherever it occurs, so the count shows how often the cast reads it. + spark.conn.create_function("counted", counted, ["BIGINT"], "DOUBLE", side_effects=True) + res = df.select( + F.expr("counted(v)").cast("int").alias("i"), + F.struct(F.expr("counted(v)").alias("x")).cast("struct").alias("s"), + F.factorial(F.expr("counted(v)") + 1).alias("f"), + F.char(F.expr("counted(v)") + 64).alias("c"), + ) + assert sorted(res.collect()) == [ + Row(i=1, s=Row(x=1), f=2, c="A"), + Row(i=3, s=Row(x=3), f=24, c="C"), + Row(i=4, s=Row(x=4), f=120, c="D"), + ] + if not USE_ACTUAL_SPARK: + assert len(calls) == 12 + + def test_string_to_binary_takes_utf8_bytes(self, spark): + df = spark.createDataFrame([("é", bytearray(b"\xff\x00\\x")), ("a\\x41", None)], "s string, b binary") + res = df.select( + F.col("s").cast("binary"), F.col("b").cast("binary"), F.array("s").cast("array").alias("a") + ) + assert res.collect() == [ + Row(s=b"\xc3\xa9", b=b"\xff\x00\\x", a=[b"\xc3\xa9"]), + Row(s=b"a\\x41", b=None, a=[b"a\\x41"]), + ] + + def test_integral_to_binary(self, spark): + df = spark.range(1) + if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3."): + assert df.select(F.col("id").cast("binary")).collect() == [Row(id=bytearray(8))] + else: + with pytest.raises(Exception, match=r"(?i)cast"): + df.select(F.col("id").cast("binary")).collect() diff --git a/tests/fast/spark/test_spark_column_names.py b/tests/fast/spark/test_spark_column_names.py new file mode 100644 index 000000000..83b7f85d6 --- /dev/null +++ b/tests/fast/spark/test_spark_column_names.py @@ -0,0 +1,828 @@ +import pytest + +_ = pytest.importorskip("duckdb.experimental.spark") + +import datetime +from decimal import Decimal + +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException +from spark_namespace.sql import functions as F +from spark_namespace.sql import types as T + +if USE_ACTUAL_SPARK: + import pyspark + + SPARK_3 = pyspark.__version__.startswith("3.") +else: + SPARK_3 = False + + +@pytest.fixture +def df(spark): + schema = T.StructType( + [ + T.StructField("a", T.LongType()), + T.StructField("b", T.DoubleType()), + T.StructField("s", T.StringType()), + T.StructField("flag", T.BooleanType()), + T.StructField("d", T.DateType()), + T.StructField("arr", T.ArrayType(T.LongType())), + T.StructField("st", T.StructType([T.StructField("x", T.LongType()), T.StructField("y", T.StringType())])), + T.StructField("CamelCase", T.LongType()), + ] + ) + return spark.createDataFrame([(1, 1.5, "x", True, datetime.date(2020, 1, 2), [1, 2], (1, "q"), 3)], schema) + + +@pytest.fixture +def df2(spark): + schema = T.StructType([T.StructField("a", T.LongType()), T.StructField("c", T.StringType())]) + return spark.createDataFrame([(1, "c1")], schema) + + +def names(df, *cols): + return df.select(*cols).columns + + +class TestColumnReferenceNames: + def test_plain_and_bound_references(self, df): + assert names(df, "a", F.col("b"), df.s, df["flag"]) == ["a", "b", "s", "flag"] + + def test_requested_case_is_kept(self, df): + assert names(df, F.col("A"), "B", F.col("camelcase"), "CAMELCASE") == ["A", "B", "camelcase", "CAMELCASE"] + assert df.select(F.col("A")).collect()[0].asDict() == {"A": 1} + + def test_bound_references_keep_requested_case_and_cast_name(self, df): + assert names(df, df["S"], df.a.cast("string"), df.a.cast("int").cast("string")) == ["S", "a", "a"] + + def test_missing_column_is_reported_as_missing(self, df): + with pytest.raises(Exception, match=r"not found|cannot be resolved"): + df.select("missing").collect() + with pytest.raises(Exception, match=r"not found|cannot be resolved"): + df.groupBy("missing").count().collect() + + def test_star(self, df): + assert names(df, "*") == df.columns + assert names(df, F.col("*")) == df.columns + + def test_struct_fields(self, df): + assert names(df, F.col("st.x"), "st.y", df.st.x, df.st["y"]) == ["x", "y", "st.x", "st.y"] + assert names(df, F.upper(df.st.y), F.upper(F.col("st.y")), F.upper("st.y")) == ["upper(st.y)"] * 3 + + def test_names_with_spaces(self, spark): + df = spark.createDataFrame([(1,)], ["my col"]) + assert names(df, F.col("my col"), F.col("my col") + 1) == ["my col", "(my col + 1)"] + + def test_backtick_quoted_names(self, spark): + df = spark.createDataFrame([("p", 2, (3,))], "a string, `a.b` int, st struct") + res = df.select(F.col("`a`"), "`a`", F.col("`a.b`"), "`a.b`", F.col("`st`.`x`"), F.upper("`a`")) + assert res.columns == ["a", "a", "a.b", "a.b", "x", "upper(a)"] + assert res.collect() == [("p", "p", 2, 2, 3, "P")] + + @pytest.mark.parametrize("name", ["a`b", "`a`b", "a``b", "`a", "a..b", "a.", "`a` "]) + def test_invalid_attribute_names_raise(self, spark, name): + # Each of these would otherwise read one of the columns a, ab or a`b. + df = spark.createDataFrame([(1, 2, 3)], ["a", "ab", "a`b"]) + for select in [lambda: df.select(F.col(name)), lambda: df.select(name)]: + with pytest.raises(AnalysisException, match=r"Syntax error in (the )?attribute name"): + select().collect() + + def test_struct_path_through_a_repeated_name_is_ambiguous(self, spark): + df = spark.createDataFrame([((1,), (2,))], "s struct, t struct").toDF("s", "s") + for name in ["s.x", "`s`.x", F.col("s.x")]: + with pytest.raises(AnalysisException, match="AMBIGUOUS_REFERENCE"): + df.select(name).collect() + + def test_star_of_a_struct_or_an_alias_through_col(self, spark): + df = spark.createDataFrame([((1, "a"), 2)], "s struct, n int") + assert df.select(F.col("s.*")).collect() == [(1, "a")] + assert df.select(F.col("`s`.*"), "n").columns == ["x", "y", "n"] + assert df.alias("q").select(F.col("q.*")).columns == ["s", "n"] + # Spark allows such a star only in select. + for make in [lambda: df.withColumn("z", F.col("s.*")), lambda: df.alias("q").select(F.col("q.*").alias("w"))]: + with pytest.raises(Exception): # noqa: B017, PT011 + make().collect() + repeated = spark.createDataFrame([((1,), (2,))], "s struct, t struct").toDF("s", "s") + for star in ["s.*", F.col("s.*")]: + with pytest.raises(AnalysisException, match="AMBIGUOUS_REFERENCE"): + repeated.select(star).collect() + + def test_backtick_quoted_names_after_a_join(self, spark): + left = spark.createDataFrame([(1, (5,))], "`k.k` int, st struct") + right = spark.createDataFrame([(1, "y"), (2, "z")], "`k.k` int, w string") + outer = left.join(right, "k.k", "outer").select("`k.k`") + assert outer.columns == ["k.k"] + assert sorted(outer.collect()) == [(1,), (2,)] + inner = left.join(right, "k.k").select("`k.k`", "`st`.`x`", "w") + assert inner.columns == ["k.k", "x", "w"] + assert inner.collect() == [(1, 5, "y")] + # A quoted name with a dot is one column name, not a struct field. + with pytest.raises(Exception, match=r"st\.x"): + left.join(right, "k.k").select("`st.x`").collect() + joined = left.join(right, "k.k") + assert joined.select("*").columns == ["k.k", "st", "w"] + assert joined.select(joined["st.x"]).collect() == [(5,)] + + def test_nested_struct_paths(self, spark): + # The field b next to a makes a path that drops a middle part read another value. + schema = "s struct>, b: struct>" + df = spark.createDataFrame([({"a": {"b": {"c": 1}}, "b": {"c": 2}},)], schema) + res = df.select( + F.col("s.a.b.c"), + "s.a.b.c", + df.s.a.b.c, + df["s"]["a"]["b"]["c"], + F.col("s")["a"]["b"]["c"], + F.col("s.a.b"), + "s.a.b", + df.s.a.b, + F.col("s.b.c"), + ) + assert res.columns == ["c", "c", "s.a.b.c", "s.a.b.c", "s.a.b.c", "b", "b", "s.a.b", "c"] + assert res.collect() == [(1, 1, 1, 1, 1, (1,), (1,), (1,), 2)] + + def test_selecting_a_generated_name(self, df): + assert names(df.select(F.abs("a")), F.col("abs(a)") + 1) == ["(abs(a) + 1)"] + assert names(df.select(df.a + 1), "(a + 1)") == ["(a + 1)"] + assert names(df.select(F.col("A")), "A") == ["A"] + + +class TestFunctionCallNames: + def test_one_argument(self, df): + assert names(df, F.abs(df.a), F.abs("a"), F.abs(F.col("A")), F.upper(df.s), F.lower("s")) == [ + "abs(a)", + "abs(a)", + "abs(A)", + "upper(s)", + "lower(s)", + ] + + def test_spark_function_name_differs_from_python_name(self, df): + assert names(df, F.sqrt("a"), F.exp("a"), F.ceil("b"), F.floor("b"), F.signum("b"), F.ucase("s")) == [ + "SQRT(a)", + "EXP(a)", + "CEIL(b)", + "FLOOR(b)", + "SIGNUM(b)", + "ucase(s)", + ] + assert names(df, F.pow("a", "b"), F.atan2("b", "b"), F.e(), F.pi()) == [ + "POWER(a, b)", + "ATAN2(b, b)", + "E()", + "PI()", + ] + assert names(df, F.curdate(), F.isnull("a"), F.isnotnull("a"), F.positive("a"), F.negative("a")) == [ + "current_date()", + "(a IS NULL)", + "(a IS NOT NULL)", + "(+ a)", + "negative(a)", + ] + + def test_function_implemented_by_another_expression(self, df): + assert names(df, F.array_max("arr"), F.to_date("s"), F.equal_null("a", "a"), F.nvl("a", "a")) == [ + "array_max(arr)", + "to_date(s)", + "equal_null(a, a)", + "nvl(a, a)", + ] + + def test_several_arguments(self, df): + assert names( + df, + F.round("b", 1), + F.greatest("a", "b"), + F.least(df.a, df.b), + F.coalesce("a", "b"), + F.concat_ws("-", "s", "s"), + F.substring("s", 1, 2), + F.lpad("s", 3, "-"), + F.date_trunc("year", "d"), + F.array_contains("arr", 1), + F.struct(df.a, F.col("s").alias("z")), + ) == [ + "round(b, 1)", + "greatest(a, b)", + "least(a, b)", + "coalesce(a, b)", + "concat_ws(-, s, s)", + "substring(s, 1, 2)", + "lpad(s, 3, -)", + "date_trunc(year, d)", + "array_contains(arr, 1)", + "struct(a, s AS z)", + ] + + def test_arguments_spark_adds(self, df): + assert names( + df, + F.round("b"), + F.bround("b"), + F.regexp_replace("s", "x", "y"), + F.sort_array("arr"), + F.split("s", ","), + F.substr("s", F.lit(1)), + ) == [ + "round(b, 0)", + "bround(b, 0)", + "regexp_replace(s, x, y, 1)", + "sort_array(arr, true)", + "split(s, ,, -1)", + "substr(s, 1, 2147483647)", + ] + + def test_special_forms(self, df): + assert names(df, F.like("s", F.lit("x%"), F.lit("!")), F.ilike("s", F.lit("x%"), F.lit("!"))) == [ + "s LIKE x%" if SPARK_3 else "s LIKE x% ESCAPE '!'", + "ilike(s, x%)", + ] + assert names(df, df.a.asc_nulls_last(), df.a.desc(), df.a.desc_nulls_first()) == [ + "a ASC NULLS LAST", + "a DESC NULLS LAST", + "a DESC NULLS FIRST", + ] + assert names( + df, + F.like("s", F.lit("x%")), + F.extract(F.lit("YEAR"), "d"), + F.datepart(F.lit("YEAR"), "d"), + F.asc("a"), + F.desc_nulls_first("a"), + F.call_function("abs", "a"), + ) == [ + "s LIKE x%", + "extract(YEAR FROM d)", + "datepart(YEAR FROM d)", + "a ASC NULLS FIRST", + "a DESC NULLS FIRST", + "abs(a)", + ] + + @pytest.mark.skipif(USE_ACTUAL_SPARK, reason="Spark's Column has no nulls_first or nulls_last") + def test_null_ordering_without_direction(self, df): + assert names(df, F.col("a").nulls_first(), F.col("a").asc().nulls_last(), F.col("a").desc().nulls_first()) == [ + "a NULLS FIRST", + "a ASC NULLS LAST", + "a DESC NULLS FIRST", + ] + + def test_nested(self, df): + assert names( + df, + F.abs(F.sqrt("a")), + F.upper(F.lower("s")), + F.round(F.col("b") * 2, 1), + F.concat(F.upper("s"), F.lit("!")), + F.abs(df.a.alias("x")), + ) == ["abs(SQRT(a))", "upper(lower(s))", "round((b * 2), 1)", "concat(upper(s), !)", "abs(a AS x)"] + + +class TestOperatorNames: + def test_arithmetic(self, df): + assert names(df, df.a + 1, df.a - df.b, df.a * 2, df.a / 2, df.a % 2, df.a**2) == [ + "(a + 1)", + "(a - b)", + "(a * 2)", + "(a / 2)", + "(a % 2)", + "POWER(a, 2)", + ] + + def test_reflected_arithmetic(self, df): + assert names(df, 1 + df.a, 2 - df.a, 2 * df.a, 2 / df.a, 2**df.a) == [ + "(a + 1)", + "(2 - a)", + "(a * 2)", + "(2 / a)", + "POWER(2, a)", + ] + + def test_precedence_is_explicit(self, df): + assert names(df, df.a + df.b * 2, (df.a + df.b) * 2) == ["(a + (b * 2))", "((a + b) * 2)"] + + def test_comparison(self, df): + assert names(df, df.a == 1, df.a != 1, df.a < 1, df.a <= 1, df.a > 1, df.a >= 1, df.a > 1) == [ + "(a = 1)", + "(NOT (a = 1))", + "(a < 1)", + "(a <= 1)", + "(a > 1)", + "(a >= 1)", + "(a > 1)", + ] + + def test_boolean(self, df): + assert names(df, (df.a > 1) & (df.b < 2), (df.a > 1) | df.flag, ~df.flag, ~(df.a > 1)) == [ + "((a > 1) AND (b < 2))", + "((a > 1) OR flag)", + "(NOT flag)", + "(NOT (a > 1))", + ] + + def test_unary_minus(self, df): + expected = ["(- a)", "(- b)"] if SPARK_3 else ["negative(a)", "negative(b)"] + assert names(df, -df.a, -F.col("b")) == expected + + +class TestLiteralNames: + def test_scalars(self, df): + assert names(df, F.lit(1), F.lit(-1), F.lit(2**40), F.lit("x"), F.lit("it's"), F.lit(True), F.lit(None)) == [ + "1", + "-1", + "1099511627776", + "x", + "it's", + "true", + "NULL", + ] + + def test_doubles_are_named_like_java(self, df): + assert names(df, F.lit(1.5), F.lit(100.0), F.lit(-2.5), F.lit(0.001), F.lit(1e-5), F.lit(12345678.0)) == [ + "1.5", + "100.0", + "-2.5", + "0.001", + "1.0E-5", + "1.2345678E7", + ] + assert names(df, F.lit(1e20), F.lit(float("inf")), F.lit(float("-inf"))) == [ + "1.0E20", + "Infinity", + "-Infinity", + ] + + def test_dates_decimals_and_bytes(self, df): + assert names( + df, + F.lit(datetime.date(2020, 1, 2)), + F.lit(datetime.datetime(2020, 1, 2, 3, 4, 5)), + F.lit(datetime.datetime(2020, 1, 2, 3, 4, 5, 120000)), + F.lit(Decimal("1.50")), + F.lit(Decimal("1E+2")), + F.lit(bytearray(b"ab")), + ) == [ + "DATE '2020-01-02'", + "TIMESTAMP '2020-01-02 03:04:05'", + "TIMESTAMP '2020-01-02 03:04:05.12'", + "1.50", + "100", + "X'6162'", + ] + + def test_timestamps_with_a_zone_are_named_by_their_local_time(self, df): + value = datetime.datetime(2024, 1, 2, 3, 4, 5, 500000, tzinfo=datetime.timezone(datetime.timedelta(hours=2))) + local = value.astimezone().replace(tzinfo=None) + text = f"TIMESTAMP '{local.isoformat(' ').rstrip('0')}'" + res = df.select(F.lit(value), F.lit(value) > datetime.datetime(2000, 1, 1)) + assert res.columns == [text, f"({text} > TIMESTAMP '2000-01-01 00:00:00')"] + assert res.collect()[0][0] == local + + def test_operands(self, df): + assert names( + df, + df.a + F.lit(1.5), + df.s == "x", + df.s == "it's", + df.a + None, + df.b > 1e20, + df.d > datetime.date(2020, 1, 1), + F.upper(F.lit("x")), + ) == [ + "(a + 1.5)", + "(s = x)", + "(s = it's)", + "(a + NULL)", + "(b > 1.0E20)", + "(d > DATE '2020-01-01')", + "upper(x)", + ] + + +class TestCastNames: + def test_cast_of_a_column_keeps_its_name(self, df): + assert names(df, F.col("a").cast("string"), F.col("a").cast(T.IntegerType()), F.col("A").cast("double")) == [ + "a", + "a", + "A", + ] + assert names(df, F.col("a").cast("int").cast("string")) == ["a"] + + def test_cast_of_an_alias_keeps_the_alias(self, df): + assert names(df, F.abs("a").alias("x").cast("string")) == ["x"] + + def test_cast_of_an_expression(self, df): + assert names( + df, + F.lit(1).cast("string"), + (df.a + 1).cast("string"), + F.abs("a").cast("string"), + F.lit(1).cast("long"), + F.lit(1).cast("integer"), + F.lit(1).cast(T.ShortType()), + F.lit(1).cast(T.DecimalType(5, 2)), + F.lit("x").cast("binary"), + F.lit("2020-01-01").cast("date"), + F.col("st.x").cast("string"), + df.st.x.cast("string"), + ) == [ + "CAST(1 AS STRING)", + "CAST((a + 1) AS STRING)", + "CAST(abs(a) AS STRING)", + "CAST(1 AS BIGINT)", + "CAST(1 AS INT)", + "CAST(1 AS SMALLINT)", + "CAST(1 AS DECIMAL(5,2))", + "CAST(x AS BINARY)", + "CAST(2020-01-01 AS DATE)", + "CAST(st.x AS STRING)", + "CAST(st.x AS STRING)", + ] + + def test_cast_to_a_complex_type_is_named_after_its_input(self, df): + assert names( + df, + F.array(F.lit(1), F.lit(2)).cast("array"), + F.array(F.lit(1), F.lit(2)).cast("array"), + F.array("a", "a").cast(T.ArrayType(T.StringType())), + F.struct(F.lit(1).alias("p")).cast("struct"), + F.lit(None).cast("array"), + F.array(F.lit(1)).cast("array").cast("array"), + F.array(F.lit(1)).cast("string"), + ) == [ + "array(1, 2)", + "array(1, 2)", + "array(a, a)", + "struct(1 AS p)", + "NULL", + "array(1)", + "CAST(array(1) AS STRING)", + ] + res = df.select(F.struct(F.col("arr").cast("array"), F.array("a").cast("array"))) + assert res.columns == ["struct(arr, array(a))"] + assert res.schema.fields[0].dataType.fieldNames() == ["col1", "col2"] + + def test_cast_of_a_struct_path_to_a_complex_type(self, spark): + df = spark.createDataFrame([(([1, 2], (3,)),)], "st struct, inner: struct>") + res = df.select(F.col("st.arr").cast("array"), F.col("st.inner").cast("struct")) + assert res.columns == ["st.arr", "st.inner"] + assert res.collect() == [([1, 2], (3,))] + + +class TestAliasNames: + def test_alias(self, df): + assert names( + df, + df.a.alias("x"), + F.abs("a").alias("y"), + F.abs("a").alias("p").alias("q"), + (df.a + 1).alias("A Plus One"), + F.col("a").alias("A"), + ) == ["x", "y", "q", "A Plus One", "A"] + + def test_alias_inside_a_name_is_quoted_unless_a_plain_word(self, df): + aliases = ["a_1", "a b", "1", "a`b"] + res = df.select(F.struct(*(F.lit(i).alias(alias) for i, alias in enumerate(aliases, 1)))) + assert res.columns == ["struct(1 AS a_1, 2 AS `a b`, 3 AS `1`, 4 AS `a``b`)"] + assert res.schema.fields[0].dataType.fieldNames() == aliases + + def test_empty_string_literal(self, df): + assert df.select(F.lit("")).collect() == [("",)] + + def test_empty_top_level_alias(self, df): + makes = [ + lambda: df.select(F.lit(1).alias("")), + lambda: df.withColumn("", F.lit(1)), + lambda: df.withColumns({"": F.lit(1)}), + lambda: df.select("a").withColumnRenamed("a", ""), + lambda: df.select("a").withColumnsRenamed({"a": ""}), + lambda: df.select("a").toDF(""), + ] + if USE_ACTUAL_SPARK: + assert all("" in make().columns for make in makes) + return + from duckdb.experimental.spark.exception import ContributionsAcceptedError + + # DuckDB has no column with an empty name. + for make in makes: + with pytest.raises(ContributionsAcceptedError, match="empty name"): + make() + + def test_empty_alias_names_a_struct_field(self, df): + res = df.select( + F.struct(F.lit(1).alias("")).alias("one"), + F.struct(F.lit(1).alias(""), F.lit("x").alias("")).alias("two"), + F.arrays_zip(F.array(F.lit(1)).alias("")).alias("zip"), + ) + assert res.schema.simpleString() == "struct,two:struct<:int,:string>,zip:array>>" + assert res.collect() == [((1,), (1, "x"), [(1,)])] + assert names(df, F.struct(F.lit(1).alias(""))) == ["struct(1 AS ``)"] + + def test_struct_with_an_empty_and_another_field_name(self, df): + if USE_ACTUAL_SPARK: + res = df.select(F.struct(F.lit(1).alias(""), F.lit(2).alias("b")).alias("s")) + assert res.schema.simpleString() == "struct>" + else: + from duckdb.experimental.spark.exception import ContributionsAcceptedError + + with pytest.raises(ContributionsAcceptedError, match="both empty and other field names"): + F.struct(F.lit(1).alias(""), F.lit(2).alias("b")) + + +class TestConditionalAndPredicateNames: + def test_when_otherwise(self, df): + assert names( + df, + F.when(df.a > 1, 1), + F.when(df.a > 1, 1).otherwise(0), + F.when(df.a > 1, "x").when(df.a > 0, "y"), + F.when(df.a > 1, "x").when(df.a > 0, "y").otherwise("z"), + F.when(df.flag, df.a).otherwise(df.b), + F.when(df.a > 1, F.lit(None)), + ) == [ + "CASE WHEN (a > 1) THEN 1 END", + "CASE WHEN (a > 1) THEN 1 ELSE 0 END", + "CASE WHEN (a > 1) THEN x WHEN (a > 0) THEN y END", + "CASE WHEN (a > 1) THEN x WHEN (a > 0) THEN y ELSE z END", + "CASE WHEN flag THEN a ELSE b END", + "CASE WHEN (a > 1) THEN NULL END", + ] + + def test_null_tests(self, df): + assert names(df, df.a.isNull(), df.a.isNotNull(), F.col("A").isNull()) == [ + "(a IS NULL)", + "(a IS NOT NULL)", + "(A IS NULL)", + ] + + def test_isin(self, df): + assert names(df, df.a.isin(1, 2), df.a.isin([1, 2]), df.s.isin("x", "y"), df.a.isin(df.a, 2)) == [ + "(a IN (1, 2))", + "(a IN (1, 2))", + "(s IN (x, y))", + "(a IN (a, 2))", + ] + + def test_string_predicates(self, df): + assert names( + df, + df.s.like("x%"), + df.s.ilike("X%"), + df.s.rlike("^x"), + df.s.contains("x"), + df.s.startswith("x"), + df.s.endswith("x"), + ) == ["s LIKE x%", "ilike(s, X%)", "RLIKE(s, ^x)", "contains(s, x)", "startswith(s, x)", "endswith(s, x)"] + + +class TestAggregateNames: + def test_agg(self, df): + assert df.groupBy("s").agg( + F.sum("a"), + F.sum(df.a), + F.count("a"), + F.count(F.lit(1)), + F.avg("a"), + F.mean("a"), + F.first("a"), + F.last("a"), + F.collect_list("a"), + F.array_agg("a"), + F.stddev("a"), + F.sum(df.a + 1), + F.sum("a") + 1, + F.sum("a").alias("total"), + ).columns == [ + "s", + "sum(a)", + "sum(a)", + "count(a)", + "count(1)", + "avg(a)", + "avg(a)", + "first(a)", + "last(a)", + "collect_list(a)", + "collect_list(a)", + "stddev(a)", + "sum((a + 1))", + "(sum(a) + 1)", + "total", + ] + + def test_aggregate_spark_names_differently_per_version(self, df): + assert df.groupBy("s").agg(F.variance("a")).columns == ["s", "var_samp(a)" if SPARK_3 else "variance(a)"] + + def test_grouping_keys(self, df): + assert df.groupBy("S").agg(F.sum("a")).columns == ["S", "sum(a)"] + assert df.groupBy(df.s).agg(F.sum("a")).columns == ["s", "sum(a)"] + assert df.groupBy(F.upper("s")).agg(F.sum("a")).columns == ["upper(s)", "sum(a)"] + + def test_grouped_data_shortcuts(self, df): + assert df.groupBy("s").count().columns == ["s", "count"] + assert df.groupBy("s").sum("a").columns == ["s", "sum(a)"] + assert df.groupBy("s").avg("a").columns == ["s", "avg(a)"] + assert df.groupBy("s").mean("a").columns == ["s", "avg(a)"] + assert df.groupBy("s").min("a", "b").columns == ["s", "min(a)", "min(b)"] + assert df.groupBy("S").count().columns == ["S", "count"] + assert df.groupBy(F.col("S")).max("a").columns == ["S", "max(a)"] + assert df.groupBy().sum("a").columns == ["sum(a)"] + + +class TestNamesAfterOtherOperations: + def test_after_joins(self, df, df2): + assert df.join(df2, "a").select(F.abs(df2.a), df2.c, F.upper(df2.c)).columns == ["abs(a)", "c", "upper(c)"] + assert df.join(df2, df.a == df2.a).select(df.a, df2.c).columns == ["a", "c"] + assert df.join(df2, "a").select("c", F.upper("c")).columns == ["c", "upper(c)"] + assert df.crossJoin(df2).select(df2.c).columns == ["c"] + + def test_qualifiers_of_aliased_dataframes_are_left_out(self, df, df2): + joined = df.alias("l").join(df2.alias("r"), "a") + assert joined.select(F.col("l.a"), F.col("r.c"), F.upper(F.col("r.c")), F.upper("r.c")).columns == [ + "a", + "c", + "upper(c)", + "upper(c)", + ] + assert joined.select(F.col("l.st.x")).columns == ["x"] + assert joined.select(F.col("l.a").cast("string"), F.col("l.st.x").cast("string")).columns == [ + "a", + "CAST(st.x AS STRING)", + ] + + def test_qualifier_that_is_also_a_column_name_is_left_out(self, spark): + struct_only = spark.createDataFrame([((-1,),)], "s struct").alias("s") + assert struct_only.select(F.abs("s.s.x"), F.abs("s.x"), F.col("s.s.x").cast("string")).columns == [ + "abs(s.x)", + "abs(s.x)", + "CAST(s.x AS STRING)", + ] + both = spark.createDataFrame([(-1, (2,))], "x int, s struct").alias("s") + res = both.select(F.abs("s.x"), F.col("s.x") + 0, F.col("s.x").cast("string")) + assert res.columns == ["abs(x)", "(x + 0)", "x"] + assert res.collect() == [(1, -1, "-1")] + # Spark 3.5.9 names a grouping key as it was written, with its qualifier. + assert both.groupBy(F.abs("s.x")).count().columns == ["abs(s.x)" if SPARK_3 else "abs(x)", "count"] + + def test_after_filter_sort_and_limit(self, df): + assert df.filter(df.a > 0).select(F.abs(df.a)).columns == ["abs(a)"] + assert df.orderBy(df.a).select(F.abs(df.a)).columns == ["abs(a)"] + assert df.limit(1).select(df.a + 1).columns == ["(a + 1)"] + + def test_with_column_uses_the_given_name(self, df): + assert df.withColumn("n", F.abs("a")).columns[-1] == "n" + assert df.select(F.abs("a")).withColumn("abs(a)", F.lit(1)).columns == ["abs(a)"] + + def test_row_fields_and_sql_expressions(self, df): + row = df.select(F.abs(df.a), df.b + 1, F.expr("a + 2"), F.expr("upper(s)")).collect()[0] + assert list(row.asDict()) == ["abs(a)", "(b + 1)", "(a + 2)", "upper(s)"] + + +def test_no_internal_relation_name_in_any_column_name(df, df2): + cols = [ + df.a, + df["b"], + -df.a, + ~df.flag, + df.a + df.b, + df.a.isin(df.a), + df.a.cast("string"), + df.a.alias("x"), + df.st.x, + df.s.like("x%"), + df.s.contains(df.s), + F.when(df.flag, df.a).otherwise(df.b), + F.abs(df.a), + F.concat_ws("-", df.s, df.s), + F.coalesce(df.a, df.b), + F.struct(df.a, df.s), + F.array(df.a, df.a), + F.array_max(df.arr), + F.asc(df.a), + F.date_diff(df.d, df.d), + F.to_date(df.s), + F.nvl2(df.a, df.a, df.a), + F.equal_null(df.a, df.a), + F.array_insert(df.arr, 1, df.a), + ] + joined = df.join(df2, df.a == df2.a).select(F.abs(df.a), F.upper(df2.c), df2.c) + for name in df.select(*cols).columns + joined.columns: + assert "unnamed_relation" not in name + + +class TestColumnRepr: + def test_repr_shows_the_name(self, df): + cols = [ + F.col("a"), + df.a, + F.col("A"), + F.col("st.x"), + F.abs("a"), + F.round("b"), + F.mean("a"), + F.upper(F.lower("s")), + F.lit(1.5), + F.lit(1e20), + F.lit(True), + F.lit(None), + df.a.alias("x"), + df.a.cast("string"), + df.a.asc_nulls_last(), + F.struct("a", "s"), + F.negative("a"), + ] + assert [repr(c) for c in cols] == [ + "Column<'a'>", + "Column<'a'>", + "Column<'A'>", + "Column<'st.x'>", + "Column<'abs(a)'>", + "Column<'round(b, 0)'>", + "Column<'avg(a)'>", + "Column<'upper(lower(s))'>", + "Column<'1.5'>", + "Column<'1.0E20'>", + "Column<'true'>", + "Column<'NULL'>", + "Column<'a AS x'>", + "Column<'CAST(a AS STRING)'>", + "Column<'a ASC NULLS LAST'>", + "Column<'struct(a, s)'>", + "Column<'negative(a)'>", + ] + assert str(F.abs("a")) == "Column<'abs(a)'>" + + +class TestSparkColumnMethods: + @pytest.mark.parametrize("method", ["between", "getItem", "getField", "substr", "eqNullSafe", "bitwiseAND"]) + def test_missing_methods_raise(self, df, method): + if USE_ACTUAL_SPARK: + assert callable(getattr(df.a, method)) + else: + from duckdb.experimental.spark.exception import ContributionsAcceptedError + + with pytest.raises(ContributionsAcceptedError, match=f"Column.{method} is not implemented"): + getattr(df.a, method) + with pytest.raises(AttributeError, match=f"Column.{method} is not implemented"): + getattr(df.a, method) + assert not hasattr(df.a, method) + assert getattr(df.a, method, None) is None + + def test_other_names_are_struct_fields(self, df): + assert df.select(df.st.x, F.col("st").y).collect()[0] == (1, "q") + + @pytest.mark.skipif(USE_ACTUAL_SPARK, reason="checks the module's list against the installed PySpark") + def test_method_list_covers_pyspark(self): + pyspark_sql = pytest.importorskip("pyspark.sql") + from duckdb.experimental.spark.sql.column import _SPARK_COLUMN_METHODS + + spark_methods = {n for n in dir(pyspark_sql.Column) if not n.startswith("_")} + assert spark_methods <= _SPARK_COLUMN_METHODS + + +@pytest.mark.skipif(USE_ACTUAL_SPARK, reason="checks every function the module implements") +def test_every_function_gives_a_spark_name(spark): + import contextlib + import inspect + + from duckdb.experimental.spark.sql import functions + from duckdb.experimental.spark.sql.column import Column + + generic = {int: 1, float: 0.5, bool: True, str: "a"} + calls = { + "call_function": ("abs", "a"), + "date_part": (F.lit("YEAR"), "d"), + "datepart": (F.lit("YEAR"), "d"), + "date_trunc": ("month", "d"), + "encode": ("s", "UTF-8"), + "extract": (F.lit("YEAR"), "d"), + "regexp_replace": ("s", "a", "b"), + "sha2": ("s", 256), + } + unnamed = [] + for name, function in list(vars(functions).items()): + if name.startswith("_") or not inspect.isfunction(function) or function.__module__ != functions.__name__: + continue + # transform raises; expr takes SQL text, whose names are DuckDB's; broadcast returns a DataFrame. + if name in {"transform", "expr", "broadcast"}: + continue + args = list(calls.get(name, ())) + if name not in calls: + for parameter in inspect.signature(function).parameters.values(): + if parameter.kind is parameter.VAR_POSITIONAL: + args += [F.col("a"), F.col("b")] + elif parameter.default is not parameter.empty: + continue + elif "Callable" in str(parameter.annotation): + args.append(lambda x: x) + elif "Column" in str(parameter.annotation): + args.append(F.col("a")) + else: + args.append(generic.get(parameter.annotation, 1)) + with pytest.warns(FutureWarning) if name == "approxCountDistinct" else contextlib.nullcontext(): + result = function(*args) + assert isinstance(result, Column), name + if result._name is None: + unnamed.append(name) + assert unnamed == [] diff --git a/tests/fast/spark/test_spark_create_dataframe.py b/tests/fast/spark/test_spark_create_dataframe.py new file mode 100644 index 000000000..c54a52630 --- /dev/null +++ b/tests/fast/spark/test_spark_create_dataframe.py @@ -0,0 +1,994 @@ +import datetime +import decimal +import math +import sys +import time +from typing import NamedTuple + +import pytest + +_ = pytest.importorskip("duckdb.experimental.spark") + +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import ParseException, PySparkTypeError, PySparkValueError, error_condition +from spark_namespace.sql import functions as sf +from spark_namespace.sql.types import ( + ArrayType, + BinaryType, + BooleanType, + ByteType, + DateType, + DayTimeIntervalType, + DecimalType, + DoubleType, + FloatType, + IntegerType, + LongType, + MapType, + NullType, + Row, + ShortType, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, +) + +import duckdb +from duckdb.experimental.spark.exception import ContributionsAcceptedError + +if USE_ACTUAL_SPARK: + import pyspark + + +def struct_of(*fields): + return StructType([StructField(name, data_type, True) for name, data_type in fields]) + + +def value_of(data_type): + return struct_of(("value", data_type)) + + +class TestSchemaString: + @pytest.mark.parametrize( + ("schema", "expected"), + [ + ("a int, b string", struct_of(("a", IntegerType()), ("b", StringType()))), + ("a: int, b: string", struct_of(("a", IntegerType()), ("b", StringType()))), + ("a INT, b STRING", struct_of(("a", IntegerType()), ("b", StringType()))), + ("A Int, B String", struct_of(("A", IntegerType()), ("B", StringType()))), + (" a int , b string ", struct_of(("a", IntegerType()), ("b", StringType()))), + ("a: byte, b: decimal( 16 , 8 ) ", struct_of(("a", ByteType()), ("b", DecimalType(16, 8)))), + ("struct", struct_of(("a", IntegerType()), ("b", StringType()))), + ("STRUCT", struct_of(("a", IntegerType()), ("b", StringType()))), + ("struct", struct_of(("a", IntegerType()), ("b", StringType()))), + ("`a b` int, `c``d` string", struct_of(("a b", IntegerType()), ("c`d", StringType()))), + ("a int, a string", struct_of(("a", IntegerType()), ("a", StringType()))), + ("string int", struct_of(("string", IntegerType()))), + ("interval: int, b: string", struct_of(("interval", IntegerType()), ("b", StringType()))), + ( + "a bigint, b double, c boolean, d date, e timestamp, f binary", + struct_of( + ("a", LongType()), + ("b", DoubleType()), + ("c", BooleanType()), + ("d", DateType()), + ("e", TimestampType()), + ("f", BinaryType()), + ), + ), + ("STRING", value_of(StringType())), + ("int", value_of(IntegerType())), + ("INT ", value_of(IntegerType())), + ("integer", value_of(IntegerType())), + ("bigint", value_of(LongType())), + ("long", value_of(LongType())), + ("tinyint", value_of(ByteType())), + ("byte", value_of(ByteType())), + ("smallint", value_of(ShortType())), + ("short", value_of(ShortType())), + ("float", value_of(FloatType())), + ("real", value_of(FloatType())), + ("double", value_of(DoubleType())), + ("boolean", value_of(BooleanType())), + ("date", value_of(DateType())), + ("timestamp", value_of(TimestampType())), + ("timestamp_ltz", value_of(TimestampType())), + ("timestamp_ntz", value_of(TimestampNTZType())), + ("binary", value_of(BinaryType())), + ("void", value_of(NullType())), + ("decimal", value_of(DecimalType(10, 0))), + ("decimal(10,2)", value_of(DecimalType(10, 2))), + ("DECIMAL( 16 , 8 )", value_of(DecimalType(16, 8))), + ("decimal(5)", value_of(DecimalType(5, 0))), + ("dec(5,1)", value_of(DecimalType(5, 1))), + ("numeric(7,3)", value_of(DecimalType(7, 3))), + ("interval day to second", value_of(DayTimeIntervalType(0, 3))), + ("array", value_of(ArrayType(IntegerType()))), + ("ARRAY", value_of(ArrayType(StringType()))), + ("array< short>", value_of(ArrayType(ShortType()))), + ("map", value_of(MapType(StringType(), IntegerType()))), + (" map ", value_of(MapType(StringType(), StringType()))), + ( + "a array>", + struct_of(("a", ArrayType(struct_of(("x", IntegerType()), ("y", StringType()))))), + ), + ( + "a map>, b struct>", + struct_of( + ("a", MapType(StringType(), ArrayType(IntegerType()))), + ("b", struct_of(("c", struct_of(("d", DecimalType(5, 2)))))), + ), + ), + ], + ) + def test_schema(self, spark, schema, expected): + df = spark.createDataFrame([], schema) + assert df.schema == expected + assert df.columns == expected.names + assert df.collect() == [] + + @pytest.mark.parametrize( + "schema", + [ + "blabla", + "bool", + "int a", + "a: int,", + "a int b int", + "a.b int", + "array", + "map", + "map>", + "struct", "'<'", "'>'"), + ("decimal(10", "'('", "end of input"), + ("a: int, b", "':'", "end of input"), + ("interval day to year", "'to'", "'to'"), + ("interval second to day", "'to'", "'to'"), + ("interval minute to hour", "'to'", "'to'"), + ("interval day to day", "'to'", "'to'"), + ], + ) + def test_syntax_error_position(self, spark, schema, spark_error, module_error): + # Spark reports the error of reading the text as a list of fields, even when it then reads it as a type; the + # module reads the text once, as a type or a list of fields by its first tokens. + with pytest.raises(ParseException) as raised: + spark.createDataFrame([], schema) + assert error_condition(raised.value) == "PARSE_SYNTAX_ERROR" + error = spark_error if USE_ACTUAL_SPARK else module_error + assert dict(raised.value.getMessageParameters()) == {"error": error, "hint": ""} + assert raised.value.getSqlState() == "42601" + + def test_unknown_type_name(self, spark): + with pytest.raises(ParseException) as raised: + spark.createDataFrame([], "a int, b strin") + assert error_condition(raised.value) == "UNSUPPORTED_DATATYPE" + assert dict(raised.value.getMessageParameters()) == {"typeName": '"STRIN"'} + assert raised.value.getSqlState() == "0A000" + + @pytest.mark.parametrize("schema", ["a strin b", "a strin, b", "a int, b strin c"]) + def test_unknown_type_name_raises_before_a_later_syntax_error(self, spark, schema): + with pytest.raises(ParseException) as raised: + spark.createDataFrame([], schema) + # Spark checks type names once the text parses; the module raises at the first one it does not know. + assert error_condition(raised.value) == ("PARSE_SYNTAX_ERROR" if USE_ACTUAL_SPARK else "UNSUPPORTED_DATATYPE") + + def test_char_type_raises_before_a_later_syntax_error(self, spark): + if USE_ACTUAL_SPARK: + with pytest.raises(ParseException): + spark.createDataFrame([], "a varchar(3) b") + else: + with pytest.raises(ContributionsAcceptedError, match="varchar"): + spark.createDataFrame([], "a varchar(3) b") + + @pytest.mark.parametrize( + ("name", "expected"), + [ + ("interval day", DayTimeIntervalType(0, 0)), + ("interval hour", DayTimeIntervalType(1, 1)), + ("interval hour to minute", DayTimeIntervalType(1, 2)), + ("interval minute to second", DayTimeIntervalType(2, 3)), + ("interval day to second", DayTimeIntervalType(0, 3)), + ], + ) + def test_interval_type_names(self, name, expected): + assert StructType().add("a", name).fields[0].dataType == expected + + def test_variant(self, spark): + if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3."): + with pytest.raises(ParseException): + spark.createDataFrame([], "a variant") + else: + assert spark.createDataFrame([], "a variant").dtypes == [("a", "variant")] + + @pytest.mark.parametrize("schema", ["s struct", "s struct"]) + def test_struct_fields_named_alike(self, spark, schema): + if USE_ACTUAL_SPARK: + row = spark.createDataFrame([((1, "x"),)], schema).collect()[0] + assert tuple(row.s) == (1, "x") + else: + from duckdb.experimental.spark.errors import PySparkNotImplementedError + + # DuckDB matches struct field names case insensitively, so a struct cannot hold both. + with pytest.raises(PySparkNotImplementedError): + spark.createDataFrame([((1, "x"),)], schema) + + @pytest.mark.parametrize( + ("schema", "spark_schema"), + [ + ("interval day", value_of(DayTimeIntervalType(0, 0))), + ("a interval hour to minute", struct_of(("a", DayTimeIntervalType(1, 2)))), + ("", StructType([])), + ("struct<>", StructType([])), + ], + ) + def test_schema_not_supported(self, spark, schema, spark_schema): + if USE_ACTUAL_SPARK: + assert spark.createDataFrame([], schema).schema == spark_schema + else: + with pytest.raises(ContributionsAcceptedError): + spark.createDataFrame([], schema) + + @pytest.mark.parametrize( + ("schema", "nullable", "metadata"), + [ + ("a int not null", False, {}), + ("a INT NOT NULL", False, {}), + ("a: int not null", False, {}), + ("struct", False, {}), + ("a int comment 'hello'", True, {"comment": "hello"}), + ('a int not null comment "it\'s"', False, {"comment": "it's"}), + ], + ) + def test_not_null_and_comment(self, spark, schema, nullable, metadata): + df = spark.createDataFrame([(1,)], schema) + assert df.collect() == [Row(a=1)] + field = df.schema["a"] + assert field.dataType == IntegerType() + if USE_ACTUAL_SPARK: + assert (field.nullable, field.metadata) == (nullable, metadata) + else: + # The module reports every field as nullable and keeps no metadata. + assert (field.nullable, field.metadata) == (True, {}) + if not nullable: + with pytest.raises(PySparkValueError): + spark.createDataFrame([(None,)], schema) + + def test_duckdb_only_types(self, spark): + schema = "a hugeint, b uuid, c ubigint, d utinyint" + if USE_ACTUAL_SPARK: + with pytest.raises(ParseException): + spark.createDataFrame([], schema) + else: + from duckdb.experimental.spark.sql.types import ( + HugeIntegerType, + UnsignedByteType, + UnsignedLongType, + UUIDType, + ) + + df = spark.createDataFrame([(2**100, None, 2**63, 255)], schema) + assert df.schema == struct_of( + ("a", HugeIntegerType()), ("b", UUIDType()), ("c", UnsignedLongType()), ("d", UnsignedByteType()) + ) + assert df.collect() == [Row(a=2**100, b=None, c=2**63, d=255)] + + @pytest.mark.parametrize("schema", ["a char(5)", "a varchar(5)"]) + def test_char_types_raise(self, spark, schema): + with pytest.raises(Exception): # noqa: B017, PT011 + spark.createDataFrame([("x",)], schema).collect() + + @pytest.mark.parametrize( + ("schema", "data", "expected"), + [ + ("a int, b string", [(1, "x"), (2, None)], [Row(a=1, b="x"), Row(a=2, b=None)]), + ("a: long, b: double", [(1, 1.5)], [Row(a=1, b=1.5)]), + ("a boolean, b float, c tinyint, d smallint", [(True, 1.5, 1, 2)], [Row(a=True, b=1.5, c=1, d=2)]), + ("struct", [(1, "a")], [Row(x=1, y="a")]), + ("a decimal(10,2)", [(decimal.Decimal("1.25"),)], [Row(a=decimal.Decimal("1.25"))]), + ( + "a date, b timestamp, c timestamp_ntz", + [(datetime.date(2024, 1, 2), datetime.datetime(2024, 1, 2, 3, 4, 5), datetime.datetime(2024, 1, 2))], + [ + Row( + a=datetime.date(2024, 1, 2), + b=datetime.datetime(2024, 1, 2, 3, 4, 5), + c=datetime.datetime(2024, 1, 2), + ) + ], + ), + ("string", ["a", None], [Row(value="a"), Row(value=None)]), + ("int", [1, 2], [Row(value=1), Row(value=2)]), + ("array", [[1, 2], None, []], [Row(value=[1, 2]), Row(value=None), Row(value=[])]), + ("map", [{"a": 1}], [Row(value={"a": 1})]), + ( + "a array, b map, c struct", + [([1, None], {"k": 2}, (3, "s"))], + [Row(a=[1, None], b={"k": 2}, c=Row(x=3, y="s"))], + ), + ("a int, b string", [{"b": "x", "a": 1}], [Row(a=1, b="x")]), + ("a int, b string", [Row(a=1, b="x")], [Row(a=1, b="x")]), + ], + ) + def test_schema_with_data(self, spark, schema, data, expected): + df = spark.createDataFrame(data, schema) + res = df.collect() + assert res == expected + assert [r.asDict(recursive=True) for r in res] == [r.asDict(recursive=True) for r in expected] + + def test_binary_with_schema(self, spark): + res = spark.createDataFrame([(b"ab",)], "a binary").collect() + assert bytes(res[0].a) == b"ab" + + def test_bigint_schema_does_not_overflow(self, spark): + df = spark.createDataFrame([2_000_000_000], "bigint") + assert df.select((sf.col("value") * 2).alias("v")).collect() == [Row(v=4_000_000_000)] + + @pytest.mark.parametrize( + ("schema", "data"), + [ + ("a int", [("x",)]), + ("a int", [(True,)]), + ("a double", [(1,)]), + ("a decimal(5,2)", [(1.5,)]), + ("a timestamp", [(datetime.date(2024, 1, 2),)]), + ("a array", [([1, "x"],)]), + ("a map", [([1],)]), + ("a struct", [(5,)]), + ("int", ["x"]), + ], + ) + def test_verify_type(self, spark, schema, data): + with pytest.raises(PySparkTypeError): + spark.createDataFrame(data, schema) + + @pytest.mark.parametrize( + ("schema", "data"), + [ + ("a tinyint", [(128,)]), + ("a int", [(2**31,)]), + ("a map", [({None: 1},)]), + ("a int not null, b int", [{"b": 1}]), + (StructType([StructField("a", IntegerType(), False)]), [(None,)]), + ], + ) + def test_verify_value(self, spark, schema, data): + with pytest.raises(PySparkValueError): + spark.createDataFrame(data, schema) + + @pytest.mark.parametrize( + ("schema", "data", "expected"), + [ + ("a string", [(1,)], [Row(a="1")]), + ("a date", [(datetime.datetime(2024, 1, 2, 3),)], [Row(a=datetime.date(2024, 1, 2))]), + ("a array", [((1, 2),)], [Row(a=[1, 2])]), + ("int", [None], [Row(value=None)]), + ], + ) + def test_verify_accepts(self, spark, schema, data, expected): + assert spark.createDataFrame(data, schema).collect() == expected + + def test_verify_schema_false(self, spark): + # Spark skips the check and stores what the JVM makes of the value; the module casts it. + spark.createDataFrame([(1,)], "a double", verifySchema=False).collect() + with pytest.raises(PySparkTypeError): + spark.createDataFrame([(1,)], "a double") + + def test_nan_in_nested_values(self, spark): + nan = float("nan") + df = spark.createDataFrame( + [([1.0, nan], {"k": nan}, (nan,), nan)], + "a array, b map, c struct, d double", + ) + row = df.collect()[0] + assert row.a[0] == 1.0 + assert math.isnan(row.a[1]) + assert math.isnan(row.b["k"]) + assert math.isnan(row.c.x) + assert math.isnan(row.d) + + def test_nan_in_inferred_nested_values(self, spark): + nan = float("nan") + row = spark.createDataFrame([([nan], {"k": nan}, Row(x=nan))]).collect()[0] + assert math.isnan(row._1[0]) + assert math.isnan(row._2["k"]) + assert math.isnan(row._3.x) + + def test_row_shorter_than_schema(self, spark): + with pytest.raises(PySparkValueError): + spark.createDataFrame([(1,)], "a int, b int").collect() + + def test_struct_type_built_with_add(self, spark): + schema = StructType().add("a", IntegerType()).add(StructField("b", StringType())) + df = spark.createDataFrame([(1, "x")], schema) + assert df.schema == struct_of(("a", IntegerType()), ("b", StringType())) + assert df.collect() == [Row(a=1, b="x")] + + def test_pandas_with_schema_string(self, spark): + pd = pytest.importorskip("pandas") + df = spark.createDataFrame(pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}), "a int, b string") + assert df.schema == struct_of(("a", IntegerType()), ("b", StringType())) + assert df.collect() == [Row(a=1, b="x"), Row(a=2, b="y")] + + +aware = datetime.datetime(2024, 1, 2, 3, 4, 5, tzinfo=datetime.timezone(datetime.timedelta(hours=2))) + + +class Point(NamedTuple): + x: int + y: str + + +class Person: + def __init__(self) -> None: + self.name = "Alice" + self.age = 1 + + +class TestInference: + @pytest.mark.parametrize( + ("value", "data_type", "collected"), + [ + (1, LongType(), 1), + (2_000_000_000, LongType(), 2_000_000_000), + (1.5, DoubleType(), 1.5), + (True, BooleanType(), True), + ("s", StringType(), "s"), + (decimal.Decimal("1.5"), DecimalType(38, 18), decimal.Decimal("1.500000000000000000")), + (datetime.date(2024, 1, 2), DateType(), datetime.date(2024, 1, 2)), + (datetime.datetime(2024, 1, 2, 3, 4, 5), TimestampType(), datetime.datetime(2024, 1, 2, 3, 4, 5)), + (aware, TimestampType(), aware.astimezone().replace(tzinfo=None)), + (datetime.timedelta(days=1, seconds=5), DayTimeIntervalType(0, 3), datetime.timedelta(days=1, seconds=5)), + ([1, 2], ArrayType(LongType()), [1, 2]), + ([1, None], ArrayType(LongType()), [1, None]), + ([[1], [2, 3]], ArrayType(ArrayType(LongType())), [[1], [2, 3]]), + ([1, "a"], ArrayType(StringType()), ["1", "a"]), + ({"a": 1}, MapType(StringType(), LongType()), {"a": 1}), + ({"a": {"b": [1]}}, MapType(StringType(), MapType(StringType(), ArrayType(LongType()))), {"a": {"b": [1]}}), + (Row(x=1, y="a"), struct_of(("x", LongType()), ("y", StringType())), Row(x=1, y="a")), + ((1, "a"), struct_of(("_1", LongType()), ("_2", StringType())), Row(_1=1, _2="a")), + ([Row(x=1), Row(x=2)], ArrayType(struct_of(("x", LongType()))), [Row(x=1), Row(x=2)]), + ], + ) + def test_infer_value(self, spark, value, data_type, collected): + df = spark.createDataFrame([(value,)]) + assert df.schema == struct_of(("_1", data_type)) + res = df.collect() + assert res == [Row(_1=collected)] + assert res[0].asDict(recursive=True) == Row(_1=collected).asDict(recursive=True) + + @pytest.mark.parametrize("value", [b"ab", bytearray(b"ab")]) + def test_infer_binary(self, spark, value): + df = spark.createDataFrame([(value,)]) + assert df.schema == struct_of(("_1", BinaryType())) + assert bytes(df.collect()[0]._1) == b"ab" + + def test_python_ints_are_bigint(self, spark): + df = spark.createDataFrame([(2_000_000_000,), (1,)], ["a"]) + assert df.schema == struct_of(("a", LongType())) + assert df.select((sf.col("a") * 2).alias("b")).collect() == [Row(b=4_000_000_000), Row(b=2)] + + @pytest.mark.parametrize( + ("rows", "data_type", "collected"), + [ + ([(None,), (1,)], LongType(), [None, 1]), + ([(1,), (None,)], LongType(), [1, None]), + ([([],), ([1],)], ArrayType(LongType()), [[], [1]]), + ([({},), ({"a": 1},)], MapType(StringType(), LongType()), [{}, {"a": 1}]), + ([(1,), ("a",)], StringType(), ["1", "a"]), + ([(True,), ("a",)], StringType(), ["true", "a"]), + ([(1.33,), ("a",)], StringType(), ["1.33", "a"]), + ], + ) + def test_infer_across_rows(self, spark, rows, data_type, collected): + df = spark.createDataFrame(rows) + assert df.schema == struct_of(("_1", data_type)) + assert [r._1 for r in df.collect()] == collected + + @pytest.mark.parametrize( + ("rows", "error", "error_class"), + [ + ([(None,), (None,)], PySparkValueError, "CANNOT_DETERMINE_TYPE"), + ([([None],)], PySparkValueError, "CANNOT_DETERMINE_TYPE"), + ([(1,), (1.5,)], PySparkTypeError, "CANNOT_MERGE_TYPE"), + ([(1.5,), (1,)], PySparkTypeError, "CANNOT_MERGE_TYPE"), + ([(1,), (True,)], PySparkTypeError, "CANNOT_MERGE_TYPE"), + ([(datetime.date(2024, 1, 2),), (datetime.datetime(2024, 1, 2),)], PySparkTypeError, "CANNOT_MERGE_TYPE"), + ([([1, 1.5],)], PySparkTypeError, "CANNOT_INFER_TYPE_FOR_FIELD"), + ([1, 2], PySparkTypeError, "CANNOT_INFER_SCHEMA_FOR_TYPE"), + (["a", "b"], PySparkTypeError, "CANNOT_INFER_SCHEMA_FOR_TYPE"), + ([], PySparkValueError, "CANNOT_INFER_EMPTY_SCHEMA"), + ], + ) + def test_infer_errors(self, spark, rows, error, error_class): + with pytest.raises(error, match=error_class): + spark.createDataFrame(rows) + + def test_infer_columns(self, spark): + df = spark.createDataFrame([(1, "a", 1.5, True)]) + assert df.schema == struct_of( + ("_1", LongType()), ("_2", StringType()), ("_3", DoubleType()), ("_4", BooleanType()) + ) + assert df.collect() == [Row(_1=1, _2="a", _3=1.5, _4=True)] + + @pytest.mark.parametrize( + ("rows", "columns", "expected"), + [ + ([Row(a=1, b="x"), Row(a=2, b="y")], ["a", "b"], [Row(a=1, b="x"), Row(a=2, b="y")]), + ([Point(1, "a")], ["x", "y"], [Row(x=1, y="a")]), + ([Person()], ["age", "name"], [Row(age=1, name="Alice")]), + ([[1, "a"], [2, "b"]], ["_1", "_2"], [Row(_1=1, _2="a"), Row(_1=2, _2="b")]), + ], + ) + def test_infer_row_kinds(self, spark, rows, columns, expected): + df = spark.createDataFrame(rows) + assert df.columns == columns + assert df.collect() == expected + + @pytest.mark.parametrize( + ("names", "rows", "columns", "expected"), + [ + (["x", "y"], [(1, "a")], ["x", "y"], [Row(x=1, y="a")]), + (["x"], [(1, "a")], ["x", "_2"], [Row(x=1, _2="a")]), + (["p", "q"], [{"b": 1, "a": "x"}], ["p", "q"], [Row(p="x", q=1)]), + (["p", "q"], [Row(a=1, b="x")], ["p", "q"], [Row(p=1, q="x")]), + ], + ) + def test_infer_with_names(self, spark, names, rows, columns, expected): + df = spark.createDataFrame(rows, names) + assert df.columns == columns + assert df.collect() == expected + + +class TestDicts: + def test_dicts(self, spark): + # Spark sorts the keys of each dict row, appends keys that later rows add, and gives NULL for a missing key. + data = [ + {"id": 1, "name": "Alice", "age": 25}, + {"id": 2, "age": 30, "name": "Bob"}, + {"age": 35, "id": 3, "name": "Charlie", "city": "New York"}, + ] + df = spark.createDataFrame(data) + assert df.columns == ["age", "id", "name", "city"] + assert df.schema == struct_of( + ("age", LongType()), ("id", LongType()), ("name", StringType()), ("city", StringType()) + ) + assert df.collect() == [ + Row(age=25, id=1, name="Alice", city=None), + Row(age=30, id=2, name="Bob", city=None), + Row(age=35, id=3, name="Charlie", city="New York"), + ] + + @pytest.mark.parametrize( + ("data", "columns", "expected"), + [ + ([{"b": 1}, {"a": 2, "b": 3}], ["b", "a"], [Row(b=1, a=None), Row(b=3, a=2)]), + ([{"a": 1, "b": 2}, {"a": 3}], ["a", "b"], [Row(a=1, b=2), Row(a=3, b=None)]), + ([{"a": None}, {"a": 1}], ["a"], [Row(a=None), Row(a=1)]), + ], + ) + def test_dicts_with_different_keys(self, spark, data, columns, expected): + df = spark.createDataFrame(data) + assert df.columns == columns + assert df.collect() == expected + + def test_dicts_with_nested_values(self, spark): + df = spark.createDataFrame([{"a": [1, 2], "m": {"k": "v"}, "r": Row(x=1)}]) + assert df.schema == struct_of( + ("a", ArrayType(LongType())), + ("m", MapType(StringType(), StringType())), + ("r", struct_of(("x", LongType()))), + ) + res = df.collect() + assert res == [Row(a=[1, 2], m={"k": "v"}, r=Row(x=1))] + assert res[0].r.x == 1 + + +class TestStruct: + @pytest.mark.parametrize( + "make_struct", + [ + lambda: sf.struct("x", "y"), + lambda: sf.struct(sf.col("x"), sf.col("y")), + lambda: sf.struct(["x", "y"]), + lambda: sf.struct([sf.col("x"), sf.col("y")]), + ], + ) + def test_struct(self, spark, make_struct): + df = spark.createDataFrame([(1, "a"), (2, "b")], ["x", "y"]) + res = df.select(make_struct().alias("s")).collect() + assert res == [Row(s=Row(x=1, y="a")), Row(s=Row(x=2, y="b"))] + assert res[0].s.y == "a" + + def test_struct_with_list_of_df_columns(self, spark): + df = spark.createDataFrame([("Alice", 2), ("Bob", 5)], ("name", "age")) + res = df.select(sf.struct([df.age, df.name]).alias("struct")).collect() + assert res == [Row(struct=Row(age=2, name="Alice")), Row(struct=Row(age=5, name="Bob"))] + assert res[0].struct.age == 2 + + def test_struct_with_alias(self, spark): + df = spark.createDataFrame([(1, "a")], ["x", "y"]) + res = df.select(sf.struct(sf.col("x"), sf.col("y").alias("z")).alias("s")).collect() + assert res[0].s.asDict() == {"x": 1, "z": "a"} + + def test_struct_with_tuple_raises(self, spark): + df = spark.createDataFrame([(1, "a")], ["x", "y"]) + with pytest.raises(PySparkTypeError): + df.select(sf.struct(("x", "y"))) + + +class TestReaderSchema: + def test_load_without_format_quotes_path(self, spark, tmp_path): + # Spark reads parquet when no format is given. + sub_dir = tmp_path / "dir with space's" + sub_dir.mkdir() + file_path = (sub_dir / "data.parquet").as_posix() + duckdb.sql("select 1 as a, 2 as b").write_parquet(file_path) + + assert spark.read.load(file_path).collect() == [Row(a=1, b=2)] + + def test_load_with_schema_string(self, spark, tmp_path): + csv_path = tmp_path / "data.csv" + csv_path.write_text("1,x\n2,y\n", newline="", encoding="utf-8") + file_path = csv_path.as_posix() + + df = spark.read.load(file_path, format="csv", schema="a int, b string") + assert df.schema == struct_of(("a", IntegerType()), ("b", StringType())) + assert df.collect() == [Row(a=1, b="x"), Row(a=2, b="y")] + + def test_csv_with_schema_string(self, spark, tmp_path): + csv_path = tmp_path / "data.csv" + csv_path.write_text("1,x\n2,y\n", newline="", encoding="utf-8") + file_path = csv_path.as_posix() + + df = spark.read.csv(file_path, schema="a bigint, b string") + assert df.schema == struct_of(("a", LongType()), ("b", StringType())) + assert df.collect() == [Row(a=1, b="x"), Row(a=2, b="y")] + + +class TestReaderSchemaMatching: + @pytest.fixture + def files(self, tmp_path): + duckdb.sql("select 1 as a, 2 as b").write_parquet((tmp_path / "ab.parquet").as_posix()) + duckdb.sql('select 1 as "Mixed", 2 as "order"').write_parquet((tmp_path / "mixed.parquet").as_posix()) + (tmp_path / "ab.json").write_text('{"a": 1, "b": 2}\n', newline="", encoding="utf-8") + (tmp_path / "mixed.json").write_text('{"Mixed": 1, "order": 2}\n', newline="", encoding="utf-8") + (tmp_path / "ab.csv").write_text("1,2\n", newline="", encoding="utf-8") + return tmp_path + + @pytest.mark.parametrize("fmt", ["parquet", "json"]) + @pytest.mark.parametrize( + ("schema", "row"), + [ + ("b int, a int", (2, 1)), + ("b int", (2,)), + ("a int, z int", (1, None)), + ], + ) + def test_by_name(self, spark, files, fmt, schema, row): + df = spark.read.load((files / f"ab.{fmt}").as_posix(), format=fmt, schema=schema) + assert [tuple(r) for r in df.collect()] == [row] + assert df.columns == [name.split()[0] for name in schema.split(", ")] + + @pytest.mark.parametrize(("fmt", "row"), [("parquet", (1, 2)), ("json", (None, 2))]) + def test_case_of_names(self, spark, files, fmt, row): + # Spark matches parquet columns ignoring case, JSON fields not. + df = spark.read.load((files / f"mixed.{fmt}").as_posix(), format=fmt, schema="mixed int, `order` int") + assert df.columns == ["mixed", "order"] + assert [tuple(r) for r in df.collect()] == [row] + + def test_without_a_format(self, spark, files): + df = spark.read.load((files / "ab.parquet").as_posix(), schema="b int, a int") + assert [tuple(r) for r in df.collect()] == [(2, 1)] + + @pytest.mark.parametrize( + ("schema", "row"), + [ + ("b int, a int", (1, 2)), + ("b int", (1,)), + ("b int, a int, c int", (1, 2, None)), + ("`order` int, a int", (1, 2)), + ], + ) + @pytest.mark.parametrize("read", ["load", "csv"]) + def test_csv_by_position(self, spark, files, read, schema, row): + path = (files / "ab.csv").as_posix() + if read == "load": + df = spark.read.load(path, format="csv", schema=schema) + else: + df = spark.read.csv(path, schema=schema) + assert [tuple(r) for r in df.collect()] == [row] + + def test_csv_load_reads_the_first_line_as_data(self, spark, tmp_path): + path = tmp_path / "head.csv" + path.write_text("id,name\n1,a\n", newline="", encoding="utf-8") + rows = spark.read.load(path.as_posix(), format="csv").collect() + assert sorted(tuple(map(str, r)) for r in rows) == [("1", "a"), ("id", "name")] + + +def _written_text(path): + if path.is_dir(): + return "".join(part.read_text(encoding="utf-8") for part in sorted(path.glob("*.csv"))) + return path.read_text(encoding="utf-8") + + +class TestCsvFlags: + @pytest.fixture + def head_csv(self, tmp_path): + path = tmp_path / "head.csv" + path.write_text("id,name\n1,a\n", newline="", encoding="utf-8") + return path.as_posix() + + @pytest.mark.parametrize("header", ["true", "TRUE", "True", True]) + def test_header_read(self, spark, head_csv, header): + df = spark.read.csv(head_csv, header=header) + assert df.columns == ["id", "name"] + assert df.count() == 1 + + @pytest.mark.parametrize("header", ["false", "False", False]) + def test_no_header_read(self, spark, head_csv, header): + df = spark.read.csv(head_csv, header=header) + assert df.columns != ["id", "name"] + assert df.count() == 2 + + @pytest.mark.parametrize("header", ["yes", 1]) + def test_header_that_is_not_a_flag(self, spark, head_csv, header): + with pytest.raises(Exception, match="header flag can be true or false"): + spark.read.csv(head_csv, header=header).collect() + + @pytest.mark.parametrize( + "option", ["inferSchema", "multiLine", "ignoreLeadingWhiteSpace", "ignoreTrailingWhiteSpace"] + ) + def test_false_flags_read(self, spark, head_csv, option): + assert spark.read.csv(head_csv, header=True, **{option: "FALSE"}).count() == 1 + + @pytest.mark.parametrize(("header", "written"), [("TRUE", True), ("true", True), (True, True), ("false", False)]) + def test_header_write(self, spark, tmp_path, header, written): + path = tmp_path / "out.csv" + spark.createDataFrame([(1, "a")], ["id", "name"]).write.csv(path.as_posix(), header=header) + assert ("id,name" in _written_text(path)) == written + + @pytest.mark.parametrize(("quote_all", "quoted"), [("TRUE", True), (True, True), ("false", False)]) + def test_quote_all_write(self, spark, tmp_path, quote_all, quoted): + path = tmp_path / "out.csv" + spark.createDataFrame([(1, "a")], ["id", "name"]).write.csv(path.as_posix(), quoteAll=quote_all) + assert ('"1","a"' in _written_text(path)) == quoted + + @pytest.mark.parametrize( + ("column", "options", "written"), + [ + ("d", {"dateFormat": "yyyy/MM/dd"}, "2024/02/03\n"), + ("d", {"dateFormat": "dd MMM yy"}, "03 Feb 24\n"), + ("d", {"dateFormat": "EEEE 'of' MMMM"}, "Saturday of February\n"), + ("t", {"timestampFormat": "yyyy/MM/dd HH:mm:ss"}, "2024/02/03 04:05:06\n"), + ("t", {"timestampFormat": "yyyy-MM-dd hh:mm:ss.SSS a"}, "2024-02-03 04:05:06.000 AM\n"), + ], + ) + def test_datetime_patterns_write(self, spark, tmp_path, column, options, written): + path = tmp_path / "out.csv" + df = spark.createDataFrame( + [(datetime.date(2024, 2, 3), datetime.datetime(2024, 2, 3, 4, 5, 6))], "d date, t timestamp" + ) + df.select(column).write.csv(path.as_posix(), **options) + assert _written_text(path) == written + + def test_datetime_pattern_without_an_exact_format(self, spark, tmp_path): + path = tmp_path / "out.csv" + df = spark.createDataFrame([(datetime.date(2024, 2, 3),)], "d date") + if USE_ACTUAL_SPARK: + df.write.csv(path.as_posix(), dateFormat="yyyy/M/d") + assert _written_text(path) == "2024/2/3\n" + else: + with pytest.raises(ContributionsAcceptedError): + df.write.csv(path.as_posix(), dateFormat="yyyy/M/d") + + def test_quotes_write(self, spark, tmp_path): + df = spark.createDataFrame([('a"b',), ("c,d",), ("e",)], ["s"]) + path = tmp_path / "default.csv" + df.write.csv(path.as_posix()) + # Spark escapes a quote with a backslash, DuckDB doubles it; each reads its own text back. + escaped = '"a\\"b"' if USE_ACTUAL_SPARK else '"a""b"' + assert _written_text(path) == f'{escaped}\n"c,d"\ne\n' + assert spark.read.csv(path.as_posix()).collect() == df.collect() + path = tmp_path / "unescaped.csv" + if USE_ACTUAL_SPARK: + df.write.csv(path.as_posix(), escapeQuotes=False) + assert _written_text(path) == 'a"b\n"c,d"\ne\n' + else: + # DuckDB quotes every value with a quote character in it. + with pytest.raises(ContributionsAcceptedError): + df.write.csv(path.as_posix(), escapeQuotes=False) + + +class TestTimestampRoundTrip: + @pytest.mark.parametrize( + "value", + [ + pytest.param( + datetime.datetime(9999, 12, 31, 12), + marks=pytest.mark.skipif( + sys.platform == "win32", + reason="outside the range of the Windows C runtime, as for PySpark on Windows", + ), + ), + datetime.datetime(1970, 1, 1), + datetime.datetime(2024, 3, 10, 2, 30), + datetime.datetime(2024, 11, 3, 1, 30), + datetime.datetime(2024, 7, 1, 12, 0, 0, 123456), + ], + ) + def test_naive_timestamp(self, spark, value): + # Spark reads a naive datetime as local time through mktime, which also settles gaps and overlaps. + expected = datetime.datetime.fromtimestamp(time.mktime(value.timetuple())).replace( + microsecond=value.microsecond + ) + (row,) = spark.createDataFrame([(value,)], "t timestamp").collect() + assert row.t == expected + assert row.t.fold == expected.fold + + @pytest.mark.skipif( + sys.platform == "win32", reason="outside the range of the Windows C runtime, as for PySpark on Windows" + ) + def test_aware_timestamp_near_the_end_of_the_calendar(self, spark): + value = datetime.datetime(9999, 12, 31, 12, tzinfo=datetime.timezone.utc) + (row,) = spark.createDataFrame([(value,)], "t timestamp").collect() + assert row.t == value.astimezone().replace(tzinfo=None) + + @pytest.mark.skipif( + sys.platform == "win32", reason="outside the range of the Windows C runtime, as for PySpark on Windows" + ) + def test_latest_instant(self, spark): + # The last microsecond of the calendar is a time like any other, not DuckDB's infinity. + value = datetime.datetime(9999, 12, 31, 23, 59, 59, 999999, tzinfo=datetime.timezone.utc) + try: + expected = value.astimezone().replace(tzinfo=None) + except OverflowError: + pytest.skip("the local time of the instant is past the end of the calendar") + (row,) = spark.createDataFrame([(value,)], "t timestamp").collect() + assert row.t == expected + + def test_ntz_timestamp_near_the_end_of_the_calendar(self, spark): + value = datetime.datetime(9999, 12, 31, 12) + assert spark.createDataFrame([(value,)], "t timestamp_ntz").collect() == [Row(t=value)] + + def test_infinite_timestamps(self, spark): + if USE_ACTUAL_SPARK: + pytest.skip("Spark has no infinite timestamps") + # DuckDB returns infinite timestamps as the ends of Python's calendar. + df = spark.sql("select 'infinity'::timestamptz a, '-infinity'::timestamptz b, 'infinity'::timestamp c") + assert df.collect() == [Row(a=datetime.datetime.max, b=datetime.datetime.min, c=datetime.datetime.max)] + + +# File text, Spark's inferred types, and where DuckDB's type detection gives another type, the module's. +_CSV_INFERENCE_CASES = [ + ("ints", "a,b\n1,2147483648\n-3,4\n", ["int", "bigint"], ["bigint", "bigint"]), + ("doubles", "a,b\n1.5,1e3\n2,-0.25\n", ["double", "double"], None), + ( + "dates", + "a,b\n2024-01-02,2024-01-02 03:04:05\n2023-12-31,2023-12-31 23:59:59.5\n", + ["date", "timestamp"], + None, + ), + ("booleans", "a,b\ntrue,True\nfalse,FALSE\n", ["boolean", "boolean"], None), + ("empty_fields", "a,b,c\n1,,x\n,2.5,\n", ["int", "double", "string"], ["bigint", "double", "string"]), + ("mixed", "a,b\n1,x\n2.5,7\n", ["double", "string"], None), + ("wide_integer", "a\n12345678901234567890\n1\n", ["decimal(20,0)"], ["hugeint"]), + ("all_empty", "a,b\n,\n,\n", ["string", "string"], None), +] + + +class TestCsvInference: + @pytest.mark.parametrize(("name", "text", "spark_types", "module_types"), _CSV_INFERENCE_CASES) + def test_strings_without_inference(self, spark, tmp_path, name, text, spark_types, module_types): + path = tmp_path / f"{name}.csv" + path.write_text(text, newline="", encoding="utf-8") + df = spark.read.csv(path.as_posix(), header=True) + lines = [line.split(",") for line in text.splitlines()[1:]] + assert [t for _, t in df.dtypes] == ["string"] * len(lines[0]) + assert [tuple(r) for r in df.collect()] == [tuple(v or None for v in line) for line in lines] + + @pytest.mark.parametrize(("name", "text", "spark_types", "module_types"), _CSV_INFERENCE_CASES) + def test_inferred_types(self, spark, tmp_path, name, text, spark_types, module_types): + path = tmp_path / f"{name}.csv" + path.write_text(text, newline="", encoding="utf-8") + df = spark.read.csv(path.as_posix(), header=True, inferSchema=True) + expected = module_types if module_types is not None and not USE_ACTUAL_SPARK else spark_types + assert [t for _, t in df.dtypes] == expected + + def test_inferred_values(self, spark, tmp_path): + path = tmp_path / "values.csv" + path.write_text( + "i,d,day,ts,b,s\n1,1.5,2024-01-02,2024-01-02 03:04:05,true,x\n,,,,,\n", newline="", encoding="utf-8" + ) + rows = spark.read.csv(path.as_posix(), header=True, inferSchema=True).collect() + assert [tuple(r) for r in rows] == [ + (1, 1.5, datetime.date(2024, 1, 2), datetime.datetime(2024, 1, 2, 3, 4, 5), True, "x"), + (None,) * 6, + ] + + def test_inferred_timestamps_are_local_times(self, spark, tmp_path): + path = tmp_path / "times.csv" + path.write_text("ts\n2024-01-02 03:04:05\n2024-07-01 12:00:00\n", newline="", encoding="utf-8") + df = spark.read.csv(path.as_posix(), header=True, inferSchema=True) + assert df.schema == struct_of(("ts", TimestampType())) + winter, summer = datetime.datetime(2024, 1, 2, 3, 4, 5), datetime.datetime(2024, 7, 1, 12) + assert [r.ts for r in df.collect()] == [winter, summer] + assert df.filter(sf.col("ts") == sf.lit(summer)).count() == 1 + local = spark.createDataFrame([(winter,), (summer,)], "ts timestamp") + assert df.join(local, "ts").count() == 2 + + +class TestNestedTimestamps: + def test_collect(self, spark): + value = datetime.datetime(2024, 3, 10, 2, 30) + local = datetime.datetime.fromtimestamp(time.mktime(value.timetuple())) + schema = ( + "a array, aa array>, s struct, " + "m map, k map, sa array>" + ) + df = spark.createDataFrame( + [([value, None], [[value]], (value, 1), {"x": value}, {value: 1}, [(value,)]), (None,) * 6], schema + ) + assert df.collect() == [ + Row(a=[local, None], aa=[[local]], s=Row(t=local, n=1), m={"x": local}, k={local: 1}, sa=[Row(t=local)]), + Row(a=None, aa=None, s=None, m=None, k=None, sa=None), + ] + assert df.first() == df.collect()[0] + assert df.take(1) == df.collect()[:1] + + def test_show(self, spark, capsys): + value = datetime.datetime(2024, 1, 2, 3, 4, 5) + df = spark.createDataFrame( + [(value, [value], (value,))], "t timestamp, a array, s struct" + ) + df.show(truncate=False) + assert capsys.readouterr().out == ( + "+-------------------+---------------------+---------------------+\n" + "|t |a |s |\n" + "+-------------------+---------------------+---------------------+\n" + "|2024-01-02 03:04:05|[2024-01-02 03:04:05]|{2024-01-02 03:04:05}|\n" + "+-------------------+---------------------+---------------------+\n\n" + ) + + +class TestMapsWithNestedKeys: + def test_struct_keys(self, spark): + df = spark.createDataFrame([({Row(a=1): 2, Row(a=None): None},), ({},), (None,)], "m map, int>") + assert df.collect() == [Row(m={Row(a=1): 2, Row(a=None): None}), Row(m={}), Row(m=None)] + + def test_struct_keys_nested(self, spark): + schema = "m map, map, struct>>, s struct, int>>" + df = spark.createDataFrame([({Row(a=1): {Row(b=2): Row(c=3)}}, Row(n={Row(a=4): 5}))], schema) + assert df.collect() == [Row(m={Row(a=1): {Row(b=2): Row(c=3)}}, s=Row(n={Row(a=4): 5}))] + + def test_array_keys_raise(self, spark): + # Python cannot hash a list, so Spark too fails to build the dict. + with pytest.raises(TypeError, match="unhashable"): + spark.createDataFrame([({(1, 2): 3},)], "m map, int>").collect() + if not USE_ACTUAL_SPARK: + with pytest.raises(TypeError, match="unhashable"): + spark.sql("select map([[1, 2]], [3]) m").collect() + + def test_struct_keys_from_sql(self, spark): + if USE_ACTUAL_SPARK: + pytest.skip("DuckDB SQL") + assert spark.sql("select map([{'a': 1}], [2]) m").collect() == [Row(m={Row(a=1): 2})] diff --git a/tests/fast/spark/test_spark_dataframe.py b/tests/fast/spark/test_spark_dataframe.py index e242092e5..ed0a00c8c 100644 --- a/tests/fast/spark/test_spark_dataframe.py +++ b/tests/fast/spark/test_spark_dataframe.py @@ -4,20 +4,79 @@ from spark_namespace import USE_ACTUAL_SPARK -from spark_namespace.errors import PySparkTypeError, PySparkValueError +from spark_namespace.errors import ( + AnalysisException, + ParseException, + PySparkTypeError, + PySparkValueError, + error_condition, +) from spark_namespace.sql.column import Column -from spark_namespace.sql.functions import col, struct, when +from spark_namespace.sql.functions import col, lit, struct, when from spark_namespace.sql.types import ( ArrayType, + BinaryType, BooleanType, + ByteType, + DateType, + DayTimeIntervalType, + DecimalType, + DoubleType, + FloatType, IntegerType, LongType, MapType, + NullType, Row, + ShortType, StringType, StructField, StructType, + TimestampNTZType, + TimestampType, +) + +requires_tree_string = pytest.mark.skipif( + not hasattr(StructType, "treeString"), reason="StructType.treeString is only introduced in PySpark 4.0.0" +) + +PERSON_SCHEMA = StructType( + [ + StructField("id", IntegerType(), True), + StructField( + "person", + StructType([StructField("name", StringType(), True), StructField("age", IntegerType(), True)]), + True, + ), + StructField("hobbies", ArrayType(StringType()), True), + ] ) +PERSON_TREE = ( + "root\n" + " |-- id: integer (nullable = true)\n" + " |-- person: struct (nullable = true)\n" + " | |-- name: string (nullable = true)\n" + " | |-- age: integer (nullable = true)\n" + " |-- hobbies: array (nullable = true)\n" + " | |-- element: string (containsNull = true)\n" +) +FLAT_TYPES = [ + ("b", ByteType()), + ("s", ShortType()), + ("i", IntegerType()), + ("l", LongType()), + ("f", FloatType()), + ("d", DoubleType()), + ("m", DecimalType(10, 2)), + ("str", StringType()), + ("bool", BooleanType()), + ("date", DateType()), + ("ts", TimestampType()), + ("ntz", TimestampNTZType()), + ("bin", BinaryType()), + ("iv", DayTimeIntervalType()), + ("null", NullType()), +] def assert_column_objects_equal(col1: Column, col2: Column): @@ -27,6 +86,13 @@ def assert_column_objects_equal(col1: Column, col2: Column): class TestDataFrame: + def test_row_is_exported_from_sql_module(self): + if USE_ACTUAL_SPARK: + from pyspark.sql import Row as SqlRow + else: + from duckdb.experimental.spark.sql import Row as SqlRow + assert SqlRow is Row + def test_dataframe_from_list_of_tuples(self, spark): # Valid address = [(1, "14851 Jeffrey Rd", "DE"), (2, "43421 Margarita St", "NY"), (3, "13111 Siemon Ave", "CA")] @@ -43,15 +109,20 @@ def test_dataframe_from_list_of_tuples(self, spark): if USE_ACTUAL_SPARK: from py4j.protocol import Py4JJavaError - - with pytest.raises(Py4JJavaError): - spark.createDataFrame(address, ["id", "address", "state"]) + from spark_namespace.errors import IllegalArgumentException + + # Spark only checks the row length when the rows are collected. + ragged = spark.createDataFrame(address, ["id", "address", "state"]) + with pytest.raises( + (IllegalArgumentException, Py4JJavaError), match="3 fields are required while 2 values are provided" + ): + ragged.collect() else: - with pytest.raises(PySparkTypeError, match="LENGTH_SHOULD_BE_THE_SAME"): + with pytest.raises(PySparkValueError, match="LENGTH_SHOULD_BE_THE_SAME"): spark.createDataFrame(address, ["id", "address", "state"]) - # Dataframe instead of list - with pytest.raises(PySparkTypeError, match="SHOULD_NOT_DATAFRAME"): + # Dataframe instead of list; Spark 4 renamed the error class from SHOULD_NOT_DATAFRAME to INVALID_TYPE. + with pytest.raises(PySparkTypeError, match="should not be a DataFrame"): spark.createDataFrame(df, ["id", "address", "state"]) # Not a list @@ -59,73 +130,53 @@ def test_dataframe_from_list_of_tuples(self, spark): spark.createDataFrame(5, ["id", "address", "test"]) # Empty list - if not USE_ACTUAL_SPARK: - # TODO: Spark raises PySparkValueError [CANNOT_INFER_EMPTY_SCHEMA] # noqa: TD002, TD003 - df = spark.createDataFrame([], ["id", "address", "test"]) - res = df.collect() - assert res == [] + with pytest.raises(PySparkValueError, match="CANNOT_INFER_EMPTY_SCHEMA"): + spark.createDataFrame([], ["id", "address", "test"]) - # Duplicate column names + # Duplicate column names: inference merges the fields by name, so both become strings address = [(1, "14851 Jeffrey Rd", "DE"), (2, "43421 Margarita St", "NY"), (3, "13111 Siemon Ave", "DE")] df = spark.createDataFrame(address, ["id", "address", "id"]) - res = df.collect() - exptected_res_str = ( - "[Row(id=1, address='14851 Jeffrey Rd', id='DE'), Row(id=2, address='43421 " - "Margarita St', id='NY'), Row(id=3, address='13111 Siemon Ave', id='DE')]" - ) - if USE_ACTUAL_SPARK: - # Spark uses string for both ID columns. DuckDB correctly infers the types. - exptected_res_str = ( - exptected_res_str.replace("id=1", "id='1'").replace("id=2", "id='2'").replace("id=3", "id='3'") - ) - assert str(res) == exptected_res_str - - # Not enough column names - if not USE_ACTUAL_SPARK: - # TODO: Spark does not raise this error # noqa: TD002, TD003 - with pytest.raises(PySparkValueError, match="number of columns in the DataFrame don't match"): - spark.createDataFrame(address, ["id", "address"]) + assert df.columns == ["id", "address", "id"] + assert df.collect() == [ + ("1", "14851 Jeffrey Rd", "DE"), + ("2", "43421 Margarita St", "NY"), + ("3", "13111 Siemon Ave", "DE"), + ] + + # Not enough column names: the remaining columns get default names + df = spark.createDataFrame(address, ["id", "address"]) + assert df.columns == ["id", "address", "_3"] # Empty column names list - # Columns are filled in with default names - # TODO: check against Spark behavior # noqa: TD002, TD003 df = spark.createDataFrame(address, []) - res = df.collect() - assert res == [ - Row(col0=1, col1="14851 Jeffrey Rd", col2="DE"), - Row(col0=2, col1="43421 Margarita St", col2="NY"), - Row(col0=3, col1="13111 Siemon Ave", col2="DE"), + assert df.columns == ["_1", "_2", "_3"] + assert df.collect() == [ + Row(_1=1, _2="14851 Jeffrey Rd", _3="DE"), + Row(_1=2, _2="43421 Margarita St", _3="NY"), + Row(_1=3, _2="13111 Siemon Ave", _3="DE"), ] # Too many column names - if not USE_ACTUAL_SPARK: - # In Spark, this leads to an IndexError - with pytest.raises(PySparkValueError, match="number of columns in the DataFrame don't match"): - spark.createDataFrame(address, ["id", "address", "one", "two", "three"]) - - # Column names is not a list (but is iterable) - if not USE_ACTUAL_SPARK: - # These things do not work in Spark or throw different errors - df = spark.createDataFrame(address, {"a": 5, "b": 6, "c": 42}) - res = df.collect() - assert res == [ - Row(a=1, b="14851 Jeffrey Rd", c="DE"), - Row(a=2, b="43421 Margarita St", c="NY"), - Row(a=3, b="13111 Siemon Ave", c="DE"), - ] + with pytest.raises(IndexError): + spark.createDataFrame(address, ["id", "address", "one", "two", "three"]) - # Column names is not a list (string, becomes a single column name) - with pytest.raises(PySparkValueError, match="number of columns in the DataFrame don't match"): - spark.createDataFrame(address, "a") + # Column names is not a list + with pytest.raises(PySparkTypeError): + spark.createDataFrame(address, {"a": 5, "b": 6, "c": 42}) - with pytest.raises(TypeError, match="must be an iterable, not int"): - spark.createDataFrame(address, 5) + with pytest.raises(PySparkTypeError): + spark.createDataFrame(address, 5) + + # A string is a schema, and "a" is not a type + with pytest.raises(ParseException): + spark.createDataFrame(address, "a") def test_dataframe(self, spark): # Create DataFrame df = spark.createDataFrame([("Scala", 25000), ("Spark", 35000), ("PHP", 21000)]) + assert df.columns == ["_1", "_2"] res = df.collect() - assert res == [Row(col0="Scala", col1=25000), Row(col0="Spark", col1=35000), Row(col0="PHP", col1=21000)] + assert res == [Row(_1="Scala", _2=25000), Row(_1="Spark", _2=35000), Row(_1="PHP", _2=21000)] @pytest.mark.skipif(USE_ACTUAL_SPARK, reason="We can't create tables with our Spark test setup") def test_writing_to_table(self, spark): @@ -252,30 +303,15 @@ def test_df_nested_struct(self, spark): df2 = spark.createDataFrame(data=structureData, schema=structureSchema) res = df2.collect() - expected_res = [ - Row( - name={"firstname": "James", "middlename": "", "lastname": "Smith"}, id="36636", gender="M", salary=3100 - ), - Row( - name={"firstname": "Michael", "middlename": "Rose", "lastname": ""}, id="40288", gender="M", salary=4300 - ), - Row( - name={"firstname": "Robert", "middlename": "", "lastname": "Williams"}, - id="42114", - gender="M", - salary=1400, - ), - Row( - name={"firstname": "Maria", "middlename": "Anne", "lastname": "Jones"}, - id="39192", - gender="F", - salary=5500, - ), - Row(name={"firstname": "Jen", "middlename": "Mary", "lastname": "Brown"}, id="", gender="F", salary=-1), + name = Row("firstname", "middlename", "lastname") + assert res == [ + Row(name=name("James", "", "Smith"), id="36636", gender="M", salary=3100), + Row(name=name("Michael", "Rose", ""), id="40288", gender="M", salary=4300), + Row(name=name("Robert", "", "Williams"), id="42114", gender="M", salary=1400), + Row(name=name("Maria", "Anne", "Jones"), id="39192", gender="F", salary=5500), + Row(name=name("Jen", "Mary", "Brown"), id="", gender="F", salary=-1), ] - if USE_ACTUAL_SPARK: - expected_res = [Row(name=Row(**r.name), id=r.id, gender=r.gender, salary=r.salary) for r in expected_res] - assert res == expected_res + assert res[0].name.firstname == "James" schema = df2.schema assert schema == StructType( [ @@ -420,13 +456,35 @@ def test_drop(self, spark): assert df.drop("two", col("three")).columns == expected assert df.drop("two", col("three"), col("missing")).columns == expected + def test_drop_ignores_case(self, spark): + df = spark.createDataFrame([(1, 2)], ["a", "b"]) + assert df.drop("A").columns == ["b"] + assert df.drop(col("A")).columns == ["b"] + assert df.drop(df["A"]).columns == ["b"] + assert df.drop(df.a).collect() == [Row(b=2)] + + def test_drop_ignores_what_is_not_a_column_reference(self, spark): + df = spark.createDataFrame([(1, 2)], ["a", "b"]) + other = spark.createDataFrame([(1, 3)], ["a", "c"]) + assert df.drop(df.a + 1).columns == ["a", "b"] + assert df.drop(df.a.alias("x")).columns == ["a", "b"] + assert df.drop(other.c).columns == ["a", "b"] + assert df.drop(other.a).columns == ["a", "b"] + + def test_drop_every_column(self, spark): + df = spark.createDataFrame([(1, 2)], ["a", "b"]) + assert df.drop("a", "b").columns == [] + assert df.drop("a", "b").schema == StructType([]) + def test_cache(self, spark): data = [(1, 2, 3, 4)] df = spark.createDataFrame(data, ["one", "two", "three", "four"]) + one = df.one cached = df.cache() - assert df is not cached - assert cached.collect() == df.collect() + assert cached is df assert cached.collect() == [Row(one=1, two=2, three=3, four=4)] + assert df.select(one).collect() == [Row(one=1)] + assert df.filter(df.two == 2).select(df.three).collect() == [Row(three=3)] def test_dtypes(self, spark): data = [("Alice", 25, 5000.0), ("Bob", 30, 6000.0)] @@ -471,52 +529,88 @@ def test_dtypes_complex_types(self, spark): def test_printSchema(self, spark, capsys): data = [("Alice", 25, 5000), ("Bob", 30, 6000)] - df = spark.createDataFrame(data, ["name", "age", "salary"]) + schema = StructType( + [StructField("name", StringType()), StructField("age", LongType()), StructField("salary", LongType())] + ) + df = spark.createDataFrame(data, schema) df.printSchema() - captured = capsys.readouterr() - output = captured.out - - assert "root" in output - assert "name" in output - assert "age" in output - assert "salary" in output - assert "string" in output or "varchar" in output.lower() - assert "int" in output.lower() or "bigint" in output.lower() + assert capsys.readouterr().out == ( + "root\n" + " |-- name: string (nullable = true)\n" + " |-- age: long (nullable = true)\n" + " |-- salary: long (nullable = true)\n" + "\n" + ) def test_printSchema_nested(self, spark, capsys): - from spark_namespace.sql.types import ArrayType, IntegerType, StringType, StructField, StructType - - schema = StructType( + df = spark.createDataFrame( [ - StructField("id", IntegerType(), True), - StructField( - "person", - StructType([StructField("name", StringType(), True), StructField("age", IntegerType(), True)]), - True, - ), - StructField("hobbies", ArrayType(StringType()), True), - ] + (1, {"name": "Alice", "age": 25}, ["reading", "coding"]), + (2, {"name": "Bob", "age": 30}, ["gaming", "music"]), + ], + PERSON_SCHEMA, ) - data = [ - (1, {"name": "Alice", "age": 25}, ["reading", "coding"]), - (2, {"name": "Bob", "age": 30}, ["gaming", "music"]), - ] - df = spark.createDataFrame(data, schema) df.printSchema() - captured = capsys.readouterr() - output = captured.out - - assert "root" in output - assert "person" in output - assert "hobbies" in output - - def test_printSchema_negative_level(self, spark): - data = [("Alice", 25)] - df = spark.createDataFrame(data, ["name", "age"]) - - with pytest.raises(PySparkValueError): - df.printSchema(level=-1) + assert capsys.readouterr().out == PERSON_TREE + "\n" + df.printSchema(1) + assert capsys.readouterr().out == ( + "root\n" + " |-- id: integer (nullable = true)\n" + " |-- person: struct (nullable = true)\n" + " |-- hobbies: array (nullable = true)\n" + "\n" + ) + df.printSchema(level=2) + assert capsys.readouterr().out == PERSON_TREE + "\n" + + @pytest.mark.parametrize("level", [0, -1]) + def test_printSchema_negative_level(self, spark, capsys, level): + df = spark.createDataFrame([(1, {"name": "Alice", "age": 25}, ["reading"])], PERSON_SCHEMA) + df.printSchema(level=level) + assert capsys.readouterr().out == PERSON_TREE + "\n" + + def test_printSchema_type_names(self, spark, capsys): + schema = StructType([StructField(name, data_type) for name, data_type in FLAT_TYPES]) + df = spark.createDataFrame([], schema) + df.printSchema() + assert capsys.readouterr().out == ( + "root\n" + " |-- b: byte (nullable = true)\n" + " |-- s: short (nullable = true)\n" + " |-- i: integer (nullable = true)\n" + " |-- l: long (nullable = true)\n" + " |-- f: float (nullable = true)\n" + " |-- d: double (nullable = true)\n" + " |-- m: decimal(10,2) (nullable = true)\n" + " |-- str: string (nullable = true)\n" + " |-- bool: boolean (nullable = true)\n" + " |-- date: date (nullable = true)\n" + " |-- ts: timestamp (nullable = true)\n" + " |-- ntz: timestamp_ntz (nullable = true)\n" + " |-- bin: binary (nullable = true)\n" + " |-- iv: interval day to second (nullable = true)\n" + " |-- null: void (nullable = true)\n" + "\n" + ) + assert df.dtypes == [ + ("b", "tinyint"), + ("s", "smallint"), + ("i", "int"), + ("l", "bigint"), + ("f", "float"), + ("d", "double"), + ("m", "decimal(10,2)"), + ("str", "string"), + ("bool", "boolean"), + ("date", "date"), + ("ts", "timestamp"), + ("ntz", "timestamp_ntz"), + ("bin", "binary"), + ("iv", "interval day to second"), + ("null", "void"), + ] + @requires_tree_string def test_treeString_basic(self, spark): data = [("Alice", 25, 5000)] df = spark.createDataFrame(data, ["name", "age", "salary"]) @@ -529,6 +623,7 @@ def test_treeString_basic(self, spark): assert "(nullable = true)" in tree assert tree.count(" |-- ") == 3 + @requires_tree_string def test_treeString_nested_struct(self, spark): from spark_namespace.sql.types import IntegerType, StringType, StructField, StructType @@ -552,6 +647,7 @@ def test_treeString_nested_struct(self, spark): assert "name:" in tree assert "age:" in tree + @requires_tree_string def test_treeString_with_level(self, spark): from spark_namespace.sql.types import IntegerType, StringType, StructField, StructType @@ -574,14 +670,23 @@ def test_treeString_with_level(self, spark): data = [(1, {"name": "Alice", "details": {"address": "123 Main St"}})] df = spark.createDataFrame(data, schema) - # Level 1 should only show top-level fields - tree_level_1 = df.schema.treeString(level=1) - assert " |-- id:" in tree_level_1 - assert " |-- person: struct" in tree_level_1 - # Should not show nested field names at level 1 - lines = tree_level_1.split("\n") - assert len([line for line in lines if line.strip()]) <= 3 + assert df.schema.treeString(maxDepth=1) == ( + "root\n |-- id: integer (nullable = true)\n |-- person: struct (nullable = true)\n" + ) + assert df.schema.treeString(2) == ( + "root\n" + " |-- id: integer (nullable = true)\n" + " |-- person: struct (nullable = true)\n" + " | |-- name: string (nullable = true)\n" + " | |-- details: struct (nullable = true)\n" + ) + full = df.schema.treeString() + assert full.endswith(" | | |-- address: string (nullable = true)\n") + assert df.schema.treeString(3) == full + assert df.schema.treeString(0) == full + assert df.schema.treeString(-1) == full + @requires_tree_string def test_treeString_array_type(self, spark): from spark_namespace.sql.types import ArrayType, StringType, StructField, StructType @@ -591,9 +696,230 @@ def test_treeString_array_type(self, spark): data = [("Alice", ["reading", "coding"])] df = spark.createDataFrame(data, schema) - tree = df.schema.treeString() + assert df.schema.treeString() == ( + "root\n" + " |-- name: string (nullable = true)\n" + " |-- hobbies: array (nullable = true)\n" + " | |-- element: string (containsNull = true)\n" + ) - assert "root\n" in tree - assert " |-- name:" in tree - assert " |-- hobbies: array<" in tree - assert "(nullable = true)" in tree + @requires_tree_string + @pytest.mark.parametrize( + ("max_depth", "expected"), + [ + pytest.param( + None, + ( + "root\n" + " |-- id: integer (nullable = true)\n" + " |-- person: struct (nullable = true)\n" + " | |-- name: string (nullable = true)\n" + " | |-- details: struct (nullable = true)\n" + " | | |-- address: string (nullable = false)\n" + " |-- arr: array (nullable = true)\n" + " | |-- element: string (containsNull = false)\n" + " |-- arr_struct: array (nullable = true)\n" + " | |-- element: struct (containsNull = true)\n" + " | | |-- x: integer (nullable = true)\n" + " |-- arr_arr: array (nullable = true)\n" + " | |-- element: array (containsNull = true)\n" + " | | |-- element: long (containsNull = true)\n" + " |-- map: map (nullable = true)\n" + " | |-- key: string\n" + " | |-- value: integer (valueContainsNull = false)\n" + " |-- map_nested: map (nullable = true)\n" + " | |-- key: string\n" + " | |-- value: struct (valueContainsNull = true)\n" + " | | |-- y: decimal(5,2) (nullable = true)\n" + " |-- a\\nb: string (nullable = true)\n" + ), + id="all", + ), + pytest.param( + 2, + ( + "root\n" + " |-- id: integer (nullable = true)\n" + " |-- person: struct (nullable = true)\n" + " | |-- name: string (nullable = true)\n" + " | |-- details: struct (nullable = true)\n" + " |-- arr: array (nullable = true)\n" + " | |-- element: string (containsNull = false)\n" + " |-- arr_struct: array (nullable = true)\n" + " | |-- element: struct (containsNull = true)\n" + " |-- arr_arr: array (nullable = true)\n" + " | |-- element: array (containsNull = true)\n" + " |-- map: map (nullable = true)\n" + " | |-- key: string\n" + " | |-- value: integer (valueContainsNull = false)\n" + " |-- map_nested: map (nullable = true)\n" + " | |-- key: string\n" + " | |-- value: struct (valueContainsNull = true)\n" + " |-- a\\nb: string (nullable = true)\n" + ), + id="2", + ), + ], + ) + def test_treeString_nesting(self, max_depth, expected): + schema = StructType( + [ + StructField("id", IntegerType()), + StructField( + "person", + StructType( + [ + StructField("name", StringType()), + StructField("details", StructType([StructField("address", StringType(), False)])), + ] + ), + ), + StructField("arr", ArrayType(StringType(), False)), + StructField("arr_struct", ArrayType(StructType([StructField("x", IntegerType())]))), + StructField("arr_arr", ArrayType(ArrayType(LongType()))), + StructField("map", MapType(StringType(), IntegerType(), False)), + StructField("map_nested", MapType(StringType(), StructType([StructField("y", DecimalType(5, 2))]))), + StructField("a\nb", StringType()), + ] + ) + assert (schema.treeString() if max_depth is None else schema.treeString(max_depth)) == expected + + +class TestRepeatedColumnNames: + @pytest.fixture + def repeated(self, spark): + return spark.createDataFrame([(1, "Ann")], ["id", "name"]).toDF("x", "x") + + def test_reference_is_ambiguous(self, repeated): + assert repeated.columns == ["x", "x"] + for make in (lambda: repeated.select("x"), lambda: repeated.x, lambda: repeated["x"]): + with pytest.raises(AnalysisException) as error: + make().collect() + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + + def test_columns_keep_their_values(self, repeated): + assert [tuple(r) for r in repeated.withColumn("z", lit(1)).collect()] == [(1, "Ann", 1)] + assert [tuple(r) for r in repeated.withColumnRenamed("x", "y").collect()] == [(1, "Ann")] + assert [tuple(r) for r in repeated.toDF("a", "b").collect()] == [(1, "Ann")] + + @pytest.fixture + def beside(self, spark): + """A repeated name beside a name equal to the name DuckDB's binder gives the repeat's second column.""" + return spark.createDataFrame([(10, 20, 30)], ["x", "x", "x_1"]) + + @pytest.mark.parametrize( + "derive", + [ + lambda df: df, + lambda df: df.limit(1), + lambda df: df.select("*"), + lambda df: df.distinct(), + lambda df: df.sort("x_1"), + lambda df: df.filter(col("x_1") > 0), + ], + ) + def test_unique_name_beside_a_repeat(self, beside, derive): + df = derive(beside) + assert df.select("x_1").collect() == [Row(x_1=30)] + assert [tuple(r) for r in df.select(df.x_1, df["x_1"], col("x_1") + 1).collect()] == [(30, 30, 31)] + assert [tuple(r) for r in df.filter(col("x_1") == 30).collect()] == [(10, 20, 30)] + assert df.groupBy().sum("x_1").collect() == [Row(**{"sum(x_1)": 30})] + assert df.groupBy("x_1").count().collect() == [Row(x_1=30, count=1)] + assert df.withColumn("y", col("x_1") * 2).select("y").collect() == [Row(y=60)] + assert df.sort(col("x_1").desc()).select("x_1").collect() == [Row(x_1=30)] + assert [tuple(r) for r in df.drop("x_1").collect()] == [(10, 20)] + + def test_unique_name_beside_a_repeat_from_sql_and_renames(self, spark): + assert spark.sql("SELECT 10 AS x, 20 AS x, 30 AS x_1").select("x_1").collect() == [Row(x_1=30)] + source = spark.createDataFrame([(10, 20, 30)], ["a", "b", "c"]) + assert source.toDF("x", "x", "x_1").select("x_1").collect() == [Row(x_1=30)] + renamed = source.toDF("x", "y", "x_1").withColumnRenamed("y", "x") + assert renamed.select("x_1").collect() == [Row(x_1=30)] + assert source.select("a", "b").toDF("x", "x_1").select("x", "x", "x_1").select("x_1").collect() == [Row(x_1=20)] + assert source.toDF("x", "X", "x_1").select("x_1").collect() == [Row(x_1=30)] + + def test_name_of_a_repeat_in_duckdb_is_unresolved(self, repeated): + for make in ( + lambda: repeated.select("x_1"), + lambda: repeated.select(col("x_1")), + lambda: repeated.filter(col("x_1") == "Ann"), + lambda: repeated["x_1"], + ): + with pytest.raises(AnalysisException) as error: + make().collect() + assert error_condition(error.value) == "UNRESOLVED_COLUMN.WITH_SUGGESTION" + + def test_join_on_a_unique_name_beside_a_repeat(self, spark, beside): + other = spark.createDataFrame([(30, "hit")], ["k", "v"]) + assert beside.join(other, beside["x_1"] == other["k"]).select("v").collect() == [Row(v="hit")] + assert beside.join(other, col("x_1") == col("k")).select("v").collect() == [Row(v="hit")] + keyed = spark.createDataFrame([(30,)], ["x_1"]) + assert [tuple(r) for r in beside.join(keyed, "x_1").select("x_1").collect()] == [(30,)] + + @pytest.mark.parametrize( + "derive", + [ + lambda j, a: j.limit(1), + lambda j, a: j.sort("x"), + lambda j, a: j.distinct(), + lambda j, a: j.select("*"), + lambda j, a: j.filter(a.x > 0), + lambda j, a: j.limit(1).filter(col("x") > 0), + ], + ) + def test_join_output_name_beside_a_repeat(self, spark, derive): + a = spark.createDataFrame([(1, 10)], ["id", "x"]) + b = spark.createDataFrame([(1, 30)], ["id", "id_1"]) + df = derive(a.join(b, a.id == b.id), a) + assert df.select("id_1").collect() == [Row(id_1=30)] + assert df.filter(col("id_1") == 30).select("x").collect() == [Row(x=10)] + assert df.groupBy().sum("id_1").collect() == [Row(**{"sum(id_1)": 30})] + crossed = a.crossJoin(b).filter(a.id == b.id) + assert crossed.select("id_1").collect() == [Row(id_1=30)] + + @pytest.fixture + def filtered_join(self, spark): + a = spark.createDataFrame([(1, 10)], ["id", "x"]).alias("a") + b = spark.createDataFrame([(1, 30)], ["id", "id_1"]).alias("b") + return a.join(b, a.id == b.id).filter(a.x > 0) + + def test_filtered_join_as_a_join_input(self, spark, filtered_join): + other = spark.createDataFrame([(100,)], ["z"]) + for joined in (filtered_join.crossJoin(other), other.crossJoin(filtered_join)): + assert joined.select("id_1", col("id_1") + 1).collect() == [Row(id_1=30, **{"(id_1 + 1)": 31})] + assert joined.groupBy().sum("id_1").collect() == [Row(**{"sum(id_1)": 30})] + assert [tuple(r) for r in filtered_join.crossJoin(other).select("*").collect()] == [(1, 10, 1, 30, 100)] + keyed = spark.createDataFrame([(30, "hit")], ["k", "v"]) + assert filtered_join.join(keyed, col("id_1") == col("k")).select("v").collect() == [Row(v="hit")] + + def test_filtered_join_on_names_as_a_join_input(self, spark): + a = spark.createDataFrame([(1, 10, 7)], ["id", "x", "k"]) + b = spark.createDataFrame([(1, 20)], ["id", "x"]) + joined = a.join(b, "id").filter(col("k") > 0).crossJoin(spark.createDataFrame([(100,)], ["z"])) + assert [tuple(r) for r in joined.collect()] == [(1, 10, 7, 20, 100)] + + def test_view_of_a_join_with_repeated_names(self, spark, filtered_join): + a = spark.createDataFrame([(1, 10)], ["id", "x"]).alias("a") + b = spark.createDataFrame([(1, 30)], ["id", "id_1"]).alias("b") + for name, df in (("repeated_join_view", a.join(b, a.id == b.id)), ("repeated_filter_view", filtered_join)): + df.createOrReplaceTempView(name) + assert spark.table(name).select("id_1").collect() == [Row(id_1=30)] + repeated = spark.createDataFrame([(10, 20)], ["x", "x"]).alias("r") + repeated.crossJoin(spark.createDataFrame([(30,)], ["x_1"])).createOrReplaceTempView("repeated_input_view") + assert spark.table("repeated_input_view").select("x_1").collect() == [Row(x_1=30)] + # Spark's view keeps the repeated names; DuckDB names a view's columns apart. + expected = ["x", "x", "x_1"] if USE_ACTUAL_SPARK else ["x", "x_2", "x_1"] + assert spark.table("repeated_input_view").columns == expected + + def test_repeat_beside_a_name_of_another_join_input(self, spark): + a = spark.createDataFrame([(10, 20)], ["x", "x"]).alias("a") + b = spark.createDataFrame([(30,)], ["x_1"]).alias("b") + y = spark.createDataFrame([(7,)], ["y"]).alias("y") + for joined in (a.crossJoin(b), b.crossJoin(a), a.crossJoin(y).crossJoin(b)): + assert joined.select("x_1").collect() == [Row(x_1=30)] + assert joined.select(col("x_1")).collect() == [Row(x_1=30)] + assert joined.filter(col("x_1") == 30).select("x_1").collect() == [Row(x_1=30)] + assert joined.groupBy(col("x_1")).count().collect() == [Row(x_1=30, count=1)] + assert joined.sort(col("x_1")).select("x_1").collect() == [Row(x_1=30)] + assert [tuple(r) for r in a.crossJoin(y).crossJoin(b).select(y.y, col("y.y")).collect()] == [(7, 7)] + assert a.join(b, col("x_1") == 30).select("x_1").collect() == [Row(x_1=30)] diff --git a/tests/fast/spark/test_spark_drop_duplicates.py b/tests/fast/spark/test_spark_drop_duplicates.py index cd658c776..49cabdfba 100644 --- a/tests/fast/spark/test_spark_drop_duplicates.py +++ b/tests/fast/spark/test_spark_drop_duplicates.py @@ -1,10 +1,19 @@ import pytest +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException, error_condition from spark_namespace.sql.types import ( Row, ) _ = pytest.importorskip("duckdb.experimental.spark") +if USE_ACTUAL_SPARK: + import pyspark + + SPARK_3 = pyspark.__version__.startswith("3.") +else: + SPARK_3 = False + class TestDataFrameDropDuplicates: @pytest.mark.parametrize("method", ["dropDuplicates", "drop_duplicates"]) @@ -84,3 +93,35 @@ def test_spark_drop_duplicates_with_keywords_cols(self, spark): dropDisDF = df.dropDuplicates(["table", "min"]).sort("table", "min") assert dropDisDF.count() == 3 + + @pytest.mark.parametrize("subset", [[], ()]) + def test_drop_duplicates_with_an_empty_subset(self, spark, subset): + # Spark compares rows on no columns at all, so they are all duplicates of the first. + df = spark.createDataFrame([(1,), (2,), (2,)], ["x"]) + assert df.dropDuplicates(subset).count() == 1 + assert df.dropDuplicates().count() == 2 + assert df.dropDuplicates(None).count() == 2 + assert df.filter(df.x > 5).dropDuplicates(subset).count() == 0 + + def test_drop_duplicates_on_a_repeated_name(self, spark): + # Spark compares every column of the name, so rows that differ in the second "id" both stay. + df = spark.createDataFrame([(1, 1), (1, 2), (1, 2), (2, 2)], ["id", "x"]).toDF("id", "id") + assert sorted(tuple(row) for row in df.dropDuplicates(["id"]).collect()) == [(1, 1), (1, 2), (2, 2)] + assert sorted(tuple(row) for row in df.dropDuplicates(["ID"]).collect()) == [(1, 1), (1, 2), (2, 2)] + + def test_drop_duplicates_after_a_join(self, spark): + # The window must compare the join's own columns, not the inputs' columns a positional key would reach. + a = spark.createDataFrame([(1,)], ["id"]) + b = spark.createDataFrame([(1, 10), (1, 20)], ["id", "v"]) + joined = a.join(b, "id") + assert [row.v for row in joined.dropDuplicates(["v"]).sort("v").collect()] == [10, 20] + assert joined.filter(joined.v > 0).dropDuplicates(["v"]).count() == 2 + assert joined.dropDuplicates(["id"]).count() == 1 + + def test_drop_duplicates_on_an_unknown_name(self, spark): + df = spark.createDataFrame([(1, 1)], ["id", "x"]) + with pytest.raises(AnalysisException) as error: + df.dropDuplicates(["nope"]).collect() + # Spark 3.5 raises the same error under a legacy class. + expected = "_LEGACY_ERROR_TEMP_1201" if SPARK_3 else "UNRESOLVED_COLUMN_AMONG_FIELD_NAMES" + assert error_condition(error.value) == expected diff --git a/tests/fast/spark/test_spark_filter.py b/tests/fast/spark/test_spark_filter.py index 8d412a2c0..a6b41ad1a 100644 --- a/tests/fast/spark/test_spark_filter.py +++ b/tests/fast/spark/test_spark_filter.py @@ -3,7 +3,6 @@ _ = pytest.importorskip("duckdb.experimental.spark") -from spark_namespace import USE_ACTUAL_SPARK from spark_namespace.errors import PySparkTypeError from spark_namespace.sql.functions import array_contains, col from spark_namespace.sql.types import ( @@ -138,11 +137,8 @@ def test_dataframe_filter(self, spark): df2 = df.filter(array_contains(df.languages, "Java")) res = df2.collect() - james_name = {"firstname": "James", "middlename": "", "lastname": "Smith"} - anna_name = {"firstname": "Anna", "middlename": "Rose", "lastname": ""} - if USE_ACTUAL_SPARK: - james_name = Row(**james_name) - anna_name = Row(**anna_name) + james_name = Row(firstname="James", middlename="", lastname="Smith") + anna_name = Row(firstname="Anna", middlename="Rose", lastname="") assert res == [ Row( name=james_name, @@ -160,11 +156,8 @@ def test_dataframe_filter(self, spark): df2 = df.filter(df.name.lastname == "Williams") res = df2.collect() - julia_name = {"firstname": "Julia", "middlename": "", "lastname": "Williams"} - mike_name = {"firstname": "Mike", "middlename": "Mary", "lastname": "Williams"} - if USE_ACTUAL_SPARK: - julia_name = Row(**julia_name) - mike_name = Row(**mike_name) + julia_name = Row(firstname="Julia", middlename="", lastname="Williams") + mike_name = Row(firstname="Mike", middlename="Mary", lastname="Williams") assert res == [ Row( name=julia_name, diff --git a/tests/fast/spark/test_spark_functions_array.py b/tests/fast/spark/test_spark_functions_array.py index 9ee2ffc22..fcab81c9f 100644 --- a/tests/fast/spark/test_spark_functions_array.py +++ b/tests/fast/spark/test_spark_functions_array.py @@ -1,11 +1,24 @@ +import datetime import platform +import re import pytest _ = pytest.importorskip("duckdb.experimental.spark") from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException, PySparkTypeError from spark_namespace.sql import functions as sf -from spark_namespace.sql.types import Row +from spark_namespace.sql.types import ArrayType, IntegerType, LongType, Row, StringType, StructField, StructType + +import duckdb + +if USE_ACTUAL_SPARK: + import pyspark + + +def _arrays(spark, rows, *columns): + return spark.createDataFrame(rows, StructType([StructField(c, ArrayType(LongType())) for c in columns])) + pytestmark = pytest.mark.skipif( platform.system() == "Emscripten", @@ -38,6 +51,15 @@ def test_array_intersect(self, spark): assert len(res) == 1 assert sorted(res[0].intersect_values) == ["a", "c"] + def test_array_intersect_of_two_element_types(self, spark): + df = spark.createDataFrame( + [([16777216, 16777217, 1, 1, None], [16777216.0, 1.0, None], [1, 2])], + "i array, f array, b array", + ) + res = df.select(sf.array_intersect("i", "f").alias("v"), sf.array_intersect("i", "b").alias("w")) + assert res.schema.simpleString() == "struct,w:array>" + assert res.collect() == [Row(v=[16777216.0, 1.0, None], w=[1])] + def test_array_union(self, spark): data = [ (["b", "a", "c"], ["c", "d", "a", "f"]), @@ -93,6 +115,13 @@ def test_get(self, spark): res = df.select(sf.get(df.data, sf.col("index") - 1).alias("r")).collect() assert res == [Row(r="a")] + def test_get_below_zero_is_null(self, spark): + df = spark.createDataFrame( + [(["a", "b", "c"], -2), (["a", "b", "c"], -4), (["a", "b", "c"], None)], ["data", "i"] + ) + res = df.select(sf.get("data", -2).alias("l"), sf.get("data", -3).alias("m"), sf.get("data", "i").alias("c")) + assert res.collect() == [Row(l=None, m=None, c=None)] * 3 + def test_flatten(self, spark): df = spark.createDataFrame([([[1, 2, 3], [4, 5], [6]],), ([None, [4, 5]],)], ["data"]) @@ -102,8 +131,8 @@ def test_flatten(self, spark): def test_array_compact(self, spark): df = spark.createDataFrame([([1, None, 2, 3],), ([4, 5, None, 4],)], ["data"]) - df.select(sf.array_compact(df.data).alias("v")).collect() - assert [Row(v=[1, 2, 3]), Row(v=[4, 5, 4])] + res = df.select(sf.array_compact(df.data).alias("v")).collect() + assert res == [Row(v=[1, 2, 3]), Row(v=[4, 5, 4])] def test_array_remove(self, spark): df = spark.createDataFrame([([1, 2, 3, 1, 1],), ([],)], ["data"]) @@ -210,6 +239,14 @@ def test_array_sort(self, spark): res = df.select(sf.array_sort(df.data).alias("r")).collect() assert res == [Row(r=[1, 2, 3, None]), Row(r=[1]), Row(r=[])] + def test_array_sort_name(self, spark): + (name,) = spark.createDataFrame([([2, 1],)], ["data"]).select(sf.array_sort("data")).columns + if USE_ACTUAL_SPARK: + # Spark names the call with the text of its default comparator. + assert name.startswith("array_sort(data, lambdafunction(") + else: + assert name == "array_sort(data)" + def test_arrays_overlap(self, spark): df = spark.createDataFrame( [(["a", "b"], ["b", "c"]), (["a"], ["b", "c"]), ([None, "c"], ["a"]), ([None, "c"], [None])], ["x", "y"] @@ -235,3 +272,375 @@ def test_arrays_zip(self, spark): ] else: assert res == [Row(zipped=[(1, 2, 3), (2, 4, 6), (3, 6, None)])] + + def test_set_operations_keep_order_and_nulls(self, spark): + df = _arrays(spark, [([1, None, 2, 2, 3], [3, None, 2, 2]), ([2, 5, 2], None)], "a", "b") + res = df.select( + sf.array_distinct("a").alias("distinct"), + sf.array_union("a", "b").alias("union"), + sf.array_union("b", "a").alias("union_ba"), + sf.array_intersect("a", "b").alias("intersect"), + ).collect() + assert res == [ + Row(distinct=[1, None, 2, 3], union=[1, None, 2, 3], union_ba=[3, None, 2, 1], intersect=[None, 2, 3]), + Row(distinct=[2, 5], union=None, union_ba=None, intersect=None), + ] + + def test_array_max_min_skip_nulls(self, spark): + df = _arrays(spark, [([3, 1, None, 2],), ([None],), (None,)], "a") + res = df.select(sf.array_max("a").alias("max"), sf.array_min("a").alias("min")).collect() + assert res == [Row(max=3, min=1), Row(max=None, min=None), Row(max=None, min=None)] + + def test_array_contains_with_nulls(self, spark): + df = _arrays(spark, [([1, None],), ([1],), ([None],), (None,)], "a") + res = df.select(sf.array_contains("a", 1).alias("one"), sf.array_contains("a", 2).alias("two")).collect() + assert res == [ + Row(one=True, two=None), + Row(one=True, two=False), + Row(one=None, two=None), + Row(one=None, two=None), + ] + + def test_array_position_of_null_array(self, spark): + df = _arrays(spark, [([1, None, 3],), (None,)], "a") + res = df.select(sf.array_position("a", 3).alias("p"), sf.array_position("a", 5).alias("q")) + assert res.schema.fields[0].dataType == LongType() + assert res.collect() == [Row(p=3, q=0), Row(p=None, q=None)] + + def test_array_remove_keeps_nulls(self, spark): + df = _arrays(spark, [([1, None, 2, 1],)], "a") + assert df.select(sf.array_remove("a", 1).alias("r")).collect() == [Row(r=[None, 2])] + + def test_array_insert_positions(self, spark): + df = _arrays(spark, [([1, 2, 3],), ([],), (None,)], "a") + positions = { + "first": 1, + "middle": 3, + "append": 4, + "pad": 6, + "last": -1, + "before_last": -2, + "front": -4, + "pad_front": -5, + } + res = df.select(*[sf.array_insert("a", p, sf.lit(9)).alias(name) for name, p in positions.items()]).collect() + assert res == [ + Row( + first=[9, 1, 2, 3], + middle=[1, 2, 9, 3], + append=[1, 2, 3, 9], + pad=[1, 2, 3, None, None, 9], + last=[1, 2, 3, 9], + before_last=[1, 2, 9, 3], + front=[9, 1, 2, 3], + pad_front=[9, None, 1, 2, 3], + ), + Row( + first=[9], + middle=[None, None, 9], + append=[None, None, None, 9], + pad=[None, None, None, None, None, 9], + last=[9], + before_last=[9, None], + front=[9, None, None, None], + pad_front=[9, None, None, None, None], + ), + Row( + first=None, middle=None, append=None, pad=None, last=None, before_last=None, front=None, pad_front=None + ), + ] + with pytest.raises(Exception, match="INVALID_INDEX_OF_ZERO"): + df.select(sf.array_insert("a", 0, sf.lit(9))).collect() + + def test_slice_bounds(self, spark): + df = _arrays(spark, [([1, 2, 3, 4],), (None,)], "a") + res = df.select( + sf.slice("a", 2, 2).alias("mid"), + sf.slice("a", -2, 2).alias("tail"), + sf.slice("a", 1, 10).alias("long"), + sf.slice("a", -5, 2).alias("before"), + sf.slice("a", 3, 0).alias("empty"), + sf.slice("a", 5, 1).alias("past"), + ).collect() + assert res == [ + Row(mid=[2, 3], tail=[3, 4], long=[1, 2, 3, 4], before=[], empty=[], past=[]), + Row(mid=None, tail=None, long=None, before=None, empty=None, past=None), + ] + with pytest.raises(Exception, match="start"): + df.select(sf.slice("a", 0, 2)).collect() + with pytest.raises(Exception, match="length"): + df.select(sf.slice("a", 1, -1)).collect() + + def test_sequence_counts_down(self, spark): + res = spark.range(1).select( + sf.sequence(sf.lit(5), sf.lit(1)).alias("down"), + sf.sequence(sf.lit(1), sf.lit(3)).alias("up"), + sf.sequence(sf.lit(2), sf.lit(2)).alias("one"), + sf.sequence(sf.lit(1), sf.lit(9), sf.lit(3)).alias("step"), + ) + assert res.collect() == [Row(down=[5, 4, 3, 2, 1], up=[1, 2, 3], one=[2], step=[1, 4, 7])] + + @pytest.mark.parametrize( + ("start", "stop", "step"), + [ + ("a", "b", "-1"), + ("b", "a", "1"), + ("a", "b", "0"), + ("d1", "d2", "INTERVAL '-1' DAY"), + ], + ) + def test_sequence_step_away_from_stop_raises(self, spark, start, stop, step): + df = spark.createDataFrame( + [(1, 5, datetime.date(2024, 1, 1), datetime.date(2024, 1, 5))], ["a", "b", "d1", "d2"] + ) + if USE_ACTUAL_SPARK: + from py4j.protocol import Py4JJavaError + from spark_namespace.errors import IllegalArgumentException + + errors = (IllegalArgumentException, Py4JJavaError) + else: + errors = duckdb.InvalidInputException + with pytest.raises(errors, match="Illegal sequence boundaries"): + df.select(sf.sequence(start, stop, sf.expr(step))).collect() + + def test_sequence_element_types(self, spark): + day = datetime.date + schema = "i int, j int, s smallint, b bigint, d1 date, d2 date" + df = spark.createDataFrame([(1, 3, 2, 3, day(2024, 1, 1), day(2024, 1, 3))], schema) + res = df.select( + sf.sequence("i", "j").alias("ints"), + sf.sequence(sf.lit(3), sf.lit(1)).alias("literals"), + sf.sequence("s", sf.lit(4).cast("smallint"), sf.lit(1).cast("smallint")).alias("smallints"), + sf.sequence("i", "b").alias("widened"), + sf.sequence("b", sf.lit(1), sf.lit(-1)).alias("narrow_step"), + sf.sequence("d1", "d2").alias("days"), + sf.sequence("d2", "d1").alias("days_down"), + sf.sequence("d1", "d2", sf.lit(datetime.timedelta(days=2))).alias("every_other_day"), + ) + assert res.schema.simpleString() == ( + "struct,literals:array,smallints:array,widened:array," + "narrow_step:array,days:array,days_down:array,every_other_day:array>" + ) + assert res.collect() == [ + Row( + ints=[1, 2, 3], + literals=[3, 2, 1], + smallints=[2, 3, 4], + widened=[1, 2, 3], + narrow_step=[3, 2, 1], + days=[day(2024, 1, 1), day(2024, 1, 2), day(2024, 1, 3)], + days_down=[day(2024, 1, 3), day(2024, 1, 2), day(2024, 1, 1)], + every_other_day=[day(2024, 1, 1), day(2024, 1, 3)], + ) + ] + + def test_sequence_of_small_integers_spans_their_range(self, spark): + df = spark.createDataFrame([(-100, 100, -30000, 30000)], "a tinyint, b tinyint, c smallint, d smallint") + res = df.select(sf.sequence("a", "b").alias("tiny"), sf.sequence("c", "d").alias("small")) + assert res.schema.simpleString() == "struct,small:array>" + row = res.collect()[0] + assert row.tiny == list(range(-100, 101)) + assert row.small == list(range(-30000, 30001)) + + def test_sequence_of_months(self, spark): + day = datetime.date + df = spark.createDataFrame([(day(2024, 1, 1), day(2024, 1, 31), day(2024, 4, 30))], "a date, b date, c date") + months = df.select(sf.sequence("a", "c", sf.expr("INTERVAL 1 MONTH")).alias("m")) + from_the_end = df.select(sf.sequence("b", "c", sf.expr("INTERVAL 1 MONTH")).alias("m")) + years = df.select(sf.sequence("a", "c", sf.expr("INTERVAL 1 YEAR")).alias("m")) + if USE_ACTUAL_SPARK: + assert months.collect() == [Row(m=[day(2024, 1, 1), day(2024, 2, 1), day(2024, 3, 1), day(2024, 4, 1)])] + spark_months = [day(2024, 1, 31), day(2024, 2, 29), day(2024, 3, 31), day(2024, 4, 30)] + assert from_the_end.collect() == [Row(m=spark_months)] + assert years.collect() == [Row(m=[day(2024, 1, 1)])] + else: + # DuckDB adds a step of months to each element in turn, which gives 2024-03-29 after 2024-02-29 where + # Spark adds multiples of it to start, so the module raises for any step with months or years. + for frame in [months, from_the_end, years]: + with pytest.raises(duckdb.InvalidInputException, match="step of months"): + frame.collect() + + def test_sequence_of_timestamps(self, spark): + moment = datetime.datetime + df = spark.createDataFrame([(moment(2024, 1, 1), moment(2024, 1, 2))], "t1 timestamp, t2 timestamp") + res = df.select(sf.sequence("t1", "t2", sf.lit(datetime.timedelta(hours=12))).alias("t")) + assert res.schema.simpleString() == "struct>" + assert res.collect() == [Row(t=[moment(2024, 1, 1), moment(2024, 1, 1, 12), moment(2024, 1, 2)])] + + def test_sequence_with_a_null_bound_is_null(self, spark): + df = spark.createDataFrame([(None, 3, 0), (1, None, 0), (None, None, 0), (2, 2, 0)], "a int, b int, s int") + res = df.select(sf.sequence("a", "b", "s").alias("zero_step"), sf.sequence("a", "b").alias("no_step")) + assert res.collect() == [ + Row(zero_step=None, no_step=None), + Row(zero_step=None, no_step=None), + Row(zero_step=None, no_step=None), + Row(zero_step=[2], no_step=[2]), + ] + dates = spark.createDataFrame([(None, datetime.date(2024, 1, 3))], "a date, b date") + assert dates.select(sf.sequence("a", "b").alias("d")).collect() == [Row(d=None)] + + def test_sequence_of_dates_by_part_of_a_day_raises(self, spark): + df = spark.createDataFrame([(datetime.date(2024, 1, 1), datetime.date(2024, 1, 3))], "d1 date, d2 date") + if USE_ACTUAL_SPARK: + from py4j.protocol import Py4JJavaError + from spark_namespace.errors import IllegalArgumentException + + errors = (IllegalArgumentException, Py4JJavaError) + else: + errors = duckdb.InvalidInputException + with pytest.raises(errors, match="day granularity"): + df.select(sf.sequence("d1", "d2", sf.lit(datetime.timedelta(hours=12)))).collect() + + def test_sequence_with_a_step_of_a_wider_type(self, spark): + df = spark.createDataFrame([(1, 3)], "i int, j int") + res = df.select(sf.sequence("i", "j", sf.lit(1).cast("bigint")).alias("s")) + if USE_ACTUAL_SPARK: + assert res.schema.simpleString() == "struct>" + assert res.collect() == [Row(s=[1, 2, 3])] + else: + # Spark gives the elements the step's type, which a cast to the type of start and stop would lose. + with pytest.raises(duckdb.InvalidInputException, match="wider integral type"): + res.collect() + + def test_sequence_with_a_zero_step_of_a_wider_type(self, spark): + df = spark.createDataFrame([(1, 0)], "i int, z bigint") + for step in ("z", sf.lit(0).cast("bigint")): + res = df.select(sf.sequence("i", "i", step).alias("s")) + if USE_ACTUAL_SPARK: + assert res.schema.simpleString() == "struct>" + assert res.collect() == [Row(s=[1])] + else: + with pytest.raises(duckdb.InvalidInputException, match="wider integral type"): + res.collect() + + def test_sequence_without_steps_to_take(self, spark): + df = spark.createDataFrame( + [(1, None)], StructType([StructField("a", LongType()), StructField("n", LongType())]) + ) + res = df.select( + sf.sequence("a", "a", sf.lit(0)).alias("zero"), + sf.sequence("a", "a", sf.lit(-1)).alias("away"), + sf.sequence("a", sf.lit(5), "n").alias("null_step"), + ) + assert res.collect() == [Row(zero=[1], away=[1], null_step=None)] + + def test_array_size_type(self, spark): + df = _arrays(spark, [([1, 2],), (None,)], "a") + res = df.select(sf.array_size("a").alias("n")) + assert res.schema.fields[0].dataType == IntegerType() + assert res.collect() == [Row(n=2), Row(n=None)] + + def test_array_join_of_null_elements_only(self, spark): + df = spark.createDataFrame([([],), ([None],), ([None, None],), (None,)], "a array") + res = df.select(sf.array_join("a", ",").alias("j"), sf.array_join("a", ",", "N").alias("r")).collect() + assert res == [Row(j="", r=""), Row(j="", r="N"), Row(j="", r="N,N"), Row(j=None, r=None)] + + def test_array_join_null_replacement_on_numbers(self, spark): + df = _arrays(spark, [([1, None, 2],), (None,)], "a") + res = df.select(sf.array_join("a", "-").alias("j"), sf.array_join("a", "-", "N").alias("r")).collect() + assert res == [Row(j="1-2", r="1-N-2"), Row(j=None, r=None)] + + def test_array_repeat_non_positive_count(self, spark): + schema = StructType([StructField("s", StringType()), StructField("n", IntegerType())]) + df = spark.createDataFrame([("x", 2), ("y", 0), ("z", -1)], schema) + res = df.select(sf.array_repeat("s", "n").alias("r"), sf.array_repeat("s", -2).alias("c")).collect() + assert res == [Row(r=["x", "x"], c=[]), Row(r=[], c=[]), Row(r=[], c=[])] + + def test_append_prepend_null_array(self, spark): + df = _arrays(spark, [(None,), ([1],)], "a") + res = df.select(sf.array_append("a", sf.lit(9)).alias("app"), sf.array_prepend("a", 0).alias("pre")).collect() + assert res == [Row(app=None, pre=None), Row(app=[1, 9], pre=[0, 1])] + + def test_null_arguments_give_null(self, spark): + schema = StructType( + [ + StructField("a", ArrayType(LongType())), + StructField("n", ArrayType(LongType())), + StructField("i", IntegerType()), + ] + ) + df = spark.createDataFrame([([1, 2, 3, 2], None, None)], schema) + res = df.select( + sf.array_position("a", None).alias("position"), + sf.array_remove("a", None).alias("remove"), + sf.array_repeat(sf.lit(5), "i").alias("repeat"), + sf.array_insert("a", "i", sf.lit(9)).alias("insert"), + sf.array_insert("n", 1, sf.lit(9)).alias("insert_into_null"), + sf.slice("n", 0, 2).alias("slice_from_zero"), + sf.slice("a", "i", sf.lit(-1)).alias("slice_null_start"), + sf.slice("a", sf.lit(0), "i").alias("slice_null_length"), + sf.arrays_overlap("a", "n").alias("overlap"), + sf.arrays_overlap("n", sf.array().cast("array")).alias("overlap_empty"), + sf.arrays_zip("a", "n").alias("zip"), + sf.array_intersect("n", "a").alias("intersect"), + ) + assert res.collect() == [Row(*[None] * 12)] + + def test_null_elements_are_kept(self, spark): + schema = StructType([StructField("a", ArrayType(LongType())), StructField("i", IntegerType())]) + df = spark.createDataFrame([([1, 2], 2)], schema) + res = df.select( + sf.array_insert("a", "i", sf.lit(None).cast("bigint")).alias("insert"), + sf.array_append("a", None).alias("append"), + sf.array_prepend("a", None).alias("prepend"), + sf.array_repeat(sf.lit(None).cast("bigint"), "i").alias("repeat"), + ) + assert res.collect() == [ + Row(insert=[1, None, 2], append=[1, 2, None], prepend=[None, 1, 2], repeat=[None, None]), + ] + + def test_arrays_zip_field_names(self, spark): + schema = StructType([StructField("a", ArrayType(LongType())), StructField("b", ArrayType(StringType()))]) + df = spark.createDataFrame([([1, 2], ["x"])], schema) + res = df.select( + sf.arrays_zip("a", "b").alias("z"), sf.arrays_zip(sf.col("a"), sf.array_compact("b")).alias("w") + ) + assert [f.dataType.elementType.fieldNames() for f in res.schema.fields] == [["a", "b"], ["a", "1"]] + assert res.collect() == [Row(z=[Row(a=1, b="x"), Row(a=2, b=None)], w=[Row(1, "x"), Row(2, None)])] + + def test_array_distinct_evaluates_its_input_once(self, spark): + rows = spark.range(20).select(sf.array_distinct(sf.array(sf.rand(), sf.lit(2.0))).alias("a")).collect() + assert {len(r.a) for r in rows} == {2} + + def test_arrays_zip_of_no_arrays(self, spark): + res = spark.range(1).select(sf.arrays_zip()) + assert res.schema.simpleString() == "struct>>" + assert res.collect() == [Row([])] + + def test_array_contains_null_literal_raises(self, spark): + df = _arrays(spark, [([1, 2],)], "a") + for value in [None, sf.lit(None)]: + with pytest.raises(AnalysisException, match=re.escape("[DATATYPE_MISMATCH.NULL_TYPE]")): + df.select(sf.array_contains("a", value)).collect() + assert df.select(sf.array_contains("a", sf.lit(None).cast("bigint")).alias("c")).collect() == [Row(c=None)] + + @pytest.mark.parametrize(("first", "second"), [("x", "i"), ("x0", "x1"), ("x1", "x0"), ("X0", "row")]) + def test_columns_named_like_lambda_parameters(self, spark, first, second): + df = _arrays(spark, [([1, 2, 2, None], [2, 3, None])], first, second) + res = df.select( + sf.array_distinct(first).alias("distinct"), + sf.array_intersect(first, second).alias("intersect"), + sf.array_union(first, second).alias("union"), + sf.array_compact(first).alias("compact"), + sf.array_join(first, ",", "-").alias("join"), + sf.arrays_zip(first, second).alias("zip"), + ) + assert res.collect() == [ + Row( + distinct=[1, 2, None], + intersect=[2, None], + union=[1, 2, None, 3], + compact=[1, 2, 2], + join="1,2,2,-", + zip=[Row(1, 2), Row(2, 3), Row(2, None), Row(None, None)], + ) + ] + removed = spark.createDataFrame([([1, 2, 3], 2)], f"{first} array, {second} bigint") + if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3."): + # PySpark 3.5 takes only a literal element. + with pytest.raises(PySparkTypeError, match="NOT_ITERABLE"): + removed.select(sf.array_remove(first, sf.col(second))) + else: + assert removed.select(sf.array_remove(first, sf.col(second)).alias("r")).collect() == [Row(r=[1, 3])] + grouped = spark.createDataFrame([(1,), (None,), (2,)], f"{first} bigint") + assert sorted(grouped.groupBy().agg(sf.collect_list(first).alias("l")).collect()[0].l) == [1, 2] diff --git a/tests/fast/spark/test_spark_functions_date.py b/tests/fast/spark/test_spark_functions_date.py index 914d33f66..ee3ad4d3f 100644 --- a/tests/fast/spark/test_spark_functions_date.py +++ b/tests/fast/spark/test_spark_functions_date.py @@ -3,12 +3,23 @@ import pytest _ = pytest.importorskip("duckdb.experimental.spark") -from datetime import date, datetime +from datetime import date, datetime, timezone +from decimal import Decimal from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException from spark_namespace.sql import functions as F from spark_namespace.sql.functions import col -from spark_namespace.sql.types import Row +from spark_namespace.sql.types import ( + DateType, + IntegerType, + Row, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, +) class TestsSparkFunctionsDate: @@ -223,10 +234,15 @@ def test_date_diff(self, spark): df = spark.createDataFrame([("2015-04-08", "2015-05-10")], ["d1", "d2"]) result_data = df.select(F.date_diff(col("d2").cast("DATE"), col("d1").cast("DATE")).alias("diff")).collect() - assert result_data[0]["diff"] == -32 + assert result_data[0]["diff"] == 32 result_data = df.select(F.date_diff(col("d1").cast("DATE"), col("d2").cast("DATE")).alias("diff")).collect() - assert result_data[0]["diff"] == 32 + assert result_data[0]["diff"] == -32 + + def test_date_diff_timestamps_count_calendar_days(self, spark): + df = spark.createDataFrame([(datetime(2015, 4, 8, 23, 0), datetime(2015, 4, 9, 1, 0))], ["t1", "t2"]) + res = df.select(F.date_diff("t2", "t1").alias("later"), F.date_diff("t1", "t2").alias("earlier")).collect() + assert res == [Row(later=1, earlier=-1)] def test_try_to_timestamp(self, spark): df = spark.createDataFrame([("1997-02-28 10:30:00",), ("2024-01-01",), ("invalid",)], ["t"]) @@ -237,7 +253,174 @@ def test_try_to_timestamp(self, spark): def test_try_to_timestamp_with_format(self, spark): df = spark.createDataFrame([("1997-02-28 10:30:00",), ("2024-01-01",), ("invalid",)], ["t"]) - res = df.select(F.try_to_timestamp(df.t, format=F.lit("%Y-%m-%d %H:%M:%S")).alias("dt")).collect() + res = df.select(F.try_to_timestamp(df.t, format=F.lit("yyyy-MM-dd HH:mm:ss")).alias("dt")).collect() assert res[0].dt == datetime(1997, 2, 28, 10, 30) assert res[1].dt is None assert res[2].dt is None + + def test_try_to_timestamp_follows_cast_rules(self, spark): + texts = [ + "2024-01-01T10:30:00", + "2024-01-01 10:30:00.123456", + "2024-01-01 10:30:00Z", + "2024-1-5", + "invalid", + None, + ] + df = spark.createDataFrame([(t,) for t in texts], ["t"]) + res = df.select(F.try_to_timestamp("t").alias("dt")).collect() + # Text with an offset is an instant, shown in the session time zone, which defaults to the local one. + utc_local = datetime(2024, 1, 1, 10, 30, tzinfo=timezone.utc).astimezone().replace(tzinfo=None) + assert [r.dt for r in res] == [ + datetime(2024, 1, 1, 10, 30), + datetime(2024, 1, 1, 10, 30, 0, 123456), + utc_local, + datetime(2024, 1, 5), + None, + None, + ] + + def test_to_timestamp_applies_offset(self, spark): + df = spark.createDataFrame([("2024-01-01 10:30:00+01:00",)], ["t"]) + columns = [ + F.to_timestamp("t").alias("ltz"), + F.to_timestamp_ltz("t").alias("ltz2"), + F.try_to_timestamp("t").alias("try_ltz"), + ] + # PySpark's to_timestamp_ntz issues a DeprecationWarning, which the test settings turn into an error. + with warnings.catch_warnings(): + if USE_ACTUAL_SPARK: + warnings.filterwarnings("ignore", category=DeprecationWarning) + res = df.select(*columns, F.to_timestamp_ntz("t").alias("ntz")).collect() + # Spark shifts the instant to the session time zone, which defaults to the local one. + local = datetime(2024, 1, 1, 9, 30, tzinfo=timezone.utc).astimezone().replace(tzinfo=None) + assert res == [Row(ltz=local, ltz2=local, try_ltz=local, ntz=datetime(2024, 1, 1, 10, 30))] + + def test_date_trunc_unknown_unit_is_null(self, spark): + df = spark.createDataFrame([(datetime(2024, 1, 31, 13, 8, 15),)], ["t"]) + res = df.select( + F.date_trunc("hh", "t").alias("hh"), + F.date_trunc("decade", "t").alias("decade"), + F.date_trunc("bogus", "t").alias("bogus"), + F.date_trunc("HOUR", "t").alias("hour"), + ).collect() + assert res == [Row(hh=None, decade=None, bogus=None, hour=datetime(2024, 1, 31, 13))] + + def test_date_trunc_sub_second(self, spark): + df = spark.createDataFrame([(datetime(2024, 1, 31, 13, 8, 15, 123456),)], ["t"]) + res = df.select( + F.date_trunc("millisecond", "t").alias("ms"), F.date_trunc("microsecond", "t").alias("us") + ).collect() + assert res == [Row(ms=datetime(2024, 1, 31, 13, 8, 15, 123000), us=datetime(2024, 1, 31, 13, 8, 15, 123456))] + + def test_date_part_fields(self, spark): + df = spark.createDataFrame([(datetime(2021, 1, 3, 13, 8, 15, 123456),)], ["ts"]) + res = df.select( + F.date_part(F.lit("DAYOFWEEK"), "ts").alias("dow"), + F.date_part(F.lit("dow_iso"), "ts").alias("dow_iso"), + F.date_part(F.lit("YEAROFWEEK"), "ts").alias("year_of_week"), + F.date_part(F.lit("WEEK"), "ts").alias("week"), + F.date_part(F.lit("QTR"), "ts").alias("quarter"), + F.date_part(F.lit("DOY"), "ts").alias("doy"), + F.date_part(F.lit("SECONDS"), "ts").alias("second"), + F.extract(F.lit("S"), "ts").alias("s"), + F.datepart(F.lit("dow"), "ts").alias("dow2"), + ).collect() + assert res == [ + Row( + dow=1, + dow_iso=7, + year_of_week=2020, + week=53, + quarter=1, + doy=3, + second=Decimal("15.123456"), + s=Decimal("15.123456"), + dow2=1, + ) + ] + assert isinstance(res[0].second, Decimal) + + def test_date_part_unknown_field(self, spark): + df = spark.createDataFrame([(datetime(2021, 1, 3, 13, 8, 15),)], ["ts"]) + with pytest.raises(AnalysisException, match="INVALID_EXTRACT_FIELD"): + df.select(F.date_part(F.lit("bogus"), "ts")).collect() + + def test_add_months_month_end(self, spark): + dates = [date(2024, 1, 31), date(2024, 2, 29), date(2023, 1, 31), date(2024, 4, 30)] + df = spark.createDataFrame([(d,) for d in dates], ["d"]) + res = df.select( + F.add_months("d", 1).alias("next"), + F.add_months("d", -12).alias("year_ago"), + F.add_months("d", 0).alias("same"), + ).collect() + assert res == [ + Row(next=date(2024, 2, 29), year_ago=date(2023, 1, 31), same=date(2024, 1, 31)), + Row(next=date(2024, 3, 29), year_ago=date(2023, 2, 28), same=date(2024, 2, 29)), + Row(next=date(2023, 2, 28), year_ago=date(2022, 1, 31), same=date(2023, 1, 31)), + Row(next=date(2024, 5, 30), year_ago=date(2023, 4, 30), same=date(2024, 4, 30)), + ] + + def test_week_numbering_at_year_boundaries(self, spark): + dates = [date(2021, 1, 3), date(2020, 12, 31), date(2024, 12, 30), date(2023, 1, 1)] + df = spark.createDataFrame([(d,) for d in dates], ["d"]) + res = df.select( + F.weekofyear("d").alias("week"), + F.dayofweek("d").alias("dow"), + F.weekday("d").alias("weekday"), + F.last_day("d").alias("last"), + ).collect() + assert res == [ + Row(week=53, dow=1, weekday=6, last=date(2021, 1, 31)), + Row(week=53, dow=5, weekday=3, last=date(2020, 12, 31)), + Row(week=1, dow=2, weekday=0, last=date(2024, 12, 31)), + Row(week=52, dow=1, weekday=6, last=date(2023, 1, 31)), + ] + + def test_timestamp_results_have_timestamp_type(self, spark): + schema = StructType([StructField("s", StringType()), StructField("t", TimestampType())]) + df = spark.createDataFrame([("2024-01-02 03:04:05", datetime(2024, 5, 6, 7, 8, 9))], schema) + res = df.select( + F.to_timestamp("s"), + F.to_timestamp("s", "yyyy-MM-dd HH:mm:ss"), + F.try_to_timestamp("s"), + F.try_to_timestamp("s", F.lit("yyyy-MM-dd HH:mm:ss")), + F.to_timestamp_ltz("s"), + F.to_timestamp_ntz("s"), + F.lit(datetime(2024, 1, 2, 3, 4, 5)), + F.date_trunc("month", "t"), + F.date_trunc("year", F.to_date("s")), + F.date_trunc("bogus", "t"), + ) + assert [f.dataType for f in res.schema.fields] == [TimestampType()] * 5 + [TimestampNTZType()] + [ + TimestampType() + ] * 4 + parsed = datetime(2024, 1, 2, 3, 4, 5) + assert res.collect() == [ + Row(*[parsed] * 7, datetime(2024, 5, 1), datetime(2024, 1, 1), None), + ] + + def test_date_parts_are_int(self, spark): + schema = StructType([StructField("t", TimestampType()), StructField("d", DateType())]) + df = spark.createDataFrame([(datetime(2024, 5, 6, 7, 8, 9), date(2024, 5, 6))], schema) + res = df.select( + F.year("d"), + F.quarter("d"), + F.month("d"), + F.day("d"), + F.dayofmonth("d"), + F.dayofweek("d"), + F.dayofyear("d"), + F.weekofyear("d"), + F.weekday("d"), + F.hour("t"), + F.minute("t"), + F.second("t"), + F.date_part(F.lit("YEAR"), "d"), + F.extract(F.lit("DOW"), "d"), + F.datepart(F.lit("HOUR"), "t"), + F.date_diff("d", F.lit(date(2024, 5, 1))), + F.unix_date("d"), + ) + assert [f.dataType for f in res.schema.fields] == [IntegerType()] * 17 + assert res.collect() == [Row(2024, 2, 5, 6, 6, 2, 127, 19, 0, 7, 8, 9, 2024, 2, 7, 5, 19849)] diff --git a/tests/fast/spark/test_spark_functions_datetime_pattern.py b/tests/fast/spark/test_spark_functions_datetime_pattern.py new file mode 100644 index 000000000..5c7742497 --- /dev/null +++ b/tests/fast/spark/test_spark_functions_datetime_pattern.py @@ -0,0 +1,246 @@ +import sys + +import pytest + +_ = pytest.importorskip("duckdb.experimental.spark") + +from datetime import date, datetime, timezone + +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import ( + IllegalArgumentException, + PySparkTypeError, + SparkRuntimeException, + SparkUpgradeException, +) +from spark_namespace.sql import functions as F + +import duckdb +from duckdb.experimental.spark.exception import ContributionsAcceptedError + +if USE_ACTUAL_SPARK: + import pyspark + from py4j.protocol import Py4JJavaError + from pyspark.errors import DateTimeException + + SPARK_3 = pyspark.__version__.startswith("3.") +else: + SPARK_3 = False + +# Spark 4 runs with ANSI mode on, so to_timestamp and to_date raise on text the pattern does not match, as +# DuckDB's strptime does, with a ConversionException for a date that does not exist; Spark 3.5 returns NULL instead. +PARSE_ERROR = (DateTimeException,) if USE_ACTUAL_SPARK else (duckdb.InvalidInputException, duckdb.ConversionException) + +PARSED = [ + ("yyyy-MM-dd", "2024-02-29", datetime(2024, 2, 29)), + ("yyyy-MM-dd HH:mm:ss", "1997-02-28 10:30:00", datetime(1997, 2, 28, 10, 30)), + ("dd/MM/yyyy", "05/01/2024", datetime(2024, 1, 5)), + ("d/M/yyyy", "5/1/2024", datetime(2024, 1, 5)), + pytest.param( + "y-M-d", + "24-1-5", + datetime(24, 1, 5), + marks=pytest.mark.skipif( + sys.platform == "win32", reason="outside the range of the Windows C runtime, as for PySpark on Windows" + ), + ), + pytest.param( + "yyy", + "024", + datetime(24, 1, 1), + marks=pytest.mark.skipif( + sys.platform == "win32", reason="outside the range of the Windows C runtime, as for PySpark on Windows" + ), + ), + ("yyyyMMddHHmmss", "20240105103000", datetime(2024, 1, 5, 10, 30)), + ("dd MMM yyyy", "05 JAN 2024", datetime(2024, 1, 5)), + ("dd LLLL yyyy", "05 january 2024", datetime(2024, 1, 5)), + ("MMMMyyyy", "May2024", datetime(2024, 5, 1)), + ("MMM d, yyyy", "Sep 5, 2024", datetime(2024, 9, 5)), + ("yyyy-MM-dd hh:mm:ss a", "2024-01-05 12:30:00 AM", datetime(2024, 1, 5, 0, 30)), + ("yyyy-MM-dd hh:mm:ss a", "2024-01-05 01:30:00 pm", datetime(2024, 1, 5, 13, 30)), + ("yyyy-MM-dd a", "2024-01-05 PM", datetime(2024, 1, 5, 12)), + ("yyyy-MM-dd HH:mm:ss.SSS", "2024-01-05 10:30:00.1", datetime(2024, 1, 5, 10, 30, 0, 100000)), + ("yyyy-MM-dd HH:mm:ssSSS", "2024-01-05 10:30:00123", datetime(2024, 1, 5, 10, 30, 0, 123000)), + ("yyyy-MM-dd HH:mm:ss.SSSSSS", "2024-01-05 10:30:00.123456", datetime(2024, 1, 5, 10, 30, 0, 123456)), + ("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'", "2024-01-05T10:30:00.123Z", datetime(2024, 1, 5, 10, 30, 0, 123000)), + ("yyyy-MM-dd'''T'''HH", "2024-01-05'T'10", datetime(2024, 1, 5, 10)), + ("yyyy%MM%dd", "2024%01%05", datetime(2024, 1, 5)), + ("yyyy年MM月dd日", "2024年01月05日", datetime(2024, 1, 5)), +] + +# Text that Spark and DuckDB's strptime both reject. The last element is Spark 3.5's answer: NULL, or an upgrade +# error where Spark 2's parser would have accepted the text. +REJECTED = [ + ("dd MMM yyyy", "05 January 2024", "upgrade"), + ("yyyy-MM-dd", "2023-02-29", "null"), + ("yyyy-MM-dd HH:mm:ss.SSS", "2024-01-05 10:30:00.1234", "null"), + ("yyyy-MM-dd hh:mm:ss a", "2024-01-05 13:30:00 PM", "null"), + ("yyyy-MM-dd'T'HH:mm:ss", "2024-01-05 10:30:00", "null"), +] + +# Text that Spark rejects and DuckDB's strptime reads: it allows fewer digits than a field's width, a year of up to +# four digits for yyyy, up to three or six fraction digits for fewer S letters, any run of spaces, tabs included, for +# a space, spaces around the text, and the words epoch, infinity and -infinity. The third element is Spark 3.5's +# answer as in REJECTED, the last DuckDB's. +ACCEPTED_BY_DUCKDB = [ + ("yyyy-MM-dd HH:mm:ss", "1997-2-28 10:30:00", "upgrade", datetime(1997, 2, 28, 10, 30)), + pytest.param( + "yyyy-MM-dd HH:mm:ss", + "97-02-28 10:30:00", + "upgrade", + datetime(97, 2, 28, 10, 30), + marks=pytest.mark.skipif( + sys.platform == "win32", reason="outside the range of the Windows C runtime, as for PySpark on Windows" + ), + ), + ("yyyy-MM-dd HH:mm:ss", "1997-02-28 10:30:00", "upgrade", datetime(1997, 2, 28, 10, 30)), + ("yyyy-MM-dd", " 2024-01-05", "upgrade", datetime(2024, 1, 5)), + ("yyyy-MM-dd", "2024-01-05 ", "upgrade", datetime(2024, 1, 5)), + ("yyyy MM", "2024\t01", "null", datetime(2024, 1, 1)), + ("yyyy-MM-dd", "epoch", "null", datetime(1970, 1, 1)), + ("yyyy-MM-dd", "infinity", "null", datetime(9999, 12, 31, 23, 59, 59, 999999)), + ("yyyy-MM-dd", "-infinity", "null", datetime(1, 1, 1)), + ("yyyy-MM-dd HH:mm:ss.S", "2024-01-05 10:30:00.12", "upgrade", datetime(2024, 1, 5, 10, 30, 0, 120000)), +] + +# Text that Spark reads and DuckDB's strptime rejects: leading zeros beyond a field's width, and literal letters in +# another case. +REJECTED_BY_DUCKDB = [ + ("d/M/yyyy", "005/01/2024", datetime(2024, 1, 5)), + ("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'", "2024-01-05t10:30:00.123z", datetime(2024, 1, 5, 10, 30, 0, 123000)), +] + + +def _first(df, column): + return df.select(column.alias("r")).collect()[0].r + + +def _assert_rejected(df, pattern, spark_3): + if SPARK_3 and spark_3 == "upgrade": + for column in [F.to_timestamp("s", pattern), F.try_to_timestamp("s", F.lit(pattern))]: + with pytest.raises((Py4JJavaError, SparkUpgradeException), match="INCONSISTENT_BEHAVIOR_CROSS_VERSION"): + _first(df, column) + return + assert _first(df, F.try_to_timestamp("s", F.lit(pattern))) is None + for column in [F.to_timestamp("s", pattern), F.to_date("s", pattern)]: + if SPARK_3: + assert _first(df, column) is None + else: + with pytest.raises(PARSE_ERROR): + _first(df, column) + + +class TestSparkDatetimePattern: + @pytest.mark.parametrize(("pattern", "text", "expected"), PARSED) + def test_parse(self, spark, pattern, text, expected): + df = spark.createDataFrame([(text,)], ["s"]) + assert _first(df, F.to_timestamp("s", pattern)) == expected + assert _first(df, F.try_to_timestamp("s", F.lit(pattern))) == expected + assert _first(df, F.to_timestamp_ntz("s", F.lit(pattern))) == expected + assert _first(df, F.to_date("s", pattern)) == expected.date() + + @pytest.mark.parametrize(("pattern", "text", "spark_3"), REJECTED) + def test_parse_rejects_what_spark_rejects(self, spark, pattern, text, spark_3): + _assert_rejected(spark.createDataFrame([(text,)], ["s"]), pattern, spark_3) + + @pytest.mark.parametrize(("pattern", "text", "spark_3", "duckdb_value"), ACCEPTED_BY_DUCKDB) + def test_parse_accepts_more_than_spark(self, spark, pattern, text, spark_3, duckdb_value): + df = spark.createDataFrame([(text,)], ["s"]) + if USE_ACTUAL_SPARK: + _assert_rejected(df, pattern, spark_3) + return + assert _first(df, F.to_timestamp("s", pattern)) == duckdb_value + assert _first(df, F.try_to_timestamp("s", F.lit(pattern))) == duckdb_value + assert _first(df, F.to_date("s", pattern)) == duckdb_value.date() + + @pytest.mark.parametrize(("pattern", "text", "spark_value"), REJECTED_BY_DUCKDB) + def test_parse_rejects_more_than_spark(self, spark, pattern, text, spark_value): + df = spark.createDataFrame([(text,)], ["s"]) + if USE_ACTUAL_SPARK: + assert _first(df, F.to_timestamp("s", pattern)) == spark_value + assert _first(df, F.try_to_timestamp("s", F.lit(pattern))) == spark_value + return + assert _first(df, F.try_to_timestamp("s", F.lit(pattern))) is None + with pytest.raises(PARSE_ERROR): + _first(df, F.to_timestamp("s", pattern)) + + def test_parse_null(self, spark): + df = spark.createDataFrame([("2024-01-05",), (None,)], ["s"]) + res = df.select(F.to_timestamp("s", "yyyy-MM-dd").alias("r")).collect() + assert [r.r for r in res] == [datetime(2024, 1, 5), None] + + @pytest.mark.parametrize( + ("pattern", "text", "expected"), + [ + ("yyyy-DDD", "2024-005", datetime(2024, 1, 5)), + # Spark shifts the instant to the session time zone, which defaults to the local one. + ( + "yyyy-MM-dd HH:mm:ssXXX", + "2024-01-05 10:30:00Z", + datetime(2024, 1, 5, 10, 30, tzinfo=timezone.utc).astimezone().replace(tzinfo=None), + ), + ("yyyy-MM-dd[ HH:mm]", "2024-01-05", datetime(2024, 1, 5)), + ("G yyyy-MM-dd", "AD 2024-01-05", datetime(2024, 1, 5)), + ("yyyy-MM-dd kk:mm", "2024-01-05 24:30", datetime(2024, 1, 5, 0, 30)), + ("yyyy-MM-dd HH:mm a", "2024-01-05 13:30 PM", datetime(2024, 1, 5, 13, 30)), + ("Hmm", "930", datetime(1970, 1, 1, 9, 30)), + # DuckDB reads the year 1900 where a pattern has none, Spark 1970. + ("HH:mm:ss", "10:30:00", datetime(1970, 1, 1, 10, 30)), + ("H", "007", datetime(1970, 1, 1, 7)), + # DuckDB reads yy as 1969 to 2068, Spark as 2000 to 2099. + ("yy-MM-dd", "99-01-05", datetime(2099, 1, 5)), + # DuckDB reads at most four year digits. + ("yyyyyy", "002024", datetime(2024, 1, 1)), + # DuckDB reads 12 as noon without an am-pm marker, Spark as midnight. + ("yyyy-MM-dd hh:mm", "2024-01-05 12:30", datetime(2024, 1, 5, 0, 30)), + # DuckDB rounds nanoseconds to microseconds, Spark truncates them. + ("yyyy-MM-dd HH:mm:ss.SSSSSSSSS", "2024-01-05 10:30:00.123456789", datetime(2024, 1, 5, 10, 30, 0, 123456)), + ("yyyy-MM-dd HH:mm:ss.SSSSSSS", "2024-01-05 10:30:00.9999999", datetime(2024, 1, 5, 10, 30, 0, 999999)), + ], + ) + def test_parse_unsupported_pattern(self, spark, pattern, text, expected): + df = spark.createDataFrame([(text,)], ["s"]) + if USE_ACTUAL_SPARK: + assert _first(df, F.to_timestamp("s", pattern)) == expected + else: + with pytest.raises(ContributionsAcceptedError, match="not supported for parsing"): + F.to_timestamp("s", pattern) + + @pytest.mark.parametrize( + "pattern", ["yyyy Q", "EEE yyyy-MM-dd", "YYYY-MM-dd", "yyyy-ww", "yyyy-MM-dd'T", "yyyy#MM"] + ) + def test_parse_invalid_pattern(self, spark, pattern): + df = spark.createDataFrame([("2024-01-05",)], ["s"]) + with pytest.raises((IllegalArgumentException, SparkUpgradeException, SparkRuntimeException)): + _first(df, F.to_timestamp("s", pattern)) + + def test_try_to_timestamp_format_column(self, spark): + df = spark.createDataFrame([("2024-01-05", "yyyy-MM-dd"), ("05/01/2024", "dd/MM/yyyy")], ["s", "f"]) + if USE_ACTUAL_SPARK: + res = df.select(F.try_to_timestamp("s", F.col("f")).alias("r")).collect() + assert [r.r for r in res] == [datetime(2024, 1, 5), datetime(2024, 1, 5)] + else: + with pytest.raises(ContributionsAcceptedError, match="literal datetime pattern"): + F.try_to_timestamp("s", F.col("f")) + + def test_to_timestamp_ltz_and_ntz_with_format(self, spark): + df = spark.createDataFrame([("2016-12-31",), ("bad",)], ["s"]) + good, bad = df.filter("s = '2016-12-31'"), df.filter("s = 'bad'") + for function in [F.to_timestamp_ltz, F.to_timestamp_ntz]: + assert _first(good, function("s", F.lit("yyyy-MM-dd"))) == datetime(2016, 12, 31) + if SPARK_3: + assert _first(bad, function("s", F.lit("yyyy-MM-dd"))) is None + else: + with pytest.raises(PARSE_ERROR): + _first(bad, function("s", F.lit("yyyy-MM-dd"))) + + def test_to_date_with_format_keeps_date_only(self, spark): + df = spark.createDataFrame([("1997-02-28 10:30:00",)], ["t"]) + assert _first(df, F.to_date("t", "yyyy-MM-dd HH:mm:ss")) == date(1997, 2, 28) + + def test_to_date_format_must_be_str(self, spark): + df = spark.createDataFrame([("2024-01-05",)], ["s"]) + for function in [F.to_date, F.to_timestamp]: + with pytest.raises(PySparkTypeError): + _first(df, function("s", F.lit("yyyy-MM-dd"))) diff --git a/tests/fast/spark/test_spark_functions_numeric.py b/tests/fast/spark/test_spark_functions_numeric.py index 7408505b8..0d6af1f0b 100644 --- a/tests/fast/spark/test_spark_functions_numeric.py +++ b/tests/fast/spark/test_spark_functions_numeric.py @@ -4,11 +4,42 @@ import math +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import PySparkTypeError, SparkRuntimeException from spark_namespace.sql import functions as sf -from spark_namespace.sql.types import Row +from spark_namespace.sql.types import ( + ArrayType, + BooleanType, + ByteType, + DoubleType, + FloatType, + IntegerType, + LongType, + Row, + ShortType, + StringType, + StructField, + StructType, +) import numpy as np +if USE_ACTUAL_SPARK: + import pyspark + +NAN = float("nan") + + +def _ansi(spark): + return USE_ACTUAL_SPARK and spark.conf.get("spark.sql.ansi.enabled") == "true" + + +def _nan_frame(spark): + """Rows k=1 NaN, k=2 1.0, k=3 NULL; the NaN comes from a literal so it does not depend on createDataFrame.""" + schema = StructType([StructField("k", LongType()), StructField("x", DoubleType())]) + df = spark.createDataFrame([(1, 0.0), (2, 1.0), (3, None)], schema) + return df.withColumn("x", sf.when(sf.col("k") == 1, sf.lit(NAN)).otherwise(sf.col("x"))) + class TestSparkFunctionsNumeric: def test_greatest(self, spark): @@ -63,6 +94,18 @@ def test_floor(self, spark): Row(floor_value=2), ] + def test_ceil_and_floor_of_integral_types_are_exact(self, spark): + df = spark.createDataFrame( + [(9007199254740993, -9223372036854775808, 7, 5), (None, None, None, None)], + "b bigint, m bigint, i int, t tinyint", + ) + res = df.select(sf.ceil("b"), sf.floor("b"), sf.ceil("m"), sf.floor("m"), sf.ceil("i"), sf.floor("t")) + assert [f.dataType for f in res.schema.fields] == [LongType()] * 6 + assert [tuple(r) for r in res.collect()] == [ + (9007199254740993, 9007199254740993, -9223372036854775808, -9223372036854775808, 7, 5), + (None,) * 6, + ] + def test_abs(self, spark): data = [ (1.1,), @@ -353,4 +396,242 @@ def test_negative(self, spark): res = df.collect() assert res[0].value == 0 assert res[1].value == -2 - assert res[2].value == -3 + assert res[2].value == 3 + + def test_negative_and_positive(self, spark): + schema = StructType([StructField("i", LongType()), StructField("d", DoubleType())]) + df = spark.createDataFrame([(2, 2.5), (-3, -3.5), (None, None)], schema) + res = df.select( + sf.negative("i").alias("ni"), + sf.negative(sf.col("d")).alias("nd"), + sf.positive("i").alias("pi"), + (-df.i).alias("op"), + ).collect() + assert res == [ + Row(ni=-2, nd=-2.5, pi=2, op=-2), + Row(ni=3, nd=3.5, pi=-3, op=3), + Row(ni=None, nd=None, pi=None, op=None), + ] + + def test_nan_literal(self, spark): + res = spark.range(1).select( + sf.lit(NAN).alias("n"), + sf.isnan(sf.lit(NAN)).alias("isnan"), + (sf.lit(NAN) == sf.lit(NAN)).alias("eq"), + (sf.lit(NAN) > sf.lit(1e308)).alias("gt"), + sf.when(sf.lit(False), 1.0).otherwise(NAN).alias("otherwise"), + ) + assert res.schema.fields[0].dataType == DoubleType() + row = res.collect()[0] + assert math.isnan(row.n) + assert math.isnan(row.otherwise) + assert (row.isnan, row.eq, row.gt) == (True, True, True) + + def test_nan_comparisons(self, spark): + df = _nan_frame(spark) + assert df.filter(sf.col("x") == NAN).select("k").collect() == [Row(k=1)] + assert df.filter(sf.col("x") != NAN).select("k").collect() == [Row(k=2)] + assert df.filter(sf.col("x").isin(NAN, 5.0)).select("k").collect() == [Row(k=1)] + res = df.select( + "k", + (sf.col("x") > 1e308).alias("gt"), + (sf.col("x") < NAN).alias("lt"), + (sf.col("x") >= NAN).alias("ge"), + sf.isnan("x").alias("isnan"), + ).orderBy("k") + assert res.collect() == [ + Row(k=1, gt=True, lt=False, ge=True, isnan=True), + Row(k=2, gt=False, lt=True, ge=False, isnan=False), + Row(k=3, gt=None, lt=None, ge=None, isnan=False), + ] + + def test_nan_is_largest(self, spark): + df = _nan_frame(spark) + assert [r.k for r in df.orderBy("x").collect()] == [3, 2, 1] + assert [r.k for r in df.orderBy(sf.desc("x")).collect()] == [1, 2, 3] + row = df.select(sf.max("x").alias("max"), sf.min("x").alias("min")).collect()[0] + assert math.isnan(row.max) + assert row.min == 1.0 + greatest = df.select("k", sf.greatest("x", sf.lit(5.0)).alias("g"), sf.least("x", sf.lit(5.0)).alias("l")) + rows = greatest.orderBy("k").collect() + assert math.isnan(rows[0].g) + assert rows[0].l == 5.0 + assert (rows[1].g, rows[1].l) == (5.0, 1.0) + + def test_isnan_of_null_is_false(self, spark): + df = spark.createDataFrame([(None,), (1.0,)], StructType([StructField("x", DoubleType())])) + assert df.select(sf.isnan("x").alias("v")).collect() == [Row(v=False), Row(v=False)] + + def test_division_by_zero(self, spark): + df = spark.createDataFrame( + [(7, 2.0)], StructType([StructField("a", LongType()), StructField("b", DoubleType())]) + ) + assert df.select((df.a / 2).alias("q"), (7 / df.b).alias("r")).collect() == [Row(q=3.5, r=3.5)] + if _ansi(spark): + with pytest.raises(Exception, match="DIVIDE_BY_ZERO"): + df.select(df.a / 0).collect() + else: + res = df.select((df.a / 0).alias("q"), (df.b / 0).alias("d"), (1 / (df.a - 7)).alias("r")).collect() + assert res == [Row(q=None, d=None, r=None)] + + def test_result_types(self, spark): + schema = StructType([StructField("d", DoubleType()), StructField("i", IntegerType())]) + df = spark.createDataFrame([(2.5, -3)], schema) + res = df.select( + sf.ceil("d").alias("c"), + sf.floor("d").alias("f"), + sf.ceil("i").alias("ci"), + sf.sign("i").alias("s"), + sf.signum("d").alias("sd"), + ) + assert [f.dataType for f in res.schema.fields] == [ + LongType(), + LongType(), + LongType(), + DoubleType(), + DoubleType(), + ] + assert res.collect() == [Row(c=3, f=2, ci=-3, s=-1.0, sd=1.0)] + + def test_logarithm_of_non_positive_is_null(self, spark): + df = spark.createDataFrame([(8.0,), (0.0,), (-1.0,), (None,)], StructType([StructField("x", DoubleType())])) + res = df.select(sf.log2("x").alias("l2"), sf.ln("x").alias("ln")).collect() + assert res[0].l2 == 3.0 + assert res[0].ln == pytest.approx(math.log(8)) + assert res[1:] == [Row(l2=None, ln=None)] * 3 + + def test_pow_with_numbers(self, spark): + df = spark.createDataFrame([(2, 3)], StructType([StructField("a", LongType()), StructField("b", LongType())])) + res = df.select(sf.pow("a", 2).alias("p"), sf.pow(2.0, "b").alias("q"), sf.pow(df.a, df.b).alias("r")).collect() + assert res == [Row(p=4.0, q=8.0, r=8.0)] + + def test_factorial_range(self, spark): + df = spark.createDataFrame( + [(5,), (0,), (20,), (21,), (-1,), (None,)], StructType([StructField("n", LongType())]) + ) + res = df.select(sf.factorial("n").alias("f")) + assert res.schema.fields[0].dataType == LongType() + assert [r.f for r in res.collect()] == [120, 1, 2432902008176640000, None, None, None] + + def test_numpy_scalars(self, spark): + scalars = { + "i8": (np.int8(1), ByteType()), + "i16": (np.int16(1), ShortType()), + "i32": (np.int32(1), IntegerType()), + "i64": (np.int64(1), LongType()), + "f32": (np.float32(1.5), FloatType()), + "f64": (np.float64(1.5), DoubleType()), + "b": (np.bool_(True), BooleanType()), + } + res = spark.range(1).select(*[sf.lit(v).alias(name) for name, (v, _) in scalars.items()]) + assert [f.dataType for f in res.schema.fields] == [t for _, t in scalars.values()] + assert res.collect() == [Row(i8=1, i16=1, i32=1, i64=1, f32=1.5, f64=1.5, b=True)] + df = spark.createDataFrame([([1, 2, 3],), ([],)], StructType([StructField("data", ArrayType(LongType()))])) + res = df.select( + sf.array_contains(df.data, np.int64(1)).alias("b"), (sf.array_size(df.data) == np.int32(3)).alias("c") + ) + assert res.collect() == [Row(b=True, c=True), Row(b=False, c=False)] + + def test_factorial_truncates_before_the_range_check(self, spark): + schema = StructType([StructField(c, DoubleType()) for c in "abcde"]) + df = spark.createDataFrame([(5.5, 20.5, -0.5, 21.0, None)], schema) + res = df.select(*(sf.factorial(c) for c in "abcde"), sf.factorial(sf.lit(5.9))) + assert res.columns == [*(f"factorial({c})" for c in "abcde"), "factorial(5.9)"] + assert [f.dataType for f in res.schema.fields] == [LongType()] * 6 + assert res.collect() == [Row(120, 2432902008176640000, 1, None, None, 120)] + + def test_sign_of_nan(self, spark): + rows = _nan_frame(spark).select(sf.sign("x").alias("s"), sf.signum("x").alias("t")).collect() + assert math.isnan(rows[0].s) + assert math.isnan(rows[0].t) + assert rows[1:] == [Row(s=1.0, t=1.0), Row(s=None, t=None)] + + def test_lit_list_keeps_nan(self, spark): + res = spark.range(1).select(sf.lit([NAN, 1.0]).alias("a"), sf.lit([[NAN], [2.0, None]]).alias("n")) + assert res.schema.simpleString() == "struct,n:array>>" + row = res.collect()[0] + assert math.isnan(row.a[0]) + assert row.a[1:] == [1.0] + assert math.isnan(row.n[0][0]) + assert row.n[1] == [2.0, None] + + def test_numpy_scalars_as_operands(self, spark): + df = spark.createDataFrame([(10,)], StructType([StructField("a", LongType())])) + res = df.select(df.a * np.float64(1e20), df.a + np.float64(1.5), df.a + np.float32(0.1), df.a + np.int32(3)) + assert res.columns == ["(a * 1.0E20)", "(a + 1.5)", "(a + 0.10000000149011612)", "(a + 3)"] + assert [f.dataType for f in res.schema.fields] == [DoubleType(), DoubleType(), DoubleType(), LongType()] + assert res.collect() == [Row(1e21, 11.5, 10.100000001490116, 13)] + + def test_numpy_scalar_literal_names(self, spark): + res = spark.range(1).select( + sf.lit(np.float64(0.1)), sf.lit(np.float32(0.1)), sf.lit(np.int8(3)), sf.lit(np.bool_(False)) + ) + assert res.columns == ["0.1", "0.1", "3", "false"] + assert [f.dataType for f in res.schema.fields] == [DoubleType(), FloatType(), ByteType(), BooleanType()] + + def test_numpy_scalar_literal_names_inside_expressions(self, spark): + # Spark casts the Python value of a NumPy number to its type and aliases the cast to the number's str(). + literals = [np.int8(5), np.int16(-5), np.int32(5), np.int64(5), np.float32(1.5), np.float64(np.inf)] + res = spark.range(1).select(*(sf.lit(v).isNull() for v in literals), sf.lit(np.int8(5)).alias("q").isNull()) + assert res.columns == [ + "(CAST(5 AS TINYINT) AS `5` IS NULL)", + "(CAST(-5 AS SMALLINT) AS `-5` IS NULL)", + "(CAST(5 AS INT) AS `5` IS NULL)", + "(CAST(5 AS BIGINT) AS `5` IS NULL)", + "(CAST(1.5 AS FLOAT) AS `1.5` IS NULL)", + "(CAST(Infinity AS DOUBLE) AS inf IS NULL)", + "(CAST(5 AS TINYINT) AS `5` AS q IS NULL)", + ] + assert spark.range(1).select(sf.struct(sf.lit(np.int8(5)))).schema.fields[0].dataType.fieldNames() == ["5"] + + @pytest.mark.parametrize( + ("array", "spark_name", "module_name", "element_type", "value"), + [ + ( + np.array([1.5, 0.1, 1e20, -0.0]), + "ARRAY(1.5D, 0.1D, 1.0E20D, -0.0D)", + "array(1.5, 0.1, 1.0E20, -0.0)", + DoubleType(), + [1.5, 0.1, 1e20, -0.0], + ), + (np.array([np.inf]), "ARRAY(CAST('Infinity' AS DOUBLE))", "array(Infinity)", DoubleType(), [np.inf]), + (np.array([1.5], dtype=np.float32), "ARRAY(CAST('1.5' AS FLOAT))", "array(1.5)", FloatType(), [1.5]), + (np.array([5, -3]), "ARRAY(5L, -3L)", "array(5, -3)", LongType(), [5, -3]), + (np.array([5], dtype=np.int32), "ARRAY(5)", "array(5)", IntegerType(), [5]), + (np.array([5], dtype=np.int16), "ARRAY(5S)", "array(5)", ShortType(), [5]), + (np.array([False, True]), "ARRAY(false, true)", "array(false, true)", BooleanType(), [False, True]), + (np.array(["a'b", "c"]), "ARRAY('a\\'b', 'c')", "array(a'b, c)", StringType(), ["a'b", "c"]), + (np.array([]), "ARRAY()", "array()", DoubleType(), []), + ], + ) + def test_numpy_array_literals(self, spark, array, spark_name, module_name, element_type, value): + res = spark.range(1).select(sf.lit(array)) + # The module names a NumPy array as the list of its values, where Spark spells each element's type. + assert res.columns == [spark_name if USE_ACTUAL_SPARK else module_name] + assert res.schema.fields[0].dataType == ArrayType(element_type) + assert res.collect() == [Row(value)] + + def test_numpy_int8_array_literal(self, spark): + res = spark.range(1).select(sf.lit(np.array([5], dtype=np.int8))) + assert res.columns == ["ARRAY(5S)" if USE_ACTUAL_SPARK else "array(5)"] + # PySpark 3.5 gives the bytes a short type. + expected = ShortType() if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3.") else ByteType() + assert res.schema.fields[0].dataType == ArrayType(expected) + + @pytest.mark.parametrize("value", [np.array([1], dtype=np.uint16), np.array([1, "a"], dtype=object)]) + def test_unsupported_numpy_arrays_raise(self, spark, value): + with pytest.raises(PySparkTypeError, match="UNSUPPORTED_NUMPY_ARRAY_SCALAR"): + spark.range(1).select(sf.lit(value)).collect() + + def test_unsupported_numpy_scalar_raises(self, spark): + # PySpark 3.5 passes the date on to Py4J, which fails on it. + expected = AttributeError if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3.") else PySparkTypeError + with pytest.raises(expected): + spark.range(1).select(sf.lit(np.datetime64("2020-01-01"))).collect() + + def test_two_dimensional_numpy_literal_raises(self, spark): + with pytest.raises(SparkRuntimeException, match=r"UNSUPPORTED_FEATURE\.LITERAL_TYPE"): + spark.range(1).select(sf.lit(np.array([[1, 2]]))).collect() + + def test_smallest_double_literal_name(self, spark): + assert spark.range(1).select(sf.lit(5e-324), sf.lit(2.5e-5)).columns == ["4.9E-324", "2.5E-5"] diff --git a/tests/fast/spark/test_spark_functions_string.py b/tests/fast/spark/test_spark_functions_string.py index ba2a540d0..948f25cab 100644 --- a/tests/fast/spark/test_spark_functions_string.py +++ b/tests/fast/spark/test_spark_functions_string.py @@ -2,9 +2,36 @@ _ = pytest.importorskip("duckdb.experimental.spark") +from decimal import Decimal + from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import PySparkTypeError from spark_namespace.sql import functions as F -from spark_namespace.sql.types import Row +from spark_namespace.sql.types import ( + ArrayType, + BinaryType, + DecimalType, + DoubleType, + IntegerType, + LongType, + Row, + StringType, + StructField, + StructType, +) + +from duckdb.experimental.spark.exception import ContributionsAcceptedError + +if USE_ACTUAL_SPARK: + import pyspark + + SPARK_3 = pyspark.__version__.startswith("3.") +else: + SPARK_3 = False + + +def _strings(spark, *columns): + return StructType([StructField(c, StringType()) for c in columns]) class TestSparkFunctionsString: @@ -355,6 +382,30 @@ def test_right(self, spark): res = df.select(F.right(df.a, df.b).alias("r")).collect() assert res == [Row(r="SQL"), Row(r=""), Row(r="")] + def test_left_and_right_of_null(self, spark): + df = spark.createDataFrame([(None, -1), (None, 0), (None, 2), ("abc", None)], "s string, n int") + res = df.select(F.left("s", "n").alias("l"), F.right("s", "n").alias("r")).collect() + assert res == [Row(l=None, r=None)] * 4 + + def test_left_and_right_at_the_integer_bounds(self, spark): + rows = [(None, -2147483648), ("abc", -2147483648), (None, 2147483647), ("abc", 2147483647)] + df = spark.createDataFrame(rows, "s string, n int") + res = df.select(F.left("s", "n").alias("l"), F.right("s", "n").alias("r")).collect() + assert res == [Row(l=None, r=None), Row(l="", r=""), Row(l=None, r=None), Row(l="abc", r="abc")] + + def test_left_and_right_of_a_fractional_length(self, spark): + # Spark casts the length to an integer, which truncates toward zero. + df = spark.createDataFrame( + [("abcde", 1.5, Decimal("2.9")), ("abcde", -0.5, Decimal("0.4"))], "s string, n double, d decimal(3,1)" + ) + res = df.select( + F.left("s", "n").alias("ln"), + F.right("s", "n").alias("rn"), + F.left("s", "d").alias("ld"), + F.right("s", "d").alias("rd"), + ).collect() + assert res == [Row(ln="a", rn="e", ld="ab", rd="de"), Row(ln="", rn="", ld="", rd="")] + def test_levenshtein(self, spark): df = spark.createDataFrame([("kitten", "sitting"), ("saturdays", "sunday")], ["a", "b"]) @@ -447,6 +498,14 @@ def test_regexp_extract_all(self, spark): res = df.select(F.regexp_extract_all("str", F.col("regexp")).alias("d")).collect() assert res == [Row(d=["100", "300"])] + def test_regexp_extract_all_of_a_group_outside_the_match(self, spark): + df = spark.createDataFrame([("a b",), ("zz",), (None,)], "s string") + res = df.select( + F.regexp_extract_all("s", F.lit("(a)|(b)"), 1).alias("first"), + F.regexp_extract_all("s", F.lit("(a)|(b)"), 2).alias("second"), + ) + assert res.collect() == [Row(first=["a", ""], second=["", "b"]), Row(first=[], second=[]), Row(None, None)] + def test_regexp_substr(self, spark): df = spark.createDataFrame([("1a 2b 14m", r"\d+")], ["str", "regexp"]) @@ -526,3 +585,312 @@ def test_instr(self, spark): Row(instr=6), Row(instr=4), ] + + def test_concat_with_null_is_null(self, spark): + df = spark.createDataFrame([("a", None), ("a", "b")], _strings(spark, "a", "b")) + assert df.select(F.concat("a", "b").alias("c")).collect() == [Row(c=None), Row(c="ab")] + arrays = StructType([StructField("x", ArrayType(LongType())), StructField("y", ArrayType(LongType()))]) + adf = spark.createDataFrame([([1, 2], [3]), ([1], None)], arrays) + assert adf.select(F.concat("x", "y").alias("c")).collect() == [Row(c=[1, 2, 3]), Row(c=None)] + assert spark.range(1).select(F.concat().alias("c")).collect() == [Row(c="")] + + def test_concat_keeps_binary_and_arrays(self, spark): + schema = "a binary, b binary, s string, l array, m array, i int" + df = spark.createDataFrame([(b"\x00A", b"B", "x", [1, 2], [3], 7), (None, b"C", None, None, [4], None)], schema) + res = df.select( + F.concat("a", "b", "b").alias("binary"), + F.concat("l", "m").alias("arrays"), + F.concat("b", "s").alias("mixed"), + F.concat("i", "s").alias("number"), + ) + assert res.schema.simpleString() == "struct,mixed:string,number:string>" + assert res.collect() == [ + Row(binary=b"\x00ABB", arrays=[1, 2, 3], mixed="Bx", number="7x"), + Row(binary=None, arrays=None, mixed=None, number=None), + ] + + def test_concat_of_one_column_has_the_type_of_a_concatenation(self, spark): + schema = "b binary, l array, i int, s string" + df = spark.createDataFrame([(b"\x00A", [1], 7, "x"), (None, None, None, None)], schema) + res = df.select(*(F.concat(c).alias(c) for c in ["b", "l", "i", "s"])) + assert res.schema.simpleString() == "struct,i:string,s:string>" + assert res.collect() == [Row(b=b"\x00A", l=[1], i="7", s="x"), Row(b=None, l=None, i=None, s=None)] + + def test_concat_reads_a_null_literal_as_a_string(self, spark): + df = spark.createDataFrame([(b"B", "x")], "b binary, s string") + res = df.select( + F.concat("s", F.lit(None)).alias("s"), + F.concat("b", F.lit(None)).alias("b"), + F.concat(F.lit(None)).alias("n"), + ) + assert res.schema.simpleString() == "struct" + assert res.collect() == [Row(s=None, b=None, n=None)] + + def test_concat_ws_skips_nulls_and_flattens_arrays(self, spark): + schema = StructType( + [ + StructField("arr", ArrayType(StringType())), + StructField("s", StringType()), + StructField("n", StringType()), + ] + ) + df = spark.createDataFrame([(["a", None, "b"], "z", None), ([], "z", None), (None, None, None)], schema) + res = df.select(F.concat_ws(",", "arr").alias("a"), F.concat_ws("-", "arr", "s", "n").alias("m")).collect() + assert res == [Row(a="a,b", m="a-b-z"), Row(a="", m="z"), Row(a="", m="")] + + def test_substring_positions(self, spark): + df = spark.createDataFrame([("hello",), (None,)], _strings(spark, "s")) + res = df.select( + F.substring("s", 0, 2).alias("zero"), + F.substring("s", 1, 2).alias("one"), + F.substring("s", -3, 2).alias("neg"), + F.substring("s", -7, 4).alias("before"), + F.substring("s", -10, 3).alias("far"), + F.substring("s", 2, -1).alias("neglen"), + F.substring("s", 3, 100).alias("long"), + ).collect() + assert res == [ + Row(zero="he", one="he", neg="ll", before="he", far="", neglen="", long="llo"), + Row(zero=None, one=None, neg=None, before=None, far=None, neglen=None, long=None), + ] + + def test_substr_with_columns(self, spark): + schema = StructType( + [StructField("s", StringType()), StructField("p", LongType()), StructField("n", LongType())] + ) + df = spark.createDataFrame([("hello", 0, 2), ("hello", -3, 2), ("hello", 2, -1)], schema) + res = df.select(F.substr("s", "p", "n").alias("a"), F.substr(df.s, df.p).alias("b")).collect() + assert res == [Row(a="he", b="hello"), Row(a="ll", b="llo"), Row(a="", b="ello")] + + def test_initcap_lowers_the_rest(self, spark): + df = spark.createDataFrame([("hELLO wORLD",), ("ÉCOLE",)], _strings(spark, "s")) + assert df.select(F.initcap("s").alias("s")).collect() == [Row(s="Hello World"), Row(s="École")] + + def test_pad_with_empty_padding(self, spark): + df = spark.createDataFrame([("spark",), (None,)], _strings(spark, "s")) + res = df.select( + F.lpad("s", 8, "").alias("a"), + F.rpad("s", 3, "").alias("b"), + F.lpad("s", 7, "ab").alias("c"), + F.lpad("s", 0, "").alias("d"), + ).collect() + assert res == [Row(a="spark", b="spa", c="abspark", d=""), Row(a=None, b=None, c=None, d=None)] + + def test_levenshtein_threshold_null(self, spark): + df = spark.createDataFrame([("kitten", "sitting"), (None, "a")], _strings(spark, "a", "b")) + res = df.select( + F.levenshtein("a", "b").alias("d"), + F.levenshtein("a", "b", 2).alias("t2"), + F.levenshtein("a", "b", 3).alias("t3"), + ) + assert res.collect() == [Row(d=3, t2=-1, t3=3), Row(d=None, t2=None, t3=None)] + + def test_split_part_edges(self, spark): + df = spark.createDataFrame([("a,b,c",)], _strings(spark, "s")) + parts = {"first": 1, "second": 2, "last": -1, "second_last": -2} + res = df.select(*[F.split_part("s", F.lit(""), F.lit(i)).alias(name) for name, i in parts.items()]).collect() + assert res == [Row(first="a,b,c", second="", last="a,b,c", second_last="")] + with pytest.raises(Exception, match="INVALID_INDEX_OF_ZERO"): + df.select(F.split_part("s", F.lit(","), F.lit(0))).collect() + + def test_char_wraps_and_handles_negatives(self, spark): + df = spark.createDataFrame( + [(65,), (256,), (512 + 66,), (-1,), (None,)], StructType([StructField("i", LongType())]) + ) + assert df.select(F.char("i").alias("c")).collect() == [ + Row(c="A"), + Row(c="\x00"), + Row(c="B"), + Row(c=""), + Row(c=None), + ] + + def test_sha2_zero_bits_is_sha256(self, spark): + df = spark.createDataFrame([("ab",)], _strings(spark, "s")) + digest = "fb8e20fc2e4c3f248c60c39bd652f3c1347298bb977b8b4d5903b85055620603" + assert df.select(F.sha2("s", 0).alias("a"), F.sha2("s", 256).alias("b")).collect() == [Row(a=digest, b=digest)] + + def test_integer_result_types(self, spark): + df = spark.createDataFrame([("hello", "help")], _strings(spark, "a", "b")) + res = df.select(F.length("a"), F.instr("a", "l"), F.levenshtein("a", "b"), F.regexp_count("a", F.lit("l"))) + assert [f.dataType for f in res.schema.fields] == [IntegerType()] * 4 + assert res.collect() == [Row(5, 3, 2, 2)] + + def test_regexp_replace_java_replacement(self, spark): + df = spark.createDataFrame([("100-200",)], _strings(spark, "s")) + res = df.select( + F.regexp_replace("s", r"(\d+)-(\d+)", "$2-$1").alias("swap"), + F.regexp_replace("s", "-", r"\$").alias("dollar"), + F.regexp_replace("s", "-", r"\\").alias("backslash"), + F.regexp_replace("s", r"(\d+)", "<$0>").alias("whole"), + ).collect() + assert res == [Row(swap="200-100", dollar="100$200", backslash="100\\200", whole="<100>-<200>")] + + def test_regexp_replace_of_an_empty_match_after_a_match(self, spark): + df = spark.createDataFrame([("aa",), ("ba",), ("abc",)], _strings(spark, "s")) + res = df.select(F.regexp_replace("s", "a*", "x").alias("a"), F.regexp_replace("s", "b*", "+").alias("b")) + if USE_ACTUAL_SPARK: + assert res.collect() == [Row(a="xx", b="+a+a+"), Row(a="xbxx", b="++a+"), Row(a="xxbxcx", b="+a++c+")] + else: + # RE2 skips an empty match that starts where the previous match ended, where Java's matcher takes it. + assert res.collect() == [Row(a="x", b="+a+a+"), Row(a="xbx", b="+a+"), Row(a="xbxcx", b="+a+c+")] + + def test_regexp_replace_reads_group_numbers_like_java(self, spark): + df = spark.createDataFrame([("abcdefghij",)], _strings(spark, "s")) + assert df.select(F.regexp_replace("s", "(a)", "$1$1").alias("r")).collect() == [Row(r="aabcdefghij")] + cases = [ + ("(a)(b)(c)(d)(e)(f)(g)(h)(i)", "$10", "a0j"), + ("(a)(b)", "$12", "a2cdefghij"), + ("(a)(b)", "$01", "acdefghij"), + ] + for pattern, replacement, spark_value in cases: + if USE_ACTUAL_SPARK: + # Java reads another digit while the group number stays within the pattern's groups. + res = df.select(F.regexp_replace("s", pattern, replacement).alias("r")) + assert res.collect() == [Row(r=spark_value)] + else: + with pytest.raises(ContributionsAcceptedError, match=r"only \$0 to \$9"): + F.regexp_replace("s", pattern, replacement) + + def test_regexp_replace_group_ten(self, spark): + df = spark.createDataFrame([("abcdefghij", "(a)(b)(c)(d)(e)(f)(g)(h)(i)(j)")], _strings(spark, "s", "p")) + ten = "(a)(b)(c)(d)(e)(f)(g)(h)(i)(j)" + if USE_ACTUAL_SPARK: + res = df.select( + F.regexp_replace("s", ten, "$10").alias("a"), F.regexp_replace("s", F.col("p"), "$10").alias("b") + ) + assert res.collect() == [Row(a="j", b="j")] + return + # DuckDB refers to groups 0 to 9 only. + for pattern in [ten, F.col("p")]: + with pytest.raises(ContributionsAcceptedError, match=r"only \$0 to \$9"): + F.regexp_replace("s", pattern, "$10") + + def test_regexp_replace_column_reads_group_numbers_like_java(self, spark): + rows = [("abcdefghij", "$10"), ("abcdefghij", "<$1>")] + df = spark.createDataFrame(rows, _strings(spark, "s", "r")) + if not USE_ACTUAL_SPARK: + with pytest.raises(ContributionsAcceptedError, match="literal replacement"): + F.regexp_replace("s", "(a)", F.col("r")) + return + nine = df.select(F.regexp_replace("s", "(a)(b)(c)(d)(e)(f)(g)(h)(i)", F.col("r")).alias("r")) + assert nine.collect() == [Row(r="a0j"), Row(r="j")] + ten = df.select(F.regexp_replace("s", "(a)(b)(c)(d)(e)(f)(g)(h)(i)(j)", F.col("r")).alias("r")) + assert ten.collect() == [Row(r="j"), Row(r="")] + + def test_find_in_set_null(self, spark): + df = spark.createDataFrame([(None, "a,b"), ("a", None), ("b", "a,b"), ("c", "a,b")], _strings(spark, "s", "l")) + assert df.select(F.find_in_set("s", "l").alias("i")).collect() == [Row(i=None), Row(i=None), Row(i=2), Row(i=0)] + + def test_octet_length_type(self, spark): + df = spark.createDataFrame([("abc",)], _strings(spark, "s")) + res = df.select(F.octet_length("s").alias("n")) + assert res.schema.fields[0].dataType == IntegerType() + assert res.collect() == [Row(n=3)] + + def test_regexp_replace_with_columns(self, spark): + rows = [("100-200", "$2-$1", "-"), ("1-2", "\\$x", "(\\d)"), ("3-4", None, "-")] + df = spark.createDataFrame(rows, _strings(spark, "s", "r", "p")) + res = df.select( + F.regexp_replace("s", F.col("p"), "+").alias("pattern"), + F.regexp_replace("s", "-", F.lit("<$0>")).alias("literal"), + ) + assert res.collect() == [ + Row(pattern="100+200", literal="100<->200"), + Row(pattern="+-+", literal="1<->2"), + Row(pattern="3+4", literal="3<->4"), + ] + if not USE_ACTUAL_SPARK: + with pytest.raises(ContributionsAcceptedError, match="literal replacement"): + F.regexp_replace("s", r"(\d+)-(\d+)", F.col("r")) + return + per_row = df.select(F.regexp_replace("s", r"(\d+)-(\d+)", F.col("r")).alias("r")) + assert per_row.collect() == [Row(r="200-100"), Row(r="$x"), Row(r=None)] + assert df.select(F.regexp_replace("s", "-", F.col("r"))).columns == ["regexp_replace(s, -, r, 1)"] + + @pytest.mark.skipif(not USE_ACTUAL_SPARK, reason="every replacement column raises in the module") + @pytest.mark.parametrize("replacement", ["x\\", "$x"]) + def test_regexp_replace_invalid_replacement_column(self, spark, replacement): + df = spark.createDataFrame([("a", replacement)], _strings(spark, "s", "r")) + with pytest.raises(Exception): # noqa: B017, PT011 + df.select(F.regexp_replace("s", "a", F.col("r"))).collect() + + def test_regexp_replace_with_a_null_replacement(self, spark): + df = spark.createDataFrame([("a",)], _strings(spark, "s")) + assert df.select(F.regexp_replace("s", "a", F.lit(None)).alias("r")).collect() == [Row(r=None)] + + @pytest.mark.parametrize(("replacement", "error"), [("x\\", "character to be escaped"), ("$x", r"only \$0 to \$9")]) + def test_regexp_replace_invalid_replacement(self, spark, replacement, error): + df = spark.createDataFrame([("a",)], _strings(spark, "s")) + with pytest.raises(Exception, match=None if USE_ACTUAL_SPARK else error): + df.select(F.regexp_replace("s", "a", replacement)).collect() + + def test_pad_with_an_empty_pad(self, spark): + df = spark.createDataFrame([("abc",)], _strings(spark, "s")) + res = df.select(F.lpad("s", 5, ""), F.rpad("s", 2, ""), F.lpad("s", -1, "")) + assert res.columns == ["lpad(s, 5, )", "rpad(s, 2, )", "lpad(s, -1, )"] + assert res.collect() == [Row("abc", "ab", "")] + + def test_pad_with_columns(self, spark): + schema = StructType( + [StructField("s", StringType()), StructField("n", IntegerType()), StructField("p", StringType())] + ) + df = spark.createDataFrame([("abc", 5, "#"), ("abc", 2, ""), ("abc", None, None)], schema) + if SPARK_3: + with pytest.raises(PySparkTypeError, match="NOT_ITERABLE"): + F.lpad("s", F.col("n"), F.col("p")) + return + res = df.select(F.lpad("s", F.col("n"), F.col("p")), F.rpad("s", F.col("n"), "#")) + assert res.columns == ["lpad(s, n, p)", "rpad(s, n, #)"] + assert res.collect() == [Row("##abc", "abc##"), Row("ab", "ab"), Row(None, None)] + + def test_pad_with_a_null_width_and_an_empty_pad(self, spark): + schema = StructType( + [StructField("s", StringType()), StructField("n", IntegerType()), StructField("p", StringType())] + ) + df = spark.createDataFrame([("abc", None, ""), ("abc", -1, ""), ("abc", 2, "")], schema) + if SPARK_3: + with pytest.raises(PySparkTypeError, match="NOT_ITERABLE"): + F.lpad("s", F.col("n"), F.col("p")) + return + res = df.select(F.lpad("s", F.col("n"), F.col("p")).alias("l"), F.rpad("s", F.col("n"), F.col("p")).alias("r")) + assert res.collect() == [Row(l=None, r=None), Row(l="", r=""), Row(l="ab", r="ab")] + + def test_split_with_a_column_limit(self, spark): + df = spark.createDataFrame([("a,b,c",)], _strings(spark, "s")) + if not USE_ACTUAL_SPARK: + with pytest.raises(ContributionsAcceptedError, match="limit is not yet supported"): + F.split("s", ",", F.lit(2)) + elif SPARK_3: + with pytest.raises(PySparkTypeError, match="NOT_ITERABLE"): + F.split("s", ",", F.lit(2)) + else: + assert df.select(F.split("s", ",", F.lit(2)).alias("p")).collect() == [Row(p=["a", "b,c"])] + + def test_substr_with_null_position_or_length(self, spark): + schema = StructType([StructField("s", StringType()), StructField("p", LongType())]) + df = spark.createDataFrame([("abcdef", None)], schema) + res = df.select( + F.substr("s", "p", F.lit(2)).alias("pos"), + F.substr("s", "p").alias("pos_only"), + F.substr("s", F.lit(2), "p").alias("len"), + ) + assert res.collect() == [Row(pos=None, pos_only=None, len=None)] + + def test_char_truncates_fractional_codes(self, spark): + schema = StructType([StructField("d", DoubleType()), StructField("m", DecimalType(5, 1))]) + df = spark.createDataFrame([(65.7, Decimal("66.9")), (-0.5, Decimal("-0.9")), (321.9, None)], schema) + assert df.select(F.char("d").alias("d"), F.char("m").alias("m")).collect() == [ + Row(d="A", m="B"), + Row(d="\x00", m="\x00"), + Row(d="A", m=None), + ] + + def test_binary_functions_take_utf8_bytes(self, spark): + df = spark.createDataFrame([("é",), ("aé€😀",), ("a\\x41",)], _strings(spark, "s")) + res = df.select(F.octet_length("s").alias("n"), F.base64("s").alias("b")) + assert res.collect() == [Row(n=2, b="w6k="), Row(n=10, b="YcOp4oKs8J+YgA=="), Row(n=5, b="YVx4NDE=")] + binary = spark.createDataFrame([(bytearray(b"\xff\x00\\x"),)], StructType([StructField("b", BinaryType())])) + assert binary.select(F.octet_length("b").alias("n"), F.base64("b").alias("b")).collect() == [ + Row(n=4, b="/wBceA==") + ] diff --git a/tests/fast/spark/test_spark_group_by.py b/tests/fast/spark/test_spark_group_by.py index fafd747dd..3c25ac35a 100644 --- a/tests/fast/spark/test_spark_group_by.py +++ b/tests/fast/spark/test_spark_group_by.py @@ -2,7 +2,6 @@ _ = pytest.importorskip("duckdb.experimental.spark") -from spark_namespace import USE_ACTUAL_SPARK from spark_namespace.sql.functions import ( any_value, approx_count_distinct, @@ -89,13 +88,11 @@ def test_group_by(self, spark): df2 = df.groupBy("department").mean("salary").sort("department") res = df2.collect() - expected_res_str = ( - "[Row(department='Finance', mean(salary)=87750.0), Row(department='Marketing', " - "mean(salary)=85500.0), Row(department='Sales', mean(salary)=85666.66666666667)]" + assert ( + str(res) + == "[Row(department='Finance', avg(salary)=87750.0), Row(department='Marketing', avg(salary)=85500.0), " + "Row(department='Sales', avg(salary)=85666.66666666667)]" ) - if USE_ACTUAL_SPARK: - expected_res_str = expected_res_str.replace("mean(", "avg(") - assert str(res) == expected_res_str df2 = df.groupBy("department", "state").sum("salary", "bonus").sort("department", "state") res = df2.collect() @@ -299,8 +296,102 @@ def test_group_by_product(self, spark): def test_group_by_skewness(self, spark): df = spark.createDataFrame([[1, "A"], [1, "A"], [2, "A"]], ["c", "group"]) res = df.groupBy("group").agg(skewness(df.c).alias("v")).collect() - # TODO: Why is this different? # noqa: TD002, TD003 - if USE_ACTUAL_SPARK: - assert pytest.approx(res[0].v) == 0.7071067811865475 - else: - assert pytest.approx(res[0].v) == 1.7320508075688699 + assert pytest.approx(res[0].v) == 0.7071067811865475 + + +@pytest.fixture +def mixed(spark): + rows = [(1, 2, 1.5, "x", True), (1, 3, None, "y", False), (2, 5, 2.5, "x", None)] + return spark.createDataFrame(rows, "a int, b long, d double, s string, f boolean") + + +def _sorted_tuples(df): + return sorted((tuple(r) for r in df.collect()), key=repr) + + +class TestGroupedDataWithoutColumns: + @pytest.mark.parametrize( + ("method", "name", "x", "y"), + [ + ("sum", "sum", (3, 7, 4.0), (1, 3, None)), + ("avg", "avg", (1.5, 3.5, 2.0), (1.0, 3.0, None)), + ("mean", "avg", (1.5, 3.5, 2.0), (1.0, 3.0, None)), + ("min", "min", (1, 2, 1.5), (1, 3, None)), + ("max", "max", (2, 5, 2.5), (1, 3, None)), + ], + ) + def test_every_numeric_column(self, mixed, method, name, x, y): + res = getattr(mixed.groupBy("s"), method)() + assert res.columns == ["s", f"{name}(a)", f"{name}(b)", f"{name}(d)"] + assert _sorted_tuples(res) == [("x", *x), ("y", *y)] + + def test_keys_are_aggregated_too(self, mixed): + res = mixed.groupBy("a").sum() + assert res.columns == ["a", "sum(a)", "sum(b)", "sum(d)"] + assert _sorted_tuples(res) == [(1, 2, 5, 1.5), (2, 2, 5, 2.5)] + assert mixed.groupBy("A").max().columns == ["A", "max(a)", "max(b)", "max(d)"] + + def test_without_keys(self, mixed): + assert mixed.groupBy().sum().collect() == [Row(**{"sum(a)": 4, "sum(b)": 10, "sum(d)": 4.0})] + assert _sorted_tuples(mixed.groupBy().max()) == [(2, 5, 2.5)] + + +class TestGroupByExpressions: + def test_expression_keys(self, spark, mixed): + res = mixed.groupBy(mixed.a + 1).count() + assert res.columns == ["(a + 1)", "count"] + assert _sorted_tuples(res) == [(2, 2), (3, 1)] + res = mixed.groupBy((col("a") * 10).alias("k")).agg(sum("b")) + assert res.columns == ["k", "sum(b)"] + assert _sorted_tuples(res) == [(10, 5), (20, 5)] + assert _sorted_tuples(mixed.groupBy(mixed.a + 1).sum()) == [(2, 2, 5, 1.5), (3, 2, 5, 2.5)] + + @pytest.mark.parametrize("derive", ["withColumn", "filter", "select"]) + def test_expression_keys_on_derived_frames(self, mixed, derive): + derived = { + "withColumn": lambda: mixed.withColumn("c", col("b") * 2), + "filter": lambda: mixed.filter(mixed.b > 0), + "select": lambda: mixed.select("a", "b"), + }[derive]() + res = derived.groupBy(derived.a + 1).agg(sum("b")) + assert res.columns == ["(a + 1)", "sum(b)"] + assert _sorted_tuples(res) == [(2, 5), (3, 5)] + + def test_expression_keys_on_a_join(self, spark, mixed): + other = spark.createDataFrame([(1, "p"), (2, "q")], "a int, w string") + joined = mixed.join(other, "a") + res = joined.groupBy(joined.w, joined.s).count() + assert res.columns == ["w", "s", "count"] + assert _sorted_tuples(res) == [("p", "x", 1), ("p", "y", 1), ("q", "x", 1)] + assert _sorted_tuples(joined.groupBy(other.w).count()) == [("p", 2), ("q", 1)] + + +class TestGroupedDataColumnNames: + @pytest.fixture + def odd_names(self, spark): + return spark.createDataFrame([("x", 5, 1, 2), ("x", 7, 1, 3)], ["k", "2024", "null", "my value"]) + + @pytest.mark.parametrize( + ("method", "values"), + [("sum", (12, 2, 5)), ("avg", (6.0, 1.0, 2.5)), ("max", (7, 1, 3)), ("min", (5, 1, 2))], + ) + def test_every_numeric_column(self, odd_names, method, values): + res = getattr(odd_names.groupBy("k"), method)() + assert res.columns == ["k", f"{method}(2024)", f"{method}(null)", f"{method}(my value)"] + assert [tuple(r) for r in res.collect()] == [("x", *values)] + + def test_named_columns(self, odd_names): + assert [tuple(r) for r in odd_names.groupBy("k").sum("2024").collect()] == [("x", 12)] + res = odd_names.groupBy("k").mean("my value") + assert res.columns == ["k", "avg(my value)"] + assert [tuple(r) for r in res.collect()] == [("x", 2.5)] + res = odd_names.groupBy("k").sum("MY VALUE") + assert res.columns == ["k", "sum(MY VALUE)"] + assert [tuple(r) for r in res.collect()] == [("x", 5)] + + def test_global(self, odd_names): + res = odd_names.groupBy().sum() + assert res.columns == ["sum(2024)", "sum(null)", "sum(my value)"] + assert [tuple(r) for r in res.collect()] == [(12, 2, 5)] + assert odd_names.groupBy().count().collect() == [Row(count=2)] + assert odd_names.groupBy("k").count().columns == ["k", "count"] diff --git a/tests/fast/spark/test_spark_join.py b/tests/fast/spark/test_spark_join.py index 5ca8ca638..4f15c1e96 100644 --- a/tests/fast/spark/test_spark_join.py +++ b/tests/fast/spark/test_spark_join.py @@ -2,11 +2,22 @@ _ = pytest.importorskip("duckdb.experimental.spark") -from spark_namespace.sql.functions import col +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException, IllegalArgumentException, error_condition +from spark_namespace.sql import functions as sf +from spark_namespace.sql.functions import col, lit from spark_namespace.sql.types import ( Row, ) +import duckdb + +if USE_ACTUAL_SPARK: + import pyspark + +# Spark raises AnalysisException for every ambiguous reference; DuckDB's binder raises its own error for some. +AMBIGUOUS_REFERENCE_ERRORS = (AnalysisException, duckdb.BinderException) + @pytest.fixture def dataframe_a(spark): @@ -406,7 +417,6 @@ def test_join_with_common_column(self, spark, dataframe_a): "name='Williams', year_joined='2010')]" ) - @pytest.mark.xfail(condition=True, reason="Selecting from a duplicate binding causes an error") def test_join_on_joined_data_error(self, spark): df = spark.createDataFrame([(2, "Alice"), (5, "Bob")]).toDF("age", "name") df2 = spark.createDataFrame([Row(height=80, name="Tom"), Row(height=85, name="Bob")], ["height", "name"]) @@ -427,3 +437,822 @@ def test_join_on_column_name_stored_in_variable(self, spark): res = df.join(df2, col1 == col2) res = res.collect() assert str(res) == "[Row(age=5, name='Bob', height=85, name='Bob')]" + + +@pytest.fixture +def people(spark): + return spark.createDataFrame([(1, "Ann", 2), (2, "Bob", 3), (3, "Cid", 1)], ["id", "name", "manager_id"]) + + +@pytest.fixture +def teams(spark): + return spark.createDataFrame([(1, "red"), (2, "blue"), (4, "gold")], ["id", "team"]) + + +class TestJoinColumnResolution: + def test_self_join_on_column_name(self, people): + df = people.join(people, "id") + assert df.columns == ["id", "name", "manager_id", "name", "manager_id"] + assert sorted(map(tuple, df.collect())) == [ + (1, "Ann", 2, "Ann", 2), + (2, "Bob", 3, "Bob", 3), + (3, "Cid", 1, "Cid", 1), + ] + + @pytest.mark.parametrize("how", ["inner", "left", "full"]) + def test_join_filtered_views_of_one_dataframe(self, people, how): + low = people.filter(col("id") < 3) + high = people.filter(col("id") > 1) + df = low.join(high, "id", how).select("id") + expected = {"inner": [2], "left": [1, 2], "full": [1, 2, 3]}[how] + assert sorted(row.id for row in df.collect()) == expected + + def test_cross_join_with_itself(self, people): + assert people.crossJoin(people).count() == 9 + assert people.join(people).count() == 9 + + def test_join_without_condition_is_a_cross_join(self, people, teams): + df = people.join(teams) + assert df.columns == ["id", "name", "manager_id", "id", "team"] + assert df.count() == 9 + + def test_self_join_with_aliases(self, people): + df = ( + people.alias("a") + .join(people.alias("b"), col("a.manager_id") == col("b.id")) + .select(col("a.name").alias("employee"), col("b.name").alias("manager")) + ) + assert sorted(df.collect()) == [ + Row(employee="Ann", manager="Bob"), + Row(employee="Bob", manager="Cid"), + Row(employee="Cid", manager="Ann"), + ] + + def test_join_two_joins_of_the_same_inputs(self, people, teams): + first = people.join(teams, "id") + second = people.join(teams, "id") + df = first.join(second, "id") + assert df.columns == ["id", "name", "manager_id", "team", "name", "manager_id", "team"] + assert sorted(map(tuple, df.collect())) == [ + (1, "Ann", 2, "red", "Ann", 2, "red"), + (2, "Bob", 3, "blue", "Bob", 3, "blue"), + ] + + def test_join_a_join_with_one_of_its_inputs(self, people, teams): + df = people.join(teams, "id").join(people, "id") + assert df.columns == ["id", "name", "manager_id", "team", "name", "manager_id"] + assert sorted(map(tuple, df.collect())) == [ + (1, "Ann", 2, "red", "Ann", 2), + (2, "Bob", 3, "blue", "Bob", 3), + ] + + def test_select_the_join_column_of_each_side(self, people, teams): + df = people.join(teams, people.id == teams.id, "left").select(people.id, teams.id, "team") + assert df.columns == ["id", "id", "team"] + assert sorted(map(tuple, df.collect())) == [ + (1, 1, "red"), + (2, 2, "blue"), + (3, None, None), + ] + + def test_column_on_both_sides_is_ambiguous(self, people, teams): + joined = people.join(teams, people.id == teams.id) + with pytest.raises(AMBIGUOUS_REFERENCE_ERRORS): + joined.select(col("id")).collect() + for reference in [lambda: joined.id, lambda: joined["id"], lambda: joined.select("id")]: + with pytest.raises(AnalysisException) as error: + reference().collect() + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + # Spark 3.5 returns a Java map. + assert dict(error.value.getMessageParameters()) == {"name": "`id`", "referenceNames": "[`id`, `id`]"} + assert error.value.getSqlState() == "42704" + + def test_column_of_a_joined_dataframe(self, people, teams): + joined = people.join(teams, people.id == teams.id) + assert sorted(joined.select(joined.name, joined["team"]).collect()) == [ + Row(name="Ann", team="red"), + Row(name="Bob", team="blue"), + ] + assert joined.filter(joined.team == "red").select("name").collect() == [Row(name="Ann")] + + @pytest.mark.parametrize( + ("how", "expected_ids"), + [("inner", [1, 2]), ("left", [1, 2, 3]), ("right", [1, 2, 4]), ("full", [1, 2, 3, 4])], + ) + def test_join_key_of_a_join_on_column_names(self, spark, people, teams, how, expected_ids): + joined = people.join(teams, "id", how) + assert sorted(row.id for row in joined.select(joined.id).collect()) == expected_ids + flags = spark.createDataFrame([(1, True), (4, False)], ["id", "flag"]) + df = joined.join(flags, joined.id == flags.id).select(joined.id, flags.flag) + assert sorted(df.collect()) == [Row(id=i, flag=i == 1) for i in [1, 4] if i in expected_ids] + + def test_with_column_on_a_column_of_a_joined_dataframe(self, spark): + visits = spark.createDataFrame([(1, "Marc", 10), (2, "Bob", 20)], ["site_id", "person", "minutes"]) + sites = spark.createDataFrame([(1, "A"), (2, "B"), (3, "C")], ["site_id", "name"]) + visits = visits.join(sites, on="site_id", how="left") + visits = visits.withColumn("minutes", visits.minutes * 2) + assert sorted(visits.collect()) == [ + Row(site_id=1, person="Marc", minutes=20, name="A"), + Row(site_id=2, person="Bob", minutes=40, name="B"), + ] + + def test_chained_joins(self, spark, people, teams): + flags = spark.createDataFrame([(1, True), (2, False)], ["id", "flag"]) + df = ( + people.join(teams, people.id == teams.id) + .join(flags, people.id == flags.id) + .select(people.id, teams.team, flags.flag) + ) + assert sorted(df.collect()) == [Row(id=1, team="red", flag=True), Row(id=2, team="blue", flag=False)] + df = people.join(teams, "id").join(flags, "id") + assert df.columns == ["id", "name", "manager_id", "team", "flag"] + assert sorted(df.collect()) == [ + Row(id=1, name="Ann", manager_id=2, team="red", flag=True), + Row(id=2, name="Bob", manager_id=3, team="blue", flag=False), + ] + + def test_join_after_with_column(self, people, teams): + doubled = people.withColumn("twice", col("id") * 2) + df = doubled.join(teams, doubled.id == teams.id).select(doubled.id, doubled.twice, teams.team) + assert sorted(df.collect()) == [Row(id=1, twice=2, team="red"), Row(id=2, twice=4, team="blue")] + + def test_with_column_after_join_with_a_shared_column_name(self, people, teams): + joined = people.join(teams, people.id == teams.id) + df = joined.withColumn("one", lit(1)) + assert df.columns == ["id", "name", "manager_id", "id", "team", "one"] + assert sorted(map(tuple, df.collect())) == [ + (1, "Ann", 2, 1, "red", 1), + (2, "Bob", 3, 2, "blue", 1), + ] + df = joined.withColumn("team", lit("x")).withColumns({"name": lit("y")}) + assert sorted(map(tuple, df.collect())) == [ + (1, "y", 2, 1, "x"), + (2, "y", 3, 2, "x"), + ] + df = joined.withColumnRenamed("team", "colour") + assert df.columns == ["id", "name", "manager_id", "id", "colour"] + df = joined.withColumn("doubled", people.manager_id * 2).select("doubled", "team") + assert sorted(df.collect()) == [Row(doubled=4, team="red"), Row(doubled=6, team="blue")] + + def test_drop_the_join_column_of_one_side(self, people, teams): + joined = people.join(teams, people.id == teams.id, "left") + assert sorted(joined.drop(teams.id).collect()) == [ + Row(id=1, name="Ann", manager_id=2, team="red"), + Row(id=2, name="Bob", manager_id=3, team="blue"), + Row(id=3, name="Cid", manager_id=1, team=None), + ] + df = joined.drop(people.id) + assert df.columns == ["name", "manager_id", "id", "team"] + assert sorted(df.collect(), key=lambda row: row.name) == [ + Row(name="Ann", manager_id=2, id=1, team="red"), + Row(name="Bob", manager_id=3, id=2, team="blue"), + Row(name="Cid", manager_id=1, id=None, team=None), + ] + + def test_reference_through_a_dataframe_on_both_sides_raises(self, people): + with pytest.raises(AMBIGUOUS_REFERENCE_ERRORS): + people.join(people, "id").select(people.name).collect() + with pytest.raises(AMBIGUOUS_REFERENCE_ERRORS): + people.join(people, people.id == people.manager_id).collect() + + def test_self_join_on_the_same_column_reference(self, people): + low = people.filter(col("id") < 3) + high = people.filter(col("id") > 1) + if USE_ACTUAL_SPARK: + # Spark rewrites an equality between the same column of both sides to left equals right. + assert people.join(people, people.id == people.id).count() == 3 + assert low.join(high, low.id == high.id).select(low.name).collect() == [Row(name="Bob")] + else: + with pytest.raises(duckdb.BinderException): + people.join(people, people.id == people.id).collect() + with pytest.raises(duckdb.BinderException): + low.join(high, low.id == high.id).collect() + + @pytest.mark.parametrize( + "derive", + [ + lambda df: df.filter(col("team") != "none"), + lambda df: df.where(col("team") != "none"), + lambda df: df.orderBy("team"), + lambda df: df.sort("team"), + lambda df: df.limit(5), + lambda df: df.distinct(), + lambda df: df.dropDuplicates(), + lambda df: df.filter(col("team") != "none").limit(5), + lambda df: df.filter(col("team") != "none").alias("x"), + ], + ) + def test_columns_of_a_frame_derived_from_a_join(self, people, teams, derive): + derived = derive(people.join(teams, "id")) + assert sorted(derived.select(derived["name"], derived.team).collect()) == [ + Row(name="Ann", team="red"), + Row(name="Bob", team="blue"), + ] + assert derived.filter(derived.team == "red").select(derived.id).collect() == [Row(id=1)] + assert sorted(derived.select(people.name, teams.team).collect()) == [ + Row(name="Ann", team="red"), + Row(name="Bob", team="blue"), + ] + + def test_columns_of_a_cached_join(self, people, teams): + cached = people.join(teams, "id").cache() + assert sorted(cached.select(cached.name, cached["team"]).collect()) == [ + Row(name="Ann", team="red"), + Row(name="Bob", team="blue"), + ] + + def test_frame_derived_from_a_join_as_a_join_input(self, spark, people, teams): + derived = people.join(teams, "id").filter(col("team") != "none") + flags = spark.createDataFrame([(1, True), (2, False)], ["flag_id", "flag"]) + expected = [Row(name="Ann", flag=True), Row(name="Bob", flag=False)] + df = derived.join(flags, col("id") == col("flag_id")).select("name", "flag") + assert sorted(df.collect()) == expected + df = derived.join(flags.withColumnRenamed("flag_id", "id"), "id").select("name", "flag") + assert sorted(df.collect()) == expected + if USE_ACTUAL_SPARK: + df = derived.join(flags, derived.id == flags.flag_id).select(derived.name, flags.flag) + assert sorted(df.collect()) == expected + else: + # The derived frame's columns resolve in its own scope, which a later join wraps in a subquery. + with pytest.raises(duckdb.BinderException): + derived.join(flags, derived.id == flags.flag_id).collect() + + +@pytest.fixture +def left(spark): + return spark.createDataFrame([(1, 2), (2, 4)], ["a", "b"]) + + +@pytest.fixture +def right(spark): + return spark.createDataFrame([(1, 3), (3, 9)], ["a", "c"]) + + +class TestDropAfterJoin: + @pytest.mark.parametrize( + ("how", "drop_left", "drop_right", "drop_name"), + [ + ("inner", [(2, 3)], [(1, 2, 3)], [(2, 3)]), + ("left", [(2, 3), (4, None)], [(1, 2, 3), (2, 4, None)], [(2, 3), (4, None)]), + ("right", [(1, 2, 3), (3, None, 9)], [(2, 3), (None, 9)], [(2, 3), (None, 9)]), + ( + "full", + [(1, 2, 3), (2, 4, None), (3, None, 9)], + [(1, 2, 3), (2, 4, None), (3, None, 9)], + [(2, 3), (4, None), (None, 9)], + ), + ], + ) + def test_drop_the_key_of_a_join_on_a_name(self, left, right, how, drop_left, drop_right, drop_name): + # The key of a join on a name is the left input's, the right input's for a right join and neither's for a + # full join, so dropping the other input's column changes nothing. + joined = left.join(right, "a", how) + + def rows(df): + return sorted(map(tuple, df.collect()), key=str) + + assert rows(joined.drop(left.a)) == drop_left + assert rows(joined.drop(right.a)) == drop_right + assert rows(joined.drop("a")) == drop_name + assert rows(joined.drop(col("A"))) == drop_name + assert joined.drop("A").columns == ["b", "c"] + + def test_drop_after_a_join_on_an_expression(self, left, right): + joined = left.join(right, left.a == right.a) + assert joined.drop(left.a).columns == ["b", "a", "c"] + assert joined.drop(left.a).collect() == [Row(b=2, a=1, c=3)] + assert joined.drop(right.a).columns == ["a", "b", "c"] + assert joined.drop(left.a, right.a).columns == ["b", "c"] + assert joined.drop("A").columns == ["b", "c"] + assert joined.drop("b", left.a).columns == ["a", "c"] + with pytest.raises(AnalysisException) as error: + joined.drop(col("a")) + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + + def test_drop_through_a_frame_derived_from_a_join(self, left, right): + assert left.join(right, "a").filter(left.a > 0).drop(left.a).columns == ["b", "c"] + + def test_drop_through_aliases(self, left): + joined = left.alias("x").join(left.alias("y"), col("x.a") == col("y.a")) + assert joined.drop(col("x.a")).columns == ["b", "a", "b"] + assert sorted(map(tuple, joined.drop(col("y.b")).collect())) == [(1, 2, 1), (2, 4, 2)] + + +class TestCacheOfAJoin: + def test_cache_keeps_join_input_references(self, people, teams): + cached = people.join(teams, "id").cache() + assert sorted(cached.select(people.name, teams.team).collect()) == [ + Row(name="Ann", team="red"), + Row(name="Bob", team="blue"), + ] + + def test_cache_keeps_duplicate_names(self, people, teams): + cached = people.join(teams, people.id == teams.id).cache() + assert cached.columns == ["id", "name", "manager_id", "id", "team"] + with pytest.raises(AMBIGUOUS_REFERENCE_ERRORS): + cached.select("id").collect() + assert sorted(cached.select(teams.id).collect()) == [Row(id=1), Row(id=2)] + + +class TestSourceFramesAreDistinct: + def test_reference_to_another_range(self, spark): + with pytest.raises((AnalysisException, duckdb.BinderException)): + spark.range(3).select(spark.range(5).id).collect() + + def test_reference_to_another_created_frame(self, spark): + first = spark.createDataFrame([(1,)], ["x"]) + second = spark.createDataFrame([(1,)], ["x"]) + with pytest.raises((AnalysisException, duckdb.BinderException)): + first.select(second.x).collect() + + def test_join_of_two_ranges(self, spark): + small, large = spark.range(2), spark.range(3) + assert sorted(small.join(large, small.id == large.id).select(small.id).collect()) == [Row(id=0), Row(id=1)] + assert sorted(small.join(large, "id").collect()) == [Row(id=0), Row(id=1)] + + def test_join_of_two_reads_of_a_file(self, spark, tmp_path): + path = str(tmp_path / "data.parquet") + spark.createDataFrame([(1, "a"), (2, "b")], ["id", "v"]).write.parquet(path) + first, second = spark.read.parquet(path), spark.read.parquet(path) + df = first.join(second, first.id == second.id).select(first.id, second.v) + assert sorted(df.collect()) == [Row(id=1, v="a"), Row(id=2, v="b")] + + +class TestDeliberateDifferences: + def test_join_on_a_name_keeps_the_column_order(self, spark): + a = spark.createDataFrame([(1, 10)], ["a", "k"]) + b = spark.createDataFrame([(2, 10)], ["b", "k"]) + joined = a.join(b, "k") + if USE_ACTUAL_SPARK: + assert joined.columns == ["k", "a", "b"] + else: + # The module keeps DuckDB's order: the left columns, then the right columns that are not keys. + assert joined.columns == ["a", "k", "b"] + assert joined.select("k", "a", "b").collect() == [Row(k=10, a=1, b=2)] + + def test_aliased_self_join_answers(self, people): + def managers(): + return people.alias("a").join(people, col("a.manager_id") == people.id).select(people.name).collect() + + if USE_ACTUAL_SPARK: + with pytest.raises(AnalysisException): + managers() + else: + # The module can tell the sides apart, so it answers with the rows of the side without the alias. + assert sorted(managers()) == [ + Row(name="Ann"), + Row(name="Bob"), + Row(name="Cid"), + ] + + +@pytest.fixture +def xs(spark): + return spark.createDataFrame([(1, "x1"), (2, "x2")], ["id", "x"]) + + +@pytest.fixture +def ys(spark): + return spark.createDataFrame([(1, "y1"), (3, "y3")], ["id", "y"]) + + +_JOIN_TYPE_KEYS = [ + ("inner", [1]), + ("cross", [1]), + ("outer", [1, 2, 3]), + ("full", [1, 2, 3]), + ("fullouter", [1, 2, 3]), + ("full_outer", [1, 2, 3]), + ("left", [1, 2]), + ("leftouter", [1, 2]), + ("left_outer", [1, 2]), + ("right", [1, 3]), + ("rightouter", [1, 3]), + ("right_outer", [1, 3]), + ("semi", [1]), + ("leftsemi", [1]), + ("left_semi", [1]), + ("anti", [2]), + ("leftanti", [2]), + ("left_anti", [2]), +] + + +class TestJoinTypes: + @pytest.mark.parametrize( + ("how", "keys"), + [(how, keys) for name, keys in _JOIN_TYPE_KEYS for how in (name, name.upper())] + [("Left_Outer", [1, 2])], + ) + def test_join_type_spellings(self, xs, ys, how, keys): + assert sorted(row.id for row in xs.join(ys, "id", how).select("id").collect()) == keys + + @pytest.mark.parametrize(("how", "rows"), [("CROSS", [(1, "x1", 1, "y1")]), ("FULL", [(1, "x1", 1, "y1")])]) + def test_join_type_with_a_condition(self, xs, ys, how, rows): + joined = xs.join(ys, xs.id == ys.id, how).filter(xs.id.isNotNull() & ys.id.isNotNull()) + assert sorted(map(tuple, joined.collect())) == rows + + def test_unknown_join_type_raises(self, xs, ys): + # Spark 4.2.0 raises an AnalysisException, Spark 3.5.9 an IllegalArgumentException. + with pytest.raises((AnalysisException, IllegalArgumentException), match="Unsupported join type"): + xs.join(ys, "id", "sideways").collect() + + +@pytest.fixture +def staff(spark): + return spark.createDataFrame([(1, "Ann"), (1, "Abe"), (2, "Bob")], ["id", "name"]) + + +@pytest.fixture +def depts(spark): + return spark.createDataFrame([(1, "red"), (2, "blue")], ["id", "team"]) + + +class TestFramesDerivedFromAJoin: + @pytest.mark.parametrize( + "derive", + [ + lambda j: j.filter(col("team") != "none").sort("team"), + lambda j: j.filter(col("team") != "none").filter(col("name") != "none"), + lambda j: j.sort("team").filter(col("team") != "none"), + lambda j: j.sort("team").sort("name"), + lambda j: j.limit(5).sort("name").filter(col("team") != "none"), + lambda j: j.filter(col("team") != "none").sort("name").limit(5).distinct(), + lambda j: j.distinct().filter(col("team") != "none").sort("name"), + lambda j: j.filter(col("team") != "none").alias("q").sort("name"), + ], + ) + def test_names_in_chains_of_steps(self, staff, depts, derive): + derived = derive(staff.join(depts, "id")) + expected = [Row(name="Abe"), Row(name="Ann"), Row(name="Bob")] + assert sorted(derived.select("name").collect()) == expected + assert sorted(derived.select(col("name")).collect()) == expected + assert sorted(derived.select(derived.name).collect()) == expected + assert sorted(derived.select(staff.name).collect()) == expected + assert sorted(derived.withColumn("n", lit(1)).select("name", "n").collect()) == [ + Row(name="Abe", n=1), + Row(name="Ann", n=1), + Row(name="Bob", n=1), + ] + assert derived.filter(col("team") == "blue").select("name").collect() == [Row(name="Bob")] + + @pytest.mark.parametrize( + "derive", + [ + lambda j: j.filter(col("team") != "none"), + lambda j: j.filter(col("team") != "none").sort("team"), + ], + ) + def test_names_in_chains_over_a_cross_join(self, staff, depts, derive): + derived = derive(staff.crossJoin(depts.withColumnRenamed("id", "dept_id"))) + assert derived.count() == 6 + assert sorted(derived.filter(col("team") == "red").select("name").collect()) == [ + Row(name="Abe"), + Row(name="Ann"), + Row(name="Bob"), + ] + assert sorted(derived.filter(col("dept_id") == 2).select(col("name"), staff.id).collect()) == [ + Row(name="Abe", id=1), + Row(name="Ann", id=1), + Row(name="Bob", id=2), + ] + + def test_filter_after_a_limit(self, staff, depts): + first = staff.join(depts, "id").orderBy("name").limit(1) + assert first.filter(col("name") == "Bob").collect() == [] + assert first.collect() == [Row(id=1, name="Abe", team="red")] + + def test_aggregates_after_a_limit_or_a_sort(self, staff, depts): + joined = staff.join(depts, "id") + assert joined.limit(1).count() == 1 + assert joined.orderBy("name").limit(2).groupBy().count().collect() == [Row(count=2)] + assert joined.orderBy("name").count() == 3 + assert sorted(joined.orderBy("name").groupBy("team").count().collect()) == [ + Row(team="blue", count=1), + Row(team="red", count=2), + ] + + def test_projection_after_a_distinct(self, staff, depts): + distinct = staff.join(depts, "id").distinct() + assert sorted(distinct.select("id").collect()) == [Row(id=1), Row(id=1), Row(id=2)] + assert sorted(distinct.select(staff.id).collect()) == [Row(id=1), Row(id=1), Row(id=2)] + assert distinct.select("team").distinct().count() == 2 + assert distinct.count() == 3 + + def test_projection_after_a_sort_keeps_the_order(self, staff, depts): + sorted_join = staff.join(depts, "id").orderBy(col("name").desc()) + assert sorted_join.select("name", col("id").alias("name_2")).collect() == [ + Row(name="Bob", name_2=2), + Row(name="Ann", name_2=1), + Row(name="Abe", name_2=1), + ] + assert sorted_join.select(col("team").alias("name")).collect() == [ + Row(name="blue"), + Row(name="red"), + Row(name="red"), + ] + + def test_repeated_names_after_a_sort_of_a_filtered_join(self, staff, depts): + joined = staff.join(depts, staff.id == depts.id) + derived = joined.filter(col("team") != "none").sort("name") + assert derived.columns == ["id", "name", "id", "team"] + assert derived.drop(depts.id).collect() == [ + Row(id=1, name="Abe", team="red"), + Row(id=1, name="Ann", team="red"), + Row(id=2, name="Bob", team="blue"), + ] + with pytest.raises(AnalysisException) as error: + derived.select("id") + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + + def test_to_df_after_a_join_with_repeated_names(self, staff, depts): + renamed = staff.join(depts, staff.id == depts.id).toDF("a", "b", "c", "d") + assert renamed.columns == ["a", "b", "c", "d"] + assert sorted(renamed.collect()) == [ + Row(a=1, b="Abe", c=1, d="red"), + Row(a=1, b="Ann", c=1, d="red"), + Row(a=2, b="Bob", c=2, d="blue"), + ] + + +class TestSemiAndAntiJoinColumns: + @pytest.fixture + def named(self, spark): + return spark.createDataFrame([(1, "c")], ["id", "name"]) + + @pytest.mark.parametrize(("how", "names"), [("left_anti", ["Bob"]), ("semi", ["Abe", "Ann"])]) + def test_name_both_inputs_have(self, staff, named, how, names): + joined = staff.join(named, "id", how) + assert sorted(row.name for row in joined.select("name").collect()) == names + assert sorted(row.name for row in joined.select(joined.name).collect()) == names + assert sorted(row.name for row in joined.withColumn("n", lit(1)).collect()) == names + + def test_name_both_inputs_have_on_a_condition(self, staff, named): + joined = staff.join(named, staff.id == named.id, "semi") + assert sorted(joined.select("name").collect()) == [Row(name="Abe"), Row(name="Ann")] + + +class TestDropThroughAnAlias: + def test_drop_a_reference_to_the_aliased_frame(self, people): + assert people.alias("q").drop(people.name).columns == ["id", "manager_id"] + assert people.alias("q").alias("r").drop(people.name).columns == ["id", "manager_id"] + assert people.alias("q").filter(col("id") > 1).drop(people.name).collect() == [ + Row(id=2, manager_id=3), + Row(id=3, manager_id=1), + ] + assert people.alias("q").drop(col("q.name")).columns == ["id", "manager_id"] + + def test_drop_a_reference_to_another_frame(self, people, teams): + assert people.drop(teams.id).columns == ["id", "name", "manager_id"] + + def test_select_a_reference_to_the_aliased_frame(self, people): + assert people.alias("q").filter(col("id") == 1).select(people.name).collect() == [Row(name="Ann")] + + +def _rows(df): + return sorted(map(tuple, df.collect()), key=str) + + +@pytest.fixture +def unmatched(spark): + # A full join of inputs without a common key, so the two id columns differ in every row. + left = spark.createDataFrame([(1, "a")], ["id", "x"]) + right = spark.createDataFrame([(2, "b")], ["id", "y"]) + return left, right, left.join(right, left.id == right.id, "full") + + +# Frames over the join with the columns id, x, id and y. +_OVER_THE_JOIN = { + "join": lambda left, right, joined: joined, + "alias": lambda left, right, joined: joined.alias("q"), + "filter": lambda left, right, joined: joined.filter(col("x").isNull() | (col("x") == "a")), + "sort": lambda left, right, joined: joined.orderBy("x"), + "limit": lambda left, right, joined: joined.limit(10), + "distinct": lambda left, right, joined: joined.distinct(), + "projection": lambda left, right, joined: joined.select(left.id, "x", right.id, "y"), +} + + +class TestRepeatedColumnNames: + @pytest.mark.parametrize("derive", _OVER_THE_JOIN.values(), ids=_OVER_THE_JOIN.keys()) + def test_frame_with_repeated_names_as_a_join_input(self, spark, unmatched, derive): + other = spark.createDataFrame([(9,)], ["c"]) + joined = derive(*unmatched).crossJoin(other) + assert joined.columns == ["id", "x", "id", "y", "c"] + rows = [(1, "a", None, None, 9), (None, None, 2, "b", 9)] + assert _rows(joined) == rows + assert _rows(joined.withColumn("n", lit(1))) == [(*row, 1) for row in rows] + assert _rows(joined.toDF("a", "b", "c", "d", "e")) == rows + assert _rows(joined.drop("c")) == [row[:4] for row in rows] + assert _rows(joined.drop(other.c)) == [row[:4] for row in rows] + assert _rows(joined.withColumnRenamed("x", "z")) == rows + + @pytest.mark.parametrize("derive", _OVER_THE_JOIN.values(), ids=_OVER_THE_JOIN.keys()) + def test_repeated_name_in_a_frame_with_repeated_names_as_a_join_input(self, spark, unmatched, derive): + other = spark.createDataFrame([(9,)], ["c"]) + joined = derive(*unmatched).crossJoin(other) + for reference in [lambda: joined.select("id"), lambda: joined.select(col("id")), lambda: joined["id"]]: + with pytest.raises(AnalysisException) as error: + reference() + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + with pytest.raises(AnalysisException) as error: + derive(*unmatched).join(other, col("id") == col("c")) + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + assert _rows(joined.select(col("x"), "c")) == [("a", 9), (None, 9)] + + @pytest.mark.parametrize( + "use", + [ + lambda df: df.select(col("id")), + lambda df: df.select(col("id") + 1), + lambda df: df.select(sf.sum("id")), + lambda df: df.filter(col("id") > 0), + lambda df: df.withColumn("n", col("id") * 2), + lambda df: df.withColumns({"n": col("id")}), + lambda df: df.groupBy("id").count(), + lambda df: df.groupBy(col("id")).count(), + lambda df: df.groupBy("x").agg(sf.max(col("id"))), + ], + ) + @pytest.mark.parametrize("derive", _OVER_THE_JOIN.values(), ids=_OVER_THE_JOIN.keys()) + def test_reference_to_a_repeated_name_is_ambiguous(self, unmatched, derive, use): + with pytest.raises(AnalysisException) as error: + use(derive(*unmatched)).collect() + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + + @pytest.mark.parametrize("derive", _OVER_THE_JOIN.values(), ids=_OVER_THE_JOIN.keys()) + def test_sort_by_a_repeated_name_raises(self, unmatched, derive): + # Spark reports these as unresolved, the module as ambiguous. + with pytest.raises(AnalysisException): + derive(*unmatched).sort(col("id")).collect() + with pytest.raises(AnalysisException): + derive(*unmatched).orderBy("id").collect() + + def test_reference_through_the_alias_of_a_frame_with_repeated_names(self, unmatched): + with pytest.raises(AnalysisException) as error: + unmatched[2].alias("q").select(col("q.id")).collect() + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + assert _rows(unmatched[2].alias("q").select(col("q.x"))) == [("a",), (None,)] + + def test_repeated_names_from_to_df(self, spark): + df = spark.createDataFrame([(1, 2)], ["a", "b"]) + renamed = df.toDF("x", "x") + for use in [lambda: renamed.select(col("x")), lambda: renamed.filter(col("x") > 0), lambda: renamed.x]: + with pytest.raises(AnalysisException) as error: + use() + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + # Spark has no column of df left after toDF; the module finds two of df.a's name. + with pytest.raises(AnalysisException): + df.toDF("a", "a").select(df.a).collect() + + def test_references_to_the_join_inputs_still_resolve(self, unmatched): + left, right, joined = unmatched + assert _rows(joined.select(left.id + 1, right.id)) == [(2, None), (None, 2)] + assert _rows(joined.filter(right.id == 2).select("y")) == [("b",)] + assert _rows(joined.withColumn("n", left.id * 10).select("n")) == [(10,), (None,)] + assert _rows(joined.groupBy(left.id).count()) == [(1, 1), (None, 1)] + assert _rows(joined.filter(col("x") == "a").select(left.id, right.id)) == [(1, None)] + + def test_integer_index_keeps_the_frame(self, unmatched): + left, right, joined = unmatched + with pytest.raises((AnalysisException, duckdb.BinderException)): + left.select(right[0]).collect() + assert left.select(left[1]).collect() == [Row(x="a")] + assert _rows(joined.select(left[0], right[1])) == [(1, None), (None, "b")] + with pytest.raises(AnalysisException) as error: + joined[0] + assert error_condition(error.value) == "AMBIGUOUS_REFERENCE" + + def test_join_on_a_name_that_repeats_in_an_input(self, spark, unmatched): + # The key is the first column of the name, as Spark resolves it. + other = spark.createDataFrame([(1, "c")], ["id", "z"]) + joined = unmatched[2].limit(10).join(other, "id") + assert joined.columns == ["id", "x", "id", "y", "z"] + assert _rows(joined) == [(1, "a", None, None, "c")] + assert _rows(joined.withColumn("n", lit(0))) == [(1, "a", None, None, "c", 0)] + + +@pytest.fixture +def abc(spark): + a = spark.createDataFrame([(1, "a1"), (2, "a2")], ["id", "x"]) + b = spark.createDataFrame([(1, "b1"), (3, "b3")], ["id", "y"]) + c = spark.createDataFrame([(1, "c1"), (2, "c2")], ["k", "z"]) + return a, b, c + + +class TestColumnsOfChainedJoins: + def test_join_on_a_name_then_on_an_expression(self, abc): + a, b, c = abc + joined = a.join(b, "id", "full").join(c, col("id") == c.k, "left").withColumn("n", lit(0)) + assert joined.columns == ["id", "x", "y", "k", "z", "n"] + assert _rows(joined) == [ + (1, "a1", "b1", 1, "c1", 0), + (2, "a2", None, 2, "c2", 0), + (3, None, "b3", None, None, 0), + ] + + def test_right_join_on_a_name(self, abc): + a, b, c = abc + joined = a.join(b, "id", "right").join(c, b.id == c.k, "left").withColumn("n", lit(0)) + assert _rows(joined) == [(1, "a1", "b1", 1, "c1", 0), (3, None, "b3", None, None, 0)] + + @pytest.mark.parametrize(("how", "rows"), [("semi", [(1, "a1", 1, "c1")]), ("anti", [(2, "a2", 2, "c2")])]) + def test_semi_and_anti_joins_as_inputs(self, abc, how, rows): + a, b, c = abc + joined = a.join(b, "id", how).join(c, a.id == c.k) + assert _rows(joined.withColumn("n", lit(0))) == [(*row, 0) for row in rows] + assert _rows(joined.toDF("p", "q", "r", "s")) == rows + + def test_join_as_the_right_input(self, abc): + a, b, c = abc + joined = c.crossJoin(a.join(b, a.id == b.id, "full")).filter(col("k") == 1).withColumn("n", lit(0)) + assert joined.columns == ["k", "z", "id", "x", "id", "y", "n"] + assert _rows(joined) == [ + (1, "c1", 1, "a1", 1, "b1", 0), + (1, "c1", 2, "a2", None, None, 0), + (1, "c1", None, None, 3, "b3", 0), + ] + assert _rows(joined.drop(b.id).select("x", "y")) == [("a1", "b1"), ("a2", None), (None, "b3")] + + @pytest.mark.parametrize(("how", "columns"), [("inner", ["id", "x", "y"]), ("right", ["ID", "x", "y"])]) + def test_join_on_a_name_in_another_case(self, spark, how, columns): + a = spark.createDataFrame([(1, "a1")], ["id", "x"]) + b = spark.createDataFrame([(1, "b1")], ["ID", "y"]) + joined = a.join(b, "id", how) + assert joined.columns == columns + assert joined.select(a.x, b.y, a.x == "a1").collect() == [Row(x="a1", y="b1", **{"(x = a1)": True})] + + +class TestNamesThroughDerivedFrames: + def test_star_keeps_repeated_names(self, spark, unmatched): + left, _, joined = unmatched + assert joined.orderBy(left.id).select("*").columns == ["id", "x", "id", "y"] + assert joined.limit(5).select(col("*")).columns == ["id", "x", "id", "y"] + assert _rows(joined.distinct().select("*")) == [(1, "a", None, None), (None, None, 2, "b")] + + def test_drop_after_a_sort_of_a_filtered_cross_join(self, spark): + a = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "v"]) + b = spark.createDataFrame([(2, "x"), (3, "y")], ["id", "w"]) + df = a.crossJoin(b).filter(a.id > 1).sort("w").drop("id") + assert df.columns == ["v", "w"] + assert df.collect() == [Row(v="b", w="x"), Row(v="b", w="y")] + + def test_drop_duplicates_keeps_repeated_names(self, unmatched): + derived = unmatched[2].limit(10) + assert derived.dropDuplicates([]).columns == ["id", "x", "id", "y"] + assert derived.dropDuplicates().columns == ["id", "x", "id", "y"] + + def test_input_columns_inside_expressions_over_a_join_with_repeated_names(self, spark, unmatched): + other = spark.createDataFrame([(9, "c")], ["c", "z"]) + joined = unmatched[2].limit(10).crossJoin(other) + assert _rows(joined.filter(other.c == 9).select("x")) == [("a",), (None,)] + assert _rows(joined.withColumn("n", other.c + 1).select("n")) == [(10,), (10,)] + assert joined.groupBy(other.z).count().collect() == [Row(z="c", count=2)] + assert _rows(joined.orderBy(other.c).select("y")) == [("b",), (None,)] + + def test_plain_input_references_through_a_limit_of_a_join(self, spark): + a = spark.createDataFrame([(1, 10, "a")], ["x", "u", "v"]) + b = spark.createDataFrame([(1, "p")], ["x", "w"]) + limited = a.join(b, a.x == b.x).orderBy(a.x).limit(1) + assert limited.withColumn("q", a.v).select("q").collect() == [Row(q="a")] + assert limited.withColumns({"q": a.u}).select("q").collect() == [Row(q=10)] + assert limited.select(a.u.alias("z"), b.w).collect() == [Row(z=10, w="p")] + + def test_qualifier_before_a_column_of_that_name(self, spark): + # The frames have a column k too; Spark resolves k.id through the alias k first. + c = spark.createDataFrame([(10, 20)], ["k", "id"]) + joined = c.alias("k").join(c.alias("m"), col("k.id") == col("m.id")) + assert joined.select(col("k.id")).collect() == [Row(id=20)] + assert joined.select("k.id", "m.k").collect() == [Row(id=20, k=10)] + + def test_names_that_differ_only_in_case_mapping(self, spark): + # Spark compares names ignoring case letter by letter, so these two names differ. + df = spark.createDataFrame([(1, 2)], ["straße", "STRASSE"]) + assert df.select("straße").collect() == [Row(straße=1)] + assert df.select(df["STRASSE"], col("straße")).collect() == [Row(STRASSE=2, straße=1)] + + +class TestQualifiedStar: + def test_star_of_a_struct_column(self, spark): + df = spark.createDataFrame([(1, "a")], ["id", "v"]).select(sf.struct("id", "v").alias("x")) + assert df.select("x.*").collect() == [Row(id=1, v="a")] + + def test_star_of_a_frame_alias(self, spark): + a = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "v"]) + b = spark.createDataFrame([(2, "x")], ["id", "w"]) + joined = a.alias("a").join(b.alias("b"), col("a.id") == col("b.id")) + assert joined.select("a.*").collect() == [Row(id=2, v="b")] + assert joined.select("b.*", "v").collect() == [Row(id=2, w="x", v="b")] + assert a.alias("q").select("q.*").columns == ["id", "v"] + + def test_star_of_a_frame_alias_with_repeated_names(self, spark, unmatched): + other = spark.createDataFrame([(9,)], ["c"]) + df = unmatched[2].alias("q").crossJoin(other).select("q.*") + assert df.columns == ["id", "x", "id", "y"] + assert _rows(df.withColumn("n", lit(1))) == [(1, "a", None, None, 1), (None, None, 2, "b", 1)] + + def test_to_arrow_keeps_repeated_names(self, spark, unmatched): + if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3."): + pytest.skip("toArrow is new in Spark 4.0") + pytest.importorskip("pyarrow") + other = spark.createDataFrame([(9,)], ["c"]) + assert unmatched[2].limit(10).crossJoin(other).toArrow().column_names == ["id", "x", "id", "y", "c"] diff --git a/tests/fast/spark/test_spark_pandas_dataframe.py b/tests/fast/spark/test_spark_pandas_dataframe.py index 4e468d291..542e3f532 100644 --- a/tests/fast/spark/test_spark_pandas_dataframe.py +++ b/tests/fast/spark/test_spark_pandas_dataframe.py @@ -1,10 +1,16 @@ +import datetime +import math + import pytest _ = pytest.importorskip("duckdb.experimental.spark") pd = pytest.importorskip("pandas") from pandas.testing import assert_frame_equal +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import PySparkTypeError, PySparkValueError from spark_namespace.sql.types import ( + ArrayType, IntegerType, Row, StringType, @@ -12,6 +18,9 @@ StructType, ) +if USE_ACTUAL_SPARK: + import pyspark + @pytest.fixture def pandasDF(spark): @@ -49,3 +58,188 @@ def test_spark_to_pandas_dataframe(self, spark, pandasDF): sparkDF = spark.createDataFrame(pandasDF) res = sparkDF.toPandas() assert_frame_equal(res, pandasDF) + + def test_categorical_column(self, spark): + pdf = pd.DataFrame({"c": pd.Categorical(["a", "b", None, "a"])}) + df = spark.createDataFrame(pdf) + assert df.dtypes == [("c", "string")] + # Spark 3.5 stores the text of the missing value, as it does for other pandas missing values. + missing = "NaN" if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3.") else None + assert [row.c for row in df.collect()] == ["a", "b", missing, "a"] + + +def _local(value): + return value.astimezone().replace(tzinfo=None) + + +class TestPandasTimestamps: + def test_naive_datetimes_are_timestamps(self, spark): + values = [datetime.datetime(2024, 1, 1, 12), datetime.datetime(2024, 7, 1, 12), None] + df = spark.createDataFrame(pd.DataFrame({"t": values, "i": [1, 2, 3]})) + assert df.dtypes == [("t", "timestamp"), ("i", "bigint")] + assert df.collect() == [Row(t=values[0], i=1), Row(t=values[1], i=2), Row(t=None, i=3)] + assert df.toPandas()["t"].tolist()[:2] == values[:2] + + @pytest.mark.parametrize("unit", ["s", "ms", "us", "ns"]) + def test_naive_datetime_units(self, spark, unit): + series = pd.Series([datetime.datetime(2024, 1, 1, 12)]).astype(f"datetime64[{unit}]") + df = spark.createDataFrame(pd.DataFrame({"t": series})) + assert df.dtypes == [("t", "timestamp")] + assert df.collect() == [Row(t=datetime.datetime(2024, 1, 1, 12))] + + @pytest.mark.parametrize("zone", ["UTC", "Asia/Tokyo"]) + def test_aware_datetimes(self, spark, zone): + series = pd.Series([datetime.datetime(2024, 1, 1, 12)]).dt.tz_localize(zone) + df = spark.createDataFrame(pd.DataFrame({"t": series})) + expected = _local(series[0].to_pydatetime()) + assert df.dtypes == [("t", "timestamp")] + assert df.collect() == [Row(t=expected)] + assert df.toPandas()["t"].tolist() == [expected] + + def test_to_pandas_gives_naive_local_times(self, spark): + value = datetime.datetime(2024, 1, 1, 12, tzinfo=datetime.timezone.utc) + df = spark.createDataFrame([(value, datetime.datetime(2024, 1, 1, 12))], "t timestamp, n timestamp_ntz") + res = df.toPandas() + assert res["t"].dt.tz is None + assert res["t"].tolist() == [_local(value)] + assert res["n"].tolist() == [datetime.datetime(2024, 1, 1, 12)] + + def test_to_pandas_after_a_join(self, spark): + left = spark.createDataFrame([(1, datetime.datetime(2024, 1, 1, 12))], "id int, t timestamp") + right = spark.createDataFrame([(1, datetime.datetime(2024, 2, 1, 12))], "id int, t timestamp") + res = left.join(right, left.id == right.id).toPandas() + assert res.values.tolist() == [[1, datetime.datetime(2024, 1, 1, 12), 1, datetime.datetime(2024, 2, 1, 12)]] + + def test_to_pandas_of_dates(self, spark): + df = spark.createDataFrame([(datetime.date(2024, 1, 2), 1), (None, 2)], "d date, i int") + res = df.toPandas() + assert str(res["d"].dtype) == "object" + assert res["d"].tolist() == [datetime.date(2024, 1, 2), None] + missing = spark.createDataFrame([(None,)], "d date").toPandas() + assert str(missing["d"].dtype) == "object" + assert missing["d"].tolist() == [None] + res = spark.createDataFrame( + [(datetime.date(1, 1, 1), {"k": datetime.date(9999, 12, 31)})], "d date, m map" + ) + assert res.toPandas().values.tolist() == [[datetime.date(1, 1, 1), {"k": datetime.date(9999, 12, 31)}]] + + def test_to_pandas_keeps_repeated_labels(self, spark): + df = spark.createDataFrame([(1, 2)], ["a", "b"]) + assert df.toDF("x", "x").toPandas().columns.tolist() == ["x", "x"] + assert df.toDF("x", "X").toPandas().columns.tolist() == ["x", "X"] + left = spark.createDataFrame([(1, "l")], ["id", "v"]) + right = spark.createDataFrame([(1, "r")], ["id", "v"]) + res = left.join(right, left.id == right.id).toPandas() + assert res.columns.tolist() == ["id", "v", "id", "v"] + assert res.values.tolist() == [[1, "l", 1, "r"]] + + +class TestPandasNullability: + @pytest.mark.parametrize( + ("data", "schema"), + [ + ({"v": [None]}, "v string not null"), + ({"a": ["x"], "v": [None]}, "a string not null, v string not null"), + ({"v": pd.Series([1, None], dtype="object")}, "v bigint not null"), + ({"v": [pd.NaT, pd.Timestamp("2024-01-01")]}, "v timestamp not null"), + ], + ) + def test_none_in_a_field_that_is_not_nullable(self, spark, data, schema): + with pytest.raises(PySparkValueError): + spark.createDataFrame(pd.DataFrame(data), schema) + + def test_values_in_fields_that_are_not_nullable(self, spark): + df = spark.createDataFrame(pd.DataFrame({"a": ["x", "y"], "v": [None, "z"]}), "a string not null, v string") + assert df.collect() == [Row(a="x", v=None), Row(a="y", v="z")] + floats = pd.DataFrame({"v": [1.0, None]}) + if USE_ACTUAL_SPARK: + # Spark keeps a NaN, which is how pandas stores a missing float. + (_, row) = spark.createDataFrame(floats, "v double not null").collect() + assert math.isnan(row.v) + else: + # DuckDB reads a NaN in a pandas float column as NULL, which the field does not allow. + for schema in ("v double not null", "v float not null"): + with pytest.raises(PySparkValueError): + spark.createDataFrame(floats, schema) + + @pytest.mark.parametrize( + "values", [pd.Series([float("nan"), "a"], dtype=object), pd.Series([pd.NA, "a"], dtype=object), [1.5, None]] + ) + def test_missing_values_in_a_string_field_that_is_not_nullable(self, spark, values): + df = pd.DataFrame({"v": values}) + if USE_ACTUAL_SPARK: + # Spark keeps the text of the missing value. + assert None not in [row.v for row in spark.createDataFrame(df, "v string not null").collect()] + else: + # DuckDB would read the missing value as NULL. + with pytest.raises(PySparkValueError): + spark.createDataFrame(df, "v string not null") + + def test_none_in_a_nested_field_that_is_not_nullable(self, spark): + with pytest.raises(PySparkValueError): + spark.createDataFrame(pd.DataFrame({"s": [{"a": None}]}), "s struct") + elements_not_null = StructType([StructField("l", ArrayType(IntegerType(), False))]) + with pytest.raises(PySparkValueError): + spark.createDataFrame(pd.DataFrame({"l": [[1, None]]}), elements_not_null) + df = spark.createDataFrame(pd.DataFrame({"s": [{"a": 1}, None]}), "s struct") + assert df.collect() == [Row(s=Row(a=1)), Row(s=None)] + + +class TestPandasValueTypes: + @pytest.mark.parametrize( + ("value", "schema"), + [ + (["001"], "v array"), + (["001"], "v array not null"), + ({"a": "001"}, "v struct"), + ({"k": "001"}, "v map"), + ([1, "x"], "v array"), + ], + ) + def test_nested_value_of_another_type(self, spark, value, schema): + with pytest.raises(PySparkTypeError): + spark.createDataFrame(pd.DataFrame({"v": [value]}), schema) + + def test_nested_values_that_fit(self, spark): + stamp = pd.Timestamp("2024-01-02 03:04:05") + df = pd.DataFrame({"v": [[1, None], None], "s": [{"a": "x"}, {"a": None}], "t": [[stamp], [None]]}) + res = spark.createDataFrame(df, "v array, s struct, t array") + assert res.collect() == [ + Row(v=[1, None], s=Row(a="x"), t=[stamp.to_pydatetime()]), + Row(v=None, s=Row(a=None), t=[None]), + ] + + def test_integers_for_a_double_field(self, spark): + df = pd.DataFrame({"v": [[1, 2]], "d": pd.Series([1], dtype=object)}) + if USE_ACTUAL_SPARK and not pyspark.__version__.startswith("3."): + # Spark 4 converts the integers; the module follows Spark 3.5's verifier. + assert spark.createDataFrame(df, "v array, d double").collect() == [Row(v=[1.0, 2.0], d=1.0)] + else: + for schema in ("v array, d string", "v string, d double"): + with pytest.raises(PySparkTypeError): + spark.createDataFrame(df, schema) + + @pytest.mark.parametrize( + ("values", "schema"), + [ + ([1.5], "a int"), + ([1.0, None], "a bigint not null"), + ([1.5], "a decimal(10,2)"), + ], + ) + def test_float_that_does_not_fit(self, spark, values, schema): + with pytest.raises(PySparkTypeError): + spark.createDataFrame(pd.DataFrame({"a": values}), schema) + + def test_whole_floats_for_an_integral_field(self, spark): + data = pd.DataFrame({"a": [2.0, None]}) + if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3."): + # Spark 3.5 rejects every float for an integral field; the module follows Spark 4. + with pytest.raises(PySparkTypeError): + spark.createDataFrame(data, "a bigint") + else: + assert spark.createDataFrame(data, "a bigint").collect() == [Row(a=2), Row(a=None)] + + def test_values_that_fit(self, spark): + assert spark.createDataFrame(pd.DataFrame({"a": [1, 2]}), "a bigint").collect() == [Row(a=1), Row(a=2)] + assert spark.createDataFrame(pd.DataFrame({"a": [1.5]}), "a double").collect() == [Row(a=1.5)] diff --git a/tests/fast/spark/test_spark_readcsv.py b/tests/fast/spark/test_spark_readcsv.py index 27cee47e4..cbca83ad4 100644 --- a/tests/fast/spark/test_spark_readcsv.py +++ b/tests/fast/spark/test_spark_readcsv.py @@ -1,20 +1,679 @@ +import datetime +import math +import sys +from decimal import Decimal + import pytest _ = pytest.importorskip("duckdb.experimental.spark") from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException, error_condition from spark_namespace.sql.types import Row +import duckdb +from duckdb.experimental.spark.exception import ContributionsAcceptedError + +if USE_ACTUAL_SPARK: + import pyspark + + SPARK_3 = pyspark.__version__.startswith("3.") +else: + SPARK_3 = False + + +def _read(spark, tmp_path, text, **options): + path = tmp_path / "data.csv" + path.write_text(text, newline="", encoding="utf-8") + return spark.read.csv(path.as_posix(), **options) + + +def _rows(df): + return [tuple(row) for row in df.collect()] + + +def _typed(spark, tmp_path, data_type, texts, **options): + """Each text read as ``data_type`` from a line of its own, in the order of the texts.""" + text = "".join(f"{i},{t}\n" for i, t in enumerate(texts)) + df = _read(spark, tmp_path, text, schema=f"k int, v {data_type}", **options) + return [row.v for row in sorted(df.collect())] + + +def _inferred(spark, tmp_path, *texts, **options): + text = "v\n" + "".join(f"{t}\n" for t in texts) + df = _read(spark, tmp_path, text, header=True, inferSchema=True, **options) + return df.dtypes[0][1], [row.v for row in df.collect()] + class TestSparkReadCSV: def test_read_csv(self, spark, tmp_path): file_path = tmp_path / "basic.csv" - file_path.write_text("1,2\n3,4\n5,6\n") + file_path.write_text("1,2\n3,4\n5,6\n", newline="", encoding="utf-8") df = spark.read.csv(file_path.as_posix()) res = df.collect() - expected_res = sorted([Row(column0=1, column1=2), Row(column0=3, column1=4), Row(column0=5, column1=6)]) + assert df.columns == ["_c0", "_c1"] + # Spark reads every column as a string unless inferSchema is set. + assert sorted(res) == [ + Row(_c0="1", _c1="2"), + Row(_c0="3", _c1="4"), + Row(_c0="5", _c1="6"), + ] + + def test_read_csv_inferring_types(self, spark, tmp_path): + file_path = tmp_path / "basic.csv" + file_path.write_text("1,2\n3,4\n", newline="", encoding="utf-8") + df = spark.read.csv(file_path.as_posix(), inferSchema=True) + assert df.columns == ["_c0", "_c1"] + assert sorted(df.collect()) == [Row(_c0=1, _c1=2), Row(_c0=3, _c1=4)] + + +# Spark reads the fields of a line by position, NULL past its end, and drops the fields past the last column. +_RAGGED = "a,b\n1,2\n3\n4,5,6\n" + + +class TestLinesOfAnotherWidth: + def test_wider_line_after_the_header(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, "a,b\n1,2,3\n", schema="a int, b int", header=True)) == [(1, 2)] + + @pytest.mark.parametrize( + ("schema", "rows"), + [ + ("a int, b int", [(1, 2), (3, None), (4, 5)]), + ("a int, b int, c int", [(1, 2, None), (3, None, None), (4, 5, 6)]), + ("a int", [(1,), (3,), (4,)]), + ], + ) + def test_with_a_schema(self, spark, tmp_path, schema, rows): + assert _rows(_read(spark, tmp_path, _RAGGED, schema=schema, header=True)) == rows + + def test_with_a_schema_and_no_header(self, spark, tmp_path): + df = _read(spark, tmp_path, _RAGGED, schema="a string, b string") + assert _rows(df) == [("a", "b"), ("1", "2"), ("3", None), ("4", "5")] + + def test_without_a_schema(self, spark, tmp_path): + df = _read(spark, tmp_path, _RAGGED, header=True) + assert df.columns == ["a", "b"] + assert _rows(df) == [("1", "2"), ("3", None), ("4", "5")] + df = _read(spark, tmp_path, _RAGGED) + assert df.columns == ["_c0", "_c1"] + assert _rows(df) == [("a", "b"), ("1", "2"), ("3", None), ("4", "5")] + + def test_first_line_narrower_than_a_later_one(self, spark, tmp_path): + df = _read(spark, tmp_path, "1\n2,3\n") + assert df.columns == ["_c0"] + assert _rows(df) == [("1",), ("2",)] + + def test_header_wider_than_the_data(self, spark, tmp_path): + df = _read(spark, tmp_path, "a,b,c\n1,2\n3,4\n", header=True) + assert df.columns == ["a", "b", "c"] + assert _rows(df) == [("1", "2", None), ("3", "4", None)] + + def test_inferring_types(self, spark, tmp_path): + df = _read(spark, tmp_path, "a,b\n1,x\n3\n", header=True, inferSchema=True) + assert df.dtypes == [("a", "int" if USE_ACTUAL_SPARK else "bigint"), ("b", "string")] + assert _rows(df) == [(1, "x"), (3, None)] + + +class TestDialect: + def test_separator_is_a_comma_unless_given(self, spark, tmp_path): + df = _read(spark, tmp_path, "a;b\n1;2\n", header=True) + assert df.columns == ["a;b"] + assert _rows(df) == [("1;2",)] + assert _rows(_read(spark, tmp_path, "a;b\n1;2\n", header=True, sep=";")) == [("1", "2")] + + def test_backslash_escape(self, spark, tmp_path): + text = 'a,b\n"p\\"q",2\n' + assert _rows(_read(spark, tmp_path, text, header=True, escape="\\")) == [('p"q', "2")] + df = _read(spark, tmp_path, text, schema="a string, b int", header=True, escape="\\") + assert _rows(df) == [('p"q', 2)] + if USE_ACTUAL_SPARK: + assert _rows(_read(spark, tmp_path, text, header=True)) == [('p"q', "2")] + else: + # DuckDB's escape is the quote, so its sniffer fails on a backslash before a quote, and its lenient + # reader reads the line as one field. + with pytest.raises(duckdb.InvalidInputException, match="Error when sniffing file"): + _read(spark, tmp_path, text, header=True) + assert _rows(_read(spark, tmp_path, text, schema="a string, b int", header=True)) == [("p\\q,2", None)] + + def test_doubled_quote(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, 'a,b\n"x""y",2\n', header=True, escape='"')) == [('x"y', "2")] + # Spark's escape is a backslash, so it reads the field as it is written. + expected = '"x""y"' if USE_ACTUAL_SPARK else 'x"y' + assert _rows(_read(spark, tmp_path, '"x""y",1\n')) == [(expected, "1")] + + def test_backslashes_are_text(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, 'C:\\temp,1\n"C:\\temp",2\n')) == [("C:\\temp", "1"), ("C:\\temp", "2")] + assert _rows(_read(spark, tmp_path, 'C:\\temp,a\\"b,x\\\n')) == [("C:\\temp", 'a\\"b', "x\\")] + rows = _rows(_read(spark, tmp_path, 'C:\\temp,1\n"C:\\temp",2\n', escape="\\")) + # DuckDB drops a backslash escape that escapes neither the quote nor itself. + kept = "C:\\temp" if USE_ACTUAL_SPARK else "C:temp" + assert rows == [(kept, "1"), (kept, "2")] + + def test_doubled_backslash_inside_quotes(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, '"a\\\\b","c\\\\"\n', escape="\\")) == [("a\\b", "c\\")] + expected = ("a\\b", "c\\") if USE_ACTUAL_SPARK else ("a\\\\b", "c\\\\") + assert _rows(_read(spark, tmp_path, '"a\\\\b","c\\\\"\n')) == [expected] + + def test_escape_follows_a_given_quote(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, "'a,b',\"c\"\n", quote="'")) == [("a,b", '"c"')] + + def test_mixed_line_ends(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, "a,b\r\n1,2\n3,4\r\n")) == [("a", "b"), ("1", "2"), ("3", "4")] + + def test_unclosed_quote(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, '"ab,2\n', schema="a string, b string")) == [("ab,2", None)] + if USE_ACTUAL_SPARK: + assert _rows(_read(spark, tmp_path, '"ab,2\n')) == [("ab,2",)] + else: + # Without a schema the width of a line comes from DuckDB's sniffer, which fails here. + with pytest.raises(duckdb.InvalidInputException, match="Error when sniffing file"): + _read(spark, tmp_path, '"ab,2\n') + rows = _rows(_read(spark, tmp_path, 'a,b\n"1,2\n3,4\n5,6\n', header=True, schema="a string, b string")) + # Spark ends an unclosed quote at the end of the line, DuckDB's lenient reader at the end of the file. + expected = [("1,2", None), ("3", "4"), ("5", "6")] if USE_ACTUAL_SPARK else [("1,2\n3,4\n5,6", None)] + assert rows == expected + + def test_wide_line_after_the_sniffer_sample(self, spark, tmp_path): + text = "a,b\n" + "i,1\n" * 3000 + "9,1,extra\n" + df = _read(spark, tmp_path, text, header=True) + assert df.columns == ["a", "b"] + assert _rows(df)[-1] == ("9", "1") + assert _rows(_read(spark, tmp_path, text, header=True, schema="a string, b int"))[-1] == ("9", 1) + + def test_header_names(self, spark, tmp_path): + # Spark names an empty header field after its position and adds the position to names that repeat. + df = _read(spark, tmp_path, "a,,A,b\n1,2,3,4\n", header=True) + assert df.columns == ["a0", "_c1", "A2", "b"] + assert _rows(df) == [("1", "2", "3", "4")] + assert _read(spark, tmp_path, "a,a,b,A\n1,2,3,4\n", header=True).columns == ["a0", "a1", "b", "A3"] + + def test_header_only(self, spark, tmp_path): + df = _read(spark, tmp_path, "a,b\n", header=True) + assert df.columns == ["a", "b"] + assert df.collect() == [] + + def test_header_line_ends(self, spark, tmp_path): + df = _read(spark, tmp_path, "a,b\r\n1,2\r\n", header=True) + assert df.columns == ["a", "b"] + assert _rows(df) == [("1", "2")] + + @pytest.mark.parametrize( + ("text", "columns", "rows"), + [ + ("\n\na,b\n1,2\n", ["a", "b"], [("a", "b"), ("1", "2")]), + (" \n\n \na,b\n1,2\n", [" "], [(None,), (" ",), ("a",), ("1",)]), + ], + ) + def test_header_after_blank_lines(self, spark, tmp_path, text, columns, rows): + df = _read(spark, tmp_path, text, header=True) + if USE_ACTUAL_SPARK: + assert (df.columns, _rows(df)) == (["a", "b"], [("1", "2")]) + else: + # The names come from the first line that is not empty, and DuckDB skips the first line as the header, + # empty or not; a line of whitespace is a row. + assert (df.columns, _rows(df)) == (columns, rows) + + def test_blank_lines(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, "a,b\n\n1,2\n\n")) == [("a", "b"), ("1", "2")] + # Spark skips a line of whitespace and an empty line; DuckDB reads the first as a row, and the second as a + # NULL row in a file of one column. + spaces = _rows(_read(spark, tmp_path, "a,b\n \n1,2\n")) + assert spaces == ([("a", "b"), ("1", "2")] if USE_ACTUAL_SPARK else [("a", "b"), (" ", None), ("1", "2")]) + one_column = _rows(_read(spark, tmp_path, "a\n\n1\n \n2\n")) + expected = [("a",), ("1",), ("2",)] if USE_ACTUAL_SPARK else [("a",), (None,), ("1",), (" ",), ("2",)] + assert one_column == expected + + def test_line_of_separators_is_a_row(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, "a,b\n,\n")) == [("a", "b"), (None, None)] + + def test_lines_of_whitespace_and_separators(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, "a,b\n ,\n1,2\n")) == [("a", "b"), (" ", None), ("1", "2")] + rows = _rows(_read(spark, tmp_path, "a,b\n ,\n1,2\n", nullValue="NA")) + assert rows == [("a", "b"), (" ", None), ("1", "2")] + assert _rows(_read(spark, tmp_path, 'x\n" "\ny\n')) == [("x",), (" ",), ("y",)] + df = _read(spark, tmp_path, " ,,\n1,2,3\n", header=True) + assert df.columns == [" ", "_c1", "_c2"] + assert _rows(df) == [("1", "2", "3")] + + def test_empty_first_field(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, "1,2\n,3\n4,5\n", schema="a string")) == [("1",), (None,), ("4",)] + assert _rows(_read(spark, tmp_path, "a\n,b\n")) == [("a",), (None,)] + assert _rows(_read(spark, tmp_path, "h\n,3\n4,5\n", header=True)) == [(None,), ("4",)] + + def test_lines_equal_to_the_header(self, spark, tmp_path): + rows = _rows(_read(spark, tmp_path, "a,b\na,b\na,b\n1,2\n", header=True)) + assert rows == [("a", "b"), ("a", "b"), ("1", "2")] + rows = _rows(_read(spark, tmp_path, "a,1\na,2\na,3\nb,4\n", header=True, schema="x string")) + assert rows == [("a",), ("a",), ("b",)] + assert _read(spark, tmp_path, "x\n" * 1000).count() == 1000 + + def test_files_of_a_glob(self, spark, tmp_path): + (tmp_path / "a.csv").write_text("h1,h2\n1,x\n2,y\n", newline="", encoding="utf-8") + (tmp_path / "b.csv").write_text("h1,h2\n3,z\n", newline="", encoding="utf-8") + path = (tmp_path / "*.csv").as_posix() + assert sorted(_rows(spark.read.csv(path, header=True))) == [("1", "x"), ("2", "y"), ("3", "z")] + df = spark.read.csv(path, header=True, inferSchema=True) + assert df.dtypes == [("h1", "int" if USE_ACTUAL_SPARK else "bigint"), ("h2", "string")] + assert sorted(_rows(df)) == [(1, "x"), (2, "y"), (3, "z")] + + def test_quoted_line_break(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, 'a,b\n"x\ny",2\n', header=True, multiLine=True)) == [("x\ny", "2")] + if USE_ACTUAL_SPARK: + # Without multiLine Spark reads each line as a record of its own. + assert _rows(_read(spark, tmp_path, '"x\ny"\n')) == [("x",), ('y"',)] + else: + with pytest.raises(Exception, match="quoted new lines"): + _rows(_read(spark, tmp_path, '"x\ny"\n')) + + +class TestNullValues: + def test_empty_fields(self, spark, tmp_path): + # An empty field is NULL, and a quoted one the empty string, which is NULL for the default nullValue only. + assert _rows(_read(spark, tmp_path, 'a,"",c\n')) == [("a", None, "c")] + assert _rows(_read(spark, tmp_path, 'a,"",,NA,"NA"\n', nullValue="NA")) == [("a", "", None, None, None)] + + def test_null_value_with_a_header(self, spark, tmp_path): + df = _read(spark, tmp_path, "a,b\n1,NA\n,2\n", header=True, nullValue="NA") + assert _rows(df) == [("1", None), (None, "2")] + df = _read(spark, tmp_path, "a,b\n1,NA\n,2\n", header=True, nullValue="NA", schema="a int, b int") + assert _rows(df) == [(1, None), (None, 2)] + + def test_header_field_equal_to_the_null_value(self, spark, tmp_path): + df = _read(spark, tmp_path, "a,NA,b\n1,NA,3\n", header=True, nullValue="NA") + assert df.columns == ["a", "_c1", "b"] + assert _rows(df) == [("1", None, "3")] + + +class TestTypedValues: + """A field of a schema's type converts with DuckDB's cast, NULL where it does not convert, as in Spark.""" + + def test_text_that_does_not_convert_is_null(self, spark, tmp_path): + # Spark keeps the other fields of the line. + assert _rows(_read(spark, tmp_path, "x,1\n2,y\n", schema="a int, b int")) == [(None, 1), (2, None)] + assert _rows(_read(spark, tmp_path, "x,y\n1,x\n", schema="x int, y string")) == [(None, "y"), (1, "x")] + + def test_integers(self, spark, tmp_path): + texts = ["1", "+1", "-0", "007", " 1", "1 ", "1.0", "1e3", "0x1F", "1_000", "2147483648"] + spark_values = [1, 1, 0, 7] + [None] * 7 + duckdb_values = [1, 1, 0, 7, 1, 1, 1, 1000, 31, 1000, None] + assert _typed(spark, tmp_path, "int", texts) == (spark_values if USE_ACTUAL_SPARK else duckdb_values) + assert _typed(spark, tmp_path, "bigint", ["2147483648", "9223372036854775808"]) == [2147483648, None] + assert _typed(spark, tmp_path, "tinyint", ["127", "128", "-128"]) == [127, None, -128] + assert _typed(spark, tmp_path, "smallint", ["32767", "32768"]) == [32767, None] + + def test_doubles(self, spark, tmp_path): + texts = ["1.5", " 1.5 ", "1e3", ".5", "5.", "1E400", "Infinity", "-Infinity", "+Infinity"] + inf = math.inf + assert _typed(spark, tmp_path, "double", texts) == [1.5, 1.5, 1000.0, 0.5, 5.0, inf, inf, -inf, inf] + assert _typed(spark, tmp_path, "double", ["1.5e", '"1,5"', "x"]) == [None] * 3 + values = _typed(spark, tmp_path, "double", ["1.5d", "1.5f", "0x1p3", "1_0", "+Inf", "inf"]) + # Java's parseDouble reads a type suffix and hexadecimal text; DuckDB reads digit separators and more + # spellings of infinity. + assert values == ([1.5, 1.5, 8.0, None, None, None] if USE_ACTUAL_SPARK else [None, None, None, 10.0, inf, inf]) + (nan,) = _typed(spark, tmp_path, "double", ["nan"]) + assert nan is None if USE_ACTUAL_SPARK else math.isnan(nan) + + def test_nan_and_infinity_texts(self, spark, tmp_path): + nan, inf, minus_inf = _typed(spark, tmp_path, "double", ["NaN", "Inf", "-Inf"]) + assert math.isnan(nan) + assert (inf, minus_inf) == (math.inf, -math.inf) + options = {"nanValue": "NA", "positiveInf": "+", "negativeInf": "-"} + values = _typed(spark, tmp_path, "double", ["NA", "+", "-", "Inf", "NaN"], **options) + assert math.isnan(values[0]) + assert math.isnan(values[4]) + assert values[1:3] == [math.inf, -math.inf] + # Spark reads Inf only as its positiveInf text; DuckDB's cast reads it whatever positiveInf is. + assert values[3] == (None if USE_ACTUAL_SPARK else math.inf) + values = _typed(spark, tmp_path, "float", ["1.5", "3.4e39", "NaN"]) + assert values[:2] == [1.5, math.inf] + assert math.isnan(values[2]) + + def test_decimals(self, spark, tmp_path): + texts = ["1.234", "1.235", "-1.235", "1e2", "+1.5", ".5", "123.456", "1234.5"] + expected = ["1.23", "1.24", "-1.24", "100.00", "1.50", "0.50", "123.46", None] + assert _typed(spark, tmp_path, "decimal(5,2)", texts) == [None if e is None else Decimal(e) for e in expected] + assert _typed(spark, tmp_path, "decimal(5,2)", ["1.5d", "abc", '"1,000"']) == [None] * 3 + # Spark leaves out the commas its default locale groups digits with, and does not trim; DuckDB trims. + spark_values, duckdb_values = [Decimal("15.00"), None], [None, Decimal("1.50")] + values = _typed(spark, tmp_path, "decimal(5,2)", ['"1,5"', " 1.5"]) + assert values == (spark_values if USE_ACTUAL_SPARK else duckdb_values) + + def test_booleans(self, spark, tmp_path): + assert _typed(spark, tmp_path, "boolean", ["true", "TRUE", "False", " true"]) == [True, True, False, None] + values = _typed(spark, tmp_path, "boolean", ["t", "yes", "1"]) + assert values == ([None] * 3 if USE_ACTUAL_SPARK else [True] * 3) + + def test_dates(self, spark, tmp_path): + texts = ["2024-01-02", "2024-1-2", "2024-01-02 12:00:00", "2024-01-02T12:00", " 2024-01-02 "] + assert _typed(spark, tmp_path, "date", texts) == [datetime.date(2024, 1, 2)] * 5 + # Spark reads a year alone and a year and month. + values = _typed(spark, tmp_path, "date", ["2024", "2024-01"]) + assert values == ([datetime.date(2024, 1, 1)] * 2 if USE_ACTUAL_SPARK else [None, None]) + + def test_invalid_dates(self, spark, tmp_path): + if SPARK_3: + pytest.skip("Spark 3.5 raises on text its legacy date parser reads differently") + texts = ["02/01/2024", "20240102", "2024-02-30", "2024-01-02Z", "1"] + assert _typed(spark, tmp_path, "date", texts) == [None] * len(texts) + # DuckDB reads slashes and years of fewer digits. + values = _typed(spark, tmp_path, "date", ["2024/01/02", "1-1-1"]) + assert values == ([None, None] if USE_ACTUAL_SPARK else [datetime.date(2024, 1, 2), datetime.date(1, 1, 1)]) + + def test_timestamps(self, spark, tmp_path): + texts = ["2024-01-02 03:04:05", "2024-01-02T03:04:05", "2024-1-2 3:4:5", " 2024-01-02 03:04:05 "] + texts += ["2024-01-02", "2024-01-02 03:04:05.123456789"] + expected = [datetime.datetime(2024, 1, 2, 3, 4, 5)] * 4 + expected += [datetime.datetime(2024, 1, 2), datetime.datetime(2024, 1, 2, 3, 4, 5, 123456)] + assert _typed(spark, tmp_path, "timestamp", texts) == expected + # Spark reads a year and month, and a date and an hour. + values = _typed(spark, tmp_path, "timestamp", ["2024-01", "2024-01-02 03"]) + spark_values = [datetime.datetime(2024, 1, 1), datetime.datetime(2024, 1, 2, 3)] + assert values == (spark_values if USE_ACTUAL_SPARK else [None, None]) + + def test_timestamps_with_a_zone(self, spark, tmp_path): + instant = datetime.datetime(2024, 1, 2, 11, tzinfo=datetime.timezone.utc).astimezone().replace(tzinfo=None) + texts = ["2024-01-02T12:00:00+01:00", "2024-01-02 12:00:00+0100", "2024-01-02 12:00:00+01"] + texts += ["2024-01-02T11:00:00Z", "2024-01-02 12:00:00 Europe/Amsterdam"] + assert _typed(spark, tmp_path, "timestamp", texts) == [instant] * len(texts) + # Spark reads an offset after GMT. + expected = instant if USE_ACTUAL_SPARK else None + assert _typed(spark, tmp_path, "timestamp", ["2024-01-02 12:00:00GMT+1"]) == [expected] + + def test_invalid_timestamps(self, spark, tmp_path): + if SPARK_3: + pytest.skip("Spark 3.5 raises on text its legacy timestamp parser reads differently") + texts = ["2024-01-02 23:59:60", "2024-01-02T", "1704196800", "x"] + assert _typed(spark, tmp_path, "timestamp", texts) == [None] * len(texts) + # DuckDB reads hour 24 as midnight of the next day. + values = _typed(spark, tmp_path, "timestamp", ["2024-01-02 24:00:00"]) + assert values == ([None] if USE_ACTUAL_SPARK else [datetime.datetime(2024, 1, 3)]) + + @pytest.mark.skipif( + sys.platform == "win32", reason="outside the range of the Windows C runtime, as for PySpark on Windows" + ) + def test_two_digit_year_timestamp(self, spark, tmp_path): + if SPARK_3: + pytest.skip("Spark 3.5 raises on text its legacy timestamp parser reads differently") + # DuckDB reads a year of two digits as it is. + values = _typed(spark, tmp_path, "timestamp", ["24-01-02"]) + assert values == ([None] if USE_ACTUAL_SPARK else [datetime.datetime(24, 1, 2)]) + + def test_timestamps_without_a_zone(self, spark, tmp_path): + texts = ["2024-01-02 03:04:05", "2024-01-02T03:04:05"] + assert _typed(spark, tmp_path, "timestamp_ntz", texts) == [datetime.datetime(2024, 1, 2, 3, 4, 5)] * 2 + # Spark rejects text with a zone and reads a year and month; DuckDB drops the zone. + values = _typed(spark, tmp_path, "timestamp_ntz", ["2024-01-02T03:04:05Z", "2024-01-02 03:04+01:00", "2024-01"]) + duckdb_values = [datetime.datetime(2024, 1, 2, 3, 4, 5), datetime.datetime(2024, 1, 2, 3, 4), None] + assert values == ([None, None, datetime.datetime(2024, 1, 1)] if USE_ACTUAL_SPARK else duckdb_values) + + def test_strings(self, spark, tmp_path): + assert _typed(spark, tmp_path, "string", ["x", " x ", '" y "']) == ["x", " x ", " y "] + + def test_binary(self, spark, tmp_path): + if SPARK_3: + with pytest.raises(AnalysisException) as error: + _typed(spark, tmp_path, "binary", ["abc"]) + assert error_condition(error.value) == "UNSUPPORTED_DATA_TYPE_FOR_DATASOURCE" + return + # DuckDB's cast reads a backslash escape of a byte. + expected = b"\\x41" if USE_ACTUAL_SPARK else b"A" + assert _typed(spark, tmp_path, "binary", ["abc", "\\x41"]) == [b"abc", expected] + + def test_intervals(self, spark, tmp_path): + expected = [None, None] if USE_ACTUAL_SPARK else [datetime.timedelta(seconds=1), datetime.timedelta(days=1)] + assert _typed(spark, tmp_path, "interval day to second", ["1", "1 day"]) == expected + + @pytest.mark.parametrize( + ("data_type", "name"), + [ + ("array", "ARRAY"), + ("struct", "STRUCT"), + ("map", "MAP"), + ("void", "VOID"), + ], + ) + def test_types_csv_cannot_hold(self, spark, tmp_path, data_type, name): + with pytest.raises(AnalysisException, match=f'column `a` of the type "{name}"') as error: + _read(spark, tmp_path, "1\n", schema=f"a {data_type}").collect() + assert error_condition(error.value) == "UNSUPPORTED_DATA_TYPE_FOR_DATASOURCE" + + def test_names_differing_in_case_only(self, spark, tmp_path): + with pytest.raises(AnalysisException, match="The column `a` already exists") as error: + _read(spark, tmp_path, "1,2\n", schema="a int, A int").collect() + assert error_condition(error.value) == "COLUMN_ALREADY_EXISTS" + + +class TestInference: + """With inferSchema the types are DuckDB's type detection's, read with DuckDB's cast.""" + + @pytest.mark.parametrize( + ("texts", "spark_type", "duckdb_type"), + [ + (["1", "-0"], "int", "bigint"), + (["1", "2147483648"], "bigint", "bigint"), + (["9223372036854775808"], "decimal(19,0)", "hugeint"), + (["1", "12345678901234567890"], "decimal(20,0)", "hugeint"), + (["123456789012345678901234567890123456789"], "double", "hugeint"), + (["1" * 41], "double", "double"), + (["5.", "1.0E1"], "decimal(2,0)", "double"), + (["1", "1.5", "1e3", "1E10", " 1", "NaN", "-Inf", "Infinity"], "double", "double"), + (["1.5d"], "double", "string"), + (["12345678901234567890", "1.5"], "double", "double"), + (["2024-01-02"], "date", "date"), + (["2024-01-02", "2024-01-02 03:04:05"], "timestamp", "string"), + (["2024-01-02T03:04:05Z"], "timestamp", "timestamp"), + (["2024-1-2"], "timestamp", "date"), + (["2024-01", "2024-01-02 03"], "timestamp", "string"), + (["true", "FALSE"], "boolean", "boolean"), + (["0x1F"], "string", "bigint"), + (["yes", "t"], "string", "boolean"), + (["2024/01/02"], "string", "date"), + (["03/02/2024"], "string", "date"), + (["12:00:00"], "timestamp", "time"), + (["+1", "007"], "int", "string"), + (["1_000", "nan", "inf", '"1,000"', "1.5x"], "string", "string"), + (["1", "true"], "string", "string"), + (["2024-01-02", "1"], "string", "string"), + (["1", ""], "int", "bigint"), + (["", ""], "string", "string"), + ], + ) + def test_types(self, spark, tmp_path, texts, spark_type, duckdb_type): + assert _inferred(spark, tmp_path, *texts)[0] == (spark_type if USE_ACTUAL_SPARK else duckdb_type) + + def test_values(self, spark, tmp_path): + assert _inferred(spark, tmp_path, "1", " 2", "3.5") == ("double", [1.0, 2.0, 3.5]) + times = ("timestamp", [datetime.datetime(2024, 1, 2, 3, 4, 5), datetime.datetime(2024, 1, 3, 4, 5, 6)]) + assert _inferred(spark, tmp_path, "2024-01-02 03:04:05", "2024-01-03T04:05:06") == times + assert _inferred(spark, tmp_path, "2024-01-02")[1] == [datetime.date(2024, 1, 2)] + + def test_values_of_other_types(self, spark, tmp_path): + df = _read(spark, tmp_path, "hex,yn,tf,sl\n0x1F,yes,t,2024/01/02\n", header=True, inferSchema=True) + if USE_ACTUAL_SPARK: + assert df.dtypes == [("hex", "string"), ("yn", "string"), ("tf", "string"), ("sl", "string")] + assert _rows(df) == [("0x1F", "yes", "t", "2024/01/02")] + else: + assert df.dtypes == [("hex", "bigint"), ("yn", "boolean"), ("tf", "boolean"), ("sl", "date")] + assert _rows(df) == [(31, True, True, datetime.date(2024, 1, 2))] + + def test_dates_in_a_detected_format(self, spark, tmp_path): + df = _read(spark, tmp_path, "d\n03/02/2024\n04/02/2024\n", header=True, inferSchema=True) + if USE_ACTUAL_SPARK: + assert _rows(df) == [("03/02/2024",), ("04/02/2024",)] + else: + assert _rows(df) == [(datetime.date(2024, 2, 3),), (datetime.date(2024, 2, 4),)] + + def test_lines_repeating_the_header(self, spark, tmp_path): + df = _read(spark, tmp_path, "v\n1\nv\n2\n", header=True, inferSchema=True) + if USE_ACTUAL_SPARK: + # Spark leaves lines that repeat the header out of the inference. + assert df.dtypes == [("v", "int")] + assert _rows(df) == [(1,), (None,), (2,)] + else: + assert df.dtypes == [("v", "string")] + assert _rows(df) == [("1",), ("v",), ("2",)] + + def test_null_and_nan_values(self, spark, tmp_path): + integer = "int" if USE_ACTUAL_SPARK else "bigint" + assert _inferred(spark, tmp_path, "NA", "1", nullValue="NA") == (integer, [None, 1]) + assert _inferred(spark, tmp_path, '""', "1", "NA", nullValue="NA")[1] == [None, 1, None] + # DuckDB's detection does not know a nanValue. + expected = "double" if USE_ACTUAL_SPARK else "string" + assert _inferred(spark, tmp_path, "NA", "1.5", nanValue="NA")[0] == expected + + def test_values_after_the_sample(self, spark, tmp_path): + """DuckDB detects the types from a sample of the lines and reads the rest as its own reader does.""" + df = _read(spark, tmp_path, "v\n" + "1\n" * 30000 + "x\n", header=True, inferSchema=True) + if USE_ACTUAL_SPARK: + assert df.dtypes == [("v", "string")] + assert _rows(df)[-1] == ("x",) + else: + assert df.dtypes == [("v", "bigint")] + with pytest.raises(duckdb.ConversionException, match='Could not convert string "x"'): + df.collect() + df = _read(spark, tmp_path, "v\n" + "1\n" * 30000 + "1.5\n", header=True, inferSchema=True) + expected = ("double", 1.5) if USE_ACTUAL_SPARK else ("bigint", 2) + assert (df.dtypes[0][1], _rows(df)[-1][0]) == expected + text = "v\n" + "2024-01-02\n" * 30000 + "2024-01-03 00:00:00\n" + df = _read(spark, tmp_path, text, header=True, inferSchema=True) + last = datetime.datetime(2024, 1, 3) + expected = ("timestamp", last) if USE_ACTUAL_SPARK else ("date", last.date()) + assert (df.dtypes[0][1], _rows(df)[-1][0]) == expected + + def test_blank_lines(self, spark, tmp_path): + rows = _rows(_read(spark, tmp_path, "a,b\n1,2\n\n3,4\n", header=True, inferSchema=True)) + assert rows == [(1, 2), (3, 4)] + # DuckDB reads an empty line of a file of one column as a NULL row, and a line of whitespace as a row. + rows = _rows(_read(spark, tmp_path, "v\n1\n\n2\n", header=True, inferSchema=True)) + assert rows == ([(1,), (2,)] if USE_ACTUAL_SPARK else [(1,), (None,), (2,)]) + df = _read(spark, tmp_path, "x,y\n1,2\n \n3,4\n", header=True, inferSchema=True) + integer = "int" if USE_ACTUAL_SPARK else "bigint" + assert df.dtypes == [("x", integer if USE_ACTUAL_SPARK else "string"), ("y", integer)] + + def test_header_of_another_width(self, spark, tmp_path): + df = _read(spark, tmp_path, "v\n1,000\n", header=True, inferSchema=True) + assert df.columns == ["v"] + assert _rows(df) == [(1,)] + + def test_header_without_data(self, spark, tmp_path): + df = _read(spark, tmp_path, "a,b\n", header=True, inferSchema=True) + assert df.dtypes == [("a", "string"), ("b", "string")] + + +class TestDatetimePatterns: + def test_date_format(self, spark, tmp_path): + df = _read(spark, tmp_path, "2024/02/03\nFeb 3\n", schema="d date", dateFormat="yyyy/MM/dd") + # Spark reads text that does not match the pattern as NULL. + assert _rows(df) == [(datetime.date(2024, 2, 3),), (None,)] + df = _read(spark, tmp_path, "2024/2/3\n", schema="d date", dateFormat="yyyy/M/d") + assert _rows(df) == [(datetime.date(2024, 2, 3),)] + + def test_timestamp_format(self, spark, tmp_path): + text = "03.02.2024 04:05:06\n" + df = _read(spark, tmp_path, text, schema="t timestamp", timestampFormat="dd.MM.yyyy HH:mm:ss") + assert _rows(df) == [(datetime.datetime(2024, 2, 3, 4, 5, 6),)] + + def test_timestamp_format_is_not_for_timestamps_without_a_zone(self, spark, tmp_path): + text = "2024-02-03T04:05:06\n" + df = _read(spark, tmp_path, text, schema="t timestamp_ntz", timestampFormat="dd.MM.yyyy HH:mm:ss") + assert _rows(df) == [(datetime.datetime(2024, 2, 3, 4, 5, 6),)] + + def test_patterns_apply_only_to_their_types(self, spark, tmp_path): + df = _read(spark, tmp_path, "2024/02/03,x\n", schema="s string, t string", dateFormat="yyyy/MM/dd") + assert _rows(df) == [("2024/02/03", "x")] + assert _read(spark, tmp_path, "2024/02/03\n", dateFormat="yyyy/MM/dd").dtypes == [("_c0", "string")] + + def test_inferring_with_a_pattern(self, spark, tmp_path): + if USE_ACTUAL_SPARK: + df = _read(spark, tmp_path, "2024/02/03\n", dateFormat="yyyy/MM/dd", inferSchema=True) + assert df.dtypes == [("_c0", "date")] + else: + with pytest.raises(ContributionsAcceptedError): + _read(spark, tmp_path, "2024/02/03\n", dateFormat="yyyy/MM/dd", inferSchema=True) + + +class TestEnforceSchema: + def test_header_that_does_not_match_the_schema(self, spark, tmp_path): + assert _rows(_read(spark, tmp_path, "a\nx\n", schema="b string", header=True)) == [("x",)] + df = _read(spark, tmp_path, "a\nx\n", schema="b string", header=True, enforceSchema="TRUE") + assert _rows(df) == [("x",)] + # With enforceSchema off Spark checks the header against the schema and fails the read. + with pytest.raises(Exception, match="CSV header does not conform to the schema"): + _read(spark, tmp_path, "a\nx\n", schema="b string", header=True, enforceSchema=False).collect() + + def test_header_of_another_width(self, spark, tmp_path): + with pytest.raises(Exception, match="Number of column in CSV header is not equal to number of fields"): + _read(spark, tmp_path, "a,b\n1,2\n", schema="a int", header=True, enforceSchema="false").collect() + + @pytest.mark.parametrize("width", ["a,c", "a,b,c"]) + def test_header_of_each_file_of_a_glob(self, spark, tmp_path, width): + (tmp_path / "1.csv").write_text("a,b\n1,2\n", newline="", encoding="utf-8") + (tmp_path / "2.csv").write_text(f"{width}\n3,4\n", newline="", encoding="utf-8") + path = (tmp_path / "*.csv").as_posix() + df = spark.read.csv(path, header=True, schema="a int, b int") + assert sorted(_rows(df)) == [(1, 2), (3, 4)] + match = "CSV header does not conform" if width == "a,c" else "Number of column in CSV header" + with pytest.raises(Exception, match=match): + spark.read.csv(path, header=True, schema="a int, b int", enforceSchema=False).collect() + + def test_header_of_files_of_a_glob_that_match(self, spark, tmp_path): + (tmp_path / "1.csv").write_text("a,b\n1,2\n", newline="", encoding="utf-8") + (tmp_path / "2.csv").write_text("A,B\n3,4\n", newline="", encoding="utf-8") + df = spark.read.csv((tmp_path / "*.csv").as_posix(), header=True, schema="a int, b int", enforceSchema=False) + assert sorted(_rows(df)) == [(1, 2), (3, 4)] + + def test_header_that_matches_the_schema_ignoring_case(self, spark, tmp_path): + df = _read(spark, tmp_path, "A,b\n1,2\n", schema="a int, B int", header=True, enforceSchema=False) + assert df.columns == ["a", "B"] + assert _rows(df) == [(1, 2)] + + +class TestLoad: + def test_csv_by_its_extension(self, spark, tmp_path): + path = (tmp_path / "data.csv").as_posix() + (tmp_path / "data.csv").write_text("1,2\n3,4\n", newline="", encoding="utf-8") + if USE_ACTUAL_SPARK: + # Spark reads a path without a format as parquet. + with pytest.raises(Exception): # noqa: B017, PT011 + spark.read.load(path, schema="a int, b int").collect() + else: + assert _rows(spark.read.load(path, schema="a int, b int")) == [(1, 2), (3, 4)] + assert _rows(spark.read.load(path)) == [("1", "2"), ("3", "4")] + + def test_header_row_is_data(self, spark, tmp_path): + (tmp_path / "data.csv").write_text("x,y\n1,x\n2,y\n", newline="", encoding="utf-8") + df = spark.read.load((tmp_path / "data.csv").as_posix(), format="csv", schema="x int, y string") + assert _rows(df) == [(None, "y"), (1, "x"), (2, "y")] + + def test_tab_separated(self, spark, tmp_path): + (tmp_path / "data.tsv").write_text("a\tb\n1\t2\n", newline="", encoding="utf-8") + path = (tmp_path / "data.tsv").as_posix() + if USE_ACTUAL_SPARK: + # Spark has no tsv format. + with pytest.raises(Exception): # noqa: B017, PT011 + spark.read.load(path, format="tsv") + else: + assert _rows(spark.read.load(path, format="tsv")) == [("a", "b"), ("1", "2")] + + def test_options(self, spark, tmp_path): + (tmp_path / "data.csv").write_text("a,b\n1,2\n", newline="", encoding="utf-8") + path = (tmp_path / "data.csv").as_posix() if USE_ACTUAL_SPARK: - # Convert all values to strings as this is how Spark reads them by default - expected_res = [Row(column0=str(row.column0), column1=str(row.column1)) for row in expected_res] - assert sorted(res) == expected_res + assert _rows(spark.read.load(path, format="csv", header=True)) == [("1", "2")] + else: + with pytest.raises(ContributionsAcceptedError): + spark.read.load(path, format="csv", header=True) + + +class TestFlags: + @pytest.mark.parametrize("option", ["ignoreLeadingWhiteSpace", "ignoreTrailingWhiteSpace"]) + def test_whitespace_is_kept_unless_asked(self, spark, tmp_path, option): + assert _rows(_read(spark, tmp_path, " a , b \n", **{option: "False"})) == [(" a ", " b ")] + if not USE_ACTUAL_SPARK: + with pytest.raises(ContributionsAcceptedError): + _read(spark, tmp_path, " a \n", **{option: True}) + + @pytest.mark.parametrize("value", [True, "TRUE", "true"]) + def test_multiline_spellings(self, spark, tmp_path, value): + assert _rows(_read(spark, tmp_path, '"x\ny"\n', multiLine=value)) == [("x\ny",)] diff --git a/tests/fast/spark/test_spark_readjson.py b/tests/fast/spark/test_spark_readjson.py index aa8d8ec57..cb677fd46 100644 --- a/tests/fast/spark/test_spark_readjson.py +++ b/tests/fast/spark/test_spark_readjson.py @@ -1,10 +1,21 @@ +import datetime +import math +from decimal import Decimal + import pytest _ = pytest.importorskip("duckdb.experimental.spark") - +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException, error_condition from spark_namespace.sql.types import Row +import duckdb +from duckdb.experimental.spark.exception import ContributionsAcceptedError + +if USE_ACTUAL_SPARK: + import pyspark + class TestSparkReadJson: def test_read_json(self, duckdb_cursor, spark, tmp_path): @@ -14,3 +25,255 @@ def test_read_json(self, duckdb_cursor, spark, tmp_path): df = spark.read.json(file_path) res = df.collect() assert res == [Row(a=42, b=True, c="this is a long string")] + + +def _read(spark, tmp_path, *lines): + path = tmp_path / "data.json" + path.write_text("\n".join(lines) + "\n", newline="", encoding="utf-8") + return spark.read.json(path.as_posix()) + + +class TestJsonInference: + def test_strings_stay_strings(self, spark, tmp_path): + df = _read( + spark, + tmp_path, + '{"d": "2024-01-01", "t": "12:00:00", "ts": "2024-01-01 12:00:00", "z": "2024-01-01T12:00:00Z"}', + '{"u": "123e4567-e89b-12d3-a456-426614174000", "n": "123"}', + ) + assert df.dtypes == [(name, "string") for name in ["d", "n", "t", "ts", "u", "z"]] + assert [tuple(row) for row in df.collect()] == [ + ("2024-01-01", None, "12:00:00", "2024-01-01 12:00:00", None, "2024-01-01T12:00:00Z"), + (None, "123", None, None, "123e4567-e89b-12d3-a456-426614174000", None), + ] + + def test_fields_sorted_by_name(self, spark, tmp_path): + df = _read(spark, tmp_path, '{"z": "x", "a": 1, "C": true}', '{"b": 1.5}') + assert df.dtypes == [("C", "boolean"), ("a", "bigint"), ("b", "double"), ("z", "string")] + assert [tuple(row) for row in df.collect()] == [(True, 1, None, "x"), (None, None, 1.5, None)] + + def test_nested_fields_sorted_by_name(self, spark, tmp_path): + df = _read(spark, tmp_path, '{"s": {"y": 1, "x": "2024-01-01"}, "l": [{"b": 1}, {"a": 2}]}') + assert df.dtypes == [("l", "array>"), ("s", "struct")] + assert df.collect() == [Row(l=[Row(a=None, b=1), Row(a=2, b=None)], s=Row(x="2024-01-01", y=1))] + + def test_null_fields_are_strings(self, spark, tmp_path): + df = _read(spark, tmp_path, '{"n": null, "l": [null], "s": {"m": null}}') + assert df.dtypes == [("l", "array"), ("n", "string"), ("s", "struct")] + + def test_values_of_mixed_types_are_strings(self, spark, tmp_path): + df = _read(spark, tmp_path, '{"v": 1, "l": [1, "x"]}', '{"v": "x", "l": [true]}') + assert df.dtypes == [("l", "array"), ("v", "string")] + assert [tuple(row) for row in df.collect()] == [(["1", "x"], "1"), (["true"], "x")] + + def test_objects_with_many_keys_are_structs(self, spark, tmp_path): + keys = [f"k{i:03d}" for i in range(250)] + df = _read(spark, tmp_path, '{"m": {' + ", ".join(f'"{k}": 1' for k in reversed(keys)) + "}}") + assert df.schema.fields[0].dataType.names == keys + + def test_integers_wider_than_a_long(self, spark, tmp_path): + df = _read(spark, tmp_path, '{"v": 12345678901234567890}') + if USE_ACTUAL_SPARK: + assert df.dtypes == [("v", "decimal(20,0)")] + else: + # DuckDB reads them as HUGEINT, as it does in CSV. + assert df.dtypes == [("v", "hugeint")] + assert int(df.collect()[0].v) == 12345678901234567890 + + +def _load(spark, tmp_path, schema, *lines): + path = tmp_path / "data.json" + path.write_text("\n".join(lines) + "\n", newline="", encoding="utf-8") + return spark.read.load(path.as_posix(), format="json", schema=schema) + + +def _field(spark, tmp_path, data_type, *values): + """Each JSON value read as ``data_type`` from a record of its own, in the order of the values.""" + lines = [f'{{"k": {i}, "v": {value}}}' for i, value in enumerate(values)] + return [row.v for row in sorted(_load(spark, tmp_path, f"k int, v {data_type}", *lines).collect())] + + +class TestJsonWithASchema: + """A field of a schema converts with DuckDB's JSON reader, NULL where a value does not convert, as in Spark.""" + + def test_integers(self, spark, tmp_path): + values = ["1", "-1", "null", '{"x": 1}', "[1]", "2147483648"] + assert _field(spark, tmp_path, "int", *values) == [1, -1, None, None, None, None] + # Spark reads an integer only from a JSON integer; DuckDB also from a fraction, a numeric string or a boolean. + values = _field(spark, tmp_path, "int", "1.7", "1.0", "1e2", '"1"', "true") + assert values == ([None] * 5 if USE_ACTUAL_SPARK else [2, 1, 100, 1, 1]) + assert _field(spark, tmp_path, "bigint", "2147483648", "9223372036854775808") == [2147483648, None] + values = ["127", "128", "255", "256", "-128", "-129"] + # Spark reads a tinyint from -128 to 255, the values above 127 as unsigned bytes. + expected = [127, -128, -1, None, -128, None] if USE_ACTUAL_SPARK else [127, None, None, None, -128, None] + assert _field(spark, tmp_path, "tinyint", *values) == expected + assert _field(spark, tmp_path, "smallint", "32767", "32768", "65535") == [32767, None, None] + + def test_doubles(self, spark, tmp_path): + assert _field(spark, tmp_path, "double", "1", "1.7", "1e2") == [1.0, 1.7, 100.0] + assert _field(spark, tmp_path, "double", '"1.5"', "true") == ([None, None] if USE_ACTUAL_SPARK else [1.5, 1.0]) + texts = ["NaN", "+INF", "-INF", "Infinity", "+Infinity", "-Infinity"] + values = _field(spark, tmp_path, "double", *(f'"{text}"' for text in texts)) + assert math.isnan(values[0]) + assert values[1:] == [math.inf, -math.inf, math.inf, math.inf, -math.inf] + values = _field(spark, tmp_path, "double", '"INF"', '"Inf"', '"nan"') + if USE_ACTUAL_SPARK: + assert values == [None, None, None] + else: + # DuckDB's cast reads more spellings of NaN and infinity. + assert values[:2] == [math.inf, math.inf] + assert math.isnan(values[2]) + assert _field(spark, tmp_path, "float", "1.5", '"+INF"') == [1.5, math.inf] + + def test_decimals(self, spark, tmp_path): + values = ["1", "1.7", '"1"', '"1,000"', '" 1"', '"abc"', "9223372036854775808"] + # Spark leaves out the commas its default locale groups digits with, and DuckDB trims a numeric string. + thousand, one = (Decimal("1000.00"), None) if USE_ACTUAL_SPARK else (None, Decimal("1.00")) + expected = [Decimal("1.00"), Decimal("1.70"), Decimal("1.00"), thousand, one, None, None] + assert _field(spark, tmp_path, "decimal(10,2)", *values) == expected + # Spark rounds the decimal text, DuckDB the double it reads. + values = _field(spark, tmp_path, "decimal(10,2)", "1.005", "true") + assert values == ([Decimal("1.01"), None] if USE_ACTUAL_SPARK else [Decimal("1.00"), Decimal("1.00")]) + + def test_booleans(self, spark, tmp_path): + assert _field(spark, tmp_path, "boolean", "true", "false") == [True, False] + values = _field(spark, tmp_path, "boolean", '"true"', "1") + assert values == ([None, None] if USE_ACTUAL_SPARK else [True, True]) + + def test_strings(self, spark, tmp_path): + values = ['"x"', "null", "true", "1", "[1]", '{"x":1}'] + assert _field(spark, tmp_path, "string", *values) == ["x", None, "true", "1", "[1]", '{"x":1}'] + + def test_text_of_numbers_and_objects(self, spark, tmp_path): + values = _field(spark, tmp_path, "string", "1.50", "1E+2", '{"b": 1, "a": [1, 2.50]}') + if USE_ACTUAL_SPARK and not pyspark.__version__.startswith("3."): + # Spark 4 keeps the text as written, Spark 3.5 writes the value again as DuckDB does. + assert values == ["1.50", "1E+2", '{"b": 1, "a": [1, 2.50]}'] + else: + assert values == ["1.5", "100.0", '{"b":1,"a":[1,2.5]}'] + + def test_dates(self, spark, tmp_path): + values = ['"2024-01-02"', '"2024-1-2"', '"2024-01-02T10:00:00"', "1"] + day = datetime.date(2024, 1, 2) + assert _field(spark, tmp_path, "date", *values) == [day, day, day, None] + # Spark reads digits as days since the epoch, as Spark 1.5 wrote dates; DuckDB reads slashes. + values = _field(spark, tmp_path, "date", '"2024/01/02"', '"1"') + assert values == ([None, datetime.date(1970, 1, 2)] if USE_ACTUAL_SPARK else [day, None]) + + def test_timestamps(self, spark, tmp_path): + values = ['"2024-01-02 10:00:00"', '"2024-01-02"', "1.7", '"abc"', "true"] + expected = [datetime.datetime(2024, 1, 2, 10), datetime.datetime(2024, 1, 2), None, None, None] + assert _field(spark, tmp_path, "timestamp", *values) == expected + # Spark reads an integer as seconds since the epoch. + # A negative epoch second cannot be converted on Windows, so the expectation is built only where it is used. + expected = ( + [datetime.datetime.fromtimestamp(1), datetime.datetime.fromtimestamp(-1)] + if USE_ACTUAL_SPARK + else [None, None] + ) + assert _field(spark, tmp_path, "timestamp", "1", "-1") == expected + instant = datetime.datetime(2024, 1, 2, 10, tzinfo=datetime.timezone.utc).astimezone().replace(tzinfo=None) + assert _field(spark, tmp_path, "timestamp", '"2024-01-02T10:00:00Z"') == [instant] + rows = _field(spark, tmp_path, "struct", '{"t": "2024-01-02 10:00:00"}') + assert [row.t for row in rows] == [datetime.datetime(2024, 1, 2, 10)] + + def test_timestamps_without_a_zone(self, spark, tmp_path): + values = _field(spark, tmp_path, "timestamp_ntz", '"2024-01-02 12:00:00"', '"2024-01-02T12:00:00Z"', "1") + # Spark rejects text with a zone, DuckDB drops the zone. + noon = datetime.datetime(2024, 1, 2, 12) + assert values == ([noon, None, None] if USE_ACTUAL_SPARK else [noon, noon, None]) + + def test_arrays(self, spark, tmp_path): + values = ["[1, null, 2]", "[]", "null", '{"a": 1}'] + assert _field(spark, tmp_path, "array", *values) == [[1, None, 2], [], None, None] + # Spark makes the array NULL when an element does not convert, DuckDB the element. + values = _field(spark, tmp_path, "array", "[1.5]", '[1, "x"]') + assert values == ([None, None] if USE_ACTUAL_SPARK else [[2], [1, None]]) + + def test_structs(self, spark, tmp_path): + # A field that does not convert is NULL, and fields are found by their name in its case. + values = ['{"a": "x", "b": 2}', "{}", '{"A": 1, "b": 2}'] + rows = _field(spark, tmp_path, "struct", *values) + assert [tuple(row) for row in rows] == [(None, 2), (None, None), (None, 2)] + rows = _field(spark, tmp_path, "array>", '[{"a": "x", "b": 1}]') + assert [[tuple(item) for item in row] for row in rows] == [[(None, 1)]] + # Spark reads a value that is not an object as NULL, DuckDB as a struct of NULL fields. + rows = _field(spark, tmp_path, "struct", "1", "[1]") + expected = [None, None] if USE_ACTUAL_SPARK else [(None, None), (None, None)] + assert [None if row is None else tuple(row) for row in rows] == expected + + def test_maps(self, spark, tmp_path): + values = ['{"a": 1}', '{"a": null}', "[1]"] + assert _field(spark, tmp_path, "map", *values) == [{"a": 1}, {"a": None}, None] + # Spark makes the map NULL when a value does not convert, DuckDB the value. + expected = None if USE_ACTUAL_SPARK else {"a": 1, "b": None} + assert _field(spark, tmp_path, "map", '{"a": 1, "b": "x"}') == [expected] + + def test_names_with_a_slash_or_a_tilde(self, spark, tmp_path): + df = _load(spark, tmp_path, "s struct<`a/b`:int,`c~d`:int>, `e.f` int", '{"s": {"a/b": 1, "c~d": 2}, "e.f": 3}') + assert [(tuple(row.s), row[1]) for row in df.collect()] == [((1, 2), 3)] + + def test_names_are_case_sensitive(self, spark, tmp_path): + df = _load(spark, tmp_path, "a int, s struct", '{"A": 1, "s": {"Q": 1}}') + assert [(row.a, tuple(row.s)) for row in df.collect()] == [(None, (None,))] + + def test_names_differing_in_case_only(self, spark, tmp_path): + with pytest.raises(AnalysisException, match="The column `a` already exists") as error: + _load(spark, tmp_path, "a int, A int", '{"a": 1}').collect() + assert error_condition(error.value) == "COLUMN_ALREADY_EXISTS" + + def test_schema_of_one_field(self, spark, tmp_path): + rows = _load(spark, tmp_path, "a int", '{"a": 1}', '{"a": 2}', "7").collect() + assert rows == [Row(a=1), Row(a=2), Row(a=None)] + lines = [f'{{"k{i}": 1}}' for i in range(30)] + ['{"a": 9}'] + assert _load(spark, tmp_path, "a int", *lines).collect()[-1] == Row(a=9) + + def test_malformed_lines(self, spark, tmp_path): + rows = _load(spark, tmp_path, "a int", '{"a": 1}', "{bad", '{"a": 3}').collect() + assert rows == [Row(a=1), Row(a=None), Row(a=3)] + + def test_binary(self, spark, tmp_path): + # Spark reads base64 text as binary, DuckDB the text's bytes. + expected = None if USE_ACTUAL_SPARK else b"abc" + assert _field(spark, tmp_path, "binary", '"abc"') == [expected] + + def test_json_by_its_extension(self, spark, tmp_path): + path = tmp_path / "data.json" + path.write_text('{"b": 1, "a": "x"}\n', newline="", encoding="utf-8") + if USE_ACTUAL_SPARK: + # Spark reads a path without a format as parquet. + with pytest.raises(Exception): # noqa: B017, PT011 + spark.read.load(path.as_posix()).collect() + else: + assert spark.read.load(path.as_posix()).collect() == [Row(a="x", b=1)] + assert spark.read.load(path.as_posix(), schema="b int").collect() == [Row(b=1)] + + +class TestJsonInferenceOfKeys: + def test_keys_differing_in_case_only(self, spark, tmp_path): + for lines in (['{"a": 1, "A": 2}'], ['{"a": 1}', '{"A": "x"}']): + expected = AnalysisException if USE_ACTUAL_SPARK else duckdb.NotImplementedException + with pytest.raises(expected): + _read(spark, tmp_path, *lines).collect() + + def test_keys_of_one_type_differing_in_case_only(self, spark, tmp_path): + if USE_ACTUAL_SPARK: + # Spark keeps the first name, and the field of the other name reads as NULL. + df = _read(spark, tmp_path, '{"a": 1}', '{"A": 2}') + assert [tuple(row) for row in df.collect()] == [(1,), (None,)] + else: + with pytest.raises(duckdb.NotImplementedException): + _read(spark, tmp_path, '{"a": 1}', '{"A": 2}') + + def test_empty_file(self, spark, tmp_path): + path = tmp_path / "data.json" + path.write_text("", newline="", encoding="utf-8") + assert spark.read.json(path.as_posix()).columns == [] + + def test_lines_that_are_not_objects(self, spark, tmp_path): + if USE_ACTUAL_SPARK: + df = _read(spark, tmp_path, '{"a": 1}', "[1,2]") + assert df.columns == ["_corrupt_record", "a"] + else: + with pytest.raises(ContributionsAcceptedError): + _read(spark, tmp_path, '{"a": 1}', "[1,2]") diff --git a/tests/fast/spark/test_spark_session.py b/tests/fast/spark/test_spark_session.py index d36f4f08a..1fdef83df 100644 --- a/tests/fast/spark/test_spark_session.py +++ b/tests/fast/spark/test_spark_session.py @@ -1,10 +1,9 @@ +import uuid + import pytest from spark_namespace import USE_ACTUAL_SPARK -from spark_namespace.sql.types import Row - -from duckdb.experimental.spark.exception import ( - ContributionsAcceptedError, -) +from spark_namespace.errors import AnalysisException, IllegalArgumentException, PySparkNotImplementedError +from spark_namespace.sql.types import Row, StructType _ = pytest.importorskip("duckdb.experimental.spark") from spark_namespace.sql import SparkSession @@ -84,10 +83,99 @@ def test_range(self, spark): assert res_2 == [Row(id=3), Row(id=5), Row(id=7), Row(id=9)] assert res_3 == [Row(id=3), Row(id=4), Row(id=5)] - if not USE_ACTUAL_SPARK: - with pytest.raises(ContributionsAcceptedError): - # partition size is not supported - spark.range(0, 10, 2, 2) + @pytest.mark.parametrize( + ("args", "expected"), + [ + ((5,), [0, 1, 2, 3, 4]), + ((0,), []), + ((-3,), []), + ((1, 5), [1, 2, 3, 4]), + ((5, 1), []), + ((3, 3), []), + ((1, 10, 3), [1, 4, 7]), + ((10, 1, -3), [10, 7, 4]), + ((1, 5, 1, 2), [1, 2, 3, 4]), + ], + ) + def test_range_forms(self, spark, args, expected): + df = spark.range(*args) + assert df.columns == ["id"] + assert df.schema.simpleString() == "struct" + assert [r.id for r in df.collect()] == expected + + def test_range_keywords(self, spark): + assert spark.range(start=2, end=4, step=1, numPartitions=3).collect() == [Row(id=2), Row(id=3)] + + def test_range_step_zero(self, spark): + with pytest.raises(IllegalArgumentException, match="step \\(0\\) cannot be 0"): + spark.range(1, 5, 0).collect() def test_udf(self, spark): udf_registration = spark.udf # noqa: F841 + + +class TestStatementWithoutResult: + # Spark 3.5 creates a Hive table unless the statement names a data source. + TABLE_SOURCE = " USING parquet" if USE_ACTUAL_SPARK else "" + + @pytest.fixture + def table_name(self, spark): + name = f"no_result_{uuid.uuid4().hex}" + yield name + spark.sql(f"DROP TABLE IF EXISTS {name}") + + @pytest.fixture + def statements(self, spark, table_name): + created = spark.sql(f"CREATE TABLE {table_name} (a INT)" + self.TABLE_SOURCE) + inserted = spark.sql(f"INSERT INTO {table_name} VALUES (1), (2)") + return [created, inserted] + + def test_statement_runs(self, spark, table_name, statements): + assert sorted(spark.table(table_name).collect()) == [Row(a=1), Row(a=2)] + + def test_schema(self, statements): + for df in statements: + assert df.schema == StructType([]) + assert df.columns == [] + assert df.dtypes == [] + + @pytest.mark.skipif(not USE_ACTUAL_SPARK, reason="DuckDB has no relation without columns") + def test_rows(self, spark, statements): + other = spark.createDataFrame([(1,)], ["a"]) + for df in statements: + assert df.count() == 0 + assert df.collect() == [] + assert df.union(df).collect() == [] + assert other.crossJoin(df).collect() == [] + + @pytest.mark.skipif(USE_ACTUAL_SPARK, reason="Spark represents a DataFrame without columns") + @pytest.mark.parametrize( + "use", + [ + lambda df, other: df.count(), + lambda df, other: df.collect(), + lambda df, other: df.show(), + lambda df, other: df.toPandas(), + lambda df, other: df.union(df), + lambda df, other: df.filter("true"), + lambda df, other: df.join(other), + lambda df, other: other.crossJoin(df), + lambda df, other: "a" in df, + ], + ) + def test_only_the_schema_is_available(self, spark, statements, use): + other = spark.createDataFrame([(1,)], ["a"]) + for df in statements: + with pytest.raises(PySparkNotImplementedError): + use(df, other) + + @pytest.mark.parametrize("operation", ["union", "unionByName", "intersect", "intersectAll", "exceptAll"]) + def test_set_operation_with_a_frame_without_columns(self, spark, statements, operation): + # Spark raises because the column counts differ; DuckDB has no relation without columns. + expected = AnalysisException if USE_ACTUAL_SPARK else PySparkNotImplementedError + other = spark.createDataFrame([(1,)], ["a"]) + for df in statements: + with pytest.raises(expected): + getattr(other, operation)(df).collect() + with pytest.raises(expected): + getattr(df, operation)(other).collect() diff --git a/tests/fast/spark/test_spark_show.py b/tests/fast/spark/test_spark_show.py new file mode 100644 index 000000000..92d5b1f94 --- /dev/null +++ b/tests/fast/spark/test_spark_show.py @@ -0,0 +1,1417 @@ +import datetime +from datetime import timedelta as td +from decimal import Decimal + +import pytest + +_ = pytest.importorskip("duckdb.experimental.spark") + +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import PySparkTypeError +from spark_namespace.sql.functions import col +from spark_namespace.sql.types import ( + ArrayType, + BinaryType, + BooleanType, + ByteType, + DateType, + DayTimeIntervalType, + DecimalType, + DoubleType, + FloatType, + IntegerType, + LongType, + MapType, + ShortType, + StringType, + StructField, + StructType, + TimestampNTZType, +) + +if USE_ACTUAL_SPARK: + import pyspark + + # PySpark 3.5 ends the "only showing top" and "(0 rows)" footers with an extra newline. + FOOTER_END = "\n\n" if pyspark.__version__ < "4" else "\n" +else: + FOOTER_END = "\n" + + +def F(name, data_type): + return StructField(name, data_type, True) + + +def S(*fields): + return StructType(list(fields)) + + +def printed(lines): + if lines[-1].startswith("only showing top") or lines[-1] == "(0 rows)": + return "\n".join(lines) + FOOTER_END + return "\n".join(lines) + "\n" + + +# Expected output is copied from PySpark 4.2.0 running the same calls. +SHOW_CASES = [ + pytest.param( + [(1, 2, 3, 4), (-100, -32768, -2147483648, -9223372036854775808), (None, None, None, None)], + S(F("b", ByteType()), F("s", ShortType()), F("i", IntegerType()), F("l", LongType())), + lambda df: df.show(), + ( + "+----+------+-----------+--------------------+", + "| b| s| i| l|", + "+----+------+-----------+--------------------+", + "| 1| 2| 3| 4|", + "|-100|-32768|-2147483648|-9223372036854775808|", + "|NULL| NULL| NULL| NULL|", + "+----+------+-----------+--------------------+", + "", + ), + id="ints", + ), + pytest.param( + [ + (Decimal("1.50"), Decimal("1.5"), Decimal("12345")), + (Decimal("-0.05"), Decimal("0"), Decimal("0")), + (Decimal("0"), Decimal("-123.000000000000000001"), None), + ], + S(F("d1", DecimalType(10, 2)), F("d2", DecimalType(38, 18)), F("d3", DecimalType(5, 0))), + lambda df: df.show(), + ( + "+-----+--------------------+-----+", + "| d1| d2| d3|", + "+-----+--------------------+-----+", + "| 1.50|1.500000000000000000|12345|", + "|-0.05|0.000000000000000000| 0|", + "| 0.00|-123.000000000000...| NULL|", + "+-----+--------------------+-----+", + "", + ), + id="decimals", + ), + pytest.param( + [(Decimal("0E-10"),), (Decimal("0.0000000001"),)], + S(F("d", DecimalType(10, 10))), + lambda df: df.show(), + ( + "+------------+", + "| d|", + "+------------+", + "|0.0000000000|", + "|0.0000000001|", + "+------------+", + "", + ), + id="decimal_small_scale", + ), + pytest.param( + [("",), ("a",), ("abcdefghijklmnopqrstuvwxyz",), ("a\nb\tc",), ("日本語",), ("é",), (None,)], + S(F("s", StringType())), + lambda df: df.show(), + ( + "+--------------------+", + "| s|", + "+--------------------+", + "| |", + "| a|", + "|abcdefghijklmnopq...|", + "| a\\nb\\tc|", + "| 日本語|", + "| é|", + "| NULL|", + "+--------------------+", + "", + ), + id="strings", + ), + pytest.param( + [("日本語日本語",), ("ab",)], + S(F("s", StringType())), + lambda df: df.show(), + ( + "+------------+", + "| s|", + "+------------+", + "|日本語日本語|", + "| ab|", + "+------------+", + "", + ), + id="cjk_width", + ), + pytest.param( + [(True,), (False,), (None,)], + S(F("b", BooleanType())), + lambda df: df.show(), + ( + "+-----+", + "| b|", + "+-----+", + "| true|", + "|false|", + "| NULL|", + "+-----+", + "", + ), + id="booleans", + ), + pytest.param( + [(datetime.date(2024, 1, 15),), (datetime.date(1, 1, 1),), (datetime.date(9999, 12, 31),), (None,)], + S(F("d", DateType())), + lambda df: df.show(), + ( + "+----------+", + "| d|", + "+----------+", + "|2024-01-15|", + "|0001-01-01|", + "|9999-12-31|", + "| NULL|", + "+----------+", + "", + ), + id="dates", + ), + pytest.param( + [(datetime.datetime(2024, 1, 15, 10, 30),), (datetime.datetime(2024, 1, 15, 10, 30, 0, 123),), (None,)], + S(F("t", TimestampNTZType())), + lambda df: df.show(), + ( + "+--------------------+", + "| t|", + "+--------------------+", + "| 2024-01-15 10:30:00|", + "|2024-01-15 10:30:...|", + "| NULL|", + "+--------------------+", + "", + ), + id="timestamps_ntz", + ), + pytest.param( + [(bytearray(b"hello"),), (bytearray(b""),), (bytearray(b"\x00\xff\x10"),), (None,)], + S(F("b", BinaryType())), + lambda df: df.show(), + ( + "+----------------+", + "| b|", + "+----------------+", + "|[68 65 6C 6C 6F]|", + "| []|", + "| [00 FF 10]|", + "| NULL|", + "+----------------+", + "", + ), + id="binary", + ), + pytest.param( + [([1, 2, 3], ["a", "b"], [[1, 2], [3]]), ([], [], []), ([1, None], [None], [None, [None]]), (None, None, None)], + S( + F("a", ArrayType(IntegerType())), + F("s", ArrayType(StringType())), + F("n", ArrayType(ArrayType(IntegerType()))), + ), + lambda df: df.show(), + ( + "+---------+------+--------------+", + "| a| s| n|", + "+---------+------+--------------+", + "|[1, 2, 3]|[a, b]| [[1, 2], [3]]|", + "| []| []| []|", + "|[1, NULL]|[NULL]|[NULL, [NULL]]|", + "| NULL| NULL| NULL|", + "+---------+------+--------------+", + "", + ), + id="arrays", + ), + pytest.param( + [([1.0, 2.5, 1e10], [Decimal("1.50")], [datetime.date(2024, 1, 2)], [True, False], [bytearray(b"hi")])], + S( + F("d", ArrayType(DoubleType())), + F("m", ArrayType(DecimalType(10, 2))), + F("t", ArrayType(DateType())), + F("b", ArrayType(BooleanType())), + F("y", ArrayType(BinaryType())), + ), + lambda df: df.show(), + ( + "+------------------+------+------------+-------------+---------+", + "| d| m| t| b| y|", + "+------------------+------+------------+-------------+---------+", + "|[1.0, 2.5, 1.0E10]|[1.50]|[2024-01-02]|[true, false]|[[68 69]]|", + "+------------------+------+------------+-------------+---------+", + "", + ), + id="arrays_of_other", + ), + pytest.param( + [({"a": 1, "b": 2},), ({},), ({"a": None},), (None,)], + S(F("m", MapType(StringType(), IntegerType()))), + lambda df: df.show(), + ( + "+----------------+", + "| m|", + "+----------------+", + "|{a -> 1, b -> 2}|", + "| {}|", + "| {a -> NULL}|", + "| NULL|", + "+----------------+", + "", + ), + id="maps", + ), + pytest.param( + [({1: ["x", "y"]}, {"k": {"a": 1}})], + S( + F("m", MapType(IntegerType(), ArrayType(StringType()))), + F("n", MapType(StringType(), MapType(StringType(), IntegerType()))), + ), + lambda df: df.show(), + ( + "+-------------+---------------+", + "| m| n|", + "+-------------+---------------+", + "|{1 -> [x, y]}|{k -> {a -> 1}}|", + "+-------------+---------------+", + "", + ), + id="map_nested", + ), + pytest.param( + [ + ((1, "x"), (1, (2, None)), ([1, 2], {"a": 1})), + ((None, None), (None, None), (None, None)), + (None, None, None), + ], + S( + F("s", S(F("a", IntegerType()), F("b", StringType()))), + F("n", S(F("a", IntegerType()), F("b", S(F("c", IntegerType()), F("d", StringType()))))), + F("c", S(F("a", ArrayType(IntegerType())), F("m", MapType(StringType(), IntegerType())))), + ), + lambda df: df.show(), + ( + "+------------+--------------+------------------+", + "| s| n| c|", + "+------------+--------------+------------------+", + "| {1, x}|{1, {2, NULL}}|{[1, 2], {a -> 1}}|", + "|{NULL, NULL}| {NULL, NULL}| {NULL, NULL}|", + "| NULL| NULL| NULL|", + "+------------+--------------+------------------+", + "", + ), + id="structs", + ), + pytest.param( + [([(1, "a"), (2, None)],), ([None],)], + S(F("a", ArrayType(S(F("x", IntegerType()), F("y", StringType()))))), + lambda df: df.show(), + ( + "+-------------------+", + "| a|", + "+-------------------+", + "|[{1, a}, {2, NULL}]|", + "| [NULL]|", + "+-------------------+", + "", + ), + id="array_of_struct", + ), + pytest.param( + [((bytearray(b"hi"), 1.0),)], + S(F("s", S(F("b", BinaryType()), F("d", DoubleType())))), + lambda df: df.show(), + ( + "+--------------+", + "| s|", + "+--------------+", + "|{[68 69], 1.0}|", + "+--------------+", + "", + ), + id="struct_binary", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(), + ( + "+---------------------+---+", + "|long_column_name_here| n|", + "+---------------------+---+", + "| abcdefghijklmnopq...| 1|", + "| short| 22|", + "| exactly twenty chars|333|", + "+---------------------+---+", + "", + ), + id="truncate_default", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate=True), + ( + "+---------------------+---+", + "|long_column_name_here| n|", + "+---------------------+---+", + "| abcdefghijklmnopq...| 1|", + "| short| 22|", + "| exactly twenty chars|333|", + "+---------------------+---+", + "", + ), + id="truncate_true", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate=False), + ( + "+--------------------------+---+", + "|long_column_name_here |n |", + "+--------------------------+---+", + "|abcdefghijklmnopqrstuvwxyz|1 |", + "|short |22 |", + "|exactly twenty chars |333|", + "+--------------------------+---+", + "", + ), + id="truncate_false", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate=0), + ( + "+--------------------------+---+", + "|long_column_name_here |n |", + "+--------------------------+---+", + "|abcdefghijklmnopqrstuvwxyz|1 |", + "|short |22 |", + "|exactly twenty chars |333|", + "+--------------------------+---+", + "", + ), + id="truncate_0", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate=5), + ( + "+---------------------+---+", + "|long_column_name_here| n|", + "+---------------------+---+", + "| ab...| 1|", + "| short| 22|", + "| ex...|333|", + "+---------------------+---+", + "", + ), + id="truncate_5", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate=4), + ( + "+---------------------+---+", + "|long_column_name_here| n|", + "+---------------------+---+", + "| a...| 1|", + "| s...| 22|", + "| e...|333|", + "+---------------------+---+", + "", + ), + id="truncate_4", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate=3), + ( + "+---------------------+---+", + "|long_column_name_here| n|", + "+---------------------+---+", + "| abc| 1|", + "| sho| 22|", + "| exa|333|", + "+---------------------+---+", + "", + ), + id="truncate_3", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate=1), + ( + "+---------------------+---+", + "|long_column_name_here| n|", + "+---------------------+---+", + "| a| 1|", + "| s| 2|", + "| e| 3|", + "+---------------------+---+", + "", + ), + id="truncate_1", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate=-1), + ( + "+--------------------------+---+", + "|long_column_name_here |n |", + "+--------------------------+---+", + "|abcdefghijklmnopqrstuvwxyz|1 |", + "|short |22 |", + "|exactly twenty chars |333|", + "+--------------------------+---+", + "", + ), + id="truncate_neg", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate=5.7), + ( + "+---------------------+---+", + "|long_column_name_here| n|", + "+---------------------+---+", + "| ab...| 1|", + "| short| 22|", + "| ex...|333|", + "+---------------------+---+", + "", + ), + id="truncate_float", + ), + pytest.param( + [("abcdefghijklmnopqrstuvwxyz", 1), ("short", 22), ("exactly twenty chars", 333)], + S(F("long_column_name_here", StringType()), F("n", IntegerType())), + lambda df: df.show(truncate="5"), + ( + "+---------------------+---+", + "|long_column_name_here| n|", + "+---------------------+---+", + "| ab...| 1|", + "| short| 22|", + "| ex...|333|", + "+---------------------+---+", + "", + ), + id="truncate_str_digit", + ), + pytest.param( + [(["abcdefghij", "klmnopqrst"],)], + S(F("a", ArrayType(StringType()))), + lambda df: df.show(), + ( + "+--------------------+", + "| a|", + "+--------------------+", + "|[abcdefghij, klmn...|", + "+--------------------+", + "", + ), + id="truncate_array", + ), + pytest.param( + [(i, f"r{i}") for i in range(25)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(), + ( + "+---+---+", + "| id| v|", + "+---+---+", + "| 0| r0|", + "| 1| r1|", + "| 2| r2|", + "| 3| r3|", + "| 4| r4|", + "| 5| r5|", + "| 6| r6|", + "| 7| r7|", + "| 8| r8|", + "| 9| r9|", + "| 10|r10|", + "| 11|r11|", + "| 12|r12|", + "| 13|r13|", + "| 14|r14|", + "| 15|r15|", + "| 16|r16|", + "| 17|r17|", + "| 18|r18|", + "| 19|r19|", + "+---+---+", + "only showing top 20 rows", + ), + id="n_default_25_rows", + ), + pytest.param( + [(i, f"r{i}") for i in range(5)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(n=2), + ( + "+---+---+", + "| id| v|", + "+---+---+", + "| 0| r0|", + "| 1| r1|", + "+---+---+", + "only showing top 2 rows", + ), + id="n_2", + ), + pytest.param( + [(i, f"r{i}") for i in range(5)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(n=1), + ( + "+---+---+", + "| id| v|", + "+---+---+", + "| 0| r0|", + "+---+---+", + "only showing top 1 row", + ), + id="n_1", + ), + pytest.param( + [(i, f"r{i}") for i in range(5)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(n=0), + ( + "+---+---+", + "| id| v|", + "+---+---+", + "+---+---+", + "only showing top 0 rows", + ), + id="n_0", + ), + pytest.param( + [(i, f"r{i}") for i in range(5)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(n=-1), + ( + "+---+---+", + "| id| v|", + "+---+---+", + "+---+---+", + "only showing top 0 rows", + ), + id="n_neg", + ), + pytest.param( + [(i, f"r{i}") for i in range(5)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(n=5), + ( + "+---+---+", + "| id| v|", + "+---+---+", + "| 0| r0|", + "| 1| r1|", + "| 2| r2|", + "| 3| r3|", + "| 4| r4|", + "+---+---+", + "", + ), + id="n_exact", + ), + pytest.param( + [(i, f"r{i}") for i in range(5)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(n=10), + ( + "+---+---+", + "| id| v|", + "+---+---+", + "| 0| r0|", + "| 1| r1|", + "| 2| r2|", + "| 3| r3|", + "| 4| r4|", + "+---+---+", + "", + ), + id="n_more", + ), + pytest.param( + [(i, f"r{i}") for i in range(5)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(3, False), + ( + "+---+---+", + "|id |v |", + "+---+---+", + "|0 |r0 |", + "|1 |r1 |", + "|2 |r2 |", + "+---+---+", + "only showing top 3 rows", + ), + id="n_positional", + ), + pytest.param( + [], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(), + ( + "+---+---+", + "| id| v|", + "+---+---+", + "+---+---+", + "", + ), + id="empty", + ), + pytest.param( + [(0, "r0"), (1, "r1"), (2, "r2"), (None, "abcdefghijklmnopqrstuvwxyz")], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(vertical=True), + ( + "-RECORD 0-------------------", + " id | 0 ", + " v | r0 ", + "-RECORD 1-------------------", + " id | 1 ", + " v | r1 ", + "-RECORD 2-------------------", + " id | 2 ", + " v | r2 ", + "-RECORD 3-------------------", + " id | NULL ", + " v | abcdefghijklmnopq... ", + "", + ), + id="vertical", + ), + pytest.param( + [(i, f"r{i}") for i in range(5)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(n=2, vertical=True), + ( + "-RECORD 0--", + " id | 0 ", + " v | r0 ", + "-RECORD 1--", + " id | 1 ", + " v | r1 ", + "only showing top 2 rows", + ), + id="vertical_n", + ), + pytest.param( + [(i, f"r{i}") for i in range(5)], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(n=1, vertical=True), + ( + "-RECORD 0--", + " id | 0 ", + " v | r0 ", + "only showing top 1 row", + ), + id="vertical_n1", + ), + pytest.param( + [], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(vertical=True), + ("(0 rows)",), + id="vertical_empty", + ), + pytest.param( + [(1, "abcdefghijklmnopqrstuvwxyz")], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(vertical=True, truncate=False), + ( + "-RECORD 0-------------------------", + " id | 1 ", + " v | abcdefghijklmnopqrstuvwxyz ", + "", + ), + id="vertical_truncate_false", + ), + pytest.param( + [(1, "abcdefghijklmnopqrstuvwxyz")], + S(F("id", LongType()), F("v", StringType())), + lambda df: df.show(vertical=True, truncate=3), + ( + "-RECORD 0--", + " id | 1 ", + " v | abc ", + "", + ), + id="vertical_truncate_3", + ), + pytest.param( + [(1, "x")], + S(F("a_rather_long_field_name", LongType()), F("v", StringType())), + lambda df: df.show(vertical=True), + ( + "-RECORD 0-----------------------", + " a_rather_long_field_name | 1 ", + " v | x ", + "", + ), + id="vertical_long_name", + ), + pytest.param( + [(1,)], + S(F("a\nb", LongType())), + lambda df: df.show(), + ( + "+----+", + "|a\\nb|", + "+----+", + "| 1|", + "+----+", + "", + ), + id="header_escape", + ), + pytest.param( + [ + (td(days=1, hours=2, minutes=3, seconds=4),), + (td(seconds=4.5),), + (td(0),), + (-td(days=1, seconds=1),), + (td(days=-1, hours=23),), + (td(microseconds=1),), + (td(days=12345),), + (None,), + ], + S(F("i", DayTimeIntervalType())), + lambda df: df.show(truncate=False), + ( + "+------------------------------------------+", + "|i |", + "+------------------------------------------+", + "|INTERVAL '1 02:03:04' DAY TO SECOND |", + "|INTERVAL '0 00:00:04.5' DAY TO SECOND |", + "|INTERVAL '0 00:00:00' DAY TO SECOND |", + "|INTERVAL '-1 00:00:01' DAY TO SECOND |", + "|INTERVAL '-0 01:00:00' DAY TO SECOND |", + "|INTERVAL '0 00:00:00.000001' DAY TO SECOND|", + "|INTERVAL '12345 00:00:00' DAY TO SECOND |", + "|NULL |", + "+------------------------------------------+", + "", + ), + id="intervals_untruncated", + ), + pytest.param( + [ + (td(days=1, hours=2, minutes=3, seconds=4),), + (td(seconds=4.5),), + (td(0),), + (-td(days=1, seconds=1),), + (td(days=-1, hours=23),), + (td(microseconds=1),), + (td(days=12345),), + (None,), + ], + S(F("i", DayTimeIntervalType())), + lambda df: df.show(), + ( + "+--------------------+", + "| i|", + "+--------------------+", + "|INTERVAL '1 02:03...|", + "|INTERVAL '0 00:00...|", + "|INTERVAL '0 00:00...|", + "|INTERVAL '-1 00:0...|", + "|INTERVAL '-0 01:0...|", + "|INTERVAL '0 00:00...|", + "|INTERVAL '12345 0...|", + "| NULL|", + "+--------------------+", + "", + ), + id="intervals", + ), + pytest.param( + [("a\rb\fc\bd\x0be\x07f\\g",), ("q\"uote's",)], + S(F("s", StringType())), + lambda df: df.show(truncate=False), + ( + "+------------------+", + "|s |", + "+------------------+", + "|a\\rb\\fc\\bd\\ve\\af\\g|", + "|q\"uote's |", + "+------------------+", + "", + ), + id="escapes_untruncated", + ), + pytest.param( + [("a\rb\fc\bd\x0be\x07f\\g",), ("q\"uote's",)], + S(F("s", StringType())), + lambda df: df.show(), + ( + "+------------------+", + "| s|", + "+------------------+", + "|a\\rb\\fc\\bd\\ve\\af\\g|", + "| q\"uote's|", + "+------------------+", + "", + ), + id="escapes", + ), + pytest.param( + [("\ud55c\uae00",), ("\uff21\uff22",), ("\uff71\uff72",), ("\U0001f600x",), ("\u3000",), ("\u00e9\u0301",)], + S(F("s", StringType())), + lambda df: df.show(truncate=False), + ( + "+----+", + "|s |", + "+----+", + "|한글|", + "|\uff21\uff22|", + "|\uff71\uff72 |", + "|😀x |", + "|\u3000 |", + "|é́ |", + "+----+", + "", + ), + id="widths_untruncated", + ), + pytest.param( + [("\ud55c\uae00",), ("\uff21\uff22",), ("\uff71\uff72",), ("\U0001f600x",), ("\u3000",), ("\u00e9\u0301",)], + S(F("s", StringType())), + lambda df: df.show(), + ( + "+----+", + "| s|", + "+----+", + "|한글|", + "|\uff21\uff22|", + "| \uff71\uff72|", + "| 😀x|", + "| \u3000|", + "| é́|", + "+----+", + "", + ), + id="widths", + ), + pytest.param( + [(1,)], + S(F("日本", LongType())), + lambda df: df.show(truncate=False), + ( + "+----+", + "|日本|", + "+----+", + "|1 |", + "+----+", + "", + ), + id="cjk_header_untruncated", + ), + pytest.param( + [(1,)], + S(F("日本", LongType())), + lambda df: df.show(), + ( + "+----+", + "|日本|", + "+----+", + "| 1|", + "+----+", + "", + ), + id="cjk_header", + ), + pytest.param( + [({"a": (1, None)},)], + S(F("m", MapType(StringType(), S(F("x", IntegerType()), F("y", StringType()))))), + lambda df: df.show(truncate=False), + ( + "+----------------+", + "|m |", + "+----------------+", + "|{a -> {1, NULL}}|", + "+----------------+", + "", + ), + id="map_null_key_struct_untruncated", + ), + pytest.param( + [({"a": (1, None)},)], + S(F("m", MapType(StringType(), S(F("x", IntegerType()), F("y", StringType()))))), + lambda df: df.show(), + ( + "+----------------+", + "| m|", + "+----------------+", + "|{a -> {1, NULL}}|", + "+----------------+", + "", + ), + id="map_null_key_struct", + ), + pytest.param( + [(["a,b", "", None, "x\ny"], {"k k": "v, v"})], + S(F("a", ArrayType(StringType())), F("m", MapType(StringType(), StringType()))), + lambda df: df.show(truncate=False), + ( + "+-------------------+-------------+", + "|a |m |", + "+-------------------+-------------+", + "|[a,b, , NULL, x\\ny]|{k k -> v, v}|", + "+-------------------+-------------+", + "", + ), + id="nested_strings_untruncated", + ), + pytest.param( + [(["a,b", "", None, "x\ny"], {"k k": "v, v"})], + S(F("a", ArrayType(StringType())), F("m", MapType(StringType(), StringType()))), + lambda df: df.show(), + ( + "+-------------------+-------------+", + "| a| m|", + "+-------------------+-------------+", + "|[a,b, , NULL, x\\ny]|{k k -> v, v}|", + "+-------------------+-------------+", + "", + ), + id="nested_strings", + ), + pytest.param( + [ + (x,) + for x in [ + 1e-3, + 9.999e-4, + 1e7 - 0.5, + 1e7, + 5e-324, + 1.7976931348623157e308, + 2.0**53, + 0.30000000000000004, + 1 / 3, + 100.5, + 1e21, + 1e-7, + 123e-20, + ] + ], + S(F("d", DoubleType())), + lambda df: df.show(truncate=False), + ( + "+----------------------+", + "|d |", + "+----------------------+", + "|0.001 |", + "|9.999E-4 |", + "|9999999.5 |", + "|1.0E7 |", + "|4.9E-324 |", + "|1.7976931348623157E308|", + "|9.007199254740992E15 |", + "|0.30000000000000004 |", + "|0.3333333333333333 |", + "|100.5 |", + "|1.0E21 |", + "|1.0E-7 |", + "|1.23E-18 |", + "+----------------------+", + "", + ), + id="float_edges_untruncated", + ), + pytest.param( + [ + (x,) + for x in [ + 1e-3, + 9.999e-4, + 1e7 - 0.5, + 1e7, + 5e-324, + 1.7976931348623157e308, + 2.0**53, + 0.30000000000000004, + 1 / 3, + 100.5, + 1e21, + 1e-7, + 123e-20, + ] + ], + S(F("d", DoubleType())), + lambda df: df.show(), + ( + "+--------------------+", + "| d|", + "+--------------------+", + "| 0.001|", + "| 9.999E-4|", + "| 9999999.5|", + "| 1.0E7|", + "| 4.9E-324|", + "|1.797693134862315...|", + "|9.007199254740992E15|", + "| 0.30000000000000004|", + "| 0.3333333333333333|", + "| 100.5|", + "| 1.0E21|", + "| 1.0E-7|", + "| 1.23E-18|", + "+--------------------+", + "", + ), + id="float_edges", + ), + pytest.param( + [(x,) for x in [0.3, 1 / 3, 16777216.0, 1e-45, 3.4028235e38, 1.17549435e-38, 9999999.0, 0.001, 2.5]], + S(F("f", FloatType())), + lambda df: df.show(truncate=False), + ( + "+-------------+", + "|f |", + "+-------------+", + "|0.3 |", + "|0.33333334 |", + "|1.6777216E7 |", + "|1.4E-45 |", + "|3.4028235E38 |", + "|1.1754944E-38|", + "|9999999.0 |", + "|0.001 |", + "|2.5 |", + "+-------------+", + "", + ), + id="float32_edges_untruncated", + ), + pytest.param( + [(x,) for x in [0.3, 1 / 3, 16777216.0, 1e-45, 3.4028235e38, 1.17549435e-38, 9999999.0, 0.001, 2.5]], + S(F("f", FloatType())), + lambda df: df.show(), + ( + "+-------------+", + "| f|", + "+-------------+", + "| 0.3|", + "| 0.33333334|", + "| 1.6777216E7|", + "| 1.4E-45|", + "| 3.4028235E38|", + "|1.1754944E-38|", + "| 9999999.0|", + "| 0.001|", + "| 2.5|", + "+-------------+", + "", + ), + id="float32_edges", + ), + pytest.param( + [ + (x,) + for x in [ + 1.0, + 0.1, + 100.0, + 1e7, + 9999999.0, + 12345678.9, + 1e10, + 1e20, + 0.001, + 0.0001, + 1.5e-5, + -0.0, + 123456789.123, + 1.23456789012e-10, + float("inf"), + float("-inf"), + None, + ] + ], + S(F("d", DoubleType())), + lambda df: df.show(truncate=False), + ( + "+-----------------+", + "|d |", + "+-----------------+", + "|1.0 |", + "|0.1 |", + "|100.0 |", + "|1.0E7 |", + "|9999999.0 |", + "|1.23456789E7 |", + "|1.0E10 |", + "|1.0E20 |", + "|0.001 |", + "|1.0E-4 |", + "|1.5E-5 |", + "|-0.0 |", + "|1.23456789123E8 |", + "|1.23456789012E-10|", + "|Infinity |", + "|-Infinity |", + "|NULL |", + "+-----------------+", + "", + ), + id="doubles_untruncated", + ), + pytest.param( + [ + (x,) + for x in [ + 1.0, + 0.1, + 100.0, + 1e7, + 9999999.0, + 12345678.9, + 1e10, + 1e20, + 0.001, + 0.0001, + 1.5e-5, + -0.0, + 123456789.123, + 1.23456789012e-10, + float("inf"), + float("-inf"), + None, + ] + ], + S(F("d", DoubleType())), + lambda df: df.show(), + ( + "+-----------------+", + "| d|", + "+-----------------+", + "| 1.0|", + "| 0.1|", + "| 100.0|", + "| 1.0E7|", + "| 9999999.0|", + "| 1.23456789E7|", + "| 1.0E10|", + "| 1.0E20|", + "| 0.001|", + "| 1.0E-4|", + "| 1.5E-5|", + "| -0.0|", + "| 1.23456789123E8|", + "|1.23456789012E-10|", + "| Infinity|", + "| -Infinity|", + "| NULL|", + "+-----------------+", + "", + ), + id="doubles", + ), + pytest.param( + [(x,) for x in [1.0, 1.1, 0.1, 3.4e38, 16777217.0, 1e-3, 1e-4, 1e7, -2.5, None]], + S(F("f", FloatType())), + lambda df: df.show(truncate=False), + ( + "+-----------+", + "|f |", + "+-----------+", + "|1.0 |", + "|1.1 |", + "|0.1 |", + "|3.4E38 |", + "|1.6777216E7|", + "|0.001 |", + "|1.0E-4 |", + "|1.0E7 |", + "|-2.5 |", + "|NULL |", + "+-----------+", + "", + ), + id="floats_untruncated", + ), + pytest.param( + [(x,) for x in [1.0, 1.1, 0.1, 3.4e38, 16777217.0, 1e-3, 1e-4, 1e7, -2.5, None]], + S(F("f", FloatType())), + lambda df: df.show(), + ( + "+-----------+", + "| f|", + "+-----------+", + "| 1.0|", + "| 1.1|", + "| 0.1|", + "| 3.4E38|", + "|1.6777216E7|", + "| 0.001|", + "| 1.0E-4|", + "| 1.0E7|", + "| -2.5|", + "| NULL|", + "+-----------+", + "", + ), + id="floats", + ), + pytest.param( + [(Decimal("12345678901234567890.123456789012345678"),)], + S(F("d", DecimalType(38, 18))), + lambda df: df.show(truncate=False), + ( + "+---------------------------------------+", + "|d |", + "+---------------------------------------+", + "|12345678901234567890.123456789012345678|", + "+---------------------------------------+", + "", + ), + id="big_decimal_untruncated", + ), + pytest.param( + [(Decimal("12345678901234567890.123456789012345678"),)], + S(F("d", DecimalType(38, 18))), + lambda df: df.show(), + ( + "+--------------------+", + "| d|", + "+--------------------+", + "|12345678901234567...|", + "+--------------------+", + "", + ), + id="big_decimal", + ), +] + + +class TestDataFrameShow: + @pytest.mark.parametrize(("data", "schema", "show", "expected"), SHOW_CASES) + def test_show(self, spark, capsys, data, schema, show, expected): + df = spark.createDataFrame(data, schema) + show(df) + assert capsys.readouterr().out == printed(expected) + + def test_show_nan(self, spark, capsys): + df = spark.createDataFrame([("NaN",), ("1.5",)], ["s"]) + df.select(col("s").cast("double").alias("d"), col("s").cast("float").alias("f")).show() + assert capsys.readouterr().out == printed( + ("+---+---+", "| d| f|", "+---+---+", "|NaN|NaN|", "|1.5|1.5|", "+---+---+", "") + ) + + def test_show_no_arguments_match_defaults(self, spark, capsys): + df = spark.createDataFrame([(i, "x" * i) for i in range(30)], S(F("i", LongType()), F("s", StringType()))) + df.show() + default = capsys.readouterr().out + df.show(20, True, False) + assert capsys.readouterr().out == default + + @pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"n": "2"}, id="n_str"), + pytest.param({"n": True}, id="n_bool"), + pytest.param({"n": 2.0}, id="n_float"), + pytest.param({"vertical": "yes"}, id="vertical_str"), + pytest.param({"vertical": 1}, id="vertical_int"), + pytest.param({"truncate": "x"}, id="truncate_str"), + ], + ) + def test_show_argument_types(self, spark, capsys, kwargs): + df = spark.createDataFrame([(1, "a")], S(F("id", LongType()), F("v", StringType()))) + with pytest.raises(PySparkTypeError, match=f"`{next(iter(kwargs))}`"): + df.show(**kwargs) + assert capsys.readouterr().out == "" + + def test_show_truncate_none(self, spark): + df = spark.createDataFrame([(1, "a")], S(F("id", LongType()), F("v", StringType()))) + with pytest.raises(TypeError, match="NoneType"): + df.show(truncate=None) + + def test_repr(self, spark): + df = spark.createDataFrame( + [], + S( + F("id", LongType()), + F("v", StringType()), + F("a", ArrayType(IntegerType())), + F("m", MapType(StringType(), DecimalType(10, 2))), + ), + ) + assert repr(df) == "DataFrame[id: bigint, v: string, a: array, m: map]" + + +@pytest.fixture +def kathmandu(spark): + if USE_ACTUAL_SPARK: + previous = spark.conf.get("spark.sql.session.timeZone") + spark.conf.set("spark.sql.session.timeZone", "Asia/Kathmandu") + yield + spark.conf.set("spark.sql.session.timeZone", previous) + else: + # The module's session time zone is DuckDB's TimeZone setting. + (previous,) = spark.conn.execute("SELECT current_setting('TimeZone')").fetchone() + spark.conn.execute("SET TimeZone = 'Asia/Kathmandu'") + yield + spark.conn.execute(f"SET TimeZone = '{previous}'") + + +class TestShowInTheSessionTimeZone: + def test_timestamps(self, spark, capsys, kathmandu): + instant = datetime.datetime(2024, 1, 1, 12, tzinfo=datetime.timezone.utc) + df = spark.createDataFrame( + [(instant, [instant], (instant, 1), {"k": instant})], + "t timestamp, a array, s struct, m map", + ) + df.show(truncate=False) + # Kathmandu is 5 hours and 45 minutes ahead of UTC. + separator = "+-------------------+---------------------+------------------------+--------------------------+" + assert capsys.readouterr().out == printed( + ( + separator, + "|t |a |s |m |", + separator, + "|2024-01-01 17:45:00|[2024-01-01 17:45:00]|{2024-01-01 17:45:00, 1}|{k -> 2024-01-01 17:45:00}|", + separator, + "", + ) + ) + # collect gives the local time of the process, as Spark does. + assert df.collect()[0].t == instant.astimezone().replace(tzinfo=None) + + def test_nanosecond_timestamps(self, spark, capsys, kathmandu): + if USE_ACTUAL_SPARK: + pytest.skip("Spark has no TIMESTAMPTZ_NS") + df = spark.sql("SELECT TIMESTAMPTZ '2024-01-01 12:00:00.123456789+00'::TIMESTAMPTZ_NS AS t") + df.show(truncate=False) + assert capsys.readouterr().out == printed( + ( + "+--------------------------+", + "|t |", + "+--------------------------+", + "|2024-01-01 17:45:00.123456|", + "+--------------------------+", + "", + ) + ) + # DuckDB truncates nanoseconds to microseconds. + utc = datetime.datetime(2024, 1, 1, 12, 0, 0, 123456, tzinfo=datetime.timezone.utc) + assert df.collect()[0].t == utc diff --git a/tests/fast/spark/test_spark_to_csv.py b/tests/fast/spark/test_spark_to_csv.py index c4f7b79f7..1425031fa 100644 --- a/tests/fast/spark/test_spark_to_csv.py +++ b/tests/fast/spark/test_spark_to_csv.py @@ -3,8 +3,8 @@ import os import pytest -from conftest import getTimeSeriesData from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import IllegalArgumentException import pandas as pd from duckdb import InvalidInputException, read_csv @@ -54,7 +54,7 @@ def test_basic_to_csv(self, pandas_df_ints, spark, tmp_path): df.write.csv(temp_file_name) - csv_rel = spark.read.csv(temp_file_name) + csv_rel = spark.read.csv(temp_file_name, inferSchema=True) assert df.collect() == csv_rel.collect() @@ -65,7 +65,7 @@ def test_to_csv_sep(self, pandas_df_ints, spark, tmp_path): df.write.csv(temp_file_name, sep=",") - csv_rel = spark.read.csv(temp_file_name, sep=",") + csv_rel = spark.read.csv(temp_file_name, inferSchema=True, sep=",") assert df.collect() == csv_rel.collect() def test_to_csv_na_rep(self, spark, tmp_path): @@ -76,7 +76,7 @@ def test_to_csv_na_rep(self, spark, tmp_path): df.write.csv(temp_file_name, nullValue="test") - csv_rel = spark.read.csv(temp_file_name, nullValue="test") + csv_rel = spark.read.csv(temp_file_name, inferSchema=True, nullValue="test") assert df.collect() == csv_rel.collect() def test_to_csv_header(self, spark, tmp_path): @@ -87,7 +87,7 @@ def test_to_csv_header(self, spark, tmp_path): df.write.csv(temp_file_name) - csv_rel = spark.read.csv(temp_file_name) + csv_rel = spark.read.csv(temp_file_name, inferSchema=True) assert df.collect() == csv_rel.collect() def test_to_csv_quotechar(self, spark, tmp_path): @@ -99,7 +99,7 @@ def test_to_csv_quotechar(self, spark, tmp_path): df.write.csv(temp_file_name, quote="'", sep=",") - csv_rel = spark.read.csv(temp_file_name, sep=",", quote="'") + csv_rel = spark.read.csv(temp_file_name, inferSchema=True, sep=",", quote="'") assert df.collect() == csv_rel.collect() def test_to_csv_escapechar(self, spark, tmp_path): @@ -116,33 +116,26 @@ def test_to_csv_escapechar(self, spark, tmp_path): df = spark.createDataFrame(pandas_df) df.write.csv(temp_file_name, quote='"', escape="!") - csv_rel = spark.read.csv(temp_file_name, quote='"', escape="!") + csv_rel = spark.read.csv(temp_file_name, inferSchema=True, quote='"', escape="!") assert df.collect() == csv_rel.collect() def test_to_csv_date_format(self, spark, tmp_path): temp_file_name = os.path.join(tmp_path, "temp_file.csv") # noqa: PTH118 - pandas_df = pd.DataFrame(getTimeSeriesData()) - dt_index = pandas_df.index - pandas_df = pd.DataFrame({"A": dt_index, "B": dt_index.shift(1)}, index=dt_index) + df = spark.createDataFrame([(datetime.date(2024, 2, 3),), (datetime.date(1999, 12, 31),)], "d date") - df = spark.createDataFrame(pandas_df) - - df.write.csv(temp_file_name, dateFormat="%Y%m%d") + df.write.csv(temp_file_name, dateFormat="yyyyMMdd") - csv_rel = spark.read.csv(temp_file_name, dateFormat="%Y%m%d") + csv_rel = spark.read.csv(temp_file_name, schema="d date", dateFormat="yyyyMMdd") assert df.collect() == csv_rel.collect() def test_to_csv_timestamp_format(self, spark, tmp_path): temp_file_name = os.path.join(tmp_path, "temp_file.csv") # noqa: PTH118 - data = [datetime.time(hour=23, minute=1, second=34, microsecond=234345)] - pandas_df = pd.DataFrame({"0": pd.Series(data=data, dtype="object")}) - - df = spark.createDataFrame(pandas_df) + df = spark.createDataFrame([(datetime.datetime(2024, 2, 3, 23, 1, 34),)], "t timestamp") - df.write.csv(temp_file_name, timestampFormat="%m/%d/%Y") + df.write.csv(temp_file_name, timestampFormat="MM/dd/yyyy HH:mm:ss") - csv_rel = spark.read.csv(temp_file_name, timestampFormat="%m/%d/%Y") + csv_rel = spark.read.csv(temp_file_name, schema="t timestamp", timestampFormat="MM/dd/yyyy HH:mm:ss") assert df.collect() == csv_rel.collect() @@ -151,24 +144,23 @@ def test_to_csv_quoting_off(self, pandas_df_strings, spark, tmp_path): df = spark.createDataFrame(pandas_df_strings) df.write.csv(temp_file_name, quoteAll=None) - csv_rel = spark.read.csv(temp_file_name) + csv_rel = spark.read.csv(temp_file_name, inferSchema=True) assert df.collect() == csv_rel.collect() def test_to_csv_quoting_on(self, pandas_df_strings, spark, tmp_path): temp_file_name = os.path.join(tmp_path, "temp_file.csv") # noqa: PTH118 df = spark.createDataFrame(pandas_df_strings) - df.write.csv(temp_file_name, quoteAll="force") + df.write.csv(temp_file_name, quoteAll=True) - csv_rel = spark.read.csv(temp_file_name) + csv_rel = spark.read.csv(temp_file_name, inferSchema=True) assert df.collect() == csv_rel.collect() - def test_to_csv_quoting_quote_all(self, pandas_df_strings, spark, tmp_path): + @pytest.mark.parametrize("quote_all", ["force", csv.QUOTE_ALL]) + def test_to_csv_quoting_not_a_flag(self, pandas_df_strings, spark, tmp_path, quote_all): temp_file_name = os.path.join(tmp_path, "temp_file.csv") # noqa: PTH118 df = spark.createDataFrame(pandas_df_strings) - df.write.csv(temp_file_name, quoteAll=csv.QUOTE_ALL) - - csv_rel = spark.read.csv(temp_file_name) - assert df.collect() == csv_rel.collect() + with pytest.raises(IllegalArgumentException, match="quoteAll flag can be true or false"): + df.write.csv(temp_file_name, quoteAll=quote_all) def test_to_csv_encoding_incorrect(self, pandas_df_strings, spark, tmp_path): temp_file_name = os.path.join(tmp_path, "temp_file.csv") # noqa: PTH118 @@ -182,7 +174,7 @@ def test_to_csv_encoding_correct(self, pandas_df_strings, spark, tmp_path): temp_file_name = os.path.join(tmp_path, "temp_file.csv") # noqa: PTH118 df = spark.createDataFrame(pandas_df_strings) df.write.csv(temp_file_name, encoding="UTF-8") - csv_rel = spark.read.csv(temp_file_name) + csv_rel = spark.read.csv(temp_file_name, inferSchema=True) assert df.collect() == csv_rel.collect() @pytest.mark.skipif( diff --git a/tests/fast/spark/test_spark_type_names.py b/tests/fast/spark/test_spark_type_names.py new file mode 100644 index 000000000..aed90ef6a --- /dev/null +++ b/tests/fast/spark/test_spark_type_names.py @@ -0,0 +1,317 @@ +import pytest + +_ = pytest.importorskip("duckdb.experimental.spark") + +from spark_namespace.errors import PySparkTypeError, PySparkValueError, error_condition +from spark_namespace.sql.types import ( + ArrayType, + BinaryType, + BooleanType, + ByteType, + DateType, + DayTimeIntervalType, + DecimalType, + DoubleType, + FloatType, + IntegerType, + LongType, + MapType, + NullType, + ShortType, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, +) + + +class TestDataTypeNames: + # Expected values are copied from PySpark 4.2.0 and 3.5.9, which agree on all of them. + @pytest.mark.parametrize( + ("data_type", "type_name", "simple_string", "json_value", "json", "representation"), + [ + pytest.param( + ByteType(), + "byte", + "tinyint", + "byte", + '"byte"', + "ByteType()", + id="byte", + ), + pytest.param( + ShortType(), + "short", + "smallint", + "short", + '"short"', + "ShortType()", + id="short", + ), + pytest.param( + IntegerType(), + "integer", + "int", + "integer", + '"integer"', + "IntegerType()", + id="int", + ), + pytest.param( + LongType(), + "long", + "bigint", + "long", + '"long"', + "LongType()", + id="long", + ), + pytest.param( + FloatType(), + "float", + "float", + "float", + '"float"', + "FloatType()", + id="float", + ), + pytest.param( + DoubleType(), + "double", + "double", + "double", + '"double"', + "DoubleType()", + id="double", + ), + pytest.param( + DecimalType(10, 2), + "decimal", + "decimal(10,2)", + "decimal(10,2)", + '"decimal(10,2)"', + "DecimalType(10,2)", + id="decimal", + ), + pytest.param( + DecimalType(), + "decimal", + "decimal(10,0)", + "decimal(10,0)", + '"decimal(10,0)"', + "DecimalType(10,0)", + id="decimal_default", + ), + pytest.param( + StringType(), + "string", + "string", + "string", + '"string"', + "StringType()", + id="string", + ), + pytest.param( + BooleanType(), + "boolean", + "boolean", + "boolean", + '"boolean"', + "BooleanType()", + id="bool", + ), + pytest.param( + DateType(), + "date", + "date", + "date", + '"date"', + "DateType()", + id="date", + ), + pytest.param( + TimestampType(), + "timestamp", + "timestamp", + "timestamp", + '"timestamp"', + "TimestampType()", + id="ts", + ), + pytest.param( + TimestampNTZType(), + "timestamp_ntz", + "timestamp_ntz", + "timestamp_ntz", + '"timestamp_ntz"', + "TimestampNTZType()", + id="ntz", + ), + pytest.param( + BinaryType(), + "binary", + "binary", + "binary", + '"binary"', + "BinaryType()", + id="binary", + ), + pytest.param( + NullType(), + "void", + "void", + "void", + '"void"', + "NullType()", + id="null", + ), + pytest.param( + DayTimeIntervalType(), + "daytimeinterval", + "interval day to second", + "interval day to second", + '"interval day to second"', + "DayTimeIntervalType(0, 3)", + id="interval", + ), + pytest.param( + DayTimeIntervalType(0, 0), + "daytimeinterval", + "interval day", + "interval day", + '"interval day"', + "DayTimeIntervalType(0, 0)", + id="interval_day", + ), + pytest.param( + DayTimeIntervalType(1, 2), + "daytimeinterval", + "interval hour to minute", + "interval hour to minute", + '"interval hour to minute"', + "DayTimeIntervalType(1, 2)", + id="interval_h_m", + ), + pytest.param( + ArrayType(IntegerType()), + "array", + "array", + {"type": "array", "elementType": "integer", "containsNull": True}, + '{"containsNull":true,"elementType":"integer","type":"array"}', + "ArrayType(IntegerType(), True)", + id="array", + ), + pytest.param( + ArrayType(IntegerType(), False), + "array", + "array", + {"type": "array", "elementType": "integer", "containsNull": False}, + '{"containsNull":false,"elementType":"integer","type":"array"}', + "ArrayType(IntegerType(), False)", + id="array_nn", + ), + pytest.param( + MapType(StringType(), LongType()), + "map", + "map", + {"type": "map", "keyType": "string", "valueType": "long", "valueContainsNull": True}, + '{"keyType":"string","type":"map","valueContainsNull":true,"valueType":"long"}', + "MapType(StringType(), LongType(), True)", + id="map", + ), + pytest.param( + StructType([StructField("a", IntegerType()), StructField("b", StringType(), False)]), + "struct", + "struct", + { + "type": "struct", + "fields": [ + {"name": "a", "type": "integer", "nullable": True, "metadata": {}}, + {"name": "b", "type": "string", "nullable": False, "metadata": {}}, + ], + }, + '{"fields":[{"metadata":{},"name":"a","nullable":true,"type":"integer"},{"metadata":{},"name":"b","nullable":false,"type":"string"}],"type":"struct"}', + "StructType([StructField('a', IntegerType(), True), StructField('b', StringType(), False)])", + id="struct", + ), + ], + ) + def test_data_type_names(self, data_type, type_name, simple_string, json_value, json, representation): + assert data_type.typeName() == type_name + assert data_type.simpleString() == simple_string + assert data_type.jsonValue() == json_value + assert data_type.json() == json + assert repr(data_type) == representation + + def test_struct_field_names(self): + field = StructField("x", LongType(), False) + assert field.simpleString() == "x:bigint" + assert field.jsonValue() == {"name": "x", "type": "long", "nullable": False, "metadata": {}} + assert field.json() == '{"metadata":{},"name":"x","nullable":false,"type":"long"}' + assert repr(field) == "StructField('x', LongType(), False)" + with pytest.raises(PySparkTypeError, match="INVALID_TYPENAME_CALL") as error: + field.typeName() + assert isinstance(error.value, TypeError) + assert error_condition(error.value) == "INVALID_TYPENAME_CALL" + + @pytest.mark.parametrize( + ("type_name", "expected"), + [ + ("string", StringType()), + ("integer", IntegerType()), + ("long", LongType()), + ("byte", ByteType()), + ("timestamp", TimestampType()), + ("timestamp_ntz", TimestampNTZType()), + ("decimal", DecimalType()), + ("decimal(5,2)", DecimalType(5, 2)), + ("interval day to second", DayTimeIntervalType()), + ], + ) + def test_struct_type_add_reads_type_names(self, type_name, expected): + schema = StructType().add("f", type_name, False) + assert schema == StructType([StructField("f", expected, False)]) + + @pytest.mark.parametrize("type_name", ["int", "bigint", "array", "struct", "bogus"]) + def test_struct_type_add_rejects_other_strings(self, type_name): + with pytest.raises(PySparkValueError, match="CANNOT_PARSE_DATATYPE"): + StructType().add("f", type_name) + + def test_nested_json_value(self): + schema = StructType( + [ + StructField("a", ArrayType(StructType([StructField("x", DecimalType(5, 2), False)]), False)), + StructField("m", MapType(StringType(), ArrayType(LongType()), False)), + ] + ) + assert schema.jsonValue() == { + "type": "struct", + "fields": [ + { + "name": "a", + "type": { + "type": "array", + "elementType": { + "type": "struct", + "fields": [{"name": "x", "type": "decimal(5,2)", "nullable": False, "metadata": {}}], + }, + "containsNull": False, + }, + "nullable": True, + "metadata": {}, + }, + { + "name": "m", + "type": { + "type": "map", + "keyType": "string", + "valueType": {"type": "array", "elementType": "long", "containsNull": True}, + "valueContainsNull": False, + }, + "nullable": True, + "metadata": {}, + }, + ], + } + assert schema.simpleString() == "struct>,m:map>>" diff --git a/tests/fast/spark/test_spark_udf.py b/tests/fast/spark/test_spark_udf.py index cee0f256e..dcc36ea85 100644 --- a/tests/fast/spark/test_spark_udf.py +++ b/tests/fast/spark/test_spark_udf.py @@ -1,7 +1,25 @@ +import datetime +from decimal import Decimal + import pytest _ = pytest.importorskip("duckdb.experimental.spark") +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import ParseException, PySparkTypeError +from spark_namespace.sql import functions as F +from spark_namespace.sql.types import IntegerType, Row, StringType + +import duckdb +import numpy as np + +if USE_ACTUAL_SPARK: + import pyspark + + SPARK_3 = pyspark.__version__.startswith("3.") +else: + SPARK_3 = False + class TestSparkUDF: def test_udf_register(self, spark): @@ -10,3 +28,320 @@ def to_upper_fn(s: str) -> str: spark.udf.register("to_upper_fn", to_upper_fn) assert spark.sql("select to_upper_fn('quack') as vl").collect()[0].vl == "QUACK" + + def test_udf_register_with_return_type(self, spark): + spark.udf.register("int_x2", lambda i: i * 2, IntegerType()) + spark.udf.register("str_x2", lambda s: s * 2, StringType()) + row = spark.sql("select int_x2(21) as i, str_x2('ab') as s").collect()[0] + assert (row.i, row.s) == (42, "abab") + + def test_udf_register_with_return_type_string(self, spark): + spark.udf.register("int_plus", lambda i: i + 1, "int") + spark.udf.register("arr_of", lambda i: [i, i], "array") + res = spark.sql("select int_plus(41) as i, arr_of(2) as a") + assert res.schema.simpleString() == "struct>" + assert res.collect()[0].asDict() == {"i": 42, "a": [2, 2]} + + @pytest.mark.parametrize( + ("type_string", "value", "spark_type", "expected"), + [ + ("VARCHAR", "x", "string", "x"), + ("json", '{"a":1}', "string", '{"a":1}'), + ("STRUCT(a INT)", {"a": 1}, "struct", Row(a=1)), + ], + ) + def test_udf_return_type_in_duckdb_syntax(self, spark, type_string, value, spark_type, expected): + if USE_ACTUAL_SPARK: + with pytest.raises(ParseException): + spark.udf.register("duckdb_typed", lambda _: value, type_string) + return + # DuckDB's type syntax is read where Spark's grammar rejects the text, as in Column.cast. + registered = spark.udf.register("duckdb_typed", lambda _: value, type_string) + res = spark.range(1).select(registered("id").alias("v")) + assert registered.returnType.simpleString() == spark_type + assert res.schema.simpleString() == f"struct" + assert res.collect() == [Row(v=expected)] + + @pytest.mark.parametrize("type_string", ["VARCHAR(10)[]", "CHAR(3)[]", "INTERVAL[]"]) + def test_udf_return_type_starting_with_a_type_the_module_lacks(self, spark, type_string): + if USE_ACTUAL_SPARK: + with pytest.raises(ParseException): + spark.udf.register("lacking_typed", lambda _: None, type_string) + return + from duckdb.experimental.spark.exception import ContributionsAcceptedError + + # Spark's grammar reads the first type before the text turns out to be DuckDB's, and the module lacks it. + with pytest.raises(ContributionsAcceptedError, match="not supported"): + spark.udf.register("lacking_typed", lambda _: None, type_string) + + def test_udf_in_sql_ignores_annotations(self, spark): + def describe(x: int, *, unused: int = 0) -> str: + return repr(x) + + spark.udf.register("describe_annotated", describe, "string") + # Spark passes the argument as it is and takes only the positional parameters. + row = spark.sql("select describe_annotated(1.5) as d, describe_annotated('3') as s").collect()[0] + assert (row.d, row.s) == ("Decimal('1.5')", "'3'") + + def test_udf_without_a_signature_in_sql(self, spark): + spark.udf.register("largest", max, "int") + assert spark.sql("select largest(1, 3, 2) as v").collect() == [Row(v=3)] + + def test_udf_replaces_a_temporary_macro_in_sql(self, spark): + if not USE_ACTUAL_SPARK: + spark.sql("CREATE TEMP MACRO udf_over_macro(a) AS a") + spark.udf.register("udf_over_macro", lambda v: v + 1, "int") + assert spark.sql("select udf_over_macro(1) as v").collect() == [Row(v=2)] + + def test_udf_register_defaults_to_string_and_returns_the_function(self, spark): + def plus_one(v): + return "none" if v is None else v + 1 + + registered = spark.udf.register("plus_one", plus_one) + assert registered.returnType == StringType() + res = spark.sql("select plus_one(1) as a, plus_one(NULL) as b") + assert res.schema.simpleString() == "struct" + assert res.collect()[0].asDict() == {"a": "2", "b": "none"} + + def test_udf_register_result_is_callable_on_columns(self, spark): + registered = spark.udf.register("times_two", lambda v: None if v is None else v * 2, "bigint") + res = spark.range(2).select(registered("id"), registered(F.col("id") + 1).alias("next")) + assert res.columns == ["times_two(id)", "next"] + assert res.collect() == [Row(0, 2), Row(2, 4)] + + def test_udf_register_result_carries_aggregates(self, spark): + registered = spark.udf.register("plus_ten", lambda v: v + 10, "bigint") + assert spark.range(3).select(registered(F.max("id")).alias("m")).collect() == [Row(m=12)] + + def test_udf_register_replaces_a_function_of_the_same_name(self, spark): + spark.udf.register("twice", lambda v: v * 2, "bigint") + spark.udf.register("twice", lambda v: v * 3, "bigint") + assert spark.sql("select twice(2) as v").collect() == [Row(v=6)] + + def test_earlier_callable_keeps_its_function(self, spark): + first = spark.udf.register("same_name", lambda v: 1, "int") + second = spark.udf.register("same_name", lambda v: 2, "int") + res = spark.range(1).select(first("id").alias("a"), second("id").alias("b")).collect() + assert res == [Row(a=1, b=2)] + assert spark.sql("select same_name(1) as v").collect() == [Row(v=2)] + + def test_udf_struct_arguments_are_rows(self, spark): + schema = "s struct, m map, l array>, ms map>" + df = spark.createDataFrame([((2, "a"), {"k": 1}, [(3,)], {"m": (4,)}), (None, None, None, None)], schema) + + def describe(s, m, l, ms): # noqa: E741 + if s is None: + return "none" + kinds = " ".join(type(v).__name__ for v in (s, m, l[0], ms["m"])) + return f"{kinds} {s['x']} {s['y']} {l[0]['z']} {ms['m']['w']}" + + registered = spark.udf.register("describe_values", describe, "string") + # Spark gives a struct as a Row, also inside an array or a map, and a map as a dict. + expected = [Row(v="Row dict Row Row 2 a 3 4"), Row(v="none")] + assert df.select(registered("s", "m", "l", "ms").alias("v")).collect() == expected + if not USE_ACTUAL_SPARK: + # DuckDB calls the function SQL text names with its own values, a struct as a dict. + expected = [Row(v="dict dict dict dict 2 a 3 4"), Row(v="none")] + assert df.select(F.call_function("describe_values", "s", "m", "l", "ms").alias("v")).collect() == expected + df.createOrReplaceTempView("udf_struct_arguments") + assert spark.sql("select describe_values(s, m, l, ms) as v from udf_struct_arguments").collect() == expected + + def test_udf_timestamp_argument_is_a_local_time(self, spark): + df = spark.createDataFrame([(datetime.datetime(2024, 1, 1, 12),)], "t timestamp") + registered = spark.udf.register("show_time", lambda t: f"{t!r}", "string") + assert df.select(registered("t").alias("v")).collect() == [Row(v="datetime.datetime(2024, 1, 1, 12, 0)")] + + def test_udf_variadic_arguments_in_sql(self, spark): + registered = spark.udf.register("count_values", lambda *values: len(values), "int") + assert spark.sql("select count_values(1, 2) as v").collect() == [Row(v=2)] + assert spark.range(1).select(registered("id", "id").alias("v")).collect() == [Row(v=2)] + + @pytest.mark.skipif(USE_ACTUAL_SPARK, reason="spark.sql takes DuckDB SQL, with its transactions") + def test_register_inside_a_transaction(self, spark): + spark.sql("BEGIN TRANSACTION") + spark.sql("CREATE TABLE udf_in_transaction AS SELECT 1 AS a") + spark.udf.register("plus_one_in_transaction", lambda v: v + 1, "int") + assert spark.sql("SELECT plus_one_in_transaction(a) AS v FROM udf_in_transaction").collect() == [Row(v=2)] + spark.sql("COMMIT") + + @pytest.mark.skipif(USE_ACTUAL_SPARK, reason="spark.sql takes DuckDB SQL, with its schemas and macros") + def test_sql_name_beside_a_macro_of_another_schema(self, spark): + spark.sql("CREATE SCHEMA udf_other_schema") + spark.sql("CREATE MACRO udf_other_schema.udf_beside(a) AS a") + spark.udf.register("udf_beside", lambda v: v + 1, "int") + assert spark.sql("SELECT udf_beside(1) AS v").collect() == [Row(v=2)] + + def test_udf_optional_arguments_in_sql(self, spark): + def plus(a, b=10): + return a + b + + spark.udf.register("plus_optional", plus, "int") + if USE_ACTUAL_SPARK: + assert spark.sql("select plus_optional(1) as a, plus_optional(1, 2) as b").collect() == [Row(a=11, b=3)] + return + # DuckDB takes every parameter of a Python function, defaults included. + assert spark.sql("select plus_optional(1, 2) as b").collect() == [Row(b=3)] + with pytest.raises(duckdb.BinderException): + spark.sql("select plus_optional(1) as a") + + def test_udf_named_like_a_builtin(self, spark): + df = spark.createDataFrame([("a b",)], "t string") + reverse = spark.udf.register("reverse", lambda v: "u", "string") + concat = spark.udf.register("concat", lambda a, b: "c", "string") + res = df.select(reverse("t"), F.reverse("t"), concat("t", "t"), F.concat("t", "t")) + assert res.columns == ["reverse(t)", "reverse(t)", "concat(t, t)", "concat(t, t)"] + assert res.collect() == [("u", "b a", "c", "a ba b")] + # Spark 3.5 runs the function SQL text names by a builtin's name, Spark 4 and DuckDB the builtin. + assert spark.sql("select reverse('a b') as v").collect() == [Row(v="u" if SPARK_3 else "b a")] + + def test_registering_a_registered_udf(self, spark): + df = spark.createDataFrame([("a",)], "t string") + registered = spark.udf.register("exclaim", lambda v: v + "!", "string") + again = spark.udf.register("exclaim_again", registered) + assert again.returnType == StringType() + assert df.select(again("t").alias("v")).collect() == [Row(v="a!")] + assert spark.sql("select exclaim_again('x') as v").collect() == [Row(v="x!")] + with pytest.raises(PySparkTypeError, match="CANNOT_SPECIFY_RETURN_TYPE_FOR_UDF"): + spark.udf.register("exclaim_typed", registered, "int") + + @pytest.mark.parametrize( + ("value", "type_string", "spark_value", "duckdb_value"), + [ + (1, "double", None, 1.0), + (1, "float", None, 1.0), + (1.5, "int", None, 2), + ("1", "int", None, 1), + (True, "int", None, duckdb.TypeMismatchException), + (1, "boolean", None, True), + (Decimal("1.5"), "double", None, 1.5), + (1, "decimal(10,2)", None, Decimal("1.00")), + (1, "string", "1", "1"), + (1.5, "string", "1.5", "1.5"), + (True, "string", "true", duckdb.TypeMismatchException), + ([1, "a"], "string", "[1, a]", "[1, a]"), + ({"a": 1}, "string", "{a=1}", "{'a': 1}"), + (2**40, "int", 0, duckdb.InvalidInputException), + (300, "tinyint", 44, duckdb.InvalidInputException), + ([1.5, 2], "array", [None, 2], [2, 2]), + ({"a": 1.5}, "map", {"a": None}, {"a": 2}), + ((1.5, "x"), "struct", Row(a=None, b="x"), Row(a=2, b="x")), + ({"b": "x", "a": 1}, "struct", Row(a=1, b="x"), Row(a=1, b="x")), + (datetime.datetime(2024, 1, 1, 5), "date", datetime.date(2024, 1, 1), datetime.date(2024, 1, 1)), + ], + ) + def test_udf_result_of_another_python_type(self, spark, value, type_string, spark_value, duckdb_value): + # Spark gives NULL for a result of another Python type, cuts an integer to a narrower type and gives Java's + # text for a string; DuckDB converts the value as it converts any Python value, or raises. + registered = spark.udf.register("typed_result", lambda _: value, type_string) + res = spark.range(1).select(registered("id").alias("v")) + if not USE_ACTUAL_SPARK and isinstance(duckdb_value, type) and issubclass(duckdb_value, duckdb.Error): + with pytest.raises(duckdb_value): + res.collect() + return + assert res.collect() == [Row(v=spark_value if USE_ACTUAL_SPARK else duckdb_value)] + + @pytest.mark.parametrize( + ("value", "spark_text", "duckdb_text"), + [(b"ab", "[B@", "ab"), (datetime.date(2024, 1, 1), "java.util.GregorianCalendar[", "2024-01-01")], + ) + def test_udf_string_result_of_bytes_or_a_date(self, spark, value, spark_text, duckdb_text): + registered = spark.udf.register("text_result", lambda _: value, "string") + (row,) = spark.range(1).select(registered("id").alias("v")).collect() + if USE_ACTUAL_SPARK: + # Spark gives the text of Java's toString, which names an object. + assert row.v.startswith(spark_text) + else: + assert row.v == duckdb_text + + def test_udf_timestamp_result_is_a_local_time(self, spark): + zone = datetime.timezone(datetime.timedelta(hours=2)) + noon = datetime.datetime(2024, 1, 1, 12) + local = spark.udf.register("local_noon", lambda _: noon, "timestamp") + aware = spark.udf.register("aware_noon", lambda _: noon.replace(tzinfo=zone), "timestamp") + res = spark.range(1).select(local("id").alias("l"), aware("id").alias("a")).collect() + assert res == [Row(l=noon, a=noon.replace(tzinfo=zone).astimezone().replace(tzinfo=None))] + nested = spark.udf.register("nested_noon", lambda _: [noon], "array") + if USE_ACTUAL_SPARK: + assert spark.range(1).select(nested("id").alias("t")).collect() == [Row(t=[noon])] + else: + # DuckDB converts a datetime to a timestamp with a time zone only if it has one. + with pytest.raises(duckdb.ConversionException): + spark.range(1).select(nested("id")).collect() + + def test_udf_arguments_are_those_of_the_call(self, spark): + def plus(a: str, b=0): + return a + b + + def count(*arguments): + return len(arguments) + + df = spark.createDataFrame([(1, 2)], "x int, y int") + added = spark.udf.register("plus", plus, "int") + counted = spark.udf.register("count_arguments", count, "int") + seven = spark.udf.register("seven", lambda: 7, "int") + res = df.select( + added("x").alias("one"), + added("x", "y").alias("both"), + counted("x", "y").alias("two"), + seven().alias("none"), + ) + # Spark ignores the annotation, which DuckDB would cast the argument to. + assert res.collect() == [Row(one=1, both=3, two=2, none=7)] + + def test_udf_timestamp_ntz_and_interval_results(self, spark): + zone = datetime.timezone(datetime.timedelta(hours=2)) + noon = datetime.datetime(2024, 1, 1, 12, tzinfo=zone) + wall_clock = spark.udf.register("wall_clock", lambda _: noon, "timestamp_ntz") + span = spark.udf.register("span", lambda _: datetime.timedelta(days=1, seconds=5), "interval day to second") + res = spark.range(1).select(wall_clock("id").alias("t"), span("id").alias("d")) + assert res.collect() == [Row(t=datetime.datetime(2024, 1, 1, 12), d=datetime.timedelta(days=1, seconds=5))] + registered = spark.udf.register("unreadable", lambda _: 5, "interval day to second") + with pytest.raises(Exception): # noqa: B017, PT011 + spark.range(1).select(registered("id")).collect() + registered = spark.udf.register("date_as_timestamp_ntz", lambda _: datetime.date(2024, 1, 1), "timestamp_ntz") + if USE_ACTUAL_SPARK: + with pytest.raises(Exception): # noqa: B017, PT011 + spark.range(1).select(registered("id")).collect() + else: + # DuckDB converts a date to a timestamp at midnight. + res = spark.range(1).select(registered("id").alias("t")).collect() + assert res == [Row(t=datetime.datetime(2024, 1, 1))] + + def test_udf_struct_result(self, spark): + class Point: + def __init__(self) -> None: + self.b = "x" + self.a = 1 + + for value in [Row(1, "x"), (1, "x"), {"b": "x", "a": 1}]: + registered = spark.udf.register("struct_result", lambda _, value=value: value, "a int, b string") + assert spark.range(1).select(registered("id").alias("v")).collect() == [Row(v=Row(a=1, b="x"))] + # Spark reads a Row made from keywords by position, as a tuple. + registered = spark.udf.register("struct_result", lambda _: Row(b=1, a=2), "a int, b int") + assert spark.range(1).select(registered("id").alias("v")).collect() == [Row(v=Row(a=1, b=2))] + registered = spark.udf.register("struct_result", lambda _: Point(), "a int, b string") + if USE_ACTUAL_SPARK: + # Spark reads an object's attributes by field name. + assert spark.range(1).select(registered("id").alias("v")).collect() == [Row(v=Row(a=1, b="x"))] + else: + with pytest.raises(duckdb.NotImplementedException): + spark.range(1).select(registered("id")).collect() + registered = spark.udf.register("struct_result", lambda _: "x", "a string") + error = "UNEXPECTED_TUPLE_WITH_STRUCT" if USE_ACTUAL_SPARK else "can't be cast to the destination type STRUCT" + with pytest.raises(Exception, match=error): + spark.range(1).select(registered("id")).collect() + + def test_udf_result_spark_cannot_read_raises(self, spark): + registered = spark.udf.register("date_as_timestamp", lambda _: datetime.date(2024, 1, 1), "timestamp") + with pytest.raises(Exception, match="date"): + spark.range(1).select(registered("id")).collect() + registered = spark.udf.register("numpy_result", lambda _: np.int64(3), "bigint") + with pytest.raises(Exception): # noqa: B017, PT011 + spark.range(1).select(registered("id")).collect() + registered = spark.udf.register("numpy_result", lambda _: np.float64(1.5), "double") + if USE_ACTUAL_SPARK: + with pytest.raises(Exception): # noqa: B017, PT011 + spark.range(1).select(registered("id")).collect() + else: + # A NumPy float is a Python float, which DuckDB converts. + assert spark.range(1).select(registered("id").alias("v")).collect() == [Row(v=1.5)] diff --git a/tests/fast/spark/test_spark_union.py b/tests/fast/spark/test_spark_union.py index 588c7ecd6..987f40f0c 100644 --- a/tests/fast/spark/test_spark_union.py +++ b/tests/fast/spark/test_spark_union.py @@ -1,12 +1,19 @@ +import datetime import platform +from decimal import Decimal import pytest _ = pytest.importorskip("duckdb.experimental.spark") -from spark_namespace.sql.functions import col +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException, error_condition +from spark_namespace.sql.functions import col, lit from spark_namespace.sql.types import Row +if USE_ACTUAL_SPARK: + import pyspark + @pytest.fixture def df(spark): @@ -69,3 +76,219 @@ def test_merge_without_duplicates(self, df, df2): Row(employee_name="Michael", department="Sales", state="NY", salary=86000, age=56, bonus=20000), Row(employee_name="Robert", department="Sales", state="CA", salary=81000, age=30, bonus=23000), ] + + def test_union_of_an_integer_and_a_string_column(self, spark): + numbers = spark.createDataFrame([(1,)], ["v"]) + strings = spark.createDataFrame([("x",)], ["v"]) + if USE_ACTUAL_SPARK and not pyspark.__version__.startswith("3."): + # Spark 4's ANSI mode casts the string to the integer type. + with pytest.raises(Exception) as error: # noqa: PT011 + numbers.union(strings).collect() + assert error_condition(error.value) == "CAST_INVALID_INPUT" + else: + # Spark 3.5 and DuckDB widen both to a string. + union = numbers.union(strings) + assert union.dtypes == [("v", "string")] + assert sorted(union.collect()) == [Row(v="1"), Row(v="x")] + + +@pytest.fixture +def ab(spark): + return spark.createDataFrame([((1, 2),)], "s struct") + + +@pytest.fixture +def ba(spark): + return spark.createDataFrame([((3, 4),)], "s struct") + + +class TestUnionByPosition: + def test_struct_fields_by_position(self, ab, ba): + union = ab.union(ba) + assert union.schema.simpleString() == "struct>" + assert union.collect() == [Row(s=Row(a=1, b=2)), Row(s=Row(a=3, b=4))] + union = ba.union(ab) + assert union.schema.simpleString() == "struct>" + assert union.collect() == [Row(s=Row(b=3, a=4)), Row(s=Row(b=1, a=2))] + + def test_struct_fields_of_other_names(self, spark, ab): + cd = spark.createDataFrame([((3, 4),)], "s struct") + assert ab.union(cd).collect() == [Row(s=Row(a=1, b=2)), Row(s=Row(a=3, b=4))] + + @pytest.mark.parametrize( + ("first", "second", "values", "expected"), + [ + ("array>", "array>", [(3, 4)], [Row(a=3, b=4)]), + ("map>", "map>", {"k": (3, 4)}, {"k": Row(a=3, b=4)}), + ("struct>", "struct>", ((3, 4),), Row(t=Row(a=3, b=4))), + ], + ) + def test_nested_struct_fields_by_position(self, spark, first, second, values, expected): + left = spark.createDataFrame([], f"s {first}") + right = spark.createDataFrame([(values,)], f"s {second}") + assert left.union(right).collect() == [Row(s=expected)] + + def test_union_by_name_matches_struct_fields_by_name(self, ab, ba): + assert ab.unionByName(ba).collect() == [Row(s=Row(a=1, b=2)), Row(s=Row(a=4, b=3))] + + @pytest.mark.parametrize( + "second", + ["struct", "struct", "struct", "array>"], + ) + def test_struct_fields_that_do_not_match_by_position(self, spark, second): + first = "array>" if second.startswith("array") else "struct" + left = spark.createDataFrame([], f"s {first}") + right = spark.createDataFrame([], f"s {second}") + with pytest.raises(AnalysisException) as error: + left.union(right).collect() + assert error_condition(error.value) == "INCOMPATIBLE_COLUMN_TYPE" + + def test_struct_fields_of_the_same_names_widen(self, spark): + narrow = spark.createDataFrame([((1, "x"),)], "s struct") + wide = spark.createDataFrame([((2, "y"),)], "s struct") + union = narrow.union(wide) + assert union.schema.simpleString() == "struct>" + assert union.collect() == [Row(s=Row(a=1, b="x")), Row(s=Row(a=2, b="y"))] + + def test_number_of_columns(self, spark): + one = spark.createDataFrame([(1,)], ["a"]) + two = spark.createDataFrame([(1, 2)], ["a", "b"]) + for left, right in [(one, two), (two, one)]: + with pytest.raises(AnalysisException) as error: + left.union(right).collect() + assert error_condition(error.value) == "NUM_COLUMNS_MISMATCH" + + @pytest.mark.parametrize("operation", ["union", "intersect", "exceptAll"]) + def test_struct_field_names_in_another_case(self, spark, operation): + lower = spark.createDataFrame([((1,),)], "s struct") + upper = spark.createDataFrame([((1,),)], "s struct") + assert getattr(lower, operation)(upper).count() == {"union": 2, "intersect": 1, "exceptAll": 0}[operation] + + @pytest.mark.parametrize("operation", ["intersect", "intersectAll", "exceptAll"]) + def test_other_set_operations_reject_other_struct_field_names(self, spark, ab, ba, operation): + with pytest.raises(AnalysisException) as error: + getattr(ab, operation)(ba).collect() + assert error_condition(error.value) == "INCOMPATIBLE_COLUMN_TYPE" + same = spark.createDataFrame([((1, 2),)], "s struct") + assert getattr(ab, operation)(same).count() == (0 if operation == "exceptAll" else 1) + + +_VALUES = { + "boolean": True, + "int": 1, + "tinyint": 1, + "bigint": 1, + "double": 1.5, + "decimal(10,2)": Decimal("1.50"), + "string": "1", + "binary": bytearray(b"1"), + "date": datetime.date(2024, 1, 2), + "timestamp": datetime.datetime(2024, 1, 2, 3, 4, 5), + "timestamp_ntz": datetime.datetime(2024, 1, 2, 3, 4, 5), + "interval day to second": datetime.timedelta(days=1), + "array": [1], + "array": [True], + "array": ["1"], + "map": {"a": 1}, + "map": {"a": True}, + "struct": (1,), + "struct": (True,), +} + + +def _column(spark, type_string): + return spark.createDataFrame([(_VALUES[type_string],)], f"v {type_string}") + + +class TestSetOperationTypes: + @pytest.mark.parametrize( + ("first", "second"), + [ + ("boolean", "int"), + ("boolean", "double"), + ("boolean", "decimal(10,2)"), + ("boolean", "binary"), + ("boolean", "date"), + ("boolean", "timestamp"), + ("date", "int"), + ("timestamp", "bigint"), + ("binary", "int"), + ("interval day to second", "int"), + ("array", "int"), + ("array", "array"), + ("map", "map"), + ("struct", "struct"), + ("string", "array"), + ("string", "struct"), + ("string", "map"), + ], + ) + @pytest.mark.parametrize("operation", ["union", "unionByName", "intersect", "intersectAll", "exceptAll"]) + def test_types_of_other_families_raise(self, spark, first, second, operation): + for left, right in [(first, second), (second, first)]: + with pytest.raises(AnalysisException) as error: + getattr(_column(spark, left), operation)(_column(spark, right)).collect() + assert error_condition(error.value) == "INCOMPATIBLE_COLUMN_TYPE" + + def test_type_of_another_family_in_a_later_column(self, spark): + left = spark.createDataFrame([(1, True)], "a int, v boolean") + right = spark.createDataFrame([(1, 2)], "a int, v int") + with pytest.raises(AnalysisException, match="second column") as error: + left.union(right).collect() + assert error_condition(error.value) == "INCOMPATIBLE_COLUMN_TYPE" + + def test_boolean_and_string(self, spark): + for left, right in [("boolean", "string"), ("string", "boolean")]: + if USE_ACTUAL_SPARK and not pyspark.__version__.startswith("3."): + # Spark 4's ANSI mode casts the string to a boolean. + assert _column(spark, left).union(_column(spark, right)).dtypes == [("v", "boolean")] + else: + with pytest.raises(AnalysisException) as error: + _column(spark, left).union(_column(spark, right)).collect() + assert error_condition(error.value) == "INCOMPATIBLE_COLUMN_TYPE" + + @pytest.mark.parametrize( + ("first", "second", "widened"), + [ + ("int", "double", "double"), + ("int", "decimal(10,2)", "decimal(12,2)"), + ("tinyint", "bigint", "bigint"), + ("date", "timestamp", "timestamp"), + ("date", "timestamp_ntz", "timestamp_ntz"), + ("timestamp", "timestamp_ntz", "timestamp"), + ], + ) + def test_types_of_one_family_widen(self, spark, first, second, widened): + for left, right in [(first, second), (second, first)]: + union = _column(spark, left).union(_column(spark, right)) + assert union.dtypes == [("v", widened)] + assert len(union.collect()) == 2 + + @pytest.mark.parametrize( + ("type_string", "value"), + [ + ("map", {"k": 1}), + ("array>", [{"k": 1}]), + ("struct>", ({"k": 1},)), + ], + ) + def test_maps(self, spark, type_string, value): + df = spark.createDataFrame([(1, value)], f"i int, v {type_string}") + nulls = spark.range(1).select(lit(None).alias("i"), lit(None).alias("v")) + assert df.union(df).count() == 2 + assert df.unionByName(df).count() == 2 + assert df.dropDuplicates(["i"]).count() == 1 + compare = [ + *(getattr(df, operation) for operation in ("intersect", "intersectAll", "exceptAll")), + nulls.intersect, + nulls.exceptAll, + lambda other: other.distinct(), + lambda other: other.dropDuplicates(), + lambda other: other.dropDuplicates(["v"]), + lambda other: other.union(other).distinct(), + ] + for operation in compare: + with pytest.raises(AnalysisException) as error: + operation(df).collect() + assert error_condition(error.value) == "UNSUPPORTED_FEATURE.SET_OPERATION_ON_MAP_TYPE" + assert dict(error.value.getMessageParameters())["colName"] == "`v`" diff --git a/tests/fast/spark/test_spark_union_by_name.py b/tests/fast/spark/test_spark_union_by_name.py index 2e4e998ad..afb70515d 100644 --- a/tests/fast/spark/test_spark_union_by_name.py +++ b/tests/fast/spark/test_spark_union_by_name.py @@ -3,10 +3,13 @@ _ = pytest.importorskip("duckdb.experimental.spark") +from spark_namespace.errors import AnalysisException, error_condition from spark_namespace.sql.types import ( Row, ) +import duckdb + @pytest.fixture def df1(spark): @@ -67,3 +70,134 @@ def test_union_by_name_allow_missing_cols_rev(self, df1, df2): Row(name="Maria", id=24), ] assert res == expected + + def test_union_by_name_allow_missing_cols_ignores_case(self, spark): + left = spark.createDataFrame([(1, 2)], ["a", "B"]) + rel = left.unionByName(spark.createDataFrame([(3, 4)], ["b", "c"]), allowMissingColumns=True) + assert rel.columns == ["a", "B", "c"] + assert rel.collect() == [Row(a=1, B=2, c=None), Row(a=None, B=3, c=4)] + rel = left.unionByName(spark.createDataFrame([(5, 6)], ["A", "z"]), allowMissingColumns=True) + assert rel.columns == ["a", "B", "z"] + assert rel.collect() == [Row(a=1, B=2, z=None), Row(a=5, B=None, z=6)] + + @pytest.mark.parametrize(("left", "right"), [(["a"], ["a", "b"]), (["a"], ["b", "a"])]) + def test_union_by_name_with_extra_columns_in_the_other(self, spark, left, right): + first = spark.createDataFrame([tuple(range(len(left)))], left) + second = spark.createDataFrame([tuple(range(len(right)))], right) + with pytest.raises(AnalysisException) as error: + first.unionByName(second).collect() + assert error_condition(error.value) == "NUM_COLUMNS_MISMATCH" + + def test_union_by_name_with_extra_columns_in_the_first(self, spark): + first = spark.createDataFrame([(1, 2)], ["a", "b"]) + second = spark.createDataFrame([(3,)], ["a"]) + # Spark raises AnalysisException; DuckDB's binder does not find the column. + with pytest.raises((AnalysisException, duckdb.BinderException)): + first.unionByName(second).collect() + + def test_union_by_name_of_names_with_a_backtick_or_a_dot(self, spark): + df = spark.createDataFrame([(1, 2, 3)], "k int, `a``b` int, `c.d` int") + other = spark.createDataFrame([(3, 4, 5)], "`C.D` int, `A``B` int, K int") + rel = df.unionByName(other) + assert rel.columns == ["k", "a`b", "c.d"] + assert [tuple(r) for r in rel.collect()] == [(1, 2, 3), (5, 4, 3)] + rel = df.unionByName(df.select("k"), allowMissingColumns=True) + assert [tuple(r) for r in rel.collect()] == [(1, 2, 3), (1, None, None)] + + +class TestUnionByNameOfStructFields: + @pytest.mark.parametrize( + ("first", "second", "values", "field", "fields"), + [ + ("struct", "struct", ((1,), (2,)), "a", "b"), + ("struct", "struct", ((1, 2), (3,)), "b", "a"), + ("struct>", "struct>", (((1,),), ((2,),)), "a", "b"), + ("array>", "array>", ([(1,)], [(2,)]), "a", "b"), + ], + ) + def test_missing_struct_field_raises(self, spark, first, second, values, field, fields): + left = spark.createDataFrame([(values[0],)], f"s {first}") + right = spark.createDataFrame([(values[1],)], f"s {second}") + with pytest.raises(AnalysisException) as error: + left.unionByName(right).collect() + assert error_condition(error.value) == "FIELD_NOT_FOUND" + assert dict(error.value.getMessageParameters()) == {"fieldName": f"`{field}`", "fields": f"`{fields}`"} + + def test_extra_struct_field_in_the_other_raises(self, spark): + left = spark.createDataFrame([((1,),)], "s struct") + right = spark.createDataFrame([((1, 2),)], "s struct") + with pytest.raises(AnalysisException) as error: + left.unionByName(right).collect() + assert error_condition(error.value) == "INCOMPATIBLE_COLUMN_TYPE" + + def test_missing_struct_fields_are_filled_where_missing_columns_are_allowed(self, spark): + a = spark.createDataFrame([((1,),)], "s struct") + b = spark.createDataFrame([((2,),)], "s struct") + ab = spark.createDataFrame([((3, 4),)], "s struct") + union = a.unionByName(b, allowMissingColumns=True) + assert union.schema.simpleString() == "struct>" + assert union.collect() == [Row(s=Row(a=1, b=None)), Row(s=Row(a=None, b=2))] + assert ab.unionByName(a, allowMissingColumns=True).collect() == [Row(s=Row(a=3, b=4)), Row(s=Row(a=1, b=None))] + + @pytest.mark.parametrize("allow_missing", [False, True]) + @pytest.mark.parametrize( + ("first", "second", "values", "expected"), + [ + ( + "map>", + "map>", + ({"k": (1, 2)}, {"k": (3, 4)}), + [{"k": Row(a=1, b=2)}, {"k": Row(a=3, b=4)}], + ), + ( + "map>", + "map>", + ({"k": (1,)}, {"k": (2,)}), + [{"k": Row(a=1)}, {"k": Row(a=2)}], + ), + ( + "map,int>", + "map,int>", + ({(1, 2): 1}, {(3, 4): 2}), + [{Row(a=1, b=2): 1}, {Row(a=3, b=4): 2}], + ), + ( + "array>>", + "array>>", + ([{"k": (1, 2)}], [{"k": (3, 4)}]), + [[{"k": Row(a=1, b=2)}], [{"k": Row(a=3, b=4)}]], + ), + ( + "map>>", + "map>>", + ({"k": [(1, 2)]}, {"k": [(3, 4)]}), + [{"k": [Row(a=1, b=2)]}, {"k": [Row(a=3, b=4)]}], + ), + ], + ) + def test_struct_fields_inside_a_map_pair_by_position(self, spark, first, second, values, expected, allow_missing): + left = spark.createDataFrame([(values[0],)], f"v {first}") + right = spark.createDataFrame([(values[1],)], f"v {second}") + union = left.unionByName(right, allowMissingColumns=allow_missing) + assert union.schema.simpleString() == f"struct" + assert union.collect() == [Row(v=value) for value in expected] + + @pytest.mark.parametrize("allow_missing", [False, True]) + @pytest.mark.parametrize( + ("first", "second", "values"), + [ + ("map>", "map>", ({"k": (1,)}, {"k": (2,)})), + ("map>", "map>", ({"k": (1, 2)}, {"k": (3,)})), + ], + ) + def test_struct_fields_inside_a_map_of_another_shape(self, spark, first, second, values, allow_missing): + left = spark.createDataFrame([(values[0],)], f"v {first}") + right = spark.createDataFrame([(values[1],)], f"v {second}") + with pytest.raises(AnalysisException) as error: + left.unionByName(right, allowMissingColumns=allow_missing).collect() + assert error_condition(error.value) == "INCOMPATIBLE_COLUMN_TYPE" + + def test_struct_fields_inside_an_array_pair_by_name(self, spark): + left = spark.createDataFrame([([(1, 2)],)], "v array>") + right = spark.createDataFrame([([(3, 4)],)], "v array>") + assert left.unionByName(right).collect() == [Row(v=[Row(a=1, b=2)]), Row(v=[Row(a=4, b=3)])] diff --git a/tests/fast/spark/test_spark_with_column.py b/tests/fast/spark/test_spark_with_column.py index 4ea62fe11..a9c414057 100644 --- a/tests/fast/spark/test_spark_with_column.py +++ b/tests/fast/spark/test_spark_with_column.py @@ -2,7 +2,6 @@ _ = pytest.importorskip("duckdb.experimental.spark") -from spark_namespace import USE_ACTUAL_SPARK from spark_namespace.sql.functions import col, lit @@ -18,7 +17,7 @@ def test_with_column(self, spark): columns = ["firstname", "middlename", "lastname", "dob", "gender", "salary"] df = spark.createDataFrame(data=data, schema=columns) - assert df.schema["salary"].dataType.typeName() == ("long" if USE_ACTUAL_SPARK else "integer") + assert df.schema["salary"].dataType.typeName() == "long" # The type of 'salary' has been cast to Bigint new_df = df.withColumn("salary", col("salary").cast("BIGINT")) diff --git a/tests/fast/spark/test_spark_with_columns.py b/tests/fast/spark/test_spark_with_columns.py index 244d40a3b..5d24fe574 100644 --- a/tests/fast/spark/test_spark_with_columns.py +++ b/tests/fast/spark/test_spark_with_columns.py @@ -3,8 +3,9 @@ _ = pytest.importorskip("duckdb.experimental.spark") -from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException, error_condition from spark_namespace.sql.functions import col, lit +from spark_namespace.sql.types import Row class TestWithColumns: @@ -19,7 +20,7 @@ def test_with_columns(self, spark): columns = ["firstname", "middlename", "lastname", "dob", "gender", "salary"] df = spark.createDataFrame(data=data, schema=columns) - assert df.schema["salary"].dataType.typeName() == ("long" if USE_ACTUAL_SPARK else "integer") + assert df.schema["salary"].dataType.typeName() == "long" # The type of 'salary' has been cast to Bigint new_df = df.withColumns({"salary": col("salary").cast("BIGINT")}) @@ -43,3 +44,31 @@ def test_with_columns(self, spark): df2 = df.drop("salary") assert "salary" not in df2.columns + + def test_replace_several_existing_columns(self, spark): + df = spark.createDataFrame([(1, 2)], ["a", "b"]) + res = df.withColumns({"a": lit(0), "b": lit(5)}) + assert res.columns == ["a", "b"] + assert res.collect() == [(0, 5)] + res = df.withColumns({"b": lit(5), "c": lit(7), "a": lit(0)}) + assert res.columns == ["a", "b", "c"] + assert res.collect() == [(0, 5, 7)] + + def test_replace_column_named_in_another_case(self, spark): + df = spark.createDataFrame([(1, 2)], ["A", "b"]) + res = df.withColumns({"a": lit(0)}) + assert res.columns == ["a", "b"] + assert res.collect() == [(0, 2)] + + @pytest.mark.parametrize("names", [("x", "X"), ("B", "b")]) + def test_names_differing_only_in_case_raise(self, spark, names): + df = spark.createDataFrame([("a", 1)], "a string, b int") + with pytest.raises(AnalysisException) as error: + df.withColumns({names[0]: lit(1), names[1]: lit(2)}) + assert error_condition(error.value) == "COLUMN_ALREADY_EXISTS" + assert dict(error.value.getMessageParameters()) == {"columnName": f"`{names[1].lower()}`"} + assert error.value.getSqlState() == "42711" + + def test_name_in_another_case_replaces_the_column(self, spark): + df = spark.createDataFrame([("a", 1)], "a string, b int") + assert df.withColumns({"A": lit(1)}).collect() == [Row(A=1, b=1)] diff --git a/tests/fast/spark/test_spark_with_columns_renamed.py b/tests/fast/spark/test_spark_with_columns_renamed.py index 8c24062ba..44579cf99 100644 --- a/tests/fast/spark/test_spark_with_columns_renamed.py +++ b/tests/fast/spark/test_spark_with_columns_renamed.py @@ -53,3 +53,9 @@ def test_with_columns_renamed(self, spark): if not USE_ACTUAL_SPARK: with pytest.raises(ValueError, match=re.escape("DataFrame does not contain column(s): unknown col")): df2.withColumnsRenamed({"unknown col": "new name", "full name": "name"}) + + def test_rename_column_named_in_another_case(self, spark): + df = spark.createDataFrame([(1, 2)], ["A", "b"]) + res = df.withColumnsRenamed({"a": "x", "B": "y"}) + assert res.columns == ["x", "y"] + assert res.collect() == [(1, 2)] diff --git a/tests/fast/spark/test_spark_writer.py b/tests/fast/spark/test_spark_writer.py new file mode 100644 index 000000000..51f9af561 --- /dev/null +++ b/tests/fast/spark/test_spark_writer.py @@ -0,0 +1,218 @@ +import datetime +import re + +import pytest + +_ = pytest.importorskip("duckdb.experimental.spark") + +from spark_namespace import USE_ACTUAL_SPARK +from spark_namespace.errors import AnalysisException, IllegalArgumentException, error_condition + +from duckdb.experimental.spark.exception import ContributionsAcceptedError + +if USE_ACTUAL_SPARK: + import pyspark + + SPARK_3 = pyspark.__version__.startswith("3.") +else: + SPARK_3 = False + + +def _written(path): + """The text of a CSV output, which Spark writes as a directory of part files.""" + if path.is_dir(): + return "".join(part.read_text(encoding="utf-8") for part in sorted(path.glob("part-*"))) + return path.read_text(encoding="utf-8") + + +def _one(df): + return df.coalesce(1) if USE_ACTUAL_SPARK else df + + +def _offset(value): + """The offset of the local time zone at ``value`` as Spark's XXX pattern letters write it, and as DuckDB does.""" + total = int(value.astimezone().utcoffset().total_seconds()) // 60 + hours, minutes = divmod(abs(total), 60) + sign = "+" if total >= 0 else "-" + spark_offset = "Z" if total == 0 else f"{sign}{hours:02d}:{minutes:02d}" + return spark_offset if USE_ACTUAL_SPARK else f"{sign}{hours:02d}" + (f":{minutes:02d}" if minutes else "") + + +@pytest.fixture(params=["csv", "parquet"]) +def kind(request): + return request.param + + +def _write(df, kind, path, **options): + getattr(df.write, kind)(path.as_posix(), **options) + + +def _read_back(spark, kind, path): + return sorted(tuple(row) for row in getattr(spark.read, kind)(path.as_posix()).collect()) + + +class TestSaveModes: + @pytest.mark.parametrize("mode", [None, "error", "errorifexists", "default", "ERROR"]) + def test_existing_path_fails_the_write(self, spark, tmp_path, kind, mode): + path = tmp_path / "out" + _write(spark.createDataFrame([(1,)], "a int"), kind, path) + with pytest.raises(AnalysisException, match=r'already exists\. Set mode as "overwrite"') as error: + _write(spark.createDataFrame([(2,)], "a int"), kind, path, mode=mode) + assert error_condition(error.value) == "PATH_ALREADY_EXISTS" + assert _read_back(spark, kind, path) == [("1",) if kind == "csv" else (1,)] + + def test_existing_file_fails_the_write(self, spark, tmp_path, kind): + path = tmp_path / "out" + path.write_text("x\n", newline="", encoding="utf-8") + with pytest.raises(AnalysisException, match=re.escape(f"Path file:{path} already exists")) as error: + _write(spark.createDataFrame([(1,)], "a int"), kind, path) + assert error_condition(error.value) == "PATH_ALREADY_EXISTS" + + @pytest.mark.parametrize(("mode", "first"), [("ignore", True), ("Ignore", True), ("overwrite", False)]) + def test_existing_path_kept_or_replaced(self, spark, tmp_path, kind, mode, first): + path = tmp_path / "out" + _write(spark.createDataFrame([(1,)], "a int"), kind, path) + _write(spark.createDataFrame([(2,)], "a int"), kind, path, mode=mode) + value = 1 if first else 2 + assert _read_back(spark, kind, path) == [(str(value),) if kind == "csv" else (value,)] + + @pytest.mark.parametrize("mode", ["append", "ignore", "overwrite", "error"]) + def test_new_path(self, spark, tmp_path, kind, mode): + path = tmp_path / "out" + _write(spark.createDataFrame([(1,)], "a int"), kind, path, mode=mode) + assert _read_back(spark, kind, path) == [("1",) if kind == "csv" else (1,)] + + def test_append_to_an_existing_path(self, spark, tmp_path, kind): + path = tmp_path / "out" + _write(spark.createDataFrame([(1,)], "a int"), kind, path) + if USE_ACTUAL_SPARK: + _write(spark.createDataFrame([(2,)], "a int"), kind, path, mode="append") + assert _read_back(spark, kind, path) == ([("1",), ("2",)] if kind == "csv" else [(1,), (2,)]) + else: + # DuckDB writes one file, which it cannot add to. + with pytest.raises(ContributionsAcceptedError): + _write(spark.createDataFrame([(2,)], "a int"), kind, path, mode="append") + + def test_file_url(self, spark, tmp_path, kind): + url = (tmp_path / "out").as_uri() + _write(spark.createDataFrame([(1,)], "a int"), kind, tmp_path / "out") + with pytest.raises(AnalysisException) as error: + getattr(spark.createDataFrame([(2,)], "a int").write, kind)(url) + assert error_condition(error.value) == "PATH_ALREADY_EXISTS" + getattr(spark.createDataFrame([(2,)], "a int").write, kind)(url, mode="ignore") + assert _read_back(spark, kind, tmp_path / "out") == [("1",) if kind == "csv" else (1,)] + + def test_invalid_mode(self, spark, tmp_path, kind): + # Spark 3.5 has no error class for it. + expected, condition = (IllegalArgumentException, None) if SPARK_3 else (AnalysisException, "INVALID_SAVE_MODE") + with pytest.raises(expected, match="bogus") as error: + _write(spark.createDataFrame([(1,)], "a int"), kind, tmp_path / "out", mode="bogus") + assert error_condition(error.value) == condition + + +class TestRepeatedColumnNames: + @pytest.mark.parametrize("names", [["x", "x"], ["x", "X"]]) + def test_repeated_names_fail_the_write(self, spark, tmp_path, kind, names): + df = spark.createDataFrame([(1, 2)], names) + path = tmp_path / "out" + if USE_ACTUAL_SPARK and kind == "csv" and not SPARK_3: + # Spark 4 writes the repeated names into the header line. + _write(df, kind, path, header=True) + assert _written(path).splitlines()[0] == ",".join(names) + return + with pytest.raises(AnalysisException) as error: + _write(df, kind, path) + assert error_condition(error.value) == "COLUMN_ALREADY_EXISTS" + assert dict(error.value.getMessageParameters()) == {"columnName": "`x`"} + assert not path.exists() + + @pytest.mark.parametrize("names", [["x", "x"], ["x", "X"]]) + def test_repeated_names_fail_to_save_a_table(self, spark, names): + with pytest.raises(AnalysisException) as error: + spark.createDataFrame([(1, 2)], names).write.saveAsTable("repeated_names_table") + assert error_condition(error.value) == "COLUMN_ALREADY_EXISTS" + + +class TestCsvOutput: + def test_quotes_and_backslashes(self, spark, tmp_path): + values = ["C:\\temp", "a,b\\c", 'q"\\', "x\\", "", None] + path = tmp_path / "out" + _one(spark.createDataFrame([(v,) for v in values], "s string")).write.csv(path.as_posix()) + if USE_ACTUAL_SPARK: + # Spark escapes a quote and a backslash inside quotes with a backslash. + assert _written(path) == 'C:\\temp\n"a,b\\\\c"\n"q\\"\\\\"\nx\\\n""\n' + else: + # DuckDB doubles a quote inside quotes, and writes an empty line for a row that is NULL in its only column. + assert _written(path) == 'C:\\temp\n"a,b\\c"\n"q""\\"\nx\\\n""\n\n' + rows = [row[0] for row in spark.read.csv(path.as_posix()).collect()] + assert sorted(row for row in rows if row is not None) == sorted(v for v in values if v) + # The empty string is NULL, and DuckDB reads its empty line as a NULL row of the one column. + assert rows.count(None) == (1 if USE_ACTUAL_SPARK else 2) + + def test_whitespace(self, spark, tmp_path): + df = _one(spark.createDataFrame([(" a ",), ("\tb\x01",), ("\xa0c\xa0",)], "s string")) + path = tmp_path / "kept" + df.write.csv(path.as_posix(), ignoreLeadingWhiteSpace=False, ignoreTrailingWhiteSpace="FALSE") + assert _written(path) == " a \n\tb\x01\n\xa0c\xa0\n" + path = tmp_path / "default" + df.write.csv(path.as_posix()) + # Spark trims strings when it writes them, DuckDB writes them as they are. + assert _written(path) == ("a\nb\n\xa0c\xa0\n" if USE_ACTUAL_SPARK else " a \n\tb\x01\n\xa0c\xa0\n") + if not USE_ACTUAL_SPARK: + for option in ("ignoreLeadingWhiteSpace", "ignoreTrailingWhiteSpace"): + with pytest.raises(ContributionsAcceptedError): + df.write.csv((tmp_path / option).as_posix(), **{option: "true"}) + + def test_escape_quotes_spellings(self, spark, tmp_path): + df = spark.createDataFrame([('a"b',)], "s string") + for i, value in enumerate([True, "TRUE", "true"]): + path = tmp_path / f"out{i}" + df.write.csv(path.as_posix(), escapeQuotes=value) + assert _written(path) == ('"a\\"b"\n' if USE_ACTUAL_SPARK else '"a""b"\n') + + def test_datetimes(self, spark, tmp_path): + moment = datetime.datetime(2024, 2, 3, 4, 5, 6, 789000) + schema = "t timestamp, n timestamp_ntz, d date" + df = spark.createDataFrame([(moment, moment, datetime.date(2024, 2, 3))], schema) + path = tmp_path / "default" + df.write.csv(path.as_posix()) + if USE_ACTUAL_SPARK: + expected = f"2024-02-03T04:05:06.789{_offset(moment)},2024-02-03T04:05:06.789,2024-02-03\n" + else: + # DuckDB writes its own text of a timestamp, which Spark's default pattern does not. + expected = f"2024-02-03 04:05:06.789{_offset(moment)},2024-02-03 04:05:06.789,2024-02-03\n" + assert _written(path) == expected + path = tmp_path / "patterns" + df.write.csv(path.as_posix(), timestampFormat="dd.MM.yyyy HH:mm", dateFormat="dd/MM/yyyy") + # Spark's timestampFormat is for timestamps with a time zone only, DuckDB's for every timestamp. + without_zone = "2024-02-03T04:05:06.789" if USE_ACTUAL_SPARK else "03.02.2024 04:05" + assert _written(path) == f"03.02.2024 04:05,{without_zone},03/02/2024\n" + + def test_offset_pattern(self, spark, tmp_path): + moment = datetime.datetime(2024, 7, 1, 12) + path = tmp_path / "out" + df = spark.createDataFrame([(moment,)], "t timestamp") + if USE_ACTUAL_SPARK: + df.write.csv(path.as_posix(), timestampFormat="HH:mmXXX") + assert _written(path) == f"12:00{_offset(moment)}\n" + else: + with pytest.raises(ContributionsAcceptedError): + df.write.csv(path.as_posix(), timestampFormat="HH:mmXXX") + + @pytest.mark.parametrize( + ("schema", "value", "name"), + [("a array", [1], "ARRAY"), ("a struct", (1,), "STRUCT"), ("a void", None, "VOID")], + ) + def test_types_csv_cannot_hold(self, spark, tmp_path, schema, value, name): + df = spark.createDataFrame([(value,)], schema) + with pytest.raises(AnalysisException, match=f'column `a` of the type "{name}"') as error: + df.write.csv((tmp_path / "out").as_posix()) + assert error_condition(error.value) == "UNSUPPORTED_DATA_TYPE_FOR_DATASOURCE" + + def test_save_mode_comes_before_the_types(self, spark, tmp_path): + df = spark.createDataFrame([([1],)], "a array") + with pytest.raises(Exception, match="bogus"): + df.write.csv((tmp_path / "new").as_posix(), mode="bogus") + spark.createDataFrame([(1,)], "a int").write.csv((tmp_path / "out").as_posix()) + df.write.csv((tmp_path / "out").as_posix(), mode="ignore") + assert _written(tmp_path / "out") == "1\n" diff --git a/tests/fast/test_expression.py b/tests/fast/test_expression.py index fb0d77012..f93c8f889 100644 --- a/tests/fast/test_expression.py +++ b/tests/fast/test_expression.py @@ -22,6 +22,16 @@ ) +@pytest.fixture +def schema_join(): + con = duckdb.connect() + con.execute("create schema s1") + con.execute("create table s1.t1 as select 1 as id, 2 as x") + con.execute("create table t2 as select 1 as id, 3 as y") + yield con.table("s1.t1").join(con.table("t2"), "t1.id = t2.id") + con.close() + + @pytest.fixture def filter_rel(): con = duckdb.connect() @@ -635,6 +645,70 @@ def test_star_expression(self): res = rel2.fetchall() assert res == [(2,)] + @pytest.mark.parametrize( + "excluded", + [ + "b.id", + ColumnExpression("b", "id"), + ColumnExpression("b.id"), + ], + ) + def test_star_expression_exclude_qualified(self, excluded): + con = duckdb.connect() + a = con.sql("select 1 as id, 10 as x").set_alias("a") + b = con.sql("select 1 as id, 20 as y").set_alias("b") + joined = a.join(b, "a.id = b.id") + + rel = joined.select(StarExpression(exclude=[excluded])) + assert rel.columns == ["id", "x", "y"] + assert rel.fetchall() == [(1, 10, 20)] + + @pytest.mark.parametrize( + ("excluded", "columns"), + [ + (ColumnExpression("s1", "t1", "id"), ["x", "id", "y"]), + (ColumnExpression("memory", "s1", "t1", "id"), ["x", "id", "y"]), + (ColumnExpression("memory.s1.t1.id"), ["x", "id", "y"]), + ("memory.s1.t1.id", ["x", "id", "y"]), + (ColumnExpression("t2", "id"), ["id", "x", "y"]), + ], + ) + def test_star_expression_exclude_schema_and_catalog(self, schema_join, excluded, columns): + assert schema_join.select(StarExpression(exclude=[excluded])).columns == columns + + @pytest.mark.parametrize("excluded", [ColumnExpression("zz", "id"), ColumnExpression("s9", "t1", "id")]) + def test_star_expression_exclude_unknown_qualifier(self, schema_join, excluded): + with pytest.raises(duckdb.BinderException, match="in EXCLUDE list not found"): + schema_join.select(StarExpression(exclude=[excluded])) + + @pytest.mark.parametrize("names", [(), ("c", "s", "t", "col", "x"), ("c.s.t.col.x",)]) + def test_star_expression_exclude_name_part_count(self, names): + with pytest.raises(ValueError, match="1 to 4 name parts"): + StarExpression(exclude=[ColumnExpression(*names)]) + + @pytest.mark.parametrize( + ("dotted", "parts"), + [ + ("a", ("a",)), + ("a.b.c.d", ("a", "b", "c", "d")), + ("a.b.c.d.e", ("a", "b", "c", "d", "e")), + ('"a.b".c', ("a.b", "c")), + ], + ) + def test_column_expression_keeps_every_name_part(self, dotted, parts): + assert str(ColumnExpression(dotted)) == str(ColumnExpression(*parts)) + + @pytest.mark.parametrize("column", ["s.a.b.c", "t.s.a.b.c"]) + def test_column_expression_nested_struct_path(self, column): + con = duckdb.connect() + rel = con.sql("select {'a': {'b': {'c': 1}}, 'b': {'c': 2}} as s").set_alias("t") + assert rel.select(ColumnExpression(column)).fetchall() == [(1,)] + + def test_column_expression_quoted_dot(self): + con = duckdb.connect() + rel = con.sql('select 1 as "b.a", 2 as b') + assert rel.select(ColumnExpression('"b.a"')).fetchall() == [(1,)] + def test_struct_expression(self): con = duckdb.connect() diff --git a/tests/fast/test_optional_dependencies.py b/tests/fast/test_optional_dependencies.py new file mode 100644 index 000000000..2cacc71f3 --- /dev/null +++ b/tests/fast/test_optional_dependencies.py @@ -0,0 +1,229 @@ +import importlib.util +import subprocess +import sys +import textwrap + +import pytest + +OPTIONAL_PACKAGES = ( + "numpy", + "pandas", + "pyarrow", + "polars", + "fsspec", + "pytz", + "typing_extensions", + "pyspark", + "IPython", + "adbc_driver_manager", +) + +# The child starts with -S and runs site.main() after installing the finder, so no .pth file imports a package first. +BLOCK_IMPORTS = """ +import sys +from importlib.abc import MetaPathFinder + +BLOCKED = frozenset(sys.argv[1].split(",")) + + +class BlockImports(MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.partition(".")[0] in BLOCKED: + raise ModuleNotFoundError(f"No module named {fullname!r}", name=fullname) + return None + + +sys.meta_path.insert(0, BlockImports()) + +import site + +site.main() +""" + + +def run_blocked(cwd, script, blocked=OPTIONAL_PACKAGES): + # An empty working directory keeps modules next to the caller's off the path that -c puts first. + proc = subprocess.run( + [sys.executable, "-S", "-c", BLOCK_IMPORTS + textwrap.dedent(script), ",".join(blocked)], + capture_output=True, + text=True, + timeout=60, + check=False, + cwd=cwd, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + + +@pytest.mark.parametrize("package", OPTIONAL_PACKAGES) +def test_blocked_package_does_not_import(tmp_path, package): + if importlib.util.find_spec(package) is None: + pytest.skip(f"{package} is not installed, so blocking it proves nothing") + run_blocked( + tmp_path, + f""" + try: + import {package} + except ModuleNotFoundError: + pass + else: + raise AssertionError("{package} imported") + """, + ) + + +def test_duckdb_runs_a_query(tmp_path): + run_blocked( + tmp_path, + """ + import duckdb + + assert duckdb.sql("SELECT 42 AS answer").fetchall() == [(42,)] + """, + ) + + +def test_spark_basic_use(tmp_path): + run_blocked( + tmp_path, + """ + from duckdb.experimental.spark.sql import SparkSession + from duckdb.experimental.spark.sql import functions as F + from duckdb.experimental.spark.sql.types import Row + + spark = SparkSession.builder.getOrCreate() + df = spark.createDataFrame([(1, "a", 10), (2, "b", 20), (3, "a", 30)], ["id", "k", "v"]) + + assert df.select("id", "k").filter(F.col("id") > 1).orderBy("id").collect() == [ + Row(id=2, k="b"), + Row(id=3, k="a"), + ] + assert df.groupBy("k").agg(F.sum("v").alias("total")).orderBy("k").collect() == [ + Row(k="a", total=40), + Row(k="b", total=20), + ] + names = spark.createDataFrame([("a", "first")], ["k", "name"]) + assert df.join(names, "k").orderBy("id").select("id", "name").collect() == [ + Row(id=1, name="first"), + Row(id=3, name="first"), + ] + assert spark.sql("SELECT 1 AS one").collect() == [Row(one=1)] + """, + ) + + +@pytest.mark.parametrize( + ("method", "package", "blocked"), + [ + ("toPandas", "pandas", OPTIONAL_PACKAGES), + ("toPandas", "pandas", ("pandas",)), + ("toArrow", "pyarrow", OPTIONAL_PACKAGES), + ("toArrow", "pyarrow", ("pyarrow",)), + ], +) +def test_spark_conversion_names_the_missing_package(tmp_path, method, package, blocked): + run_blocked( + tmp_path, + f""" + from duckdb.experimental.spark.sql import SparkSession + + df = SparkSession.builder.getOrCreate().createDataFrame([(1, "a")], ["id", "k"]) + try: + df.{method}() + except ImportError as error: + assert error.name == {package!r}, repr(error) + assert {package!r} in str(error), repr(error) + else: + raise AssertionError("{method}() did not raise ImportError") + """, + blocked, + ) + + +def test_spark_timestamps(tmp_path): + run_blocked( + tmp_path, + """ + import datetime + + from duckdb.experimental.spark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + value = datetime.datetime(2024, 1, 2, 3, 4, 5) + schema = "id long, ts timestamp, l array, s struct" + df = spark.createDataFrame([(1, value, [value], (value,))], schema) + row = df.collect()[0] + assert (row.ts, row.l, row.s.t) == (value, [value], value), row + assert df.first().ts == value + df.show() + (row,) = spark.sql("SELECT TIMESTAMPTZ '2024-01-02 03:04:05+00'::TIMESTAMPTZ_NS AS t").collect() + assert row.t == datetime.datetime(2024, 1, 2, 3, 4, 5, tzinfo=datetime.timezone.utc), row + """, + ) + + +def test_spark_timestamps_to_pandas(tmp_path): + # pandas before 3.0 cannot be imported without pytz. + pytest.importorskip("pandas", minversion="3.0.0") + run_blocked( + tmp_path, + """ + import datetime + + from duckdb.experimental.spark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + value = datetime.datetime(2024, 1, 2, 3, 4, 5) + pdf = spark.createDataFrame([(1, value)], ["id", "ts"]).toPandas() + assert pdf["ts"].tolist() == [value], pdf + """, + tuple(package for package in OPTIONAL_PACKAGES if package not in ("numpy", "pandas")), + ) + + +@pytest.mark.parametrize("blocked", [OPTIONAL_PACKAGES, ("numpy",)]) +def test_spark_udf_needs_numpy(tmp_path, blocked): + """The binding's Python UDF path needs NumPy to register a function.""" + run_blocked( + tmp_path, + """ + import duckdb + from duckdb.experimental.spark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + try: + spark.udf.register("plus_one", lambda n: n + 1, "int") + except duckdb.InvalidInputException as error: + assert "'numpy' is required" in str(error), repr(error) + else: + raise AssertionError("registering a UDF without numpy did not raise") + """, + blocked, + ) + + +def test_spark_udf_with_a_timestamp_needs_pytz(tmp_path): + """The binding's Python UDF path needs pytz to pass a timestamp to a function.""" + if importlib.util.find_spec("numpy") is None: + pytest.skip("a UDF needs numpy") + run_blocked( + tmp_path, + """ + import datetime + + import duckdb + from duckdb.experimental.spark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + plus_one = spark.udf.register("plus_one", lambda n: n + 1, "int") + assert spark.createDataFrame([(1,)], "n int").select(plus_one("n")).collect()[0][0] == 2 + year = spark.udf.register("year_of", lambda t: t.year, "int") + df = spark.createDataFrame([(datetime.datetime(2024, 1, 2),)], "t timestamp") + try: + df.select(year("t")).collect() + except duckdb.InvalidInputException as error: + assert "'pytz' failed to import" in str(error), repr(error) + else: + raise AssertionError("a UDF of a timestamp without pytz did not raise") + """, + ("pytz",), + ) diff --git a/tests/fast/test_parameter_list.py b/tests/fast/test_parameter_list.py index 833c09129..972f6da3b 100644 --- a/tests/fast/test_parameter_list.py +++ b/tests/fast/test_parameter_list.py @@ -29,6 +29,43 @@ def test_explicit_nan_param(self): res = con.execute("select isnan(cast(? as double))", (float("nan"),)) assert res.fetchone()[0] + @pytest.mark.parametrize( + ("param", "query"), + [ + ([float("nan")], "select isnan(?[1])"), + ({"a": float("nan")}, "select isnan(?.a)"), + ({"key": ["k"], "value": [float("nan")]}, "select isnan(?['k'])"), + ([[1.0, float("nan")]], "select isnan(?[1][2])"), + ([{"a": [float("nan")]}], "select isnan(?[1].a[1])"), + ], + ) + def test_nested_nan_param(self, param, query): + con = duckdb.connect() + assert con.execute(query, [param]).fetchone() == (True,) + + @pytest.mark.parametrize( + ("value", "query"), + [ + (duckdb.Value(float("nan"), duckdb.sqltype("DOUBLE")), "select isnan(?)"), + (duckdb.Value({"k": float("nan")}, duckdb.sqltype("MAP(VARCHAR, DOUBLE)")), "select isnan(?['k'])"), + (duckdb.Value({"a": float("nan")}, duckdb.sqltype("STRUCT(a DOUBLE)")), "select isnan(?.a)"), + (duckdb.Value([float("nan"), 1.0], duckdb.sqltype("DOUBLE[2]")), "select isnan(?[1])"), + (duckdb.Value({float("nan"): 1}, duckdb.sqltype("MAP(DOUBLE, INTEGER)")), "select isnan(map_keys(?)[1])"), + ( + duckdb.Value((float("nan"), [float("nan")]), duckdb.sqltype("STRUCT(a DOUBLE, b DOUBLE[])")), + "select isnan(?.b[1])", + ), + ], + ) + def test_nan_typed_value_param(self, value, query): + con = duckdb.connect() + assert con.execute(query, [value]).fetchone() == (True,) + + def test_nested_nan_typed_param(self): + con = duckdb.connect() + value = duckdb.Value((float("nan"), 1.0), duckdb.sqltype("STRUCT(a DOUBLE, b DOUBLE)")) + assert con.execute("select isnan(?.a)", [value]).fetchone() == (True,) + def test_string_parameter(self, duckdb_cursor): conn = duckdb.connect() conn.execute("create table orders (o_orderdate date)") diff --git a/tests/fast/test_type.py b/tests/fast/test_type.py index 17cec9e63..9d2b5b8cb 100644 --- a/tests/fast/test_type.py +++ b/tests/fast/test_type.py @@ -85,6 +85,13 @@ def test_struct_type(self): type = duckdb.struct_type([BIGINT, BOOLEAN]) assert str(type) == "STRUCT(v1 BIGINT, v2 BOOLEAN)" + @pytest.mark.parametrize("fields", [{}, []]) + def test_empty_struct_type(self, fields): + type = duckdb.struct_type(fields) + assert type == duckdb.sql("select {} as s").types[0] + assert duckdb.execute("select ?", [duckdb.Value({}, type)]).fetchall() == [({},)] + assert duckdb.sql(f"select {{}}::{type}").fetchall() == [({},)] + def test_incomplete_struct_type(self): with pytest.raises( duckdb.InvalidInputException, match="Could not convert empty dictionary to a duckdb STRUCT type" diff --git a/tests/spark_namespace/errors.py b/tests/spark_namespace/errors.py index 74d10e3fe..334a28099 100644 --- a/tests/spark_namespace/errors.py +++ b/tests/spark_namespace/errors.py @@ -4,3 +4,8 @@ from pyspark.errors import * else: from duckdb.experimental.spark.errors import * + + +def error_condition(error: Exception) -> str | None: + """The error condition of a Spark exception; Spark 4 deprecates getErrorClass, and Spark 3.5 has only it.""" + return error.getCondition() if hasattr(error, "getCondition") else error.getErrorClass()