IMPALA-14523: JDBC Scan Parallelism with Shared Cursor Fetching

Impala’s current JDBC DataSourceScan creates a separate JDBC
connection per scanner thread, but since JDBC cursors are
inherently sequential, this limits effective parallelism and
introduces unnecessary connection and JNI overhead. As a result,
even with MT_DOP enabled, scans do not scale efficiently and can
suffer from poor throughput for large datasets.

This patch introduces a shared JDBC connection model per fragment,
allowing multiple scanner threads to cooperatively fetch from a
single cursor. A new SharedJdbcConnection abstraction is added and
owned by DataSourceScanPlanNode, ensuring the connection is opened
and closed exactly once. Multiple DataSourceScanNode instances
invoke FetchBatch() concurrently, where cursor movement is
serialized while row materialization happens in parallel, and EOS
is tracked via an atomic flag to avoid redundant fetches.

Planner and scheduler changes ensure all JDBC scan ranges are
pinned to a single executor using a new is_data_source_scan flag,
and one virtual scan range is generated per MT_DOP instance. On the
Java side, batch fetching with minimal locking is introduced, and
JdbcDataSource is updated for safe concurrent access and improved
materialization, reducing overhead and improving scan throughput.

For instance, for query:
WITH recent_orders AS (
  SELECT o.o_orderkey, o.o_custkey, o.o_totalprice, o.o_orderdate
  FROM tpch_jdbc.orders o WHERE o.o_orderdate >= '1995-01-01'
)
SELECT c.c_custkey, c.c_name, c.c_mktsegment,
       COUNT(ro.o_orderkey) AS num_orders,
       SUM(ro.o_totalprice) AS total_spent,
       AVG(ro.o_totalprice) AS avg_order_value,
       RANK() OVER (ORDER BY SUM(ro.o_totalprice) DESC) AS spend_rank,
       CASE
         WHEN SUM(ro.o_totalprice) > 500000 THEN 'VIP'
         WHEN SUM(ro.o_totalprice) > 100000 THEN 'Premium'
         ELSE 'Regular'
       END AS customer_tier,
       s.s_name AS supplier_name,
       s.s_acctbal AS supplier_balance
FROM tpch_jdbc.customer c
JOIN recent_orders ro
     ON c.c_custkey = ro.o_custkey
JOIN tpch_jdbc.supplier s
     ON c.c_nationkey = s.s_nationkey
GROUP BY c.c_custkey, c.c_name, c.c_mktsegment, s.s_name, s.s_acctbal
ORDER BY total_spent DESC
LIMIT 5;

It takes around 260 sec with mt_dop =1 while with mt_dop =7,
it takes close to 77 sec.

Also added JdbcJniWaitTime, JdbcCursorFetchTime,
JdbcLockWaitTime and JdbcJniCallCount for better analysis
of jdbc queries, in query profile.

Testing:
Some planner tests are included in jdbc-parallel.test and benchmark
tests are added in test_ext_data_sources.py.

Generated by:
Some codes and comments are generated with
the help of claude-4.6-sonnet-medium-thinking.

Change-Id: I3b25b99f5cb77d32c111ba37c0a01378ffdc1107
Reviewed-on: http://gerrit.cloudera.org:8080/24133
Tested-by: Impala Public Jenkins <impala-public-jenkins@cloudera.com>
Reviewed-by: Zoltan Borok-Nagy <boroknagyz@cloudera.com>

Author: pranavyl

be/src/exec/data-source-scan-node.cc

@@ -37,6 +37,7 @@
 #include "util/jni-util.h"
 #include "util/periodic-counter-updater.h"
 #include "util/runtime-profile-counters.h"
+#include "util/time.h"
 #include "util/ubsan.h"
 
 #include "common/names.h"
@@ -52,6 +53,26 @@ namespace impala {
 PROFILE_DEFINE_COUNTER(NumExternalDataSourceGetNext, DEBUG, TUnit::UNIT,
     "The total number of calls to ExternalDataSource::GetNext()");
 
+PROFILE_DEFINE_TIMER(JdbcJniWaitTime, UNSTABLE,
+    "Aggregate wall-clock time this scanner thread spent blocked inside the full "
+    "JNI/Java JDBC round-trip (Thrift serialization + fetchLock wait + JDBC cursor "
+    "fetch + Thrift deserialization). Use JdbcCursorFetchTime and JdbcLockWaitTime "
+    "for the Java-reported sub-components.");
+
+PROFILE_DEFINE_TIMER(JdbcCursorFetchTime, UNSTABLE,
+    "Aggregate time this scanner thread held the Java-side fetchLock while iterating "
+    "the JDBC cursor, as reported by the Java layer. Directly shows JDBC cursor "
+    "utilisation; divide by JdbcJniWaitTime to get the cursor-active fraction.");
+
+PROFILE_DEFINE_TIMER(JdbcLockWaitTime, UNSTABLE,
+    "Aggregate time this scanner thread waited to acquire the Java-side fetchLock "
+    "before gaining access to the JDBC cursor, as reported by the Java layer. "
+    "High values indicate cursor contention between concurrent scanner threads.");
+
+PROFILE_DEFINE_COUNTER(JdbcJniCallCount, DEBUG, TUnit::UNIT,
+    "Number of JNI GetNext() calls issued to the JDBC data source by "
+    "this scanner thread.");
+
 // $0 = num expected cols, $1 = actual num columns
 const string ERROR_NUM_COLUMNS = "Data source returned unexpected number of columns. "
     "Expected $0 but received $1. This likely indicates a problem with the data source "
@@ -73,9 +94,80 @@ const string ERROR_MEM_LIMIT_EXCEEDED = "DataSourceScanNode::$0() failed to allo
 // Size of an encoded TIMESTAMP
 const size_t TIMESTAMP_SIZE = sizeof(int64_t) + sizeof(int32_t);
 
-DataSourceScanNode::DataSourceScanNode(ObjectPool* pool, const ScanPlanNode& pnode,
-    const DescriptorTbl& descs)
+// -----------------------------------------------------------------------
+// SharedJdbcConnection
+// -----------------------------------------------------------------------
+
+Status SharedJdbcConnection::Open(const string& jar_path, const string& class_name,
+    const string& api_version, const string& init_string,
+    const extdatasource::TOpenParams& params, extdatasource::TOpenResult* result) {
+  ref_count_.fetch_add(1, std::memory_order_relaxed);
+
+  std::unique_lock<std::mutex> lk(open_mu_);
+  if (!open_done_) {
+    // First caller: initialize the executor and open the Java data source.
+    // All N C++ scanner threads share this single connection.
+    Status s = executor_.Init(jar_path, class_name, api_version, init_string);
+    if (s.ok()) s = executor_.Open(params, result);
+    if (s.ok()) s = Status(result->status);
+    if (s.ok()) scan_handle_ = result->scan_handle;
+    open_status_ = s;
+    open_done_ = true;
+    open_cv_.notify_all();
+    return s;
+  }
+
+  // Subsequent callers: wait for the first caller to finish, then reuse.
+  open_cv_.wait(lk, [this] { return open_done_; });
+  result->__set_scan_handle(scan_handle_);
+  return open_status_;
+}
+
+Status SharedJdbcConnection::FetchBatch(extdatasource::TGetNextResult* result) {
+  // Fast path: another thread already exhausted the stream.
+  if (eos_.load(std::memory_order_acquire)) {
+    result->__set_eos(true);
+    return Status::OK();
+  }
+
+  extdatasource::TGetNextParams params;
+  params.__set_scan_handle(scan_handle_);
+
+  // Multiple threads may call GetNext() concurrently. Serialization is provided by
+  // the Java-side fetchLock in JdbcRecordIterator::fetchBatch(), so only one thread
+  // holds the JDBC cursor at a time. Batch conversion happens in parallel in C++.
+  RETURN_IF_ERROR(executor_.GetNext(params, result));
+  if (result->eos) eos_.store(true, std::memory_order_release);
+  return Status(result->status);
+}
+
+Status SharedJdbcConnection::Close(const extdatasource::TCloseParams& params,
+    extdatasource::TCloseResult* result) {
+  // Only the last active instance closes the shared JDBC connection.
+  if (ref_count_.fetch_sub(1, std::memory_order_acq_rel) == 1) {
+    // Guard against the case where Open() was never called or failed: only invoke
+    // executor_.Close() if the Java data source was successfully opened.
+    std::unique_lock<std::mutex> lk(open_mu_);
+    bool should_close = open_done_ && open_status_.ok();
+    lk.unlock();
+    if (should_close) return executor_.Close(params, result);
+  }
+  return Status::OK();
+}
+
+// DataSourceScanPlanNode
+Status DataSourceScanPlanNode::CreateExecNode(RuntimeState* state,
+    ExecNode** node) const {
+  *node = state->obj_pool()->Add(
+      new DataSourceScanNode(state->obj_pool(), *this, state->desc_tbl()));
+  return Status::OK();
+}
+
+//DataSourceScanNode
+DataSourceScanNode::DataSourceScanNode(ObjectPool* pool,
+    const DataSourceScanPlanNode& pnode, const DescriptorTbl& descs)
     : ScanNode(pool, pnode, descs),
+      plan_node_(pnode),
       data_src_node_(pnode.tnode_->data_source_node),
       tuple_idx_(0),
       num_rows_(0),
@@ -90,14 +182,17 @@ Status DataSourceScanNode::Prepare(RuntimeState* state) {
   tuple_desc_ = state->desc_tbl().GetTupleDescriptor(data_src_node_.tuple_id);
   DCHECK(tuple_desc_ != NULL);
 
-  data_source_executor_.reset(new ExternalDataSourceExecutor());
-  RETURN_IF_ERROR(data_source_executor_->Init(data_src_node_.data_source.hdfs_location,
-      data_src_node_.data_source.class_name, data_src_node_.data_source.api_version,
-      data_src_node_.init_string));
+  // Point at the plan node's shared connection. The object is owned by the plan node
+  // and lives for the lifetime of the fragment; no construction needed here.
+  shared_conn_ = &plan_node_.shared_conn;
 
   cols_next_val_idx_.resize(tuple_desc_->slots().size(), 0);
   num_ext_data_source_get_next_ =
       PROFILE_NumExternalDataSourceGetNext.Instantiate(runtime_profile_);
+  jdbc_jni_wait_timer_ = PROFILE_JdbcJniWaitTime.Instantiate(runtime_profile_);
+  jdbc_cursor_fetch_timer_ = PROFILE_JdbcCursorFetchTime.Instantiate(runtime_profile_);
+  jdbc_lock_wait_timer_ = PROFILE_JdbcLockWaitTime.Instantiate(runtime_profile_);
+  jdbc_jni_call_count_ = PROFILE_JdbcJniCallCount.Instantiate(runtime_profile_);
   return Status::OK();
 }
 
@@ -107,7 +202,8 @@ Status DataSourceScanNode::Open(RuntimeState* state) {
   RETURN_IF_ERROR(ExecNode::Open(state));
   RETURN_IF_CANCELLED(state);
 
-  // Prepare the schema for TOpenParams.row_schema
+  // Build TOpenParams. All N scanner threads share the single Java data source stream,
+  // with fetch serialization provided by the Java-side fetchLock in JdbcRecordIterator.
   vector<extdatasource::TColumnDesc> cols;
   for (const SlotDescriptor* slot: tuple_desc_->slots()) {
     extdatasource::TColumnDesc col;
@@ -128,10 +224,18 @@ Status DataSourceScanNode::Open(RuntimeState* state) {
   params.__set_batch_size(FLAGS_data_source_batch_size);
   params.__set_predicates(data_src_node_.accepted_predicates);
   params.__set_clean_dbcp_ds_cache(state->query_options().clean_dbcp_ds_cache);
+
+  // The first caller across all N scanner threads initializes the connection;
+  // all others block on open_cv_ until it is ready and then reuse it.
   TOpenResult result;
-  RETURN_IF_ERROR(data_source_executor_->Open(params, &result));
-  RETURN_IF_ERROR(Status(result.status));
-  scan_handle_ = result.scan_handle;
+  RETURN_IF_ERROR(shared_conn_->Open(
+      data_src_node_.data_source.hdfs_location,
+      data_src_node_.data_source.class_name,
+      data_src_node_.data_source.api_version,
+      data_src_node_.init_string,
+      params, &result));
+  scan_handle_ = shared_conn_->scan_handle();
+
   return GetNextInputBatch();
 }
 
@@ -161,15 +265,30 @@ Status DataSourceScanNode::ValidateRowBatchSize() {
 }
 
 Status DataSourceScanNode::GetNextInputBatch() {
-  input_batch_.reset(new TGetNextResult());
   next_row_idx_ = 0;
-  // Reset all the indexes into the column value arrays to 0
   Ubsan::MemSet(cols_next_val_idx_.data(), 0, sizeof(int) * cols_next_val_idx_.size());
-  TGetNextParams params;
-  params.__set_scan_handle(scan_handle_);
+
+  input_batch_.reset(new TGetNextResult());
+  // Short-circuit: another scanner thread already reached EOS on the shared stream.
+  if (shared_conn_->eos()) {
+    input_batch_->__set_eos(true);
+    return Status::OK();
+  }
+  // FetchBatch() serializes the JNI call via the Java-side fetchLock; this scan node
+  // then materializes the returned batch in parallel while others may be fetching.
   COUNTER_ADD(num_ext_data_source_get_next_, 1);
-  RETURN_IF_ERROR(data_source_executor_->GetNext(params, input_batch_.get()));
-  RETURN_IF_ERROR(Status(input_batch_->status));
+  COUNTER_ADD(jdbc_jni_call_count_, 1);
+  {
+    SCOPED_TIMER(jdbc_jni_wait_timer_);
+    RETURN_IF_ERROR(shared_conn_->FetchBatch(input_batch_.get()));
+  }
+  if (input_batch_->__isset.cursor_fetch_time_ns) {
+    COUNTER_ADD(jdbc_cursor_fetch_timer_, input_batch_->cursor_fetch_time_ns);
+  }
+  if (input_batch_->__isset.lock_wait_time_ns) {
+    COUNTER_ADD(jdbc_lock_wait_timer_, input_batch_->lock_wait_time_ns);
+  }
+
   RETURN_IF_ERROR(ValidateRowBatchSize());
   if (!InputBatchHasNext() && !input_batch_->eos) {
     // The data source should have set eos, but if it didn't we should just log a
@@ -416,7 +535,9 @@ void DataSourceScanNode::Close(RuntimeState* state) {
   TCloseParams params;
   params.__set_scan_handle(scan_handle_);
   TCloseResult result;
-  Status status = data_source_executor_->Close(params, &result);
+  // shared_conn_->Close() only calls executor_.Close() when the last active
+  // DataSourceScanNode in the fragment decrements the reference count to zero.
+  Status status = shared_conn_->Close(params, &result);
   if (!status.ok()) state->LogError(status.msg());
   ScanNode::Close(state);
 }

be/src/exec/data-source-scan-node.h

@@ -19,55 +19,128 @@
 #ifndef IMPALA_EXEC_DATA_SOURCE_SCAN_NODE_H_
 #define IMPALA_EXEC_DATA_SOURCE_SCAN_NODE_H_
 
+#include <atomic>
+#include <condition_variable>
+#include <mutex>
 #include <string>
+#include <vector>
 #include <boost/scoped_ptr.hpp>
 
 #include "common/global-types.h"
-#include "exec/scan-node.h"
+#include "common/status.h"
 #include "exec/external-data-source-executor.h"
+#include "exec/scan-node.h"
 #include "runtime/descriptors.h"
 #include "runtime/mem-pool.h"
 
 #include "gen-cpp/ExternalDataSource_types.h"
 
 namespace impala {
 
+class DataSourceScanNode;
 class Tuple;
 
+/// Shared JDBC connection used by all DataSourceScanNode instances within the same
+/// fragment. A single ExternalDataSourceExecutor is opened once and then called
+/// concurrently by N MT_DOP scanner threads. Fetch serialization is provided by the
+/// Java-side fetchLock in JdbcRecordIterator; once a thread returns from FetchBatch()
+/// it materializes its batch in parallel with other threads.
+class SharedJdbcConnection {
+ public:
+  SharedJdbcConnection() {}
+
+  /// Opens the JDBC connection. Thread-safe: the first caller initializes the executor
+  /// and opens the Java data source; all subsequent callers block until the first caller
+  /// finishes and then reuse the same connection.
+  Status Open(const std::string& jar_path, const std::string& class_name,
+      const std::string& api_version, const std::string& init_string,
+      const extdatasource::TOpenParams& params, extdatasource::TOpenResult* result);
+
+  /// Fetches the next batch from the shared JDBC connection. Multiple threads may
+  /// call this concurrently; serialization is handled by the Java-side fetchLock.
+  /// Short-circuits immediately if eos() is already true.
+  Status FetchBatch(extdatasource::TGetNextResult* result);
+
+  /// Decrements the open reference count. The last caller (ref_count_ reaches 0)
+  /// closes the Java data source.
+  Status Close(const extdatasource::TCloseParams& params,
+      extdatasource::TCloseResult* result);
+
+  const std::string& scan_handle() const { return scan_handle_; }
+
+  /// Returns true once the JDBC stream has been fully consumed. All scanner threads
+  /// check this before calling FetchBatch() to avoid redundant JNI calls.
+  bool eos() const { return eos_.load(std::memory_order_acquire); }
+
+ private:
+  ExternalDataSourceExecutor executor_;
+  std::string scan_handle_;
+
+  /// Guards open_done_ / open_status_: used in Open() for one-time initialization
+  /// and in Close() to check whether executor_.Close() should be called.
+  std::mutex open_mu_;
+  std::condition_variable open_cv_;
+  bool open_done_ = false;
+  Status open_status_;
+
+  /// Reference count: incremented in Open(), decremented in Close(). When it reaches
+  /// zero, the last DSSN instance calls executor_.Close().
+  std::atomic<int> ref_count_{0};
+
+  /// Set to true when the JDBC stream is exhausted. Threads check this before
+  /// crossing the JNI bridge to avoid unnecessary calls after EOS.
+  std::atomic<bool> eos_{false};
+};
+
+/// Per-fragment plan node for data source scans. Created once per fragment (shared
+/// across all DataSourceScanNode instances), it owns the SharedJdbcConnection so that
+/// all N MT_DOP instances share a single JDBC connection with parallel batch conversion.
+class DataSourceScanPlanNode : public ScanPlanNode {
+ public:
+  virtual Status CreateExecNode(RuntimeState* state, ExecNode** node) const override;
+
+  /// Owned by the plan node; shared by all DataSourceScanNode instances in the fragment.
+  /// Mutable so that Open() (called from exec nodes) can initialize it on first call.
+  mutable SharedJdbcConnection shared_conn;
+};
+
 /// Scan node for external data sources. The external data source jar is loaded
-/// in Prepare() (via an ExternalDataSourceExecutor), and then the data source
-/// is called to receive row batches when necessary. This node converts the
-/// rows stored in a thrift structure to RowBatches. The external data source is
-/// closed in Close().
+/// in Prepare() (via SharedJdbcConnection), and then the data source is called to
+/// receive row batches when necessary. N MT_DOP instances of this node share a single
+/// SharedJdbcConnection; fetch serialization is handled by the Java-side fetchLock
+/// while batch conversion/materialization runs in parallel across all N threads.
 class DataSourceScanNode : public ScanNode {
  public:
-  DataSourceScanNode(
-      ObjectPool* pool, const ScanPlanNode& pnode, const DescriptorTbl& descs);
+  DataSourceScanNode(ObjectPool* pool, const DataSourceScanPlanNode& pnode,
+      const DescriptorTbl& descs);
 
   ~DataSourceScanNode();
 
-  /// Load the data source library and create the ExternalDataSourceExecutor.
+  /// Initialize the shared JDBC connection (once per fragment) and profile counters.
   virtual Status Prepare(RuntimeState* state) override;
 
-  /// Open the data source and initialize the first row batch.
+  /// Opens the data source via the shared connection and fetches the first batch.
   virtual Status Open(RuntimeState* state) override;
 
-  /// Fill the next row batch, calls GetNext() on the external scanner.
+  /// Fill the next row batch from the external scanner.
   virtual Status GetNext(RuntimeState* state, RowBatch* row_batch, bool* eos) override;
 
   /// NYI
   virtual Status Reset(RuntimeState* state, RowBatch* row_batch) override;
 
-  /// Close the scanner, and report errors.
+  /// Close the scanner via the shared connection.
   virtual void Close(RuntimeState* state) override;
 
  protected:
   /// Write debug string of this into out.
   virtual void DebugString(int indentation_level, std::stringstream* out) const override;
 
  private:
-  /// Used to call the external data source.
-  boost::scoped_ptr<ExternalDataSourceExecutor> data_source_executor_;
+  /// Reference to the per-fragment plan node; owns shared_conn.
+  const DataSourceScanPlanNode& plan_node_;
+
+  /// Non-owning pointer to the plan node's shared JDBC connection.
+  SharedJdbcConnection* shared_conn_ = nullptr;
 
   /// Thrift structure describing the data source scan node.
   const TDataSourceScanNode data_src_node_;
@@ -78,7 +151,8 @@ class DataSourceScanNode : public ScanNode {
   /// Tuple index in tuple row.
   int tuple_idx_;
 
-  /// The opaque handle returned by the data source for the scan.
+  /// The opaque handle returned by the data source for the scan (shared across all
+  /// instances — all call getNext() on the same Java object using the same handle).
   std::string scan_handle_;
 
   /// The current result from calling GetNext() on the data source. Contains the
@@ -101,11 +175,25 @@ class DataSourceScanNode : public ScanNode {
   /// The total number of calls to ExternalDataSource::GetNext().
   RuntimeProfile::Counter* num_ext_data_source_get_next_;
 
+  /// Wall-clock time this thread spent blocked in the full JNI/Java JDBC round-trip.
+  RuntimeProfile::Counter* jdbc_jni_wait_timer_ = nullptr;
+
+  /// Time this thread held the Java-side fetchLock while iterating the JDBC cursor,
+  /// as reported by the Java layer. Reflects pure JDBC cursor utilisation.
+  RuntimeProfile::Counter* jdbc_cursor_fetch_timer_ = nullptr;
+
+  /// Time this thread waited to acquire the Java-side fetchLock, as reported by the
+  /// Java layer. High values indicate cursor contention across scanner threads.
+  RuntimeProfile::Counter* jdbc_lock_wait_timer_ = nullptr;
+
+  /// Number of JNI GetNext() calls issued by this scan node.
+  RuntimeProfile::Counter* jdbc_jni_call_count_ = nullptr;
+
   /// Materializes the next row (next_row_idx_) into tuple. 'local_tz' is used as the
   /// local time-zone for materializing 'TYPE_TIMESTAMP' slots.
   Status MaterializeNextRow(const Timezone* local_tz, MemPool* mem_pool, Tuple* tuple);
 
-  /// Gets the next batch from the data source, stored in input_batch_.
+  /// Fetches the next batch from the data source via the shared JDBC connection.
   Status GetNextInputBatch();
 
   /// Validate row_batch_ contains the correct number of columns and that columns