[WIP][SPARK-57462][PYTHON][SQL] Add Arrow support for nanosecond-precision timestamp types in PySpark - #58694
Open
stevomitric wants to merge 4 commits into
Open
Conversation
… timestamp types in PySpark
### What changes were proposed in this pull request?
Follow-up to SPARK-57462, which added the PySpark `TimestampNTZNanosType` / `TimestampLTZNanosType`
classes and the classic (Py4J) value path but deliberately deferred the Arrow / pandas value path:
`to_arrow_type` rejected these types, and `DataFrame.toPandas`, `SparkSession.createDataFrame` from a
pandas `DataFrame`, and the whole Spark Connect data path raised
`UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION`.
This change wires the two types through the Python Arrow path. The JVM side already supports them
(`ArrowUtils` maps them to an Arrow `Timestamp(NANOSECOND)` field plus a
`SPARK::timestampNanos::precision` metadata tag, with native `ArrowWriter` writers and
`ArrowColumnVector` accessors), mirroring the shipped `TimeType`, so the change is Python-only:
- `python/pyspark/sql/pandas/types.py`: `to_arrow_type` maps the types to `pa.timestamp("ns", tz=...)`
(LTZ carries the session time zone, NTZ does not); the `ns` -> `us` truncation in
`_check_arrow_array_timestamps_localize` is skipped for a nanosecond target type;
`_to_corrected_pandas_type` maps them to `datetime64[ns]`; and the read / write pandas converters
handle them exactly like `TimestampType` (LTZ, time-zone localized) and `TimestampNTZType` (NTZ).
`from_arrow_type` is intentionally left unchanged, so a plain Arrow `timestamp[ns]` still infers
microsecond `TimestampType` -- type inference is unchanged.
- `python/pyspark/sql/conversion.py`: `LocalDataToArrowConversion` and `ArrowTableToRowsConversion`
(the Spark Connect createDataFrame / collect paths) handle the two types like the existing
microsecond timestamp types.
- Removes the deterministic-rejection guards added as placeholders in SPARK-57462 (in
`pandas/conversion.py`, `connect/session.py`, `connect/dataframe.py`, and the two converter
factories) and the now-unused `_first_timestamp_nanos_type` helper.
The values are carried as an Arrow `timestamp[ns]` (a 64-bit count of nanoseconds since the epoch),
so full nanosecond precision is preserved (pandas `datetime64[ns]`); values outside the
`datetime64[ns]` range (roughly the years 1677 to 2262) are out of range for this path, consistent
with the JVM `ArrowWriter`, which throws on overflow.
### Why are the changes needed?
Without this, a nanosecond-typed column cannot be read into pandas (`toPandas`) or written from a
pandas `DataFrame`, and -- because the Spark Connect data path is entirely Arrow-based -- cannot be
collected at all over Spark Connect. This is the missing piece that makes the nanosecond timestamp
types usable from Spark Connect clients.
### Does this PR introduce _any_ user-facing change?
Yes, when the `spark.sql.timestampNanosTypes.enabled` preview flag is on (it is off by default).
`DataFrame.toPandas`, `SparkSession.createDataFrame` from a pandas `DataFrame`, and the Spark Connect
data path now accept `TimestampNTZNanosType` / `TimestampLTZNanosType` and preserve full nanosecond
precision, where they previously raised `UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION`. There is no
behavior change for any other type, or when the flag is off.
### How was this patch tested?
- Rewrote the classic `test_timestamp_nanos_type_arrow_conversion_unsupported` test into a positive
round-trip `test_timestamp_nanos_type_arrow_conversion` that asserts a value with nine
fractional-second digits survives `toPandas` (pandas `Timestamp.nanosecond == 789`) and a
pandas -> `createDataFrame` -> Spark -> pandas round-trip, for both the NTZ and LTZ types.
- Un-skipped the now-supported nanosecond data-path tests in the Spark Connect parity suite
(`connect/test_parity_types.py`) and replaced the "unsupported" Connect test with a positive
`test_timestamp_nanos_type_connect_data_path`; the classic-only Py4J-UDF and map-key tests remain
skipped there.
- Ran the new test against a live `SparkSession` and confirmed it passes, and confirmed the existing
`TimestampType` / `TimestampNTZType` / `TimeType` Arrow round-trips are unchanged.
### Was this patch authored or co-authored using generative AI tooling?
Yes, co-authored using Claude (Opus).
Co-authored-by: Isaac <no-reply@databricks.com>
…row support Follow-up review fixes on top of the previous commit: - Localize a naive Arrow ``timestamp[ns]`` to the session time zone when the target type is ``TimestampLTZNanosType`` (not only ``TimestampType``) in ``_check_arrow_array_timestamps_localize``. Without this, ``createDataFrame`` from a ``pyarrow.Table`` under a non-UTC ``spark.sql.session.timeZone`` read the value as UTC, so the instant was silently off by the session offset. - ``ArrowTableToRowsConversion``: reduce a nanosecond ``pandas.Timestamp`` (which ``pyarrow.Array.to_pylist`` yields for a ``timestamp[ns]`` column) to a microsecond ``datetime.datetime``, so ``DataFrame.collect()`` over Spark Connect matches the classic ``collect()`` boundary. ``toPandas`` remains the lossless nanosecond path. - Preserve nanoseconds for a ``TimestampLTZNanosType`` nested in an Array/Map/Struct on the ``createDataFrame``-from-pandas path; the nested converter previously truncated to microseconds via ``to_pydatetime()``, inconsistent with the nested NTZ path. - Join the ``isinstance()`` tuples in ``conversion.py`` so ``ruff format`` leaves them unchanged (they were hand-split and would have failed the format check), and shorten the over-long skip messages in the Spark Connect parity suite. - Qualify the ``TimestampNTZNanosType`` docstring: full nanosecond precision holds on the Arrow path; with Arrow disabled, ``toPandas`` falls back to ``collect()`` and truncates to microseconds. - Tests: add ``test_timestamp_nanos_type_arrow_conversion_non_utc`` covering ``createDataFrame`` from a ``pyarrow.Table`` under a non-UTC session time zone for both the NTZ and LTZ types. Co-authored-by: Isaac <no-reply@databricks.com>
…fields + lint Two CI fixes on top of the previous commit: - Tag the Arrow field with the nanosecond precision (``SPARK::timestampNanos::precision``, matching the JVM ``ArrowUtils.timestampNanosPrecisionKey``) whenever ``to_arrow_type`` / ``to_arrow_schema`` builds a field for ``TimestampNTZNanosType`` / ``TimestampLTZNanosType``, at every position (top-level, struct field, array element, map key, map value). Without the tag the JVM ``fromArrowField`` reconstructs the maximum precision (9), so ``createDataFrame`` over Spark Connect with an explicit precision-7 or -8 schema failed with ``INVALID_COLUMN_OR_FIELD_DATA_TYPE`` (``TIMESTAMP_LTZ(9)`` vs the required ``TIMESTAMP_LTZ(7)``); Connect reconstructs the schema from the Arrow field, unlike the classic path which uses the passed schema. Non-nanosecond fields are untagged and unchanged. - Collapse two ``assertEqual`` calls in the Connect parity suite onto single lines so ``ruff format`` (dev/lint-python) leaves them unchanged. Co-authored-by: Isaac <no-reply@databricks.com>
…, guard nanos map keys over Connect Two findings from a spark-dev auto-review pass: - Remove the redundant function-local `from pyspark.sql.types import AnyTimestampNanoType` in `_with_timestamp_nanos_precision`; `AnyTimestampNanoType` is already imported at module scope. - Mirror the classic `collect()` nanosecond-map-key guard onto the Spark Connect `DataFrame`. A nanosecond timestamp used as a map key collapses to a single microsecond `datetime.datetime` when Arrow rows are turned into Python dicts, dropping keys that differ only below a microsecond. `ArrowTableToRowsConversion` (the Connect collect path) did this silently, whereas the classic path raises `TIMESTAMP_NANOS_PYTHON_MAP_KEY`. Connect `collect()` / `toLocalIterator()` now apply the same guard (via `_first_timestamp_nanos_map_key_type`), failing deterministically instead of silently dropping entries, and `test_timestamp_nanos_type_map_key_collision` is un-skipped in the Connect parity suite. Co-authored-by: Isaac <no-reply@databricks.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
Follow-up to SPARK-57462, which added the PySpark
TimestampNTZNanosType/TimestampLTZNanosTypeclasses and the classic (Py4J) value path but deliberately deferred the Arrow / pandas value path:to_arrow_typerejected these types, andDataFrame.toPandas,SparkSession.createDataFramefrom a pandasDataFrame, and the whole Spark Connect data path raisedUNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION.This change wires the two types through the Python Arrow path. The JVM side already supports them (
ArrowUtilsmaps them to an ArrowTimestamp(NANOSECOND)field plus aSPARK::timestampNanos::precisionmetadata tag, with nativeArrowWriterwriters andArrowColumnVectoraccessors).The values are carried as an Arrow
timestamp[ns](a 64-bit count of nanoseconds since the epoch), so full nanosecond precision is preserved (pandasdatetime64[ns]); values outside thedatetime64[ns]range are out of range for this path, consistent with the JVMArrowWriter, which throws on overflow.Why are the changes needed?
Without this, a nanosecond-typed column cannot be read into pandas (
toPandas) or written from a pandasDataFrame, and because the Spark Connect data path is entirely Arrow-based, cannot be collected at all over Spark Connect.Does this PR introduce any user-facing change?
Yes, when the
spark.sql.timestampNanosTypes.enabledpreview flag is on (it is off by default).DataFrame.toPandas,SparkSession.createDataFramefrom a pandasDataFrame, and the Spark Connect data path now acceptTimestampNTZNanosType/TimestampLTZNanosTypeand preserve full nanosecond precision, where they previously raisedUNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION.How was this patch tested?
test_timestamp_nanos_type_arrow_conversion_unsupportedtest into a positive round-triptest_timestamp_nanos_type_arrow_conversionthat asserts a value with nine fractional-second digits survivestoPandas(pandasTimestamp.nanosecond == 789) and a pandas ->createDataFrame-> Spark -> pandas round-trip, for both the NTZ and LTZ types.connect/test_parity_types.py) and replaced the "unsupported" Connect test with a positivetest_timestamp_nanos_type_connect_data_path; the classic-only Py4J-UDF and map-key tests remain skipped there.SparkSessionand confirmed it passes, and confirmed the existingTimestampType/TimestampNTZType/TimeTypeArrow round-trips are unchanged.Was this patch authored or co-authored using generative AI tooling?
Co-Authored-By: Claude Opus 4.8