Skip to content

Commit

Permalink
apacheGH-37705: [Java] Extra input methods for binary writers (apache…
Browse files Browse the repository at this point in the history
…#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: apache#37705

Authored-by: James Duong <duong.james@gmail.com>
Signed-off-by: David Li <li.davidm96@gmail.com>
  • Loading branch information
jduo authored and loicalleyne committed Nov 13, 2023
1 parent 2b8bf1a commit 5be4ac0
Show file tree
Hide file tree
Showing 3 changed files with 189 additions and 0 deletions.
18 changes: 18 additions & 0 deletions java/vector/src/main/codegen/templates/AbstractFieldWriter.java
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,24 @@ public void write(${name}Holder holder) {
}
</#if>

<#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}");
}
</#if>

</#list></#list>

public void writeNull() {
Expand Down
33 changes: 33 additions & 0 deletions java/vector/src/main/codegen/templates/ComplexWriters.java
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,28 @@ public void writeNull() {
vector.setValueCount(idx()+1);
}
</#if>
<#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);
}
</#if>
}
<@pp.changeOutputFile name="/org/apache/arrow/vector/complex/writer/${eName}Writer.java" />
Expand Down Expand Up @@ -223,6 +245,17 @@ public interface ${eName}Writer extends BaseWriter {
@Deprecated
public void writeBigEndianBytesTo${minor.class}(byte[] value);
</#if>
<#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);
</#if>

}

</#list>
Expand Down
Original file line number Diff line number Diff line change
@@ -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);
}
}
}

0 comments on commit 5be4ac0

Please sign in to comment.