Repository: mlflow/mlflow

Issue:
[BUG] input_example DataFrame type is always "split" no matter the format of the saved example
### 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

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

### MLflow version

- Client: 2.4.1

### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: windows
- **Python version**: 3.8.16


### Describe the problem

When input_example is passed to a logged model and parsed as a dataframe, the info.pandas_orient is always set to `split` ([link](https://github.com/mlflow/mlflow/blob/master/mlflow/models/utils.py#L211)), regardless of the actual orient of the data. Since it looks like `index` and `columns` are now being deleted in default cases ([link](https://github.com/mlflow/mlflow/blob/master/mlflow/models/utils.py#L179)), calling the resulting dataframe `split` leads to errors when deserializing the json dataframe. The example below is based directly on the `input_example` example snippet from the transformers.log_model documentation.

My suggestion would be to default to `values` and use split specifically when `index` and `columns` are provided, but of course worth discussing.

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell

```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
from mlflow.models.signature import infer_signature
from mlflow.transformers import generate_signature_output
from transformers import pipeline
import mlflow
import pandas as pd

en_to_de = pipeline("translation_en_to_de")

data = "MLflow is great!"
output = generate_signature_output(en_to_de, data)
signature = infer_signature(data, output)

with mlflow.start_run() as run:
    mlflow.transformers.log_model(
        transformers_model=en_to_de,
        artifact_path="english_to_german_translator",
        signature=signature,
        input_example=data,
    )

# pulled directly from MLmodel to save space
example_info = {
		"type": "dataframe",
		"pandas_orient": "split"
	}

model_uri = f"runs:/{run.info.run_id}/english_to_german_translator"
with open(f"{model_uri}/input_example.json", 'r') as sample_input_file:
    if example_info['type'] == 'dataframe':
        sample_input = pd.read_json(
            sample_input_file,
            orient=example_info['pandas_orient'],
            dtype=False
        )

# KeyError: 'columns'
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
Traceback (most recent call last):
  File "..\resources\mlflow_unit\get_formatted_example.py", line 28, in <module>
    sample_input = pd.read_json(
  File "C:\Users\wamartin\AppData\Local\Miniconda\4.7.12.1\envs\scoringscripttest\lib\site-packages\pandas\io\json\_json.py", line 784, in read_json
    return json_reader.read()
  File "C:\Users\wamartin\AppData\Local\Miniconda\4.7.12.1\envs\scoringscripttest\lib\site-packages\pandas\io\json\_json.py", line 975, in read
    obj = self._get_object_parser(self.data)
  File "C:\Users\wamartin\AppData\Local\Miniconda\4.7.12.1\envs\scoringscripttest\lib\site-packages\pandas\io\json\_json.py", line 1001, in _get_object_parser
    obj = FrameParser(json, **kwargs).parse()
  File "C:\Users\wamartin\AppData\Local\Miniconda\4.7.12.1\envs\scoringscripttest\lib\site-packages\pandas\io\json\_json.py", line 1134, in parse
    self._parse()
  File "C:\Users\wamartin\AppData\Local\Miniconda\4.7.12.1\envs\scoringscripttest\lib\site-packages\pandas\io\json\_json.py", line 1330, in _parse
    for col in decoded["columns"]
KeyError: 'columns'
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### 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
- [X] `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?

- [X] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
@wamartin-aml sounds like a good idea :) LMK when the PR is ready!
<!-- assign-maintainer -->
@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.

Resolving patch:
diff --git a/mlflow/models/utils.py b/mlflow/models/utils.py
--- a/mlflow/models/utils.py
+++ b/mlflow/models/utils.py
@@ -208,10 +208,11 @@ def _handle_dataframe_input(df):
                     f"but got '{type(input_example)}'",
                 )
             self.data = _handle_dataframe_input(self._inference_data)
+            orient = "split" if "columns" in self.data else "values"
             self.info = {
                 "artifact_path": example_filename,
                 "type": "dataframe",
-                "pandas_orient": "split",
+                "pandas_orient": orient,
             }
 
     def save(self, parent_dir_path: str):


Test patch:
diff --git a/tests/models/test_model_input_examples.py b/tests/models/test_model_input_examples.py
--- a/tests/models/test_model_input_examples.py
+++ b/tests/models/test_model_input_examples.py
@@ -50,6 +50,11 @@ def pandas_df_with_all_types():
     return df
 
 
+@pytest.fixture
+def df_without_columns():
+    return pd.DataFrame({0: [1, 2, 3], 1: [4, 5, 6], 2: [7, 8, 9]})
+
+
 @pytest.fixture
 def df_with_nan():
     return pd.DataFrame(
@@ -167,12 +172,66 @@ def test_input_examples(pandas_df_with_all_types, dict_of_ndarrays):
         assert example == parsed_df.to_dict(orient="records")[0]
 
 
+def test_pandas_orients_for_input_examples(
+    pandas_df_with_all_types, df_without_columns, dict_of_ndarrays
+):
+    # test setting example with data frame with all supported data types
+    with TempDir() as tmp:
+        example = _Example(pandas_df_with_all_types)
+        example.save(tmp.path())
+        filename = example.info["artifact_path"]
+        assert example.info["type"] == "dataframe"
+        assert example.info["pandas_orient"] == "split"
+        with open(tmp.path(filename)) as f:
+            data = json.load(f)
+            dataframe = pd.read_json(
+                json.dumps(data), orient=example.info["pandas_orient"], precise_float=True
+            )
+            assert (
+                (
+                    pandas_df_with_all_types.drop(columns=["binary"])
+                    == dataframe.drop(columns=["binary"])
+                )
+                .all()
+                .all()
+            )
+
+    with TempDir() as tmp:
+        example = _Example(df_without_columns)
+        example.save(tmp.path())
+        filename = example.info["artifact_path"]
+        assert example.info["type"] == "dataframe"
+        assert example.info["pandas_orient"] == "values"
+        with open(tmp.path(filename)) as f:
+            data = json.load(f)
+            assert set(data.keys()) == {"data"}
+            # NOTE: when no column names are provided (i.e. values orient),
+            # saving an example adds a "data" key rather than directly storing the plain data
+            data = data["data"]
+            dataframe = pd.read_json(json.dumps(data), orient=example.info["pandas_orient"])
+            assert (dataframe == df_without_columns).all().all()
+
+    # pass dict with scalars
+    with TempDir() as tmp:
+        example = {"a": 1, "b": "abc"}
+        x = _Example(example)
+        x.save(tmp.path())
+        filename = x.info["artifact_path"]
+        assert x.info["type"] == "dataframe"
+        assert x.info["pandas_orient"] == "split"
+        with open(tmp.path(filename)) as f:
+            parsed_json = json.load(f)
+            dataframe = pd.read_json(json.dumps(parsed_json), orient=x.info["pandas_orient"])
+            assert (dataframe == example).all().all()
+
+
 def test_sparse_matrix_input_examples(dict_of_sparse_matrix):
     for example_type, input_example in dict_of_sparse_matrix.items():
         with TempDir() as tmp:
             example = _Example(input_example)
             example.save(tmp.path())
             filename = example.info["artifact_path"]
+            assert example.info["type"] == example_type
             parsed_matrix = _read_sparse_matrix_from_json(tmp.path(filename), example_type)
             np.testing.assert_array_equal(parsed_matrix.toarray(), input_example.toarray())
 
@@ -184,9 +243,12 @@ def test_input_examples_with_nan(df_with_nan, dict_of_ndarrays_with_nans):
         example = _Example(df_with_nan)
         example.save(tmp.path())
         filename = example.info["artifact_path"]
+        assert example.info["type"] == "dataframe"
+        assert example.info["pandas_orient"] == "split"
         with open(tmp.path(filename)) as f:
             data = json.load(f)
             assert set(data.keys()) == {"columns", "data"}
+            pd.read_json(json.dumps(data), orient=example.info["pandas_orient"])
 
         parsed_df = dataframe_from_raw_json(tmp.path(filename), schema=sig.inputs)
 
@@ -210,6 +272,7 @@ def test_input_examples_with_nan(df_with_nan, dict_of_ndarrays_with_nans):
             example = _Example(input_example)
             example.save(tmp.path())
             filename = example.info["artifact_path"]
+            assert example.info["type"] == "ndarray"
             parsed_ary = _read_tensor_input_from_json(tmp.path(filename), schema=sig.inputs)
             assert np.array_equal(parsed_ary, input_example, equal_nan=True)
 
diff --git a/tests/transformers/test_transformers_model_export.py b/tests/transformers/test_transformers_model_export.py
--- a/tests/transformers/test_transformers_model_export.py
+++ b/tests/transformers/test_transformers_model_export.py
@@ -2000,6 +2000,9 @@ def test_infer_signature_from_example_only(
             assert set(saved_example[0].keys()).intersection(example.keys()) == set(
                 saved_example[0].keys()
             )
+        assert model.saved_input_example_info["type"] == "dataframe"
+        orient = "split" if pipeline_name == "zero_shot_pipeline" else "values"
+        assert model.saved_input_example_info["pandas_orient"] == orient
     else:
         assert model.saved_input_example_info is None