[CORE] Add reusable JNI input adapters for backend runtimes

cpp/core/compute/Runtime.h

@@ -36,6 +36,7 @@
 
 namespace gluten {
 
+struct JniInputIteratorContext;
 class ResultIterator;
 
 struct SparkTaskInfo {
@@ -190,6 +191,8 @@ class Runtime : public std::enable_shared_from_this<Runtime> {
     return dumper_.get();
   }
 
+  virtual std::unique_ptr<ColumnarBatchIterator> createJniInputIterator(const JniInputIteratorContext& context);
+
   ObjectHandle saveObject(std::shared_ptr<void> obj) {
     return objStore_->save(obj);
   }

cpp/core/jni/JniCommon.cc

@@ -67,6 +67,12 @@ gluten::Runtime* gluten::getRuntime(JNIEnv* env, jobject runtimeAware) {
   return ctx;
 }
 
+std::unique_ptr<gluten::ColumnarBatchIterator> gluten::Runtime::createJniInputIterator(
+    const JniInputIteratorContext& context) {
+  return std::make_unique<JniColumnarBatchIterator>(
+      context.env, context.jColumnarBatchIterator, this, context.iteratorIndex);
+}
+
 std::unique_ptr<gluten::JniColumnarBatchIterator>
 gluten::makeJniColumnarBatchIterator(JNIEnv* env, jobject jColumnarBatchItr, gluten::Runtime* runtime) {
   return std::make_unique<JniColumnarBatchIterator>(env, jColumnarBatchItr, runtime);

cpp/core/jni/JniCommon.h

@@ -151,6 +151,10 @@ static T* jniCastOrThrow(jlong handle) {
 }
 namespace gluten {
 
+class StreamReader;
+
+std::shared_ptr<StreamReader> makeShuffleStreamReader(JNIEnv* env, jobject jShuffleStreamReader);
+
 class JniCommonState {
  public:
   virtual ~JniCommonState() = default;
@@ -286,6 +290,12 @@ DEFINE_SAFE_GET_PRIMITIVE_ARRAY_FUNCTIONS(kLong, jlongArray, Long)
 DEFINE_SAFE_GET_PRIMITIVE_ARRAY_FUNCTIONS(kFloat, jfloatArray, Float)
 DEFINE_SAFE_GET_PRIMITIVE_ARRAY_FUNCTIONS(kDouble, jdoubleArray, Double)
 
+struct JniInputIteratorContext {
+  JNIEnv* env;
+  jobject jColumnarBatchIterator;
+  int32_t iteratorIndex;
+};
+
 class JniColumnarBatchIterator : public ColumnarBatchIterator {
  public:
   explicit JniColumnarBatchIterator(

cpp/core/jni/JniWrapper.cc

@@ -273,6 +273,14 @@ class ShuffleStreamReader : public StreamReader {
 
 } // namespace
 
+namespace gluten {
+
+std::shared_ptr<StreamReader> makeShuffleStreamReader(JNIEnv* env, jobject jShuffleStreamReader) {
+  return std::make_shared<::ShuffleStreamReader>(env, jShuffleStreamReader);
+}
+
+} // namespace gluten
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -566,7 +574,7 @@ Java_org_apache_gluten_vectorized_PlanEvaluatorJniWrapper_nativeCreateKernelWith
     inputIters.reserve(itersLen);
     for (int idx = 0; idx < itersLen; idx++) {
       jobject iter = env->GetObjectArrayElement(batchItrArray, idx);
-      auto arrayIter = std::make_unique<JniColumnarBatchIterator>(env, iter, ctx, idx);
+      auto arrayIter = ctx->createJniInputIterator({env, iter, idx});
       auto resultIter = std::make_shared<ResultIterator>(std::move(arrayIter));
       inputIters.push_back(std::move(resultIter));
     }

cpp/core/tests/CMakeLists.txt

@@ -17,3 +17,4 @@ add_test_case(round_robin_partitioner_test SOURCES RoundRobinPartitionerTest.cc)
 add_test_case(object_store_test SOURCES ObjectStoreTest.cc)
 add_test_case(memory_allocator_test SOURCES MemoryAllocatorTest.cc)
 add_test_case(ffor_codec_test SOURCES FForCodecTest.cc)
+add_test_case(jni_input_iterator_test SOURCES JniInputIteratorTest.cc)

cpp/core/tests/JniInputIteratorTest.cc

@@ -0,0 +1,65 @@
+/*
+ * 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.
+ */
+
+#include "jni/JniCommon.h"
+
+#include <gtest/gtest.h>
+
+using namespace gluten;
+
+namespace {
+
+class TestColumnarBatchIterator final : public ColumnarBatchIterator {
+ public:
+  std::shared_ptr<ColumnarBatch> next() override {
+    return nullptr;
+  }
+};
+
+class TestRuntime final : public Runtime {
+ public:
+  TestRuntime() : Runtime("jni-input-iterator-test", nullptr, nullptr, {}) {}
+
+  std::unique_ptr<ColumnarBatchIterator> createJniInputIterator(const JniInputIteratorContext& context) override {
+    env = context.env;
+    iterator = context.jColumnarBatchIterator;
+    iteratorIndex = context.iteratorIndex;
+    return std::make_unique<TestColumnarBatchIterator>();
+  }
+
+  JNIEnv* env{nullptr};
+  jobject iterator{nullptr};
+  int32_t iteratorIndex{-1};
+};
+
+TEST(JniInputIterator, DispatchesThroughRuntime) {
+  static int envMarker;
+  static int iteratorMarker;
+  auto* env = reinterpret_cast<JNIEnv*>(&envMarker);
+  auto iterator = reinterpret_cast<jobject>(&iteratorMarker);
+
+  TestRuntime runtime;
+  Runtime* baseRuntime = &runtime;
+  auto result = baseRuntime->createJniInputIterator({env, iterator, 7});
+
+  ASSERT_NE(result, nullptr);
+  EXPECT_EQ(runtime.env, env);
+  EXPECT_EQ(runtime.iterator, iterator);
+  EXPECT_EQ(runtime.iteratorIndex, 7);
+}
+
+} // namespace