fix: add date types to ArrowHelper (#141)

cloudquery/cloudquery#15431

Author: spieden

Diff for: lib/src/main/java/io/cloudquery/helper/ArrowHelper.java

@@ -22,6 +22,7 @@
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.BigIntVector;
 import org.apache.arrow.vector.BitVector;
+import org.apache.arrow.vector.DateDayVector;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.FixedSizeBinaryVector;
 import org.apache.arrow.vector.Float4Vector;
@@ -139,6 +140,10 @@ private static void setVectorData(FieldVector vector, Object data) {
       jsonVector.setSafe(0, (byte[]) data);
       return;
     }
+    if (vector instanceof DateDayVector dayDateVector) {
+      dayDateVector.set(0, (int) data);
+      return;
+    }
 
     throw new IllegalArgumentException("Unsupported vector type: " + vector.getClass());
   }

Diff for: lib/src/test/java/io/cloudquery/helper/ArrowHelperTest.java

@@ -15,8 +15,10 @@
 import io.cloudquery.schema.Resource;
 import io.cloudquery.schema.Table;
 import java.io.IOException;
+import java.time.LocalDate;
 import java.util.List;
 import java.util.Map;
+import org.apache.arrow.vector.types.DateUnit;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -41,7 +43,11 @@ public class ArrowHelperTest {
                       .primaryKey(true)
                       .build(),
                   Column.builder().name("string_column2").type(ArrowType.Utf8.INSTANCE).build(),
-                  Column.builder().name("boolean_column").type(ArrowType.Bool.INSTANCE).build()))
+                  Column.builder().name("boolean_column").type(ArrowType.Bool.INSTANCE).build(),
+                  Column.builder()
+                      .name("date_days_column")
+                      .type(new ArrowType.Date(DateUnit.DAY))
+                      .build()))
           .build();
 
   @Test
@@ -69,6 +75,9 @@ public void testToArrowSchema() {
             CQ_EXTENSION_PRIMARY_KEY,
             "false"));
 
+    assertEquals(arrowSchema.getFields().get(2).getName(), "boolean_column");
+    assertEquals(arrowSchema.getFields().get(3).getName(), "date_days_column");
+
     assertEquals(
         arrowSchema.getCustomMetadata(),
         Map.of(
@@ -84,7 +93,8 @@ public void testFromArrowSchema() {
     List<Field> fields =
         List.of(
             Field.nullable("string_column1", ArrowType.Utf8.INSTANCE),
-            Field.nullable("string_column2", ArrowType.Utf8.INSTANCE));
+            Field.nullable("string_column2", ArrowType.Utf8.INSTANCE),
+            Field.nullable("date_days_column", new ArrowType.Date(DateUnit.DAY)));
 
     Schema schema = new Schema(fields, Map.of(CQ_TABLE_NAME, "table1"));
 
@@ -120,6 +130,7 @@ public void testRoundTripResourceEncoding() throws Exception {
     Resource resource = Resource.builder().table(TEST_TABLE).build();
     resource.set("string_column1", "test_data");
     resource.set("string_column2", "test_data2");
+    resource.set("date_days_column", (int) LocalDate.parse("2023-11-24").toEpochDay());
     resource.set("boolean_column", true);
 
     Assertions.assertDoesNotThrow(