[SPARK-57462][PYTHON][SQL] Add PySpark support for nanosecond-precision timestamp types

common/utils/src/main/resources/error/error-conditions.json

@@ -8117,6 +8117,12 @@
     ],
     "sqlState" : "42K0E"
   },
+  "TIMESTAMP_NANOS_PYTHON_MAP_KEY" : {
+    "message" : [
+      "Cannot convert a map with <type> keys to Python. Python `datetime.datetime` is microsecond-resolution and cannot represent the sub-microsecond digits, so two distinct nanosecond keys can collapse to a single entry. Avoid nanosecond-precision timestamp types as map keys on the Python conversion path."
+    ],
+    "sqlState" : "22000"
+  },
   "TRAILING_COMMA_IN_SELECT" : {
     "message" : [
       "Trailing comma detected in SELECT clause. Remove the trailing comma before the FROM clause."

python/docs/source/reference/pyspark.sql/data_types.rst

@@ -50,6 +50,8 @@ Data Types
     VariantType
     TimestampType
     TimestampNTZType
+    TimestampNTZNanosType
+    TimestampLTZNanosType
     DayTimeIntervalType
     YearMonthIntervalType
     CalendarIntervalType

python/pyspark/errors/error-conditions.json

@@ -465,6 +465,11 @@
       "Timeout timestamp (<timestamp>) cannot be earlier than the current watermark (<watermark>)."
     ]
   },
+  "INVALID_TIMESTAMP_PRECISION": {
+    "message": [
+      "The seconds precision <precision> of <type> is invalid. Expected an integer in [7, 9] for nanosecond precision; use precision 6 or parameterless <type> for the standard microsecond type."
+    ]
+  },
   "INVALID_TYPE": {
     "message": [
       "Argument `<arg_name>` should not be a <arg_type>."
@@ -1009,6 +1014,11 @@
       "<test_class_path> doesn't exist. Spark sql test classes are not compiled."
     ]
   },
+  "TIMESTAMP_NANOS_PYTHON_MAP_KEY": {
+    "message": [
+      "Cannot convert a map with <type> keys to Python. Python `datetime.datetime` is microsecond-resolution and cannot represent the sub-microsecond digits, so two distinct nanosecond keys can collapse to a single entry. Avoid nanosecond-precision timestamp types as map keys on the Python conversion path."
+    ]
+  },
   "TOO_MANY_VALUES": {
     "message": [
       "Expected <expected> values for `<item>`, got <actual>."

python/pyspark/sql/classic/dataframe.py

@@ -70,6 +70,7 @@
 from pyspark.sql.types import (
     Row,
     StructType,
+    _first_timestamp_nanos_map_key_type,
     _parse_datatype_json_string,
 )
 from pyspark.sql.utils import get_active_spark_context, to_java_array, to_scala_map
@@ -506,13 +507,29 @@ def _converter(parameter: Union[str, list, float, int, Column]) -> Any:
     def count(self) -> int:
         return int(self._jdf.count())
 
+    def _check_timestamp_nanos_map_key(self) -> None:
+        # SPARK-57462: a nanosecond timestamp map key collapses to a single microsecond
+        # datetime.datetime on the Python side, dropping an entry. The value conversion runs in a
+        # background serve thread, whose exception would surface to collect() as a normal EOF
+        # (empty/partial result, not an error), so reject the schema up front here instead.
+        key_type = _first_timestamp_nanos_map_key_type(self.schema)
+        if key_type is not None:
+            from pyspark.errors import PySparkTypeError
+
+            raise PySparkTypeError(
+                errorClass="TIMESTAMP_NANOS_PYTHON_MAP_KEY",
+                messageParameters={"type": key_type.simpleString()},
+            )
+
     def collect(self) -> List[Row]:
+        self._check_timestamp_nanos_map_key()
         with SCCallSiteSync(self._sc):
             sock_info = self._jdf.collectToPython()
         with _load_from_socket(sock_info, BatchedSerializer(CPickleSerializer())) as stream:
             return list(stream)
 
     def toLocalIterator(self, prefetchPartitions: bool = False) -> Iterator[Row]:
+        self._check_timestamp_nanos_map_key()
         with SCCallSiteSync(self._sc):
             sock_info = self._jdf.toPythonIterator(prefetchPartitions)
         return _local_iterator_from_socket(sock_info, BatchedSerializer(CPickleSerializer()))
@@ -529,6 +546,7 @@ def take(self, num: int) -> List[Row]:
         return self.limit(num).collect()
 
     def tail(self, num: int) -> List[Row]:
+        self._check_timestamp_nanos_map_key()
         with SCCallSiteSync(self._sc):
             sock_info = self._jdf.tailToPython(num)
         with _load_from_socket(sock_info, BatchedSerializer(CPickleSerializer())) as stream:

python/pyspark/sql/connect/dataframe.py

@@ -1999,6 +1999,15 @@ 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

python/pyspark/sql/connect/session.py

@@ -517,6 +517,15 @@ 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",

python/pyspark/sql/connect/types.py

@@ -43,6 +43,8 @@
     StringType,
     StructField,
     StructType,
+    TimestampLTZNanosType,
+    TimestampNTZNanosType,
     TimestampNTZType,
     TimestampType,
     TimeType,
@@ -159,6 +161,10 @@ def pyspark_types_to_proto_types(data_type: DataType) -> pb2.DataType:
         ret.timestamp.CopyFrom(pb2.DataType.Timestamp())
     elif isinstance(data_type, TimestampNTZType):
         ret.timestamp_ntz.CopyFrom(pb2.DataType.TimestampNTZ())
+    elif isinstance(data_type, TimestampNTZNanosType):
+        ret.timestamp_ntz_nanos.precision = data_type.precision
+    elif isinstance(data_type, TimestampLTZNanosType):
+        ret.timestamp_ltz_nanos.precision = data_type.precision
     elif isinstance(data_type, DayTimeIntervalType):
         ret.day_time_interval.start_field = data_type.startField
         ret.day_time_interval.end_field = data_type.endField
@@ -251,6 +257,19 @@ def proto_schema_to_pyspark_data_type(schema: pb2.DataType) -> DataType:
         return TimestampType()
     elif schema.HasField("timestamp_ntz"):
         return TimestampNTZType()
+    elif schema.HasField("timestamp_ntz_nanos"):
+        # `precision` is optional on the wire; per types.proto it defaults to 9 when omitted.
+        return (
+            TimestampNTZNanosType(schema.timestamp_ntz_nanos.precision)
+            if schema.timestamp_ntz_nanos.HasField("precision")
+            else TimestampNTZNanosType()
+        )
+    elif schema.HasField("timestamp_ltz_nanos"):
+        return (
+            TimestampLTZNanosType(schema.timestamp_ltz_nanos.precision)
+            if schema.timestamp_ltz_nanos.HasField("precision")
+            else TimestampLTZNanosType()
+        )
     elif schema.HasField("day_time_interval"):
         start: Optional[int] = (
             schema.day_time_interval.start_field

python/pyspark/sql/conversion.py

@@ -32,6 +32,7 @@
 )
 from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
 from pyspark.sql.types import (
+    AnyTimestampNanoType,
     ArrayType,
     BinaryType,
     BooleanType,
@@ -542,6 +543,11 @@ def _need_converter(
         elif isinstance(dataType, (TimestampType, TimestampNTZType)):
             # 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
@@ -595,6 +601,18 @@ 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:
@@ -1177,6 +1195,11 @@ def _need_converter(dataType: DataType) -> bool:
         elif isinstance(dataType, (TimestampType, TimestampNTZType)):
             # 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):
@@ -1212,6 +1235,19 @@ 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
 

python/pyspark/sql/pandas/conversion.py

@@ -320,11 +320,18 @@ def _to_pandas(self, **kwargs: Any) -> "PandasDataFrameLike":
 
         assert isinstance(self, DataFrame)
 
-        from pyspark.sql.pandas.types import _create_converter_to_pandas
+        from pyspark.sql.pandas.types import (
+            _create_converter_to_pandas,
+            _reject_timestamp_nanos_conversion,
+        )
         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
 
         (
@@ -632,6 +639,16 @@ 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. Covers a
+        # bare atomic DataType schema as well as a StructType (a DDL string / list of names cannot
+        # carry these types without an explicit DataType, so 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

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,
@@ -77,6 +78,45 @@
 metadata_key = b"SPARK::metadata::json"
 
 
+def _contains_timestamp_nanos(dt: DataType) -> bool:
+    """True if ``dt`` is, or structurally contains, a nanosecond-capable timestamp type.
+
+    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``), callers use this to fail deterministically; see
+    :func:`_reject_timestamp_nanos_conversion`.
+    """
+    if isinstance(dt, AnyTimestampNanoType):
+        return True
+    elif isinstance(dt, ArrayType):
+        return _contains_timestamp_nanos(dt.elementType)
+    elif isinstance(dt, MapType):
+        return _contains_timestamp_nanos(dt.keyType) or _contains_timestamp_nanos(dt.valueType)
+    elif isinstance(dt, StructType):
+        return any(_contains_timestamp_nanos(f.dataType) for f in dt.fields)
+    elif isinstance(dt, UserDefinedType):
+        return _contains_timestamp_nanos(dt.sqlType())
+    else:
+        return False
+
+
+def _reject_timestamp_nanos_conversion(schema: DataType) -> None:
+    """Raise if ``schema`` involves a nanosecond timestamp type, for Arrow/pandas value paths.
+
+    Keeps the not-yet-supported Arrow/pandas/Connect data paths failing deterministically instead
+    of silently producing wrong values, consistent with :func:`to_arrow_type`, which already
+    rejects these types with the same error condition.
+    """
+    from pyspark.errors import PySparkTypeError
+
+    if _contains_timestamp_nanos(schema):
+        raise PySparkTypeError(
+            errorClass="UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION",
+            messageParameters={"data_type": str(schema)},
+        )
+
+
 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")}

python/pyspark/sql/tests/connect/test_connect_plan.py

@@ -54,7 +54,10 @@
         WriteOperation,
     )
     from pyspark.sql.connect.readwriter import DataFrameReader
-    from pyspark.sql.connect.types import pyspark_types_to_proto_types
+    from pyspark.sql.connect.types import (
+        proto_schema_to_pyspark_data_type,
+        pyspark_types_to_proto_types,
+    )
     from pyspark.sql.types import (
         ArrayType,
         DoubleType,
@@ -63,6 +66,8 @@
         StringType,
         StructField,
         StructType,
+        TimestampLTZNanosType,
+        TimestampNTZNanosType,
     )
 
 
@@ -818,6 +823,50 @@ def test_to(self):
         new_plan = df.to(schema)._plan.to_proto(self.connect)
         self.assertEqual(pyspark_types_to_proto_types(schema), new_plan.root.to_schema.schema)
 
+    def test_timestamp_nanos_datatype_conversion(self):
+        # SPARK-57462: the nanosecond timestamp types round-trip through the DataType proto.
+        for dt in [
+            TimestampNTZNanosType(7),
+            TimestampNTZNanosType(8),
+            TimestampNTZNanosType(9),
+            TimestampLTZNanosType(7),
+            TimestampLTZNanosType(8),
+            TimestampLTZNanosType(9),
+        ]:
+            with self.subTest(dt=repr(dt)):
+                self.assertEqual(
+                    dt, proto_schema_to_pyspark_data_type(pyspark_types_to_proto_types(dt))
+                )
+
+        schema = StructType(
+            [
+                StructField("ntz", TimestampNTZNanosType(9), True),
+                StructField("arr", ArrayType(TimestampLTZNanosType(7), True), False),
+                StructField(
+                    "map",
+                    MapType(TimestampNTZNanosType(8), TimestampLTZNanosType(9), True),
+                    True,
+                ),
+            ]
+        )
+        self.assertEqual(
+            schema, proto_schema_to_pyspark_data_type(pyspark_types_to_proto_types(schema))
+        )
+
+        # `precision` is optional on the wire; types.proto documents 9 as the default.
+        self.assertEqual(
+            TimestampNTZNanosType(9),
+            proto_schema_to_pyspark_data_type(
+                proto.DataType(timestamp_ntz_nanos=proto.DataType.TimestampNTZNanos())
+            ),
+        )
+        self.assertEqual(
+            TimestampLTZNanosType(9),
+            proto_schema_to_pyspark_data_type(
+                proto.DataType(timestamp_ltz_nanos=proto.DataType.TimestampLTZNanos())
+            ),
+        )
+
     def test_write_operation(self):
         wo = WriteOperation(self.connect.readTable("name")._plan)
         wo.mode = "overwrite"