stevomitric commented on code in PR #58418:
URL: https://github.com/apache/spark/pull/58418#discussion_r3944404621
##########
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:
added a probe UDF asserting the arg is a datetime
--
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]
