Skip to content

Commit

Permalink
GH-37705: [Java] Extra input methods for VarChar writers
Browse files Browse the repository at this point in the history 
Add write() methods for Text and String types.
Ensure these methods are part of the writer interfaces and not
just the Impls.
  • Loading branch information
@jduo
jduo committed Sep 26, 2023
1 parent 2895af4 commit dfd280f
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 17 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
14 changes: 13 additions & 1 deletion java/vector/src/main/codegen/templates/ComplexWriters.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,11 +120,18 @@ 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) {
vector.setSafe(idx(), new Text(value));
vector.setValueCount(idx()+1);
}
</#if>
<#if minor.class?starts_with("Decimal")>
Expand Down Expand Up @@ -256,6 +263,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 dfd280f

Please sign in to comment.