[AIR][Serve] Add support for multi-modal array input
diff --git a/doc/source/serve/http-adapters.md b/doc/source/serve/http-adapters.md
index 56fb75605c4dd..129dbb77bad8a 100644
--- a/doc/source/serve/http-adapters.md
+++ b/doc/source/serve/http-adapters.md
@@ -141,6 +141,6 @@ Here is a list of adapters and please feel free to [contribute more](https://git
 
 ```{eval-rst}
 .. automodule:: ray.serve.http_adapters
-    :members: json_to_ndarray, image_to_ndarray, starlette_request, json_request, pandas_read_json
+    :members: json_to_ndarray, image_to_ndarray, starlette_request, json_request, pandas_read_json, json_to_multi_ndarray
 
 ```
\ No newline at end of file
diff --git a/python/ray/serve/http_adapters.py b/python/ray/serve/http_adapters.py
index 948288985ec1b..88bddd81102af 100644
--- a/python/ray/serve/http_adapters.py
+++ b/python/ray/serve/http_adapters.py
@@ -52,6 +52,12 @@ def json_to_ndarray(payload: NdArray) -> np.ndarray:
     return arr
 
 
+@PublicAPI(stability="beta")
+def json_to_multi_ndarray(payload: Dict[str, NdArray]) -> Dict[str, np.ndarray]:
+    """Accepts a JSON of shape {str_key: NdArray} and converts it to dict of arrays."""
+    return {key: json_to_ndarray(arr_obj) for key, arr_obj in payload.items()}
+
+
 @PublicAPI(stability="beta")
 def starlette_request(
     request: starlette.requests.Request,
diff --git a/python/ray/serve/model_wrappers.py b/python/ray/serve/model_wrappers.py
index 479261dfab198..815b6fea11843 100644
--- a/python/ray/serve/model_wrappers.py
+++ b/python/ray/serve/model_wrappers.py
@@ -1,3 +1,4 @@
+from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Tuple, Type, Union
 import numpy as np
 
@@ -45,18 +46,73 @@ def collate_array(
     def unpack(output_arr):
         if isinstance(output_arr, list):
             return output_arr
-        assert isinstance(
-            output_arr, np.ndarray
-        ), f"The output should be np.ndarray but Serve got {type(output_arr)}."
-        assert len(output_arr) == batch_size, (
-            f"The output array should have shape of ({batch_size}, ...) "
-            f"but Serve got {output_arr.shape}"
-        )
+        if not isinstance(output_arr, np.ndarray):
+            raise TypeError(
+                f"The output should be np.ndarray but Serve got {type(output_arr)}."
+            )
+        if len(output_arr) != batch_size:
+            raise ValueError(
+                f"The output array should have shape of ({batch_size}, ...) "
+                f"because the input has {batch_size} entries "
+                f"but Serve got {output_arr.shape}"
+            )
         return [arr.squeeze(axis=0) for arr in np.split(output_arr, batch_size, axis=0)]
 
     return batched, unpack
 
 
+def collate_dict_array(
+    input_list: List[Dict[str, np.ndarray]]
+) -> Tuple[
+    Dict[str, np.ndarray],
+    Callable[[Dict[str, np.ndarray]], List[Dict[str, np.ndarray]]],
+]:
+    batch_size = len(input_list)
+
+    # Check all input has the same dict keys.
+    input_keys = [set(item.keys()) for item in input_list]
+    batch_has_same_keys = input_keys.count(input_keys[0]) == batch_size
+    if not batch_has_same_keys:
+        raise ValueError(
+            f"The input batch contains dictionary of different keys: {input_keys}"
+        )
+
+    # Turn list[dict[str, array]] to dict[str, List[array]]
+    key_to_list = defaultdict(list)
+    for single_dict in input_list:
+        for key, arr in single_dict.items():
+            key_to_list[key].append(arr)
+
+    # Turn dict[str, List[array]] to dict[str, array]
+    batched_dict = {}
+    unpack_dict = {}
+    for key, list_of_arr in key_to_list.items():
+        arr, unpack_func = collate_array(list_of_arr)
+        batched_dict[key] = arr
+        unpack_dict[key] = unpack_func
+
+    def unpack(output_dict: Dict[str, np.ndarray]):
+        # short circuit behavior, assume users already unpacked the output for us.
+        if isinstance(output_dict, list):
+            return output_dict
+
+        if not isinstance(output_dict, Dict):
+            raise TypeError(
+                f"The output should be a dictionary but Serve got {type(output_dict)}."
+            )
+
+        split_list_of_dict = [{} for _ in range(batch_size)]
+        for key, arr_unpack_func in unpack_dict.items():
+            arr_list = arr_unpack_func(output_dict[key])
+            # in place update each dictionary with the split array chunk.
+            for item, arr in zip(split_list_of_dict, arr_list):
+                item[key] = arr
+
+        return split_list_of_dict
+
+    return batched_dict, unpack
+
+
 @require_packages(["pandas"])
 def collate_dataframe(
     input_list: List["pd.DataFrame"],
@@ -69,13 +125,17 @@ def collate_dataframe(
     def unpack(output_df):
         if isinstance(output_df, list):
             return output_df
-        assert isinstance(
-            output_df, pd.DataFrame
-        ), f"The output should be a Pandas DataFrame but Serve got {type(output_df)}"
-        assert len(output_df) % batch_size == 0, (
-            f"The output dataframe should have length divisible by {batch_size}, "
-            f"but Serve got length {len(output_df)}."
-        )
+        if not isinstance(output_df, pd.DataFrame):
+            raise TypeError(
+                "The output should be a Pandas DataFrame but Serve got "
+                f"{type(output_df)}"
+            )
+        if len(output_df) % batch_size != 0:
+            raise ValueError(
+                f"The output dataframe should have length divisible by {batch_size}, "
+                f"because the input from {batch_size} different requests "
+                f"but Serve got length {len(output_df)}."
+            )
         return [df.reset_index(drop=True) for df in np.split(output_df, batch_size)]
 
     return batched, unpack
diff --git a/python/ray/serve/tests/test_http_adapters.py b/python/ray/serve/tests/test_http_adapters.py
index bd23804f4a135..c7d661ca979df 100644
--- a/python/ray/serve/tests/test_http_adapters.py
+++ b/python/ray/serve/tests/test_http_adapters.py
@@ -8,6 +8,7 @@
 
 from ray.serve.http_adapters import (
     NdArray,
+    json_to_multi_ndarray,
     json_to_ndarray,
     image_to_ndarray,
     pandas_read_json,
@@ -51,6 +52,12 @@ def test_json_to_ndarray():
     )
 
 
+def test_json_to_multi_ndarray():
+    assert json_to_multi_ndarray(
+        {"a": NdArray(array=[1]), "b": NdArray(array=[3])}
+    ) == {"a": np.array(1), "b": np.array(3)}
+
+
 def test_image_to_ndarray():
     buffer = io.BytesIO()
     arr = (np.random.rand(100, 100, 3) * 255).astype("uint8")
diff --git a/python/ray/serve/tests/test_model_wrappers.py b/python/ray/serve/tests/test_model_wrappers.py
index d99dbf54abf36..62d3f3f10f730 100644
--- a/python/ray/serve/tests/test_model_wrappers.py
+++ b/python/ray/serve/tests/test_model_wrappers.py
@@ -10,6 +10,7 @@
     ModelWrapperDeployment,
     collate_array,
     collate_dataframe,
+    collate_dict_array,
 )
 from ray.air.checkpoint import Checkpoint
 from ray.air.predictor import DataBatchType, Predictor
@@ -34,11 +35,28 @@ def test_array(self):
     def test_array_error(self):
         list_of_arr = [np.array([i]) for i in range(4)]
         _, unpack = collate_array(list_of_arr)
-        with pytest.raises(AssertionError, match="output array should have shape of"):
+        with pytest.raises(ValueError, match="output array should have shape of"):
             unpack(np.arange(2))
-        with pytest.raises(AssertionError, match="output should be np.ndarray but"):
+        with pytest.raises(TypeError, match="output should be np.ndarray but"):
             unpack("string")
 
+    def test_dict_array(self):
+        list_of_dicts = [
+            {"a": np.array([1, 2]), "b": np.array(3)},
+            {"a": np.array([3, 4]), "b": np.array(4)},
+        ]
+        batched_dict = {"a": np.array([[1, 2], [3, 4]]), "b": np.array([3, 4])}
+
+        batched, unpack = collate_dict_array(list_of_dicts)
+        assert batched.keys() == batched_dict.keys()
+        for key in batched.keys():
+            assert np.array_equal(batched[key], batched_dict[key])
+
+        for original, unpacked in zip(list_of_dicts, unpack(batched)):
+            assert original.keys() == unpacked.keys()
+            for key in original.keys():
+                assert np.array_equal(original[key], unpacked[key])
+
     def test_dataframe(self):
         list_of_dfs = [pd.DataFrame({"a": [i, i], "b": [i, i]}) for i in range(4)]
         batched_df = pd.DataFrame(

<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?
Ray AIR's data batch type accepts no just simple array or dataframe, but also a dictionary of array. This PR adds support for that in Serve's model wrapper, as well as a new http adapter accepting that. Most code re-uses existing array implementation.
<!-- Please give a short summary of the change and the problem this solves. -->

## Related issue number

<!-- For example: "Closes #1234" -->

## Checks

- [x] I've run `scripts/format.sh` to lint the changes in this PR.
- [x] I've included any doc changes needed for https://docs.ray.io/en/master/.
- [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(
