From 57df50135213da10e6e2f5dfaea3b9f8d1801d9c Mon Sep 17 00:00:00 2001 From: James Duong Date: Tue, 19 Sep 2023 11:10:37 -0700 Subject: [PATCH] GH-37705: [Java] Extra input methods for binary writers (#37791) ### Rationale for this change ByteBuffer and byte[] are commonly used to hold binary data. The current writers require working with ArrowBuf objects which need to be populated by copying from these types, then copying into the vector. ### What changes are included in this PR? Add methods to VarBinary and LargeVarBinary writers to take in common binary parameters - byte[] and ByteBuffer. The writer now sets these objects on the Vectors directly. ### Are these changes tested? Yes. * Closes: #37705 Authored-by: James Duong Signed-off-by: David Li --- .../templates/AbstractFieldWriter.java | 18 +++ .../codegen/templates/ComplexWriters.java | 33 +++++ .../complex/writer/TestSimpleWriter.java | 138 ++++++++++++++++++ 3 files changed, 189 insertions(+) create mode 100644 java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java diff --git a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java index 1f80f25266b57..5e6580b6131c1 100644 --- a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java +++ b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java @@ -124,6 +124,24 @@ public void write(${name}Holder holder) { } + <#if minor.class?ends_with("VarBinary")> + public void writeTo${minor.class}(byte[] value) { + fail("${name}"); + } + + public void writeTo${minor.class}(byte[] value, int offset, int length) { + fail("${name}"); + } + + public void writeTo${minor.class}(ByteBuffer value) { + fail("${name}"); + } + + public void writeTo${minor.class}(ByteBuffer value, int offset, int length) { + fail("${name}"); + } + + public void writeNull() { diff --git a/java/vector/src/main/codegen/templates/ComplexWriters.java b/java/vector/src/main/codegen/templates/ComplexWriters.java index 0b1e321afb70e..4ae4c4f75f208 100644 --- a/java/vector/src/main/codegen/templates/ComplexWriters.java +++ b/java/vector/src/main/codegen/templates/ComplexWriters.java @@ -180,6 +180,28 @@ public void writeNull() { vector.setValueCount(idx()+1); } + + <#if minor.class?ends_with("VarBinary")> + public void writeTo${minor.class}(byte[] value) { + vector.setSafe(idx(), value); + vector.setValueCount(idx() + 1); + } + + public void writeTo${minor.class}(byte[] value, int offset, int length) { + vector.setSafe(idx(), value, offset, length); + vector.setValueCount(idx() + 1); + } + + public void writeTo${minor.class}(ByteBuffer value) { + vector.setSafe(idx(), value, 0, value.remaining()); + vector.setValueCount(idx() + 1); + } + + public void writeTo${minor.class}(ByteBuffer value, int offset, int length) { + vector.setSafe(idx(), value, offset, length); + vector.setValueCount(idx() + 1); + } + } <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/writer/${eName}Writer.java" /> @@ -223,6 +245,17 @@ public interface ${eName}Writer extends BaseWriter { @Deprecated public void writeBigEndianBytesTo${minor.class}(byte[] value); + +<#if minor.class?ends_with("VarBinary")> + public void writeTo${minor.class}(byte[] value); + + public void writeTo${minor.class}(byte[] value, int offset, int length); + + public void writeTo${minor.class}(ByteBuffer value); + + public void writeTo${minor.class}(ByteBuffer value, int offset, int length); + + } diff --git a/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java new file mode 100644 index 0000000000000..7c06509b23c87 --- /dev/null +++ b/java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java @@ -0,0 +1,138 @@ +/* + * 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.arrow.vector.complex.writer; + +import java.nio.ByteBuffer; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.LargeVarBinaryVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.complex.impl.LargeVarBinaryWriterImpl; +import org.apache.arrow.vector.complex.impl.VarBinaryWriterImpl; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +public class TestSimpleWriter { + + private BufferAllocator allocator; + + @Before + public void init() { + allocator = new RootAllocator(Integer.MAX_VALUE); + } + + @After + public void terminate() throws Exception { + allocator.close(); + } + + @Test + public void testWriteByteArrayToVarBinary() { + try (VarBinaryVector vector = new VarBinaryVector("test", allocator); + VarBinaryWriterImpl writer = new VarBinaryWriterImpl(vector)) { + byte[] input = new byte[] { 0x01, 0x02 }; + writer.writeToVarBinary(input); + byte[] result = vector.get(0); + Assert.assertArrayEquals(input, result); + } + } + + @Test + public void testWriteByteArrayWithOffsetToVarBinary() { + try (VarBinaryVector vector = new VarBinaryVector("test", allocator); + VarBinaryWriterImpl writer = new VarBinaryWriterImpl(vector)) { + byte[] input = new byte[] { 0x01, 0x02 }; + writer.writeToVarBinary(input, 1, 1); + byte[] result = vector.get(0); + Assert.assertArrayEquals(new byte[] { 0x02 }, result); + } + } + + @Test + public void testWriteByteBufferToVarBinary() { + try (VarBinaryVector vector = new VarBinaryVector("test", allocator); + VarBinaryWriterImpl writer = new VarBinaryWriterImpl(vector)) { + byte[] input = new byte[] { 0x01, 0x02 }; + ByteBuffer buffer = ByteBuffer.wrap(input); + writer.writeToVarBinary(buffer); + byte[] result = vector.get(0); + Assert.assertArrayEquals(input, result); + } + } + + @Test + public void testWriteByteBufferWithOffsetToVarBinary() { + try (VarBinaryVector vector = new VarBinaryVector("test", allocator); + VarBinaryWriterImpl writer = new VarBinaryWriterImpl(vector)) { + byte[] input = new byte[] { 0x01, 0x02 }; + ByteBuffer buffer = ByteBuffer.wrap(input); + writer.writeToVarBinary(buffer, 1, 1); + byte[] result = vector.get(0); + Assert.assertArrayEquals(new byte[] { 0x02 }, result); + } + } + + @Test + public void testWriteByteArrayToLargeVarBinary() { + try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator); + LargeVarBinaryWriterImpl writer = new LargeVarBinaryWriterImpl(vector)) { + byte[] input = new byte[] { 0x01, 0x02 }; + writer.writeToLargeVarBinary(input); + byte[] result = vector.get(0); + Assert.assertArrayEquals(input, result); + } + } + + @Test + public void testWriteByteArrayWithOffsetToLargeVarBinary() { + try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator); + LargeVarBinaryWriterImpl writer = new LargeVarBinaryWriterImpl(vector)) { + byte[] input = new byte[] { 0x01, 0x02 }; + writer.writeToLargeVarBinary(input, 1, 1); + byte[] result = vector.get(0); + Assert.assertArrayEquals(new byte[] { 0x02 }, result); + } + } + + @Test + public void testWriteByteBufferToLargeVarBinary() { + try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator); + LargeVarBinaryWriterImpl writer = new LargeVarBinaryWriterImpl(vector)) { + byte[] input = new byte[] { 0x01, 0x02 }; + ByteBuffer buffer = ByteBuffer.wrap(input); + writer.writeToLargeVarBinary(buffer); + byte[] result = vector.get(0); + Assert.assertArrayEquals(input, result); + } + } + + @Test + public void testWriteByteBufferWithOffsetToLargeVarBinary() { + try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator); + LargeVarBinaryWriterImpl writer = new LargeVarBinaryWriterImpl(vector)) { + byte[] input = new byte[] { 0x01, 0x02 }; + ByteBuffer buffer = ByteBuffer.wrap(input); + writer.writeToLargeVarBinary(buffer, 1, 1); + byte[] result = vector.get(0); + Assert.assertArrayEquals(new byte[] { 0x02 }, result); + } + } +}