[SPARK-57462][PYTHON][SQL] Add PySpark support for nanosecond-precision timestamp types - #58418
[SPARK-57462][PYTHON][SQL] Add PySpark support for nanosecond-precision timestamp types#58418stevomitric wants to merge 9 commits into
Conversation
bba625a to
b8598c7
Compare
…on timestamp types ### What changes were proposed in this pull request? Exposes the nanosecond-capable timestamp types `TIMESTAMP_NTZ(p)` / `TIMESTAMP_LTZ(p)` (`p` in [7, 9]) in PySpark. Before this change `python/pyspark/sql/types.py` defined only the microsecond singletons, so any DataFrame whose schema contained one of these types failed on the Python side even though the JVM and the Spark Connect protocol already supported them. - `python/pyspark/sql/types.py`: adds `TimestampNTZNanosType(precision)` and `TimestampLTZNanosType(precision)` (plus the shared `AnyTimestampNanoType` base, which is not exported), with precision validation, `simpleString` / `jsonValue` / `__repr__`, and `toInternal` / `fromInternal`. Registers the parameterized JSON type names, mirroring `DataType.parseDataType` in `sql/api`: precision 6 maps to the standard microsecond type, 7-9 to the nanosecond types, and any other precision is rejected. Also registers the two types in `_acceptable_types` (used by `_make_type_verifier`, so `createDataFrame` accepts `datetime.datetime` values) and in `_get_jvm_type_name` (so `printSchema()` renders `timestamp_ntz(9)` rather than a name derived from the class). - `python/pyspark/sql/connect/types.py`: converts the two types to and from the Connect `DataType` proto in both directions, treating an omitted `precision` as 9 per types.proto. - `sql/api/.../types/ops/TimestampNanosTypeApiOps.scala`: implements the Types Framework Python-interop hooks (`needConversionInPython`, `makeFromJava`) so values round-trip over Py4J. Without these, `EvaluatePython.makeFromJava` fell through to its catch-all and silently produced NULL for every nanosecond column. - `sql/core/.../EvaluatePython.scala`: adds the reverse direction, converting the internal `TimestampNanosVal` to epoch microseconds. Without it the value reached the pickler as a raw `TimestampNanosVal`, which has no registered pickler. - `python/pyspark/errors/error-conditions.json`: adds `INVALID_TIMESTAMP_PRECISION`, worded to match the JVM error condition of the same name. The external Python value is `datetime.datetime`, which is microsecond-resolution, so the Py4J protocol carries epoch microseconds and sub-microsecond digits are truncated at the Python boundary in both directions. This mirrors the shipped `TimeType` behaviour and is the microsecond-only Python/UDF limitation already documented by SPARK-57808; the stored value keeps full precision. Type inference is unchanged: a bare `datetime.datetime` still infers microsecond `TimestampType`, and the nanosecond types are reachable only through an explicit schema. Arrow and pandas value conversion (`toPandas`, `createDataFrame` from pandas, and therefore the Spark Connect data path) is deliberately not included here and remains follow-up work; `to_arrow_type` continues to reject these types. ### Why are the changes needed? Without Python type classes, a `TIMESTAMP(9)` column cannot be read or written from PySpark at all: `proto_schema_to_pyspark_data_type` raised `UNSUPPORTED_OPERATION` for the Connect schema, and reading `df.schema` failed because the parameterized JSON type name had no Python parser. This is the last missing client for the umbrella SPARK-56822. ### Does this PR introduce _any_ user-facing change? Yes. `TimestampNTZNanosType` and `TimestampLTZNanosType` are new public types in `pyspark.sql.types`. They are only reachable via an explicit schema or a nanosecond-typed query result, and the server keeps them behind the `spark.sql.timestampNanosTypes.enabled` preview flag, so no existing behaviour changes. ### How was this patch tested? New tests in `python/pyspark/sql/tests/test_types.py`: - `DataTypeTests`: precision validation (7-9 accepted; -1/0/5/6/10 rejected with `INVALID_TIMESTAMP_PRECISION`), string representations including `printSchema()` rendering, equality / hashing / pickling, JSON parsing across precisions including the 6 -> microsecond mapping and the rejected precisions, nested array/map/struct JSON round-trip, and `toInternal` / `fromInternal` agreement with the microsecond types. - `DataTypeVerificationTests`: accepted and rejected values for both types. - `TypesTestsMixin.test_timestamp_nanos_type`: DDL parse agreement with the JVM, plus a `createDataFrame` / `collect` round-trip with nulls, and a check that a value stored at nanosecond precision still renders 9 fractional digits server-side while truncating to microseconds when collected as a `datetime`. New test in `python/pyspark/sql/tests/connect/test_connect_plan.py`: DataType proto round-trip for both types across precisions, nested, and with `precision` omitted. Co-authored-by: Isaac <no-reply@databricks.com>
b8598c7 to
f8b74b1
Compare
… nanosecond timestamp types ### What changes were proposed in this pull request? Adds two tests to `python/pyspark/sql/tests/test_types.py` that drive a nanosecond timestamp value *into* a classic Python UDF as an argument, exercising the JVM -> Python direction of `EvaluatePython.toJava` through the UDF-input caller: - `test_timestamp_nanos_type_python_udf_input`: `udf(lambda x: x, TimestampNTZNanosType(9))` over a nanosecond column, covering the `TimestampNanosVal -> epochMicros` scalar arm. - `test_timestamp_nanos_type_map_key_python_udf_input`: a map with nanosecond keys fed into a UDF, reaching the map-key rejection branch (`TIMESTAMP_NANOS_PYTHON_MAP_KEY`) in `EvaluatePython.toJava`. The UDF returns a non-map type, so the result collect does not re-trip the earlier Python-side guard in `classic/dataframe.py`. ### Why are the changes needed? Addresses review feedback: the existing nanosecond UDF test only *returns* a nanosecond value, and the map-key collision test uses `collect()`, which trips the Python-side guard in `classic/dataframe.py` before reaching the JVM `toJava` throw. The `toJava` scalar arm reached via UDF input, and the map-key throw branch (whose own comment cites the Python-UDF input path), were therefore unexercised. ### Does this PR introduce _any_ user-facing change? No. Test-only. ### How was this patch tested? New session tests in `python/pyspark/sql/tests/test_types.py`, run under CI (both require a live SparkSession). `useArrow=False` forces the classic Py4J path. Co-authored-by: Isaac <no-reply@databricks.com>
420f06f to
7e3668c
Compare
…timestamp types Addresses review feedback: 4.4.0 is the next feature release (branch-4.x is at 4.4.0.dev0) and matches the other recent versionadded:: 4.4.0 entries. master's version.py (5.0.0.dev0) is a longer-horizon placeholder, not the release this ships in. Co-authored-by: Isaac <no-reply@databricks.com>
…TAMP_NANOS_PYTHON_MAP_KEY Addresses review feedback: the Python-side guard filled the error message with the whole DataFrame schema (self.schema.simpleString()), so it rendered e.g. "a map with struct<m:map<timestamp_ntz(9),int>> keys...". It now reports the offending map's key type via a new _first_timestamp_nanos_map_key_type helper (replacing the bool _has_timestamp_nanos_map_key), mirroring the JVM EvaluatePython.toJava twin which reports mt.keyType.sql. Co-authored-by: Isaac <no-reply@databricks.com>
uros-b
left a comment
There was a problem hiding this comment.
Thank you @stevomitric! Adding @HyukjinKwon @Yicong-Huang for PySpark
uros-b
left a comment
There was a problem hiding this comment.
Actually, leaving a few more comments here
HyukjinKwon
left a comment
There was a problem hiding this comment.
0 blocking, 6 non-blocking, 1 nit.
Core implementation looks correct and thorough; the open items are test-assertion tightening, a public typeName() gap, and two message/comment accuracy fixes -- all already raised in review and awaiting the author.
Already raised in existing discussion (7)
- typeName() returns the mangled default 'timestampntznanos'/'timestampltznanos' for the new public types instead of a clean name like the peer TimestampNTZType ("timestamp_ntz"); add an instance typeName() returning simpleString(). -- existing discussion
- The map-key-collision, UDF-map-key, and preview-flag-off tests use assertRaises(Exception); switch to check_error with the specific error class so they actually assert the intended condition. -- existing discussion
- test_timestamp_nanos_type_python_udf_input uses lambda x: x, which cannot prove toJava produced a datetime (the round-trip masks a raw-Long bug); assert isinstance(x, datetime.datetime) inside the UDF. -- existing discussion
- The new Arrow/pandas/Connect rejection paths have no test asserting UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION; add a Connect collect test and a classic toPandas assertion. -- existing discussion
- _contains_timestamp_nanos duplicates _has_type(dt, AnyTimestampNanoType); call _has_type directly and drop the helper. -- existing discussion
- _reject_timestamp_nanos_conversion reports str(schema) (the whole struct) in the UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION message; report the offending leaf type, consistent with to_arrow_type and the just-fixed map-key error. -- existing discussion
- The createDataFrame comment claims a DDL string cannot carry these types; classic createDataFrame parses the DDL to a StructType first, so the guard does fire -- only a plain list of names lacks type info. Correct the comment. -- existing discussion
Verification
Independently confirmed the value contract rather than restating the PR: TimestampNTZNanosType.fromInternal / TimestampLTZNanosType.toInternal are line-for-line the micros TimestampNTZType/TimestampType logic (epoch-micros, UTC-grid vs local-zone), so the round-trip truncates to micros without shifting the instant; makeFromJava calls checkTimestampNanosTypesEnabled() at converter-build time, which test_timestamp_nanos_type_preview_flag_off exercises; and the map-key guard's structural walk matches _has_type(dt.keyType, AnyTimestampNanoType).
…or messages, and tests Addresses a second round of review feedback: - typeName() on the nanosecond types now returns simpleString() (e.g. timestamp_ntz(9)) instead of the class-derived "timestampntznanos", matching the JVM type name. - _reject_timestamp_nanos_conversion reports the offending leaf type (like to_arrow_type) instead of the whole schema, via a shared _first_timestamp_nanos_type helper; this removes _contains_timestamp_nanos, which duplicated _has_type(dt, AnyTimestampNanoType). - Add bare "timestamp_ltz" to _all_mappable_types (the LTZ spelling of the default TimestampType, matching DataTypeAstBuilder), alongside the existing "timestamp_ntz". - Correct the pandas createDataFrame comment: createDataFrame pre-parses a DDL-string schema into a DataType, so the isinstance(schema, DataType) guard already covers DDL strings. - Tighten the nanos tests to assert specific error conditions (TIMESTAMP_NANOS_PYTHON_MAP_KEY, FEATURE_NOT_ENABLED, UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION) rather than bare Exception; prove the Python-UDF-input path actually receives a datetime.datetime; and add classic Arrow-rejection coverage (toPandas / pandas createDataFrame) plus a Connect collect rejection. Co-authored-by: Isaac <no-reply@databricks.com>
…sertSchemaEqual Since typeName() now embeds the precision (e.g. timestamp_ntz(9)), compare_datatypes_ignore_nullable already distinguishes precisions at the typeName equality check, so the AnyTimestampNanoType branch (and its now-inaccurate comment claiming the type name does not distinguish precisions) is dead. Remove the branch and its unused import; behavior is unchanged -- same precision stays equal under ignoreNullable, and different precisions / micros-vs-nanos stay distinct. Co-authored-by: Isaac <no-reply@databricks.com>
…timestamp_nanos_map_key_type In the MapType branch, once _has_type(dt.keyType, AnyTimestampNanoType) is False the whole key subtree provably carries no nanosecond type, so a map nested inside the key cannot have a nanosecond key either -- the _first_timestamp_nanos_map_key_type(dt.keyType) operand of the or is always None. Recurse only into the value type. Behavior is unchanged (verified across nested map/array/struct and key-is-a-map cases). Co-authored-by: Isaac <no-reply@databricks.com>
|
Thanks @uros-b for and @HyukjinKwon for the review! @HyukjinKwon, addressed all items, see individual comments that Uros posted. |
The sole failing check was the K8s integration test, which failed while building the Spark Docker image (docker buildx 404 on moby/buildkit:buildx-stable-1) -- unrelated to this PySpark change. Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
uros-b
left a comment
There was a problem hiding this comment.
Thank you for addressing all comments @stevomitric!
What changes were proposed in this pull request?
Exposes the nanosecond-capable timestamp types
TIMESTAMP_NTZ(p)/TIMESTAMP_LTZ(p)(pin [7, 9]) in PySpark. Before this changepython/pyspark/sql/types.pydefined only the microsecond singletons, so any DataFrame whose schema contained one of these types failed on the Python side even though the JVM and the Spark Connect protocol already supported them.python/pyspark/sql/types.py: addsTimestampNTZNanosType(precision)andTimestampLTZNanosType(precision)(plus the sharedAnyTimestampNanoTypebase, which is not exported), with precision validation,simpleString/jsonValue/__repr__, andtoInternal/fromInternal. Registers the parameterized JSON type names, mirroringDataType.parseDataTypeinsql/api: precision 6 maps to the standard microsecond type, 7-9 to the nanosecond types, and any other precision is rejected. Also registers the two types in_acceptable_types(used by_make_type_verifier, socreateDataFrameacceptsdatetime.datetimevalues) and in_get_jvm_type_name(soprintSchema()renderstimestamp_ntz(9)rather than a name derived from the class).python/pyspark/sql/connect/types.py: converts the two types to and from the ConnectDataTypeproto in both directions, treating an omittedprecisionas 9 per types.proto.sql/api/.../types/ops/TimestampNanosTypeApiOps.scala: implements the Types Framework Python-interop hooks (needConversionInPython,makeFromJava) so values round-trip over Py4J. Without these,EvaluatePython.makeFromJavafell through to its catch-all and silently produced NULL for every nanosecond column.sql/core/.../EvaluatePython.scala: adds the reverse direction, converting the internalTimestampNanosValto epoch microseconds. Without it the value reached the pickler as a rawTimestampNanosVal, which has no registered pickler.python/pyspark/errors/error-conditions.json: addsINVALID_TIMESTAMP_PRECISION, worded to match the JVM error condition of the same name.The external Python value is
datetime.datetime, which is microsecond-resolution, so the Py4J protocol carries epoch microseconds and sub-microsecond digits are truncated at the Python boundary in both directions. This mirrors the shippedTimeTypebehaviour and is the microsecond-only Python/UDF limitation already documented by SPARK-57808; the stored value keeps full precision. Type inference is unchanged: a baredatetime.datetimestill infers microsecondTimestampType, and the nanosecond types are reachable only through an explicit schema.Arrow and pandas value conversion (
toPandas,createDataFramefrom pandas, and therefore the Spark Connect data path) is deliberately not included here and remains follow-up work;to_arrow_typecontinues to reject these types.Why are the changes needed?
Without Python type classes, a
TIMESTAMP(9)column cannot be read or written from PySpark at all:proto_schema_to_pyspark_data_typeraisedUNSUPPORTED_OPERATIONfor the Connect schema, and readingdf.schemafailed because the parameterized JSON type name had no Python parser. This is the last missing client for the umbrella SPARK-56822.Does this PR introduce any user-facing change?
Yes.
TimestampNTZNanosTypeandTimestampLTZNanosTypeare new public types inpyspark.sql.types. They are only reachable via an explicit schema or a nanosecond-typed query result, and the server keeps them behind thespark.sql.timestampNanosTypes.enabledpreview flag, so no existing behaviour changes.How was this patch tested?
New tests in
python/pyspark/sql/tests/test_types.py:DataTypeTests: precision validation (7-9 accepted; -1/0/5/6/10 rejected withINVALID_TIMESTAMP_PRECISION), string representations includingprintSchema()rendering, equality / hashing / pickling, JSON parsing across precisions including the 6 -> microsecond mapping and the rejected precisions, nested array/map/struct JSON round-trip, andtoInternal/fromInternalagreement with the microsecond types.DataTypeVerificationTests: accepted and rejected values for both types.TypesTestsMixin.test_timestamp_nanos_type: DDL parse agreement with the JVM, plus acreateDataFrame/collectround-trip with nulls, and a check that a value stored at nanosecond precision still renders 9 fractional digits server-side while truncating to microseconds when collected as adatetime.New test in
python/pyspark/sql/tests/connect/test_connect_plan.py: DataType proto round-trip for both types across precisions, nested, and withprecisionomitted.Was this patch authored or co-authored using generative AI tooling?
Co-authored: Claude Opus 5