From 819665541bbea6033786a136c28cce79d0719cd2 Mon Sep 17 00:00:00 2001 From: Sepuri Sai Krishna Date: Fri, 14 Aug 2026 12:04:20 +0530 Subject: [PATCH] [FLINK-40346][model-triton] Preserve null array elements when deserializing Triton responses (#28947) (cherry picked from commit 969ea74468d0027d8a71d9ac2c727426af500b2b) Generated-by: Claude Code (Opus 5) --- .../flink/model/triton/TritonTypeMapper.java | 66 +++++++ .../model/triton/TritonTypeMapperTest.java | 175 ++++++++++++++++++ 2 files changed, 241 insertions(+) diff --git a/flink-models/flink-model-triton/src/main/java/org/apache/flink/model/triton/TritonTypeMapper.java b/flink-models/flink-model-triton/src/main/java/org/apache/flink/model/triton/TritonTypeMapper.java index 7b31b6c294e8d7..9813342bc36d43 100644 --- a/flink-models/flink-model-triton/src/main/java/org/apache/flink/model/triton/TritonTypeMapper.java +++ b/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. * diff --git a/flink-models/flink-model-triton/src/test/java/org/apache/flink/model/triton/TritonTypeMapperTest.java b/flink-models/flink-model-triton/src/test/java/org/apache/flink/model/triton/TritonTypeMapperTest.java index b811595d68d3b9..d5e04d2b7e3b3e 100644 --- a/flink-models/flink-model-triton/src/test/java/org/apache/flink/model/triton/TritonTypeMapperTest.java +++ b/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. + 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 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)); + } }