Merge pull request #68 from influxdata/crepererum/py_typing_optional

feat: support `Optional[T]` in Python

Author: crepererum

guests/python/src/lib.rs

@@ -4,7 +4,7 @@
 //! [CPython]: https://www.python.org/
 //! [`pyo3`]: https://pyo3.rs/
 use std::any::Any;
-use std::ops::ControlFlow;
+use std::ops::{ControlFlow, Range};
 use std::sync::Arc;
 
 use arrow::datatypes::DataType;
@@ -25,6 +25,9 @@ mod error;
 mod inspect;
 mod signature;
 
+/// Supported Python version range.
+const PYTHON_VERSION_RANGE: Range<(u8, u8, u8)> = (3, 14, 0)..(3, 15, 0);
+
 /// A test UDF that demonstrate that we can call Python.
 #[derive(Debug)]
 struct PythonScalarUDF {
@@ -267,10 +270,21 @@ fn root() -> Option<Vec<u8>> {
 ///   <no Python frame>
 /// ```
 ///
+/// This also checks if the running Python version is supported.
+///
 ///
 /// [Python Standard Library]: https://docs.python.org/3/library/index.html
 fn init_python() {
     Python::initialize();
+
+    Python::attach(|py| {
+        let version_info = py.version_info();
+        let version_tuple = (version_info.major, version_info.minor, version_info.patch);
+        assert!(
+            PYTHON_VERSION_RANGE.contains(&version_tuple),
+            "Unsupported python version: {version_tuple:?}, supported range is {PYTHON_VERSION_RANGE:?}",
+        );
+    });
 }
 
 /// Generate UDFs from given Python string.

guests/python/src/signature.rs

@@ -51,6 +51,9 @@ pub(crate) enum PythonType {
 /// - <https://docs.python.org/3/library/stdtypes.html#the-null-object>
 ///
 /// Nullable types are represented using a union with another type, e.g. `int | None`.
+///
+/// There used to be an older representation too: `typing.Optional[int]`. As of Python 3.14, this results in the same
+/// representation as `int | None`. See <https://docs.python.org/3.14/whatsnew/3.14.html#typing>. So we support both.
 #[derive(Debug)]
 pub(crate) struct PythonNullableType {
     /// Python type.

host/tests/integration_tests/python/runtime/null_handling.rs

@@ -26,7 +26,7 @@ def add(x: int, y: int) -> int:
 }
 
 #[tokio::test(flavor = "multi_thread")]
-async fn test_first_arg_optional() {
+async fn test_first_arg_optional_union() {
     const CODE: &str = "
 def add(x: int | None, y: int) -> int:
     if x is None:
@@ -42,7 +42,25 @@ def add(x: int | None, y: int) -> int:
 }
 
 #[tokio::test(flavor = "multi_thread")]
-async fn test_second_arg_optional() {
+async fn test_first_arg_optional_old_alias() {
+    const CODE: &str = "
+from typing import Optional
+
+def add(x: Optional[int], y: int) -> int:
+    if x is None:
+        x = 9
+    assert y is not None
+    return x + y
+";
+
+    assert_eq!(
+        xy_null_test(CODE).await.as_ref(),
+        &Int64Array::from_iter([None, Some(29), None, Some(44)]) as &dyn Array,
+    );
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_second_arg_optional_union() {
     const CODE: &str = "
 def add(x: int, y: int | None) -> int:
     assert x is not None
@@ -58,7 +76,25 @@ def add(x: int, y: int | None) -> int:
 }
 
 #[tokio::test(flavor = "multi_thread")]
-async fn test_both_args_optional() {
+async fn test_second_arg_optional_old_alias() {
+    const CODE: &str = "
+from typing import Optional
+
+def add(x: int, y: Optional[int]) -> int:
+    assert x is not None
+    if y is None:
+        y = 90
+    return x + y
+";
+
+    assert_eq!(
+        xy_null_test(CODE).await.as_ref(),
+        &Int64Array::from_iter([None, None, Some(93), Some(44)]) as &dyn Array,
+    );
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_both_args_optional_union() {
     const CODE: &str = "
 def add(x: int | None, y: int | None) -> int:
     if x is None:
@@ -75,7 +111,26 @@ def add(x: int | None, y: int | None) -> int:
 }
 
 #[tokio::test(flavor = "multi_thread")]
-async fn test_optional_passthrough() {
+async fn test_both_args_optional_old_alias() {
+    const CODE: &str = "
+from typing import Optional
+
+def add(x: Optional[int], y: Optional[int]) -> int:
+    if x is None:
+        x = 9
+    if y is None:
+        y = 90
+    return x + y
+";
+
+    assert_eq!(
+        xy_null_test(CODE).await.as_ref(),
+        &Int64Array::from_iter([Some(99), Some(29), Some(93), Some(44)]) as &dyn Array,
+    );
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_optional_passthrough_union() {
     const CODE: &str = "
 def add(x: int | None, y: int | None) -> int | None:
     if x is None or y is None:
@@ -90,7 +145,24 @@ def add(x: int | None, y: int | None) -> int | None:
 }
 
 #[tokio::test(flavor = "multi_thread")]
-async fn test_optional_flip() {
+async fn test_optional_passthrough_old_alias() {
+    const CODE: &str = "
+from typing import Optional
+
+def add(x: Optional[int], y: Optional[int]) -> Optional[int]:
+    if x is None or y is None:
+        return None
+    return x + y
+";
+
+    assert_eq!(
+        xy_null_test(CODE).await.as_ref(),
+        &Int64Array::from_iter([None, None, None, Some(44)]) as &dyn Array,
+    );
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_optional_flip_union() {
     const CODE: &str = "
 def add(x: int | None, y: int | None) -> int | None:
     if x is None:
@@ -112,6 +184,31 @@ def add(x: int | None, y: int | None) -> int | None:
     );
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_optional_flip_old_alias() {
+    const CODE: &str = "
+from typing import Optional
+
+def add(x: Optional[int], y: Optional[int]) -> Optional[int]:
+    if x is None:
+        x = 9
+    else:
+        return None
+
+    if y is None:
+        y = 90
+    else:
+        return None
+
+    return x + y
+";
+
+    assert_eq!(
+        xy_null_test(CODE).await.as_ref(),
+        &Int64Array::from_iter([Some(99), None, None, None]) as &dyn Array,
+    );
+}
+
 async fn xy_null_test(code: &str) -> ArrayRef {
     let udf = python_scalar_udf(code).await.unwrap();
     udf.invoke_with_args(ScalarFunctionArgs {