From f52977e0fc092c58923cc568fbbd2ee40b92d230 Mon Sep 17 00:00:00 2001 From: James Duong Date: Tue, 26 Sep 2023 10:35:11 -0700 Subject: [PATCH] GH-37705: [Java] Extra input methods for VarChar writers Add write() methods for Text and String types. Ensure these methods are part of the writer interfaces and not just the Impls. --- .../templates/AbstractFieldWriter.java | 10 +++ .../codegen/templates/ComplexWriters.java | 21 ++++- .../complex/writer/TestSimpleWriter.java | 81 +++++++++++++++---- 3 files changed, 94 insertions(+), 18 deletions(-) diff --git a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java index 5e6580b6131c1..bb4ee45eaa073 100644 --- a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java +++ b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java @@ -142,6 +142,16 @@ public void write(${name}Holder holder) { } + <#if minor.class?ends_with("VarChar")> + public void write${minor.class}(${friendlyType} value) { + fail("${name}"); + } + + public void write${minor.class}(String value) { + 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 4ae4c4f75f208..51a52a6e3070d 100644 --- a/java/vector/src/main/codegen/templates/ComplexWriters.java +++ b/java/vector/src/main/codegen/templates/ComplexWriters.java @@ -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(); + + +public ${eName}WriterImpl(${name}Vector vector) { this.vector = vector; } @@ -120,11 +124,19 @@ public void write(Nullable${minor.class}Holder h) { } - <#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 minor.class?starts_with("Decimal")> @@ -256,6 +268,11 @@ public interface ${eName}Writer extends BaseWriter { public void writeTo${minor.class}(ByteBuffer value, int offset, int length); +<#if minor.class?ends_with("VarChar")> + public void write${minor.class}(${friendlyType} value); + + public void write${minor.class}(String value); + } 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 index 7c06509b23c87..ef918b13fb691 100644 --- 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 @@ -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; @@ -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); @@ -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); @@ -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); @@ -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); @@ -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); @@ -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); @@ -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); @@ -125,9 +130,9 @@ 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); @@ -135,4 +140,48 @@ public void testWriteByteBufferWithOffsetToLargeVarBinary() { 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); + } + } }