Repository: mlflow/mlflow

Issue:
[BUG] Batch inference with mlflow.pyfunc.spark_udf fails against timestamp input data types
### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

1.27.0

### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
Databricks ML Runtime 11.0



### Describe the problem

I'm trying to score a table using `mlflow.pyfunc.spark_udf`. The table has the schema shown below, with timestamp columns:
![image](https://user-images.githubusercontent.com/2358483/177837792-b77e3d53-a6ce-4f94-b80e-25149ee73475.png)

I have some wrapper logic around `mlflow.pyfunc.spark_udf` :

```python
import mlflow
from pyspark.sql.functions import struct


def predict_batch(spark_session, model_uri, input_table_name, output_table_name):
    """
    Apply the model at the specified URI for batch inference on the table with name input_table_name,
    writing results to the table with name output_table_name
    """
    table = spark_session.table(input_table_name)
    predict = mlflow.pyfunc.spark_udf(spark_session, model_uri, result_type="string", env_manager="conda")
    output_df = table.withColumn("prediction", predict(struct(*table.columns)))
    output_df.display()
    output_df.write.format("delta").mode("overwrite").saveAsTable(output_table_name)
```

When I try to call this function, e.g:
```python
predict_batch(spark, model_uri, input_table_name, output_table_name)
```

I get the following exception:
```
org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 0.0 failed 4 times, most recent failure: Lost task 0.3 in stage 0.0 (TID 3) (ip-10-68-149-252.us-west-2.compute.internal executor driver): org.apache.spark.api.python.PythonException: 'TypeError: Object of type Timestamp is not JSON serializable'. Full traceback below:
Traceback (most recent call last):
  File "/databricks/python/lib/python3.9/site-packages/mlflow/pyfunc/__init__.py", line 1274, in udf
    os.kill(scoring_server_proc.pid, signal.SIGTERM)
  File "/databricks/python/lib/python3.9/site-packages/mlflow/pyfunc/__init__.py", line 1116, in _predict_row_batch
    result = predict_fn(pdf)
  File "/databricks/python/lib/python3.9/site-packages/mlflow/pyfunc/__init__.py", line 1247, in batch_predict_fn
    return client.invoke(pdf)
  File "/databricks/python/lib/python3.9/site-packages/mlflow/pyfunc/scoring_server/client.py", line 39, in invoke
    post_data = json.dumps(scoring_server._get_jsonable_obj(data, pandas_orient="split"))
  File "/usr/lib/python3.9/json/__init__.py", line 231, in dumps
    return _default_encoder.encode(obj)
  File "/usr/lib/python3.9/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/usr/lib/python3.9/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/usr/lib/python3.9/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type Timestamp is not JSON serializable

	at org.apache.spark.api.python.BasePythonRunner$ReaderIterator.handlePythonException(PythonRunner.scala:696)
	at org.apache.spark.sql.execution.python.PythonArrowOutput$$anon$1.read(PythonArrowOutput.scala:101)
	at org.apache.spark.sql.execution.python.PythonArrowOutput$$anon$1.read(PythonArrowOutput.scala:50)
	at org.apache.spark.api.python.BasePythonRunner$ReaderIterator.hasNext(PythonRunner.scala:649)
	at org.apache.spark.InterruptibleIterator.hasNext(InterruptibleIterator.scala:37)
	at scala.collection.Iterator$$anon$11.hasNext(Iterator.scala:491)
	at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:460)
...
```

However, performing predictions locally using `mlflow.pyfunc.predict` doesn't hit the same JSON serialization issue:
```
import pandas as pd
pdf = spark.read.table("nyc_taxi_data_train").toPandas()
p.predict(pdf)
```

### Tracking information

_No response_

### Code to reproduce issue

See above

### Other info / logs

_No response_

### 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
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
- [ ] `area/projects`: MLproject format, project running backends
- [X] `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
[BUG] Batch inference with mlflow.pyfunc.spark_udf fails against timestamp input data types
### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

1.27.0

### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
Databricks ML Runtime 11.0



### Describe the problem

I'm trying to score a table using `mlflow.pyfunc.spark_udf`. The table has the schema shown below, with timestamp columns:
![image](https://user-images.githubusercontent.com/2358483/177837792-b77e3d53-a6ce-4f94-b80e-25149ee73475.png)

I have some wrapper logic around `mlflow.pyfunc.spark_udf` :

```python
import mlflow
from pyspark.sql.functions import struct


def predict_batch(spark_session, model_uri, input_table_name, output_table_name):
    """
    Apply the model at the specified URI for batch inference on the table with name input_table_name,
    writing results to the table with name output_table_name
    """
    table = spark_session.table(input_table_name)
    predict = mlflow.pyfunc.spark_udf(spark_session, model_uri, result_type="string", env_manager="conda")
    output_df = table.withColumn("prediction", predict(struct(*table.columns)))
    output_df.display()
    output_df.write.format("delta").mode("overwrite").saveAsTable(output_table_name)
```

When I try to call this function, e.g:
```python
predict_batch(spark, model_uri, input_table_name, output_table_name)
```

I get the following exception:
```
org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 0.0 failed 4 times, most recent failure: Lost task 0.3 in stage 0.0 (TID 3) (ip-10-68-149-252.us-west-2.compute.internal executor driver): org.apache.spark.api.python.PythonException: 'TypeError: Object of type Timestamp is not JSON serializable'. Full traceback below:
Traceback (most recent call last):
  File "/databricks/python/lib/python3.9/site-packages/mlflow/pyfunc/__init__.py", line 1274, in udf
    os.kill(scoring_server_proc.pid, signal.SIGTERM)
  File "/databricks/python/lib/python3.9/site-packages/mlflow/pyfunc/__init__.py", line 1116, in _predict_row_batch
    result = predict_fn(pdf)
  File "/databricks/python/lib/python3.9/site-packages/mlflow/pyfunc/__init__.py", line 1247, in batch_predict_fn
    return client.invoke(pdf)
  File "/databricks/python/lib/python3.9/site-packages/mlflow/pyfunc/scoring_server/client.py", line 39, in invoke
    post_data = json.dumps(scoring_server._get_jsonable_obj(data, pandas_orient="split"))
  File "/usr/lib/python3.9/json/__init__.py", line 231, in dumps
    return _default_encoder.encode(obj)
  File "/usr/lib/python3.9/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/usr/lib/python3.9/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/usr/lib/python3.9/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type Timestamp is not JSON serializable

	at org.apache.spark.api.python.BasePythonRunner$ReaderIterator.handlePythonException(PythonRunner.scala:696)
	at org.apache.spark.sql.execution.python.PythonArrowOutput$$anon$1.read(PythonArrowOutput.scala:101)
	at org.apache.spark.sql.execution.python.PythonArrowOutput$$anon$1.read(PythonArrowOutput.scala:50)
	at org.apache.spark.api.python.BasePythonRunner$ReaderIterator.hasNext(PythonRunner.scala:649)
	at org.apache.spark.InterruptibleIterator.hasNext(InterruptibleIterator.scala:37)
	at scala.collection.Iterator$$anon$11.hasNext(Iterator.scala:491)
	at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:460)
...
```

However, performing predictions locally using `mlflow.pyfunc.predict` doesn't hit the same JSON serialization issue:
```
import pandas as pd
pdf = spark.read.table("nyc_taxi_data_train").toPandas()
p.predict(pdf)
```

### Tracking information

_No response_

### Code to reproduce issue

See above

### Other info / logs

_No response_

### 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
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
- [ ] `area/projects`: MLproject format, project running backends
- [X] `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:
~I realized that the model I'm performing inference with (the one trained by https://github.com/mlflow/mlp-regression-template) does some feature engineering in its MLP `split` step to remove timestamp columns at training time, so I need to apply the same transformations at inference time. But still I think ^this bug may be worth fixing~

Never mind - the model trained by the MLP regression template includes the timestamp columns in its schema, so this issue currently blocks doing distributed batch inference using the MLP regression model. We could separately make the MLP regression template model batch-inference-friendly by dropping timestamp columns prior to model training
Hey @smurching, we could define an encoder for `pandas.Timestamp` to avoid the error:

```python
import pandas as pd

import json


class PandasTimeStampEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, pd.Timestamp):
            return str(obj)
        return json.JSONEncoder.default(self, obj)


try:
    print(json.dumps({"timestamp": pd.Timestamp("2017-01-01T12")}))
except Exception as e:
    print("Error", e)

print(json.dumps({"timestamp": pd.Timestamp("2017-01-01T12")}, cls=PandasTimeStampEncoder))
```

Output:

```
Error Object of type Timestamp is not JSON serializable
{"timestamp": "2017-01-01 12:00:00"}
```

This might cause an issue if the model contains a transformer that takes a `timestamp` column and converts it to some numeric columns.

> Hey @smurching, we could define an encoder for `pandas.Timestamp` to avoid the error:
> 
> ```python
> import pandas as pd
> 
> import json
> 
> 
> class PandasTimeStampEncoder(json.JSONEncoder):
>     def default(self, obj):
>         if isinstance(obj, pd.Timestamp):
>             return str(obj)
>         return json.JSONEncoder.default(self, obj)
> 
> 
> try:
>     print(json.dumps({"timestamp": pd.Timestamp("2017-01-01T12")}))
> except Exception as e:
>     print("Error", e)
> 
> print(json.dumps({"timestamp": pd.Timestamp("2017-01-01T12")}, cls=PandasTimeStampEncoder))
> ```
> 
> Output:
> 
> ```
> Error Object of type Timestamp is not JSON serializable
> {"timestamp": "2017-01-01 12:00:00"}
> ```
> 
> This might cause an issue if the model contains a transformer that takes a `timestamp` column and converts it to some numeric columns.

Shall we add the encoder into MLflow inference server code ?
@harupy @WeichenXu123 thanks for taking a look -

> This might cause an issue if the model contains a transformer that takes a timestamp column and converts it to some numeric columns.

To confirm my understanding, are you saying this might cause an issue because the timestamp (`"2017-01-01 12:00:00"`) would then be deserialized as a string rather than a timestamp, and so the model might not be able to perform inference on it? I wonder (but not sure) if this could be solved by adding logic in the inference server to deserialize timestamp types, as I think @WeichenXu123 is suggesting above. How do people usually get around this issue when working with `pandas`?
@smurching Yes, that's what I meant.

> How do people usually get around this issue when working with pandas?

I think they just extract some numerical features (e.g. `month`, `is_holiday`) from a datetime column, then drop it?
~I realized that the model I'm performing inference with (the one trained by https://github.com/mlflow/mlp-regression-template) does some feature engineering in its MLP `split` step to remove timestamp columns at training time, so I need to apply the same transformations at inference time. But still I think ^this bug may be worth fixing~

Never mind - the model trained by the MLP regression template includes the timestamp columns in its schema, so this issue currently blocks doing distributed batch inference using the MLP regression model. We could separately make the MLP regression template model batch-inference-friendly by dropping timestamp columns prior to model training
Hey @smurching, we could define an encoder for `pandas.Timestamp` to avoid the error:

```python
import pandas as pd

import json


class PandasTimeStampEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, pd.Timestamp):
            return str(obj)
        return json.JSONEncoder.default(self, obj)


try:
    print(json.dumps({"timestamp": pd.Timestamp("2017-01-01T12")}))
except Exception as e:
    print("Error", e)

print(json.dumps({"timestamp": pd.Timestamp("2017-01-01T12")}, cls=PandasTimeStampEncoder))
```

Output:

```
Error Object of type Timestamp is not JSON serializable
{"timestamp": "2017-01-01 12:00:00"}
```

This might cause an issue if the model contains a transformer that takes a `timestamp` column and converts it to some numeric columns.

> Hey @smurching, we could define an encoder for `pandas.Timestamp` to avoid the error:
> 
> ```python
> import pandas as pd
> 
> import json
> 
> 
> class PandasTimeStampEncoder(json.JSONEncoder):
>     def default(self, obj):
>         if isinstance(obj, pd.Timestamp):
>             return str(obj)
>         return json.JSONEncoder.default(self, obj)
> 
> 
> try:
>     print(json.dumps({"timestamp": pd.Timestamp("2017-01-01T12")}))
> except Exception as e:
>     print("Error", e)
> 
> print(json.dumps({"timestamp": pd.Timestamp("2017-01-01T12")}, cls=PandasTimeStampEncoder))
> ```
> 
> Output:
> 
> ```
> Error Object of type Timestamp is not JSON serializable
> {"timestamp": "2017-01-01 12:00:00"}
> ```
> 
> This might cause an issue if the model contains a transformer that takes a `timestamp` column and converts it to some numeric columns.

Shall we add the encoder into MLflow inference server code ?
@harupy @WeichenXu123 thanks for taking a look -

> This might cause an issue if the model contains a transformer that takes a timestamp column and converts it to some numeric columns.

To confirm my understanding, are you saying this might cause an issue because the timestamp (`"2017-01-01 12:00:00"`) would then be deserialized as a string rather than a timestamp, and so the model might not be able to perform inference on it? I wonder (but not sure) if this could be solved by adding logic in the inference server to deserialize timestamp types, as I think @WeichenXu123 is suggesting above. How do people usually get around this issue when working with `pandas`?
@smurching Yes, that's what I meant.

> How do people usually get around this issue when working with pandas?

I think they just extract some numerical features (e.g. `month`, `is_holiday`) from a datetime column, then drop it?

Resolving patch:
diff --git a/examples/spark_udf/spark_udf_datetime.py b/examples/spark_udf/spark_udf_datetime.py
new file mode 100644
--- /dev/null
+++ b/examples/spark_udf/spark_udf_datetime.py
@@ -0,0 +1,66 @@
+import random
+import datetime
+
+import mlflow
+from pyspark.sql import SparkSession
+from sklearn.datasets import load_iris
+from sklearn.neighbors import KNeighborsClassifier
+from sklearn.compose import ColumnTransformer
+from sklearn.preprocessing import FunctionTransformer
+from sklearn.pipeline import Pipeline
+
+
+def print_with_title(title, *args):
+    print(f"\n===== {title} =====\n")
+    for a in args:
+        print(a)
+
+
+def extract_month(df):
+    print_with_title("extract_month input", df.head(), df.dtypes)
+    transformed = df.assign(month=df["timestamp"].dt.month)
+    print_with_title("extract_month output", transformed.head(), transformed.dtypes)
+    return transformed
+
+
+def main():
+    X, y = load_iris(as_frame=True, return_X_y=True)
+    X = X.assign(
+        timestamp=[datetime.datetime(2022, random.randint(1, 12), 1) for _ in range(len(X))]
+    )
+    print_with_title("Ran input", X.head(30), X.dtypes)
+
+    signature = mlflow.models.infer_signature(X, y)
+    print_with_title("Signature", signature)
+
+    month_extractor = FunctionTransformer(extract_month, validate=False)
+    timestamp_remover = ColumnTransformer(
+        [("selector", "passthrough", X.columns.drop("timestamp"))], remainder="drop"
+    )
+    model = Pipeline(
+        [
+            ("month_extractor", month_extractor),
+            ("timestamp_remover", timestamp_remover),
+            ("knn", KNeighborsClassifier()),
+        ]
+    )
+    model.fit(X, y)
+
+    with mlflow.start_run():
+        model_info = mlflow.sklearn.log_model(model, "model", signature=signature)
+
+    with SparkSession.builder.getOrCreate() as spark:
+        infer_spark_df = spark.createDataFrame(X.sample(n=10, random_state=42))
+        print_with_title(
+            "Inference input",
+            infer_spark_df._jdf.showString(5, 20, False),  # numRows, truncate, vertical
+            infer_spark_df._jdf.schema().treeString(),
+        )
+
+        pyfunc_udf = mlflow.pyfunc.spark_udf(spark, model_info.model_uri, env_manager="conda")
+        result = infer_spark_df.select(pyfunc_udf(*X.columns).alias("predictions")).toPandas()
+        print_with_title("Inference result", result)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/mlflow/pyfunc/scoring_server/client.py b/mlflow/pyfunc/scoring_server/client.py
--- a/mlflow/pyfunc/scoring_server/client.py
+++ b/mlflow/pyfunc/scoring_server/client.py
@@ -1,8 +1,10 @@
 import requests
 import time
-from mlflow.pyfunc import scoring_server
 import json
 
+from mlflow.pyfunc import scoring_server
+from mlflow.utils.proto_json_utils import _DateTimeEncoder
+
 
 class ScoringServerClient:
     def __init__(self, host, port):
@@ -36,7 +38,10 @@ def invoke(self, data):
         Invoke inference on input data. The input data must be pandas dataframe or json instance.
         """
         content_type = scoring_server.CONTENT_TYPE_JSON
-        post_data = json.dumps(scoring_server._get_jsonable_obj(data, pandas_orient="split"))
+        post_data = json.dumps(
+            scoring_server._get_jsonable_obj(data, pandas_orient="split"),
+            cls=_DateTimeEncoder,
+        )
 
         response = requests.post(
             url=self.url_prefix + "/invocations",
diff --git a/mlflow/utils/proto_json_utils.py b/mlflow/utils/proto_json_utils.py
--- a/mlflow/utils/proto_json_utils.py
+++ b/mlflow/utils/proto_json_utils.py
@@ -220,7 +220,6 @@ def _dataframe_from_json(
                 dtypes = dict(zip(schema.input_names(), schema.numpy_types()))
         else:
             dtypes = dict(zip(schema.input_names(), schema.pandas_types()))
-
         df = pd.read_json(
             path_or_str,
             orient=pandas_orient,
@@ -356,3 +355,14 @@ def cast_schema_type(input_data):
             )
 
     return data
+
+
+# Reference: https://stackoverflow.com/a/12126976
+class _DateTimeEncoder(json.JSONEncoder):
+    def default(self, o):
+        import pandas as pd
+
+        if isinstance(o, (datetime.datetime, datetime.date, datetime.time, pd.Timestamp)):
+            return o.isoformat()
+
+        return super().default(o)


Test patch:
diff --git a/tests/examples/test_examples.py b/tests/examples/test_examples.py
--- a/tests/examples/test_examples.py
+++ b/tests/examples/test_examples.py
@@ -193,6 +193,7 @@ def test_mlflow_run_example(directory, params, tmpdir):
         ("evaluation", ["python", "evaluate_with_custom_metrics.py"]),
         ("evaluation", ["python", "evaluate_with_custom_metrics_comprehensive.py"]),
         ("diviner", ["python", "train.py"]),
+        ("spark_udf", ["python", "spark_udf_datetime.py"]),
     ],
 )
 def test_command_example(directory, command):
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
@@ -1,6 +1,8 @@
 import os
 import sys
 import time
+import random
+import datetime
 from typing import Iterator
 import threading
 from unittest import mock
@@ -27,6 +29,9 @@
 import sklearn.datasets as datasets
 from collections import namedtuple
 from sklearn.neighbors import KNeighborsClassifier
+from sklearn.compose import ColumnTransformer
+from sklearn.preprocessing import FunctionTransformer
+from sklearn.pipeline import Pipeline
 
 from pyspark.sql.functions import pandas_udf
 from pyspark.sql.functions import col, struct
@@ -485,3 +490,67 @@ def run_job():
     # assert ping failed, i.e. the server process is killed successfully.
     with pytest.raises(Exception):  # pylint: disable=pytest-raises-without-match
         client.ping()
+
+
+def test_spark_udf_datetime_with_model_schema(spark):
+    X, y = datasets.load_iris(as_frame=True, return_X_y=True)
+    X = X.assign(
+        timestamp=[datetime.datetime(2022, random.randint(1, 12), 1) for _ in range(len(X))]
+    )
+
+    month_extractor = FunctionTransformer(
+        lambda df: df.assign(month=df["timestamp"].map(lambda d: d.month)), validate=False
+    )
+    timestamp_remover = ColumnTransformer(
+        [("selector", "passthrough", X.columns.drop("timestamp"))], remainder="drop"
+    )
+    model = Pipeline(
+        [
+            ("month_extractor", month_extractor),
+            ("timestamp_remover", timestamp_remover),
+            ("knn", KNeighborsClassifier()),
+        ]
+    )
+    model.fit(X, y)
+
+    timestamp_dtype = {"timestamp": "datetime64[ns]"}
+    with mlflow.start_run():
+        signature = mlflow.models.infer_signature(X.astype(timestamp_dtype), y)
+        model_info = mlflow.sklearn.log_model(model, "model", signature=signature)
+
+    inference_sample = X.sample(n=10, random_state=42)
+    infer_spark_df = spark.createDataFrame(inference_sample.astype(timestamp_dtype))
+    pyfunc_udf = mlflow.pyfunc.spark_udf(spark, model_info.model_uri, env_manager="conda")
+    result = infer_spark_df.select(pyfunc_udf(*X.columns).alias("predictions")).toPandas()
+    np.testing.assert_almost_equal(result.to_numpy().squeeze(), model.predict(inference_sample))
+
+
+def test_spark_udf_string_datetime_with_model_schema(spark):
+    X, y = datasets.load_iris(as_frame=True, return_X_y=True)
+    X = X.assign(timestamp=["2022-{:02d}-01".format(random.randint(1, 12)) for _ in range(len(X))])
+
+    month_extractor = FunctionTransformer(
+        lambda df: df.assign(month=df["timestamp"].str.extract(r"^2022-0?(\d{1,2})-").astype(int)),
+        validate=False,
+    )
+    timestamp_remover = ColumnTransformer(
+        [("selector", "passthrough", X.columns.drop("timestamp"))], remainder="drop"
+    )
+    model = Pipeline(
+        [
+            ("month_extractor", month_extractor),
+            ("timestamp_remover", timestamp_remover),
+            ("knn", KNeighborsClassifier()),
+        ]
+    )
+    model.fit(X, y)
+
+    with mlflow.start_run():
+        signature = mlflow.models.infer_signature(X, y)
+        model_info = mlflow.sklearn.log_model(model, "model", signature=signature)
+
+    inference_sample = X.sample(n=10, random_state=42)
+    infer_spark_df = spark.createDataFrame(inference_sample)
+    pyfunc_udf = mlflow.pyfunc.spark_udf(spark, model_info.model_uri, env_manager="conda")
+    result = infer_spark_df.select(pyfunc_udf(*X.columns).alias("predictions")).toPandas()
+    np.testing.assert_almost_equal(result.to_numpy().squeeze(), model.predict(inference_sample))
diff --git a/tests/utils/test_proto_json_utils.py b/tests/utils/test_proto_json_utils.py
--- a/tests/utils/test_proto_json_utils.py
+++ b/tests/utils/test_proto_json_utils.py
@@ -1,5 +1,6 @@
 import base64
 import json
+import datetime
 import numpy as np
 import pandas as pd
 import pytest
@@ -20,6 +21,7 @@
     _stringify_all_experiment_ids,
     parse_tf_serving_input,
     _dataframe_from_json,
+    _DateTimeEncoder,
 )
 
 # Prevent pytest from trying to collect TestMessage as a test class:
@@ -527,3 +529,34 @@ def test_dataframe_from_json():
             source.to_json(orient="records"), pandas_orient="records", schema=tensor_schema
         )
     )
+
+    schema = Schema([ColSpec("datetime", "datetime")])
+    parsed = _dataframe_from_json(
+        '{"datetime": ["2022-01-01T00:00:00", "2022-01-02T03:04:05", "00:00:00"]}',
+        pandas_orient="records",
+        schema=schema,
+    )
+    expected = pd.DataFrame(
+        {
+            "datetime": [
+                pd.Timestamp("2022-01-01T00:00:00"),
+                pd.Timestamp("2022-01-02T03:04:05"),
+                pd.Timestamp("00:00:00"),
+            ]
+        },
+    )
+    pd.testing.assert_frame_equal(parsed, expected)
+
+
+@pytest.mark.parametrize(
+    "dt, expected",
+    [
+        (datetime.datetime(2022, 1, 1), '"2022-01-01T00:00:00"'),
+        (datetime.datetime(2022, 1, 2, 3, 4, 5), '"2022-01-02T03:04:05"'),
+        (datetime.date(2022, 1, 1), '"2022-01-01"'),
+        (datetime.time(0, 0, 0), '"00:00:00"'),
+        (pd.Timestamp(2022, 1, 1), '"2022-01-01T00:00:00"'),
+    ],
+)
+def test_datetime_encoder(dt, expected):
+    assert json.dumps(dt, cls=_DateTimeEncoder) == expected