From 228a511e1dd36202b8e2cff8081ef17f803a2420 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 6 Nov 2026 23:18:54 -0800 Subject: [PATCH 1/3] [SPARK-61102][CONNECT][RUST] Match NumPy longdouble literal class checks --- crates/pyspark-rs/src/functions.rs | 20 +++++++++--- python/tests/test_dropin_offline.py | 47 ++++++++++++++++++++++++++++- 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/crates/pyspark-rs/src/functions.rs b/crates/pyspark-rs/src/functions.rs index 74da2f4f..233d42e2 100644 --- a/crates/pyspark-rs/src/functions.rs +++ b/crates/pyspark-rs/src/functions.rs @@ -333,9 +333,10 @@ fn numpy_ndarray_ndim(obj: &Bound<'_, PyAny>) -> PyResult> { /// The literal of a NumPy scalar, typed as PySpark's `LiteralExpression._infer_type` types it: a /// bool, a signed int of 8 to 64 bits (tinyint to bigint), or a float32 or float64. A str or bytes /// scalar is `None`, for the string and binary conversions to take. Any other NumPy scalar (an -/// unsigned int, float16, longdouble, complex, datetime64 or timedelta64) is -/// UNSUPPORTED_DATA_TYPE, named by its Python type. NumPy is looked up in `sys.modules`, as a -/// NumPy scalar exists only when NumPy is imported. +/// unsigned int, float16, complex, datetime64 or timedelta64) is UNSUPPORTED_DATA_TYPE, +/// named by its Python type. A distinct longdouble scalar whose dtype maps to double fails the +/// reference's scalar-class assertion. NumPy is looked up in `sys.modules`, as a NumPy +/// scalar exists only when NumPy is imported. fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult> { let py = obj.py(); let modules = py.import("sys")?.getattr("modules")?; @@ -359,7 +360,18 @@ fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult LiteralExpression::Integer(obj.extract::()? as i32), ("i", 8) => LiteralExpression::Long(obj.extract::()?), ("f", 4) => LiteralExpression::Float(obj.extract::()? as f32), - ("f", 8) => LiteralExpression::Double(obj.extract::()?), + ("f", 8) => { + // On some platforms longdouble has an eight-byte dtype equal to float64, but + // remains a distinct scalar class. PySpark infers DoubleType then asserts the + // accepted classes before converting its value to float. + if !obj.is_instance_of::() + && !obj.is_instance(&numpy.getattr("float32")?)? + && !obj.is_instance(&numpy.getattr("float64")?)? + { + return Err(pyo3::exceptions::PyAssertionError::new_err(())); + } + LiteralExpression::Double(obj.extract::()?) + } _ => { let type_name = obj .get_type() diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 05be95c1..f0182860 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -9596,7 +9596,6 @@ def test_numpy_scalars_of_other_types_are_unsupported_as_pyspark(mock_connect): (np.uint8(5), "uint8"), (np.uint64(5), "uint64"), (np.float16(1.5), "float16"), - (np.longdouble(1.5), "longdouble"), (np.complex128(1 + 2j), "complex128"), (np.datetime64("2020-01-01"), "datetime64"), (np.timedelta64(1, "D"), "timedelta64"), @@ -9611,6 +9610,51 @@ def test_numpy_scalars_of_other_types_are_unsupported_as_pyspark(mock_connect): assert error.value.getMessageParameters() == {"dtype": "uint8"} +def test_numpy_longdouble_scalar_matches_dtype_and_class_checks_as_pyspark(): + np = pytest.importorskip("numpy") + from pyspark.errors import PySparkTypeError + from pyspark.sql import functions as F + + value = np.longdouble(1.5) + if value.dtype == np.dtype("float64"): + # A dtype mapping to DoubleType does not suffice: PySpark also asserts that its + # scalar is a float, np.float32 or np.float64. Distinct eight-byte longdouble + # classes (such as on macOS arm64) fail that assertion without a message. + if isinstance(value, (float, np.float32, np.float64)): + assert repr(F.lit(value)) == "Column<'1.5'>" + else: + with pytest.raises(AssertionError) as error: + F.lit(value) + assert error.value.args == () + else: + # Wider longdouble dtypes do not map to a Spark type at all. + with pytest.raises(PySparkTypeError) as error: + F.lit(value) + assert error.value.getCondition() == "UNSUPPORTED_DATA_TYPE" + assert error.value.getMessageParameters() == {"data_type": type(value).__name__} + + +def test_numpy_longdouble_array_matches_dtype_and_element_checks_as_pyspark(): + np = pytest.importorskip("numpy") + from pyspark.errors import PySparkTypeError + from pyspark.sql import functions as F + + value = np.array([1.5], dtype=np.longdouble) + if value.dtype == np.dtype("float64"): + if isinstance(value[0], (float, np.float32, np.float64)): + assert repr(F.lit(value)).lower() == "column<'cast(array(1.5) as array)'>" + else: + # A mapped array dtype still builds its literals through each scalar's checks. + with pytest.raises(AssertionError) as error: + F.lit(value) + assert error.value.args == () + else: + with pytest.raises(PySparkTypeError) as error: + F.lit(value) + assert error.value.getCondition() == "UNSUPPORTED_NUMPY_ARRAY_SCALAR" + assert error.value.getMessageParameters() == {"dtype": value.dtype.name} + + # --------------------------------------------------------------------------- 4.3.0 gaps @@ -9988,6 +10032,7 @@ def test_lit_of_a_numpy_array_is_an_array_of_its_type_as_pyspark(): for value, expected in ( (np.array([1, 2]), "cast(array(1, 2) as array)"), (np.array([1.5], dtype=np.float32), "cast(array(1.5) as array)"), + (np.array([1.5], dtype=np.float64), "cast(array(1.5) as array)"), (np.array(["x"]), "cast(array(x) as array)"), ): assert repr(F.lit(value)).lower() == f"column<'{expected}'>" From 8c216dfd97c74a9e4d5f2eb81b39eb814cea4dce Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 6 Nov 2026 23:29:44 -0800 Subject: [PATCH 2/3] [SPARK-61102][CONNECT][RUST] Apply scalar class checks consistently to numeric aliases --- crates/pyspark-rs/src/functions.rs | 50 +++++++++++++++------- python/tests/test_dropin_offline.py | 66 ++++++++++++++++++++++------- 2 files changed, 86 insertions(+), 30 deletions(-) diff --git a/crates/pyspark-rs/src/functions.rs b/crates/pyspark-rs/src/functions.rs index 233d42e2..14ac379f 100644 --- a/crates/pyspark-rs/src/functions.rs +++ b/crates/pyspark-rs/src/functions.rs @@ -334,9 +334,10 @@ fn numpy_ndarray_ndim(obj: &Bound<'_, PyAny>) -> PyResult> { /// bool, a signed int of 8 to 64 bits (tinyint to bigint), or a float32 or float64. A str or bytes /// scalar is `None`, for the string and binary conversions to take. Any other NumPy scalar (an /// unsigned int, float16, complex, datetime64 or timedelta64) is UNSUPPORTED_DATA_TYPE, -/// named by its Python type. A distinct longdouble scalar whose dtype maps to double fails the -/// reference's scalar-class assertion. NumPy is looked up in `sys.modules`, as a NumPy -/// scalar exists only when NumPy is imported. +/// named by its Python type. A supported dtype with a scalar class outside PySpark's accepted +/// classes raises a bare `AssertionError`, such as a distinct `longlong`, or an eight-byte +/// `longdouble` on macOS arm64 and Windows. Wider `longdouble` dtypes are unsupported. +/// NumPy is looked up in `sys.modules`, as a NumPy scalar exists only when NumPy is imported. fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult> { let py = obj.py(); let modules = py.import("sys")?.getattr("modules")?; @@ -353,6 +354,36 @@ fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult &["int8"], + ("i", 2) => &["int16", "int8"], + ("i", 4) => &["int32", "int16", "int8"], + ("i", 8) => &["int64", "int32", "int16", "int8"], + ("f", 4) => &["float32"], + ("f", 8) => &["float64", "float32"], + _ => &[], + }; + if !scalar_classes.is_empty() { + let mut accepted = match kind.as_str() { + "i" => obj.is_instance_of::(), + "f" => obj.is_instance_of::(), + _ => false, + }; + if !accepted { + for class in scalar_classes { + if obj.is_instance(&numpy.getattr(*class)?)? { + accepted = true; + break; + } + } + } + if !accepted { + return Err(pyo3::exceptions::PyAssertionError::new_err(())); + } + } let literal = match (kind.as_str(), itemsize) { ("U", _) | ("S", _) => return Ok(None), ("i", 1) => LiteralExpression::Byte(obj.extract::()? as i32), @@ -360,18 +391,7 @@ fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult LiteralExpression::Integer(obj.extract::()? as i32), ("i", 8) => LiteralExpression::Long(obj.extract::()?), ("f", 4) => LiteralExpression::Float(obj.extract::()? as f32), - ("f", 8) => { - // On some platforms longdouble has an eight-byte dtype equal to float64, but - // remains a distinct scalar class. PySpark infers DoubleType then asserts the - // accepted classes before converting its value to float. - if !obj.is_instance_of::() - && !obj.is_instance(&numpy.getattr("float32")?)? - && !obj.is_instance(&numpy.getattr("float64")?)? - { - return Err(pyo3::exceptions::PyAssertionError::new_err(())); - } - LiteralExpression::Double(obj.extract::()?) - } + ("f", 8) => LiteralExpression::Double(obj.extract::()?), _ => { let type_name = obj .get_type() diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index f0182860..749d322c 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -9616,18 +9616,18 @@ def test_numpy_longdouble_scalar_matches_dtype_and_class_checks_as_pyspark(): from pyspark.sql import functions as F value = np.longdouble(1.5) + if isinstance(value, (float, np.float32, np.float64)): + pytest.skip("longdouble aliases a supported float scalar class here") if value.dtype == np.dtype("float64"): # A dtype mapping to DoubleType does not suffice: PySpark also asserts that its # scalar is a float, np.float32 or np.float64. Distinct eight-byte longdouble - # classes (such as on macOS arm64) fail that assertion without a message. - if isinstance(value, (float, np.float32, np.float64)): - assert repr(F.lit(value)) == "Column<'1.5'>" - else: - with pytest.raises(AssertionError) as error: - F.lit(value) - assert error.value.args == () + # classes (such as on macOS arm64 and Windows) fail without a message. + with pytest.raises(AssertionError) as error: + F.lit(value) + assert error.value.args == () else: - # Wider longdouble dtypes do not map to a Spark type at all. + # Linux CI normally has a wider longdouble, which exercises the unsupported + # dtype path; the distinct integer aliases below exercise the class check there. with pytest.raises(PySparkTypeError) as error: F.lit(value) assert error.value.getCondition() == "UNSUPPORTED_DATA_TYPE" @@ -9640,14 +9640,13 @@ def test_numpy_longdouble_array_matches_dtype_and_element_checks_as_pyspark(): from pyspark.sql import functions as F value = np.array([1.5], dtype=np.longdouble) + if isinstance(value[0], (float, np.float32, np.float64)): + pytest.skip("longdouble aliases a supported float scalar class here") if value.dtype == np.dtype("float64"): - if isinstance(value[0], (float, np.float32, np.float64)): - assert repr(F.lit(value)).lower() == "column<'cast(array(1.5) as array)'>" - else: - # A mapped array dtype still builds its literals through each scalar's checks. - with pytest.raises(AssertionError) as error: - F.lit(value) - assert error.value.args == () + # A mapped array dtype still builds its literals through each scalar's checks. + with pytest.raises(AssertionError) as error: + F.lit(value) + assert error.value.args == () else: with pytest.raises(PySparkTypeError) as error: F.lit(value) @@ -9655,6 +9654,43 @@ def test_numpy_longdouble_array_matches_dtype_and_element_checks_as_pyspark(): assert error.value.getMessageParameters() == {"dtype": value.dtype.name} +@pytest.mark.parametrize("name", [ + "byte", "short", "intc", "int_", "intp", "longlong", "single", "double", +]) +@pytest.mark.parametrize("as_array", [False, True], ids=["scalar", "array"]) +def test_numpy_numeric_aliases_match_reference_scalar_classes(name, as_array): + np = pytest.importorskip("numpy") + from pyspark.sql import functions as F + + scalar_type = getattr(np, name) + value = scalar_type(1.5 if np.issubdtype(scalar_type, np.floating) else 5) + kind, size = value.dtype.kind, value.dtype.itemsize + if kind == "i": + # These are the classes asserted by LiteralExpression for the inferred width. + accepted_classes = (int,) + tuple( + getattr(np, f"int{width}") for width in (8, 16, 32, 64) if width <= size * 8 + ) + spark_type = {1: "tinyint", 2: "smallint", 4: "int", 8: "bigint"}[size] + rendered = "5" + else: + accepted_classes = (float, np.float32) + ((np.float64,) if size == 8 else ()) + spark_type = {4: "float", 8: "double"}[size] + rendered = "1.5" + literal = np.array([value], dtype=scalar_type) if as_array else value + if not isinstance(value, accepted_classes): + # Distinct longlong scalars have int64 dtype but fail the class assertion on + # Linux/macOS; platforms where it aliases int64 exercise the success case. + with pytest.raises(AssertionError) as error: + F.lit(literal) + assert error.value.args == () + elif as_array: + assert repr(F.lit(literal)).lower() == ( + f"column<'cast(array({rendered}) as array<{spark_type}>)'>" + ) + else: + assert repr(F.lit(literal)) == f"Column<'{rendered}'>" + + # --------------------------------------------------------------------------- 4.3.0 gaps From 07b5866057d37db1b8d8f871d9102e1ceaa4867b Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 6 Nov 2026 23:37:57 -0800 Subject: [PATCH 3/3] Cover optional NumPy long scalar aliases --- python/tests/test_dropin_offline.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 749d322c..7e2c0aba 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -9655,14 +9655,17 @@ def test_numpy_longdouble_array_matches_dtype_and_element_checks_as_pyspark(): @pytest.mark.parametrize("name", [ - "byte", "short", "intc", "int_", "intp", "longlong", "single", "double", + "byte", "short", "intc", "int_", "intp", "long", "longlong", "single", "double", ]) @pytest.mark.parametrize("as_array", [False, True], ids=["scalar", "array"]) def test_numpy_numeric_aliases_match_reference_scalar_classes(name, as_array): np = pytest.importorskip("numpy") from pyspark.sql import functions as F - scalar_type = getattr(np, name) + # Read real scalar exports without NumPy's deprecated builtin-alias lookup warnings. + scalar_type = vars(np).get(name) + if scalar_type is None or not issubclass(scalar_type, np.generic): + pytest.skip(f"numpy.{name} is not a NumPy scalar class in this NumPy") value = scalar_type(1.5 if np.issubdtype(scalar_type, np.floating) else 5) kind, size = value.dtype.kind, value.dtype.itemsize if kind == "i":