Fix casts and unit tests

Signed-off-by: Andy Qin <[email protected]>

Author: q-andy

common/src/main/java/org/opensearch/ml/common/connector/functions/postprocess/RemoteMlCommonsPassthroughPostProcessFunction.java

@@ -15,7 +15,7 @@
 import org.opensearch.ml.common.output.model.ModelTensor;
 
 /**
- * A post-processing function for calling a remote ml commons instancethat preserves the original neural sparse response structure
+ * A post-processing function for calling a remote ml commons instance that preserves the original neural sparse response structure
  * to avoid double-wrapping when receiving responses from another ML-Commons instance.
  */
 public class RemoteMlCommonsPassthroughPostProcessFunction extends ConnectorPostProcessFunction<Map<String, Object>> {
@@ -26,7 +26,6 @@ public void validate(Object input) {
         }
     }
 
-    @SuppressWarnings("unchecked")
     @Override
     public List<ModelTensor> process(Map<String, Object> input, MLResultDataType dataType) {
         // Check if this is an ML-Commons response with inference_results
@@ -41,7 +40,9 @@ public List<ModelTensor> process(Map<String, Object> input, MLResultDataType dat
                     for (Map<String, Object> output : outputs) {
                         // This inner map should represent a model tensor, so we try to parse and instantiate a new one.
                         ModelTensor modelTensor = createModelTensorFromMap(output);
-                        modelTensors.add(modelTensor);
+                        if (modelTensor != null) {
+                            modelTensors.add(modelTensor);
+                        }
                     }
                 }
             }
@@ -59,8 +60,11 @@ public List<ModelTensor> process(Map<String, Object> input, MLResultDataType dat
      * Creates a ModelTensor from a Map<String, Object> representation based on the API format
      * of the /_predict API
      */
-    @SuppressWarnings("unchecked")
     private ModelTensor createModelTensorFromMap(Map<String, Object> map) {
+        if (map == null || map.isEmpty()) {
+            return null;
+        }
+
         String name = (String) map.getOrDefault(ModelTensor.NAME_FIELD, OUTPUT_FIELD);
         String result = (String) map.get(ModelTensor.RESULT_FIELD);
         Map<String, Object> dataAsMap = (Map<String, Object>) map.get(ModelTensor.DATA_AS_MAP_FIELD);
@@ -82,8 +86,7 @@ private ModelTensor createModelTensorFromMap(Map<String, Object> map) {
         long[] shape = null;
         if (map.containsKey(ModelTensor.SHAPE_FIELD)) {
             Object shapeObj = map.get(ModelTensor.SHAPE_FIELD);
-            if (shapeObj instanceof List) {
-                List<?> shapeList = (List<?>) shapeObj;
+            if (shapeObj instanceof List<?> shapeList) {
                 shape = new long[shapeList.size()];
                 for (int i = 0; i < shapeList.size(); i++) {
                     Object item = shapeList.get(i);
@@ -98,8 +101,7 @@ private ModelTensor createModelTensorFromMap(Map<String, Object> map) {
         Number[] data = null;
         if (map.containsKey(ModelTensor.DATA_FIELD)) {
             Object dataObj = map.get(ModelTensor.DATA_FIELD);
-            if (dataObj instanceof List) {
-                List<?> dataList = (List<?>) dataObj;
+            if (dataObj instanceof List<?> dataList) {
                 data = new Number[dataList.size()];
                 for (int i = 0; i < dataList.size(); i++) {
                     Object item = dataList.get(i);

common/src/test/java/org/opensearch/ml/common/connector/functions/postprocess/RemoteMlCommonsPassthroughPostProcessFunctionTest.java

@@ -56,7 +56,7 @@ public void process_MLCommonsResponse_RankFeatures() {
         assertEquals(innerDataAsMap, tensor.getDataAsMap());
 
         // Verify the nested sparse data structure
-        Map<String, Object> dataAsMap =  (Map<String, Object>) tensor.getDataAsMap();
+        Map<String, Object> dataAsMap = (Map<String, Object>) tensor.getDataAsMap();
         List<Map<String, Object>> response = (List<Map<String, Object>>) dataAsMap.get("response");
         assertEquals(1, response.size());
         assertEquals(0.35982144, (Double) response.get(0).get("hello"), 0.0001);
@@ -88,7 +88,7 @@ public void process_MLCommonsResponse_DenseVector() {
         assertEquals(1, result.size());
         ModelTensor tensor = result.get(0);
         assertEquals("sentence_embedding", tensor.getName());
-        assertEquals(3, tensor.getShape().length);
+        assertEquals(1, tensor.getShape().length);
         assertEquals(3L, tensor.getShape()[0]);
         assertEquals(3, tensor.getData().length);
         assertEquals(0.5400895, tensor.getData()[0].doubleValue(), 0.0001);