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]>
apache · Aug 19, 2022 · b11bc50 · b11bc50
1 parent 1a34a07
commit b11bc50
Show file tree

Hide file tree

Showing 3 changed files with 267 additions and 1 deletion.
diff --git a/...jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/ColumnBinderArrowTypeVisitor.java b/...jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/ColumnBinderArrowTypeVisitor.java
@@ -44,6 +44,7 @@
 import org.apache.arrow.vector.TinyIntVector;
 import org.apache.arrow.vector.VarBinaryVector;
 import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.complex.ListVector;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 
 /**
@@ -78,7 +79,7 @@ public ColumnBinder visit(ArrowType.Struct type) {
 
   @Override
   public ColumnBinder visit(ArrowType.List type) {
-    throw new UnsupportedOperationException("No column binder implemented for type " + type);
+    return new ListBinder((ListVector) vector);
   }
 
   @Override

diff --git a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/ListBinder.java b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/ListBinder.java
@@ -0,0 +1,76 @@
+/*
+ * 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.adapter.jdbc.binder;
+
+import java.lang.reflect.Array;
+import java.util.ArrayList;
+import java.util.Arrays;
+
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.impl.UnionListReader;
+import org.apache.arrow.vector.util.Text;
+
+/**
+ * A column binder for list of primitive values.
+ */
+public class ListBinder extends BaseColumnBinder<ListVector> {
+
+  private final UnionListReader listReader;
+  private final Class<?> arrayElementClass;
+  private final boolean isTextColumn;
+
+  public ListBinder(ListVector vector) {
+    this(vector, java.sql.Types.ARRAY);
+  }
+
+  /**
+   * Init ListBinder and determine type of data vector.
+   *
+   * @param vector corresponding data vector from arrow buffer for binding
+   * @param jdbcType parameter jdbc type
+   */
+  public ListBinder(ListVector vector, int jdbcType) {
+    super(vector, jdbcType);
+    listReader = vector.getReader();
+    Class<? extends FieldVector> dataVectorClass = vector.getDataVector().getClass();
+    try {
+      arrayElementClass = dataVectorClass.getMethod("getObject", Integer.TYPE).getReturnType();
+    } catch (NoSuchMethodException e) {
+      final String message = String.format("Issue to determine type for getObject method of data vector class %s ",
+              dataVectorClass.getName());
+      throw new RuntimeException(message);
+    }
+    isTextColumn = arrayElementClass.isAssignableFrom(Text.class);
+  }
+
+  @Override
+  public void bind(java.sql.PreparedStatement statement, int parameterIndex, int rowIndex)throws java.sql.SQLException {
+    listReader.setPosition(rowIndex);
+    ArrayList<?> sourceArray = (ArrayList<?>) listReader.readObject();
+    Object array;
+    if (!isTextColumn) {
+      array = Array.newInstance(arrayElementClass, sourceArray.size());
+      Arrays.setAll((Object[]) array, sourceArray::get);
+    } else {
+      array = new String[sourceArray.size()];
+      Arrays.setAll((Object[]) array, idx -> sourceArray.get(idx) != null ? sourceArray.get(idx).toString() : null);
+    }
+    statement.setObject(parameterIndex, array);
+  }
+}
diff --git a/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/JdbcParameterBinderTest.java b/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/JdbcParameterBinderTest.java
@@ -33,6 +33,7 @@
 import java.util.function.BiConsumer;
 
 import org.apache.arrow.adapter.jdbc.binder.ColumnBinder;
+import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.BaseLargeVariableWidthVector;
@@ -67,11 +68,13 @@
 import org.apache.arrow.vector.VarBinaryVector;
 import org.apache.arrow.vector.VarCharVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.complex.ListVector;
 import org.apache.arrow.vector.types.DateUnit;
 import org.apache.arrow.vector.types.FloatingPointPrecision;
 import org.apache.arrow.vector.types.TimeUnit;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -385,6 +388,91 @@ void decimal256() throws SQLException {
         Arrays.asList(new BigDecimal("120.429"), new BigDecimal("-10590.123"), new BigDecimal("0.000")));
   }
 
+  @Test
+  void listOfDouble() throws SQLException {
+    TriConsumer<ListVector, Integer, Double[]> setValue = (listVector, index, values) -> {
+      org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
+      writer.setPosition(index);
+      writer.startList();
+      Arrays.stream(values).forEach(doubleValue -> writer.float8().writeFloat8(doubleValue));
+      writer.endList();
+      listVector.setLastSet(index);
+    };
+    List<Double[]> values = Arrays.asList(new Double[]{0.0, Math.PI}, new Double[]{1.1, -352346.2, 2355.6},
+                                          new Double[]{-1024.3}, new Double[]{});
+    testListType(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), setValue, ListVector::setNull, values);
+  }
+
+  @Test
+  void listOfInt64() throws SQLException {
+    TriConsumer<ListVector, Integer, Long[]> setValue = (listVector, index, values) -> {
+      org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
+      writer.setPosition(index);
+      writer.startList();
+      Arrays.stream(values).forEach(longValue -> writer.bigInt().writeBigInt(longValue));
+      writer.endList();
+      listVector.setLastSet(index);
+    };
+    List<Long[]> values = Arrays.asList(new Long[]{1L, 2L, 3L}, new Long[]{4L, 5L},
+            new Long[]{512L, 1024L, 2048L, 4096L}, new Long[]{});
+    testListType((ArrowType) new ArrowType.Int(64, true), setValue, ListVector::setNull, values);
+  }
+
+  @Test
+  void listOfInt32() throws SQLException {
+    TriConsumer<ListVector, Integer, Integer[]> setValue = (listVector, index, values) -> {
+      org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
+      writer.setPosition(index);
+      writer.startList();
+      Arrays.stream(values).forEach(integerValue -> writer.integer().writeInt(integerValue));
+      writer.endList();
+      listVector.setLastSet(index);
+    };
+    List<Integer[]> values = Arrays.asList(new Integer[]{1, 2, 3}, new Integer[]{4, 5},
+            new Integer[]{512, 1024, 2048, 4096}, new Integer[]{});
+    testListType((ArrowType) new ArrowType.Int(32, true), setValue, ListVector::setNull, values);
+  }
+
+  @Test
+  void listOfBoolean() throws SQLException {
+    TriConsumer<ListVector, Integer, Boolean[]> setValue = (listVector, index, values) -> {
+      org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
+      writer.setPosition(index);
+      writer.startList();
+      Arrays.stream(values).forEach(booleanValue -> writer.bit().writeBit(booleanValue ? 1 : 0));
+      writer.endList();
+      listVector.setLastSet(index);
+    };
+    List<Boolean[]> values = Arrays.asList(new Boolean[]{true, false},
+            new Boolean[]{false, false}, new Boolean[]{true, true, false, true}, new Boolean[]{});
+    testListType((ArrowType) new ArrowType.Bool(), setValue, ListVector::setNull, values);
+  }
+
+  @Test
+  void listOfString() throws SQLException {
+    TriConsumer<ListVector, Integer, String[]> setValue = (listVector, index, values) -> {
+      org.apache.arrow.vector.complex.impl.UnionListWriter writer = listVector.getWriter();
+      writer.setPosition(index);
+      writer.startList();
+      Arrays.stream(values).forEach(stringValue -> {
+        if (stringValue != null) {
+          byte[] stringValueBytes = stringValue.getBytes(StandardCharsets.UTF_8);
+          try (ArrowBuf stringBuffer = allocator.buffer(stringValueBytes.length)) {
+            stringBuffer.writeBytes(stringValueBytes);
+            writer.varChar().writeVarChar(0, stringValueBytes.length, stringBuffer);
+          }
+        } else {
+          writer.varChar().writeNull();
+        }
+      });
+      writer.endList();
+      listVector.setLastSet(index);
+    };
+    List<String[]> values = Arrays.asList(new String[]{"aaaa", "b1"},
+            new String[]{"c", null, "d"}, new String[]{"e", "f", "g", "h"}, new String[]{});
+    testListType((ArrowType) new ArrowType.Utf8(), setValue, ListVector::setNull, values);
+  }
+
   @FunctionalInterface
   interface TriConsumer<T, U, V> {
     void accept(T value1, U value2, V value3);
@@ -483,4 +571,105 @@ <T, V extends FieldVector> void testSimpleType(ArrowType arrowType, int jdbcType
       assertThat(binder.next()).isFalse();
     }
   }
+
+  <T, V extends FieldVector> void testListType(ArrowType arrowType, TriConsumer<V, Integer, T> setValue,
+                          BiConsumer<V, Integer> setNull, List<T> values) throws SQLException {
+    int jdbcType = Types.ARRAY;
+    Schema schema = new Schema(Collections.singletonList(new Field("field", FieldType.nullable(
+            new ArrowType.List()), Collections.singletonList(
+            new Field("element", FieldType.notNullable(arrowType), null)
+    ))));
+    try (final MockPreparedStatement statement = new MockPreparedStatement();
+         final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+      final JdbcParameterBinder binder =
+          JdbcParameterBinder.builder(statement, root).bindAll().build();
+      assertThat(binder.next()).isFalse();
+
+      @SuppressWarnings("unchecked")
+      final V vector = (V) root.getVector(0);
+      final ColumnBinder columnBinder = ColumnBinder.forVector(vector);
+      assertThat(columnBinder.getJdbcType()).isEqualTo(jdbcType);
+
+      setValue.accept(vector, 0, values.get(0));
+      setValue.accept(vector, 1, values.get(1));
+      setNull.accept(vector, 2);
+      root.setRowCount(3);
+
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(1));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isNull();
+      assertThat(statement.getParamType(1)).isEqualTo(jdbcType);
+      assertThat(binder.next()).isFalse();
+
+      binder.reset();
+
+      setNull.accept(vector, 0);
+      setValue.accept(vector, 1, values.get(3));
+      setValue.accept(vector, 2, values.get(0));
+      setValue.accept(vector, 3, values.get(2));
+      setValue.accept(vector, 4, values.get(1));
+      root.setRowCount(5);
+
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isNull();
+      assertThat(statement.getParamType(1)).isEqualTo(jdbcType);
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(3));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(2));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(1));
+      assertThat(binder.next()).isFalse();
+    }
+
+    // Non-nullable (since some types have a specialized binder)
+    schema = new Schema(Collections.singletonList(new Field("field", FieldType.notNullable(
+            new ArrowType.List()), Collections.singletonList(
+            new Field("element", FieldType.notNullable(arrowType), null)
+    ))));
+    try (final MockPreparedStatement statement = new MockPreparedStatement();
+         final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+      final JdbcParameterBinder binder =
+          JdbcParameterBinder.builder(statement, root).bindAll().build();
+      assertThat(binder.next()).isFalse();
+
+      @SuppressWarnings("unchecked")
+      final V vector = (V) root.getVector(0);
+      setValue.accept(vector, 0, values.get(0));
+      setValue.accept(vector, 1, values.get(1));
+      root.setRowCount(2);
+
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(1));
+      assertThat(binder.next()).isFalse();
+
+      binder.reset();
+
+      setValue.accept(vector, 0, values.get(0));
+      setValue.accept(vector, 1, values.get(2));
+      setValue.accept(vector, 2, values.get(0));
+      setValue.accept(vector, 3, values.get(2));
+      setValue.accept(vector, 4, values.get(1));
+      root.setRowCount(5);
+
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(2));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(0));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(2));
+      assertThat(binder.next()).isTrue();
+      assertThat(statement.getParamValue(1)).isEqualTo(values.get(1));
+      assertThat(binder.next()).isFalse();
+    }
+  }
 }