DX-61115 cherry picked commits form arrow

README_DREMIO.md

@@ -1,8 +1,5 @@
-git checkout -b dremio_24.3_12.0 apache-arrow-12.0.1  
-
-Apply starting from the bottom and working upwards.
-
-git cherry-pick d4dc6d525cfd89aa7ca7a23e0d755f5f2e3b2c48
+git cherry-pick cf53c1e023605fd506a3617bbd9fe1788b33c5a3
+git cherry-pick 0230230a4a28d0e1e2aa23570804a14444c58e4a
 git cherry-pick 0002859096a93cd9b0630db757a277b75a483263
 git cherry-pick 5ed2f61a9fa72d011248bf9cfec7a9734d5b2aa0
 git cherry-pick c601e318a897b3a6fd5f4fd3068beda1460173a2

java/pom.xml

@@ -33,7 +33,7 @@
     <dep.junit.jupiter.version>5.9.0</dep.junit.jupiter.version>
     <dep.slf4j.version>1.7.25</dep.slf4j.version>
     <dep.guava-bom.version>31.1-jre</dep.guava-bom.version>
-    <dep.netty-bom.version>4.1.93.Final</dep.netty-bom.version>
+    <dep.netty-bom.version>4.1.96.Final</dep.netty-bom.version>
     <dep.grpc-bom.version>1.56.0</dep.grpc-bom.version>
     <dep.protobuf-bom.version>3.21.9</dep.protobuf-bom.version>
     <dep.jackson-bom.version>2.13.4</dep.jackson-bom.version>

java/vector/src/main/codegen/templates/BaseWriter.java

@@ -106,6 +106,7 @@ public interface ComplexWriter {
     void copyReader(FieldReader reader);
     StructWriter rootAsStruct();
     ListWriter rootAsList();
+    MapWriter rootAsMap(boolean keysSorted);
 
     void setPosition(int index);
     void setValueCount(int count);

java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java

@@ -19,6 +19,7 @@
 
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.MapVector;
 import org.apache.arrow.vector.complex.NonNullableStructVector;
 import org.apache.arrow.vector.complex.StateTool;
 import org.apache.arrow.vector.complex.StructVector;
@@ -32,14 +33,15 @@ public class ComplexWriterImpl extends AbstractFieldWriter implements ComplexWri
 
   private NullableStructWriter structRoot;
   private UnionListWriter listRoot;
+  private UnionMapWriter mapRoot;
   private final NonNullableStructVector container;
 
   Mode mode = Mode.INIT;
   private final String name;
   private final boolean unionEnabled;
   private final NullableStructWriterFactory nullableStructWriterFactory;
 
-  private enum Mode { INIT, STRUCT, LIST }
+  private enum Mode { INIT, STRUCT, LIST, MAP }
 
   /**
    * Constructs a new instance.
@@ -107,6 +109,9 @@ public void clear() {
       case LIST:
         listRoot.clear();
         break;
+      case MAP:
+        mapRoot.clear();
+        break;
       default:
         break;
     }
@@ -121,6 +126,9 @@ public void setValueCount(int count) {
       case LIST:
         listRoot.setValueCount(count);
         break;
+      case MAP:
+        mapRoot.setValueCount(count);
+        break;
       default:
         break;
     }
@@ -136,6 +144,9 @@ public void setPosition(int index) {
       case LIST:
         listRoot.setPosition(index);
         break;
+      case MAP:
+        mapRoot.setPosition(index);
+        break;
       default:
         break;
     }
@@ -223,5 +234,29 @@ public ListWriter rootAsList() {
     return listRoot;
   }
 
+  @Override
+  public MapWriter rootAsMap(boolean keysSorted) {
+    switch (mode) {
 
+      case INIT:
+        int vectorCount = container.size();
+        // TODO allow dictionaries in complex types
+        MapVector mapVector = container.addOrGetMap(name, keysSorted);
+        if (container.size() > vectorCount) {
+          mapVector.allocateNew();
+        }
+        mapRoot = new UnionMapWriter(mapVector);
+        mapRoot.setPosition(idx());
+        mode = Mode.MAP;
+        break;
+
+      case MAP:
+        break;
+
+      default:
+        check(Mode.INIT, Mode.STRUCT);
+    }
+
+    return mapRoot;
+  }
 }

java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java

@@ -1401,4 +1401,108 @@ public void testStructOfList() {
       Assert.assertEquals(0, size);
     }
   }
+
+  @Test
+  public void testMap() {
+    try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
+      ComplexWriter writer = new ComplexWriterImpl("root", parent);
+      MapWriter mapWriter = writer.rootAsMap(false);
+      for (int i = 0; i < COUNT; i++) {
+        mapWriter.startMap();
+        for (int j = 0; j < i % 7; j++) {
+          mapWriter.startEntry();
+          if (j % 2 == 0) {
+            mapWriter.key().integer().writeInt(j);
+            mapWriter.value().integer().writeInt(j + 1);
+          } else {
+            IntHolder keyHolder = new IntHolder();
+            keyHolder.value = j;
+            IntHolder valueHolder = new IntHolder();
+            valueHolder.value = j + 1;
+            mapWriter.key().integer().write(keyHolder);
+            mapWriter.value().integer().write(valueHolder);
+          }
+          mapWriter.endEntry();
+        }
+        mapWriter.endMap();
+      }
+      writer.setValueCount(COUNT);
+      UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
+      for (int i = 0; i < COUNT; i++) {
+        mapReader.setPosition(i);
+        for (int j = 0; j < i % 7; j++) {
+          mapReader.next();
+          assertEquals(j, mapReader.key().readInteger().intValue());
+          assertEquals(j + 1, mapReader.value().readInteger().intValue());
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testMapWithNulls() {
+    try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
+      ComplexWriter writer = new ComplexWriterImpl("root", parent);
+      MapWriter mapWriter = writer.rootAsMap(false);
+      mapWriter.startMap();
+      mapWriter.startEntry();
+      mapWriter.key().integer().writeNull();
+      mapWriter.value().integer().writeInt(1);
+      mapWriter.endEntry();
+      mapWriter.endMap();
+      writer.setValueCount(1);
+      UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
+      Assert.assertNull(mapReader.key().readInteger());
+      assertEquals(1, mapReader.value().readInteger().intValue());
+    }
+  }
+
+  @Test
+  public void testMapWithListKey() {
+    try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
+      ComplexWriter writer = new ComplexWriterImpl("root", parent);
+      MapWriter mapWriter = writer.rootAsMap(false);
+      mapWriter.startMap();
+      mapWriter.startEntry();
+      mapWriter.key().list().startList();
+      for (int i = 0; i < 3; i++) {
+        mapWriter.key().list().integer().writeInt(i);
+      }
+      mapWriter.key().list().endList();
+      mapWriter.value().integer().writeInt(1);
+      mapWriter.endEntry();
+      mapWriter.endMap();
+      writer.setValueCount(1);
+      UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
+      mapReader.key().next();
+      assertEquals(0, mapReader.key().reader().readInteger().intValue());
+      mapReader.key().next();
+      assertEquals(1, mapReader.key().reader().readInteger().intValue());
+      mapReader.key().next();
+      assertEquals(2, mapReader.key().reader().readInteger().intValue());
+      assertEquals(1, mapReader.value().readInteger().intValue());
+    }
+  }
+
+  @Test
+  public void testMapWithStructKey() {
+    try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
+      ComplexWriter writer = new ComplexWriterImpl("root", parent);
+      MapWriter mapWriter = writer.rootAsMap(false);
+      mapWriter.startMap();
+      mapWriter.startEntry();
+      mapWriter.key().struct().start();
+      mapWriter.key().struct().integer("value1").writeInt(1);
+      mapWriter.key().struct().integer("value2").writeInt(2);
+      mapWriter.key().struct().end();
+      mapWriter.value().integer().writeInt(1);
+      mapWriter.endEntry();
+      mapWriter.endMap();
+      writer.setValueCount(1);
+      UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
+      assertEquals(1, mapReader.key().reader("value1").readInteger().intValue());
+      assertEquals(2, mapReader.key().reader("value2").readInteger().intValue());
+      assertEquals(1, mapReader.value().readInteger().intValue());
+    }
+  }
 }