feat(java/driver/jni): implement executePartitioned

Closes #4242.

Author: lidavidm

.github/workflows/java.yml

@@ -233,7 +233,7 @@ jobs:
       - name: Start Dependencies
         if: matrix.os == 'Linux' && matrix.arch == 'amd64'
         run: |
-          docker compose up --detach --wait mssql-test postgres-test
+          docker compose up --detach --wait flightsql-sqlite-test mssql-test postgres-test
           cat .env | grep -v -e '^#' | grep -e '^ADBC_' | awk NF | sed 's/"//g' | tee -a $GITHUB_ENV
 
       - name: Download thirdparty driver

go/adbc/driver/flightsql/flightsql_connection.go

@@ -1145,15 +1145,15 @@ func (c *connectionImpl) ReadPartition(ctx context.Context, serializedPartition
 	var info flight.FlightInfo
 	if err := proto.Unmarshal(serializedPartition, &info); err != nil {
 		return nil, adbc.Error{
-			Msg:  err.Error(),
+			Msg:  fmt.Sprintf("[flightsql] could not unmarshal partition as FlightInfo: %v", err),
 			Code: adbc.StatusInvalidArgument,
 		}
 	}
 
 	// The driver only ever returns one endpoint.
 	if len(info.Endpoint) != 1 {
 		return nil, adbc.Error{
-			Msg:  fmt.Sprintf("Invalid partition: expected 1 endpoint, got %d", len(info.Endpoint)),
+			Msg:  fmt.Sprintf("[flightsql] invalid partition: expected 1 endpoint, got %d", len(info.Endpoint)),
 			Code: adbc.StatusInvalidArgument,
 		}
 	}

go/adbc/driver/flightsql/flightsql_statement.go

@@ -811,7 +811,7 @@ func (s *statement) ExecutePartitions(ctx context.Context) (*arrow.Schema, adbc.
 		data, err := proto.Marshal(partition)
 		if err != nil {
 			return sc, out, -1, adbc.Error{
-				Msg:  err.Error(),
+				Msg:  fmt.Sprintf("[flightsql] could not marshal partition as FlightInfo: %v", err),
 				Code: adbc.StatusInternal,
 			}
 		}

java/core/src/main/java/org/apache/arrow/adbc/core/AdbcStatement.java

@@ -24,6 +24,7 @@
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.ipc.ArrowReader;
 import org.apache.arrow.vector.types.pojo.Schema;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * A container for all state needed to execute a database query, such as the query itself,
@@ -231,19 +232,21 @@ public String toString() {
 
   /** The partitions of a result set. */
   class PartitionResult {
-    private final Schema schema;
+    private final @Nullable Schema schema;
     private final long affectedRows;
     private final List<PartitionDescriptor> partitionDescriptors;
 
     public PartitionResult(
-        Schema schema, long affectedRows, List<PartitionDescriptor> partitionDescriptors) {
+        @Nullable Schema schema,
+        long affectedRows,
+        List<PartitionDescriptor> partitionDescriptors) {
       this.schema = schema;
       this.affectedRows = affectedRows;
       this.partitionDescriptors = partitionDescriptors;
     }
 
     /** Get the schema of the eventual result set. */
-    public Schema getSchema() {
+    public @Nullable Schema getSchema() {
       return schema;
     }
 

java/driver/jni-validation/src/test/java/org/apache/arrow/adbc/driver/jni/FlightSqlIntegrationTest.java

@@ -0,0 +1,98 @@
+/*
+ * 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.adbc.driver.jni;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.arrow.adbc.core.AdbcConnection;
+import org.apache.arrow.adbc.core.AdbcDatabase;
+import org.apache.arrow.adbc.core.AdbcDriver;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class FlightSqlIntegrationTest {
+  public static final String URI_ENV = "ADBC_SQLITE_FLIGHTSQL_URI";
+  static String URI = System.getenv(URI_ENV);
+
+  BufferAllocator allocator;
+  JniDriver driver;
+  AdbcDatabase db;
+  AdbcConnection conn;
+
+  @BeforeAll
+  static void beforeAll() {
+    Assumptions.assumeFalse(
+        URI == null || URI.isEmpty(),
+        String.format("Must set %s to run Flight SQL integration tests", URI_ENV));
+  }
+
+  @BeforeEach
+  void beforeEach() throws Exception {
+    allocator = new RootAllocator();
+    driver = new JniDriver(allocator);
+    Map<String, Object> parameters = new HashMap<>();
+    JniDriver.PARAM_DRIVER.set(parameters, "adbc_driver_flightsql");
+    AdbcDriver.PARAM_URI.set(parameters, URI);
+    db = driver.open(parameters);
+    conn = db.connect();
+  }
+
+  @AfterEach
+  void afterEach() throws Exception {
+    conn.close();
+    db.close();
+    allocator.close();
+  }
+
+  @Test
+  void simple() throws Exception {
+    try (var stmt = conn.createStatement()) {
+      stmt.setSqlQuery("SELECT 1 + 1 AS sum");
+      try (var reader = stmt.executeQuery()) {
+        assertThat(reader.getReader().loadNextBatch()).isTrue();
+        assertThat(reader.getReader().getVectorSchemaRoot().getVector("sum").getObject(0))
+            .isEqualTo(2L);
+      }
+    }
+  }
+
+  @Test
+  void partitioned() throws Exception {
+    try (var stmt = conn.createStatement()) {
+      stmt.setSqlQuery("SELECT 1 + 1 AS sum");
+      var partitions = stmt.executePartitioned();
+      assertThat(partitions.getPartitionDescriptors().size()).isEqualTo(1);
+      assertThat(partitions.getAffectedRows()).isEqualTo(-1);
+      // The test server doesn't give a schema.
+      assertThat(partitions.getSchema()).isNull();
+
+      try (var reader =
+          conn.readPartition(partitions.getPartitionDescriptors().get(0).getDescriptor())) {
+        assertThat(reader.loadNextBatch()).isTrue();
+        assertThat(reader.getVectorSchemaRoot().getVector("sum").getObject(0)).isEqualTo(2L);
+      }
+    }
+  }
+}

java/driver/jni/src/main/cpp/jni_wrapper.cc

@@ -34,6 +34,19 @@
 
 namespace {
 
+void ThrowJavaException(JNIEnv* env, const std::string& klass,
+                        const std::string& message) {
+  jclass exception_klass = env->FindClass(klass.c_str());
+  assert(exception_klass != nullptr);
+  jmethodID exception_ctor =
+      env->GetMethodID(exception_klass, "<init>", "(Ljava/lang/String)V");
+  assert(exception_ctor != nullptr);
+  jstring message_jni = env->NewStringUTF(message.c_str());
+  auto exc = static_cast<jthrowable>(
+      env->NewObject(exception_klass, exception_ctor, message_jni));
+  env->Throw(exc);
+}
+
 /// Internal exception.  Meant to be used with RaiseAdbcException and
 ///   CHECK_ADBC_ERROR.
 struct AdbcException {
@@ -112,19 +125,26 @@ void RaiseAdbcException(AdbcStatusCode code, const AdbcError& error) {
   } while (0)
 
 /// Require that a Java class exists or error.
-jclass RequireImplClass(JNIEnv* env, std::string_view name) {
-  static std::string kPrefix = "org/apache/arrow/adbc/driver/jni/impl/";
-  std::string full_name = kPrefix + std::string(name);
-  jclass klass = env->FindClass(full_name.c_str());
+jclass RequireClass(JNIEnv* env, const std::string& name) {
+  jclass klass = env->FindClass(name.c_str());
   if (klass == nullptr) {
+    std::string message = "[JNI] Could not find class ";
+    message += name;
     throw AdbcException{
         .code = ADBC_STATUS_INTERNAL,
-        .message = "[JNI] Could not find class " + full_name,
+        .message = std::move(message),
     };
   }
   return klass;
 }
 
+/// Require that a Java class exists or error.
+jclass RequireImplClass(JNIEnv* env, std::string_view name) {
+  static std::string kPrefix = "org/apache/arrow/adbc/driver/jni/impl/";
+  std::string full_name = kPrefix + std::string(name);
+  return RequireClass(env, full_name);
+}
+
 /// Require that a Java method exists or error.
 jmethodID RequireMethod(JNIEnv* env, jclass klass, std::string_view name,
                         std::string_view signature) {
@@ -381,6 +401,60 @@ Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_statementGetParameterSchem
   return nullptr;
 }
 
+JNIEXPORT jobject JNICALL
+Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_statementExecutePartitions(
+    JNIEnv* env, [[maybe_unused]] jclass self, jlong handle) {
+  struct AdbcError error = ADBC_ERROR_INIT;
+  auto* ptr = reinterpret_cast<struct AdbcStatement*>(static_cast<uintptr_t>(handle));
+  struct ArrowSchema schema = {};
+  struct AdbcPartitions partitions = {};
+  int64_t rows_affected = 0;
+  jobject result = nullptr;
+
+  try {
+    jclass native_result_class = RequireImplClass(env, "NativePartitionResult");
+    jmethodID native_result_ctor =
+        RequireMethod(env, native_result_class, "<init>", "(JJ)V");
+    jmethodID native_result_add_partition =
+        RequireMethod(env, native_result_class, "addPartition", "([B)V");
+
+    CHECK_ADBC_ERROR(
+        AdbcStatementExecutePartitions(ptr, &schema, &partitions, &rows_affected, &error),
+        error);
+
+    result = env->NewObject(native_result_class, native_result_ctor, rows_affected,
+                            static_cast<jlong>(reinterpret_cast<uintptr_t>(&schema)));
+    if (env->ExceptionCheck()) goto cleanupall;
+
+    for (size_t i = 0; i < partitions.num_partitions; i++) {
+      size_t length = partitions.partition_lengths[i];
+      jbyteArray partition = env->NewByteArray(static_cast<jsize>(length));
+      env->SetByteArrayRegion(partition, 0, static_cast<jsize>(length),
+                              reinterpret_cast<const jbyte*>(partitions.partitions[i]));
+      if (env->ExceptionCheck()) goto cleanupall;
+      env->CallObjectMethod(result, native_result_add_partition, partition);
+      if (env->ExceptionCheck()) goto cleanupall;
+    }
+  } catch (const AdbcException& e) {
+    e.ThrowJavaException(env);
+  }
+
+  // We can't release schema, but we copied out the partitions
+  if (partitions.release != nullptr) {
+    partitions.release(&partitions);
+  }
+  return result;
+
+cleanupall:
+  if (schema.release != nullptr) {
+    schema.release(&schema);
+  }
+  if (partitions.release != nullptr) {
+    partitions.release(&partitions);
+  }
+  return nullptr;
+}
+
 JNIEXPORT jobject JNICALL
 Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_statementExecuteQuery(
     JNIEnv* env, [[maybe_unused]] jclass self, jlong handle) {
@@ -979,6 +1053,95 @@ Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_connectionRollback(
   }
 }
 
+JNIEXPORT jobject JNICALL
+Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_connectionReadPartition(
+    JNIEnv* env, [[maybe_unused]] jclass self, jlong handle, jobject partition) {
+  struct AdbcError error = ADBC_ERROR_INIT;
+  auto* conn = reinterpret_cast<struct AdbcConnection*>(static_cast<uintptr_t>(handle));
+  struct ArrowArrayStream out = {};
+  size_t serialized_length = 0;
+  const uint8_t* serialized_partition = nullptr;
+  std::vector<uint8_t> allocated_partition;
+
+  try {
+    jclass bb_class = RequireClass(env, "java/nio/ByteBuffer");
+    jmethodID bb_remaining = RequireMethod(env, bb_class, "remaining", "()I");
+
+    if (!env->IsInstanceOf(partition, bb_class)) {
+      ThrowJavaException(env, "java/lang/IllegalArgumentException",
+                         "Partition must be a ByteBuffer");
+      return nullptr;
+    }
+    jint remaining = env->CallIntMethod(partition, bb_remaining);
+    if (remaining < 0) {
+      ThrowJavaException(env, "java/lang/IllegalArgumentException",
+                         "ByteBuffer remaining() must be non-negative");
+      return nullptr;
+    }
+    serialized_length = static_cast<size_t>(remaining);
+
+    // fast path (if direct buffer)
+    void* buf = env->GetDirectBufferAddress(partition);
+    if (buf) {
+      serialized_partition = static_cast<const uint8_t*>(buf);
+    }
+
+    // middle path (backing array)
+    if (!serialized_partition) {
+      jmethodID bb_has_array = RequireMethod(env, bb_class, "hasArray", "()Z");
+      jmethodID bb_array = RequireMethod(env, bb_class, "array", "()[B");
+      jmethodID bb_array_offset = RequireMethod(env, bb_class, "arrayOffset", "()I");
+      jboolean has_array = env->CallBooleanMethod(partition, bb_has_array);
+      if (env->ExceptionCheck()) return nullptr;
+      if (has_array) {
+        jint array_offset = env->CallIntMethod(partition, bb_array_offset);
+        if (env->ExceptionCheck()) return nullptr;
+
+        auto array =
+            reinterpret_cast<jbyteArray>(env->CallObjectMethod(partition, bb_array));
+        if (env->ExceptionCheck()) return nullptr;
+
+        assert(serialized_length <= static_cast<size_t>(env->GetArrayLength(array)));
+        allocated_partition.resize(serialized_length);
+        env->GetByteArrayRegion(array, array_offset,
+                                static_cast<jsize>(serialized_length),
+                                reinterpret_cast<jbyte*>(allocated_partition.data()));
+        serialized_partition = allocated_partition.data();
+      }
+    }
+
+    // slow path (copy)
+    if (!serialized_partition) {
+      jmethodID bb_get = RequireMethod(env, bb_class, "get", "([B)Ljava/nio/ByteBuffer;");
+      jbyteArray temp = env->NewByteArray(static_cast<jsize>(serialized_length));
+      if (!temp) {
+        ThrowJavaException(env, "java/lang/OutOfMemoryError",
+                           "Failed to allocate byte array for partition");
+        return nullptr;
+      }
+
+      env->CallVoidMethod(partition, bb_get, temp);
+      if (env->ExceptionCheck()) return nullptr;
+
+      allocated_partition.resize(serialized_length);
+      env->GetByteArrayRegion(temp, 0, static_cast<jsize>(serialized_length),
+                              reinterpret_cast<jbyte*>(allocated_partition.data()));
+      serialized_partition = allocated_partition.data();
+    }
+
+    assert(serialized_partition != nullptr);
+
+    CHECK_ADBC_ERROR(AdbcConnectionReadPartition(conn, serialized_partition,
+                                                 serialized_length, &out, &error),
+                     error);
+
+    return MakeNativeQueryResult(env, -1, &out);
+  } catch (const AdbcException& e) {
+    e.ThrowJavaException(env);
+  }
+  return nullptr;
+}
+
 JNIEXPORT jbyteArray JNICALL
 Java_org_apache_arrow_adbc_driver_jni_impl_NativeAdbc_databaseGetOptionBytes(
     JNIEnv* env, [[maybe_unused]] jclass self, jlong handle, jstring key) {

java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java

@@ -17,6 +17,7 @@
 
 package org.apache.arrow.adbc.driver.jni;
 
+import java.nio.ByteBuffer;
 import org.apache.arrow.adbc.core.AdbcConnection;
 import org.apache.arrow.adbc.core.AdbcException;
 import org.apache.arrow.adbc.core.AdbcStatement;
@@ -254,6 +255,11 @@ public void setCurrentDbSchema(String dbSchema) throws AdbcException {
     setOption(JniDriver.CURRENT_DB_SCHEMA, dbSchema);
   }
 
+  @Override
+  public ArrowReader readPartition(ByteBuffer descriptor) throws AdbcException {
+    return JniLoader.INSTANCE.connectionReadPartition(handle, descriptor).importStream(allocator);
+  }
+
   @Override
   public void close() {
     handle.close();