From c3601a97a0718ae47726e6c134cbed4b98bd1a36 Mon Sep 17 00:00:00 2001
From: Maksim Yegorov <997437+myegorov@users.noreply.github.com>
Date: Tue, 5 Nov 2024 19:07:06 -0500
Subject: [PATCH] GH-44344: [Java] fix VectorSchemaRoot.getTransferPair for
 NullVector (#44631)

### Rationale for this change

Do not throw [UnsupportedOperationException("Tried to get allocator from NullVector")](https://github.com/apache/arrow/blob/release-18.0.0-rc0/java/vector/src/main/java/org/apache/arrow/vector/NullVector.java#L160) from [VectorSchemaRoot.slice()](https://github.com/apache/arrow/blob/release-18.0.0-rc0/java/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java#L341) when slicing a VSR containing a NullVector or ZeroVector. Details in https://github.com/apache/arrow/issues/44344

### Are these changes tested?

Added unit test that would trigger an UnsupportedOperationException on the legacy path.
* GitHub Issue: #44344

Authored-by: Maksim Yegorov <59841139+maksimyego-db@users.noreply.github.com>
Signed-off-by: David Li <li.davidm96@gmail.com>
---
 .../org/apache/arrow/vector/NullVector.java   |  9 ++++++-
 .../org/apache/arrow/vector/ValueVector.java  |  6 +++++
 .../arrow/vector/TestSplitAndTransfer.java    | 25 +++++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/NullVector.java b/java/vector/src/main/java/org/apache/arrow/vector/NullVector.java
index 227ca716f6391..6bfe540d232fc 100644
--- a/java/vector/src/main/java/org/apache/arrow/vector/NullVector.java
+++ b/java/vector/src/main/java/org/apache/arrow/vector/NullVector.java
@@ -155,9 +155,16 @@ public boolean allocateNewSafe() {
   @Override
   public void reAlloc() {}
 
+  /*
+   * IMPORTANT NOTE
+   * It's essential that NullVector (and ZeroVector) do not require BufferAllocator for any data storage.
+   * However, some methods of the parent interface may require passing in a BufferAllocator, even if null.
+   *
+   * @return null
+   */
   @Override
   public BufferAllocator getAllocator() {
-    throw new UnsupportedOperationException("Tried to get allocator from NullVector");
+    return null;
   }
 
   @Override
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ValueVector.java b/java/vector/src/main/java/org/apache/arrow/vector/ValueVector.java
index 724941aa2a1e8..0a45409eb9860 100644
--- a/java/vector/src/main/java/org/apache/arrow/vector/ValueVector.java
+++ b/java/vector/src/main/java/org/apache/arrow/vector/ValueVector.java
@@ -80,6 +80,12 @@ public interface ValueVector extends Closeable, Iterable<ValueVector> {
    */
   void reAlloc();
 
+  /**
+   * Get the allocator associated with the vector. CAVEAT: Some ValueVector subclasses (e.g.
+   * NullVector) do not require an allocator for data storage and may return null.
+   *
+   * @return Returns nullable allocator.
+   */
   BufferAllocator getAllocator();
 
   /**
diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java b/java/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java
index a3f25bc5207b6..6aace956214ff 100644
--- a/java/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java
+++ b/java/vector/src/test/java/org/apache/arrow/vector/TestSplitAndTransfer.java
@@ -198,6 +198,31 @@ public void testWithEmptyVector() {
     toDUV.clear();
   }
 
+  @Test
+  public void testWithNullVector() {
+    int valueCount = 123;
+    int startIndex = 10;
+    NullVector fromNullVector = new NullVector("nullVector");
+    fromNullVector.setValueCount(valueCount);
+    TransferPair transferPair = fromNullVector.getTransferPair(fromNullVector.getAllocator());
+    transferPair.splitAndTransfer(startIndex, valueCount - startIndex);
+    NullVector toNullVector = (NullVector) transferPair.getTo();
+
+    assertEquals(valueCount - startIndex, toNullVector.getValueCount());
+    // no allocations to clear for NullVector
+  }
+
+  @Test
+  public void testWithZeroVector() {
+    ZeroVector fromZeroVector = new ZeroVector("zeroVector");
+    TransferPair transferPair = fromZeroVector.getTransferPair(fromZeroVector.getAllocator());
+    transferPair.splitAndTransfer(0, 0);
+    ZeroVector toZeroVector = (ZeroVector) transferPair.getTo();
+
+    assertEquals(0, toZeroVector.getValueCount());
+    // no allocations to clear for ZeroVector
+  }
+
   @Test /* VarCharVector */
   public void test() throws Exception {
     try (final VarCharVector varCharVector = new VarCharVector("myvector", allocator)) {