feat: respect batchSize/workerThreads/blockingThreads configurations for native_iceberg_compat scan (apache#1587)

Author: wForget

common/src/main/java/org/apache/comet/parquet/Native.java

@@ -256,7 +256,10 @@ public static native long initRecordBatchReader(
       byte[] filter,
       byte[] requiredSchema,
       byte[] dataSchema,
-      String sessionTimezone);
+      String sessionTimezone,
+      int batchSize,
+      int workerThreads,
+      int blockingThreads);
 
   // arrow native version of read batch
   /**

common/src/main/java/org/apache/comet/parquet/NativeBatchReader.java

@@ -353,6 +353,20 @@ public void init() throws URISyntaxException, IOException {
       }
     }
 
+    int batchSize =
+        conf.getInt(
+            CometConf.COMET_BATCH_SIZE().key(),
+            (Integer) CometConf.COMET_BATCH_SIZE().defaultValue().get());
+    int workerThreads =
+        conf.getInt(
+            CometConf.COMET_WORKER_THREADS().key(),
+            (Integer) CometConf.COMET_WORKER_THREADS().defaultValue().get());
+    ;
+    int blockingThreads =
+        conf.getInt(
+            CometConf.COMET_BLOCKING_THREADS().key(),
+            (Integer) CometConf.COMET_BLOCKING_THREADS().defaultValue().get());
+    ;
     this.handle =
         Native.initRecordBatchReader(
             filePath,
@@ -362,7 +376,10 @@ public void init() throws URISyntaxException, IOException {
             nativeFilter,
             serializedRequestedArrowSchema,
             serializedDataArrowSchema,
-            timeZoneId);
+            timeZoneId,
+            batchSize,
+            workerThreads,
+            blockingThreads);
     isInitialized = true;
   }
 

native/core/src/execution/planner.rs

@@ -175,6 +175,11 @@ impl PhysicalPlanner {
         }
     }
 
+    /// Return session context of this planner.
+    pub fn session_ctx(&self) -> &Arc<SessionContext> {
+        &self.session_ctx
+    }
+
     /// get DataFusion PartitionedFiles from a Spark FilePartition
     fn get_partitioned_files(
         &self,

native/core/src/parquet/mod.rs

@@ -54,8 +54,9 @@ use crate::parquet::parquet_support::prepare_object_store;
 use arrow::array::{Array, RecordBatch};
 use arrow::buffer::{Buffer, MutableBuffer};
 use datafusion::datasource::listing::PartitionedFile;
-use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::execution::SendableRecordBatchStream;
 use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::{SessionConfig, SessionContext};
 use futures::{poll, StreamExt};
 use jni::objects::{JBooleanArray, JByteArray, JLongArray, JPrimitiveArray, JString, ReleaseMode};
 use jni::sys::jstring;
@@ -650,21 +651,29 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
     required_schema: jbyteArray,
     data_schema: jbyteArray,
     session_timezone: jstring,
+    batch_size: jint,
+    worker_threads: jint,
+    blocking_threads: jint,
 ) -> jlong {
     try_unwrap_or_throw(&e, |mut env| unsafe {
-        let task_ctx = TaskContext::default();
+        let session_config = SessionConfig::new().with_batch_size(batch_size as usize);
+        let planer =
+            PhysicalPlanner::new(Arc::new(SessionContext::new_with_config(session_config)));
+        let session_ctx = planer.session_ctx();
 
         let path: String = env
             .get_string(&JString::from_raw(file_path))
             .unwrap()
             .into();
 
         let runtime = tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(worker_threads as usize)
+            .max_blocking_threads(blocking_threads as usize)
             .enable_all()
             .build()?;
 
         let (object_store_url, object_store_path) =
-            prepare_object_store(task_ctx.runtime_env(), path.clone())?;
+            prepare_object_store(session_ctx.runtime_env(), path.clone())?;
 
         let required_schema_array = JByteArray::from_raw(required_schema);
         let required_schema_buffer = env.convert_byte_array(&required_schema_array)?;
@@ -674,8 +683,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
         let data_schema_buffer = env.convert_byte_array(&data_schema_array)?;
         let data_schema = Arc::new(deserialize_schema(data_schema_buffer.as_bytes())?);
 
-        let planer = PhysicalPlanner::default();
-
         let data_filters = if !filter.is_null() {
             let filter_array = JByteArray::from_raw(filter);
             let filter_buffer = env.convert_byte_array(&filter_array)?;
@@ -708,7 +715,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
         )?;
 
         let partition_index: usize = 0;
-        let batch_stream = Some(scan.execute(partition_index, Arc::new(task_ctx))?);
+        let batch_stream = Some(scan.execute(partition_index, session_ctx.task_ctx())?);
 
         let ctx = BatchContext {
             runtime,

spark/src/main/scala/org/apache/comet/parquet/CometParquetFileFormat.scala

@@ -232,6 +232,11 @@ object CometParquetFileFormat extends Logging with ShimSQLConf {
     hadoopConf.setBoolean(
       CometConf.COMET_EXCEPTION_ON_LEGACY_DATE_TIMESTAMP.key,
       CometConf.COMET_EXCEPTION_ON_LEGACY_DATE_TIMESTAMP.get())
+    hadoopConf.setInt(CometConf.COMET_BATCH_SIZE.key, CometConf.COMET_BATCH_SIZE.get())
+    hadoopConf.setInt(CometConf.COMET_WORKER_THREADS.key, CometConf.COMET_WORKER_THREADS.get())
+    hadoopConf.setInt(
+      CometConf.COMET_BLOCKING_THREADS.key,
+      CometConf.COMET_BLOCKING_THREADS.get())
   }
 
   def getDatetimeRebaseSpec(