diff --git a/src/main/java/tech/streamfusion/arrow/writers/ArrayWriter.java b/src/main/java/tech/streamfusion/arrow/writers/ArrayWriter.java index f10e188a..ddbb3001 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/ArrayWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/ArrayWriter.java @@ -57,8 +57,12 @@ public abstract class ArrayWriter extends ArrowFieldWriter { if (!isNullAt(in, ordinal)) { ((ListVector) getValueVector()).startNewValue(getCount()); ArrayData array = readArray(in, ordinal); - for (int i = 0; i < array.size(); i++) { - elementWriter.write(array, i); + if (elementWriter instanceof IntWriter.IntWriterForArray integers) { + integers.writeArray(array); + } else { + for (int i = 0; i < array.size(); i++) { + elementWriter.write(array, i); + } } ((ListVector) getValueVector()).endValue(getCount(), array.size()); } diff --git a/src/main/java/tech/streamfusion/arrow/writers/ArrowFieldWriter.java b/src/main/java/tech/streamfusion/arrow/writers/ArrowFieldWriter.java index 4782e85b..6bdb9cb6 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/ArrowFieldWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/ArrowFieldWriter.java @@ -60,6 +60,10 @@ public abstract class ArrowFieldWriter { count += 1; } + protected void advanceCount(int amount) { + count += amount; + } + /** Finishes the writing of the current row batch. */ public void finish() { valueVector.setValueCount(count); diff --git a/src/main/java/tech/streamfusion/arrow/writers/IntWriter.java b/src/main/java/tech/streamfusion/arrow/writers/IntWriter.java index 7ddb2091..3e02b199 100644 --- a/src/main/java/tech/streamfusion/arrow/writers/IntWriter.java +++ b/src/main/java/tech/streamfusion/arrow/writers/IntWriter.java @@ -82,6 +82,18 @@ public abstract class IntWriter extends ArrowFieldWriter { super(intVector); } + public void writeArray(ArrayData array) { + var vector = (IntVector) getValueVector(); + int start = getCount(); + int end = Math.addExact(start, array.size()); + while (vector.getValueCapacity() < end) vector.reAlloc(); + for (int i = 0; i < array.size(); i++) { + if (array.isNullAt(i)) vector.setNull(start + i); + else vector.set(start + i, array.getInt(i)); + } + advanceCount(array.size()); + } + @Override boolean isNullAt(ArrayData in, int ordinal) { return in.isNullAt(ordinal);