[core] Expose columnar data in RowColumnVector/ArrayColumnVector/MapC…

…olumnVector (#3207)
apache · Apr 16, 2024 · c05cbca · c05cbca
1 parent 0bd955c
commit c05cbca
Show file tree

Hide file tree

Showing 10 changed files with 84 additions and 0 deletions.
diff --git a/paimon-common/src/main/java/org/apache/paimon/data/columnar/ArrayColumnVector.java b/paimon-common/src/main/java/org/apache/paimon/data/columnar/ArrayColumnVector.java
@@ -23,4 +23,6 @@
 /** Array column vector. */
 public interface ArrayColumnVector extends ColumnVector {
     InternalArray getArray(int i);
+
+    ColumnVector getColumnVector();
 }
diff --git a/paimon-common/src/main/java/org/apache/paimon/data/columnar/MapColumnVector.java b/paimon-common/src/main/java/org/apache/paimon/data/columnar/MapColumnVector.java
@@ -23,4 +23,8 @@
 /** Map column vector. */
 public interface MapColumnVector extends ColumnVector {
     InternalMap getMap(int i);
+
+    ColumnVector getKeyColumnVector();
+
+    ColumnVector getValueColumnVector();
 }
diff --git a/paimon-common/src/main/java/org/apache/paimon/data/columnar/RowColumnVector.java b/paimon-common/src/main/java/org/apache/paimon/data/columnar/RowColumnVector.java
@@ -23,4 +23,6 @@
 /** Row column vector. */
 public interface RowColumnVector extends ColumnVector {
     InternalRow getRow(int i);
+
+    VectorizedColumnBatch getBatch();
 }
diff --git a/paimon-common/src/main/java/org/apache/paimon/data/columnar/heap/HeapArrayVector.java b/paimon-common/src/main/java/org/apache/paimon/data/columnar/heap/HeapArrayVector.java
@@ -84,4 +84,9 @@ public InternalArray getArray(int i) {
         long length = lengths[i];
         return new ColumnarArray(child, (int) offset, (int) length);
     }
+
+    @Override
+    public ColumnVector getColumnVector() {
+        return child;
+    }
 }
diff --git a/paimon-common/src/main/java/org/apache/paimon/data/columnar/heap/HeapMapVector.java b/paimon-common/src/main/java/org/apache/paimon/data/columnar/heap/HeapMapVector.java
@@ -72,4 +72,14 @@ public InternalMap getMap(int i) {
         long length = lengths[i];
         return new ColumnarMap(keys, values, (int) offset, (int) length);
     }
+
+    @Override
+    public ColumnVector getKeyColumnVector() {
+        return keys;
+    }
+
+    @Override
+    public ColumnVector getValueColumnVector() {
+        return values;
+    }
 }
diff --git a/paimon-common/src/main/java/org/apache/paimon/data/columnar/heap/HeapRowVector.java b/paimon-common/src/main/java/org/apache/paimon/data/columnar/heap/HeapRowVector.java
@@ -45,6 +45,11 @@ public ColumnarRow getRow(int i) {
         return columnarRow;
     }
 
+    @Override
+    public VectorizedColumnBatch getBatch() {
+        return new VectorizedColumnBatch(fields);
+    }
+
     @Override
     public void reset() {
         super.reset();

diff --git a/paimon-common/src/main/java/org/apache/paimon/utils/VectorMappingUtils.java b/paimon-common/src/main/java/org/apache/paimon/utils/VectorMappingUtils.java
@@ -39,6 +39,7 @@
 import org.apache.paimon.data.columnar.RowColumnVector;
 import org.apache.paimon.data.columnar.ShortColumnVector;
 import org.apache.paimon.data.columnar.TimestampColumnVector;
+import org.apache.paimon.data.columnar.VectorizedColumnBatch;
 import org.apache.paimon.types.ArrayType;
 import org.apache.paimon.types.BigIntType;
 import org.apache.paimon.types.BinaryType;
@@ -333,6 +334,12 @@ public InternalArray getArray(int i) {
                 public boolean isNullAt(int i) {
                     return partition.isNullAt(index);
                 }
+
+                @Override
+                public ColumnVector getColumnVector() {
+                    throw new UnsupportedOperationException(
+                            "Doesn't support getting ColumnVector.");
+                }
             };
         }
 
@@ -348,6 +355,18 @@ public InternalMap getMap(int i) {
                 public boolean isNullAt(int i) {
                     return partition.isNullAt(index);
                 }
+
+                @Override
+                public ColumnVector getKeyColumnVector() {
+                    throw new UnsupportedOperationException(
+                            "Doesn't support getting key ColumnVector.");
+                }
+
+                @Override
+                public ColumnVector getValueColumnVector() {
+                    throw new UnsupportedOperationException(
+                            "Doesn't support getting value ColumnVector.");
+                }
             };
         }
 
@@ -363,6 +382,18 @@ public InternalMap getMap(int i) {
                 public boolean isNullAt(int i) {
                     return partition.isNullAt(index);
                 }
+
+                @Override
+                public ColumnVector getKeyColumnVector() {
+                    throw new UnsupportedOperationException(
+                            "Doesn't support getting key ColumnVector.");
+                }
+
+                @Override
+                public ColumnVector getValueColumnVector() {
+                    throw new UnsupportedOperationException(
+                            "Doesn't support getting value ColumnVector.");
+                }
             };
         }
 
@@ -378,6 +409,11 @@ public InternalRow getRow(int i) {
                 public boolean isNullAt(int i) {
                     return partition.isNullAt(index);
                 }
+
+                @Override
+                public VectorizedColumnBatch getBatch() {
+                    throw new UnsupportedOperationException("Doesn't support getting batch.");
+                }
             };
         }
 

diff --git a/paimon-format/src/main/java/org/apache/paimon/format/orc/reader/OrcArrayColumnVector.java b/paimon-format/src/main/java/org/apache/paimon/format/orc/reader/OrcArrayColumnVector.java
@@ -44,4 +44,9 @@ public InternalArray getArray(int i) {
         long length = hiveVector.lengths[i];
         return new ColumnarArray(paimonVector, (int) offset, (int) length);
     }
+
+    @Override
+    public ColumnVector getColumnVector() {
+        return paimonVector;
+    }
 }
diff --git a/paimon-format/src/main/java/org/apache/paimon/format/orc/reader/OrcMapColumnVector.java b/paimon-format/src/main/java/org/apache/paimon/format/orc/reader/OrcMapColumnVector.java
@@ -46,4 +46,14 @@ public InternalMap getMap(int i) {
         long length = hiveVector.lengths[i];
         return new ColumnarMap(keyPaimonVector, valuePaimonVector, (int) offset, (int) length);
     }
+
+    @Override
+    public ColumnVector getKeyColumnVector() {
+        return keyPaimonVector;
+    }
+
+    @Override
+    public ColumnVector getValueColumnVector() {
+        return valuePaimonVector;
+    }
 }
diff --git a/paimon-format/src/main/java/org/apache/paimon/format/orc/reader/OrcRowColumnVector.java b/paimon-format/src/main/java/org/apache/paimon/format/orc/reader/OrcRowColumnVector.java
@@ -45,4 +45,9 @@ public OrcRowColumnVector(StructColumnVector hiveVector, RowType type) {
     public ColumnarRow getRow(int i) {
         return new ColumnarRow(batch, i);
     }
+
+    @Override
+    public VectorizedColumnBatch getBatch() {
+        return batch;
+    }
 }