diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 593a700f82acf..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() ) @@ -1999,15 +2016,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 @@ -2231,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/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..2085a3a57cd31 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,9 @@ 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 +598,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 +778,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 +791,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 +1177,9 @@ 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 +1215,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,24 +1328,33 @@ 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: 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 - elif isinstance(dataType, TimestampNTZType): + elif isinstance(dataType, (TimestampNTZType, TimestampNTZNanosType)): def convert_timestamp_ntz(value: Any) -> Any: if value is None: 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/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..062c2d3614115 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, @@ -76,25 +79,26 @@ # 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 _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 _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.""" + 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]]: @@ -164,6 +168,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): @@ -178,6 +191,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): @@ -190,6 +204,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", @@ -200,6 +215,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): @@ -219,7 +235,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 ] @@ -307,7 +325,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 ] @@ -570,14 +590,25 @@ 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): + 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 @@ -919,6 +950,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 +1069,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 +1258,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 +1274,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 +1405,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,6 +1606,21 @@ def convert_struct(value: Any) -> Any: return convert_struct + 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 diff --git a/python/pyspark/sql/tests/connect/test_parity_types.py b/python/pyspark/sql/tests/connect/test_parity_types.py index 57ff0c7ba87b2..4ff0f90cafdce 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,35 @@ 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 -- 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: nanosecond timestamp types are pending Connect/Arrow support.") - 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: 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: nanosecond timestamp types are pending Connect/Arrow support.") + @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() - # 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..8f097a9548af8 100644 --- a/python/pyspark/sql/tests/test_types.py +++ b/python/pyspark/sql/tests/test_types.py @@ -2354,36 +2354,77 @@ 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)"}, - ) + 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_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 - # 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)"}, - ) + 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") diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py index 27cc977fe44c6..c8458577ebaf6 100644 --- a/python/pyspark/sql/types.py +++ b/python/pyspark/sql/types.py @@ -573,10 +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 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-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 """ @@ -3084,34 +3088,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: ...