Skip to content

[WIP][SPARK-57462][PYTHON][SQL] Add Arrow support for nanosecond-precision timestamp types in PySpark #58694

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stevomitric wants to merge 4 commits into
base: master
Choose a base branch
from
Open

[WIP][SPARK-57462][PYTHON][SQL] Add Arrow support for nanosecond-precision timestamp types in PySpark #58694

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
93347e9
[SPARK-XXXXX][PYTHON][SQL] Add Arrow support for nanosecond-precision…
stevomitric Sep 9, 2026
62a1cd2
[SPARK-XXXXX][PYTHON][SQL] Address review for nanosecond-precision Ar…
stevomitric Sep 10, 2026
9438c4d
[SPARK-XXXXX][PYTHON][SQL] Fix CI: nanosecond precision tag on Arrow …
stevomitric Sep 10, 2026
8957999
[SPARK-XXXXX][PYTHON][SQL] Address auto-review: drop redundant import…
stevomitric Sep 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 18 additions & 9 deletions python/pyspark/sql/connect/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand All @@ -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()
)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand Down
9 changes: 0 additions & 9 deletions python/pyspark/sql/connect/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
58 changes: 17 additions & 41 deletions python/pyspark/sql/conversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,8 @@
StringType,
StructField,
StructType,
TimestampLTZNanosType,
TimestampNTZNanosType,
TimestampNTZType,
TimestampType,
TimeType,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
21 changes: 1 addition & 20 deletions python/pyspark/sql/pandas/conversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

(
Expand Down Expand Up @@ -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
Expand Down
Loading