diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 5af823a3f24b2..1ec2bda6160cc 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/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 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." diff --git a/python/docs/source/reference/pyspark.sql/data_types.rst b/python/docs/source/reference/pyspark.sql/data_types.rst index ffb24c68445dd..8a863758f0b2b 100644 --- a/python/docs/source/reference/pyspark.sql/data_types.rst +++ b/python/docs/source/reference/pyspark.sql/data_types.rst @@ -50,6 +50,8 @@ Data Types VariantType TimestampType TimestampNTZType + TimestampNTZNanosType + TimestampLTZNanosType DayTimeIntervalType YearMonthIntervalType CalendarIntervalType diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index ae35e98a5e5c0..66d2ec3e37ac9 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -465,6 +465,11 @@ "Timeout timestamp () cannot be earlier than the current watermark ()." ] }, + "INVALID_TIMESTAMP_PRECISION": { + "message": [ + "The seconds precision of is invalid. Expected an integer in [7, 9] for nanosecond precision; use precision 6 or parameterless for the standard microsecond type." + ] + }, "INVALID_TYPE": { "message": [ "Argument `` should not be a ." @@ -1009,6 +1014,11 @@ " doesn't exist. Spark sql test classes are not compiled." ] }, + "TIMESTAMP_NANOS_PYTHON_MAP_KEY": { + "message": [ + "Cannot convert a map with 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 values for ``, got ." diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py index aea6869f455f8..ddef85e3008ef 100644 --- a/python/pyspark/sql/classic/dataframe.py +++ b/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: diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 95d6cddd50ba6..593a700f82acf 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/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 diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index bd6f32ab5ac05..a518fdd327eac 100644 --- a/python/pyspark/sql/connect/session.py +++ b/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", diff --git a/python/pyspark/sql/connect/types.py b/python/pyspark/sql/connect/types.py index 78270ca4c2280..3bc5aeb2cfe9c 100644 --- a/python/pyspark/sql/connect/types.py +++ b/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 diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index 0585a730fb087..39b1c7fe39860 100644 --- a/python/pyspark/sql/conversion.py +++ b/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 diff --git a/python/pyspark/sql/pandas/conversion.py b/python/pyspark/sql/pandas/conversion.py index 3bdaa67b91e64..2edd1c3449f20 100644 --- a/python/pyspark/sql/pandas/conversion.py +++ b/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 diff --git a/python/pyspark/sql/pandas/types.py b/python/pyspark/sql/pandas/types.py index c4facb3e3a8b4..1e6b6b57c62ff 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, @@ -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")} diff --git a/python/pyspark/sql/tests/connect/test_connect_plan.py b/python/pyspark/sql/tests/connect/test_connect_plan.py index 60f0add9a52a8..d3b660f6ccd86 100644 --- a/python/pyspark/sql/tests/connect/test_connect_plan.py +++ b/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" diff --git a/python/pyspark/sql/tests/connect/test_parity_types.py b/python/pyspark/sql/tests/connect/test_parity_types.py index 42346c4fc4a17..63d90075cbefb 100644 --- a/python/pyspark/sql/tests/connect/test_parity_types.py +++ b/python/pyspark/sql/tests/connect/test_parity_types.py @@ -22,6 +22,34 @@ 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.") + 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.") + 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.") + def test_timestamp_nanos_type_map_key_python_udf_input(self): + super().test_timestamp_nanos_type_map_key_python_udf_input() + @unittest.skip("Spark Connect does not support RDD but the tests depend on them.") def test_apply_schema(self): super().test_apply_schema() diff --git a/python/pyspark/sql/tests/test_types.py b/python/pyspark/sql/tests/test_types.py index e391fcb441ea9..daf2cb53366ee 100644 --- a/python/pyspark/sql/tests/test_types.py +++ b/python/pyspark/sql/tests/test_types.py @@ -61,6 +61,8 @@ StringType, StructField, StructType, + TimestampLTZNanosType, + TimestampNTZNanosType, TimestampNTZType, TimestampType, TimeType, @@ -2215,6 +2217,126 @@ def test_daytime_interval_type(self): for n, (a, e) in enumerate(zip(actual, expected)): self.assertEqual(a, e, "%s does not match with %s" % (exprs[n], expected[n])) + def test_timestamp_nanos_type(self): + from pyspark.sql.types import _parse_datatype_string + + # SPARK-57462: createDataFrame / collect with an explicit nanosecond timestamp schema. + # The types are behind a preview flag on the server; it is on under tests, but set it + # explicitly rather than relying on that default. + with self.sql_conf({"spark.sql.timestampNanosTypes.enabled": True}): + schema = StructType( + [ + StructField("ntz", TimestampNTZNanosType(9), True), + StructField("ltz", TimestampLTZNanosType(7), True), + ] + ) + # The JVM DDL parser and the Python JSON reader must agree on the type names. + self.assertEqual( + schema, + _parse_datatype_string("ntz timestamp_ntz(9), ltz timestamp_ltz(7)"), + ) + + # datetime.datetime is microsecond-resolution, so values cross the Python boundary at + # microsecond precision; naive values round-trip exactly (see the class docstrings). + ts = datetime.datetime(2020, 1, 2, 3, 4, 5, 123456) + df = self.spark.createDataFrame([(ts, ts), (None, None)], schema) + self.assertEqual(schema, df.schema) + + rows = df.collect() + self.assertEqual(2, len(rows)) + self.assertEqual(ts, rows[0].ntz) + self.assertEqual(ts, rows[0].ltz) + self.assertIsNone(rows[1].ntz) + self.assertIsNone(rows[1].ltz) + + # The server keeps the full precision: the microsecond truncation above is a property + # of datetime.datetime, not of the stored value. + nanos = self.spark.sql( + "SELECT CAST('2020-01-02 03:04:05.123456789' AS TIMESTAMP_NTZ(9)) AS ts" + ) + self.assertEqual(TimestampNTZNanosType(9), nanos.schema["ts"].dataType) + self.assertEqual( + "2020-01-02 03:04:05.123456789", + nanos.select(F.col("ts").cast("string")).first()[0], + ) + # ... and the same value truncates to microseconds when collected as a datetime. + self.assertEqual(datetime.datetime(2020, 1, 2, 3, 4, 5, 123456), nanos.first().ts) + + def test_timestamp_nanos_type_preview_flag_off(self): + # SPARK-57462: with the preview flag off, the classic explicit-schema createDataFrame + # path (which goes through EvaluatePython.makeFromJava, not the row encoder) must not + # execute; the eager guard on makeFromJava enforces that. + schema = StructType([StructField("ts", TimestampNTZNanosType(9))]) + data = [(datetime.datetime(2020, 1, 1),)] + with self.sql_conf({"spark.sql.timestampNanosTypes.enabled": False}): + with self.assertRaises(Exception): + self.spark.createDataFrame(data, schema).collect() + + def test_timestamp_nanos_type_python_udf(self): + # SPARK-57462: a Python UDF with a nanosecond return type exercises makeFromJava + # (Python -> JVM). useArrow=False forces the classic Py4J path; the Arrow-based UDF path + # is not yet implemented for these types. The value round-trips at microsecond resolution. + from pyspark.sql.functions import udf + + with self.sql_conf({"spark.sql.timestampNanosTypes.enabled": True}): + value = datetime.datetime(2021, 6, 7, 8, 9, 10, 123456) + nanos_udf = udf(lambda _: value, returnType=TimestampLTZNanosType(9), useArrow=False) + row = self.spark.range(1).select(nanos_udf("id").alias("ts")).first() + self.assertEqual(value, row.ts) + + def test_timestamp_nanos_type_map_key_collision(self): + # SPARK-57462: two nanosecond keys that differ only below a microsecond collapse to the + # same microsecond-resolution Python key. Rather than silently drop a map entry, the + # conversion fails deterministically. + with self.sql_conf({"spark.sql.timestampNanosTypes.enabled": True}): + df = self.spark.sql( + "SELECT map(" + "CAST('2020-01-01 00:00:00.123456700' AS TIMESTAMP_NTZ(9)), 1, " + "CAST('2020-01-01 00:00:00.123456800' AS TIMESTAMP_NTZ(9)), 2) AS m" + ) + with self.assertRaises(Exception): + df.collect() + + def test_timestamp_nanos_type_python_udf_input(self): + # SPARK-57462: a Python UDF that takes a nanosecond column as an *argument* exercises the + # JVM -> Python direction of EvaluatePython.toJava (the TimestampNanosVal -> epochMicros + # arm) through the UDF-input caller, which the collect-based tests above do not reach. + # useArrow=False forces the classic Py4J path; the value round-trips at microsecond + # resolution (datetime.datetime is microsecond-precision). + from pyspark.sql.functions import udf + + with self.sql_conf({"spark.sql.timestampNanosTypes.enabled": True}): + value = datetime.datetime(2021, 6, 7, 8, 9, 10, 123456) + df = self.spark.createDataFrame( + [(value,)], StructType([StructField("ts", TimestampNTZNanosType(9))]) + ) + identity_udf = udf(lambda x: x, returnType=TimestampNTZNanosType(9), useArrow=False) + row = df.select(identity_udf("ts").alias("out")).first() + self.assertEqual(value, row.out) + + def test_timestamp_nanos_type_map_key_python_udf_input(self): + # SPARK-57462: feeding a map with nanosecond keys into a Python UDF reaches the map-key + # rejection branch in EvaluatePython.toJava (the "Python UDF input path" its own comment + # cites). The collect-based collision test does not reach it -- collect() trips the earlier + # Python-side guard in classic/dataframe.py first. The UDF returns a non-map type so the + # result collect does not re-trip that Python-side guard; the failure must come from the + # nanosecond-map-key input conversion. + from pyspark.sql.functions import udf + + with self.sql_conf({"spark.sql.timestampNanosTypes.enabled": True}): + df = self.spark.sql( + "SELECT map(" + "CAST('2020-01-01 00:00:00.123456700' AS TIMESTAMP_NTZ(9)), 1, " + "CAST('2020-01-01 00:00:00.123456800' AS TIMESTAMP_NTZ(9)), 2) AS m" + ) + size_udf = udf( + lambda m: 0 if m is None else len(m), + returnType=IntegerType(), + useArrow=False, + ) + with self.assertRaises(Exception): + df.select(size_udf("m").alias("n")).collect() + def test_yearmonth_interval_type_constructor(self): self.assertEqual(YearMonthIntervalType().simpleString(), "interval year to month") self.assertEqual( @@ -2953,7 +3075,7 @@ def test_cal_interval_in_collect(self): self.spark.sql("SELECT make_interval(100, 11, 1, 1, 12, 30, 01.001001)").first()[0] -class DataTypeTests(unittest.TestCase): +class DataTypeTests(unittest.TestCase, PySparkErrorTestUtils): # regression test for SPARK-6055 def test_data_type_eq(self): lt = LongType() @@ -2987,6 +3109,161 @@ def test_varchar_type(self): self.assertEqual(v1, v3) self.assertFalse(v1 is v3) + def test_timestamp_nanos_type_precision(self): + for cls in [TimestampNTZNanosType, TimestampLTZNanosType]: + with self.subTest(cls=cls.__name__): + # The default precision is the maximum, 9 (nanoseconds). + self.assertEqual(9, cls().precision) + for p in [7, 8, 9]: + self.assertEqual(p, cls(p).precision) + # Precision 6 and below is the standard microsecond type's territory, and 10 is + # past nanoseconds; both are rejected the same way the JVM side rejects them. + for p in [-1, 0, 5, 6, 10]: + with self.assertRaises(PySparkValueError) as pe: + cls(p) + self.check_error( + exception=pe.exception, + errorClass="INVALID_TIMESTAMP_PRECISION", + messageParameters={ + "precision": str(p), + "type": cls._sqlTypeName.upper(), + }, + ) + # Non-integer precision must be rejected rather than silently accepted by the + # range check (7.5 and NaN are both between the bounds under `<` / `>`). + for p in [7.5, float("nan"), "8"]: + with self.assertRaises(PySparkValueError) as pe: + cls(p) + self.check_error( + exception=pe.exception, + errorClass="INVALID_TIMESTAMP_PRECISION", + messageParameters={ + "precision": repr(p), + "type": cls._sqlTypeName.upper(), + }, + ) + + def test_timestamp_nanos_type_string_representations(self): + # simpleString / jsonValue must match the JVM `typeName` so a schema round-trips. + self.assertEqual("timestamp_ntz(9)", TimestampNTZNanosType(9).simpleString()) + self.assertEqual("timestamp_ntz(7)", TimestampNTZNanosType(7).jsonValue()) + self.assertEqual('"timestamp_ntz(8)"', TimestampNTZNanosType(8).json()) + self.assertEqual("timestamp_ltz(9)", TimestampLTZNanosType(9).simpleString()) + self.assertEqual("timestamp_ltz(7)", TimestampLTZNanosType(7).jsonValue()) + self.assertEqual('"timestamp_ltz(8)"', TimestampLTZNanosType(8).json()) + self.assertEqual("TimestampNTZNanosType(9)", repr(TimestampNTZNanosType(9))) + self.assertEqual("TimestampLTZNanosType(7)", repr(TimestampLTZNanosType(7))) + # printSchema() / treeString() must render the precision rather than the name derived + # from the class, so these have to count as parameterized types in _get_jvm_type_name. + self.assertEqual("timestamp_ntz(9)", DataType._get_jvm_type_name(TimestampNTZNanosType(9))) + self.assertEqual("timestamp_ltz(7)", DataType._get_jvm_type_name(TimestampLTZNanosType(7))) + self.assertIn( + "|-- ts: timestamp_ntz(8) (nullable = true)", + StructType([StructField("ts", TimestampNTZNanosType(8))]).treeString(), + ) + + def test_timestamp_nanos_type_equality(self): + self.assertEqual(TimestampNTZNanosType(9), TimestampNTZNanosType(9)) + self.assertEqual(TimestampNTZNanosType(), TimestampNTZNanosType(9)) + self.assertNotEqual(TimestampNTZNanosType(9), TimestampNTZNanosType(7)) + # NTZ and LTZ are distinct types at the same precision, and neither equals the microsecond + # type whose name they parameterize. + self.assertNotEqual(TimestampNTZNanosType(9), TimestampLTZNanosType(9)) + self.assertNotEqual(TimestampNTZNanosType(9), TimestampNTZType()) + self.assertNotEqual(TimestampLTZNanosType(9), TimestampType()) + # Distinct types must not collide as dict keys / in sets. + self.assertEqual( + 3, + len( + { + TimestampNTZNanosType(9), + TimestampNTZNanosType(7), + TimestampLTZNanosType(9), + TimestampNTZNanosType(9), + } + ), + ) + for t in [TimestampNTZNanosType(8), TimestampLTZNanosType(8)]: + self.assertEqual(t, pickle.loads(pickle.dumps(t))) + + def test_timestamp_nanos_type_from_json(self): + from pyspark.sql.types import _parse_datatype_json_value + + # Mirrors DataType.parseDataType in sql/api: 7-9 are the nanosecond types, 6 is the + # standard microsecond type, everything else is rejected. + for name, expected in [ + ("timestamp_ntz(7)", TimestampNTZNanosType(7)), + ("timestamp_ntz(8)", TimestampNTZNanosType(8)), + ("timestamp_ntz(9)", TimestampNTZNanosType(9)), + ("timestamp_ltz(7)", TimestampLTZNanosType(7)), + ("timestamp_ltz( 9 )", TimestampLTZNanosType(9)), + ("timestamp_ntz(6)", TimestampNTZType()), + ("timestamp_ltz(6)", TimestampType()), + ("timestamp_ntz", TimestampNTZType()), + ]: + with self.subTest(name=name): + self.assertEqual(expected, _parse_datatype_json_value(name)) + + for name, sql_type in [ + ("timestamp_ntz(5)", "TIMESTAMP_NTZ"), + ("timestamp_ntz(10)", "TIMESTAMP_NTZ"), + ("timestamp_ltz(0)", "TIMESTAMP_LTZ"), + ]: + with self.subTest(name=name): + with self.assertRaises(PySparkValueError) as pe: + _parse_datatype_json_value(name) + self.check_error( + exception=pe.exception, + errorClass="INVALID_TIMESTAMP_PRECISION", + messageParameters={ + "precision": name[name.index("(") + 1 : -1], + "type": sql_type, + }, + ) + + # The type-name regexes are fully anchored (fullmatch), matching the JVM extractor, so + # trailing junk is not silently accepted as a valid nanosecond type. + for name in ["timestamp_ntz(9)garbage", "timestamp_ltz(9) ", "xtimestamp_ntz(9)"]: + with self.subTest(name=name): + with self.assertRaises(PySparkValueError): + _parse_datatype_json_value(name) + + def test_timestamp_nanos_type_nested_json_round_trip(self): + from pyspark.sql.types import _parse_datatype_json_string + + schema = StructType( + [ + StructField("ntz", TimestampNTZNanosType(9)), + StructField("ltz", TimestampLTZNanosType(7), False), + StructField("arr", ArrayType(TimestampNTZNanosType(8))), + StructField("map", MapType(TimestampLTZNanosType(9), TimestampNTZNanosType(7))), + StructField("nested", StructType([StructField("a", TimestampLTZNanosType(8))])), + ] + ) + self.assertEqual(schema, _parse_datatype_json_string(schema.json())) + self.assertEqual( + "struct," + "map:map,nested:struct>", + schema.simpleString(), + ) + + def test_timestamp_nanos_type_internal_conversion(self): + # The external Python value is datetime.datetime, so the internal representation is epoch + # microseconds -- identical to the microsecond types, whose conversion these mirror. + naive = datetime.datetime(2020, 1, 2, 3, 4, 5, 123456) + ntz = TimestampNTZNanosType(9) + self.assertEqual(TimestampNTZType().toInternal(naive), ntz.toInternal(naive)) + self.assertEqual(naive, ntz.fromInternal(ntz.toInternal(naive))) + + aware = datetime.datetime(2020, 1, 2, 3, 4, 5, 123456, tzinfo=datetime.timezone.utc) + ltz = TimestampLTZNanosType(7) + self.assertEqual(TimestampType().toInternal(aware), ltz.toInternal(aware)) + + for t in [ntz, ltz]: + self.assertTrue(t.needConversion()) + self.assertIsNone(t.toInternal(None)) + self.assertIsNone(t.fromInternal(None)) + # regression test for SPARK-10392 def test_datetype_equal_zero(self): dt = DateType() @@ -3149,6 +3426,8 @@ def __init__(self, **kwargs): (datetime.time(1, 0, 0), TimeType()), (datetime.datetime(2000, 1, 2, 3, 4), DateType()), (datetime.datetime(2000, 1, 2, 3, 4), TimestampType()), + (datetime.datetime(2000, 1, 2, 3, 4), TimestampNTZNanosType(9)), + (datetime.datetime(2000, 1, 2, 3, 4), TimestampLTZNanosType(7)), # Array ([], ArrayType(IntegerType())), (["1", None], ArrayType(StringType(), containsNull=True)), @@ -3213,6 +3492,8 @@ def __init__(self, **kwargs): ("2000-01-02", DateType(), TypeError), ("23:59:59", TimeType(), TypeError), (946811040, TimestampType(), TypeError), + (946811040, TimestampNTZNanosType(9), TypeError), + ("2000-01-02 03:04:05.123456789", TimestampLTZNanosType(9), TypeError), # Array (["1", None], ArrayType(StringType(), containsNull=False), ValueError), ([1, "2"], ArrayType(IntegerType()), TypeError), diff --git a/python/pyspark/sql/tests/test_utils.py b/python/pyspark/sql/tests/test_utils.py index e4f9420f3bc3d..773d0bd0f2da8 100644 --- a/python/pyspark/sql/tests/test_utils.py +++ b/python/pyspark/sql/tests/test_utils.py @@ -43,6 +43,7 @@ StringType, StructField, StructType, + TimestampNTZNanosType, ) from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( @@ -1872,6 +1873,16 @@ def test_assert_schema_equal_with_decimal_types(self): with self.assertRaises(PySparkAssertionError): assertSchemaEqual(s1, s2) + def test_assert_schema_equal_with_timestamp_nanos_types(self): + """Test assertSchemaEqual with nanosecond timestamp types of different precision + (SPARK-57462): the precision must be compared even under ignoreNullable, like decimal.""" + s1 = StructType([StructField("ts", TimestampNTZNanosType(9), True)]) + # Same precision - should pass, including with the ignoreNullable default. + assertSchemaEqual(s1, StructType([StructField("ts", TimestampNTZNanosType(9), False)])) + # Different precision - should fail rather than be treated as equal. + with self.assertRaises(PySparkAssertionError): + assertSchemaEqual(s1, StructType([StructField("ts", TimestampNTZNanosType(7), True)])) + class UtilsTests(UtilsTestsMixin, ReusedSQLTestCase): pass diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py index e777325d5c356..47d5a251c77f7 100644 --- a/python/pyspark/sql/types.py +++ b/python/pyspark/sql/types.py @@ -22,6 +22,7 @@ import decimal import json import math +import operator import os import re import sys @@ -90,6 +91,8 @@ "TimeType", "TimestampType", "TimestampNTZType", + "TimestampNTZNanosType", + "TimestampLTZNanosType", "DecimalType", "DoubleType", "Geography", @@ -228,6 +231,7 @@ def _get_jvm_type_name(cls, dataType: "DataType") -> str: DayTimeIntervalType, YearMonthIntervalType, TimeType, + AnyTimestampNanoType, ), ): return dataType.simpleString() @@ -484,6 +488,142 @@ def fromInternal(self, ts: int) -> datetime.datetime: ) +class AnyTimestampNanoType(DatetimeType): + """ + Super class of the nanosecond-capable timestamp data types + :class:`TimestampNTZNanosType` and :class:`TimestampLTZNanosType`. + + .. versionadded:: 4.4.0 + """ + + MIN_PRECISION: int = 7 + MAX_PRECISION: int = 9 + DEFAULT_PRECISION: int = 9 + # Precision of the standard microsecond timestamp types, which the parameterized DDL / JSON + # type names also accept (``timestamp_ntz(6)`` / ``timestamp_ltz(6)``). + MICROS_PRECISION: int = 6 + + # Set by each subclass to the SQL type name used in the DDL / JSON representation, e.g. + # "timestamp_ntz". Also used, upper-cased, in the invalid-precision error message. + _sqlTypeName: str = "" + + def __init__(self, precision: int = DEFAULT_PRECISION): + # Reject non-integer precision (e.g. 7.5 or float("nan")), which would otherwise slip + # through the range comparison below. operator.index accepts any integer-like value + # (including a NumPy integer) and rejects the rest with a TypeError. + try: + precision = operator.index(precision) + except TypeError: + raise PySparkValueError( + errorClass="INVALID_TIMESTAMP_PRECISION", + messageParameters={ + "precision": repr(precision), + "type": self._sqlTypeName.upper(), + }, + ) + if precision < self.MIN_PRECISION or precision > self.MAX_PRECISION: + raise PySparkValueError( + errorClass="INVALID_TIMESTAMP_PRECISION", + messageParameters={ + "precision": str(precision), + "type": self._sqlTypeName.upper(), + }, + ) + self.precision = precision + + def needConversion(self) -> bool: + return True + + def simpleString(self) -> str: + return "%s(%d)" % (self._sqlTypeName, self.precision) + + def jsonValue(self) -> str: + return "%s(%d)" % (self._sqlTypeName, self.precision) + + def __repr__(self) -> str: + return "%s(%d)" % (type(self).__name__, self.precision) + + +class TimestampNTZNanosType(AnyTimestampNanoType): + """Timestamp (datetime.datetime) data type without timezone information, with + nanosecond-capable fractional-second precision (7 to 9 digits). + + Parameters + ---------- + precision : int, optional + Number of digits of fractional seconds, one of 7, 8 or 9 (default: 9). + + Notes + ----- + These types are behind the ``spark.sql.timestampNanosTypes.enabled`` preview flag (disabled + by default); using them while it is off raises an error. + + ``datetime.datetime`` is microsecond-resolution, so values crossing the Python boundary as + ``datetime.datetime`` -- :meth:`DataFrame.collect`, :meth:`DataFrame.toLocalIterator`, and + Python UDF arguments -- are truncated to microseconds, as are ``datetime.datetime`` values + supplied to :meth:`SparkSession.createDataFrame` from Python lists/rows. The value stored by + Spark keeps full precision; only this Python boundary is microsecond-resolution. A ``map`` + 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. + + .. versionadded:: 4.4.0 + """ + + _sqlTypeName = "timestamp_ntz" + + def toInternal(self, dt: datetime.datetime) -> int: + # Mirrors TimestampNTZType.toInternal: the value is on the UTC grid and carries no zone. + if dt is not None: + seconds = calendar.timegm(dt.timetuple()) + return int(seconds) * 1000000 + dt.microsecond + + def fromInternal(self, ts: int) -> datetime.datetime: + if ts is not None: + # using int to avoid precision loss in float + return datetime.datetime.fromtimestamp(ts // 1000000, datetime.timezone.utc).replace( + microsecond=ts % 1000000, tzinfo=None + ) + + +class TimestampLTZNanosType(AnyTimestampNanoType): + """Timestamp (datetime.datetime) data type with local timezone semantics, with + nanosecond-capable fractional-second precision (7 to 9 digits). + + Parameters + ---------- + precision : int, optional + Number of digits of fractional seconds, one of 7, 8 or 9 (default: 9). + + Notes + ----- + Carries the same microsecond-only Python boundary as :class:`TimestampNTZNanosType`; see the + notes there. + + .. versionadded:: 4.4.0 + """ + + _sqlTypeName = "timestamp_ltz" + + def toInternal(self, dt: datetime.datetime) -> int: + # Mirrors TimestampType.toInternal: an aware value is converted through UTC, a naive one + # is interpreted in the local time zone. + if dt is not None: + seconds = ( + calendar.timegm(dt.utctimetuple()) if dt.tzinfo else time.mktime(dt.timetuple()) + ) + return int(seconds) * 1000000 + dt.microsecond + + def fromInternal(self, ts: int) -> datetime.datetime: + if ts is not None: + # using int to avoid precision loss in float + return datetime.datetime.fromtimestamp(ts // 1000000).replace(microsecond=ts % 1000000) + + class DecimalType(FractionalType): """Decimal (decimal.Decimal) data type. @@ -2286,6 +2426,8 @@ def fromWKB(cls, wkb: bytes, srid: int) -> "Geometry": _INTERVAL_DAYTIME = re.compile(r"interval (day|hour|minute|second)( to (day|hour|minute|second))?") _INTERVAL_YEARMONTH = re.compile(r"interval (year|month)( to (year|month))?") _TIME = re.compile(r"time\(\s*(\d+)\s*\)") +_TIMESTAMP_NTZ_PRECISION = re.compile(r"timestamp_ntz\(\s*(\d+)\s*\)") +_TIMESTAMP_LTZ_PRECISION = re.compile(r"timestamp_ltz\(\s*(\d+)\s*\)") _GEOMETRY = re.compile(r"^geometry$") _GEOMETRY_CRS = re.compile(r"geometry\(\s*([\w]+:-?[\w]+)\s*\)") _GEOGRAPHY = re.compile(r"^geography$") @@ -2418,6 +2560,18 @@ def _parse_datatype_json_string(json_string: str) -> DataType: return _parse_datatype_json_value(json.loads(json_string)) +def _parse_parameterized_timestamp_type(precision: int, ntz: bool) -> DataType: + """Maps a parameterized ``timestamp_ntz(p)`` / ``timestamp_ltz(p)`` type name to a type. + + Mirrors ``DataType.parseDataType`` in ``sql/api``: precision 6 denotes the standard + microsecond type, 7 to 9 the nanosecond-capable types, and any other precision is rejected + with ``INVALID_TIMESTAMP_PRECISION`` (raised by the nanosecond type's constructor). + """ + if precision == AnyTimestampNanoType.MICROS_PRECISION: + return TimestampNTZType() if ntz else TimestampType() + return TimestampNTZNanosType(precision) if ntz else TimestampLTZNanosType(precision) + + def _parse_datatype_json_value( # type: ignore[return] json_value: Union[dict, str], fieldPath: str = "", @@ -2434,6 +2588,10 @@ def _parse_datatype_json_value( # type: ignore[return] return DecimalType(int(m.group(1)), int(m.group(2))) elif m := _TIME.match(json_value): return TimeType(int(m.group(1))) + elif m := _TIMESTAMP_NTZ_PRECISION.fullmatch(json_value): + return _parse_parameterized_timestamp_type(int(m.group(1)), ntz=True) + elif m := _TIMESTAMP_LTZ_PRECISION.fullmatch(json_value): + return _parse_parameterized_timestamp_type(int(m.group(1)), ntz=False) elif m := _INTERVAL_DAYTIME.match(json_value): inverted_fields = DayTimeIntervalType._inverted_fields first_field = inverted_fields.get(m.group(1)) @@ -2883,6 +3041,38 @@ def _has_type(dt: DataType, dts: Union[type, Tuple[type, ...]]) -> bool: return False +def _first_timestamp_nanos_map_key_type(dt: DataType) -> Optional["DataType"]: + """Return the key type of the first map (depth-first) whose key carries a nanosecond timestamp + type, or ``None`` if ``dt`` contains no such map. + + Such keys cannot be represented on the Python conversion path: a ``datetime.datetime`` key is + microsecond-resolution, so two nanosecond keys can collapse to one map entry (SPARK-57462). + Callers reject this schema shape up front rather than let an entry be silently dropped. Unlike + a scalar / array / struct-field nanosecond value, which converts fine (truncated to micros), + only the map-key position is unsafe, so this is narrower than ``_has_type``. The returned key + type is reported in the error message, mirroring the JVM ``EvaluatePython.toJava`` twin, which + reports ``mt.keyType.sql``. + """ + if isinstance(dt, MapType): + if _has_type(dt.keyType, AnyTimestampNanoType): + return dt.keyType + return _first_timestamp_nanos_map_key_type( + dt.keyType + ) or _first_timestamp_nanos_map_key_type(dt.valueType) + elif isinstance(dt, ArrayType): + return _first_timestamp_nanos_map_key_type(dt.elementType) + elif isinstance(dt, StructType): + for field in dt.fields: + found = _first_timestamp_nanos_map_key_type(field.dataType) + if found is not None: + return found + return None + elif isinstance(dt, UserDefinedType): + return _first_timestamp_nanos_map_key_type(dt.sqlType()) + else: + return None + + @overload def _merge_type(a: StructType, b: StructType, name: Optional[str] = None) -> StructType: ... @@ -3079,6 +3269,8 @@ def convert_struct(obj: Any) -> Optional[Tuple]: TimeType: (datetime.time,), TimestampType: (datetime.datetime,), TimestampNTZType: (datetime.datetime,), + TimestampNTZNanosType: (datetime.datetime,), + TimestampLTZNanosType: (datetime.datetime,), DayTimeIntervalType: (datetime.timedelta,), ArrayType: (list, tuple, array), MapType: (dict,), diff --git a/python/pyspark/testing/utils.py b/python/pyspark/testing/utils.py index e3af3fe40c904..f5c261740a6da 100644 --- a/python/pyspark/testing/utils.py +++ b/python/pyspark/testing/utils.py @@ -41,7 +41,7 @@ from pyspark.sql import Row from pyspark.sql.dataframe import DataFrame from pyspark.sql.functions import col, when -from pyspark.sql.types import StructField, StructType, VariantVal +from pyspark.sql.types import AnyTimestampNanoType, StructField, StructType, VariantVal __all__ = ["assertDataFrameEqual", "assertSchemaEqual"] @@ -701,6 +701,11 @@ def compare_datatypes_ignore_nullable(dt1: Any, dt2: Any): elif dt1.typeName() == "decimal": # Fix for SPARK-51062: Compare precision and scale for decimal types return dt1.precision == dt2.precision and dt1.scale == dt2.scale + elif isinstance(dt1, AnyTimestampNanoType): + # SPARK-57462: the nanosecond timestamp types carry a fractional-second + # precision, like decimal above; the type name alone does not distinguish + # timestamp_ntz(7) from timestamp_ntz(9), so compare the precision too. + return dt1.precision == dt2.precision elif dt1.typeName() == "struct": return compare_schemas_ignore_nullable(dt1, dt2) else: diff --git a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeApiOps.scala b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeApiOps.scala index 76bbf0c6f4a83..8431bebd117ed 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeApiOps.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/types/ops/TimestampNanosTypeApiOps.scala @@ -78,6 +78,34 @@ abstract class TimestampNanosTypeApiOps extends TypeApiOps with DataTypeErrorsBa // column to STRING_TYPE for consistency, mirroring the reference TimeType ops. override def thriftTypeName: Option[String] = Some("STRING_TYPE") + // ==================== Python Interop ==================== + + // The external Python value is `datetime.datetime`, which is microsecond-resolution, so PySpark + // represents these types as epoch microseconds (TimestampNTZNanosType / TimestampLTZNanosType in + // pyspark/sql/types.py) and sub-microsecond digits never cross the Py4J boundary in either + // direction. That is the documented Python/UDF microsecond-only limitation (SPARK-57808). + // Lossless Arrow/pandas value conversion for these types (toPandas) is a pending PySpark + // follow-up, so it is not yet an alternative that preserves the extra digits. + override def needConversionInPython: Option[Boolean] = Some(true) + + // Python hands us epoch microseconds; rebuild the internal value with a zero sub-microsecond + // remainder. The reverse direction is EvaluatePython.toJava, which yields `epochMicros`. + // + // The gate is enforced eagerly, when the converter is built, so that the classic PySpark + // explicit-schema path (SparkSession.applySchemaToPythonRDD, which calls + // EvaluatePython.makeFromJava rather than the guarded getEncoder) and Python UDF nanosecond + // return types cannot execute with spark.sql.timestampNanosTypes.enabled = false. Mirrors the + // guard on getEncoder below. + override def makeFromJava: Option[Any => Any] = { + DataTypeErrors.checkTimestampNanosTypesEnabled() + Some((obj: Any) => + nullSafeConvert(obj) { + case c: Long => TimestampNanosVal.fromParts(c, 0.toShort) + // Py4J serializes values between MIN_INT and MAX_INT as Ints, not Longs + case c: Int => TimestampNanosVal.fromParts(c.toLong, 0.toShort) + }) + } + // ==================== Row Encoding ==================== // Honor the spark.sql.timestampNanosTypes.enabled gate just like the legacy diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala index 3b9c2a3e69cd2..ec46d144ba414 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala @@ -24,7 +24,7 @@ import scala.jdk.CollectionConverters._ import net.razorvine.pickle.{IObjectPickler, Opcodes, Pickler} -import org.apache.spark.SparkIllegalArgumentException +import org.apache.spark.{SparkIllegalArgumentException, SparkRuntimeException} import org.apache.spark.api.python.SerDeUtil import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow @@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.types.ops.TypeApiOps import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData, GenericArrayData, MapData, STUtils} import org.apache.spark.sql.types._ -import org.apache.spark.unsafe.types.{BinaryView, UTF8String, VariantVal} +import org.apache.spark.unsafe.types.{BinaryView, TimestampNanosVal, UTF8String, VariantVal} object EvaluatePython { @@ -46,6 +46,21 @@ object EvaluatePython { TypeApiOps(dt).flatMap(_.needConversionInPython) .getOrElse(needConversionInPythonDefault(dt)) + /** + * True if `dt` is, or structurally contains, a nanosecond-capable timestamp type. Used to gate + * the map-key collision check in [[toJava]] to maps whose key type can actually truncate when + * handed to Python (directly, or as a nanos field nested in a struct/array/map key). + */ + private def typeCarriesTimestampNanos(dt: DataType): Boolean = dt match { + case _: AnyTimestampNanoType => true + case ArrayType(elementType, _) => typeCarriesTimestampNanos(elementType) + case MapType(keyType, valueType, _) => + typeCarriesTimestampNanos(keyType) || typeCarriesTimestampNanos(valueType) + case StructType(fields) => fields.exists(f => typeCarriesTimestampNanos(f.dataType)) + case udt: UserDefinedType[_] => typeCarriesTimestampNanos(udt.sqlType) + case _ => false + } + private def needConversionInPythonDefault(dt: DataType): Boolean = dt match { case DateType | TimestampType | TimestampNTZType | VariantType | _: DayTimeIntervalType | _: GeometryType | _: GeographyType => true @@ -100,6 +115,21 @@ object EvaluatePython { }) values + case (map: MapData, mt: MapType) if typeCarriesTimestampNanos(mt.keyType) => + // A nanosecond timestamp key is handed to Python as epoch micros (datetime.datetime is + // microsecond-resolution), so two keys that differ only below a microsecond collapse to + // the same Python key and one entry is silently dropped. A post-hoc size check is not + // reliable: composite keys (e.g. a struct with a nanos field and a binary field) stay + // distinct in the JVM map -- BytesWrapper uses identity equality -- yet still collapse + // once pickled to Python. Scalar precision loss is tolerable, but silently changing a + // map's cardinality is not, so reject nanosecond-keyed maps on this path outright. This + // covers the Python UDF input path, where the task failure propagates; the collect path + // is guarded earlier, in Python, since a failure in the background serve thread would not + // surface (see DataFrame.collect). + throw new SparkRuntimeException( + errorClass = "TIMESTAMP_NANOS_PYTHON_MAP_KEY", + messageParameters = Map("type" -> mt.keyType.sql)) + case (map: MapData, mt: MapType) => sizeAcc.foreach(_.add(PickledSizeAccumulator.PER_VALUE_OVERHEAD)) val jmap = new java.util.HashMap[Any, Any](map.numElements()) @@ -136,6 +166,16 @@ object EvaluatePython { bytes } + case (v: TimestampNanosVal, _: AnyTimestampNanoType) => + // `datetime.datetime` is microsecond-resolution, so the sub-microsecond remainder cannot + // be represented on the Python side: hand over epoch micros and let the Python type's + // fromInternal build the datetime. Without this case the value would fall through as a + // raw TimestampNanosVal, which has no registered pickler. See the Python Interop section + // of TimestampNanosTypeApiOps for the reverse direction and SPARK-57808 for the + // documented microsecond-only Python/UDF boundary. + sizeAcc.foreach(_.addLeaf(v.epochMicros)) + v.epochMicros + case (other, _) => sizeAcc.foreach(_.addLeaf(other)) other