Skip to content

Commit b11bc50

Browse files
igor-suhorukovigor-suhorukov
authored
ARROW-17430: [Java] ListBinder to bind Arrow List type to DB column (#13906)
Typical real life Arrow datasets contain List type vectors of primitive type. This PR introduce ListBinder mapping of primitive types lists to java.sql.Types.ARRAY Lead-authored-by: Igor Suhorukov <[email protected]> Co-authored-by: igor.suhorukov <[email protected]> Signed-off-by: David Li <[email protected]>
1 parent 1a34a07 commit b11bc50

3 files changed

Lines changed: 267 additions & 1 deletion

File tree

‎java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/ColumnBinderArrowTypeVisitor.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
import org.apache.arrow.vector.TinyIntVector;
4545
import org.apache.arrow.vector.VarBinaryVector;
4646
import org.apache.arrow.vector.VarCharVector;
47+
import org.apache.arrow.vector.complex.ListVector;
4748
import org.apache.arrow.vector.types.pojo.ArrowType;
4849

4950
/**
@@ -78,7 +79,7 @@ public ColumnBinder visit(ArrowType.Struct type) {
7879

7980
@Override
8081
public ColumnBinder visit(ArrowType.List type) {
81-
throw new UnsupportedOperationException("No column binder implemented for type " + type);
82+
return new ListBinder((ListVector) vector);
8283
}
8384

8485
@Override

‎java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/ListBinder.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.apache.arrow.adapter.jdbc.binder;
19+
20+
import java.lang.reflect.Array;
21+
import java.util.ArrayList;
22+
import java.util.Arrays;
23+
24+
import org.apache.arrow.vector.FieldVector;
25+
import org.apache.arrow.vector.complex.ListVector;
26+
import org.apache.arrow.vector.complex.impl.UnionListReader;
27+
import org.apache.arrow.vector.util.Text;
28+
29+
/**
30+
* A column binder for list of primitive values.
31+
*/
32+
public class ListBinder extends BaseColumnBinder<ListVector> {
33+
34+
private final UnionListReader listReader;
35+
private final Class<?> arrayElementClass;
36+
private final boolean isTextColumn;
37+
38+
public ListBinder(ListVector vector) {
39+
this(vector, java.sql.Types.ARRAY);
40+
}
41+
42+
/**
43+
* Init ListBinder and determine type of data vector.
44+
*
45+
* @param vector corresponding data vector from arrow buffer for binding
46+
* @param jdbcType parameter jdbc type
47+
*/
48+
public ListBinder(ListVector vector, int jdbcType) {
49+
super(vector, jdbcType);
50+
listReader = vector.getReader();
51+
Class<? extends FieldVector> dataVectorClass = vector.getDataVector().getClass();
52+
try {
53+
arrayElementClass = dataVectorClass.getMethod("getObject", Integer.TYPE).getReturnType();
54+
} catch (NoSuchMethodException e) {
55+
final String message = String.format("Issue to determine type for getObject method of data vector class %s ",
56+
dataVectorClass.getName());
57+
throw new RuntimeException(message);
58+
}
59+
isTextColumn = arrayElementClass.isAssignableFrom(Text.class);
60+
}
61+
62+
@Override
63+
public void bind(java.sql.PreparedStatement statement, int parameterIndex, int rowIndex)throws java.sql.SQLException {
64+
listReader.setPosition(rowIndex);
65+
ArrayList<?> sourceArray = (ArrayList<?>) listReader.readObject();
66+
Object array;
67+
if (!isTextColumn) {
68+
array = Array.newInstance(arrayElementClass, sourceArray.size());
69+
Arrays.setAll((Object[]) array, sourceArray::get);
70+
} else {
71+
array = new String[sourceArray.size()];
72+
Arrays.setAll((Object[]) array, idx -> sourceArray.get(idx) != null ? sourceArray.get(idx).toString() : null);
73+
}
74+
statement.setObject(parameterIndex, array);
75+
}
76+
}

‎java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/JdbcParameterBinderTest.java‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import java.util.function.BiConsumer;
3434

3535
import org.apache.arrow.adapter.jdbc.binder.ColumnBinder;
36+
import org.apache.arrow.memory.ArrowBuf;
3637
import org.apache.arrow.memory.BufferAllocator;
3738
import org.apache.arrow.memory.RootAllocator;
3839
import org.apache.arrow.vector.BaseLargeVariableWidthVector;
@@ -67,11 +68,13 @@
6768
import org.apache.arrow.vector.VarBinaryVector;
6869
import org.apache.arrow.vector.VarCharVector;
6970
import org.apache.arrow.vector.VectorSchemaRoot;
71+
import org.apache.arrow.vector.complex.ListVector;
7072
import org.apache.arrow.vector.types.DateUnit;
7173
import org.apache.arrow.vector.types.FloatingPointPrecision;
7274
import org.apache.arrow.vector.types.TimeUnit;
7375
import org.apache.arrow.vector.types.pojo.ArrowType;
7476
import org.apache.arrow.vector.types.pojo.Field;
77+
import org.apache.arrow.vector.types.pojo.FieldType;
7578
import org.apache.arrow.vector.types.pojo.Schema;
7679
import org.junit.jupiter.api.AfterEach;
7780
import org.junit.jupiter.api.BeforeEach;
@@ -385,6 +388,91 @@ void decimal256() throws SQLException {
385388
Arrays.asList(new BigDecimal("120.429"), new BigDecimal("-10590.123"), new BigDecimal("0.000")));
386389
}
387390

391+
@Test
392+
void listOfDouble() throws SQLException {
393+
TriConsumer<ListVector, Integer, Double[]> setValue = (listVector, index, values) -> {
394+
org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
395+
writer.setPosition(index);
396+
writer.startList();
397+
Arrays.stream(values).forEach(doubleValue -> writer.float8().writeFloat8(doubleValue));
398+
writer.endList();
399+
listVector.setLastSet(index);
400+
};
401+
List<Double[]> values = Arrays.asList(new Double[]{0.0, Math.PI}, new Double[]{1.1, -352346.2, 2355.6},
402+
new Double[]{-1024.3}, new Double[]{});
403+
testListType(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), setValue, ListVector::setNull, values);
404+
}
405+
406+
@Test
407+
void listOfInt64() throws SQLException {
408+
TriConsumer<ListVector, Integer, Long[]> setValue = (listVector, index, values) -> {
409+
org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
410+
writer.setPosition(index);
411+
writer.startList();
412+
Arrays.stream(values).forEach(longValue -> writer.bigInt().writeBigInt(longValue));
413+
writer.endList();
414+
listVector.setLastSet(index);
415+
};
416+
List<Long[]> values = Arrays.asList(new Long[]{1L, 2L, 3L}, new Long[]{4L, 5L},
417+
new Long[]{512L, 1024L, 2048L, 4096L}, new Long[]{});
418+
testListType((ArrowType) new ArrowType.Int(64, true), setValue, ListVector::setNull, values);
419+
}
420+
421+
@Test
422+
void listOfInt32() throws SQLException {
423+
TriConsumer<ListVector, Integer, Integer[]> setValue = (listVector, index, values) -> {
424+
org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
425+
writer.setPosition(index);
426+
writer.startList();
427+
Arrays.stream(values).forEach(integerValue -> writer.integer().writeInt(integerValue));
428+
writer.endList();
429+
listVector.setLastSet(index);
430+
};
431+
List<Integer[]> values = Arrays.asList(new Integer[]{1, 2, 3}, new Integer[]{4, 5},
432+
new Integer[]{512, 1024, 2048, 4096}, new Integer[]{});
433+
testListType((ArrowType) new ArrowType.Int(32, true), setValue, ListVector::setNull, values);
434+
}
435+
436+
@Test
437+
void listOfBoolean() throws SQLException {
438+
TriConsumer<ListVector, Integer, Boolean[]> setValue = (listVector, index, values) -> {
439+
org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
440+
writer.setPosition(index);
441+
writer.startList();
442+
Arrays.stream(values).forEach(booleanValue -> writer.bit().writeBit(booleanValue ? 1 : 0));
443+
writer.endList();
444+
listVector.setLastSet(index);
445+
};
446+
List<Boolean[]> values = Arrays.asList(new Boolean[]{true, false},
447+
new Boolean[]{false, false}, new Boolean[]{true, true, false, true}, new Boolean[]{});
448+
testListType((ArrowType) new ArrowType.Bool(), setValue, ListVector::setNull, values);
449+
}
450+
451+
@Test
452+
void listOfString() throws SQLException {
453+
TriConsumer<ListVector, Integer, String[]> setValue = (listVector, index, values) -> {
454+
org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
455+
writer.setPosition(index);
456+
writer.startList();
457+
Arrays.stream(values).forEach(stringValue -> {
458+
if (stringValue != null) {
459+
byte[] stringValueBytes = stringValue.getBytes(StandardCharsets.UTF_8);
460+
try (ArrowBuf stringBuffer = allocator.buffer(stringValueBytes.length)) {
461+
stringBuffer.writeBytes(stringValueBytes);
462+
writer.varChar().writeVarChar(0, stringValueBytes.length, stringBuffer);
463+
}
464+
} else {
465+
writer.varChar().writeNull();
466+
}
467+
});
468+
writer.endList();
469+
listVector.setLastSet(index);
470+
};
471+
List<String[]> values = Arrays.asList(new String[]{"aaaa", "b1"},
472+
new String[]{"c", null, "d"}, new String[]{"e", "f", "g", "h"}, new String[]{});
473+
testListType((ArrowType) new ArrowType.Utf8(), setValue, ListVector::setNull, values);
474+
}
475+
388476
@FunctionalInterface
389477
interface TriConsumer<T, U, V> {
390478
void accept(T value1, U value2, V value3);
@@ -483,4 +571,105 @@ <T, V extends FieldVector> void testSimpleType(ArrowType arrowType, int jdbcType
483571
assertThat(binder.next()).isFalse();
484572
}
485573
}
574+
575+
<T, V extends FieldVector> void testListType(ArrowType arrowType, TriConsumer<V, Integer, T> setValue,
576+
BiConsumer<V, Integer> setNull, List<T> values) throws SQLException {
577+
int jdbcType = Types.ARRAY;
578+
Schema schema = new Schema(Collections.singletonList(new Field("field", FieldType.nullable(
579+
new ArrowType.List()), Collections.singletonList(
580+
new Field("element", FieldType.notNullable(arrowType), null)
581+
))));
582+
try (final MockPreparedStatement statement = new MockPreparedStatement();
583+
final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
584+
final JdbcParameterBinder binder =
585+
JdbcParameterBinder.builder(statement, root).bindAll().build();
586+
assertThat(binder.next()).isFalse();
587+
588+
@SuppressWarnings("unchecked")
589+
final V vector = (V) root.getVector(0);
590+
final ColumnBinder columnBinder = ColumnBinder.forVector(vector);
591+
assertThat(columnBinder.getJdbcType()).isEqualTo(jdbcType);
592+
593+
setValue.accept(vector, 0, values.get(0));
594+
setValue.accept(vector, 1, values.get(1));
595+
setNull.accept(vector, 2);
596+
root.setRowCount(3);
597+
598+
assertThat(binder.next()).isTrue();
599+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
600+
assertThat(binder.next()).isTrue();
601+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(1));
602+
assertThat(binder.next()).isTrue();
603+
assertThat(statement.getParamValue(1)).isNull();
604+
assertThat(statement.getParamType(1)).isEqualTo(jdbcType);
605+
assertThat(binder.next()).isFalse();
606+
607+
binder.reset();
608+
609+
setNull.accept(vector, 0);
610+
setValue.accept(vector, 1, values.get(3));
611+
setValue.accept(vector, 2, values.get(0));
612+
setValue.accept(vector, 3, values.get(2));
613+
setValue.accept(vector, 4, values.get(1));
614+
root.setRowCount(5);
615+
616+
assertThat(binder.next()).isTrue();
617+
assertThat(statement.getParamValue(1)).isNull();
618+
assertThat(statement.getParamType(1)).isEqualTo(jdbcType);
619+
assertThat(binder.next()).isTrue();
620+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(3));
621+
assertThat(binder.next()).isTrue();
622+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
623+
assertThat(binder.next()).isTrue();
624+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(2));
625+
assertThat(binder.next()).isTrue();
626+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(1));
627+
assertThat(binder.next()).isFalse();
628+
}
629+
630+
// Non-nullable (since some types have a specialized binder)
631+
schema = new Schema(Collections.singletonList(new Field("field", FieldType.notNullable(
632+
new ArrowType.List()), Collections.singletonList(
633+
new Field("element", FieldType.notNullable(arrowType), null)
634+
))));
635+
try (final MockPreparedStatement statement = new MockPreparedStatement();
636+
final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
637+
final JdbcParameterBinder binder =
638+
JdbcParameterBinder.builder(statement, root).bindAll().build();
639+
assertThat(binder.next()).isFalse();
640+
641+
@SuppressWarnings("unchecked")
642+
final V vector = (V) root.getVector(0);
643+
setValue.accept(vector, 0, values.get(0));
644+
setValue.accept(vector, 1, values.get(1));
645+
root.setRowCount(2);
646+
647+
assertThat(binder.next()).isTrue();
648+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
649+
assertThat(binder.next()).isTrue();
650+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(1));
651+
assertThat(binder.next()).isFalse();
652+
653+
binder.reset();
654+
655+
setValue.accept(vector, 0, values.get(0));
656+
setValue.accept(vector, 1, values.get(2));
657+
setValue.accept(vector, 2, values.get(0));
658+
setValue.accept(vector, 3, values.get(2));
659+
setValue.accept(vector, 4, values.get(1));
660+
root.setRowCount(5);
661+
662+
assertThat(binder.next()).isTrue();
663+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
664+
assertThat(binder.next()).isTrue();
665+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(2));
666+
assertThat(binder.next()).isTrue();
667+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
668+
assertThat(binder.next()).isTrue();
669+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(2));
670+
assertThat(binder.next()).isTrue();
671+
assertThat(statement.getParamValue(1)).isEqualTo(values.get(1));
672+
assertThat(binder.next()).isFalse();
673+
}
674+
}
486675
}

0 commit comments

Comments
 (0)