Repository: mlflow/mlflow

Issue:
[BUG] Error when returning a null value for array types in spark_udf
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

- Client: 2.2.2


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Ubuntu AWS
- **Python version**: 3.9
- **yarn version, if running the dev UI**:


### Describe the problem

When returning an array type in a model and using spark_udf, I get an error message.

### Tracking information

_No response_

### Code to reproduce issue

Here is a minimal reproducible example:

```python
import mlflow
import numpy as np
import pandas as pd
import pyspark.sql.functions as F


class MyMLFLowModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        model_input["B"] = list(np.random.rand(len(model_input), 100))
        # Insert null values
        model_input["B"] = [v if i % 2 == 1 else None for i, v in zip(model_input["A"], model_input["B"])]
        return model_input

with mlflow.start_run():
    model = MyMLFLowModel()
    model_info = mlflow.pyfunc.log_model("model", python_model=model)

model_udf = mlflow.pyfunc.spark_udf(spark, model_uri=model_info.model_uri, result_type="B array<double>")


test_df = spark.createDataFrame(pd.DataFrame(np.arange(10), columns=['A']))
pred_df = test_df.withColumn("prediction", model_udf(F.struct("*")))
pred_df.toPandas()
```

It works fine when creating my own UDF:

```python
from typing import Any, Union, Iterator, Tuple
import numpy as np
import pandas as pd
import pyspark.sql.functions as F

@F.pandas_udf("B array<double>")
def custom_udf(
    iterator: Iterator[Tuple[Union[pd.Series, pd.DataFrame], ...]]
) -> Iterator[pd.DataFrame]:
    for input_batch in iterator:
        if isinstance(input_batch, (pd.Series, pd.DataFrame)):
            # UDF is called with only one argument
            row_batch_args = (input_batch,)
        else:
            row_batch_args = input_batch

        if len(row_batch_args[0]) > 0:
            pred = model_loaded.predict(*row_batch_args)
            yield pred

test_df = spark.createDataFrame(pd.DataFrame(np.arange(10), columns=['A']))
pred_df = test_df.withColumn("prediction", custom_udf(F.struct("*"))).select("A", "prediction.*")
pred_df.toPandas()
```


### Stack trace

```python
Traceback (most recent call last):
  File "pyarrow/array.pxi", line 915, in pyarrow.lib.Array.from_pandas
  File "pyarrow/array.pxi", line 312, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 83, in pyarrow.lib._ndarray_to_array
  File "pyarrow/error.pxi", line 99, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: Can only convert 1-dimensional array values
```

### Other info / logs

The bug is probably cause by [this line](https://github.com/mlflow/mlflow/blob/master/mlflow/pyfunc/__init__.py#L1125) which basically create a nan singleton when the array is null. 

Also, there is a current PR that will touch this code but does not remove the bug: https://github.com/mlflow/mlflow/pull/8934/files#diff-592813f9fb198dda0982794f2221a408f316709190b3bf02c0e77eb2c1107075R907-R910

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
See https://github.com/mlflow/mlflow/pull/8934/commits/d8f76719dabf10f7ff5e8f248bd149849f4e97fa
You can help test it . 
@WeichenXu123 Do we need to fix this issue in #8934? If not, please separate it.
I will file a separate PR .
I create a temp branch for @freud14-tm testing. https://github.com/WeichenXu123/mlflow/tree/infer-spark-udf-return-type-test

Separate PR will be filed later, after https://github.com/mlflow/mlflow/pull/8934 is merged.
Hi @WeichenXu123 , your branch seems to work fine with my example. Could expand it with np.nan as well as None for the null value?
Something like this maybe?

```python
    def isnull(x):
        return x is None or (isinstance(x, float) and bool(np.isnan(x)))

    if array_dim == 1:
        return [
            np.array(v, dtype=np_type) if not isnull(v) else None
            for v in values
        ]
    else:
        return [
            list(np.array(v, dtype=np_type)) if not isnull(v) else None
            for v in values
        ]
```
@freud14-tm 

Curious, using nan to represent a Null array type value looks weird, why not normalize them to None in your model PyFunc code ?

Resolving patch:
diff --git a/mlflow/pyfunc/__init__.py b/mlflow/pyfunc/__init__.py
--- a/mlflow/pyfunc/__init__.py
+++ b/mlflow/pyfunc/__init__.py
@@ -923,6 +923,12 @@ def _parse_spark_datatype(datatype: str):
     return udf(lambda x: x, returnType=datatype).returnType
 
 
+def _is_none_or_nan(value):
+    # The condition `isinstance(value, float)` is needed to avoid error
+    # from `np.isnan(value)` if value is a non-numeric type.
+    return value is None or isinstance(value, float) and np.isnan(value)
+
+
 def _convert_array_values(values, elem_type, array_dim, spark_primitive_type_to_np_type):
     """
     Convert list or numpy array values to spark dataframe column values.
@@ -935,10 +941,12 @@ def _convert_array_values(values, elem_type, array_dim, spark_primitive_type_to_
             error_code=INVALID_PARAMETER_VALUE,
         )
 
+    # For array type result values, if provided value is None or NaN, regard it as a null array.
+    # see https://github.com/mlflow/mlflow/issues/8986
     if array_dim == 1:
-        return [np.array(v, dtype=np_type) for v in values]
+        return [None if _is_none_or_nan(v) else np.array(v, dtype=np_type) for v in values]
     else:
-        return [list(np.array(v, dtype=np_type)) for v in values]
+        return [None if _is_none_or_nan(v) else list(np.array(v, dtype=np_type)) for v in values]
 
 
 def _get_spark_primitive_types():


Test patch:
diff --git a/tests/pyfunc/test_spark.py b/tests/pyfunc/test_spark.py
--- a/tests/pyfunc/test_spark.py
+++ b/tests/pyfunc/test_spark.py
@@ -1014,3 +1014,24 @@ def predict(self, context, model_input, params=None):
         )
         res = good_data.withColumn("res", udf("a")).select("res").toPandas()
         assert res["res"][0] == [("str", 0, 1, 0.0, 0.1, True)]
+
+
+def test_spark_udf_return_nullable_array_field(spark):
+    class TestModel(PythonModel):
+        def predict(self, context, model_input):
+            values = [np.array([1.0, np.nan])] * (len(model_input) - 2) + [None, np.nan]
+            return pd.DataFrame({"a": values})
+
+    with mlflow.start_run():
+        mlflow_info = mlflow.pyfunc.log_model(
+            "model",
+            python_model=TestModel(),
+        )
+        udf = mlflow.pyfunc.spark_udf(
+            spark,
+            mlflow_info.model_uri,
+            result_type="a array<double>",
+        )
+        data1 = spark.range(3).repartition(1)
+        result = data1.select(udf("id").alias("res")).select("res.a").toPandas()
+        assert list(result["a"]) == [[1.0, None], None, None]