mbutrovich commented on code in PR #54:
URL: https://github.com/apache/datafusion-spark/pull/54#discussion_r4565088611
##########
crates/spark/scripts/pyspark_slt_runner.py:
##########
@@ -0,0 +1,1587 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under
one or more contributor license agreements
+
+"""Run the datafusion-spark sqllogictest files against Apache Spark 4.2.
+
+Each record in crates/sqllogictest/test_files/**/*.slt is run on a
+local PySpark session, and Spark's result is compared with the expected result
+in the file. The expected results come from DataFusion, so every difference
+reported here is a query where DataFusion and Spark disagree, or where the
+test expects the wrong value.
+
+This is a manual tool and does not run in CI. See README.md for setup.
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import math
+import os
+import re
+import struct
+import sys
+import tempfile
+import time
+import traceback
+from collections import Counter
+from dataclasses import dataclass, field
+from decimal import ROUND_HALF_UP, Decimal, localcontext
+from pathlib import Path
+from typing import Optional
+
+SPARK_VERSION = "4.2"
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+REPO_ROOT = SCRIPT_DIR.parents[2]
Review Comment:
`crates/spark/Cargo.toml` has no `include` or `exclude`, so `scripts/`
(about 2,300 lines) would ship in the `datafusion-spark` crates.io package. In
the packaged crate, `SCRIPT_DIR.parents[2]` is not the repository root, so
`DEFAULT_TEST_DIR` and the paths in reports would be wrong. Could we move the
runner out of the published crate? `crates/sqllogictest` has `publish = false`
and holds the `.slt` files the runner reads, so `crates/sqllogictest/scripts/`
looks like a natural home.
##########
crates/spark/scripts/README.md:
##########
@@ -0,0 +1,140 @@
+<!--
+SPDX-License-Identifier: Apache-2.0
+SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under
one or more contributor license agreements
+-->
+
+# Running the Spark sqllogictests on Apache Spark
+
+`pyspark_slt_runner.py` runs the `.slt` files in
+[`crates/sqllogictest/test_files`](../../sqllogictest/test_files)
+on Apache Spark 4.2 and reports every query where Spark's result differs from
+the expected result in the file.
+
+The expected results in those files come from DataFusion, and nothing else
+checks them against Spark. A difference means that DataFusion and Spark
+disagree, or that the test expects the wrong value. This is a manual tool. It
+does not run in CI.
+
+## Setup
+
+You need Python 3.10 or later and Java 17 or 21. Install PySpark 4.2.0 into a
+virtual environment:
+
+```shell
+python3 -m venv /tmp/pyspark-4.2
+/tmp/pyspark-4.2/bin/pip install -r crates/spark/scripts/requirements.txt
+```
+
+The runner ignores `SPARK_HOME`, so it always uses the Spark bundled with this
+PySpark.
+
+## Usage
+
+From the repository root:
+
+```shell
+# All files
+/tmp/pyspark-4.2/bin/python crates/spark/scripts/pyspark_slt_runner.py
+
+# One file or directory
+/tmp/pyspark-4.2/bin/python crates/spark/scripts/pyspark_slt_runner.py \
+ crates/sqllogictest/test_files/math/abs.slt
+
+# Also list skipped and passing records
+/tmp/pyspark-4.2/bin/python crates/spark/scripts/pyspark_slt_runner.py
--verbose
+```
+
+Each difference is printed with its `file:line`, the query, the expected result
+and Spark's result. A summary of every outcome and the most common reasons for
+skipping comes at the end. The runner exits with status 1 if it found any
+differences, and with status 2 if it stopped early because Spark failed.
+
+| Outcome | Meaning
|
+| --------------------------------------------------- |
------------------------------------------------------------------------ |
+| results differ | Spark returns
different rows |
+| Spark raises an error | The file expects a
result, but Spark raises an error |
+| DataFusion expects an error, Spark returns a result | The file expects an
error, but Spark does not raise one |
+| results match, column types differ | The rows match, but a
column's type letter (`I`, `R`, `T`, ...) does not |
+| same numbers, printed with a different scale | For example `7` and
`7.0`, usually a `DOUBLE` against a `DECIMAL` |
+| results differ only in floating-point digits | The values differ by a
relative error below 1e-12 |
+| same rows in a different order | Usually a query with
no `ORDER BY` |
+| the runner failed | A bug in this runner,
not Spark behavior |
+
+Spark reads a literal such as `1.5` as a `DECIMAL` and DataFusion reads it as a
+`DOUBLE`, which explains many of the scale differences.
+
+The runner collects at most 1000 rows per query, because some tests expect an
+error from a query that Spark answers with billions of rows. It skips a
+`stack()` of more than 1000 rows, such as `stack(2000000000, 1)`, because Spark
+builds all of those rows at once and runs out of memory. If Spark fails anyway,
+the runner skips the rest of that file and starts a new Spark session for the
+next one.
+
+## How queries are run
+
+Each file starts with ANSI mode off and the session time zone set to UTC, which
+are DataFusion's defaults. These settings in the files are passed on to Spark:
+
+- `datafusion.execution.enable_ansi_mode` sets `spark.sql.ansi.enabled`
+- `datafusion.execution.time_zone` sets `spark.sql.session.timeZone`
+- `datafusion.spark.map_key_dedup_policy` sets `spark.sql.mapKeyDedupPolicy`
+
+Spark 4.2 parses most of the SQL as written, including `::` casts. The runner
+makes only these changes:
+
+- `CREATE TABLE` becomes `CREATE TEMPORARY VIEW`, and `DROP TABLE` becomes
+ `DROP VIEW`
+- PostgreSQL type names that Spark does not accept are renamed, for example
+ `FLOAT8` to `DOUBLE` and `TEXT` to `STRING`
+- `column1`, `column2`, ... become `col1`, `col2`, ..., which is how Spark
names
+ `VALUES` columns
+- `spark_cast(x, 'type')`, this crate's emulation of Spark's `CAST`, becomes
+ `CAST(x AS type)`
+- `arrow_cast(x, 'Type')` becomes `CAST(x AS type)` when the Arrow type has the
+ same values as a Spark type. Types that differ only in their physical layout
+ map to one Spark type: `Utf8`, `LargeUtf8` and `Utf8View` are all `STRING`,
+ `LargeList(Int64)` and `FixedSizeList(3, Int64)` are `ARRAY<BIGINT>`, and a
+ `Dictionary` is its value type. Arrow reads a number cast to a date or
+ timestamp as a count of the type's unit, so those casts are translated only
+ for literals: `arrow_cast(1703512245, 'Timestamp(s)')` becomes
+ `timestamp_seconds(1703512245)`. In a UTC session, a timestamp without a
+ time zone becomes Spark's `TIMESTAMP`, which holds the same values and is the
+ type functions such as `unix_micros` accept. A date cast to an integer
becomes
+ `unix_date(x)`, because Arrow counts days since the epoch and Spark's `CAST`
+ returns NULL.
+
+These records are skipped and counted, with the reason, in the summary:
+
+- queries Spark cannot parse, such as `[1, 2]` array literals
+- queries that call a function Spark does not have, such as `arrow_typeof`
+- `arrow_cast` to a type Spark does not have, such as an unsigned integer,
+ `Float16`, or a nanosecond timestamp, and `arrow_cast` of a non-literal to a
+ date, time or timestamp
+- queries whose string literals contain a backslash. Spark reads the backslash
+ as an escape character and DataFusion does not, so the two engines would see
+ different strings.
+- `DECIMAL` types with a precision above 38
+- queries that use a table whose `CREATE TABLE` Spark could not run
+- `EXPLAIN` and `DESCRIBE`, whose output differs between the engines
+- results PySpark cannot return to Python, such as a map with array keys, and
+ interval values, which this runner does not format
+- records marked `skipif spark`, or `onlyif` another engine
+
+## Values
+
+Spark's results are printed the way DataFusion's sqllogictest engine prints
+them, so they can be compared as text. For example, a `DOUBLE` keeps 15 decimal
+places, a `DECIMAL` keeps its scale and an empty string is `(empty)`. A Spark
+`TIMESTAMP` is printed in the session time zone. An expected timestamp with an
+offset, such as `2024-07-15T00:00:00Z`, is converted to that time zone before
+the comparison.
+
+## Tests
+
+The runner's parsing, translation and formatting have unit tests that need
+neither Spark nor Java:
+
+```shell
+python3 -m unittest discover -s crates/spark/scripts
Review Comment:
The unit tests take about 20 ms and need neither Spark nor Java (I ran them:
44 tests, OK), but nothing runs them. The runner depends on `.slt` syntax and
on `arrow_cast` spellings that other PRs keep changing, so the translation will
rot without a check. Could we add a small job to `.github/workflows` that runs
`python3 -m unittest discover -s crates/spark/scripts`? Running Spark itself
can stay manual.
##########
crates/spark/scripts/pyspark_slt_runner.py:
##########
@@ -0,0 +1,1587 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under
one or more contributor license agreements
+
+"""Run the datafusion-spark sqllogictest files against Apache Spark 4.2.
+
+Each record in crates/sqllogictest/test_files/**/*.slt is run on a
+local PySpark session, and Spark's result is compared with the expected result
+in the file. The expected results come from DataFusion, so every difference
+reported here is a query where DataFusion and Spark disagree, or where the
+test expects the wrong value.
+
+This is a manual tool and does not run in CI. See README.md for setup.
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import math
+import os
+import re
+import struct
+import sys
+import tempfile
+import time
+import traceback
+from collections import Counter
+from dataclasses import dataclass, field
+from decimal import ROUND_HALF_UP, Decimal, localcontext
+from pathlib import Path
+from typing import Optional
+
+SPARK_VERSION = "4.2"
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+REPO_ROOT = SCRIPT_DIR.parents[2]
+DEFAULT_TEST_DIR = REPO_ROOT / "crates" / "sqllogictest" / "test_files"
+
+# The label that `skipif` and `onlyif` directives use for this runner.
+LABEL = "spark"
+
+# DataFusion runs these files with ANSI mode off and no time zone (UTC), so
+# every file starts from the same Spark settings.
+SPARK_DEFAULTS = {
+ "spark.sql.ansi.enabled": "false",
+ "spark.sql.session.timeZone": "UTC",
+ "spark.sql.mapKeyDedupPolicy": "EXCEPTION",
+ # Spark 4.2 rejects the TIME type unless this internal setting is on.
+ "spark.sql.timeType.enabled": "true",
+}
+
+# DataFusion settings that have a Spark equivalent.
+SETTINGS = {
+ "datafusion.execution.enable_ansi_mode": "spark.sql.ansi.enabled",
+ "datafusion.execution.time_zone": "spark.sql.session.timeZone",
+ "datafusion.spark.map_key_dedup_policy": "spark.sql.mapKeyDedupPolicy",
+}
+
+# Type names DataFusion accepts and Spark does not.
+TYPE_ALIASES = {
+ "BYTEA": "BINARY",
+ "FLOAT4": "FLOAT",
+ "FLOAT8": "DOUBLE",
+ "INT2": "SMALLINT",
+ "INT4": "INT",
+ "INT8": "BIGINT",
+ "TEXT": "STRING",
+}
+
+# Spark's maximum decimal precision.
+MAX_DECIMAL_PRECISION = 38
+
+# The most rows collected from Spark for one query. Some tests expect an error
+# for a query that Spark answers with billions of rows, such as
stack(2000000000, 1).
+MAX_ROWS = 1000
+
+# The most result lines printed for one difference.
+MAX_REPORT_LINES = 20
+
+
+class Skip(Exception):
+ """The record cannot be run on Spark. The message says why."""
+
+
+class Unsupported(Exception):
+ """Spark returned a value of a type this runner does not format."""
+
+
+# ---------------------------------------------------------------------------
+# Parsing .slt files
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class Record:
+ """A `statement` or `query` record from an .slt file."""
+
+ path: Path
+ line: int
+ kind: str # "statement" or "query"
+ sql: str = ""
+ expect_error: bool = False
+ types: str = ""
+ sort_mode: str = "nosort"
+ expected: list[str] = field(default_factory=list)
+ skip_reason: Optional[str] = None
+
+
+def parse_slt(path: Path) -> list[Record]:
+ """Parse the `statement` and `query` records of an .slt file."""
+ lines = path.read_text().split("\n")
+ records: list[Record] = []
+ conditions: list[tuple[str, str]] = []
+ i = 0
+ while i < len(lines):
+ words = lines[i].split()
+ if not words or words[0].startswith("#"):
+ i += 1
+ continue
+ if words[0] in ("skipif", "onlyif") and len(words) > 1:
+ conditions.append((words[0], words[1]))
+ i += 1
+ continue
+ if words[0] not in ("statement", "query"):
+ # Directives such as `include` and `halt`, which these files do
not use.
+ conditions = []
+ i += 1
+ continue
+
+ record = Record(path=path, line=i + 1, kind=words[0])
+ args = words[1:]
+ if args and args[0] == "error":
+ record.expect_error = True
+ elif record.kind == "query" and args:
+ record.types = args[0]
+ if len(args) > 1 and args[1] in ("nosort", "rowsort", "valuesort"):
+ record.sort_mode = args[1]
+ for directive, label in conditions:
+ if directive == "skipif" and label == LABEL:
+ record.skip_reason = f"skipif {label}"
+ elif directive == "onlyif" and label != LABEL:
+ record.skip_reason = f"onlyif {label}"
+ conditions = []
+
+ # The SQL runs until a blank line or the `----` separator.
+ i += 1
+ sql = []
+ while i < len(lines) and lines[i].strip() and lines[i] != "----":
+ if not lines[i].lstrip().startswith("#"):
+ sql.append(lines[i])
+ i += 1
+ record.sql = "\n".join(sql)
+
+ # Expected results (or an expected error message) run until an empty
+ # line. A result line may start with `#`, so it is not a comment here,
+ # and it may be a single space, the engine's empty first row before a
+ # multi-line value.
+ if i < len(lines) and lines[i] == "----":
+ i += 1
+ while i < len(lines) and lines[i] != "":
+ record.expected.append(lines[i])
+ i += 1
+ records.append(record)
+ return records
+
+
+# ---------------------------------------------------------------------------
+# Translating DataFusion SQL to Spark SQL
+#
+# Spark 4.2 parses most of these files as written, including `::` casts.
+# Only a few renames are applied here, and `arrow_cast` and `spark_cast`
+# become a CAST (see below). Anything else Spark cannot parse, and any function
+# Spark does not have (such as `arrow_typeof`), is skipped.
+# ---------------------------------------------------------------------------
+
+_TOKEN_RE = re.compile(
+ r"""
+ (?P<string>'(?:[^']|'')*')
+ | (?P<quoted>`[^`]*`|"[^"]*")
+ | (?P<comment>--[^\n]*)
+ | (?P<word>[A-Za-z_][A-Za-z0-9_]*)
+ | (?P<number>\d+(?:\.\d*)?(?:[eE][+-]?\d+)?|\.\d+(?:[eE][+-]?\d+)?)
+ | (?P<cast>::)
+ | (?P<space>\s+)
+ | (?P<other>.)
+ """,
+ re.VERBOSE | re.DOTALL,
+)
+
+_COLUMN_RE = re.compile(r"column(\d+)", re.IGNORECASE)
+
+
+def _tokenize(sql: str) -> list[tuple[str, str]]:
+ return [(m.lastgroup, m.group()) for m in _TOKEN_RE.finditer(sql)]
+
+
+def _next_significant(tokens: list[tuple[str, str]], start: int) -> int:
+ """Return the index of the next token that is not whitespace or a
comment."""
+ i = start
+ while i < len(tokens) and tokens[i][0] in ("space", "comment"):
+ i += 1
+ return i
+
+
+def _translate_type(tokens: list[tuple[str, str]], i: int) -> str:
+ """Translate the type name at tokens[i]."""
+ text = tokens[i][1]
+ upper = text.upper()
+ if upper in TYPE_ALIASES:
+ return TYPE_ALIASES[upper]
+ j = _next_significant(tokens, i + 1)
+ has_args = j < len(tokens) and tokens[j][1] == "("
+ if upper == "VARCHAR" and not has_args:
+ return "STRING"
+ if upper in ("DECIMAL", "DEC", "NUMERIC") and has_args:
+ k = _next_significant(tokens, j + 1)
+ if k < len(tokens) and tokens[k][0] == "number":
+ if int(tokens[k][1]) > MAX_DECIMAL_PRECISION:
+ raise Skip(f"DECIMAL precision above {MAX_DECIMAL_PRECISION}")
+ return text
+
+
+# ---------------------------------------------------------------------------
+# arrow_cast
+#
+# `arrow_cast(expr, 'Type')` casts to an Arrow type. Many Arrow types differ
+# from a Spark type only in their physical layout: Utf8, LargeUtf8 and Utf8View
+# are all Spark's STRING, and a dictionary is an encoding of its value type. So
+# the call becomes a CAST to that Spark type. A type Spark does not have, such
+# as an unsigned integer, skips the record.
+#
+# Arrow interprets a number cast to a date, time or timestamp as a count of the
+# type's unit (days, seconds, nanoseconds, ...), and Spark does not, so those
+# casts are only translated for literal arguments.
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class ArrowType:
+ """A parsed Arrow type, such as `Decimal128(10, 2)` or `List(Int32)`."""
+
+ name: str
+ args: tuple = ()
+
+
+@dataclass(frozen=True)
+class ArrowField:
+ """A named field in a `Struct` or `Map` type, such as `"a": Int64`."""
+
+ name: str
+ type: object
+
+
+_ARROW_TOKEN_RE = re.compile(
+ r"""\s*(?:
+ (?P<string>"(?:[^"\\]|\\.)*")
+ | (?P<number>-?\d+)
+ | (?P<word>[A-Za-zµ_][A-Za-z0-9µ_-]*)
+ | (?P<punct>[(),:])
+ )""",
+ re.VERBOSE,
+)
+
+
+def parse_arrow_type(text: str):
+ """Parse an Arrow type as DataFusion's arrow_cast spells it."""
+ tokens = []
+ pos, text = 0, text.strip()
+ while pos < len(text):
+ m = _ARROW_TOKEN_RE.match(text, pos)
+ if not m:
+ raise ValueError(text)
+ tokens.append((m.lastgroup, m.group(m.lastgroup)))
+ pos = m.end()
+ i = 0
+
+ def peek() -> Optional[str]:
+ return tokens[i][1] if i < len(tokens) else None
+
+ def expect(value: str) -> None:
+ nonlocal i
+ if peek() != value:
+ raise ValueError(text)
+ i += 1
+
+ def value():
+ nonlocal i
+ if i >= len(tokens):
+ raise ValueError(text)
+ kind, token = tokens[i]
+ i += 1
+ if kind == "number":
+ return int(token)
+ if kind == "string":
+ name = token[1:-1]
+ if peek() == ":":
+ i += 1
+ return ArrowField(name, value())
+ return name
+ if kind == "word":
+ # Nullability does not change the Spark type.
+ if token in ("non-null", "nullable"):
+ return value()
+ args = []
+ if peek() == "(":
+ i += 1
+ while peek() != ")":
+ args.append(value())
+ if peek() == ",":
+ i += 1
+ expect(")")
+ return ArrowType(token, tuple(args))
+ raise ValueError(text)
+
+ parsed = value()
+ if i != len(tokens) or not isinstance(parsed, ArrowType):
+ raise ValueError(text)
+ return parsed
+
+
+_ARROW_SIMPLE_TYPES = {
+ "Boolean": "BOOLEAN",
+ "Int8": "TINYINT",
+ "Int16": "SMALLINT",
+ "Int32": "INT",
+ "Int64": "BIGINT",
+ "Float32": "FLOAT",
+ "Float64": "DOUBLE",
+ "Utf8": "STRING",
+ "LargeUtf8": "STRING",
+ "Utf8View": "STRING",
+ "Binary": "BINARY",
+ "LargeBinary": "BINARY",
+ "BinaryView": "BINARY",
+ "FixedSizeBinary": "BINARY",
+ "Date32": "DATE",
+ "Date64": "DATE",
+}
+
+_ARROW_LIST_TYPES = ("List", "LargeList", "ListView", "LargeListView",
"FixedSizeList")
+
+_ARROW_DECIMAL_TYPES = ("Decimal32", "Decimal64", "Decimal128", "Decimal256")
+
+_ARROW_UNITS = {
+ "Second": "s",
+ "s": "s",
+ "Millisecond": "ms",
+ "ms": "ms",
+ "Microsecond": "us",
+ "µs": "us",
+ "us": "us",
+ "Nanosecond": "ns",
+ "ns": "ns",
+}
+
+# Spark's TIME(p) precision for each Arrow time unit. Spark has no nanoseconds.
+_TIME_PRECISION = {"s": 0, "ms": 3, "us": 6}
+
+_UTC_NAMES = {"UTC", "Etc/UTC", "Z", "+00:00", "+00", "+0000", "GMT"}
+
+
+def _arrow_unit(t: ArrowType) -> str:
+ if not t.args or not isinstance(t.args[0], ArrowType) or t.args[0].name
not in _ARROW_UNITS:
+ raise Skip("arrow_cast to a type this runner cannot parse")
+ return _ARROW_UNITS[t.args[0].name]
+
+
+def _arrow_time_zone(t: ArrowType) -> Optional[str]:
+ """The time zone of an Arrow Timestamp: `None`, `Some("UTC")` or
`"UTC"`."""
+ if len(t.args) < 2:
+ return None
+ tz = t.args[1]
+ if isinstance(tz, str):
+ return tz
+ if isinstance(tz, ArrowType) and tz.name == "Some" and len(tz.args) == 1:
+ return tz.args[0]
+ if isinstance(tz, ArrowType) and tz.name == "None":
+ return None
+ raise Skip("arrow_cast to a type this runner cannot parse")
+
+
+def _same_time_zone(a: str, b: str) -> bool:
+ return a == b or (a in _UTC_NAMES and b in _UTC_NAMES)
+
+
+def spark_type_for_arrow(t: ArrowType) -> str:
+ """The Spark type with the same values as an Arrow type, or raise Skip."""
+ name = t.name
+ if name in _ARROW_SIMPLE_TYPES:
+ return _ARROW_SIMPLE_TYPES[name]
+ if name in ("UInt8", "UInt16", "UInt32", "UInt64"):
+ raise Skip("arrow_cast to an unsigned integer, which Spark does not
have")
+ if name in _ARROW_DECIMAL_TYPES:
+ precision, scale = t.args
+ if precision > MAX_DECIMAL_PRECISION:
+ raise Skip(f"DECIMAL precision above {MAX_DECIMAL_PRECISION}")
+ if scale < 0:
+ raise Skip("arrow_cast to a DECIMAL with a negative scale")
+ return f"DECIMAL({precision}, {scale})"
+ if name in _ARROW_LIST_TYPES:
+ return f"ARRAY<{spark_type_for_arrow(t.args[-1])}>"
+ if name == "Dictionary":
+ return spark_type_for_arrow(t.args[1])
+ if name == "Struct":
+ fields = [f"`{f.name}`: {spark_type_for_arrow(f.type)}" for f in
t.args]
+ return f"STRUCT<{', '.join(fields)}>"
+ if name == "Map":
+ entries = t.args[0]
+ key, value = entries.type.args
+ return f"MAP<{spark_type_for_arrow(key.type)},
{spark_type_for_arrow(value.type)}>"
+ if name in ("Time32", "Time64"):
+ unit = _arrow_unit(t)
+ if unit not in _TIME_PRECISION:
+ raise Skip("arrow_cast to nanoseconds, which Spark does not have")
+ return f"TIME({_TIME_PRECISION[unit]})"
+ if name == "Timestamp":
+ unit = _arrow_unit(t)
+ if unit != "us" or _arrow_time_zone(t) is not None:
+ raise Skip("arrow_cast to a nested timestamp that is not
TIMESTAMP_NTZ")
+ return "TIMESTAMP_NTZ"
+ raise Skip(f"arrow_cast to {name}, which Spark does not have")
+
+
+def _strip_parens(tokens: list[tuple[str, str]]) -> list[tuple[str, str]]:
+ """Drop whitespace, comments and parentheses that wrap the whole
argument."""
+ tokens = [t for t in tokens if t[0] not in ("space", "comment")]
+ while len(tokens) >= 2 and tokens[0][1] == "(" and tokens[-1][1] == ")":
+ depth = 0
+ for k, (_, text) in enumerate(tokens):
+ depth += {"(": 1, ")": -1}.get(text, 0)
+ if depth == 0 and k < len(tokens) - 1:
+ return tokens
+ tokens = tokens[1:-1]
+ return tokens
+
+
+def _literal(tokens: list[tuple[str, str]]):
+ """Return ("null", None), ("string", text) or ("integer", n) for a literal
argument."""
+ tokens = _strip_parens(tokens)
+ if len(tokens) == 1 and tokens[0][0] == "word" and tokens[0][1].upper() ==
"NULL":
+ return ("null", None)
+ if len(tokens) == 1 and tokens[0][0] == "string":
+ return ("string", tokens[0][1][1:-1].replace("''", "'"))
+ # A typed literal, such as DATE '2024-01-01' or TIMESTAMP '2024-01-01
10:00:00'
+ if (
+ len(tokens) == 2
+ and tokens[0][0] == "word"
+ and tokens[0][1].upper() in ("DATE", "TIMESTAMP", "TIMESTAMP_NTZ")
+ and tokens[1][0] == "string"
+ ):
+ return ("typed", tokens[1][1][1:-1].replace("''", "'"))
+ sign = 1
+ if len(tokens) == 2 and tokens[0][1] in ("-", "+"):
+ sign = -1 if tokens[0][1] == "-" else 1
+ tokens = tokens[1:]
+ if len(tokens) == 1 and tokens[0][0] == "number" and
tokens[0][1].isdigit():
+ return ("integer", sign * int(tokens[0][1]))
+ return (None, None)
+
+
+# Fractional-second digits each timestamp unit keeps.
+_UNIT_DIGITS = {"s": 0, "ms": 3, "us": 6, "ns": 9}
+
+_FRACTION_RE = re.compile(r":\d{2}\.(\d+)")
+
+_EPOCH_FUNCTIONS = {"s": "timestamp_seconds", "ms": "timestamp_millis", "us":
"timestamp_micros"}
+
+_MILLIS_PER_DAY = 86_400_000
+
+
+def _temporal_cast(t: ArrowType, arg: list[tuple[str, str]], arg_sql: str,
time_zone: str) -> str:
+ """Translate arrow_cast to a date, time or timestamp, which depends on the
argument."""
+ kind, value = _literal(arg)
+ if kind == "typed":
+ # DataFusion's TIMESTAMP literal has no time zone and Spark's is in the
+ # session time zone, so they are the same value only in UTC.
+ if time_zone not in _UTC_NAMES:
+ raise Skip(f"arrow_cast of a typed literal to {t.name} in a
session time zone other than UTC")
+ kind = "string"
+ arg_sql = "'" + value.replace("'", "''") + "'"
+ if kind is None:
+ raise Skip(f"arrow_cast of an expression to {t.name}, which Arrow and
Spark convert differently")
+ if t.name in ("Date32", "Date64"):
+ if kind == "integer":
+ days = value
+ if t.name == "Date64":
+ if value % _MILLIS_PER_DAY:
+ raise Skip("arrow_cast to Date64 of milliseconds that are
not whole days")
+ days = value // _MILLIS_PER_DAY
+ return f"date_from_unix_date({days})"
+ return f"CAST({arg_sql} AS DATE)"
+ if t.name in ("Time32", "Time64"):
+ spark_type = spark_type_for_arrow(t)
+ if kind == "integer":
+ raise Skip(f"arrow_cast of a number to {t.name}, which Spark
cannot cast")
+ return f"CAST({arg_sql} AS {spark_type})"
+
+ # Timestamp
+ unit = _arrow_unit(t)
+ tz = _arrow_time_zone(t)
+ if tz is not None and not _same_time_zone(tz, time_zone):
+ raise Skip("arrow_cast to a timestamp in a time zone other than the
session's")
+ # In a UTC session, Spark's TIMESTAMP holds the same values as an Arrow
+ # timestamp without a time zone, and functions such as unix_micros accept
+ # only TIMESTAMP, not TIMESTAMP_NTZ.
+ utc_session = time_zone in _UTC_NAMES
+ spark_type = "TIMESTAMP" if tz is not None or utc_session else
"TIMESTAMP_NTZ"
+ if kind == "null":
+ return f"CAST(NULL AS {spark_type})"
+ if kind == "string":
+ m = _FRACTION_RE.search(value)
+ digits = len(m.group(1).rstrip("0")) if m else 0
+ if digits > 6:
+ raise Skip("arrow_cast to nanoseconds, which Spark does not have")
+ if digits > _UNIT_DIGITS[unit]:
+ raise Skip("arrow_cast that truncates a timestamp to seconds or
milliseconds")
+ return f"CAST({arg_sql} AS {spark_type})"
+ # An integer counts units since the epoch.
+ if unit == "ns":
+ if value % 1000:
+ raise Skip("arrow_cast to nanoseconds, which Spark does not have")
+ value, unit = value // 1000, "us"
+ if tz is None and not utc_session:
+ # Spark would print the instant in the session time zone, and Arrow
+ # prints a timestamp without a time zone as UTC.
+ raise Skip("arrow_cast of a number to a timestamp in a session time
zone other than UTC")
+ return f"{_EPOCH_FUNCTIONS[unit]}({value})"
+
+
+# Brackets that nest, so a comma inside `[1, 2]` or `f(a, b)` does not split
arguments.
+_NESTING = {"(": 1, "[": 1, "{": 1, ")": -1, "]": -1, "}": -1}
+
+
+# Spark functions that return a DATE.
+_DATE_FUNCTIONS = {
+ "add_months",
+ "current_date",
+ "curdate",
+ "date_add",
+ "date_from_unix_date",
+ "date_sub",
+ "dateadd",
+ "last_day",
+ "make_date",
+ "next_day",
+ "to_date",
+ "trunc",
+}
+
+_INTEGER_TYPES = ("TINYINT", "SMALLINT", "INT", "BIGINT")
+
+
+def _is_date(sql: str) -> bool:
+ """Whether a translated expression is visibly a DATE: a DATE literal, a
cast
+ to DATE, or a call to a function that returns one."""
+ tokens = _strip_parens(_tokenize(sql))
+ if len(tokens) == 2 and tokens[0][1].upper() == "DATE" and tokens[1][0] ==
"string":
+ return True
+ if len(tokens) < 3 or tokens[0][0] != "word" or tokens[1][1] != "(":
+ return False
+ if _matching_paren(tokens, 1) != len(tokens) - 1:
+ return False
+ name = tokens[0][1].lower()
+ if name == "cast":
+ inner = tokens[2:-1]
+ return len(inner) >= 2 and inner[-1][1].upper() == "DATE" and
inner[-2][1].upper() == "AS"
+ return name in _DATE_FUNCTIONS
+
+
+def _matching_paren(tokens: list[tuple[str, str]], open_index: int) -> int:
+ depth = 0
+ for k in range(open_index, len(tokens)):
+ if tokens[k][0] == "other":
+ depth += _NESTING.get(tokens[k][1], 0)
+ if depth == 0:
+ return k
+ raise Skip("unbalanced parentheses")
+
+
+# A Spark type name, such as `timestamp`, `decimal(10, 2)` or `array<int>`.
+_SPARK_TYPE_RE = re.compile(r"[A-Za-z][A-Za-z0-9_<>(),:`\s]*")
+
+
+def _call_args(tokens: list[tuple[str, str]], open_index: int, close: int) ->
list:
+ """Split the tokens between a call's parentheses on top-level commas."""
+ args, start, depth = [], open_index + 1, 0
+ for k in range(open_index + 1, close):
+ if tokens[k][0] == "other":
+ if tokens[k][1] == "," and depth == 0:
+ args.append(tokens[start:k])
+ start = k + 1
+ depth += _NESTING.get(tokens[k][1], 0)
+ args.append(tokens[start:close])
+ return args
+
+
+def _type_string(name: str, args: list) -> str:
+ """The type argument of arrow_cast or spark_cast, which must be a string
literal."""
+ if len(args) != 2:
+ raise Skip(f"{name} without exactly two arguments")
+ type_arg = _strip_parens(args[1])
+ if len(type_arg) != 1 or type_arg[0][0] != "string":
+ raise Skip(f"{name} to a type that is not a string literal")
+ return type_arg[0][1][1:-1].replace("''", "'")
+
+
+def _arrow_cast(arg: list, arg_sql: str, type_string: str, time_zone: str) ->
str:
+ try:
+ arrow_type = parse_arrow_type(type_string)
+ except (ValueError, TypeError, AttributeError):
+ raise Skip("arrow_cast to a type this runner cannot parse")
+ if arrow_type.name in ("Date32", "Date64", "Time32", "Time64",
"Timestamp"):
+ return _temporal_cast(arrow_type, arg, arg_sql, time_zone)
+ try:
+ spark_type = spark_type_for_arrow(arrow_type)
+ except (ValueError, TypeError, AttributeError, IndexError):
+ raise Skip("arrow_cast to a type this runner cannot parse")
+ if spark_type in _INTEGER_TYPES and _is_date(arg_sql):
+ # Arrow casts a date to an integer as days since the epoch.
+ # Spark's CAST returns NULL, and unix_date returns the days.
+ days = f"unix_date({arg_sql})"
+ return days if spark_type == "INT" else f"CAST({days} AS {spark_type})"
+ return f"CAST({arg_sql} AS {spark_type})"
+
+
+def rewrite_casts(sql: str, time_zone: str = "UTC") -> str:
+ """Replace each `arrow_cast(expr, 'Type')` and `spark_cast(expr, 'type')`
+ with a Spark CAST, or raise Skip.
+
+ `spark_cast` is this crate's emulation of Spark's CAST, and its second
+ argument is already a Spark type name.
+ """
+ tokens = _tokenize(sql)
+ out = []
+ i = 0
+ while i < len(tokens):
+ kind, text = tokens[i]
+ name = text.lower()
+ open_index = _next_significant(tokens, i + 1)
+ if not (
+ kind == "word"
+ and name in ("arrow_cast", "spark_cast")
+ and open_index < len(tokens)
+ and tokens[open_index][1] == "("
+ ):
+ out.append(text)
+ i += 1
+ continue
+ close = _matching_paren(tokens, open_index)
+ args = _call_args(tokens, open_index, close)
+ type_string = _type_string(name, args)
+ arg_sql = rewrite_casts("".join(t for _, t in args[0]),
time_zone).strip()
+ if name == "arrow_cast":
+ out.append(_arrow_cast(args[0], arg_sql, type_string, time_zone))
+ else:
+ if not _SPARK_TYPE_RE.fullmatch(type_string.strip()):
+ raise Skip("spark_cast to a type that is not a Spark type
name")
+ out.append(f"CAST({arg_sql} AS {type_string.strip()})")
+ i = close + 1
+ return "".join(out)
+
+
+def translate(sql: str, time_zone: str = "UTC") -> str:
+ """Translate the parts of DataFusion's SQL that Spark spells
differently."""
+ tokens = _tokenize(rewrite_casts(sql, time_zone))
+ out = []
+ # The function name before each open parenthesis, to find CAST(x AS type).
+ calls: list[Optional[str]] = []
+ last: Optional[tuple[str, str]] = None
+ expect_type = False
+ for i, (kind, text) in enumerate(tokens):
+ if kind == "string" and "\\" in text:
+ raise Skip(
+ "a string literal contains a backslash, which Spark reads as
an "
+ "escape and DataFusion does not"
+ )
+ if kind in ("space", "comment"):
+ out.append(text)
+ continue
+ if kind == "word":
+ if expect_type:
+ text = _translate_type(tokens, i)
+ else:
+ m = _COLUMN_RE.fullmatch(text)
+ if m:
+ # Spark names VALUES columns col1, col2, ...
+ text = f"col{m.group(1)}"
+ expect_type = text.upper() == "AS" and bool(calls) and calls[-1]
in (
+ "CAST",
+ "TRY_CAST",
+ )
+ elif kind == "cast":
+ expect_type = True
+ else:
+ expect_type = False
+ if text == "(":
+ calls.append(last[1].upper() if last and last[0] == "word"
else None)
+ elif text == ")" and calls:
+ calls.pop()
+ last = (kind, text)
+ out.append(text)
+ return "".join(out)
+
+
+_NAME = r"[A-Za-z_][A-Za-z0-9_]*"
+_CREATE_RE = re.compile(
+
rf"\s*CREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+(?:IF\s+NOT\s+EXISTS\s+)?({_NAME})\s*",
+ re.IGNORECASE,
+)
+_DROP_RE = re.compile(
+ rf"\s*DROP\s+(?:TABLE|VIEW)\s+(?:IF\s+EXISTS\s+)?({_NAME})\s*;?\s*$",
+ re.IGNORECASE,
+)
+_AS_RE = re.compile(r"AS\b\s*", re.IGNORECASE)
+
+
+def _split_top_level(text: str) -> list[str]:
+ """Split on commas that are not inside parentheses."""
+ parts, depth, start = [], 0, 0
+ for i, ch in enumerate(text):
+ if ch == "(":
+ depth += 1
+ elif ch == ")":
+ depth -= 1
+ elif ch == "," and depth == 0:
+ parts.append(text[start:i])
+ start = i + 1
+ parts.append(text[start:])
+ return [p.strip() for p in parts]
+
+
+def translate_statement(sql: str, time_zone: str = "UTC") -> str:
+ """Translate a statement, turning DataFusion tables into Spark temporary
views."""
+ m = _DROP_RE.match(sql)
+ if m:
+ return f"DROP VIEW IF EXISTS {m.group(1)}"
+ m = _CREATE_RE.match(sql)
+ if not m:
+ return translate(sql, time_zone)
+ name, rest = m.group(1), sql[m.end() :]
+ columns = None
+ if rest.startswith("("):
+ depth = 0
+ for end, ch in enumerate(rest):
+ depth += {"(": 1, ")": -1}.get(ch, 0)
+ if depth == 0:
+ break
+ columns = _split_top_level(rest[1:end])
+ rest = rest[end + 1 :].lstrip()
+ m = _AS_RE.match(rest)
+ if not m:
+ raise Skip("CREATE TABLE without AS")
+ query = rest[m.end() :].rstrip().rstrip(";")
+ if columns is None:
+ return translate(f"CREATE OR REPLACE TEMPORARY VIEW {name} AS
{query}", time_zone)
+ names, casts = [], []
+ for column in columns:
+ parts = column.split(None, 1)
+ if len(parts) != 2:
+ raise Skip(f"unsupported column definition: {column}")
+ names.append(parts[0])
+ casts.append(f"CAST({parts[0]} AS {parts[1]}) AS {parts[0]}")
+ return translate(
+ f"CREATE OR REPLACE TEMPORARY VIEW {name} AS SELECT {', '.join(casts)}
"
+ f"FROM ({query}) AS __slt_values({', '.join(names)})",
+ time_zone,
+ )
+
+
+_SET_RE = re.compile(
+ r"\s*SET\s+([A-Za-z0-9_.]+)\s*(?:=|\s+TO\s+)\s*(.*?)\s*;?\s*$",
+ re.IGNORECASE | re.DOTALL,
+)
+_RESET_RE = re.compile(r"\s*RESET\s+([A-Za-z0-9_.]+)\s*;?\s*$", re.IGNORECASE)
+
+
+def parse_setting(sql: str) -> Optional[tuple[str, Optional[str]]]:
+ """Parse `SET key = value` or `RESET key`. A reset has the value None."""
+ m = _SET_RE.match(sql)
+ if m:
+ value = m.group(2)
+ if len(value) >= 2 and value[0] == value[-1] and value[0] in "'\"":
+ value = value[1:-1]
+ return m.group(1).lower(), value
+ m = _RESET_RE.match(sql)
+ if m:
+ return m.group(1).lower(), None
+ return None
+
+
+# ---------------------------------------------------------------------------
+# Formatting Spark values the way DataFusion's sqllogictest engine prints them
+#
+# Top-level values follow datafusion/sqllogictest/src/engines/conversion.rs and
+# datafusion_engine/normalize.rs. Values nested in arrays, maps and structs
+# follow arrow-rs's ArrayFormatter.
+# ---------------------------------------------------------------------------
+
+_SIMPLE_TYPES = {
+ "BinaryType": "binary",
+ "BooleanType": "boolean",
+ "ByteType": "int",
+ "CharType": "string",
+ "DateType": "date",
+ "DoubleType": "double",
+ "FloatType": "float",
+ "IntegerType": "int",
+ "LongType": "int",
+ "NullType": "null",
+ "ShortType": "int",
+ "StringType": "string",
+ "TimeType": "time",
+ "TimestampNTZType": "timestamp_ntz",
+ "TimestampType": "timestamp",
+ "VarcharType": "string",
+}
+
+# The type letters DataFusion's sqllogictest engine uses.
+_TYPE_LETTERS = {
+ "boolean": "B",
+ "date": "D",
+ "decimal": "R",
+ "double": "R",
+ "float": "R",
+ "int": "I",
+ "string": "T",
+ "time": "D",
+ "timestamp": "P",
+ "timestamp_ntz": "P",
+}
+
+
+def type_tree(data_type) -> tuple:
+ """Describe a PySpark DataType as nested tuples, so formatting needs no
PySpark."""
+ name = type(data_type).__name__
+ if name == "ArrayType":
+ return ("array", type_tree(data_type.elementType))
+ if name == "MapType":
+ return ("map", type_tree(data_type.keyType),
type_tree(data_type.valueType))
+ if name == "StructType":
+ fields = tuple((f.name, type_tree(f.dataType)) for f in
data_type.fields)
+ return ("struct", fields)
+ if name == "DecimalType":
+ return ("decimal", data_type.precision, data_type.scale)
+ return (_SIMPLE_TYPES.get(name, name),)
+
+
+def type_letter(tree: tuple) -> str:
+ return _TYPE_LETTERS.get(tree[0], "?")
+
+
+def _contains_timestamp(tree: tuple) -> bool:
+ if tree[0] == "timestamp":
+ return True
+ if tree[0] == "struct":
+ return any(_contains_timestamp(t) for _, t in tree[1])
+ return any(isinstance(t, tuple) and _contains_timestamp(t) for t in
tree[1:])
+
+
+def _shortest_f32(value: float) -> str:
+ """The shortest decimal text that round-trips to the same 32-bit float."""
+ for precision in range(1, 10):
+ text = f"{value:.{precision}g}"
+ try:
+ if struct.unpack("<f", struct.pack("<f", float(text)))[0] == value:
+ return text
+ except OverflowError:
+ pass
+ return repr(value)
+
+
+def _round_float(value: float, single: bool, digits: int) -> str:
+ """DataFusion's float_to_str: round, drop trailing zeros, no exponent."""
+ if math.isnan(value):
+ return "NaN"
+ if math.isinf(value):
+ return "Infinity" if value > 0 else "-Infinity"
+ text = _shortest_f32(value) if single else repr(value)
+ with localcontext() as ctx:
+ ctx.prec = 1000
+ rounded = Decimal(text).quantize(Decimal(1).scaleb(-digits),
ROUND_HALF_UP)
+ text = format(rounded.normalize(), "f")
+ # BigDecimal has no negative zero.
+ return "0" if text == "-0" else text
+
+
+def _ryu(value: float, single: bool) -> str:
+ """Format a float the way the Rust `ryu` crate does (used by arrow-rs)."""
+ if math.isnan(value):
+ return "NaN"
+ if math.isinf(value):
+ return "inf" if value > 0 else "-inf"
+ if value == 0:
+ return "-0.0" if math.copysign(1, value) < 0 else "0.0"
+ sign, digit_tuple, exponent = Decimal(_shortest_f32(value) if single else
repr(value)).as_tuple()
+ digits = "".join(map(str, digit_tuple)).rstrip("0")
+ exponent += len(digit_tuple) - len(digits)
+ length = len(digits)
+ kk = length + exponent
+ limit, low = (13, -6) if single else (16, -5)
+ if 0 <= exponent and kk <= limit:
+ text = digits + "0" * exponent + ".0"
+ elif 0 < kk <= limit:
+ text = digits[:kk] + "." + digits[kk:]
+ elif low < kk <= 0:
+ text = "0." + "0" * -kk + digits
+ elif length == 1:
+ text = f"{digits}e{kk - 1}"
+ else:
+ text = f"{digits[0]}.{digits[1:]}e{kk - 1}"
+ return ("-" if sign else "") + text
+
+
+def _plain_decimal(value: Decimal, scale: int) -> str:
+ with localcontext() as ctx:
+ ctx.prec = 1000
+ return format(value.quantize(Decimal(1).scaleb(-scale)), "f")
+
+
+def _fraction(microsecond: int) -> str:
+ """chrono prints fractional seconds in groups of 3 digits, and none for
zero."""
+ if microsecond == 0:
+ return ""
+ if microsecond % 1000 == 0:
+ return f".{microsecond // 1000:03d}"
+ return f".{microsecond:06d}"
+
+
+def _naive_timestamp(value: datetime.datetime) -> str:
+ return (
+ f"{value.year:04d}-{value.month:02d}-{value.day:02d}T"
+ f"{value.hour:02d}:{value.minute:02d}:{value.second:02d}"
+ f"{_fraction(value.microsecond)}"
+ )
+
+
+def zone(name: str) -> datetime.tzinfo:
+ """Resolve a Spark/DataFusion time zone name or offset such as +05:30."""
+ if name.upper() in ("UTC", "Z", "GMT", "+00:00"):
+ return datetime.timezone.utc
+ m = re.fullmatch(r"([+-])(\d{1,2}):?(\d{2})", name)
+ if m:
+ offset = datetime.timedelta(hours=int(m.group(2)),
minutes=int(m.group(3)))
+ return datetime.timezone(-offset if m.group(1) == "-" else offset)
+ from zoneinfo import ZoneInfo
+
+ return ZoneInfo(name)
+
+
+def _local_timestamp(value: datetime.datetime, time_zone: str) -> str:
+ """Spark TIMESTAMP values, which PySpark returns as naive UTC (TZ=UTC)."""
+ utc = value.replace(tzinfo=datetime.timezone.utc)
+ return
_naive_timestamp(utc.astimezone(zone(time_zone)).replace(tzinfo=None))
+
+
+def format_nested(value, tree: tuple, time_zone: str) -> str:
+ """Format a value the way arrow-rs's ArrayFormatter displays it."""
+ if value is None:
+ return "NULL"
+ kind = tree[0]
+ if kind == "boolean":
+ return "true" if value else "false"
+ if kind == "int":
+ return str(value)
+ if kind in ("float", "double"):
+ return _ryu(value, kind == "float")
+ if kind == "decimal":
+ return _plain_decimal(value, tree[2])
+ if kind == "string":
+ return value
+ if kind == "binary":
+ return bytes(value).hex()
+ if kind == "date":
+ return value.isoformat()
+ if kind == "timestamp":
+ return _local_timestamp(value, time_zone)
+ if kind == "timestamp_ntz":
+ return _naive_timestamp(value)
+ if kind == "time":
+ return (
+ f"{value.hour:02d}:{value.minute:02d}:{value.second:02d}"
+ f"{_fraction(value.microsecond)}"
+ )
+ if kind == "array":
+ return "[" + ", ".join(format_nested(v, tree[1], time_zone) for v in
value) + "]"
+ if kind == "map":
+ # A dict from PySpark, or a list of (key, value) rows from map_entries.
+ pairs = value.items() if isinstance(value, dict) else ((e[0], e[1])
for e in value)
+ entries = (
+ f"{format_nested(k, tree[1], time_zone)}: {format_nested(v,
tree[2], time_zone)}"
+ for k, v in pairs
+ )
+ return "{" + ", ".join(entries) + "}"
+ if kind == "struct":
+ entries = (
+ f"{name}: {format_nested(value[i], t, time_zone)}"
+ for i, (name, t) in enumerate(tree[1])
+ )
+ return "{" + ", ".join(entries) + "}"
+ if kind == "null":
+ return "NULL"
+ raise Unsupported(kind)
+
+
+def format_cell(value, tree: tuple, time_zone: str) -> str:
+ """Format a top-level value the way DataFusion's sqllogictest engine
prints it."""
+ if value is None:
+ return "NULL"
+ kind = tree[0]
+ if kind == "string":
+ if value == "":
+ return "(empty)"
+ return value.rstrip("\n").replace("\x00", "\\0")
+ if kind == "float":
+ return _round_float(value, single=True, digits=12)
+ if kind == "double":
+ # The engine keeps 15 decimal places for the Spark test files.
+ return _round_float(value, single=False, digits=15)
+ return format_nested(value, tree, time_zone)
+
+
+def format_rows(rows, trees: list[tuple], time_zone: str) -> list[str]:
+ """Format rows as result lines, splitting a multi-line last cell like the
engine."""
+ lines = []
+ for row in rows:
+ cells = [format_cell(value, tree, time_zone) for value, tree in
zip(row, trees)]
+ parts = cells[-1].split("\n") if cells else []
+ if len(parts) < 2:
+ lines.append(" ".join(cells).rstrip())
+ continue
+ lines.append(" ".join(cells[:-1]).rstrip())
+ for n, part in enumerate(parts, 1):
+ content = part.lstrip()
+ indent = "-" * (len(part) - len(content))
+ lines.append(f"{n:02d}){indent}{content}".rstrip())
+ return lines
+
+
+_ZONED_TIMESTAMP_RE = re.compile(
+
r"(\d{4})-(\d{2})-(\d{2})T(\d{2}):(\d{2}):(\d{2})(?:\.(\d{1,9}))?(Z|[+-]\d{2}:\d{2})"
+)
+
+
+def localize_timestamps(line: str, time_zone: str) -> str:
+ """Rewrite DataFusion timestamps with an offset as wall-clock time in
time_zone.
+
+ Spark's TIMESTAMP has no zone of its own, so it is printed in the session
+ time zone. DataFusion prints a zoned timestamp in its type's zone with an
+ offset (`2024-07-15T00:00:00Z`). This puts both on the same footing.
+ """
+
+ def replace(m: re.Match) -> str:
+ fraction = m.group(7) or ""
+ if len(fraction) > 6 and int(fraction[6:]):
+ return m.group(0)
+ micro = int(fraction[:6].ljust(6, "0")) if fraction else 0
+ offset = datetime.timezone.utc if m.group(8) == "Z" else
zone(m.group(8))
+ parts = [int(m.group(g)) for g in range(1, 7)]
+ value = datetime.datetime(*parts, micro, tzinfo=offset)
+ return
_naive_timestamp(value.astimezone(zone(time_zone)).replace(tzinfo=None))
+
+ return _ZONED_TIMESTAMP_RE.sub(replace, line)
+
+
+# ---------------------------------------------------------------------------
+# Comparing results
+# ---------------------------------------------------------------------------
+
+
+def _sorted(lines: list[str], sort_mode: str) -> list[str]:
+ lines = [line.rstrip() for line in lines]
+ if sort_mode == "rowsort":
+ return sorted(lines)
+ if sort_mode == "valuesort":
+ return sorted(value for line in lines for value in line.split())
+ return lines
+
+
+def results_match(expected: list[str], actual: list[str], sort_mode: str) ->
bool:
+ return _sorted(expected, sort_mode) == _sorted(actual, sort_mode)
+
+
+def numeric_difference(expected: list[str], actual: list[str], sort_mode: str)
-> Optional[str]:
+ """Classify results that differ only in their numbers.
+
+ Returns "scale" if the numbers are equal but printed differently (`7` and
+ `7.0`), "close" if some also differ in their last floating-point digits,
and
+ None if anything else differs.
+ """
+ expected, actual = _sorted(expected, sort_mode), _sorted(actual, sort_mode)
+ if len(expected) != len(actual):
+ return None
+ kind = "scale"
+ for e_line, a_line in zip(expected, actual):
+ e_values, a_values = e_line.split(" "), a_line.split(" ")
+ if len(e_values) != len(a_values):
+ return None
+ for e, a in zip(e_values, a_values):
+ if e == a:
+ continue
+ try:
+ e_float, a_float = float(e), float(a)
+ except ValueError:
+ return None
+ if e_float == a_float:
+ continue
+ if not math.isclose(e_float, a_float, rel_tol=1e-12,
abs_tol=1e-15):
+ return None
+ kind = "close"
+ return kind
+
+
+# ---------------------------------------------------------------------------
+# Running records on Spark
+# ---------------------------------------------------------------------------
+
+
+class SparkError(Exception):
+ """Spark rejected or failed a query. `kind` is parse, function, table or
error."""
+
+ def __init__(self, kind: str, message: str):
+ super().__init__(message)
+ self.kind = kind
+ self.message = message
+
+
+class SparkRunner:
+ """A local Spark session reset to SPARK_DEFAULTS before each file."""
+
+ def __init__(self):
+ # A SPARK_HOME for another Spark version breaks the pip-installed
+ # PySpark with "'JavaPackage' object is not callable".
+ os.environ.pop("SPARK_HOME", None)
+ # PySpark returns TIMESTAMP values in the Python process's time zone.
+ os.environ["TZ"] = "UTC"
+ if hasattr(time, "tzset"):
+ time.tzset()
+
+ install = f"pip install -r {SCRIPT_DIR / 'requirements.txt'}"
+ try:
+ import pyspark
+ except ImportError:
+ raise SystemExit(f"error: PySpark is not installed. Install it
with `{install}`.")
+
+ if not pyspark.__version__.startswith(SPARK_VERSION + "."):
+ raise SystemExit(
+ f"error: this runner needs PySpark {SPARK_VERSION}, found "
+ f"{pyspark.__version__}. Install it with `{install}`."
+ )
+ from pyspark.logger import PySparkLogger
+
+ # PySpark logs every failed query to stderr as JSON. The runner reports
+ # them itself.
+ for name in ("SQLQueryContextLogger", "DataFrameQueryContextLogger"):
+ PySparkLogger.getLogger(name).disabled = True
+ self._start()
+
+ def _start(self) -> None:
+ from pyspark.sql import SparkSession
+
+ builder = (
+ SparkSession.builder.master("local[1]")
+ .appName("datafusion-pyspark-slt-runner")
+ .config("spark.ui.enabled", "false")
+ .config("spark.ui.showConsoleProgress", "false")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.driver.bindAddress", "127.0.0.1")
+ .config("spark.sql.shuffle.partitions", "1")
+ .config("spark.sql.warehouse.dir",
tempfile.mkdtemp(prefix="pyspark-slt-"))
+ # Spark logs every failed query; the runner reports them instead.
+ .config("spark.log.level", "FATAL")
+ )
+ for key, value in SPARK_DEFAULTS.items():
+ builder = builder.config(key, value)
+ self.spark = builder.getOrCreate()
+ self.spark.sparkContext.setLogLevel("FATAL")
+ self.version = self.spark.version
+ self.time_zone = SPARK_DEFAULTS["spark.sql.session.timeZone"]
+ # Set when a query runs the JVM out of memory, which leaves it
unusable.
+ self.broken = False
+
+ def restart(self) -> None:
+ """Replace the session with a new JVM, for example after an
OutOfMemoryError."""
+ from pyspark import SparkContext
+ from pyspark.sql import SparkSession
+
+ # A JVM that ran out of memory can hang in a clean stop, so end the
process.
+ process = getattr(SparkContext._gateway, "proc", None)
+ if process is not None:
+ process.kill()
+ process.wait()
+ # PySpark keeps the old gateway and session in class attributes, and
+ # would reuse them instead of launching a new JVM.
+ SparkContext._gateway = None
+ SparkContext._jvm = None
+ SparkContext._active_spark_context = None
+ SparkSession._instantiatedSession = None
+ SparkSession._activeSession = None
+ self._start()
+
+ def reset(self) -> None:
+ for key, value in SPARK_DEFAULTS.items():
+ self.spark.conf.set(key, value)
+ self.time_zone = SPARK_DEFAULTS["spark.sql.session.timeZone"]
+ for table in self.spark.catalog.listTables():
+ if table.isTemporary:
+ self.spark.catalog.dropTempView(table.name)
+
+ def set(self, key: str, value: Optional[str]) -> None:
+ from pyspark.errors import PySparkException
+
+ if value is None:
+ value = SPARK_DEFAULTS[key]
+ elif key == "spark.sql.ansi.enabled":
+ value = value.lower()
+ try:
+ self.spark.conf.set(key, value)
+ except PySparkException as e:
+ raise SparkError("error", _first_line(e))
+ if key == "spark.sql.session.timeZone":
+ self.time_zone = value
+
+ def run(self, sql: str):
+ """Run SQL and return (schema, rows), or raise SparkError."""
+ from py4j.protocol import Py4JJavaError
+ from pyspark.errors import ParseException, PySparkException
+
+ try:
+ df = self.spark.sql(sql)
+ schema = df.schema
+ return schema, _maps_as_entries(df, schema).limit(MAX_ROWS +
1).collect()
+ except TypeError as e:
+ # PySpark returns each MAP as a dict, which cannot have array keys.
+ raise SparkError("convert", f"PySpark cannot convert the result
({e})")
Review Comment:
This `except TypeError` also covers `_maps_as_entries(df, schema)` on line
1233, so a `TypeError` from the runner's own code is reported as a skip
("PySpark cannot convert the result") instead of `runner-error`. The README
promises the opposite for runner bugs. Could we build the DataFrame outside the
`try` and wrap only `.collect()`?
##########
crates/spark/scripts/pyspark_slt_runner.py:
##########
@@ -0,0 +1,1587 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under
one or more contributor license agreements
+
+"""Run the datafusion-spark sqllogictest files against Apache Spark 4.2.
+
+Each record in crates/sqllogictest/test_files/**/*.slt is run on a
+local PySpark session, and Spark's result is compared with the expected result
+in the file. The expected results come from DataFusion, so every difference
+reported here is a query where DataFusion and Spark disagree, or where the
+test expects the wrong value.
+
+This is a manual tool and does not run in CI. See README.md for setup.
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import math
+import os
+import re
+import struct
+import sys
+import tempfile
+import time
+import traceback
+from collections import Counter
+from dataclasses import dataclass, field
+from decimal import ROUND_HALF_UP, Decimal, localcontext
+from pathlib import Path
+from typing import Optional
+
+SPARK_VERSION = "4.2"
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+REPO_ROOT = SCRIPT_DIR.parents[2]
+DEFAULT_TEST_DIR = REPO_ROOT / "crates" / "sqllogictest" / "test_files"
+
+# The label that `skipif` and `onlyif` directives use for this runner.
+LABEL = "spark"
+
+# DataFusion runs these files with ANSI mode off and no time zone (UTC), so
+# every file starts from the same Spark settings.
+SPARK_DEFAULTS = {
+ "spark.sql.ansi.enabled": "false",
+ "spark.sql.session.timeZone": "UTC",
+ "spark.sql.mapKeyDedupPolicy": "EXCEPTION",
+ # Spark 4.2 rejects the TIME type unless this internal setting is on.
+ "spark.sql.timeType.enabled": "true",
+}
+
+# DataFusion settings that have a Spark equivalent.
+SETTINGS = {
+ "datafusion.execution.enable_ansi_mode": "spark.sql.ansi.enabled",
+ "datafusion.execution.time_zone": "spark.sql.session.timeZone",
+ "datafusion.spark.map_key_dedup_policy": "spark.sql.mapKeyDedupPolicy",
+}
+
+# Type names DataFusion accepts and Spark does not.
+TYPE_ALIASES = {
+ "BYTEA": "BINARY",
+ "FLOAT4": "FLOAT",
+ "FLOAT8": "DOUBLE",
+ "INT2": "SMALLINT",
+ "INT4": "INT",
+ "INT8": "BIGINT",
+ "TEXT": "STRING",
+}
+
+# Spark's maximum decimal precision.
+MAX_DECIMAL_PRECISION = 38
+
+# The most rows collected from Spark for one query. Some tests expect an error
+# for a query that Spark answers with billions of rows, such as
stack(2000000000, 1).
+MAX_ROWS = 1000
+
+# The most result lines printed for one difference.
+MAX_REPORT_LINES = 20
+
+
+class Skip(Exception):
+ """The record cannot be run on Spark. The message says why."""
+
+
+class Unsupported(Exception):
+ """Spark returned a value of a type this runner does not format."""
+
+
+# ---------------------------------------------------------------------------
+# Parsing .slt files
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class Record:
+ """A `statement` or `query` record from an .slt file."""
+
+ path: Path
+ line: int
+ kind: str # "statement" or "query"
+ sql: str = ""
+ expect_error: bool = False
+ types: str = ""
+ sort_mode: str = "nosort"
+ expected: list[str] = field(default_factory=list)
+ skip_reason: Optional[str] = None
+
+
+def parse_slt(path: Path) -> list[Record]:
+ """Parse the `statement` and `query` records of an .slt file."""
+ lines = path.read_text().split("\n")
+ records: list[Record] = []
+ conditions: list[tuple[str, str]] = []
+ i = 0
+ while i < len(lines):
+ words = lines[i].split()
+ if not words or words[0].startswith("#"):
+ i += 1
+ continue
+ if words[0] in ("skipif", "onlyif") and len(words) > 1:
+ conditions.append((words[0], words[1]))
+ i += 1
+ continue
+ if words[0] not in ("statement", "query"):
+ # Directives such as `include` and `halt`, which these files do
not use.
+ conditions = []
+ i += 1
+ continue
+
+ record = Record(path=path, line=i + 1, kind=words[0])
+ args = words[1:]
+ if args and args[0] == "error":
+ record.expect_error = True
+ elif record.kind == "query" and args:
+ record.types = args[0]
+ if len(args) > 1 and args[1] in ("nosort", "rowsort", "valuesort"):
+ record.sort_mode = args[1]
+ for directive, label in conditions:
+ if directive == "skipif" and label == LABEL:
+ record.skip_reason = f"skipif {label}"
+ elif directive == "onlyif" and label != LABEL:
+ record.skip_reason = f"onlyif {label}"
+ conditions = []
+
+ # The SQL runs until a blank line or the `----` separator.
+ i += 1
+ sql = []
+ while i < len(lines) and lines[i].strip() and lines[i] != "----":
+ if not lines[i].lstrip().startswith("#"):
+ sql.append(lines[i])
+ i += 1
+ record.sql = "\n".join(sql)
+
+ # Expected results (or an expected error message) run until an empty
+ # line. A result line may start with `#`, so it is not a comment here,
+ # and it may be a single space, the engine's empty first row before a
+ # multi-line value.
+ if i < len(lines) and lines[i] == "----":
+ i += 1
+ while i < len(lines) and lines[i] != "":
+ record.expected.append(lines[i])
+ i += 1
+ records.append(record)
+ return records
+
+
+# ---------------------------------------------------------------------------
+# Translating DataFusion SQL to Spark SQL
+#
+# Spark 4.2 parses most of these files as written, including `::` casts.
+# Only a few renames are applied here, and `arrow_cast` and `spark_cast`
+# become a CAST (see below). Anything else Spark cannot parse, and any function
+# Spark does not have (such as `arrow_typeof`), is skipped.
+# ---------------------------------------------------------------------------
+
+_TOKEN_RE = re.compile(
+ r"""
+ (?P<string>'(?:[^']|'')*')
+ | (?P<quoted>`[^`]*`|"[^"]*")
+ | (?P<comment>--[^\n]*)
+ | (?P<word>[A-Za-z_][A-Za-z0-9_]*)
+ | (?P<number>\d+(?:\.\d*)?(?:[eE][+-]?\d+)?|\.\d+(?:[eE][+-]?\d+)?)
+ | (?P<cast>::)
+ | (?P<space>\s+)
+ | (?P<other>.)
+ """,
+ re.VERBOSE | re.DOTALL,
+)
+
+_COLUMN_RE = re.compile(r"column(\d+)", re.IGNORECASE)
+
+
+def _tokenize(sql: str) -> list[tuple[str, str]]:
+ return [(m.lastgroup, m.group()) for m in _TOKEN_RE.finditer(sql)]
+
+
+def _next_significant(tokens: list[tuple[str, str]], start: int) -> int:
+ """Return the index of the next token that is not whitespace or a
comment."""
+ i = start
+ while i < len(tokens) and tokens[i][0] in ("space", "comment"):
+ i += 1
+ return i
+
+
+def _translate_type(tokens: list[tuple[str, str]], i: int) -> str:
+ """Translate the type name at tokens[i]."""
+ text = tokens[i][1]
+ upper = text.upper()
+ if upper in TYPE_ALIASES:
+ return TYPE_ALIASES[upper]
+ j = _next_significant(tokens, i + 1)
+ has_args = j < len(tokens) and tokens[j][1] == "("
+ if upper == "VARCHAR" and not has_args:
+ return "STRING"
+ if upper in ("DECIMAL", "DEC", "NUMERIC") and has_args:
+ k = _next_significant(tokens, j + 1)
+ if k < len(tokens) and tokens[k][0] == "number":
+ if int(tokens[k][1]) > MAX_DECIMAL_PRECISION:
+ raise Skip(f"DECIMAL precision above {MAX_DECIMAL_PRECISION}")
+ return text
+
+
+# ---------------------------------------------------------------------------
+# arrow_cast
+#
+# `arrow_cast(expr, 'Type')` casts to an Arrow type. Many Arrow types differ
+# from a Spark type only in their physical layout: Utf8, LargeUtf8 and Utf8View
+# are all Spark's STRING, and a dictionary is an encoding of its value type. So
+# the call becomes a CAST to that Spark type. A type Spark does not have, such
+# as an unsigned integer, skips the record.
+#
+# Arrow interprets a number cast to a date, time or timestamp as a count of the
+# type's unit (days, seconds, nanoseconds, ...), and Spark does not, so those
+# casts are only translated for literal arguments.
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class ArrowType:
+ """A parsed Arrow type, such as `Decimal128(10, 2)` or `List(Int32)`."""
+
+ name: str
+ args: tuple = ()
+
+
+@dataclass(frozen=True)
+class ArrowField:
+ """A named field in a `Struct` or `Map` type, such as `"a": Int64`."""
+
+ name: str
+ type: object
+
+
+_ARROW_TOKEN_RE = re.compile(
+ r"""\s*(?:
+ (?P<string>"(?:[^"\\]|\\.)*")
+ | (?P<number>-?\d+)
+ | (?P<word>[A-Za-zµ_][A-Za-z0-9µ_-]*)
+ | (?P<punct>[(),:])
+ )""",
+ re.VERBOSE,
+)
+
+
+def parse_arrow_type(text: str):
+ """Parse an Arrow type as DataFusion's arrow_cast spells it."""
+ tokens = []
+ pos, text = 0, text.strip()
+ while pos < len(text):
+ m = _ARROW_TOKEN_RE.match(text, pos)
+ if not m:
+ raise ValueError(text)
+ tokens.append((m.lastgroup, m.group(m.lastgroup)))
+ pos = m.end()
+ i = 0
+
+ def peek() -> Optional[str]:
+ return tokens[i][1] if i < len(tokens) else None
+
+ def expect(value: str) -> None:
+ nonlocal i
+ if peek() != value:
+ raise ValueError(text)
+ i += 1
+
+ def value():
+ nonlocal i
+ if i >= len(tokens):
+ raise ValueError(text)
+ kind, token = tokens[i]
+ i += 1
+ if kind == "number":
+ return int(token)
+ if kind == "string":
+ name = token[1:-1]
+ if peek() == ":":
+ i += 1
+ return ArrowField(name, value())
+ return name
+ if kind == "word":
+ # Nullability does not change the Spark type.
+ if token in ("non-null", "nullable"):
+ return value()
+ args = []
+ if peek() == "(":
+ i += 1
+ while peek() != ")":
+ args.append(value())
+ if peek() == ",":
+ i += 1
+ expect(")")
+ return ArrowType(token, tuple(args))
+ raise ValueError(text)
+
+ parsed = value()
+ if i != len(tokens) or not isinstance(parsed, ArrowType):
+ raise ValueError(text)
+ return parsed
+
+
+_ARROW_SIMPLE_TYPES = {
+ "Boolean": "BOOLEAN",
+ "Int8": "TINYINT",
+ "Int16": "SMALLINT",
+ "Int32": "INT",
+ "Int64": "BIGINT",
+ "Float32": "FLOAT",
+ "Float64": "DOUBLE",
+ "Utf8": "STRING",
+ "LargeUtf8": "STRING",
+ "Utf8View": "STRING",
+ "Binary": "BINARY",
+ "LargeBinary": "BINARY",
+ "BinaryView": "BINARY",
+ "FixedSizeBinary": "BINARY",
+ "Date32": "DATE",
+ "Date64": "DATE",
+}
+
+_ARROW_LIST_TYPES = ("List", "LargeList", "ListView", "LargeListView",
"FixedSizeList")
+
+_ARROW_DECIMAL_TYPES = ("Decimal32", "Decimal64", "Decimal128", "Decimal256")
+
+_ARROW_UNITS = {
+ "Second": "s",
+ "s": "s",
+ "Millisecond": "ms",
+ "ms": "ms",
+ "Microsecond": "us",
+ "µs": "us",
+ "us": "us",
+ "Nanosecond": "ns",
+ "ns": "ns",
+}
+
+# Spark's TIME(p) precision for each Arrow time unit. Spark has no nanoseconds.
+_TIME_PRECISION = {"s": 0, "ms": 3, "us": 6}
+
+_UTC_NAMES = {"UTC", "Etc/UTC", "Z", "+00:00", "+00", "+0000", "GMT"}
+
+
+def _arrow_unit(t: ArrowType) -> str:
+ if not t.args or not isinstance(t.args[0], ArrowType) or t.args[0].name
not in _ARROW_UNITS:
+ raise Skip("arrow_cast to a type this runner cannot parse")
+ return _ARROW_UNITS[t.args[0].name]
+
+
+def _arrow_time_zone(t: ArrowType) -> Optional[str]:
+ """The time zone of an Arrow Timestamp: `None`, `Some("UTC")` or
`"UTC"`."""
+ if len(t.args) < 2:
+ return None
+ tz = t.args[1]
+ if isinstance(tz, str):
+ return tz
+ if isinstance(tz, ArrowType) and tz.name == "Some" and len(tz.args) == 1:
+ return tz.args[0]
+ if isinstance(tz, ArrowType) and tz.name == "None":
+ return None
+ raise Skip("arrow_cast to a type this runner cannot parse")
+
+
+def _same_time_zone(a: str, b: str) -> bool:
+ return a == b or (a in _UTC_NAMES and b in _UTC_NAMES)
+
+
+def spark_type_for_arrow(t: ArrowType) -> str:
+ """The Spark type with the same values as an Arrow type, or raise Skip."""
+ name = t.name
+ if name in _ARROW_SIMPLE_TYPES:
+ return _ARROW_SIMPLE_TYPES[name]
+ if name in ("UInt8", "UInt16", "UInt32", "UInt64"):
+ raise Skip("arrow_cast to an unsigned integer, which Spark does not
have")
+ if name in _ARROW_DECIMAL_TYPES:
+ precision, scale = t.args
+ if precision > MAX_DECIMAL_PRECISION:
+ raise Skip(f"DECIMAL precision above {MAX_DECIMAL_PRECISION}")
+ if scale < 0:
+ raise Skip("arrow_cast to a DECIMAL with a negative scale")
+ return f"DECIMAL({precision}, {scale})"
+ if name in _ARROW_LIST_TYPES:
+ return f"ARRAY<{spark_type_for_arrow(t.args[-1])}>"
+ if name == "Dictionary":
+ return spark_type_for_arrow(t.args[1])
+ if name == "Struct":
+ fields = [f"`{f.name}`: {spark_type_for_arrow(f.type)}" for f in
t.args]
+ return f"STRUCT<{', '.join(fields)}>"
+ if name == "Map":
+ entries = t.args[0]
+ key, value = entries.type.args
+ return f"MAP<{spark_type_for_arrow(key.type)},
{spark_type_for_arrow(value.type)}>"
+ if name in ("Time32", "Time64"):
+ unit = _arrow_unit(t)
+ if unit not in _TIME_PRECISION:
+ raise Skip("arrow_cast to nanoseconds, which Spark does not have")
+ return f"TIME({_TIME_PRECISION[unit]})"
+ if name == "Timestamp":
+ unit = _arrow_unit(t)
+ if unit != "us" or _arrow_time_zone(t) is not None:
+ raise Skip("arrow_cast to a nested timestamp that is not
TIMESTAMP_NTZ")
+ return "TIMESTAMP_NTZ"
+ raise Skip(f"arrow_cast to {name}, which Spark does not have")
+
+
+def _strip_parens(tokens: list[tuple[str, str]]) -> list[tuple[str, str]]:
+ """Drop whitespace, comments and parentheses that wrap the whole
argument."""
+ tokens = [t for t in tokens if t[0] not in ("space", "comment")]
+ while len(tokens) >= 2 and tokens[0][1] == "(" and tokens[-1][1] == ")":
+ depth = 0
+ for k, (_, text) in enumerate(tokens):
+ depth += {"(": 1, ")": -1}.get(text, 0)
+ if depth == 0 and k < len(tokens) - 1:
+ return tokens
+ tokens = tokens[1:-1]
+ return tokens
+
+
+def _literal(tokens: list[tuple[str, str]]):
+ """Return ("null", None), ("string", text) or ("integer", n) for a literal
argument."""
+ tokens = _strip_parens(tokens)
+ if len(tokens) == 1 and tokens[0][0] == "word" and tokens[0][1].upper() ==
"NULL":
+ return ("null", None)
+ if len(tokens) == 1 and tokens[0][0] == "string":
+ return ("string", tokens[0][1][1:-1].replace("''", "'"))
+ # A typed literal, such as DATE '2024-01-01' or TIMESTAMP '2024-01-01
10:00:00'
+ if (
+ len(tokens) == 2
+ and tokens[0][0] == "word"
+ and tokens[0][1].upper() in ("DATE", "TIMESTAMP", "TIMESTAMP_NTZ")
+ and tokens[1][0] == "string"
+ ):
+ return ("typed", tokens[1][1][1:-1].replace("''", "'"))
+ sign = 1
+ if len(tokens) == 2 and tokens[0][1] in ("-", "+"):
+ sign = -1 if tokens[0][1] == "-" else 1
+ tokens = tokens[1:]
+ if len(tokens) == 1 and tokens[0][0] == "number" and
tokens[0][1].isdigit():
+ return ("integer", sign * int(tokens[0][1]))
+ return (None, None)
+
+
+# Fractional-second digits each timestamp unit keeps.
+_UNIT_DIGITS = {"s": 0, "ms": 3, "us": 6, "ns": 9}
+
+_FRACTION_RE = re.compile(r":\d{2}\.(\d+)")
+
+_EPOCH_FUNCTIONS = {"s": "timestamp_seconds", "ms": "timestamp_millis", "us":
"timestamp_micros"}
+
+_MILLIS_PER_DAY = 86_400_000
+
+
+def _temporal_cast(t: ArrowType, arg: list[tuple[str, str]], arg_sql: str,
time_zone: str) -> str:
+ """Translate arrow_cast to a date, time or timestamp, which depends on the
argument."""
+ kind, value = _literal(arg)
+ if kind == "typed":
+ # DataFusion's TIMESTAMP literal has no time zone and Spark's is in the
+ # session time zone, so they are the same value only in UTC.
+ if time_zone not in _UTC_NAMES:
+ raise Skip(f"arrow_cast of a typed literal to {t.name} in a
session time zone other than UTC")
+ kind = "string"
+ arg_sql = "'" + value.replace("'", "''") + "'"
+ if kind is None:
+ raise Skip(f"arrow_cast of an expression to {t.name}, which Arrow and
Spark convert differently")
+ if t.name in ("Date32", "Date64"):
+ if kind == "integer":
+ days = value
+ if t.name == "Date64":
+ if value % _MILLIS_PER_DAY:
+ raise Skip("arrow_cast to Date64 of milliseconds that are
not whole days")
+ days = value // _MILLIS_PER_DAY
+ return f"date_from_unix_date({days})"
+ return f"CAST({arg_sql} AS DATE)"
+ if t.name in ("Time32", "Time64"):
+ spark_type = spark_type_for_arrow(t)
+ if kind == "integer":
+ raise Skip(f"arrow_cast of a number to {t.name}, which Spark
cannot cast")
+ return f"CAST({arg_sql} AS {spark_type})"
+
+ # Timestamp
+ unit = _arrow_unit(t)
+ tz = _arrow_time_zone(t)
+ if tz is not None and not _same_time_zone(tz, time_zone):
+ raise Skip("arrow_cast to a timestamp in a time zone other than the
session's")
+ # In a UTC session, Spark's TIMESTAMP holds the same values as an Arrow
+ # timestamp without a time zone, and functions such as unix_micros accept
+ # only TIMESTAMP, not TIMESTAMP_NTZ.
+ utc_session = time_zone in _UTC_NAMES
+ spark_type = "TIMESTAMP" if tz is not None or utc_session else
"TIMESTAMP_NTZ"
+ if kind == "null":
+ return f"CAST(NULL AS {spark_type})"
+ if kind == "string":
+ m = _FRACTION_RE.search(value)
+ digits = len(m.group(1).rstrip("0")) if m else 0
+ if digits > 6:
+ raise Skip("arrow_cast to nanoseconds, which Spark does not have")
+ if digits > _UNIT_DIGITS[unit]:
+ raise Skip("arrow_cast that truncates a timestamp to seconds or
milliseconds")
+ return f"CAST({arg_sql} AS {spark_type})"
+ # An integer counts units since the epoch.
+ if unit == "ns":
+ if value % 1000:
+ raise Skip("arrow_cast to nanoseconds, which Spark does not have")
+ value, unit = value // 1000, "us"
+ if tz is None and not utc_session:
+ # Spark would print the instant in the session time zone, and Arrow
+ # prints a timestamp without a time zone as UTC.
+ raise Skip("arrow_cast of a number to a timestamp in a session time
zone other than UTC")
+ return f"{_EPOCH_FUNCTIONS[unit]}({value})"
+
+
+# Brackets that nest, so a comma inside `[1, 2]` or `f(a, b)` does not split
arguments.
+_NESTING = {"(": 1, "[": 1, "{": 1, ")": -1, "]": -1, "}": -1}
+
+
+# Spark functions that return a DATE.
+_DATE_FUNCTIONS = {
+ "add_months",
+ "current_date",
+ "curdate",
+ "date_add",
+ "date_from_unix_date",
+ "date_sub",
+ "dateadd",
+ "last_day",
+ "make_date",
+ "next_day",
+ "to_date",
+ "trunc",
+}
+
+_INTEGER_TYPES = ("TINYINT", "SMALLINT", "INT", "BIGINT")
+
+
+def _is_date(sql: str) -> bool:
+ """Whether a translated expression is visibly a DATE: a DATE literal, a
cast
+ to DATE, or a call to a function that returns one."""
+ tokens = _strip_parens(_tokenize(sql))
+ if len(tokens) == 2 and tokens[0][1].upper() == "DATE" and tokens[1][0] ==
"string":
+ return True
+ if len(tokens) < 3 or tokens[0][0] != "word" or tokens[1][1] != "(":
+ return False
+ if _matching_paren(tokens, 1) != len(tokens) - 1:
+ return False
+ name = tokens[0][1].lower()
+ if name == "cast":
+ inner = tokens[2:-1]
+ return len(inner) >= 2 and inner[-1][1].upper() == "DATE" and
inner[-2][1].upper() == "AS"
+ return name in _DATE_FUNCTIONS
+
+
+def _matching_paren(tokens: list[tuple[str, str]], open_index: int) -> int:
+ depth = 0
+ for k in range(open_index, len(tokens)):
+ if tokens[k][0] == "other":
+ depth += _NESTING.get(tokens[k][1], 0)
+ if depth == 0:
+ return k
+ raise Skip("unbalanced parentheses")
+
+
+# A Spark type name, such as `timestamp`, `decimal(10, 2)` or `array<int>`.
+_SPARK_TYPE_RE = re.compile(r"[A-Za-z][A-Za-z0-9_<>(),:`\s]*")
+
+
+def _call_args(tokens: list[tuple[str, str]], open_index: int, close: int) ->
list:
+ """Split the tokens between a call's parentheses on top-level commas."""
+ args, start, depth = [], open_index + 1, 0
+ for k in range(open_index + 1, close):
+ if tokens[k][0] == "other":
+ if tokens[k][1] == "," and depth == 0:
+ args.append(tokens[start:k])
+ start = k + 1
+ depth += _NESTING.get(tokens[k][1], 0)
+ args.append(tokens[start:close])
+ return args
+
+
+def _type_string(name: str, args: list) -> str:
+ """The type argument of arrow_cast or spark_cast, which must be a string
literal."""
+ if len(args) != 2:
+ raise Skip(f"{name} without exactly two arguments")
+ type_arg = _strip_parens(args[1])
+ if len(type_arg) != 1 or type_arg[0][0] != "string":
+ raise Skip(f"{name} to a type that is not a string literal")
+ return type_arg[0][1][1:-1].replace("''", "'")
+
+
+def _arrow_cast(arg: list, arg_sql: str, type_string: str, time_zone: str) ->
str:
+ try:
+ arrow_type = parse_arrow_type(type_string)
+ except (ValueError, TypeError, AttributeError):
+ raise Skip("arrow_cast to a type this runner cannot parse")
+ if arrow_type.name in ("Date32", "Date64", "Time32", "Time64",
"Timestamp"):
+ return _temporal_cast(arrow_type, arg, arg_sql, time_zone)
+ try:
+ spark_type = spark_type_for_arrow(arrow_type)
+ except (ValueError, TypeError, AttributeError, IndexError):
+ raise Skip("arrow_cast to a type this runner cannot parse")
+ if spark_type in _INTEGER_TYPES and _is_date(arg_sql):
+ # Arrow casts a date to an integer as days since the epoch.
+ # Spark's CAST returns NULL, and unix_date returns the days.
+ days = f"unix_date({arg_sql})"
+ return days if spark_type == "INT" else f"CAST({days} AS {spark_type})"
+ return f"CAST({arg_sql} AS {spark_type})"
+
+
+def rewrite_casts(sql: str, time_zone: str = "UTC") -> str:
+ """Replace each `arrow_cast(expr, 'Type')` and `spark_cast(expr, 'type')`
+ with a Spark CAST, or raise Skip.
+
+ `spark_cast` is this crate's emulation of Spark's CAST, and its second
+ argument is already a Spark type name.
+ """
+ tokens = _tokenize(sql)
+ out = []
+ i = 0
+ while i < len(tokens):
+ kind, text = tokens[i]
+ name = text.lower()
+ open_index = _next_significant(tokens, i + 1)
+ if not (
+ kind == "word"
+ and name in ("arrow_cast", "spark_cast")
+ and open_index < len(tokens)
+ and tokens[open_index][1] == "("
+ ):
+ out.append(text)
+ i += 1
+ continue
+ close = _matching_paren(tokens, open_index)
+ args = _call_args(tokens, open_index, close)
+ type_string = _type_string(name, args)
+ arg_sql = rewrite_casts("".join(t for _, t in args[0]),
time_zone).strip()
+ if name == "arrow_cast":
+ out.append(_arrow_cast(args[0], arg_sql, type_string, time_zone))
+ else:
+ if not _SPARK_TYPE_RE.fullmatch(type_string.strip()):
+ raise Skip("spark_cast to a type that is not a Spark type
name")
+ out.append(f"CAST({arg_sql} AS {type_string.strip()})")
+ i = close + 1
+ return "".join(out)
+
+
+def translate(sql: str, time_zone: str = "UTC") -> str:
+ """Translate the parts of DataFusion's SQL that Spark spells
differently."""
+ tokens = _tokenize(rewrite_casts(sql, time_zone))
+ out = []
+ # The function name before each open parenthesis, to find CAST(x AS type).
+ calls: list[Optional[str]] = []
+ last: Optional[tuple[str, str]] = None
+ expect_type = False
+ for i, (kind, text) in enumerate(tokens):
+ if kind == "string" and "\\" in text:
+ raise Skip(
+ "a string literal contains a backslash, which Spark reads as
an "
+ "escape and DataFusion does not"
+ )
+ if kind in ("space", "comment"):
+ out.append(text)
+ continue
+ if kind == "word":
+ if expect_type:
+ text = _translate_type(tokens, i)
+ else:
+ m = _COLUMN_RE.fullmatch(text)
+ if m:
+ # Spark names VALUES columns col1, col2, ...
+ text = f"col{m.group(1)}"
+ expect_type = text.upper() == "AS" and bool(calls) and calls[-1]
in (
+ "CAST",
+ "TRY_CAST",
+ )
+ elif kind == "cast":
+ expect_type = True
+ else:
+ expect_type = False
+ if text == "(":
+ calls.append(last[1].upper() if last and last[0] == "word"
else None)
+ elif text == ")" and calls:
+ calls.pop()
+ last = (kind, text)
+ out.append(text)
+ return "".join(out)
+
+
+_NAME = r"[A-Za-z_][A-Za-z0-9_]*"
+_CREATE_RE = re.compile(
+
rf"\s*CREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+(?:IF\s+NOT\s+EXISTS\s+)?({_NAME})\s*",
+ re.IGNORECASE,
+)
+_DROP_RE = re.compile(
+ rf"\s*DROP\s+(?:TABLE|VIEW)\s+(?:IF\s+EXISTS\s+)?({_NAME})\s*;?\s*$",
+ re.IGNORECASE,
+)
+_AS_RE = re.compile(r"AS\b\s*", re.IGNORECASE)
+
+
+def _split_top_level(text: str) -> list[str]:
+ """Split on commas that are not inside parentheses."""
+ parts, depth, start = [], 0, 0
+ for i, ch in enumerate(text):
+ if ch == "(":
+ depth += 1
+ elif ch == ")":
+ depth -= 1
+ elif ch == "," and depth == 0:
+ parts.append(text[start:i])
+ start = i + 1
+ parts.append(text[start:])
+ return [p.strip() for p in parts]
+
+
+def translate_statement(sql: str, time_zone: str = "UTC") -> str:
+ """Translate a statement, turning DataFusion tables into Spark temporary
views."""
+ m = _DROP_RE.match(sql)
+ if m:
+ return f"DROP VIEW IF EXISTS {m.group(1)}"
+ m = _CREATE_RE.match(sql)
+ if not m:
+ return translate(sql, time_zone)
+ name, rest = m.group(1), sql[m.end() :]
+ columns = None
+ if rest.startswith("("):
+ depth = 0
+ for end, ch in enumerate(rest):
+ depth += {"(": 1, ")": -1}.get(ch, 0)
+ if depth == 0:
+ break
+ columns = _split_top_level(rest[1:end])
+ rest = rest[end + 1 :].lstrip()
+ m = _AS_RE.match(rest)
+ if not m:
+ raise Skip("CREATE TABLE without AS")
+ query = rest[m.end() :].rstrip().rstrip(";")
+ if columns is None:
+ return translate(f"CREATE OR REPLACE TEMPORARY VIEW {name} AS
{query}", time_zone)
+ names, casts = [], []
+ for column in columns:
+ parts = column.split(None, 1)
+ if len(parts) != 2:
+ raise Skip(f"unsupported column definition: {column}")
+ names.append(parts[0])
+ casts.append(f"CAST({parts[0]} AS {parts[1]}) AS {parts[0]}")
+ return translate(
+ f"CREATE OR REPLACE TEMPORARY VIEW {name} AS SELECT {', '.join(casts)}
"
+ f"FROM ({query}) AS __slt_values({', '.join(names)})",
+ time_zone,
+ )
+
+
+_SET_RE = re.compile(
+ r"\s*SET\s+([A-Za-z0-9_.]+)\s*(?:=|\s+TO\s+)\s*(.*?)\s*;?\s*$",
+ re.IGNORECASE | re.DOTALL,
+)
+_RESET_RE = re.compile(r"\s*RESET\s+([A-Za-z0-9_.]+)\s*;?\s*$", re.IGNORECASE)
+
+
+def parse_setting(sql: str) -> Optional[tuple[str, Optional[str]]]:
+ """Parse `SET key = value` or `RESET key`. A reset has the value None."""
+ m = _SET_RE.match(sql)
+ if m:
+ value = m.group(2)
+ if len(value) >= 2 and value[0] == value[-1] and value[0] in "'\"":
+ value = value[1:-1]
+ return m.group(1).lower(), value
+ m = _RESET_RE.match(sql)
+ if m:
+ return m.group(1).lower(), None
+ return None
+
+
+# ---------------------------------------------------------------------------
+# Formatting Spark values the way DataFusion's sqllogictest engine prints them
+#
+# Top-level values follow datafusion/sqllogictest/src/engines/conversion.rs and
+# datafusion_engine/normalize.rs. Values nested in arrays, maps and structs
+# follow arrow-rs's ArrayFormatter.
+# ---------------------------------------------------------------------------
+
+_SIMPLE_TYPES = {
+ "BinaryType": "binary",
+ "BooleanType": "boolean",
+ "ByteType": "int",
+ "CharType": "string",
+ "DateType": "date",
+ "DoubleType": "double",
+ "FloatType": "float",
+ "IntegerType": "int",
+ "LongType": "int",
+ "NullType": "null",
+ "ShortType": "int",
+ "StringType": "string",
+ "TimeType": "time",
+ "TimestampNTZType": "timestamp_ntz",
+ "TimestampType": "timestamp",
+ "VarcharType": "string",
+}
+
+# The type letters DataFusion's sqllogictest engine uses.
+_TYPE_LETTERS = {
+ "boolean": "B",
+ "date": "D",
+ "decimal": "R",
+ "double": "R",
+ "float": "R",
+ "int": "I",
+ "string": "T",
+ "time": "D",
+ "timestamp": "P",
+ "timestamp_ntz": "P",
+}
+
+
+def type_tree(data_type) -> tuple:
+ """Describe a PySpark DataType as nested tuples, so formatting needs no
PySpark."""
+ name = type(data_type).__name__
+ if name == "ArrayType":
+ return ("array", type_tree(data_type.elementType))
+ if name == "MapType":
+ return ("map", type_tree(data_type.keyType),
type_tree(data_type.valueType))
+ if name == "StructType":
+ fields = tuple((f.name, type_tree(f.dataType)) for f in
data_type.fields)
+ return ("struct", fields)
+ if name == "DecimalType":
+ return ("decimal", data_type.precision, data_type.scale)
+ return (_SIMPLE_TYPES.get(name, name),)
+
+
+def type_letter(tree: tuple) -> str:
+ return _TYPE_LETTERS.get(tree[0], "?")
+
+
+def _contains_timestamp(tree: tuple) -> bool:
+ if tree[0] == "timestamp":
+ return True
+ if tree[0] == "struct":
+ return any(_contains_timestamp(t) for _, t in tree[1])
+ return any(isinstance(t, tuple) and _contains_timestamp(t) for t in
tree[1:])
+
+
+def _shortest_f32(value: float) -> str:
+ """The shortest decimal text that round-trips to the same 32-bit float."""
+ for precision in range(1, 10):
+ text = f"{value:.{precision}g}"
+ try:
+ if struct.unpack("<f", struct.pack("<f", float(text)))[0] == value:
+ return text
+ except OverflowError:
+ pass
+ return repr(value)
+
+
+def _round_float(value: float, single: bool, digits: int) -> str:
+ """DataFusion's float_to_str: round, drop trailing zeros, no exponent."""
+ if math.isnan(value):
+ return "NaN"
+ if math.isinf(value):
+ return "Infinity" if value > 0 else "-Infinity"
+ text = _shortest_f32(value) if single else repr(value)
+ with localcontext() as ctx:
+ ctx.prec = 1000
+ rounded = Decimal(text).quantize(Decimal(1).scaleb(-digits),
ROUND_HALF_UP)
+ text = format(rounded.normalize(), "f")
+ # BigDecimal has no negative zero.
+ return "0" if text == "-0" else text
+
+
+def _ryu(value: float, single: bool) -> str:
+ """Format a float the way the Rust `ryu` crate does (used by arrow-rs)."""
+ if math.isnan(value):
+ return "NaN"
+ if math.isinf(value):
+ return "inf" if value > 0 else "-inf"
+ if value == 0:
+ return "-0.0" if math.copysign(1, value) < 0 else "0.0"
+ sign, digit_tuple, exponent = Decimal(_shortest_f32(value) if single else
repr(value)).as_tuple()
+ digits = "".join(map(str, digit_tuple)).rstrip("0")
+ exponent += len(digit_tuple) - len(digits)
+ length = len(digits)
+ kk = length + exponent
+ limit, low = (13, -6) if single else (16, -5)
+ if 0 <= exponent and kk <= limit:
+ text = digits + "0" * exponent + ".0"
+ elif 0 < kk <= limit:
+ text = digits[:kk] + "." + digits[kk:]
+ elif low < kk <= 0:
+ text = "0." + "0" * -kk + digits
+ elif length == 1:
+ text = f"{digits}e{kk - 1}"
+ else:
+ text = f"{digits[0]}.{digits[1:]}e{kk - 1}"
+ return ("-" if sign else "") + text
+
+
+def _plain_decimal(value: Decimal, scale: int) -> str:
+ with localcontext() as ctx:
+ ctx.prec = 1000
+ return format(value.quantize(Decimal(1).scaleb(-scale)), "f")
+
+
+def _fraction(microsecond: int) -> str:
+ """chrono prints fractional seconds in groups of 3 digits, and none for
zero."""
+ if microsecond == 0:
+ return ""
+ if microsecond % 1000 == 0:
+ return f".{microsecond // 1000:03d}"
+ return f".{microsecond:06d}"
+
+
+def _naive_timestamp(value: datetime.datetime) -> str:
+ return (
+ f"{value.year:04d}-{value.month:02d}-{value.day:02d}T"
+ f"{value.hour:02d}:{value.minute:02d}:{value.second:02d}"
+ f"{_fraction(value.microsecond)}"
+ )
+
+
+def zone(name: str) -> datetime.tzinfo:
+ """Resolve a Spark/DataFusion time zone name or offset such as +05:30."""
+ if name.upper() in ("UTC", "Z", "GMT", "+00:00"):
+ return datetime.timezone.utc
+ m = re.fullmatch(r"([+-])(\d{1,2}):?(\d{2})", name)
+ if m:
+ offset = datetime.timedelta(hours=int(m.group(2)),
minutes=int(m.group(3)))
+ return datetime.timezone(-offset if m.group(1) == "-" else offset)
+ from zoneinfo import ZoneInfo
+
+ return ZoneInfo(name)
+
+
+def _local_timestamp(value: datetime.datetime, time_zone: str) -> str:
+ """Spark TIMESTAMP values, which PySpark returns as naive UTC (TZ=UTC)."""
+ utc = value.replace(tzinfo=datetime.timezone.utc)
+ return
_naive_timestamp(utc.astimezone(zone(time_zone)).replace(tzinfo=None))
+
+
+def format_nested(value, tree: tuple, time_zone: str) -> str:
+ """Format a value the way arrow-rs's ArrayFormatter displays it."""
+ if value is None:
+ return "NULL"
+ kind = tree[0]
+ if kind == "boolean":
+ return "true" if value else "false"
+ if kind == "int":
+ return str(value)
+ if kind in ("float", "double"):
+ return _ryu(value, kind == "float")
+ if kind == "decimal":
+ return _plain_decimal(value, tree[2])
+ if kind == "string":
+ return value
+ if kind == "binary":
+ return bytes(value).hex()
+ if kind == "date":
+ return value.isoformat()
+ if kind == "timestamp":
+ return _local_timestamp(value, time_zone)
+ if kind == "timestamp_ntz":
+ return _naive_timestamp(value)
+ if kind == "time":
+ return (
+ f"{value.hour:02d}:{value.minute:02d}:{value.second:02d}"
+ f"{_fraction(value.microsecond)}"
+ )
+ if kind == "array":
+ return "[" + ", ".join(format_nested(v, tree[1], time_zone) for v in
value) + "]"
+ if kind == "map":
+ # A dict from PySpark, or a list of (key, value) rows from map_entries.
+ pairs = value.items() if isinstance(value, dict) else ((e[0], e[1])
for e in value)
+ entries = (
+ f"{format_nested(k, tree[1], time_zone)}: {format_nested(v,
tree[2], time_zone)}"
+ for k, v in pairs
+ )
+ return "{" + ", ".join(entries) + "}"
+ if kind == "struct":
+ entries = (
+ f"{name}: {format_nested(value[i], t, time_zone)}"
+ for i, (name, t) in enumerate(tree[1])
+ )
+ return "{" + ", ".join(entries) + "}"
+ if kind == "null":
+ return "NULL"
+ raise Unsupported(kind)
+
+
+def format_cell(value, tree: tuple, time_zone: str) -> str:
+ """Format a top-level value the way DataFusion's sqllogictest engine
prints it."""
+ if value is None:
+ return "NULL"
+ kind = tree[0]
+ if kind == "string":
+ if value == "":
+ return "(empty)"
+ return value.rstrip("\n").replace("\x00", "\\0")
+ if kind == "float":
+ return _round_float(value, single=True, digits=12)
+ if kind == "double":
+ # The engine keeps 15 decimal places for the Spark test files.
+ return _round_float(value, single=False, digits=15)
+ return format_nested(value, tree, time_zone)
+
+
+def format_rows(rows, trees: list[tuple], time_zone: str) -> list[str]:
+ """Format rows as result lines, splitting a multi-line last cell like the
engine."""
+ lines = []
+ for row in rows:
+ cells = [format_cell(value, tree, time_zone) for value, tree in
zip(row, trees)]
+ parts = cells[-1].split("\n") if cells else []
+ if len(parts) < 2:
+ lines.append(" ".join(cells).rstrip())
+ continue
+ lines.append(" ".join(cells[:-1]).rstrip())
+ for n, part in enumerate(parts, 1):
+ content = part.lstrip()
+ indent = "-" * (len(part) - len(content))
+ lines.append(f"{n:02d}){indent}{content}".rstrip())
+ return lines
+
+
+_ZONED_TIMESTAMP_RE = re.compile(
+
r"(\d{4})-(\d{2})-(\d{2})T(\d{2}):(\d{2}):(\d{2})(?:\.(\d{1,9}))?(Z|[+-]\d{2}:\d{2})"
+)
+
+
+def localize_timestamps(line: str, time_zone: str) -> str:
+ """Rewrite DataFusion timestamps with an offset as wall-clock time in
time_zone.
+
+ Spark's TIMESTAMP has no zone of its own, so it is printed in the session
+ time zone. DataFusion prints a zoned timestamp in its type's zone with an
+ offset (`2024-07-15T00:00:00Z`). This puts both on the same footing.
+ """
+
+ def replace(m: re.Match) -> str:
+ fraction = m.group(7) or ""
+ if len(fraction) > 6 and int(fraction[6:]):
+ return m.group(0)
+ micro = int(fraction[:6].ljust(6, "0")) if fraction else 0
+ offset = datetime.timezone.utc if m.group(8) == "Z" else
zone(m.group(8))
+ parts = [int(m.group(g)) for g in range(1, 7)]
+ value = datetime.datetime(*parts, micro, tzinfo=offset)
+ return
_naive_timestamp(value.astimezone(zone(time_zone)).replace(tzinfo=None))
+
+ return _ZONED_TIMESTAMP_RE.sub(replace, line)
+
+
+# ---------------------------------------------------------------------------
+# Comparing results
+# ---------------------------------------------------------------------------
+
+
+def _sorted(lines: list[str], sort_mode: str) -> list[str]:
+ lines = [line.rstrip() for line in lines]
+ if sort_mode == "rowsort":
+ return sorted(lines)
+ if sort_mode == "valuesort":
+ return sorted(value for line in lines for value in line.split())
+ return lines
+
+
+def results_match(expected: list[str], actual: list[str], sort_mode: str) ->
bool:
+ return _sorted(expected, sort_mode) == _sorted(actual, sort_mode)
+
+
+def numeric_difference(expected: list[str], actual: list[str], sort_mode: str)
-> Optional[str]:
+ """Classify results that differ only in their numbers.
+
+ Returns "scale" if the numbers are equal but printed differently (`7` and
+ `7.0`), "close" if some also differ in their last floating-point digits,
and
+ None if anything else differs.
+ """
+ expected, actual = _sorted(expected, sort_mode), _sorted(actual, sort_mode)
+ if len(expected) != len(actual):
+ return None
+ kind = "scale"
+ for e_line, a_line in zip(expected, actual):
+ e_values, a_values = e_line.split(" "), a_line.split(" ")
+ if len(e_values) != len(a_values):
+ return None
+ for e, a in zip(e_values, a_values):
+ if e == a:
+ continue
+ try:
+ e_float, a_float = float(e), float(a)
+ except ValueError:
+ return None
+ if e_float == a_float:
+ continue
+ if not math.isclose(e_float, a_float, rel_tol=1e-12,
abs_tol=1e-15):
+ return None
+ kind = "close"
+ return kind
+
+
+# ---------------------------------------------------------------------------
+# Running records on Spark
+# ---------------------------------------------------------------------------
+
+
+class SparkError(Exception):
+ """Spark rejected or failed a query. `kind` is parse, function, table or
error."""
+
+ def __init__(self, kind: str, message: str):
+ super().__init__(message)
+ self.kind = kind
+ self.message = message
+
+
+class SparkRunner:
+ """A local Spark session reset to SPARK_DEFAULTS before each file."""
+
+ def __init__(self):
+ # A SPARK_HOME for another Spark version breaks the pip-installed
+ # PySpark with "'JavaPackage' object is not callable".
+ os.environ.pop("SPARK_HOME", None)
+ # PySpark returns TIMESTAMP values in the Python process's time zone.
+ os.environ["TZ"] = "UTC"
+ if hasattr(time, "tzset"):
+ time.tzset()
+
+ install = f"pip install -r {SCRIPT_DIR / 'requirements.txt'}"
+ try:
+ import pyspark
+ except ImportError:
+ raise SystemExit(f"error: PySpark is not installed. Install it
with `{install}`.")
+
+ if not pyspark.__version__.startswith(SPARK_VERSION + "."):
+ raise SystemExit(
+ f"error: this runner needs PySpark {SPARK_VERSION}, found "
+ f"{pyspark.__version__}. Install it with `{install}`."
+ )
+ from pyspark.logger import PySparkLogger
+
+ # PySpark logs every failed query to stderr as JSON. The runner reports
+ # them itself.
+ for name in ("SQLQueryContextLogger", "DataFrameQueryContextLogger"):
+ PySparkLogger.getLogger(name).disabled = True
+ self._start()
+
+ def _start(self) -> None:
+ from pyspark.sql import SparkSession
+
+ builder = (
+ SparkSession.builder.master("local[1]")
+ .appName("datafusion-pyspark-slt-runner")
+ .config("spark.ui.enabled", "false")
+ .config("spark.ui.showConsoleProgress", "false")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.driver.bindAddress", "127.0.0.1")
+ .config("spark.sql.shuffle.partitions", "1")
+ .config("spark.sql.warehouse.dir",
tempfile.mkdtemp(prefix="pyspark-slt-"))
+ # Spark logs every failed query; the runner reports them instead.
+ .config("spark.log.level", "FATAL")
+ )
+ for key, value in SPARK_DEFAULTS.items():
+ builder = builder.config(key, value)
+ self.spark = builder.getOrCreate()
+ self.spark.sparkContext.setLogLevel("FATAL")
+ self.version = self.spark.version
+ self.time_zone = SPARK_DEFAULTS["spark.sql.session.timeZone"]
+ # Set when a query runs the JVM out of memory, which leaves it
unusable.
+ self.broken = False
+
+ def restart(self) -> None:
+ """Replace the session with a new JVM, for example after an
OutOfMemoryError."""
+ from pyspark import SparkContext
+ from pyspark.sql import SparkSession
+
+ # A JVM that ran out of memory can hang in a clean stop, so end the
process.
+ process = getattr(SparkContext._gateway, "proc", None)
+ if process is not None:
+ process.kill()
+ process.wait()
+ # PySpark keeps the old gateway and session in class attributes, and
+ # would reuse them instead of launching a new JVM.
+ SparkContext._gateway = None
+ SparkContext._jvm = None
+ SparkContext._active_spark_context = None
+ SparkSession._instantiatedSession = None
+ SparkSession._activeSession = None
+ self._start()
+
+ def reset(self) -> None:
+ for key, value in SPARK_DEFAULTS.items():
+ self.spark.conf.set(key, value)
+ self.time_zone = SPARK_DEFAULTS["spark.sql.session.timeZone"]
+ for table in self.spark.catalog.listTables():
+ if table.isTemporary:
+ self.spark.catalog.dropTempView(table.name)
+
+ def set(self, key: str, value: Optional[str]) -> None:
+ from pyspark.errors import PySparkException
+
+ if value is None:
+ value = SPARK_DEFAULTS[key]
+ elif key == "spark.sql.ansi.enabled":
+ value = value.lower()
+ try:
+ self.spark.conf.set(key, value)
+ except PySparkException as e:
+ raise SparkError("error", _first_line(e))
+ if key == "spark.sql.session.timeZone":
+ self.time_zone = value
+
+ def run(self, sql: str):
+ """Run SQL and return (schema, rows), or raise SparkError."""
+ from py4j.protocol import Py4JJavaError
+ from pyspark.errors import ParseException, PySparkException
+
+ try:
+ df = self.spark.sql(sql)
+ schema = df.schema
+ return schema, _maps_as_entries(df, schema).limit(MAX_ROWS +
1).collect()
+ except TypeError as e:
+ # PySpark returns each MAP as a dict, which cannot have array keys.
+ raise SparkError("convert", f"PySpark cannot convert the result
({e})")
+ except ParseException as e:
+ raise SparkError("parse", f"Spark cannot parse it
({_condition(e)})")
+ except PySparkException as e:
+ if "OutOfMemoryError" in str(e):
+ self.broken = True
+ raise SparkError("memory", "Spark ran out of memory")
+ m = re.search(r"(\w+Type)\.fromInternal is not implemented",
str(e))
+ if m:
+ raise SparkError("convert", f"Spark returned a {m.group(1)}
value")
+ condition = _condition(e)
+ if condition == "UNRESOLVED_ROUTINE":
+ name = _parameter(e, "routineName") or "?"
+ raise SparkError("function", f"Spark has no function {name}")
+ if condition == "TABLE_OR_VIEW_NOT_FOUND":
+ raise SparkError("table", "needs a table that Spark could not
create")
+ message = str(e).strip().split("\n")[0]
+ raise SparkError("error", message)
+ except Py4JJavaError as e:
+ message = str(e.java_exception)
+ if "OutOfMemoryError" in message:
+ self.broken = True
+ raise SparkError("memory", "Spark ran out of memory")
+ raise SparkError("error", message.strip().split("\n")[0])
+
+
+def _maps_as_entries(df, schema):
+ """Collect top-level MAP columns as arrays of entries.
+
+ PySpark returns a MAP as a dict built from a Java HashMap, which loses
+ Spark's entry order and cannot have keys such as arrays.
+ """
+ is_map = [type(f.dataType).__name__ == "MapType" for f in schema.fields]
+ if not any(is_map):
+ return df
+ from pyspark.sql import functions as F
+
+ names = [f"_c{i}" for i in range(len(is_map))]
+ columns = [F.map_entries(n).alias(n) if m else F.col(n) for n, m in
zip(names, is_map)]
+ return df.toDF(*names).select(columns)
+
+
+def _condition(e) -> Optional[str]:
+ getter = getattr(e, "getCondition", None) or getattr(e, "getErrorClass",
None)
+ return getter() if getter else None
+
+
+def _parameter(e, name: str) -> Optional[str]:
+ getter = getattr(e, "getMessageParameters", None)
+ parameters = (getter() if getter else None) or {}
+ return parameters.get(name)
+
+
+@dataclass
+class Outcome:
+ record: Record
+ status: str
+ detail: str = ""
+ spark_sql: Optional[str] = None
+ actual: list[str] = field(default_factory=list)
+
+
+# Statuses that mean DataFusion and Spark disagree, in report order.
+DIVERGENT = {
+ "mismatch": "results differ",
+ "spark-error": "Spark raises an error",
+ "no-spark-error": "DataFusion expects an error, Spark returns a result",
+ "type": "results match, column types differ",
+ "scale": "same numbers, printed with a different scale",
+ "close": "results differ only in floating-point digits",
+ "order": "same rows in a different order",
+ "runner-error": "the runner failed",
+}
+
+
+# Spark's stack() builds all of its rows at once, so stack(2000000000, 1) runs
+# the JVM out of memory before the MAX_ROWS limit applies. A row count above
+# the INT range is a BIGINT, which Spark rejects before running anything.
+_LARGE_STACK_RE = re.compile(r"\bstack\s*\(\s*(\d+)", re.IGNORECASE)
+_MAX_INT = 2**31 - 1
+
+# Statements whose output is a plan or schema in each engine's own format.
+_ENGINE_SPECIFIC_RE = re.compile(r"\s*(EXPLAIN|DESCRIBE|DESC|SHOW)\b",
re.IGNORECASE)
+
+
+def run_record(runner: SparkRunner, record: Record) -> Outcome:
+ if record.skip_reason:
+ return Outcome(record, "skip", record.skip_reason)
+ setting = parse_setting(record.sql) if record.kind == "statement" else None
+ if setting is not None:
+ key, value = setting
+ if key not in SETTINGS:
+ return Outcome(record, "skip", f"no Spark equivalent for {key}")
+ try:
+ runner.set(SETTINGS[key], value)
+ except SparkError as e:
+ if record.expect_error:
+ return Outcome(record, "pass")
+ return Outcome(record, "spark-error", e.message)
+ if record.expect_error:
+ return Outcome(record, "no-spark-error")
+ return Outcome(record, "pass")
+ m = _ENGINE_SPECIFIC_RE.match(record.sql)
+ if m:
+ return Outcome(record, "skip", f"{m.group(1).upper()} output differs
between engines")
+ try:
+ if record.kind == "statement":
+ spark_sql = translate_statement(record.sql, runner.time_zone)
+ else:
+ spark_sql = translate(record.sql, runner.time_zone)
+ except Skip as e:
+ return Outcome(record, "skip", str(e))
+ if any("<slt:ignore>" in line for line in record.expected):
+ return Outcome(record, "skip", "uses <slt:ignore>")
+ m = _LARGE_STACK_RE.search(spark_sql)
+ if m and MAX_ROWS < int(m.group(1)) <= _MAX_INT:
+ return Outcome(record, "skip", f"stack() of more than {MAX_ROWS} rows")
+
+ try:
+ schema, rows = runner.run(spark_sql)
+ except SparkError as e:
+ if e.kind != "error":
+ return Outcome(record, "skip", e.message, spark_sql)
+ if record.expect_error:
+ return Outcome(record, "pass", spark_sql=spark_sql)
+ return Outcome(record, "spark-error", e.message, spark_sql)
Review Comment:
For a record that expects an error, any Spark error counts as `pass`,
whatever the error is. The message in `statement error <message>` is never
read. So the `pass` count includes records where Spark fails for a different
reason than DataFusion does. Could the README say so next to the outcome table?
A one-sentence note is enough, so readers don't treat `pass` as "same error".
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
