diff --git a/src/main/java/tech/streamfusion/arrow/ArrowWriter.java b/src/main/java/tech/streamfusion/arrow/ArrowWriter.java index e5e0477a..517cf193 100644 --- a/src/main/java/tech/streamfusion/arrow/ArrowWriter.java +++ b/src/main/java/tech/streamfusion/arrow/ArrowWriter.java @@ -65,6 +65,14 @@ public final class ArrowWriter { rowCount++; } + /** Caller must reserve top-level fixed-width capacity for every row in the batch. */ + public void writeReserved(IN row) { + for (int i = 0; i < fieldWriters.length; i++) { + fieldWriters[i].writeReserved(row, i); + } + rowCount++; + } + /** Finishes the writing of the current row batch. */ public void finish() { root.setRowCount(rowCount); diff --git a/src/main/java/tech/streamfusion/arrow/writers/ArrowFieldWriter.java b/src/main/java/tech/streamfusion/arrow/writers/ArrowFieldWriter.java index 6bdb9cb6..99f649f0 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/ArrowFieldWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/ArrowFieldWriter.java @@ -60,6 +60,16 @@ public abstract class ArrowFieldWriter { count += 1; } + /** Caller must have reserved capacity for this field's next value. */ + public void writeReserved(IN row, int ordinal) { + doWriteReserved(row, ordinal); + count += 1; + } + + protected void doWriteReserved(IN row, int ordinal) { + doWrite(row, ordinal); + } + protected void advanceCount(int amount) { count += amount; } diff --git a/src/main/java/tech/streamfusion/arrow/writers/BigIntWriter.java b/src/main/java/tech/streamfusion/arrow/writers/BigIntWriter.java index a509573c..ce806829 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/BigIntWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/BigIntWriter.java @@ -22,6 +22,7 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.table.data.ArrayData; import org.apache.flink.table.data.RowData; +import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.BigIntVector; /** {@link ArrowFieldWriter} for BigInt. */ @@ -60,6 +61,17 @@ public abstract class BigIntWriter extends ArrowFieldWriter { /** {@link BigIntWriter} for {@link RowData} input. */ public static final class BigIntWriterForRow extends BigIntWriter { + @Override + protected void doWriteReserved(RowData in, int ordinal) { + var vector = (BigIntVector) getValueVector(); + if (in.isNullAt(ordinal)) { + BitVectorHelper.unsetBit(vector.getValidityBuffer(), getCount()); + } else { + vector.set(getCount(), in.getLong(ordinal)); + } + } + + private BigIntWriterForRow(BigIntVector bigIntVector) { super(bigIntVector); } diff --git a/src/main/java/tech/streamfusion/arrow/writers/DoubleWriter.java b/src/main/java/tech/streamfusion/arrow/writers/DoubleWriter.java index dad79909..96633915 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/DoubleWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/DoubleWriter.java @@ -22,6 +22,7 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.table.data.ArrayData; import org.apache.flink.table.data.RowData; +import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.Float8Vector; /** {@link ArrowFieldWriter} for Double. */ @@ -60,6 +61,17 @@ public abstract class DoubleWriter extends ArrowFieldWriter { /** {@link DoubleWriter} for {@link RowData} input. */ public static final class DoubleWriterForRow extends DoubleWriter { + @Override + protected void doWriteReserved(RowData in, int ordinal) { + var vector = (Float8Vector) getValueVector(); + if (in.isNullAt(ordinal)) { + BitVectorHelper.unsetBit(vector.getValidityBuffer(), getCount()); + } else { + vector.set(getCount(), in.getDouble(ordinal)); + } + } + + private DoubleWriterForRow(Float8Vector doubleVector) { super(doubleVector); } diff --git a/src/main/java/tech/streamfusion/arrow/writers/FloatWriter.java b/src/main/java/tech/streamfusion/arrow/writers/FloatWriter.java index 33694a35..96ce0afd 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/FloatWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/FloatWriter.java @@ -22,6 +22,7 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.table.data.ArrayData; import org.apache.flink.table.data.RowData; +import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.Float4Vector; /** {@link ArrowFieldWriter} for Float. */ @@ -60,6 +61,17 @@ public abstract class FloatWriter extends ArrowFieldWriter { /** {@link FloatWriter} for {@link RowData} input. */ public static final class FloatWriterForRow extends FloatWriter { + @Override + protected void doWriteReserved(RowData in, int ordinal) { + var vector = (Float4Vector) getValueVector(); + if (in.isNullAt(ordinal)) { + BitVectorHelper.unsetBit(vector.getValidityBuffer(), getCount()); + } else { + vector.set(getCount(), in.getFloat(ordinal)); + } + } + + private FloatWriterForRow(Float4Vector floatVector) { super(floatVector); } diff --git a/src/main/java/tech/streamfusion/arrow/writers/IntWriter.java b/src/main/java/tech/streamfusion/arrow/writers/IntWriter.java index 3e02b199..8e004446 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/IntWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/IntWriter.java @@ -22,6 +22,7 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.table.data.ArrayData; import org.apache.flink.table.data.RowData; +import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.IntVector; /** {@link ArrowFieldWriter} for Int. */ @@ -60,6 +61,17 @@ public abstract class IntWriter extends ArrowFieldWriter { /** {@link IntWriter} for {@link RowData} input. */ public static final class IntWriterForRow extends IntWriter { + @Override + protected void doWriteReserved(RowData in, int ordinal) { + var vector = (IntVector) getValueVector(); + if (in.isNullAt(ordinal)) { + BitVectorHelper.unsetBit(vector.getValidityBuffer(), getCount()); + } else { + vector.set(getCount(), in.getInt(ordinal)); + } + } + + private IntWriterForRow(IntVector intVector) { super(intVector); } diff --git a/src/main/java/tech/streamfusion/arrow/writers/SmallIntWriter.java b/src/main/java/tech/streamfusion/arrow/writers/SmallIntWriter.java index 05ed5ebb..551bdb01 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/SmallIntWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/SmallIntWriter.java @@ -22,6 +22,7 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.table.data.ArrayData; import org.apache.flink.table.data.RowData; +import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.SmallIntVector; /** {@link ArrowFieldWriter} for SmallInt. */ @@ -60,6 +61,17 @@ public abstract class SmallIntWriter extends ArrowFieldWriter { /** {@link SmallIntWriter} for {@link RowData} input. */ public static final class SmallIntWriterForRow extends SmallIntWriter { + @Override + protected void doWriteReserved(RowData in, int ordinal) { + var vector = (SmallIntVector) getValueVector(); + if (in.isNullAt(ordinal)) { + BitVectorHelper.unsetBit(vector.getValidityBuffer(), getCount()); + } else { + vector.set(getCount(), in.getShort(ordinal)); + } + } + + private SmallIntWriterForRow(SmallIntVector intVector) { super(intVector); } diff --git a/src/main/java/tech/streamfusion/arrow/writers/TinyIntWriter.java b/src/main/java/tech/streamfusion/arrow/writers/TinyIntWriter.java index 61ddd6ea..1bdca166 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/TinyIntWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/TinyIntWriter.java @@ -22,6 +22,7 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.table.data.ArrayData; import org.apache.flink.table.data.RowData; +import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.TinyIntVector; /** {@link ArrowFieldWriter} for TinyInt. */ @@ -60,6 +61,17 @@ public abstract class TinyIntWriter extends ArrowFieldWriter { /** {@link TinyIntWriter} for {@link RowData} input. */ public static final class TinyIntWriterForRow extends TinyIntWriter { + @Override + protected void doWriteReserved(RowData in, int ordinal) { + var vector = (TinyIntVector) getValueVector(); + if (in.isNullAt(ordinal)) { + BitVectorHelper.unsetBit(vector.getValidityBuffer(), getCount()); + } else { + vector.set(getCount(), in.getByte(ordinal)); + } + } + + private TinyIntWriterForRow(TinyIntVector tinyIntVector) { super(tinyIntVector); } diff --git a/src/main/java/tech/streamfusion/operator/RowDataToArrowOperator.java b/src/main/java/tech/streamfusion/operator/RowDataToArrowOperator.java index 3f292647..64ed97ab 100644 --- a/src/main/java/tech/streamfusion/operator/RowDataToArrowOperator.java +++ b/src/main/java/tech/streamfusion/operator/RowDataToArrowOperator.java @@ -100,7 +100,8 @@ public class RowDataToArrowOperator extends FlinkStreamOperator kinds.setSafe(i, org.apache.flink.types.RowKind.INSERT.toByteValue()); } } - writer.write(row); + if (batchSize > 0) writer.writeReserved(row); + else writer.write(row); if (kinds != null) kinds.setSafe(pendingRows, row.getRowKind().toByteValue()); pendingRows++; } catch (RuntimeException | Error failure) { diff --git a/src/test/java/tech/streamfusion/operator/TransposeOperatorsTest.java b/src/test/java/tech/streamfusion/operator/TransposeOperatorsTest.java index 848c4198..135e8e06 100644 --- a/src/test/java/tech/streamfusion/operator/TransposeOperatorsTest.java +++ b/src/test/java/tech/streamfusion/operator/TransposeOperatorsTest.java @@ -197,6 +197,67 @@ class TransposeOperatorsTest { } } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void reservedPrimitiveEntryOwnsReusedRowsAndPreservesBitsAcrossBatches(boolean objectReuse) + throws Exception { + var schema = RowType.of(new org.apache.flink.table.types.logical.TinyIntType(), + new org.apache.flink.table.types.logical.SmallIntType(), new IntType(), new BigIntType(), + new org.apache.flink.table.types.logical.FloatType(), + new org.apache.flink.table.types.logical.DoubleType()); + List expected = List.of( + GenericRowData.of(Byte.MIN_VALUE, Short.MIN_VALUE, Integer.MIN_VALUE, Long.MIN_VALUE, + Float.intBitsToFloat(0x7fc00023), Double.longBitsToDouble(0x7ff8000000000023L)), + new GenericRowData(6), + GenericRowData.of((byte) 0, (short) 0, 0, 0L, -0.0f, -0.0d), + GenericRowData.of(Byte.MAX_VALUE, Short.MAX_VALUE, Integer.MAX_VALUE, Long.MAX_VALUE, + Float.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY), + GenericRowData.of((byte) -1, (short) -1, -1, -1L, + Float.intBitsToFloat(0xffc00023), Double.longBitsToDouble(0xfff8000000000023L)), + GenericRowData.of((byte) 1, (short) 1, 1, 1L, Float.MIN_VALUE, Double.MIN_VALUE)); + var kinds = org.apache.flink.types.RowKind.values(); + for (int i = 0; i < expected.size(); i++) expected.get(i).setRowKind(kinds[i % kinds.length]); + List retained = new ArrayList<>(); + try (var harness = new OneInputStreamOperatorTestHarness( + new RowDataToArrowOperator(schema, 2, false, null))) { + if (objectReuse) harness.getExecutionConfig().enableObjectReuse(); + harness.setup(new ArrowBatchSerializer()); + harness.open(); + var reused = new GenericRowData(6); + for (var row : expected) { + for (int field = 0; field < 6; field++) reused.setField(field, row.getField(field)); + reused.setRowKind(row.getRowKind()); + harness.processElement(new StreamRecord<>(reused)); + for (int field = 0; field < 6; field++) reused.setField(field, null); + } + List batches = values(harness); + assertEquals(3, batches.size()); + for (var batch : batches) { + try (var root = batch.root()) { + retained.addAll(RowDataArrowConverter.read(root, schema)); + } + } + } + assertEquals(expected.size(), retained.size()); + for (int i = 0; i < expected.size(); i++) { + var input = expected.get(i); + var output = retained.get(i); + assertEquals(input.getRowKind(), output.getRowKind()); + for (int field = 0; field < 6; field++) { + assertEquals(input.isNullAt(field), output.isNullAt(field)); + } + if (input.isNullAt(0)) continue; + assertEquals(input.getByte(0), output.getByte(0)); + assertEquals(input.getShort(1), output.getShort(1)); + assertEquals(input.getInt(2), output.getInt(2)); + assertEquals(input.getLong(3), output.getLong(3)); + assertEquals(Float.floatToRawIntBits(input.getFloat(4)), + Float.floatToRawIntBits(output.getFloat(4))); + assertEquals(Double.doubleToRawLongBits(input.getDouble(5)), + Double.doubleToRawLongBits(output.getDouble(5))); + } + } + @Test void rowToArrowOwnsBufferedRowsFromReusingUpstream() throws Exception { try (OneInputStreamOperatorTestHarness harness = @@ -224,6 +285,25 @@ class TransposeOperatorsTest { } } + @Test + void growingWriterStillExpandsBeyondInitialFixedWidthCapacity() { + var schema = RowType.of(new IntType()); + try (var allocator = new org.apache.arrow.memory.RootAllocator(); + var root = org.apache.arrow.vector.VectorSchemaRoot.create( + tech.streamfusion.arrow.ArrowConversion.toArrowSchema(schema), allocator)) { + var writer = tech.streamfusion.arrow.ArrowConversion.createRowDataArrowWriter(root, schema, 1); + int rows = ((org.apache.arrow.vector.IntVector) root.getVector(0)).getValueCapacity() + 2; + for (int i = 0; i < rows; i++) writer.write(GenericRowData.of(i % 7 == 0 ? null : i)); + writer.finish(); + var output = RowDataArrowConverter.read(root, schema); + assertEquals(rows, output.size()); + for (int i = 0; i < rows; i++) { + assertEquals(i % 7 == 0, output.get(i).isNullAt(0)); + if (i % 7 != 0) assertEquals(i, output.get(i).getInt(0)); + } + } + } + @ParameterizedTest @ValueSource(booleans = {false, true}) void exitRowsKeepNestedValuesAfterArrowBatchCloses(boolean objectReuse) throws Exception {