From 93347e9c36e7c525b74f02c84e36b9377baf6e78 Mon Sep 17 00:00:00 2001 From: Stevo Mitric Date: Wed, 9 Sep 2026 12:14:22 +0000 Subject: [PATCH 1/4] [SPARK-XXXXX][PYTHON][SQL] Add Arrow support for nanosecond-precision 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 --- python/pyspark/sql/connect/dataframe.py | 9 --- python/pyspark/sql/connect/session.py | 9 --- python/pyspark/sql/conversion.py | 53 ++++----------- python/pyspark/sql/pandas/conversion.py | 21 +----- python/pyspark/sql/pandas/types.py | 60 ++++++++++------- .../sql/tests/connect/test_parity_types.py | 54 +++++++-------- python/pyspark/sql/tests/test_types.py | 67 +++++++++++-------- python/pyspark/sql/types.py | 38 ++--------- 8 files changed, 115 insertions(+), 196 deletions(-) diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 593a700f82acf..95d6cddd50ba6 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -1999,15 +1999,6 @@ def toPandas(self) -> "PandasDataFrameLike": return self._to_pandas() def _to_pandas(self, **kwargs: Any) -> "PandasDataFrameLike": - # SPARK-57462: the Arrow-based nanosecond timestamp value path is a pending follow-up. - # Connect overrides _to_pandas and goes straight to client.to_pandas, so the guard on the - # classic PandasConversionMixin is not reached; reject here too so the behavior is - # deterministic (and consistent with classic toPandas and to_arrow_type) rather than - # emitting unhandled pandas nanosecond / wrong-time-zone values. - from pyspark.sql.pandas.types import _reject_timestamp_nanos_conversion - - _reject_timestamp_nanos_conversion(self.schema) - query = self._plan.to_proto(self._session.client) pdf, ei = self._session.client.to_pandas(query, self._plan.observations, **kwargs) self._execution_info = ei diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index a518fdd327eac..bd6f32ab5ac05 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -517,15 +517,6 @@ def createDataFrame( }, ) - # SPARK-57462: building a DataFrame over Spark Connect goes through Arrow, whose - # nanosecond timestamp value conversion is a pending follow-up. Reject an explicit - # nanosecond-typed schema here, right after resolution, so the empty-input and NumPy fast - # paths (which never build an Arrow converter) fail deterministically too. - if _schema is not None: - from pyspark.sql.pandas.types import _reject_timestamp_nanos_conversion - - _reject_timestamp_nanos_conversion(_schema) - if isinstance(data, np.ndarray) and data.ndim not in [1, 2]: raise PySparkValueError( errorClass="INVALID_NDARRAY_DIMENSION", diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index fc03acf026436..25ce2814ae67a 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -56,6 +56,8 @@ StringType, StructField, StructType, + TimestampLTZNanosType, + TimestampNTZNanosType, TimestampNTZType, TimestampType, TimeType, @@ -540,14 +542,11 @@ def _need_converter( return True elif isinstance(dataType, BinaryType): return True - elif isinstance(dataType, (TimestampType, TimestampNTZType)): + elif isinstance( + dataType, (TimestampType, TimestampNTZType, AnyTimestampNanoType) + ): # Always truncate return True - elif isinstance(dataType, AnyTimestampNanoType): - # Needs a converter so _create_converter is built (and eagerly rejects) for every - # direct caller -- Arrow UDF return values, Python data-source writes -- not only the - # LocalDataToArrowConversion.convert path. - return True elif isinstance(dataType, DecimalType): # Convert Decimal('NaN') to None # Rescale Decimal values @@ -601,18 +600,6 @@ def _create_converter( else: return lambda value: value - if isinstance(dataType, AnyTimestampNanoType): - # SPARK-57462: the Arrow-based value path for the nanosecond timestamp types is a - # pending follow-up. Reject eagerly, when the converter is built, so building a - # DataFrame from Arrow / returning nanoseconds from an Arrow UDF fails deterministically - # rather than mis-encoding the value. Consistent with to_arrow_type. - from pyspark.errors import PySparkTypeError - - raise PySparkTypeError( - errorClass="UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION", - messageParameters={"data_type": str(dataType)}, - ) - if isinstance(dataType, NullType): def convert_null(value: Any) -> Any: @@ -793,7 +780,7 @@ def convert_binary(value: Any) -> Any: return convert_binary - elif isinstance(dataType, TimestampType): + elif isinstance(dataType, (TimestampType, TimestampLTZNanosType)): def convert_timestamp(value: Any) -> Any: if value is None: @@ -806,7 +793,7 @@ def convert_timestamp(value: Any) -> Any: return convert_timestamp - elif isinstance(dataType, TimestampNTZType): + elif isinstance(dataType, (TimestampNTZType, TimestampNTZNanosType)): def convert_timestamp_ntz(value: Any) -> Any: if value is None: @@ -1192,14 +1179,11 @@ def _need_converter(dataType: DataType) -> bool: return True elif isinstance(dataType, BinaryType): return True - elif isinstance(dataType, (TimestampType, TimestampNTZType)): + elif isinstance( + dataType, (TimestampType, TimestampNTZType, AnyTimestampNanoType) + ): # Always remove the time zone info for now return True - elif isinstance(dataType, AnyTimestampNanoType): - # Needs a converter so _create_converter is built (and eagerly rejects) for every - # direct caller -- Connect collect, batched Arrow UDF inputs, foreachPartition, and - # Python data-source reads -- not only the ArrowTableToRowsConversion.convert path. - return True elif isinstance(dataType, UserDefinedType): return True elif isinstance(dataType, VariantType): @@ -1235,19 +1219,6 @@ def _create_converter( else: return lambda value: value - if isinstance(dataType, AnyTimestampNanoType): - # SPARK-57462: the Arrow-based value path for the nanosecond timestamp types is a - # pending follow-up. Reject eagerly, when the converter is built (all callers build - # converters up front), so it is not data-dependent and cannot leak a raw - # Arrow-derived value (a nanosecond-precision, possibly timezone-aware - # pandas.Timestamp). Consistent with to_arrow_type, which already rejects these types. - from pyspark.errors import PySparkTypeError - - raise PySparkTypeError( - errorClass="UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION", - messageParameters={"data_type": str(dataType)}, - ) - if isinstance(dataType, NullType): return lambda value: None @@ -1361,7 +1332,7 @@ def convert_binary(value: Any) -> Any: return convert_binary - elif isinstance(dataType, TimestampType): + elif isinstance(dataType, (TimestampType, TimestampLTZNanosType)): def convert_timestamp(value: Any) -> Any: if value is None: @@ -1372,7 +1343,7 @@ def convert_timestamp(value: Any) -> Any: return convert_timestamp - elif isinstance(dataType, TimestampNTZType): + elif isinstance(dataType, (TimestampNTZType, TimestampNTZNanosType)): def convert_timestamp_ntz(value: Any) -> Any: if value is None: diff --git a/python/pyspark/sql/pandas/conversion.py b/python/pyspark/sql/pandas/conversion.py index cf64d53943239..3bdaa67b91e64 100644 --- a/python/pyspark/sql/pandas/conversion.py +++ b/python/pyspark/sql/pandas/conversion.py @@ -320,18 +320,11 @@ def _to_pandas(self, **kwargs: Any) -> "PandasDataFrameLike": assert isinstance(self, DataFrame) - from pyspark.sql.pandas.types import ( - _create_converter_to_pandas, - _reject_timestamp_nanos_conversion, - ) + from pyspark.sql.pandas.types import _create_converter_to_pandas from pyspark.sql.pandas.utils import require_minimum_pandas_version require_minimum_pandas_version() - # Arrow/pandas value conversion for the nanosecond timestamp types is a pending follow-up; - # fail deterministically here rather than fall back to a lossy / wrong-timezone result. - _reject_timestamp_nanos_conversion(self.schema) - import pandas as pd ( @@ -639,18 +632,6 @@ def createDataFrame( # type: ignore[misc] selfcheck = arrowSafeTypeConversion == "true" infer_pandas_dict_as_map = inferPandasDictAsMap == "true" - # Building a DataFrame from pandas/PyArrow data goes through Arrow, whose nanosecond - # timestamp value conversion is a pending follow-up; fail deterministically for an explicit - # nanosecond-typed schema rather than silently mis-handle the values. createDataFrame has - # already parsed a DDL-string schema into a DataType by this point, so this guard covers - # both an explicit DataType and a DDL string (bare atomic or StructType). Only a plain list - # of column names carries no type information -- and so cannot name these types -- and is - # left to the server to gate. - if isinstance(schema, DataType): - from pyspark.sql.pandas.types import _reject_timestamp_nanos_conversion - - _reject_timestamp_nanos_conversion(schema) - if type(data).__name__ == "Table": # `data` is a PyArrow Table from pyspark.sql.pandas.utils import require_minimum_pyarrow_version diff --git a/python/pyspark/sql/pandas/types.py b/python/pyspark/sql/pandas/types.py index bdef2b805ed7f..7b718ff6865f6 100644 --- a/python/pyspark/sql/pandas/types.py +++ b/python/pyspark/sql/pandas/types.py @@ -30,6 +30,7 @@ from pyspark.errors import PySparkTypeError, PySparkValueError, UnsupportedOperationException from pyspark.loose_version import LooseVersion from pyspark.sql.types import ( + AnyTimestampNanoType, ArrayType, BinaryType, BooleanType, @@ -53,6 +54,8 @@ StringType, StructField, StructType, + TimestampLTZNanosType, + TimestampNTZNanosType, TimestampNTZType, TimestampType, TimeType, @@ -77,26 +80,6 @@ metadata_key = b"SPARK::metadata::json" -def _reject_timestamp_nanos_conversion(schema: DataType) -> None: - """Raise if ``schema`` involves a nanosecond timestamp type, for Arrow/pandas value paths. - - The Arrow / pandas value conversion for :class:`TimestampNTZNanosType` / - :class:`TimestampLTZNanosType` is not implemented yet (planned follow-up). Rather than let these - paths silently mis-handle the value (wrong time zone for LTZ, or a leaked ``pandas.Timestamp``), - fail deterministically here, consistent with :func:`to_arrow_type`, which already rejects these - types with the same error condition and reports the offending leaf type. - """ - from pyspark.errors import PySparkTypeError - from pyspark.sql.types import _first_timestamp_nanos_type - - offending = _first_timestamp_nanos_type(schema) - if offending is not None: - raise PySparkTypeError( - errorClass="UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION", - messageParameters={"data_type": str(offending)}, - ) - - def to_arrow_metadata(metadata: Optional[Dict[str, Any]] = None) -> Optional[Dict[bytes, bytes]]: if metadata is not None and len(metadata) > 0: return {metadata_key: json.dumps(metadata).encode("utf-8")} @@ -164,6 +147,15 @@ def to_arrow_type( arrow_type = pa.timestamp("us", tz=timezone) elif isinstance(dt, TimestampNTZType): arrow_type = pa.timestamp("us", tz=None) + elif isinstance(dt, TimestampLTZNanosType): + # Nanosecond-precision timestamps interchange as Arrow Timestamp(NANOSECOND), matching the + # JVM ArrowUtils mapping. This is the LTZ (timezone-aware) variant, so the session timezone + # is attached exactly like TimestampType. The precision (7-9) is carried out of band by the + # Spark schema; the Arrow unit is always nanoseconds. + assert timezone is not None + arrow_type = pa.timestamp("ns", tz=timezone) + elif isinstance(dt, TimestampNTZNanosType): + arrow_type = pa.timestamp("ns", tz=None) elif isinstance(dt, DayTimeIntervalType): arrow_type = pa.duration("us") elif isinstance(dt, TimeType): @@ -570,7 +562,14 @@ def _check_arrow_array_timestamps_localize( ] ) - if types.is_timestamp(a.type) and truncate and a.type.unit == "ns": + if ( + types.is_timestamp(a.type) + and truncate + and a.type.unit == "ns" + and not isinstance(dt, AnyTimestampNanoType) + ): + # Nanosecond timestamps are floored to microseconds for the microsecond-precision Spark + # types, but a nanosecond-precision target type keeps its full resolution. a = pc.floor_temporal(a, unit="microsecond") if types.is_timestamp(a.type) and a.type.tz is None and isinstance(dt, TimestampType): @@ -919,6 +918,10 @@ def _to_corrected_pandas_type(dt: DataType) -> Optional[Any]: return np.dtype("datetime64[ns]") else: return np.dtype("datetime64[us]") + elif isinstance(dt, AnyTimestampNanoType): + # Nanosecond-precision timestamps always map to datetime64[ns] regardless of pandas + # version -- it is the only pandas resolution that preserves the sub-microsecond digits. + return np.dtype("datetime64[ns]") elif isinstance(dt, DayTimeIntervalType): if LooseVersion(pd.__version__) < "3.0.0": return np.dtype("timedelta64[ns]") @@ -1034,7 +1037,10 @@ def correct_dtype(pser: pd.Series) -> pd.Series: else: return pser.astype(pandas_type, copy=False) - elif isinstance(data_type, TimestampType): + elif isinstance(data_type, (TimestampType, TimestampLTZNanosType)): + # TimestampLTZNanosType is the nanosecond-precision, timezone-aware timestamp; it is + # localized to the session timezone exactly like TimestampType. Its pandas_type is + # datetime64[ns] (see _to_corrected_pandas_type), so full resolution is preserved. assert timezone is not None def correct_dtype(pser: pd.Series) -> pd.Series: @@ -1220,7 +1226,7 @@ def convert_struct_as_dict(value: Any) -> Any: messageParameters={"var": str(_struct_in_pandas)}, ) - elif isinstance(dt, TimestampType): + elif isinstance(dt, (TimestampType, TimestampLTZNanosType)): assert timezone is not None local_tz: Union[datetime.tzinfo, str] = ( @@ -1236,7 +1242,7 @@ def convert_timestamp(value: Any) -> Any: return convert_timestamp - elif isinstance(dt, TimestampNTZType): + elif isinstance(dt, (TimestampNTZType, TimestampNTZNanosType)): def convert_timestamp_ntz(value: Any) -> Any: return pd.Timestamp(value) @@ -1367,7 +1373,9 @@ def _create_converter_from_pandas( """ import pandas as pd - if isinstance(data_type, TimestampType): + if isinstance(data_type, (TimestampType, TimestampLTZNanosType)): + # TimestampLTZNanosType is timezone-aware like TimestampType; the vectorized internal + # conversion preserves the datetime64[ns] resolution, so nanoseconds survive to Arrow. assert timezone is not None def correct_timestamp(pser: pd.Series) -> pd.Series: @@ -1566,7 +1574,7 @@ def convert_struct(value: Any) -> Any: return convert_struct - elif isinstance(dt, TimestampType): + elif isinstance(dt, (TimestampType, TimestampLTZNanosType)): assert timezone is not None def convert_timestamp(value: Any) -> Any: diff --git a/python/pyspark/sql/tests/connect/test_parity_types.py b/python/pyspark/sql/tests/connect/test_parity_types.py index 57ff0c7ba87b2..f6ac9c24664ee 100644 --- a/python/pyspark/sql/tests/connect/test_parity_types.py +++ b/python/pyspark/sql/tests/connect/test_parity_types.py @@ -15,6 +15,7 @@ # limitations under the License. # +import datetime import unittest from pyspark.sql.tests.test_types import TypesTestsMixin @@ -22,53 +23,44 @@ class TypesParityTests(TypesTestsMixin, ReusedConnectTestCase): - # SPARK-57462: nanosecond timestamp types are not yet supported over Spark Connect, whose - # data path goes through Arrow (to_arrow_type / ArrowTableToRowsConversion). These inherited - # tests build or collect nanosecond data and are covered by the classic (non-Connect) suite; - # pending the Arrow follow-up they are skipped here. - @unittest.skip("SPARK-57462: nanosecond timestamp types are pending Connect/Arrow support.") - def test_timestamp_nanos_type(self): - super().test_timestamp_nanos_type() - - @unittest.skip("SPARK-57462: nanosecond timestamp types are pending Connect/Arrow support.") - def test_timestamp_nanos_type_preview_flag_off(self): - super().test_timestamp_nanos_type_preview_flag_off() - - @unittest.skip("SPARK-57462: nanosecond timestamp types are pending Connect/Arrow support.") + # SPARK-57462 follow-up: the nanosecond timestamp value path now works over Spark Connect, + # whose data path goes through Arrow (to_arrow_type / ArrowTableToRowsConversion). The + # data-path tests -- test_timestamp_nanos_type, test_timestamp_nanos_type_preview_flag_off and + # test_timestamp_nanos_type_arrow_conversion -- are therefore inherited and run here. The + # nanosecond tests below exercise classic-only mechanisms that do not apply to Connect: the + # Py4J (useArrow=False) UDF path and the classic collect map-key guard in classic/dataframe.py. + @unittest.skip("SPARK-57462: uses the classic Py4J (useArrow=False) UDF path, not Connect.") def test_timestamp_nanos_type_python_udf(self): super().test_timestamp_nanos_type_python_udf() - @unittest.skip("SPARK-57462: nanosecond timestamp types are pending Connect/Arrow support.") + @unittest.skip("SPARK-57462: the collect map-key guard is classic-only (classic/dataframe.py).") def test_timestamp_nanos_type_map_key_collision(self): super().test_timestamp_nanos_type_map_key_collision() - @unittest.skip("SPARK-57462: nanosecond timestamp types are pending Connect/Arrow support.") + @unittest.skip("SPARK-57462: uses the classic Py4J (useArrow=False) UDF input path, not Connect.") def test_timestamp_nanos_type_python_udf_input(self): super().test_timestamp_nanos_type_python_udf_input() - @unittest.skip("SPARK-57462: nanosecond timestamp types are pending Connect/Arrow support.") + @unittest.skip("SPARK-57462: uses the classic Py4J (useArrow=False) UDF input path, not Connect.") def test_timestamp_nanos_type_map_key_python_udf_input(self): super().test_timestamp_nanos_type_map_key_python_udf_input() - # The classic Arrow-rejection test asserts the classic PySparkTypeError / message parameters; - # Connect's data path rejects at a different layer, so it is covered by the Connect-specific - # test below rather than by inheriting the classic one. - @unittest.skip("SPARK-57462: nanosecond timestamp types are pending Connect/Arrow support.") - def test_timestamp_nanos_type_arrow_conversion_unsupported(self): - super().test_timestamp_nanos_type_arrow_conversion_unsupported() - - def test_timestamp_nanos_type_connect_data_path_unsupported(self): - # SPARK-57462: the Spark Connect data path goes through Arrow (to_arrow_type / - # ArrowTableToRowsConversion), so collecting a nanosecond value must raise the documented - # UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION rather than mis-handle it. Connect collect is - # otherwise safe only because to_table supplies the Spark schema; this locks that in. + def test_timestamp_nanos_type_connect_data_path(self): + # SPARK-57462 follow-up: the Spark Connect data path goes through Arrow, so collecting a + # nanosecond value now succeeds. collect() yields microsecond-resolution datetime.datetime + # (the Python boundary), while toPandas keeps full nanosecond precision (datetime64[ns]). + import pandas as pd + with self.sql_conf({"spark.sql.timestampNanosTypes.enabled": True}): df = self.spark.sql( "SELECT CAST('2020-01-02 03:04:05.123456789' AS TIMESTAMP_NTZ(9)) AS ts" ) - with self.assertRaises(Exception) as pe: - df.collect() - self.assertIn("UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION", str(pe.exception)) + self.assertEqual( + datetime.datetime(2020, 1, 2, 3, 4, 5, 123456), df.collect()[0].ts + ) + self.assertEqual( + pd.Timestamp("2020-01-02 03:04:05.123456789"), df.toPandas()["ts"][0] + ) @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") def test_apply_schema(self): diff --git a/python/pyspark/sql/tests/test_types.py b/python/pyspark/sql/tests/test_types.py index 1f053f0fa4826..7128f9f255f6e 100644 --- a/python/pyspark/sql/tests/test_types.py +++ b/python/pyspark/sql/tests/test_types.py @@ -2354,36 +2354,47 @@ def test_timestamp_nanos_type_map_key_python_udf_input(self): # SparkRuntimeException); assert that condition rather than any failure. self.assertIn("TIMESTAMP_NANOS_PYTHON_MAP_KEY", str(pe.exception)) - def test_timestamp_nanos_type_arrow_conversion_unsupported(self): - # SPARK-57462: Arrow/pandas value conversion for the nanosecond timestamp types is a pending - # follow-up; until then the classic read (toPandas) and write (createDataFrame from a pandas - # DataFrame) paths must reject an explicit nanosecond schema deterministically, naming the - # offending leaf type, rather than silently mis-handle the value. (The Connect data path is - # asserted separately in the parity suite.) + def test_timestamp_nanos_type_arrow_conversion(self): + # SPARK-57462 follow-up: the Arrow / pandas value path carries the nanosecond timestamp + # types as an Arrow timestamp[ns], so -- unlike the microsecond-resolution + # datetime.datetime boundary used by collect() / Python UDFs -- DataFrame.toPandas and + # createDataFrame from a pandas DataFrame preserve full nanosecond precision (pandas + # datetime64[ns]). The session time zone is pinned so the timezone-aware LTZ value is + # deterministic. import pandas as pd - with self.sql_conf({"spark.sql.timestampNanosTypes.enabled": True}): - schema = StructType([StructField("ts", TimestampNTZNanosType(9))]) - value = datetime.datetime(2020, 1, 2, 3, 4, 5, 123456) - - # Read path: DataFrame.toPandas(). - df = self.spark.createDataFrame([(value,)], schema) - with self.assertRaises(PySparkTypeError) as pe: - df.toPandas() - self.check_error( - exception=pe.exception, - errorClass="UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION", - messageParameters={"data_type": "TimestampNTZNanosType(9)"}, - ) - - # Write path: createDataFrame from a pandas DataFrame with an explicit nanos schema. - with self.assertRaises(PySparkTypeError) as pe: - self.spark.createDataFrame(pd.DataFrame({"ts": [value]}), schema) - self.check_error( - exception=pe.exception, - errorClass="UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION", - messageParameters={"data_type": "TimestampNTZNanosType(9)"}, - ) + with self.sql_conf( + { + "spark.sql.timestampNanosTypes.enabled": True, + "spark.sql.session.timeZone": "UTC", + "spark.sql.execution.arrow.pyspark.enabled": True, + } + ): + # Read path: toPandas keeps the sub-microsecond digits. + pdf = self.spark.sql( + "SELECT CAST('2020-01-02 03:04:05.123456789' AS TIMESTAMP_NTZ(9)) AS ts" + ).toPandas() + self.assertEqual("datetime64[ns]", str(pdf["ts"].dtype)) + self.assertEqual(pd.Timestamp("2020-01-02 03:04:05.123456789"), pdf["ts"][0]) + # The sub-microsecond digits survive; datetime.datetime could not carry them. + self.assertEqual(789, pdf["ts"][0].nanosecond) + + # Write path: a datetime64[ns] pandas column round-trips its nanoseconds back to Spark, + # for both the NTZ and the timezone-aware LTZ nanosecond types. + ns_string = "2020-01-02 03:04:05.123456789" + in_pdf = pd.DataFrame({"ts": pd.to_datetime(pd.Series([ns_string]))}) + self.assertEqual("datetime64[ns]", str(in_pdf["ts"].dtype)) + for nanos_type in (TimestampNTZNanosType(9), TimestampLTZNanosType(9)): + schema = StructType([StructField("ts", nanos_type)]) + df = self.spark.createDataFrame(in_pdf, schema) + self.assertEqual(schema, df.schema) + # The stored value keeps all nine fractional digits (checked server-side). + self.assertEqual( + ns_string, + df.select(F.col("ts").cast("string")).first()[0], + ) + # ... and a full pandas -> Spark -> pandas round-trip is lossless. + self.assertEqual(pd.Timestamp(ns_string), df.toPandas()["ts"][0]) def test_yearmonth_interval_type_constructor(self): self.assertEqual(YearMonthIntervalType().simpleString(), "interval year to month") diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py index 27cc977fe44c6..abdd9efd4c0d9 100644 --- a/python/pyspark/sql/types.py +++ b/python/pyspark/sql/types.py @@ -573,10 +573,12 @@ class TimestampNTZNanosType(AnyTimestampNanoType): with keys of this type that differ only below a microsecond would collapse to one entry, so that conversion raises rather than silently dropping an entry. - Arrow- and pandas-based conversion for these types -- :meth:`DataFrame.toPandas`, - :meth:`SparkSession.createDataFrame` from a pandas ``DataFrame``, Arrow-based UDFs, and the - Spark Connect data path -- is not yet supported and raises - ``UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION``; it is planned as a follow-up. + Arrow- and pandas-based conversion -- :meth:`DataFrame.toPandas`, + :meth:`SparkSession.createDataFrame` from a pandas ``DataFrame``, and the Spark Connect data + path -- carries the value as an Arrow ``timestamp[ns]`` and preserves full nanosecond + precision (pandas ``datetime64[ns]``). Because that Arrow encoding counts nanoseconds since + the epoch in a 64-bit integer, values outside the ``datetime64[ns]`` range (roughly the years + 1677 to 2262) cannot be carried on this path. .. versionadded:: 4.4.0 """ @@ -3084,34 +3086,6 @@ def _first_timestamp_nanos_map_key_type(dt: DataType) -> Optional["DataType"]: return None -def _first_timestamp_nanos_type(dt: DataType) -> Optional["DataType"]: - """Return the first nanosecond-capable timestamp type (depth-first) that ``dt`` is or contains, - or ``None`` if it contains none. - - The Arrow / pandas / Connect value conversion for :class:`TimestampNTZNanosType` / - :class:`TimestampLTZNanosType` is not implemented yet (planned follow-up). Callers reject such a - schema up front rather than mis-handle the value; the returned leaf type feeds the error - message, consistent with :func:`~pyspark.sql.pandas.types.to_arrow_type`, which reports the - offending leaf. This is the "find the node" companion to the ``_has_type`` boolean check. - """ - if isinstance(dt, AnyTimestampNanoType): - return dt - elif isinstance(dt, ArrayType): - return _first_timestamp_nanos_type(dt.elementType) - elif isinstance(dt, MapType): - return _first_timestamp_nanos_type(dt.keyType) or _first_timestamp_nanos_type(dt.valueType) - elif isinstance(dt, StructType): - for field in dt.fields: - found = _first_timestamp_nanos_type(field.dataType) - if found is not None: - return found - return None - elif isinstance(dt, UserDefinedType): - return _first_timestamp_nanos_type(dt.sqlType()) - else: - return None - - @overload def _merge_type(a: StructType, b: StructType, name: Optional[str] = None) -> StructType: ... From 62a1cd24776b097bd0204acc2827e8dde8f74205 Mon Sep 17 00:00:00 2001 From: Stevo Mitric Date: Thu, 10 Sep 2026 12:23:40 +0000 Subject: [PATCH 2/4] [SPARK-XXXXX][PYTHON][SQL] Address review for nanosecond-precision Arrow 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 --- python/pyspark/sql/conversion.py | 17 +++++++---- python/pyspark/sql/pandas/types.py | 27 ++++++++++++++--- .../sql/tests/connect/test_parity_types.py | 6 ++-- python/pyspark/sql/tests/test_types.py | 30 +++++++++++++++++++ python/pyspark/sql/types.py | 14 +++++---- 5 files changed, 75 insertions(+), 19 deletions(-) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index 25ce2814ae67a..2085a3a57cd31 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -542,9 +542,7 @@ def _need_converter( return True elif isinstance(dataType, BinaryType): return True - elif isinstance( - dataType, (TimestampType, TimestampNTZType, AnyTimestampNanoType) - ): + elif isinstance(dataType, (TimestampType, TimestampNTZType, AnyTimestampNanoType)): # Always truncate return True elif isinstance(dataType, DecimalType): @@ -1179,9 +1177,7 @@ def _need_converter(dataType: DataType) -> bool: return True elif isinstance(dataType, BinaryType): return True - elif isinstance( - dataType, (TimestampType, TimestampNTZType, AnyTimestampNanoType) - ): + elif isinstance(dataType, (TimestampType, TimestampNTZType, AnyTimestampNanoType)): # Always remove the time zone info for now return True elif isinstance(dataType, UserDefinedType): @@ -1339,6 +1335,11 @@ def convert_timestamp(value: Any) -> Any: return None else: assert isinstance(value, datetime.datetime) + # A timestamp[ns] value materializes as a pandas.Timestamp; collect()'s + # datetime.datetime boundary is microsecond resolution (toPandas is the + # lossless path), so drop any sub-microsecond digits to match classic collect(). + if hasattr(value, "to_pydatetime"): + value = value.to_pydatetime(warn=False) return value.astimezone().replace(tzinfo=None) return convert_timestamp @@ -1350,6 +1351,10 @@ def convert_timestamp_ntz(value: Any) -> Any: return None else: assert isinstance(value, datetime.datetime) + # See convert_timestamp: reduce a nanosecond pandas.Timestamp to a microsecond + # datetime.datetime so collect() matches the classic path. + if hasattr(value, "to_pydatetime"): + value = value.to_pydatetime(warn=False) return value return convert_timestamp_ntz diff --git a/python/pyspark/sql/pandas/types.py b/python/pyspark/sql/pandas/types.py index 7b718ff6865f6..8782c83426eb8 100644 --- a/python/pyspark/sql/pandas/types.py +++ b/python/pyspark/sql/pandas/types.py @@ -572,11 +572,15 @@ def _check_arrow_array_timestamps_localize( # types, but a nanosecond-precision target type keeps its full resolution. a = pc.floor_temporal(a, unit="microsecond") - if types.is_timestamp(a.type) and a.type.tz is None and isinstance(dt, TimestampType): + if ( + types.is_timestamp(a.type) + and a.type.tz is None + and isinstance(dt, (TimestampType, TimestampLTZNanosType)) + ): assert timezone is not None - # Only localize timestamps that will become Spark TimestampType columns. - # Do not localize timestamps that will become Spark TimestampNTZType columns. + # Localize naive Arrow timestamps whose target is a Spark local-time-zone timestamp + # (TimestampType or the nanosecond LTZ type); leave NTZ / NTZ-nanos targets naive. return pc.assume_timezone(a, timezone) if types.is_list(a.type): # Return the ListArray as-is if it contains no nested fields or timestamps @@ -1574,7 +1578,22 @@ def convert_struct(value: Any) -> Any: return convert_struct - elif isinstance(dt, (TimestampType, TimestampLTZNanosType)): + elif isinstance(dt, TimestampLTZNanosType): + assert timezone is not None + + def convert_timestamp_ltz_nanos(value: Any) -> Any: + if isinstance(value, datetime.datetime) and value.tzinfo is not None: + ts = pd.Timestamp(value) + else: + ts = pd.Timestamp(value).tz_localize(timezone) + # Keep the tz-aware pandas.Timestamp so a nested LTZ nanosecond value retains its + # sub-microsecond digits through the Arrow build (mirrors the nested NTZ path, + # which is identity). to_pydatetime() would truncate to microseconds. + return ts + + return convert_timestamp_ltz_nanos + + elif isinstance(dt, TimestampType): assert timezone is not None def convert_timestamp(value: Any) -> Any: diff --git a/python/pyspark/sql/tests/connect/test_parity_types.py b/python/pyspark/sql/tests/connect/test_parity_types.py index f6ac9c24664ee..7d202cabff286 100644 --- a/python/pyspark/sql/tests/connect/test_parity_types.py +++ b/python/pyspark/sql/tests/connect/test_parity_types.py @@ -33,15 +33,15 @@ class TypesParityTests(TypesTestsMixin, ReusedConnectTestCase): def test_timestamp_nanos_type_python_udf(self): super().test_timestamp_nanos_type_python_udf() - @unittest.skip("SPARK-57462: the collect map-key guard is classic-only (classic/dataframe.py).") + @unittest.skip("SPARK-57462: classic-only collect map-key guard (classic/dataframe.py).") def test_timestamp_nanos_type_map_key_collision(self): super().test_timestamp_nanos_type_map_key_collision() - @unittest.skip("SPARK-57462: uses the classic Py4J (useArrow=False) UDF input path, not Connect.") + @unittest.skip("SPARK-57462: classic-only Py4J UDF input path (useArrow=False), not Connect.") def test_timestamp_nanos_type_python_udf_input(self): super().test_timestamp_nanos_type_python_udf_input() - @unittest.skip("SPARK-57462: uses the classic Py4J (useArrow=False) UDF input path, not Connect.") + @unittest.skip("SPARK-57462: classic-only Py4J map-key UDF input path, not Connect.") def test_timestamp_nanos_type_map_key_python_udf_input(self): super().test_timestamp_nanos_type_map_key_python_udf_input() diff --git a/python/pyspark/sql/tests/test_types.py b/python/pyspark/sql/tests/test_types.py index 7128f9f255f6e..8f097a9548af8 100644 --- a/python/pyspark/sql/tests/test_types.py +++ b/python/pyspark/sql/tests/test_types.py @@ -2396,6 +2396,36 @@ def test_timestamp_nanos_type_arrow_conversion(self): # ... and a full pandas -> Spark -> pandas round-trip is lossless. self.assertEqual(pd.Timestamp(ns_string), df.toPandas()["ts"][0]) + def test_timestamp_nanos_type_arrow_conversion_non_utc(self): + # SPARK-57462 follow-up: a naive Arrow timestamp[ns] column ingested as the timezone-aware + # LTZ nanosecond type under a non-UTC session time zone must be localized to that zone + # (assume_timezone), exactly like TimestampType -- otherwise it is silently read as UTC and + # the instant is off by the session offset. This exercises the createDataFrame from a + # pyarrow.Table path (createDataFrame from a pandas DataFrame uses a different, already + # covered converter). NTZ is the control: it stays a naive wall clock either way. + import pandas as pd + import pyarrow as pa + + ns_string = "2020-06-15 12:30:00.123456789" + with self.sql_conf( + { + "spark.sql.timestampNanosTypes.enabled": True, + "spark.sql.session.timeZone": "America/New_York", + "spark.sql.execution.arrow.pyspark.enabled": True, + } + ): + table = pa.table({"ts": pa.array([pd.Timestamp(ns_string)], type=pa.timestamp("ns"))}) + for nanos_type in (TimestampNTZNanosType(9), TimestampLTZNanosType(9)): + schema = StructType([StructField("ts", nanos_type)]) + df = self.spark.createDataFrame(table, schema) + self.assertEqual(schema, df.schema) + # Rendered in the session time zone the wall clock is unchanged: NTZ carries no + # zone, and LTZ interpreted the naive input in the session zone (not UTC -- which + # would shift it by the session offset). + self.assertEqual(ns_string, df.select(F.col("ts").cast("string")).first()[0]) + # Nanoseconds survive the Arrow round-trip. + self.assertEqual(789, df.toPandas()["ts"][0].nanosecond) + def test_yearmonth_interval_type_constructor(self): self.assertEqual(YearMonthIntervalType().simpleString(), "interval year to month") self.assertEqual( diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py index abdd9efd4c0d9..c8458577ebaf6 100644 --- a/python/pyspark/sql/types.py +++ b/python/pyspark/sql/types.py @@ -573,12 +573,14 @@ class TimestampNTZNanosType(AnyTimestampNanoType): with keys of this type that differ only below a microsecond would collapse to one entry, so that conversion raises rather than silently dropping an entry. - Arrow- and pandas-based conversion -- :meth:`DataFrame.toPandas`, - :meth:`SparkSession.createDataFrame` from a pandas ``DataFrame``, and the Spark Connect data - path -- carries the value as an Arrow ``timestamp[ns]`` and preserves full nanosecond - precision (pandas ``datetime64[ns]``). Because that Arrow encoding counts nanoseconds since - the epoch in a 64-bit integer, values outside the ``datetime64[ns]`` range (roughly the years - 1677 to 2262) cannot be carried on this path. + Arrow-based conversion -- :meth:`DataFrame.toPandas` and + :meth:`SparkSession.createDataFrame` from a pandas ``DataFrame`` with Arrow enabled + (``spark.sql.execution.arrow.pyspark.enabled``), including the Spark Connect data path -- + carries the value as an Arrow ``timestamp[ns]`` and preserves full nanosecond precision + (pandas ``datetime64[ns]``). Because that Arrow encoding counts nanoseconds since the epoch in + a 64-bit integer, values outside the ``datetime64[ns]`` range (roughly the years 1677 to 2262) + cannot be carried on this path. With Arrow disabled, :meth:`DataFrame.toPandas` falls back to + :meth:`DataFrame.collect` and, like it, truncates to microseconds. .. versionadded:: 4.4.0 """ From 9438c4df027f9c9051477d3d568ad3da9d5ff3bd Mon Sep 17 00:00:00 2001 From: Stevo Mitric Date: Thu, 10 Sep 2026 15:11:58 +0000 Subject: [PATCH 3/4] [SPARK-XXXXX][PYTHON][SQL] Fix CI: nanosecond precision tag on Arrow 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 --- python/pyspark/sql/pandas/types.py | 34 +++++++++++++++++-- .../sql/tests/connect/test_parity_types.py | 8 ++--- 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/python/pyspark/sql/pandas/types.py b/python/pyspark/sql/pandas/types.py index 8782c83426eb8..32cbd27f315d8 100644 --- a/python/pyspark/sql/pandas/types.py +++ b/python/pyspark/sql/pandas/types.py @@ -79,6 +79,29 @@ # Should keep in line with org.apache.spark.sql.util.ArrowUtils.metadataKey metadata_key = b"SPARK::metadata::json" +# Should keep in line with org.apache.spark.sql.util.ArrowUtils.timestampNanosPrecisionKey. The +# nanosecond timestamp types map to an Arrow Timestamp(NANOSECOND) field whose declared precision +# (7-9) cannot be recovered from the Arrow type alone; the JVM ArrowUtils.fromArrowField reads the +# precision from this field-metadata key (defaulting to the maximum when it is absent). Tagging it +# here keeps the precision through a Python -> Arrow -> JVM round trip (e.g. createDataFrame over +# Spark Connect, which reconstructs the schema from the Arrow field rather than a passed schema). +timestamp_nanos_precision_key = b"SPARK::timestampNanos::precision" + + +def _with_timestamp_nanos_precision( + dt: DataType, metadata: Optional[Dict[bytes, bytes]] +) -> Optional[Dict[bytes, bytes]]: + """Merge the nanosecond-precision tag into an Arrow field's metadata for a nanosecond + timestamp type, mirroring the JVM's ``toPrecisionTaggedArrowField``; other types are + unchanged.""" + from pyspark.sql.types import AnyTimestampNanoType + + if isinstance(dt, AnyTimestampNanoType): + merged = dict(metadata) if metadata else {} + merged[timestamp_nanos_precision_key] = str(dt.precision).encode("utf-8") + return merged + return metadata + def to_arrow_metadata(metadata: Optional[Dict[str, Any]] = None) -> Optional[Dict[bytes, bytes]]: if metadata is not None and len(metadata) > 0: @@ -170,6 +193,7 @@ def to_arrow_type( prefers_large_types=prefers_large_types, ), nullable=dt.containsNull, + metadata=_with_timestamp_nanos_precision(dt.elementType, None), ) arrow_type = pa.list_(field) elif isinstance(dt, MapType): @@ -182,6 +206,7 @@ def to_arrow_type( prefers_large_types=prefers_large_types, ), nullable=False, + metadata=_with_timestamp_nanos_precision(dt.keyType, None), ) value_field = pa.field( "value", @@ -192,6 +217,7 @@ def to_arrow_type( prefers_large_types=prefers_large_types, ), nullable=dt.valueContainsNull, + metadata=_with_timestamp_nanos_precision(dt.valueType, None), ) arrow_type = pa.map_(key_field, value_field) elif isinstance(dt, StructType): @@ -211,7 +237,9 @@ def to_arrow_type( prefers_large_types=prefers_large_types, ), nullable=field.nullable, - metadata=to_arrow_metadata(field.metadata), + metadata=_with_timestamp_nanos_precision( + field.dataType, to_arrow_metadata(field.metadata) + ), ) for field in dt ] @@ -299,7 +327,9 @@ def to_arrow_schema( prefers_large_types=prefers_large_types, ), nullable=field.nullable, - metadata=to_arrow_metadata(field.metadata), + metadata=_with_timestamp_nanos_precision( + field.dataType, to_arrow_metadata(field.metadata) + ), ) for field in schema ] diff --git a/python/pyspark/sql/tests/connect/test_parity_types.py b/python/pyspark/sql/tests/connect/test_parity_types.py index 7d202cabff286..059c3888566ea 100644 --- a/python/pyspark/sql/tests/connect/test_parity_types.py +++ b/python/pyspark/sql/tests/connect/test_parity_types.py @@ -55,12 +55,8 @@ def test_timestamp_nanos_type_connect_data_path(self): df = self.spark.sql( "SELECT CAST('2020-01-02 03:04:05.123456789' AS TIMESTAMP_NTZ(9)) AS ts" ) - self.assertEqual( - datetime.datetime(2020, 1, 2, 3, 4, 5, 123456), df.collect()[0].ts - ) - self.assertEqual( - pd.Timestamp("2020-01-02 03:04:05.123456789"), df.toPandas()["ts"][0] - ) + self.assertEqual(datetime.datetime(2020, 1, 2, 3, 4, 5, 123456), df.collect()[0].ts) + self.assertEqual(pd.Timestamp("2020-01-02 03:04:05.123456789"), df.toPandas()["ts"][0]) @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") def test_apply_schema(self): From 895799954627b32a0f658841a10984ebde5fd135 Mon Sep 17 00:00:00 2001 From: Stevo Mitric Date: Thu, 10 Sep 2026 17:33:33 +0000 Subject: [PATCH 4/4] [SPARK-XXXXX][PYTHON][SQL] Address auto-review: drop redundant import, 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 --- python/pyspark/sql/connect/dataframe.py | 18 ++++++++++++++++++ python/pyspark/sql/pandas/types.py | 2 -- .../sql/tests/connect/test_parity_types.py | 11 +++-------- 3 files changed, 21 insertions(+), 10 deletions(-) diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 95d6cddd50ba6..10f8d6153ab2a 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -1959,6 +1959,21 @@ def __dir__(self) -> List[str]: attrs.update(self.columns) return sorted(attrs) + def _check_timestamp_nanos_map_key(self, schema: StructType) -> None: + # SPARK-57462: a nanosecond timestamp used as a map key collapses to a single microsecond + # datetime.datetime when Arrow rows are turned into Python dicts, dropping entries that + # differ only below a microsecond. Reject such a schema up front -- matching the classic + # DataFrame's guard -- rather than silently dropping map entries. + from pyspark.errors import PySparkTypeError + from pyspark.sql.types import _first_timestamp_nanos_map_key_type + + key_type = _first_timestamp_nanos_map_key_type(schema) + if key_type is not None: + raise PySparkTypeError( + errorClass="TIMESTAMP_NANOS_PYTHON_MAP_KEY", + messageParameters={"type": key_type.simpleString()}, + ) + def collect(self) -> List[Row]: table, schema = self._to_table() @@ -1969,6 +1984,8 @@ def collect(self) -> List[Row]: assert schema is not None and isinstance(schema, StructType) + self._check_timestamp_nanos_map_key(schema) + return ArrowTableToRowsConversion.convert( table, schema, binary_as_bytes=self._get_binary_as_bytes() ) @@ -2222,6 +2239,7 @@ def is_cached(self) -> bool: return self.storageLevel != StorageLevel.NONE def toLocalIterator(self, prefetchPartitions: bool = False) -> Iterator[Row]: + self._check_timestamp_nanos_map_key(self.schema) query = self._plan.to_proto(self._session.client) binary_as_bytes = self._get_binary_as_bytes() diff --git a/python/pyspark/sql/pandas/types.py b/python/pyspark/sql/pandas/types.py index 32cbd27f315d8..062c2d3614115 100644 --- a/python/pyspark/sql/pandas/types.py +++ b/python/pyspark/sql/pandas/types.py @@ -94,8 +94,6 @@ def _with_timestamp_nanos_precision( """Merge the nanosecond-precision tag into an Arrow field's metadata for a nanosecond timestamp type, mirroring the JVM's ``toPrecisionTaggedArrowField``; other types are unchanged.""" - from pyspark.sql.types import AnyTimestampNanoType - if isinstance(dt, AnyTimestampNanoType): merged = dict(metadata) if metadata else {} merged[timestamp_nanos_precision_key] = str(dt.precision).encode("utf-8") diff --git a/python/pyspark/sql/tests/connect/test_parity_types.py b/python/pyspark/sql/tests/connect/test_parity_types.py index 059c3888566ea..4ff0f90cafdce 100644 --- a/python/pyspark/sql/tests/connect/test_parity_types.py +++ b/python/pyspark/sql/tests/connect/test_parity_types.py @@ -25,18 +25,13 @@ class TypesParityTests(TypesTestsMixin, ReusedConnectTestCase): # SPARK-57462 follow-up: the nanosecond timestamp value path now works over Spark Connect, # whose data path goes through Arrow (to_arrow_type / ArrowTableToRowsConversion). The - # data-path tests -- test_timestamp_nanos_type, test_timestamp_nanos_type_preview_flag_off and - # test_timestamp_nanos_type_arrow_conversion -- are therefore inherited and run here. The - # nanosecond tests below exercise classic-only mechanisms that do not apply to Connect: the - # Py4J (useArrow=False) UDF path and the classic collect map-key guard in classic/dataframe.py. + # data-path tests -- and the collect map-key collision guard, now mirrored onto the Connect + # DataFrame -- are therefore inherited and run here. The nanosecond tests still skipped below + # exercise the classic-only Py4J (useArrow=False) UDF path, which does not apply to Connect. @unittest.skip("SPARK-57462: uses the classic Py4J (useArrow=False) UDF path, not Connect.") def test_timestamp_nanos_type_python_udf(self): super().test_timestamp_nanos_type_python_udf() - @unittest.skip("SPARK-57462: classic-only collect map-key guard (classic/dataframe.py).") - def test_timestamp_nanos_type_map_key_collision(self): - super().test_timestamp_nanos_type_map_key_collision() - @unittest.skip("SPARK-57462: classic-only Py4J UDF input path (useArrow=False), not Connect.") def test_timestamp_nanos_type_python_udf_input(self): super().test_timestamp_nanos_type_python_udf_input()