[BP-2.3][FLINK-40346][model-triton] Preserve null array elements when deserializing Triton responses

flink-models/flink-model-triton/src/main/java/org/apache/flink/model/triton/TritonTypeMapper.java

@@ -32,11 +32,14 @@
 import org.apache.flink.table.types.logical.SmallIntType;
 import org.apache.flink.table.types.logical.TinyIntType;
 import org.apache.flink.table.types.logical.VarCharType;
+import org.apache.flink.table.types.logical.utils.LogicalTypeUtils;
 import org.apache.flink.util.Preconditions;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 
+import java.lang.reflect.Array;
+
 /** Utility class for mapping between Flink logical types and Triton data types. */
 public class TritonTypeMapper {
 
@@ -208,6 +211,15 @@ private static ArrayData deserializeArrayFromJson(JsonNode dataNode, LogicalType
 
         int size = dataNode.size();
 
+        // The primitive-backed arrays below cannot represent a JSON null: GenericArrayData
+        // reports isNullAt() == false for every position of a primitive array, so a null would
+        // silently be read back as 0, false, or the literal string "null". Use a boxed array
+        // whenever the payload actually contains a null, leaving the primitive fast path
+        // untouched for the common all-non-null case.
+        if (containsNull(dataNode)) {
+            return deserializeNullableArrayFromJson(dataNode, elementType, size);
+        }
+
         // Handle different element types with appropriate array types
         if (elementType instanceof BooleanType) {
             boolean[] array = new boolean[size];
@@ -270,6 +282,60 @@ private static ArrayData deserializeArrayFromJson(JsonNode dataNode, LogicalType
         }
     }
 
+    /** Returns whether the given JSON array contains at least one null element. */
+    private static boolean containsNull(JsonNode dataNode) {
+        for (JsonNode element : dataNode) {
+            if (element.isNull()) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    /**
+     * Deserializes a JSON array that contains at least one null element into a boxed {@link
+     * GenericArrayData}, which is the only representation able to report {@code isNullAt(pos)}.
+     *
+     * @param dataNode The JSON array node
+     * @param elementType The element type
+     * @param size The number of elements
+     * @return The deserialized ArrayData preserving null elements
+     */
+    private static ArrayData deserializeNullableArrayFromJson(
+            JsonNode dataNode, LogicalType elementType, int size) {
+        // Reject element types the module does not support before allocating anything, mirroring
+        // the trailing else branch of the primitive path above.
+        toTritonDataType(elementType);
+        // toTritonDataType recurses into an ArrayType rather than rejecting it, and the primitive
+        // path above does not support nested arrays; reject them explicitly so that the presence
+        // of a null element cannot change which element types are accepted.
+        Preconditions.checkArgument(
+                !(elementType instanceof ArrayType),
+                "Unsupported array element type: %s",
+                elementType);
+        // Writing a null into an array declared NOT NULL would violate the output schema, so fail
+        // loudly rather than substituting a value the model never produced.
+        Preconditions.checkArgument(
+                elementType.isNullable(),
+                "Received a null array element but the declared element type is NOT NULL: %s",
+                elementType);
+
+        // GenericArrayData requires a boxed array of the concrete component type: a plain Object[]
+        // would break ArrayObjectArrayConverter#toExternal, whose fast path returns the underlying
+        // array straight to the caller, where it is cast to e.g. Integer[].
+        Object[] array =
+                (Object[])
+                        Array.newInstance(
+                                LogicalTypeUtils.toInternalConversionClass(elementType), size);
+        int i = 0;
+        for (JsonNode element : dataNode) {
+            // deserializeFromJson maps a JSON null to a Java null and already covers every
+            // element type supported above, including the FloatType special case.
+            array[i++] = deserializeFromJson(element, elementType);
+        }
+        return new GenericArrayData(array);
+    }
+
     /**
      * Calculates the shape dimensions for the input data.
      *

flink-models/flink-model-triton/src/test/java/org/apache/flink/model/triton/TritonTypeMapperTest.java

@@ -17,26 +17,33 @@
 
 package org.apache.flink.model.triton;
 
+import org.apache.flink.table.api.DataTypes;
 import org.apache.flink.table.data.ArrayData;
 import org.apache.flink.table.data.GenericArrayData;
 import org.apache.flink.table.data.GenericRowData;
 import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
 import org.apache.flink.table.data.binary.BinaryStringData;
+import org.apache.flink.table.data.conversion.DataStructureConverter;
+import org.apache.flink.table.data.conversion.DataStructureConverters;
 import org.apache.flink.table.types.logical.ArrayType;
 import org.apache.flink.table.types.logical.BigIntType;
 import org.apache.flink.table.types.logical.BooleanType;
 import org.apache.flink.table.types.logical.DoubleType;
 import org.apache.flink.table.types.logical.FloatType;
 import org.apache.flink.table.types.logical.IntType;
+import org.apache.flink.table.types.logical.LogicalType;
 import org.apache.flink.table.types.logical.SmallIntType;
 import org.apache.flink.table.types.logical.TinyIntType;
 import org.apache.flink.table.types.logical.VarCharType;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import org.junit.jupiter.api.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.assertj.core.api.Assertions.within;
 
 /** Test for {@link TritonTypeMapper}. */
@@ -184,4 +191,172 @@ void testSerializeNull() {
         assertThat(jsonArray).hasSize(1);
         assertThat(jsonArray.get(0).isNull()).isTrue();
     }
+
+    @Test
+    void testDeserializeArrayWithNullStringElement() throws Exception {
+        ArrayData result =
+                (ArrayData)
+                        TritonTypeMapper.deserializeFromJson(
+                                objectMapper.readTree("[\"a\", null, \"b\"]"),
+                                new ArrayType(new VarCharType(VarCharType.MAX_LENGTH)));
+
+        assertThat(result.size()).isEqualTo(3);
+        assertThat(result.isNullAt(1)).isTrue();
+        assertThat(result.getString(0)).hasToString("a");
+        assertThat(result.getString(2)).hasToString("b");
+    }
+
+    @Test
+    void testDeserializeArrayWithNullNumericElements() throws Exception {
+        ArrayData ints =
+                (ArrayData)
+                        TritonTypeMapper.deserializeFromJson(
+                                objectMapper.readTree("[1, null, 3]"),
+                                new ArrayType(new IntType()));
+        assertThat(ints.size()).isEqualTo(3);
+        assertThat(ints.isNullAt(1)).isTrue();
+        assertThat(ints.getInt(0)).isEqualTo(1);
+        assertThat(ints.getInt(2)).isEqualTo(3);
+
+        ArrayData longs =
+                (ArrayData)
+                        TritonTypeMapper.deserializeFromJson(
+                                objectMapper.readTree("[1, null]"),
+                                new ArrayType(new BigIntType()));
+        assertThat(longs.isNullAt(1)).isTrue();
+
+        ArrayData doubles =
+                (ArrayData)
+                        TritonTypeMapper.deserializeFromJson(
+                                objectMapper.readTree("[1.5, null]"),
+                                new ArrayType(new DoubleType()));
+        assertThat(doubles.isNullAt(1)).isTrue();
+
+        ArrayData floats =
+                (ArrayData)
+                        TritonTypeMapper.deserializeFromJson(
+                                objectMapper.readTree("[1.5, null]"),
+                                new ArrayType(new FloatType()));
+        assertThat(floats.isNullAt(1)).isTrue();
+
+        ArrayData bytes =
+                (ArrayData)
+                        TritonTypeMapper.deserializeFromJson(
+                                objectMapper.readTree("[1, null]"),
+                                new ArrayType(new TinyIntType()));
+        assertThat(bytes.isNullAt(1)).isTrue();
+
+        ArrayData shorts =
+                (ArrayData)
+                        TritonTypeMapper.deserializeFromJson(
+                                objectMapper.readTree("[1, null]"),
+                                new ArrayType(new SmallIntType()));
+        assertThat(shorts.isNullAt(1)).isTrue();
+    }
+
+    @Test
+    void testDeserializeArrayWithNullBooleanElement() throws Exception {
+        ArrayData result =
+                (ArrayData)
+                        TritonTypeMapper.deserializeFromJson(
+                                objectMapper.readTree("[true, null, false]"),
+                                new ArrayType(new BooleanType()));
+
+        assertThat(result.size()).isEqualTo(3);
+        assertThat(result.isNullAt(1)).isTrue();
+        assertThat(result.getBoolean(0)).isTrue();
+        assertThat(result.getBoolean(2)).isFalse();
+    }
+
+    @Test
+    void testNullElementSurvivesSerializeDeserializeRoundTrip() throws Exception {
+        ArrayType arrayType = new ArrayType(new VarCharType(VarCharType.MAX_LENGTH));
+        RowData rowData =
+                GenericRowData.of(
+                        new GenericArrayData(
+                                new Object[] {
+                                    BinaryStringData.fromString("a"),
+                                    null,
+                                    BinaryStringData.fromString("b")
+                                }));
+
+        ArrayNode serialized = objectMapper.createArrayNode();
+        TritonTypeMapper.serializeToJsonArray(rowData, 0, arrayType, serialized);
+        assertThat(serialized.get(1).isNull()).isTrue();
+
+        ArrayData roundTripped =
+                (ArrayData) TritonTypeMapper.deserializeFromJson(serialized, arrayType);
+
+        assertThat(roundTripped.size()).isEqualTo(3);
+        assertThat(roundTripped.isNullAt(1)).isTrue();
+        assertThat(roundTripped.getString(0)).hasToString("a");
+        assertThat(roundTripped.getString(2)).hasToString("b");
+    }
+
+    @Test
+    void testDeserializeArrayRejectsNullForNotNullElementType() throws Exception {
+        JsonNode dataNode = objectMapper.readTree("[1, null]");
+        ArrayType notNullElements = new ArrayType(new IntType(false));
+
+        assertThatThrownBy(() -> TritonTypeMapper.deserializeFromJson(dataNode, notNullElements))
+                .isInstanceOf(IllegalArgumentException.class)
+                .hasMessageContaining("NOT NULL");
+    }
+
+    @Test
+    void testDeserializeArrayWithoutNullsIsUnchanged() throws Exception {
+        ArrayData result =
+                (ArrayData)
+                        TritonTypeMapper.deserializeFromJson(
+                                objectMapper.readTree("[1, 2, 3]"), new ArrayType(new IntType()));
+
+        assertThat(result.size()).isEqualTo(3);
+        assertThat(result.isNullAt(0)).isFalse();
+        assertThat(result.getInt(0)).isEqualTo(1);
+        assertThat(result.getInt(1)).isEqualTo(2);
+        assertThat(result.getInt(2)).isEqualTo(3);
+    }
+
+    @Test
+    void testNullableArrayKeepsConcreteComponentType() throws Exception {
+        // GenericArrayData requires a boxed array to carry its concrete component type; a plain
+        // Object[] backing array is not a valid representation of ARRAY<INT>.
+        assertThat(deserializeArray("[1, null, 3]", new IntType()).toObjectArray())
+                .isInstanceOf(Integer[].class);
+        assertThat(deserializeArray("[1, null]", new BigIntType()).toObjectArray())
+                .isInstanceOf(Long[].class);
+        assertThat(deserializeArray("[1, null]", new TinyIntType()).toObjectArray())
+                .isInstanceOf(Byte[].class);
+        assertThat(deserializeArray("[1, null]", new SmallIntType()).toObjectArray())
+                .isInstanceOf(Short[].class);
+        assertThat(deserializeArray("[1.5, null]", new FloatType()).toObjectArray())
+                .isInstanceOf(Float[].class);
+        assertThat(deserializeArray("[1.5, null]", new DoubleType()).toObjectArray())
+                .isInstanceOf(Double[].class);
+        assertThat(deserializeArray("[true, null]", new BooleanType()).toObjectArray())
+                .isInstanceOf(Boolean[].class);
+        assertThat(deserializeArray("[\"a\", null]", new VarCharType()).toObjectArray())
+                .isInstanceOf(StringData[].class);
+    }
+
+    @Test
+    void testNullableArraySurvivesExternalConversion() throws Exception {
+        // ArrayObjectArrayConverter#toExternal returns the backing array of a GenericArrayData
+        // directly, so an Object[] backing array would fail the cast to Integer[] below.
+        ArrayData ints = deserializeArray("[1, null, 3]", new IntType());
+
+        DataStructureConverter<Object, Object> converter =
+                DataStructureConverters.getConverter(DataTypes.ARRAY(DataTypes.INT()));
+        converter.open(TritonTypeMapperTest.class.getClassLoader());
+
+        Integer[] external = (Integer[]) converter.toExternal(ints);
+        assertThat(external).containsExactly(1, null, 3);
+    }
+
+    private GenericArrayData deserializeArray(String json, LogicalType elementType)
+            throws Exception {
+        return (GenericArrayData)
+                TritonTypeMapper.deserializeFromJson(
+                        objectMapper.readTree(json), new ArrayType(elementType));
+    }
 }