Skip to content

Commit

Permalink
apacheGH-37705: [Java] Extra input methods for VarChar writers (apach…
Browse files Browse the repository at this point in the history 
…e#37883)

### Rationale for this change
Improve the convenience of using VarCharWriter and LargeVarCharWriter interfaces.
Also allow users to avoid unnecessary overhead creating Arrow buffers when writing
String and Text data.

### What changes are included in this PR?
Add write() methods for Text and String types.
Ensure these methods are part of the writer interfaces and not just the Impls.### Are these changes tested?

### Are these changes tested?
Yes.

### Are there any user-facing changes?
No.
* Closes: apache#37705

Authored-by: James Duong <[email protected]>
Signed-off-by: David Li <[email protected]>
  • Loading branch information
@jduo @dgreiss
jduo authored and dgreiss committed Feb 17, 2024
1 parent 7f3bdd6 commit 31464cc
Show file tree
Hide file tree
Showing 3 changed files with 94 additions and 18 deletions.
10 changes: 10 additions & 0 deletions java/vector/src/main/codegen/templates/AbstractFieldWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,16 @@ public void write(${name}Holder holder) {
}
</#if>

<#if minor.class?ends_with("VarChar")>
public void write${minor.class}(${friendlyType} value) {
fail("${name}");
}

public void write${minor.class}(String value) {
fail("${name}");
}
</#if>

</#list></#list>

public void writeNull() {
Expand Down
21 changes: 19 additions & 2 deletions java/vector/src/main/codegen/templates/ComplexWriters.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,11 @@ public class ${eName}WriterImpl extends AbstractFieldWriter {

final ${name}Vector vector;

public ${eName}WriterImpl(${name}Vector vector) {
<#if minor.class?ends_with("VarChar")>
private final Text textBuffer = new Text();
</#if>

public ${eName}WriterImpl(${name}Vector vector) {
this.vector = vector;
}

Expand Down Expand Up @@ -120,11 +124,19 @@ public void write(Nullable${minor.class}Holder h) {
}
</#if>

<#if minor.class == "VarChar">
<#if minor.class?ends_with("VarChar")>
@Override
public void write${minor.class}(${friendlyType} value) {
vector.setSafe(idx(), value);
vector.setValueCount(idx()+1);
}

@Override
public void write${minor.class}(String value) {
textBuffer.set(value);
vector.setSafe(idx(), textBuffer);
vector.setValueCount(idx()+1);
}
</#if>

<#if minor.class?starts_with("Decimal")>
Expand Down Expand Up @@ -256,6 +268,11 @@ public interface ${eName}Writer extends BaseWriter {
public void writeTo${minor.class}(ByteBuffer value, int offset, int length);
</#if>

<#if minor.class?ends_with("VarChar")>
public void write${minor.class}(${friendlyType} value);

public void write${minor.class}(String value);
</#if>
}

</#list>
Expand Down
81 changes: 65 additions & 16 deletions java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,14 @@
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.LargeVarBinaryVector;
import org.apache.arrow.vector.LargeVarCharVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.complex.impl.LargeVarBinaryWriterImpl;
import org.apache.arrow.vector.complex.impl.LargeVarCharWriterImpl;
import org.apache.arrow.vector.complex.impl.VarBinaryWriterImpl;
import org.apache.arrow.vector.complex.impl.VarCharWriterImpl;
import org.apache.arrow.vector.util.Text;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
Expand All @@ -45,9 +50,9 @@ public void terminate() throws Exception {
}

@Test
public void testWriteByteArrayToVarBinary() {
public void testWriteByteArrayToVarBinary() throws Exception {
try (VarBinaryVector vector = new VarBinaryVector("test", allocator);
VarBinaryWriterImpl writer = new VarBinaryWriterImpl(vector)) {
VarBinaryWriter writer = new VarBinaryWriterImpl(vector)) {
byte[] input = new byte[] { 0x01, 0x02 };
writer.writeToVarBinary(input);
byte[] result = vector.get(0);
Expand All @@ -56,9 +61,9 @@ public void testWriteByteArrayToVarBinary() {
}

@Test
public void testWriteByteArrayWithOffsetToVarBinary() {
public void testWriteByteArrayWithOffsetToVarBinary() throws Exception {
try (VarBinaryVector vector = new VarBinaryVector("test", allocator);
VarBinaryWriterImpl writer = new VarBinaryWriterImpl(vector)) {
VarBinaryWriter writer = new VarBinaryWriterImpl(vector)) {
byte[] input = new byte[] { 0x01, 0x02 };
writer.writeToVarBinary(input, 1, 1);
byte[] result = vector.get(0);
Expand All @@ -67,9 +72,9 @@ public void testWriteByteArrayWithOffsetToVarBinary() {
}

@Test
public void testWriteByteBufferToVarBinary() {
public void testWriteByteBufferToVarBinary() throws Exception {
try (VarBinaryVector vector = new VarBinaryVector("test", allocator);
VarBinaryWriterImpl writer = new VarBinaryWriterImpl(vector)) {
VarBinaryWriter writer = new VarBinaryWriterImpl(vector)) {
byte[] input = new byte[] { 0x01, 0x02 };
ByteBuffer buffer = ByteBuffer.wrap(input);
writer.writeToVarBinary(buffer);
Expand All @@ -79,9 +84,9 @@ public void testWriteByteBufferToVarBinary() {
}

@Test
public void testWriteByteBufferWithOffsetToVarBinary() {
public void testWriteByteBufferWithOffsetToVarBinary() throws Exception {
try (VarBinaryVector vector = new VarBinaryVector("test", allocator);
VarBinaryWriterImpl writer = new VarBinaryWriterImpl(vector)) {
VarBinaryWriter writer = new VarBinaryWriterImpl(vector)) {
byte[] input = new byte[] { 0x01, 0x02 };
ByteBuffer buffer = ByteBuffer.wrap(input);
writer.writeToVarBinary(buffer, 1, 1);
Expand All @@ -91,9 +96,9 @@ public void testWriteByteBufferWithOffsetToVarBinary() {
}

@Test
public void testWriteByteArrayToLargeVarBinary() {
public void testWriteByteArrayToLargeVarBinary() throws Exception {
try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator);
LargeVarBinaryWriterImpl writer = new LargeVarBinaryWriterImpl(vector)) {
LargeVarBinaryWriter writer = new LargeVarBinaryWriterImpl(vector)) {
byte[] input = new byte[] { 0x01, 0x02 };
writer.writeToLargeVarBinary(input);
byte[] result = vector.get(0);
Expand All @@ -102,9 +107,9 @@ public void testWriteByteArrayToLargeVarBinary() {
}

@Test
public void testWriteByteArrayWithOffsetToLargeVarBinary() {
public void testWriteByteArrayWithOffsetToLargeVarBinary() throws Exception {
try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator);
LargeVarBinaryWriterImpl writer = new LargeVarBinaryWriterImpl(vector)) {
LargeVarBinaryWriter writer = new LargeVarBinaryWriterImpl(vector)) {
byte[] input = new byte[] { 0x01, 0x02 };
writer.writeToLargeVarBinary(input, 1, 1);
byte[] result = vector.get(0);
Expand All @@ -113,9 +118,9 @@ public void testWriteByteArrayWithOffsetToLargeVarBinary() {
}

@Test
public void testWriteByteBufferToLargeVarBinary() {
public void testWriteByteBufferToLargeVarBinary() throws Exception {
try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator);
LargeVarBinaryWriterImpl writer = new LargeVarBinaryWriterImpl(vector)) {
LargeVarBinaryWriter writer = new LargeVarBinaryWriterImpl(vector)) {
byte[] input = new byte[] { 0x01, 0x02 };
ByteBuffer buffer = ByteBuffer.wrap(input);
writer.writeToLargeVarBinary(buffer);
Expand All @@ -125,14 +130,58 @@ public void testWriteByteBufferToLargeVarBinary() {
}

@Test
public void testWriteByteBufferWithOffsetToLargeVarBinary() {
public void testWriteByteBufferWithOffsetToLargeVarBinary() throws Exception {
try (LargeVarBinaryVector vector = new LargeVarBinaryVector("test", allocator);
LargeVarBinaryWriterImpl writer = new LargeVarBinaryWriterImpl(vector)) {
LargeVarBinaryWriter 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);
}
}

@Test
public void testWriteStringToVarChar() throws Exception {
try (VarCharVector vector = new VarCharVector("test", allocator);
VarCharWriter writer = new VarCharWriterImpl(vector)) {
String input = "testInput";
writer.writeVarChar(input);
String result = vector.getObject(0).toString();
Assert.assertEquals(input, result);
}
}

@Test
public void testWriteTextToVarChar() throws Exception {
try (VarCharVector vector = new VarCharVector("test", allocator);
VarCharWriter writer = new VarCharWriterImpl(vector)) {
String input = "testInput";
writer.writeVarChar(new Text(input));
String result = vector.getObject(0).toString();
Assert.assertEquals(input, result);
}
}

@Test
public void testWriteStringToLargeVarChar() throws Exception {
try (LargeVarCharVector vector = new LargeVarCharVector("test", allocator);
LargeVarCharWriter writer = new LargeVarCharWriterImpl(vector)) {
String input = "testInput";
writer.writeLargeVarChar(input);
String result = vector.getObject(0).toString();
Assert.assertEquals(input, result);
}
}

@Test
public void testWriteTextToLargeVarChar() throws Exception {
try (LargeVarCharVector vector = new LargeVarCharVector("test", allocator);
LargeVarCharWriter writer = new LargeVarCharWriterImpl(vector)) {
String input = "testInput";
writer.writeLargeVarChar(new Text(input));
String result = vector.getObject(0).toString();
Assert.assertEquals(input, result);
}
}
}

0 comments on commit 31464cc

Please sign in to comment.