uros-b commented on code in PR #58418:
URL: https://github.com/apache/spark/pull/58418#discussion_r3927696742
##########
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()
Review Comment:
Map-key and flag-off tests use assertRaises(Exception)
test_timestamp_nanos_type_map_key_collision,
test_timestamp_nanos_type_map_key_python_udf_input, and
test_timestamp_nanos_type_preview_flag_off will pass on any failure. That is
how the earlier “whole schema in the error” bug slipped through.
For collect:
```
with self.assertRaises(PySparkTypeError) as pe:
df.collect()
self.check_error(
exception=pe.exception,
errorClass="TIMESTAMP_NANOS_PYTHON_MAP_KEY",
messageParameters={"type": "timestamp_ntz(9)"},
)
```
Flag-off should assert FEATURE_NOT_ENABLED. The UDF map-key path should
assert TIMESTAMP_NANOS_PYTHON_MAP_KEY (JVM SparkRuntimeException), not a
generic Exception.
##########
python/pyspark/sql/tests/test_types.py:
##########
@@ -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):
Review Comment:
UDF input test does not prove toJava produced a datetime
test_timestamp_nanos_type_python_udf_input uses lambda x: x. If toJava
handed over a raw Long (epoch micros), the UDF would return that Long,
makeFromJava would rebuild a TimestampNanosVal, and collect would still yield a
datetime. The assertion would pass either way.
Check the argument inside the UDF, for example isinstance(x,
datetime.datetime) (and that it equals the expected value).
##########
python/pyspark/sql/types.py:
##########
@@ -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"
Review Comment:
typeName() is a public API wart
TimestampNTZType.typeName() is "timestamp_ntz". These classes inherit the
default cls.__name__[:-4].lower(), so you get "timestampntznanos" /
"timestampltznanos". printSchema() is fine because _get_jvm_type_name uses
simpleString(), but df.schema["ts"].dataType.typeName() is not.
JVM typeName is timestamp_ntz($precision). Cleanest Python match is an
instance typeName() that returns simpleString(), same pattern as GeographyType.
Then compare_datatypes_ignore_nullable can keep comparing precision (or rely on
the distinct type names).
If you instead override to the bare "timestamp_ntz" / "timestamp_ltz",
assertSchemaEqual(..., ignoreNullable=True) would treat TimestampNTZType() as
equal to TimestampNTZNanosType(9) unless you also distinguish micro vs nano
there.
##########
python/pyspark/sql/pandas/types.py:
##########
@@ -77,6 +78,45 @@
metadata_key = b"SPARK::metadata::json"
+def _contains_timestamp_nanos(dt: DataType) -> bool:
Review Comment:
_contains_timestamp_nanos in pandas/types.py duplicates _has_type(...,
AnyTimestampNanoType) in types.py.
##########
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
Review Comment:
New rejection paths are untested
A lot of new code exists only to fail deterministically, and none of it is
asserted:
classic toPandas / toArrow / pandas createDataFrame with an explicit nanos
schema → UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION
Connect collect / toPandas / createDataFrame → same (parity tests are
skipped with no replacement)
Arrow Python UDF (useArrow=True) → clear error, not a pickle/EOF surprise
from_arrow_type still maps pa.timestamp('ns') to TimestampNTZType /
TimestampType. Connect collect is safe only because to_table usually supplies
the Spark schema and ArrowTableToRowsConversion then raises. toPandas needed an
extra guard because it does not go through that converter. A Connect test that
spark.sql("SELECT CAST(... AS TIMESTAMP_NTZ(9))").collect() raises the
documented error would lock this in.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
