GH-34223: [Java] Java Substrait Consumer JNI call to ACERO C++

ci/scripts/java_jni_macos_build.sh

@@ -71,6 +71,7 @@ cmake \
   -DARROW_BUILD_TESTS=${ARROW_BUILD_TESTS} \
   -DARROW_CSV=${ARROW_DATASET} \
   -DARROW_DATASET=${ARROW_DATASET} \
+  -DARROW_SUBSTRAIT=${ARROW_DATASET} \
   -DARROW_DEPENDENCY_USE_SHARED=OFF \
   -DARROW_GANDIVA=${ARROW_GANDIVA} \
   -DARROW_GANDIVA_STATIC_LIBSTDCPP=ON \

ci/scripts/java_jni_manylinux_build.sh

@@ -70,6 +70,7 @@ cmake \
   -DARROW_BUILD_TESTS=ON \
   -DARROW_CSV=${ARROW_DATASET} \
   -DARROW_DATASET=${ARROW_DATASET} \
+  -DARROW_SUBSTRAIT=${ARROW_DATASET} \
   -DARROW_DEPENDENCY_SOURCE="VCPKG" \
   -DARROW_DEPENDENCY_USE_SHARED=OFF \
   -DARROW_GANDIVA_PC_CXX_FLAGS=${GANDIVA_CXX_FLAGS} \

ci/scripts/java_jni_windows_build.sh

@@ -58,6 +58,7 @@ cmake \
   -DARROW_BUILD_TESTS=ON \
   -DARROW_CSV=${ARROW_DATASET} \
   -DARROW_DATASET=${ARROW_DATASET} \
+  -DARROW_SUBSTRAIT=${ARROW_DATASET} \
   -DARROW_DEPENDENCY_USE_SHARED=OFF \
   -DARROW_ORC=${ARROW_ORC} \
   -DARROW_PARQUET=${ARROW_PARQUET} \

docs/source/java/index.rst

@@ -37,6 +37,7 @@ on the Arrow format and other language bindings see the :doc:`parent documentati
    flight_sql
    flight_sql_jdbc_driver
    dataset
+   substrait
    cdata
    jdbc
    Reference (javadoc) <reference/index>

docs/source/java/substrait.rst

@@ -0,0 +1,120 @@
+.. 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.
+
+=========
+Substrait
+=========
+
+Java Substrait offer capabilities to ``Query`` data received as a Susbtrait
+Plan (`plain or binary format`).
+
+During this process, Substrait plans are read, executed, and ArrowReaders are
+returned for reading Schema and ArrowRecordBatches. For Substrait plan that contains
+``Local Files`` the URI per table are defined in the Substrait plan, different
+than ``Named Tables`` where is needed to define a mapping name of tables
+and theirs ArrowReader representation.
+
+.. contents::
+
+Getting Started
+===============
+
+Java Substrait API uses Acero C++ Substrait API capabilities thru JNI wrappers.
+
+.. seealso:: :doc:`../cpp/streaming_execution` for more information on Acero.
+
+Substrait Consumer
+==================
+
+Substrait Plan offer two ways to define URI for Query data:
+
+- Local Files: A fixed URI value on the plan
+- Named Table: An external configuration to define URI value
+
+Local Files:
+
+.. code-block:: json
+
+    "local_files": {
+      "items": [
+        {
+          "uri_file": "file:///tmp/opt/lineitem.parquet",
+          "parquet": {}
+        }
+      ]
+    }
+
+Named Table:
+
+.. code-block:: json
+
+    "namedTable": {
+        "names": ["LINEITEM"]
+    }
+
+Here is an example of a Java program that queries a Parquet file using Java Substrait:
+
+.. code-block:: Java
+
+    // Query: SELECT * from nation
+    String uri = "file:///data/tpch_parquet/nation.parquet";
+    ScanOptions options = new ScanOptions(/*batchSize*/ 32768);
+    try (
+        DatasetFactory datasetFactory = new FileSystemDatasetFactory(rootAllocator(), NativeMemoryPool.getDefault(),
+            FileFormat.PARQUET, uri);
+        Dataset dataset = datasetFactory.finish();
+        Scanner scanner = dataset.newScan(options);
+        ArrowReader reader = scanner.scanBatches()
+    ) {
+      // map table to reader
+      Map<String, ArrowReader> mapTableToArrowReader = new HashMap<>();
+      mapTableToArrowReader.put("NATION", reader);
+      // get binary plan
+      String sql = "SELECT * from nation";
+      String nation = "CREATE TABLE NATION (N_NATIONKEY BIGINT NOT NULL, N_NAME CHAR(25), " +
+          "N_REGIONKEY BIGINT NOT NULL, N_COMMENT VARCHAR(152))";
+      Plan plan = getPlan(sql, ImmutableList.of(nation));
+      ByteBuffer substraitPlan = ByteBuffer.allocateDirect(plan.toByteArray().length);
+      substraitPlan.put(plan.toByteArray());
+      // run query
+      try (ArrowReader arrowReader = new SubstraitAceroConsumer(rootAllocator()).runQuery(
+          substraitPlan,
+          mapTableToArrowReader
+      )) {
+        while (arrowReader.loadNextBatch()) {
+          assertEquals(arrowReader.getVectorSchemaRoot().getRowCount(), 25);
+          assertTrue(arrowReader.getVectorSchemaRoot().contentToTSVString().contains("MOROCCO"));
+        }
+      }
+    } catch (Exception e) {
+      e.printStackTrace();
+    }
+
+.. code-block:: text
+
+    // Results example:
+    FieldPath(0)	FieldPath(1)	FieldPath(2)	FieldPath(3)
+    0	ALGERIA	0	 haggle. carefully final deposits detect slyly agai
+    1	ARGENTINA	1	al foxes promise slyly according to the regular accounts. bold requests alon
+
+Substrait Producer
+==================
+
+The following options are available for producing Substrait Plans: Acero,
+Isthmus, Ibis, DuckDB, others.
+
+You can generate Substrait plans and then send them to Java Substrait for consumption.

java/dataset/CMakeLists.txt

@@ -16,6 +16,7 @@
 # under the License.
 
 find_package(ArrowDataset REQUIRED)
+find_package(ArrowSubstrait REQUIRED)
 
 include_directories(${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR}
                     ${JNI_INCLUDE_DIRS} ${JNI_HEADERS_DIR})
@@ -26,14 +27,18 @@ add_jar(arrow_java_jni_dataset_jar
         src/main/java/org/apache/arrow/dataset/file/JniWrapper.java
         src/main/java/org/apache/arrow/dataset/jni/NativeMemoryPool.java
         src/main/java/org/apache/arrow/dataset/jni/ReservationListener.java
+        src/main/java/org/apache/arrow/dataset/substrait/JniWrapper.java
         GENERATE_NATIVE_HEADERS
         arrow_java_jni_dataset_headers)
 
 add_library(arrow_java_jni_dataset SHARED src/main/cpp/jni_wrapper.cc
                                           src/main/cpp/jni_util.cc)
 set_property(TARGET arrow_java_jni_dataset PROPERTY OUTPUT_NAME "arrow_dataset_jni")
-target_link_libraries(arrow_java_jni_dataset arrow_java_jni_dataset_headers jni
-                      ArrowDataset::arrow_dataset_static)
+target_link_libraries(arrow_java_jni_dataset
+                      arrow_java_jni_dataset_headers
+                      jni
+                      ArrowDataset::arrow_dataset_static
+                      ArrowSubstrait::arrow_substrait_static)
 
 if(BUILD_TESTING)
   add_executable(arrow-java-jni-dataset-test src/main/cpp/jni_util_test.cc

java/dataset/src/main/cpp/jni_wrapper.cc

@@ -16,6 +16,7 @@
 // under the License.
 
 #include <mutex>
+#include <unordered_map>
 
 #include "arrow/array.h"
 #include "arrow/array/concatenate.h"
@@ -24,12 +25,14 @@
 #include "arrow/dataset/api.h"
 #include "arrow/dataset/file_base.h"
 #include "arrow/filesystem/localfs.h"
+#include "arrow/engine/substrait/util.h"
 #include "arrow/ipc/api.h"
 #include "arrow/util/iterator.h"
 #include "jni_util.h"
 #include "org_apache_arrow_dataset_file_JniWrapper.h"
 #include "org_apache_arrow_dataset_jni_JniWrapper.h"
 #include "org_apache_arrow_dataset_jni_NativeMemoryPool.h"
+#include "org_apache_arrow_dataset_substrait_JniWrapper.h"
 
 namespace {
 
@@ -261,6 +264,26 @@ void JNI_OnUnload(JavaVM* vm, void* reserved) {
   default_memory_pool_id = -1L;
 }
 
+/// Iterate over an object array of Tables Name on position `i` and theirs
+/// Memory Address representation on `i+1` position linearly.
+/// Return a mapping of the Table Name to Query as a Key and a RecordBatchReader
+/// as a Value.
+std::unordered_map<std::string, std::shared_ptr<arrow::Table>> ToMapTableToArrowReader(JNIEnv* env, jobjectArray& str_array) {
+  std::unordered_map<std::string, std::shared_ptr<arrow::Table>> map_table_to_record_batch_reader;
+  int length = env->GetArrayLength(str_array);
+  std::shared_ptr<arrow::Table> output_table;
+  for (int pos = 0; pos < length; pos++) {
+    auto j_string_key = reinterpret_cast<jstring>(env->GetObjectArrayElement(str_array, pos));
+    pos++;
+    auto j_string_value = reinterpret_cast<jstring>(env->GetObjectArrayElement(str_array, pos));
+    auto* arrow_stream_in = reinterpret_cast<ArrowArrayStream*>(std::stol(JStringToCString(env, j_string_value)));
+    std::shared_ptr<arrow::RecordBatchReader> readerIn = JniGetOrThrow(arrow::ImportRecordBatchReader(arrow_stream_in));
+    output_table = JniGetOrThrow(readerIn->ToTable());
+    map_table_to_record_batch_reader[JStringToCString(env, j_string_key)] = output_table;
+  }
+  return map_table_to_record_batch_reader;
+}
+
 /*
  * Class:     org_apache_arrow_dataset_jni_NativeMemoryPool
  * Method:    getDefaultMemoryPool
@@ -578,3 +601,102 @@ Java_org_apache_arrow_dataset_file_JniWrapper_writeFromScannerToFile(
   JniAssertOkOrThrow(arrow::dataset::FileSystemDataset::Write(options, scanner));
   JNI_METHOD_END()
 }
+
+/*
+ * Class:     org_apache_arrow_dataset_substrait_JniWrapper
+ * Method:    executeSerializedPlanLocalFiles
+ * Signature: (Ljava/lang/String;J)V
+ */
+JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_substrait_JniWrapper_executeSerializedPlanLocalFiles__Ljava_lang_String_2J (
+    JNIEnv* env, jobject, jstring plan, jlong c_arrow_array_stream_address_out) {
+  JNI_METHOD_START
+  auto* arrow_stream = reinterpret_cast<ArrowArrayStream*>(c_arrow_array_stream_address_out);
+  std::shared_ptr<arrow::Buffer> buffer = JniGetOrThrow(arrow::engine::SerializeJsonPlan(JStringToCString(env, plan)));
+  std::shared_ptr<arrow::RecordBatchReader> reader = JniGetOrThrow(arrow::engine::ExecuteSerializedPlan(*buffer));
+  JniAssertOkOrThrow(arrow::ExportRecordBatchReader(reader, arrow_stream));
+  JNI_METHOD_END()
+}
+
+/*
+ * Class:     org_apache_arrow_dataset_substrait_JniWrapper
+ * Method:    executeSerializedPlanLocalFiles
+ * Signature: (Ljava/nio/ByteBuffer;J)V
+ */
+JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_substrait_JniWrapper_executeSerializedPlanLocalFiles__Ljava_nio_ByteBuffer_2J (
+    JNIEnv* env, jobject, jobject plan, jlong c_arrow_array_stream_address_out) {
+  JNI_METHOD_START
+  auto* arrow_stream = reinterpret_cast<ArrowArrayStream*>(c_arrow_array_stream_address_out);
+  // mapping arrow::Buffer
+  auto *buff = reinterpret_cast<jbyte*>(env->GetDirectBufferAddress(plan));
+  int length = env->GetDirectBufferCapacity(plan);
+  std::shared_ptr<arrow::Buffer> buffer = JniGetOrThrow(arrow::AllocateBuffer(length));
+  std::memcpy(buffer->mutable_data(), buff, length);
+  // execute plan
+  std::shared_ptr<arrow::RecordBatchReader> reader = JniGetOrThrow(arrow::engine::ExecuteSerializedPlan(*buffer));
+  JniAssertOkOrThrow(arrow::ExportRecordBatchReader(reader, arrow_stream));
+  JNI_METHOD_END()
+}
+
+/*
+ * Class:     org_apache_arrow_dataset_substrait_JniWrapper
+ * Method:    executeSerializedPlanNamedTables
+ * Signature: (Ljava/lang/String;[Ljava/lang/String;J)V
+ */
+JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_substrait_JniWrapper_executeSerializedPlanNamedTables__Ljava_lang_String_2_3Ljava_lang_String_2J (
+    JNIEnv* env, jobject, jstring plan, jobjectArray table_to_memory_address_input, jlong memory_address_output) {
+  JNI_METHOD_START
+  // get mapping of table name to memory address
+  std::unordered_map<std::string, std::shared_ptr<arrow::Table>> map_table_to_reader = ToMapTableToArrowReader(env, table_to_memory_address_input);
+  // create table provider
+  arrow::engine::NamedTableProvider table_provider = [&map_table_to_reader](const std::vector<std::string>& names, const arrow::Schema&) {
+    std::shared_ptr<arrow::Table> output_table;
+    for (const auto& name : names) {
+      output_table = map_table_to_reader[name];
+    }
+    std::shared_ptr<arrow::compute::ExecNodeOptions> options =
+      std::make_shared<arrow::compute::TableSourceNodeOptions>(std::move(output_table));
+    return arrow::compute::Declaration("table_source", {}, options, "java_source");
+  };
+  arrow::engine::ConversionOptions conversion_options;
+  conversion_options.named_table_provider = std::move(table_provider);
+  // execute plan
+  std::shared_ptr<arrow::Buffer> buffer = JniGetOrThrow(arrow::engine::SerializeJsonPlan(JStringToCString(env, plan)));
+  std::shared_ptr<arrow::RecordBatchReader> readerOut = JniGetOrThrow(arrow::engine::ExecuteSerializedPlan(*buffer, nullptr, nullptr, conversion_options));
+  auto* arrow_stream_out = reinterpret_cast<ArrowArrayStream*>(memory_address_output);
+  JniAssertOkOrThrow(arrow::ExportRecordBatchReader(readerOut, arrow_stream_out));
+  JNI_METHOD_END()
+}
+
+/*
+ * Class:     org_apache_arrow_dataset_substrait_JniWrapper
+ * Method:    executeSerializedPlanNamedTables
+ * Signature: (Ljava/nio/ByteBuffer;[Ljava/lang/String;J)V
+ */
+JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_substrait_JniWrapper_executeSerializedPlanNamedTables__Ljava_nio_ByteBuffer_2_3Ljava_lang_String_2J (
+    JNIEnv* env, jobject, jobject plan, jobjectArray table_to_memory_address_input, jlong memory_address_output) {
+  JNI_METHOD_START
+  // get mapping of table name to memory address
+  std::unordered_map<std::string, std::shared_ptr<arrow::Table>> map_table_to_reader = ToMapTableToArrowReader(env, table_to_memory_address_input);
+  // create table provider
+  arrow::engine::NamedTableProvider table_provider = [&map_table_to_reader](const std::vector<std::string>& names, const arrow::Schema&) {
+    std::shared_ptr<arrow::Table> output_table;
+    for (const auto& name : names) {
+      output_table = map_table_to_reader[name];
+    }
+    std::shared_ptr<arrow::compute::ExecNodeOptions> options =
+      std::make_shared<arrow::compute::TableSourceNodeOptions>(std::move(output_table));
+    return arrow::compute::Declaration("table_source", {}, options, "java_source");
+  };
+  arrow::engine::ConversionOptions conversion_options;
+  conversion_options.named_table_provider = std::move(table_provider);
+  // mapping arrow::Buffer
+  auto *buff = reinterpret_cast<jbyte*>(env->GetDirectBufferAddress(plan));
+  int length = env->GetDirectBufferCapacity(plan);
+  std::shared_ptr<arrow::Buffer> buffer = JniGetOrThrow(arrow::AllocateBuffer(length));
+  std::memcpy(buffer->mutable_data(), buff, length);
+  // execute plan
+  std::shared_ptr<arrow::RecordBatchReader> readerOut = JniGetOrThrow(arrow::engine::ExecuteSerializedPlan(*buffer, nullptr, nullptr, conversion_options));
+  auto* arrow_stream_out = reinterpret_cast<ArrowArrayStream*>(memory_address_output);
+  JniAssertOkOrThrow(arrow::ExportRecordBatchReader(readerOut, arrow_stream_out));
+  JNI_METHOD_END()
+}