Merge pull request #159 from influxdata/crepererum/python-null

feat: Python `None`/null support

Author: crepererum

guests/python/README.md

@@ -70,6 +70,7 @@ Types are mapped to/from [Apache Arrow] as follows:
 | [`datetime`] | [`Timestamp`] w/ [`Microsecond`] and NO timezone |
 | [`float`]    | [`Float64`] |
 | [`int`]      | [`Int64`]   |
+| [`None`]     | [`Null`]   |
 | [`str`]      | [`Utf8`]    |
 | [`time`]     | [`Time64`] w/ [`Microsecond`] and NO timezone |
 | [`timedelta`]| [`Duration`]      |
@@ -226,6 +227,8 @@ There is NO other I/O available that escapes the sandbox.
 [`functools.cache`]: https://docs.python.org/3/library/functools.html#functools.cache
 [`int`]: https://docs.python.org/3/library/stdtypes.html#numeric-types-int-float-complex
 [`Int64`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#variant.Int64
+[`None`]: https://docs.python.org/3/library/constants.html#None
+[`Null`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#variant.Null
 [`time`]: https://docs.python.org/3/library/datetime.html#datetime.time
 [`Time64`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#variant.Time64
 [`timedelta`]: https://docs.python.org/3/library/datetime.html#datetime.timedelta

guests/python/src/conversion.rs

@@ -4,7 +4,7 @@ use std::{ops::ControlFlow, sync::Arc};
 use arrow::{
     array::{
         Array, ArrayRef, BinaryBuilder, BooleanBuilder, Date32Builder, DurationMicrosecondBuilder,
-        Float64Builder, Int64Builder, StringBuilder, Time64MicrosecondBuilder,
+        Float64Builder, Int64Builder, NullBuilder, StringBuilder, Time64MicrosecondBuilder,
         TimestampMicrosecondBuilder,
     },
     datatypes::{DataType, TimeUnit},
@@ -13,8 +13,8 @@ use chrono::{DateTime, Datelike, NaiveDate, TimeZone, Timelike, Utc};
 use datafusion_common::{
     cast::{
         as_binary_array, as_boolean_array, as_date32_array, as_duration_microsecond_array,
-        as_float64_array, as_int64_array, as_string_array, as_time64_microsecond_array,
-        as_timestamp_microsecond_array,
+        as_float64_array, as_int64_array, as_null_array, as_string_array,
+        as_time64_microsecond_array, as_timestamp_microsecond_array,
     },
     error::Result as DataFusionResult,
     exec_datafusion_err, exec_err,
@@ -53,6 +53,7 @@ impl PythonType {
             Self::DateTime => DataType::Timestamp(TimeUnit::Microsecond, None),
             Self::Float => DataType::Float64,
             Self::Int => DataType::Int64,
+            Self::None => DataType::Null,
             Self::Str => DataType::Utf8,
             Self::Bytes => DataType::Binary,
             Self::Date => DataType::Date32,
@@ -166,6 +167,16 @@ impl PythonType {
 
                 Ok(Box::new(it))
             }
+            Self::None => {
+                let array = as_null_array(array)?;
+
+                let val = PyNone::get(py)
+                    .into_bound_py_any(py)
+                    .map_err(|e| exec_datafusion_err!("cannot build Python None value: {e}"))?;
+                let it = std::iter::repeat_n(Some(val), array.len()).map(Ok);
+
+                Ok(Box::new(it))
+            }
             Self::Str => {
                 let array = as_string_array(array)?;
 
@@ -294,6 +305,7 @@ impl PythonType {
             Self::DateTime => Box::new(TimestampMicrosecondBuilder::with_capacity(num_rows)),
             Self::Float => Box::new(Float64Builder::with_capacity(num_rows)),
             Self::Int => Box::new(Int64Builder::with_capacity(num_rows)),
+            Self::None => Box::new(NullBuilder::new()),
             Self::Str => Box::new(StringBuilder::with_capacity(num_rows, 1024)),
             Self::Bytes => Box::new(BinaryBuilder::with_capacity(num_rows, 1024)),
             Self::Date => Box::new(Date32Builder::with_capacity(num_rows)),
@@ -467,6 +479,24 @@ impl<'py> ArrayBuilder<'py> for Int64Builder {
     }
 }
 
+impl<'py> ArrayBuilder<'py> for NullBuilder {
+    fn push(&mut self, val: Bound<'py, PyAny>) -> DataFusionResult<()> {
+        val.cast_exact::<PyNone>().map_err(|_| {
+            exec_datafusion_err!("expected `None` but got {}", py_representation(&val))
+        })?;
+        self.append_empty_value();
+        Ok(())
+    }
+
+    fn skip(&mut self) {
+        self.append_empty_value();
+    }
+
+    fn finish(&mut self) -> ArrayRef {
+        Arc::new(self.finish())
+    }
+}
+
 impl<'py> ArrayBuilder<'py> for StringBuilder {
     fn push(&mut self, val: Bound<'py, PyAny>) -> DataFusionResult<()> {
         let val: &str = val.extract().map_err(|_| {

guests/python/src/inspect.rs

@@ -1,5 +1,5 @@
 //! Inspection of Python code to extract [signature](crate::signature) information.
-use std::ffi::CString;
+use std::{collections::HashSet, ffi::CString};
 
 use datafusion_common::{DataFusionError, error::Result as DataFusionResult};
 use pyo3::{
@@ -35,6 +35,10 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PythonType {
         let type_time = mod_datetime.getattr(intern!(py, "time"))?;
         let type_timedelta = mod_datetime.getattr(intern!(py, "timedelta"))?;
 
+        // https://docs.python.org/3/library/types.html
+        let mod_types = py.import(intern!(py, "types"))?;
+        let type_none = mod_types.getattr(intern!(py, "NoneType"))?;
+
         if ob.is(type_bool) {
             Ok(Self::Bool)
         } else if ob.is(type_bytes) {
@@ -45,6 +49,8 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PythonType {
             Ok(Self::DateTime)
         } else if ob.is(type_float) {
             Ok(Self::Float)
+        } else if ob.is(&type_none) || ob.is_instance(&type_none).unwrap_or_default() {
+            Ok(Self::None)
         } else if ob.is(type_int) {
             Ok(Self::Int)
         } else if ob.is(type_str) {
@@ -85,38 +91,34 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PythonNullableType {
         // https://docs.python.org/3/library/types.html
         let mod_types = py.import(intern!(py, "types"))?;
         let type_union = mod_types.getattr(intern!(py, "UnionType"))?;
-        let type_none = mod_types.getattr(intern!(py, "NoneType"))?;
 
         if ob.is_instance(&type_union)? {
             let args = ob.getattr(intern!(py, "__args__"))?;
 
-            let n_args = args.len()?;
-            if n_args != 2 {
+            let mut args = args
+                .try_iter()?
+                .map(|arg| {
+                    let arg = arg?;
+                    arg.extract::<PythonType>()
+                })
+                .collect::<PyResult<HashSet<_>>>()?;
+
+            let nullable = args.len() > 1 && args.remove(&PythonType::None);
+
+            if args.len() != 1 {
                 return Err(PyErr::new::<PyTypeError, _>(format!(
-                    "only unions of length 2 are supported, got {n_args}"
+                    "only unions of form `T | None` are suppored, but got a union of {} distinct none-NULL types",
+                    args.len()
                 )));
             }
-            let (arg1, arg2): (Bound<'py, PyAny>, Bound<'py, PyAny>) = args.extract()?;
-
-            let inner_type = if arg1.is(&type_none) {
-                arg2
-            } else if arg2.is(&type_none) {
-                arg1
-            } else {
-                return Err(PyErr::new::<PyTypeError, _>(
-                    "only unions with None are supported",
-                ));
-            };
-
-            Ok(Self {
-                t: inner_type.extract()?,
-                nullable: true,
-            })
+            let t = args.into_iter().next().expect("just checked length");
+
+            Ok(Self { t, nullable })
         } else {
-            Ok(Self {
-                t: ob.extract()?,
-                nullable: false,
-            })
+            let t = ob.extract()?;
+            let nullable = t == PythonType::None;
+
+            Ok(Self { t, nullable })
         }
     }
 }

guests/python/src/signature.rs

@@ -8,7 +8,7 @@ use pyo3::{Py, PyAny};
 /// # Naming
 /// Since Python and Arrow use different names for the same type, we have to settle on some consistency. We chose to
 /// use the Python name in CamelCase style, so Python's `datetime` will become `DateTime`.
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub(crate) enum PythonType {
     /// Boolean.
     ///
@@ -82,6 +82,18 @@ pub(crate) enum PythonType {
     /// We map this to [`Int64`](arrow::datatypes::DataType::Int64).
     Int,
 
+    /// None/Null.
+    ///
+    /// # Python
+    /// The type is called `None`, documentation can be found here:
+    ///
+    /// - <https://docs.python.org/3/library/constants.html#None>
+    /// - <https://docs.python.org/3/library/types.html#types.NoneType>
+    ///
+    /// # Arrow
+    /// We map this to [`Null`](https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#variant.Null).
+    None,
+
     /// String.
     ///
     /// # Python

host/tests/integration_tests/python/inspection/errors.rs

@@ -80,7 +80,7 @@ def add_one(x: int | str) -> int:
         @r"
     scalar_udfs
     caused by
-    Error during planning: TypeError: only unions with None are supported
+    Error during planning: TypeError: only unions of form `T | None` are suppored, but got a union of 2 distinct none-NULL types
 
     The above exception was the direct cause of the following exception:
 
@@ -94,7 +94,7 @@ def add_one(x: int | str) -> int:
 }
 
 #[tokio::test(flavor = "multi_thread")]
-async fn test_union_type_3() {
+async fn test_union_type_2_and_none() {
     const CODE: &str = "
 def add_one(x: int | str | None) -> int:
     return x + 1
@@ -105,7 +105,82 @@ def add_one(x: int | str | None) -> int:
         @r"
     scalar_udfs
     caused by
-    Error during planning: TypeError: only unions of length 2 are supported, got 3
+    Error during planning: TypeError: only unions of form `T | None` are suppored, but got a union of 2 distinct none-NULL types
+
+    The above exception was the direct cause of the following exception:
+
+    TypeError: inspect parameter 1
+
+    The above exception was the direct cause of the following exception:
+
+    TypeError: inspect type of `add_one`
+    ",
+    );
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_union_type_2_identical() {
+    const CODE: &str = "
+def add_one(x: int | str | int) -> int:
+    return x + 1
+";
+
+    insta::assert_snapshot!(
+        err(CODE).await,
+        @r"
+    scalar_udfs
+    caused by
+    Error during planning: TypeError: only unions of form `T | None` are suppored, but got a union of 2 distinct none-NULL types
+
+    The above exception was the direct cause of the following exception:
+
+    TypeError: inspect parameter 1
+
+    The above exception was the direct cause of the following exception:
+
+    TypeError: inspect type of `add_one`
+    ",
+    );
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_union_type_2_identical_and_none() {
+    const CODE: &str = "
+def add_one(x: int | None | str | int) -> int:
+    return x + 1
+";
+
+    insta::assert_snapshot!(
+        err(CODE).await,
+        @r"
+    scalar_udfs
+    caused by
+    Error during planning: TypeError: only unions of form `T | None` are suppored, but got a union of 2 distinct none-NULL types
+
+    The above exception was the direct cause of the following exception:
+
+    TypeError: inspect parameter 1
+
+    The above exception was the direct cause of the following exception:
+
+    TypeError: inspect type of `add_one`
+    ",
+    );
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_union_type_3() {
+    const CODE: &str = "
+def add_one(x: int | str | float) -> int:
+    return x + 1
+";
+
+    insta::assert_snapshot!(
+        err(CODE).await,
+        @r"
+    scalar_udfs
+    caused by
+    Error during planning: TypeError: only unions of form `T | None` are suppored, but got a union of 3 distinct none-NULL types
 
     The above exception was the direct cause of the following exception:
 

host/tests/integration_tests/python/types/mod.rs

@@ -4,6 +4,8 @@ mod date;
 mod datetime;
 mod float;
 mod int;
+mod none;
 mod str;
 mod time;
 mod timedelta;
+mod union;