[AURON #1850] Add ArrowFieldWriters for temporal and composite types

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/arrow/FlinkArrowUtils.java

@@ -29,18 +29,26 @@
 import org.apache.arrow.vector.IntVector;
 import org.apache.arrow.vector.NullVector;
 import org.apache.arrow.vector.SmallIntVector;
+import org.apache.arrow.vector.TimeMicroVector;
+import org.apache.arrow.vector.TimeMilliVector;
+import org.apache.arrow.vector.TimeNanoVector;
+import org.apache.arrow.vector.TimeSecVector;
+import org.apache.arrow.vector.TimeStampVector;
 import org.apache.arrow.vector.TinyIntVector;
 import org.apache.arrow.vector.ValueVector;
 import org.apache.arrow.vector.VarBinaryVector;
 import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.complex.ListVector;
 import org.apache.arrow.vector.complex.MapVector;
+import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.types.DateUnit;
 import org.apache.arrow.vector.types.FloatingPointPrecision;
 import org.apache.arrow.vector.types.TimeUnit;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.auron.flink.arrow.writers.ArrayWriter;
 import org.apache.auron.flink.arrow.writers.ArrowFieldWriter;
 import org.apache.auron.flink.arrow.writers.BigIntWriter;
 import org.apache.auron.flink.arrow.writers.BooleanWriter;
@@ -49,8 +57,12 @@
 import org.apache.auron.flink.arrow.writers.DoubleWriter;
 import org.apache.auron.flink.arrow.writers.FloatWriter;
 import org.apache.auron.flink.arrow.writers.IntWriter;
+import org.apache.auron.flink.arrow.writers.MapWriter;
 import org.apache.auron.flink.arrow.writers.NullWriter;
+import org.apache.auron.flink.arrow.writers.RowWriter;
 import org.apache.auron.flink.arrow.writers.SmallIntWriter;
+import org.apache.auron.flink.arrow.writers.TimeWriter;
+import org.apache.auron.flink.arrow.writers.TimestampWriter;
 import org.apache.auron.flink.arrow.writers.TinyIntWriter;
 import org.apache.auron.flink.arrow.writers.VarBinaryWriter;
 import org.apache.auron.flink.arrow.writers.VarCharWriter;
@@ -283,6 +295,45 @@ static ArrowFieldWriter<RowData> createArrowFieldWriterForRow(ValueVector vector
             return DecimalWriter.forRow((DecimalVector) vector, decimalType.getPrecision(), decimalType.getScale());
         } else if (vector instanceof DateDayVector) {
             return DateWriter.forRow((DateDayVector) vector);
+        } else if (vector instanceof TimeSecVector
+                || vector instanceof TimeMilliVector
+                || vector instanceof TimeMicroVector
+                || vector instanceof TimeNanoVector) {
+            return TimeWriter.forRow(vector);
+        } else if (vector instanceof TimeStampVector) {
+            int precision;
+            if (fieldType instanceof LocalZonedTimestampType) {
+                precision = ((LocalZonedTimestampType) fieldType).getPrecision();
+            } else {
+                precision = ((TimestampType) fieldType).getPrecision();
+            }
+            return TimestampWriter.forRow(vector, precision);
+        } else if (vector instanceof MapVector) {
+            // MapVector extends ListVector, so this check must come before ListVector
+            MapVector mapVector = (MapVector) vector;
+            MapType mapType = (MapType) fieldType;
+            StructVector entriesVector = (StructVector) mapVector.getDataVector();
+            ArrowFieldWriter<ArrayData> keyWriter =
+                    createArrowFieldWriterForArray(entriesVector.getChild(MapVector.KEY_NAME), mapType.getKeyType());
+            ArrowFieldWriter<ArrayData> valueWriter = createArrowFieldWriterForArray(
+                    entriesVector.getChild(MapVector.VALUE_NAME), mapType.getValueType());
+            return MapWriter.forRow(mapVector, keyWriter, valueWriter);
+        } else if (vector instanceof ListVector) {
+            ListVector listVector = (ListVector) vector;
+            ArrayType arrayType = (ArrayType) fieldType;
+            ArrowFieldWriter<ArrayData> elementWriter =
+                    createArrowFieldWriterForArray(listVector.getDataVector(), arrayType.getElementType());
+            return ArrayWriter.forRow(listVector, elementWriter);
+        } else if (vector instanceof StructVector) {
+            StructVector structVector = (StructVector) vector;
+            RowType rowType = (RowType) fieldType;
+            @SuppressWarnings("unchecked")
+            ArrowFieldWriter<RowData>[] fieldsWriters = new ArrowFieldWriter[rowType.getFieldCount()];
+            for (int i = 0; i < fieldsWriters.length; i++) {
+                fieldsWriters[i] =
+                        createArrowFieldWriterForRow(structVector.getChildByOrdinal(i), rowType.getTypeAt(i));
+            }
+            return RowWriter.forRow(structVector, fieldsWriters);
         } else {
             throw new UnsupportedOperationException(
                     "Unsupported vector type: " + vector.getClass().getSimpleName());
@@ -323,6 +374,45 @@ static ArrowFieldWriter<ArrayData> createArrowFieldWriterForArray(ValueVector ve
             return DecimalWriter.forArray((DecimalVector) vector, decimalType.getPrecision(), decimalType.getScale());
         } else if (vector instanceof DateDayVector) {
             return DateWriter.forArray((DateDayVector) vector);
+        } else if (vector instanceof TimeSecVector
+                || vector instanceof TimeMilliVector
+                || vector instanceof TimeMicroVector
+                || vector instanceof TimeNanoVector) {
+            return TimeWriter.forArray(vector);
+        } else if (vector instanceof TimeStampVector) {
+            int precision;
+            if (fieldType instanceof LocalZonedTimestampType) {
+                precision = ((LocalZonedTimestampType) fieldType).getPrecision();
+            } else {
+                precision = ((TimestampType) fieldType).getPrecision();
+            }
+            return TimestampWriter.forArray(vector, precision);
+        } else if (vector instanceof MapVector) {
+            // MapVector extends ListVector, so this check must come before ListVector
+            MapVector mapVector = (MapVector) vector;
+            MapType mapType = (MapType) fieldType;
+            StructVector entriesVector = (StructVector) mapVector.getDataVector();
+            ArrowFieldWriter<ArrayData> keyWriter =
+                    createArrowFieldWriterForArray(entriesVector.getChild(MapVector.KEY_NAME), mapType.getKeyType());
+            ArrowFieldWriter<ArrayData> valueWriter = createArrowFieldWriterForArray(
+                    entriesVector.getChild(MapVector.VALUE_NAME), mapType.getValueType());
+            return MapWriter.forArray(mapVector, keyWriter, valueWriter);
+        } else if (vector instanceof ListVector) {
+            ListVector listVector = (ListVector) vector;
+            ArrayType arrayType = (ArrayType) fieldType;
+            ArrowFieldWriter<ArrayData> elementWriter =
+                    createArrowFieldWriterForArray(listVector.getDataVector(), arrayType.getElementType());
+            return ArrayWriter.forArray(listVector, elementWriter);
+        } else if (vector instanceof StructVector) {
+            StructVector structVector = (StructVector) vector;
+            RowType rowType = (RowType) fieldType;
+            @SuppressWarnings("unchecked")
+            ArrowFieldWriter<RowData>[] fieldsWriters = new ArrowFieldWriter[rowType.getFieldCount()];
+            for (int i = 0; i < fieldsWriters.length; i++) {
+                fieldsWriters[i] =
+                        createArrowFieldWriterForRow(structVector.getChildByOrdinal(i), rowType.getTypeAt(i));
+            }
+            return RowWriter.forArray(structVector, fieldsWriters);
         } else {
             throw new UnsupportedOperationException(
                     "Unsupported vector type: " + vector.getClass().getSimpleName());

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/arrow/writers/ArrayWriter.java

@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.auron.flink.arrow.writers;
+
+import java.util.Objects;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.flink.table.data.ArrayData;
+import org.apache.flink.table.data.RowData;
+
+/**
+ * {@link ArrowFieldWriter} for arrays ({@link ListVector}).
+ *
+ * <p>Holds an {@code elementWriter} that writes each array element. The element writer operates on
+ * {@link ArrayData} since array elements are accessed via {@link ArrayData} interface.
+ *
+ * @param <T> the input data type
+ */
+public abstract class ArrayWriter<T> extends ArrowFieldWriter<T> {
+
+    /** Creates an ArrayWriter that reads from {@link RowData}. */
+    public static ArrayWriter<RowData> forRow(ListVector listVector, ArrowFieldWriter<ArrayData> elementWriter) {
+        return new ArrayWriterForRow(listVector, elementWriter);
+    }
+
+    /** Creates an ArrayWriter that reads from {@link ArrayData}. */
+    public static ArrayWriter<ArrayData> forArray(ListVector listVector, ArrowFieldWriter<ArrayData> elementWriter) {
+        return new ArrayWriterForArray(listVector, elementWriter);
+    }
+
+    // ------------------------------------------------------------------------------------------
+
+    private final ArrowFieldWriter<ArrayData> elementWriter;
+
+    private ArrayWriter(ListVector listVector, ArrowFieldWriter<ArrayData> elementWriter) {
+        super(listVector);
+        this.elementWriter = Objects.requireNonNull(elementWriter);
+    }
+
+    abstract boolean isNullAt(T in, int ordinal);
+
+    abstract ArrayData readArray(T in, int ordinal);
+
+    @Override
+    public void doWrite(T in, int ordinal) {
+        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);
+            }
+            ((ListVector) getValueVector()).endValue(getCount(), array.size());
+        }
+    }
+
+    @Override
+    public void finish() {
+        super.finish();
+        elementWriter.finish();
+    }
+
+    @Override
+    public void reset() {
+        super.reset();
+        elementWriter.reset();
+    }
+
+    // ---- Inner classes for RowData and ArrayData ----
+
+    public static final class ArrayWriterForRow extends ArrayWriter<RowData> {
+        private ArrayWriterForRow(ListVector listVector, ArrowFieldWriter<ArrayData> elementWriter) {
+            super(listVector, elementWriter);
+        }
+
+        @Override
+        boolean isNullAt(RowData in, int ordinal) {
+            return in.isNullAt(ordinal);
+        }
+
+        @Override
+        ArrayData readArray(RowData in, int ordinal) {
+            return in.getArray(ordinal);
+        }
+    }
+
+    public static final class ArrayWriterForArray extends ArrayWriter<ArrayData> {
+        private ArrayWriterForArray(ListVector listVector, ArrowFieldWriter<ArrayData> elementWriter) {
+            super(listVector, elementWriter);
+        }
+
+        @Override
+        boolean isNullAt(ArrayData in, int ordinal) {
+            return in.isNullAt(ordinal);
+        }
+
+        @Override
+        ArrayData readArray(ArrayData in, int ordinal) {
+            return in.getArray(ordinal);
+        }
+    }
+}

auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/arrow/writers/MapWriter.java

@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.auron.flink.arrow.writers;
+
+import java.util.Objects;
+import org.apache.arrow.vector.complex.MapVector;
+import org.apache.arrow.vector.complex.StructVector;
+import org.apache.flink.table.data.ArrayData;
+import org.apache.flink.table.data.MapData;
+import org.apache.flink.table.data.RowData;
+
+/**
+ * {@link ArrowFieldWriter} for maps ({@link MapVector}).
+ *
+ * <p>Arrow represents maps as {@code List<Struct{key, value}>}. This writer holds separate key and
+ * value writers that operate on {@link ArrayData} (from {@link MapData#keyArray()} and {@link
+ * MapData#valueArray()}).
+ *
+ * @param <T> the input data type
+ */
+public abstract class MapWriter<T> extends ArrowFieldWriter<T> {
+
+    /** Creates a MapWriter that reads from {@link RowData}. */
+    public static MapWriter<RowData> forRow(
+            MapVector mapVector, ArrowFieldWriter<ArrayData> keyWriter, ArrowFieldWriter<ArrayData> valueWriter) {
+        return new MapWriterForRow(mapVector, keyWriter, valueWriter);
+    }
+
+    /** Creates a MapWriter that reads from {@link ArrayData}. */
+    public static MapWriter<ArrayData> forArray(
+            MapVector mapVector, ArrowFieldWriter<ArrayData> keyWriter, ArrowFieldWriter<ArrayData> valueWriter) {
+        return new MapWriterForArray(mapVector, keyWriter, valueWriter);
+    }
+
+    // ------------------------------------------------------------------------------------------
+
+    private final ArrowFieldWriter<ArrayData> keyWriter;
+    private final ArrowFieldWriter<ArrayData> valueWriter;
+
+    private MapWriter(
+            MapVector mapVector, ArrowFieldWriter<ArrayData> keyWriter, ArrowFieldWriter<ArrayData> valueWriter) {
+        super(mapVector);
+        this.keyWriter = Objects.requireNonNull(keyWriter);
+        this.valueWriter = Objects.requireNonNull(valueWriter);
+    }
+
+    abstract boolean isNullAt(T in, int ordinal);
+
+    abstract MapData readMap(T in, int ordinal);
+
+    @Override
+    public void doWrite(T in, int ordinal) {
+        if (!isNullAt(in, ordinal)) {
+            ((MapVector) getValueVector()).startNewValue(getCount());
+
+            StructVector structVector = (StructVector) ((MapVector) getValueVector()).getDataVector();
+            MapData map = readMap(in, ordinal);
+            ArrayData keys = map.keyArray();
+            ArrayData values = map.valueArray();
+            for (int i = 0; i < map.size(); i++) {
+                structVector.setIndexDefined(keyWriter.getCount());
+                keyWriter.write(keys, i);
+                valueWriter.write(values, i);
+            }
+
+            ((MapVector) getValueVector()).endValue(getCount(), map.size());
+        }
+    }
+
+    @Override
+    public void finish() {
+        super.finish();
+        keyWriter.finish();
+        valueWriter.finish();
+    }
+
+    @Override
+    public void reset() {
+        super.reset();
+        keyWriter.reset();
+        valueWriter.reset();
+    }
+
+    // ---- Inner classes for RowData and ArrayData ----
+
+    public static final class MapWriterForRow extends MapWriter<RowData> {
+        private MapWriterForRow(
+                MapVector mapVector, ArrowFieldWriter<ArrayData> keyWriter, ArrowFieldWriter<ArrayData> valueWriter) {
+            super(mapVector, keyWriter, valueWriter);
+        }
+
+        @Override
+        boolean isNullAt(RowData in, int ordinal) {
+            return in.isNullAt(ordinal);
+        }
+
+        @Override
+        MapData readMap(RowData in, int ordinal) {
+            return in.getMap(ordinal);
+        }
+    }
+
+    public static final class MapWriterForArray extends MapWriter<ArrayData> {
+        private MapWriterForArray(
+                MapVector mapVector, ArrowFieldWriter<ArrayData> keyWriter, ArrowFieldWriter<ArrayData> valueWriter) {
+            super(mapVector, keyWriter, valueWriter);
+        }
+
+        @Override
+        boolean isNullAt(ArrayData in, int ordinal) {
+            return in.isNullAt(ordinal);
+        }
+
+        @Override
+        MapData readMap(ArrayData in, int ordinal) {
+            return in.getMap(ordinal);
+        }
+    }
+}