benchmark-test: cancellation, memory-release, and reduce-path diagnos…

.gitignore

@@ -75,3 +75,5 @@ doc-tools/missing-doclet/bin/
 **/Cargo.lock
 /sandbox/plugins/analytics-backend-datafusion/target/
 /sandbox/libs/dataformat-native/rust/target
+.docker-gradle-home/
+scripts/docker-crash/

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSink.java

@@ -50,6 +50,16 @@ default void feed(VectorSchemaRoot batch, int sourceOrdinal) {
         feed(batch);
     }
 
+    /**
+     * Whether the downstream consumer has finished and will read no more batches (e.g. a reduce
+     * whose LimitExec satisfied its fetch). Producers may poll this after a {@link #feed} to stop
+     * early. Default {@code false}; best-effort — a {@code true} is monotonic, a {@code false} may
+     * race a concurrent completion.
+     */
+    default boolean isConsumerDone() {
+        return false;
+    }
+
     /**
      * Signal that no more batches will be fed. Releases resources.
      */

sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/AggregateFunctionTests.java

@@ -164,6 +164,24 @@ public void testValuesReducerIsSelf() {
         assertEquals(integer, resolve(fields.get(0), integer));
     }
 
+    // ── PERCENTILE_APPROX: state-bearing, no decomposition declared ──
+    //
+    // DataFusion's t-digest state is multi-field and doesn't fit the single-field
+    // IntermediateField shape. OpenSearchAggregateSplitRule skips the partial/final split
+    // for STATE_EXPANDING aggregates, so PERCENTILE_APPROX runs single-stage on the
+    // coordinator after a singleton gather.
+    public void testPercentileApproxHasNoDecomposition() {
+        assertFalse(AggregateFunction.PERCENTILE_APPROX.hasDecomposition());
+        assertNull(AggregateFunction.PERCENTILE_APPROX.intermediateFields());
+        assertEquals(AggregateFunction.Type.STATE_EXPANDING, AggregateFunction.PERCENTILE_APPROX.getType());
+        assertEquals(SqlKind.OTHER, AggregateFunction.PERCENTILE_APPROX.getSqlKind());
+    }
+
+    public void testPercentileApproxResolvesByName() {
+        assertSame(AggregateFunction.PERCENTILE_APPROX, AggregateFunction.fromNameOrError("percentile_approx"));
+        assertSame(AggregateFunction.PERCENTILE_APPROX, AggregateFunction.fromNameOrError("PERCENTILE_APPROX"));
+    }
+
     // ── fromSqlKind still works ──
 
     public void testFromSqlKindResolvesExistingEntries() {

sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs

@@ -323,7 +323,7 @@ pub fn create_global_runtime(
     let memory_pool = Arc::new(TrackConsumersPool::new(
         dynamic_pool,
         NonZeroUsize::new(5).unwrap(),
-    ));
+    )) as Arc<dyn datafusion::execution::memory_pool::MemoryPool>;
 
     let (cache_manager_config, custom_cache_manager) = if cache_manager_ptr != 0 {
         let mgr = unsafe { *Box::from_raw(cache_manager_ptr as *mut CustomCacheManager) };
@@ -406,6 +406,18 @@ pub fn set_min_target_partitions(value: i64) {
     crate::query_budget::set_min_target_partitions(value.max(1) as usize);
 }
 
+/// Initial target_partitions for coordinator-reduce sessions. Defaults to 4.
+static REDUCE_TARGET_PARTITIONS: std::sync::atomic::AtomicUsize =
+    std::sync::atomic::AtomicUsize::new(4);
+
+pub fn set_reduce_target_partitions(value: i64) {
+    REDUCE_TARGET_PARTITIONS.store(value.max(1).min(32) as usize, std::sync::atomic::Ordering::Release);
+}
+
+pub fn get_reduce_target_partitions() -> usize {
+    REDUCE_TARGET_PARTITIONS.load(std::sync::atomic::Ordering::Acquire)
+}
+
 /// Creates a native reader (ShardView) for the given path and files.
 ///
 /// Returns a heap-allocated pointer (as i64) to `ShardView`.
@@ -908,6 +920,7 @@ pub async unsafe fn stream_next(
             } else {
                 batch
             };
+
             let struct_array: StructArray = batch.into();
             let array_data = struct_array.into_data();
             let ffi_array = FFI_ArrowArray::new(&array_data);
@@ -978,9 +991,10 @@ fn view_needs_gc(buffers: &[arrow::buffer::Buffer], bytes_used: usize) -> bool {
 /// `stream_ptr` must be 0 or a valid pointer returned by `execute_query`.
 pub unsafe fn stream_close(stream_ptr: i64) {
     if stream_ptr != 0 {
-        // Dropping the handle drops both the stream and the query context.
-        // The context's Drop impl marks the query completed in the registry.
-        let _ = Box::from_raw(stream_ptr as *mut QueryStreamHandle);
+        native_bridge_common::log_debug!("[stream-close] dropping QueryStreamHandle ptr={:#x}", stream_ptr);
+        let handle = Box::from_raw(stream_ptr as *mut QueryStreamHandle);
+        drop(handle);
+        native_bridge_common::log_debug!("[stream-close] QueryStreamHandle dropped ptr={:#x}", stream_ptr);
     }
 }
 
@@ -1317,11 +1331,10 @@ pub(crate) fn base_schema_for_table(plan: &substrait::proto::Plan, table_name: &
 /// `create_global_runtime`.
 pub unsafe fn create_local_session(runtime_ptr: i64) -> Result<i64, DataFusionError> {
     let runtime = &*(runtime_ptr as *const DataFusionRuntime);
-    // No phantom reservation at creation time — the schema isn't known yet.
-    // register_partition_stream acquires a schema-accurate phantom once the
-    // output schema is derived from the producer plan.
     let session = LocalSession::new(&runtime.runtime_env);
-    Ok(Box::into_raw(Box::new(session)) as i64)
+    let ptr = Box::into_raw(Box::new(session)) as i64;
+    native_bridge_common::log_debug!("[local-session] OPEN ptr={:#x}", ptr);
+    Ok(ptr)
 }
 
 /// Closes a `LocalSession`. Safe to call with 0 (no-op).
@@ -1330,7 +1343,10 @@ pub unsafe fn create_local_session(runtime_ptr: i64) -> Result<i64, DataFusionEr
 /// `ptr` must be 0 or a valid pointer returned by `create_local_session`.
 pub unsafe fn close_local_session(ptr: i64) {
     if ptr != 0 {
-        let _ = Box::from_raw(ptr as *mut LocalSession);
+        let session = Box::from_raw(ptr as *mut LocalSession);
+        let phantom_size = session.phantom_size();
+        native_bridge_common::log_debug!("[local-session] CLOSE ptr={:#x} phantom_bytes={}", ptr, phantom_size);
+        drop(session);
     }
 }
 
@@ -1502,7 +1518,7 @@ pub unsafe fn sender_send(
     array_ptr: i64,
     schema_ptr: i64,
     io_handle: &tokio::runtime::Handle,
-) -> Result<(), DataFusionError> {
+) -> Result<crate::partition_stream::SendOutcome, DataFusionError> {
     let sender = &*(sender_ptr as *const PartitionStreamSender);
 
     // Take ownership of the Java-allocated FFI structs. `from_raw` reads
@@ -1525,7 +1541,7 @@ pub unsafe fn sender_send(
     let struct_array = StructArray::from(array_data);
     let batch = RecordBatch::from(struct_array);
 
-    sender.send_blocking(Ok(batch), io_handle)
+    Ok(sender.send_blocking(Ok(batch), io_handle))
 }
 
 /// Closes a partition stream sender. Dropping the sender closes the mpsc,

sandbox/plugins/analytics-backend-datafusion/rust/src/cross_rt_stream.rs

@@ -36,6 +36,12 @@ pub struct CrossRtStream {
     phantom_corrector: Option<Arc<PhantomCorrector>>,
 }
 
+impl Drop for CrossRtStream {
+    fn drop(&mut self) {
+        native_bridge_common::log_debug!("[cross-rt-stream] DROP");
+    }
+}
+
 impl CrossRtStream {
     fn new_with_tx<F, Fut>(f: F, schema: SchemaRef) -> Self
     where
@@ -77,9 +83,11 @@ impl CrossRtStream {
             tokio::pin!(stream);
             while let Some(res) = stream.next().await {
                 if tx_captured.send(res).await.is_err() {
+                    native_bridge_common::log_debug!("[cross-rt-stream] CPU task: receiver gone, exiting");
                     return;
                 }
             }
+            native_bridge_common::log_debug!("[cross-rt-stream] CPU task: stream exhausted normally");
         };
 
         let (abort_handle, join_fut) = exec.spawn_with_abort_handle(fut);

sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs

@@ -157,6 +157,11 @@ pub extern "C" fn df_set_min_target_partitions(value: i64) {
     api::set_min_target_partitions(value);
 }
 
+#[no_mangle]
+pub extern "C" fn df_set_reduce_target_partitions(value: i64) {
+    api::set_reduce_target_partitions(value);
+}
+
 /// Sets memory guard thresholds. Values are thresholds multiplied by 1000
 /// (e.g., 700 = 0.70, 850 = 0.85, 950 = 0.95).
 #[no_mangle]
@@ -542,12 +547,20 @@ pub unsafe extern "C" fn df_execute_local_plan(
         .map_err(|e| e.to_string())
 }
 
+/// Positive FFI sentinel from [`df_sender_send`] when the send was skipped because the consumer
+/// finished (receiver dropped). Rides the success half of the return contract (`>= 0` = success),
+/// so `checkResult` passes it through. MUST match `NativeBridge.SENDER_SEND_RECEIVER_DROPPED`.
+pub const SENDER_SEND_RECEIVER_DROPPED: i64 = 1;
+
 #[ffm_safe]
 #[no_mangle]
 pub unsafe extern "C" fn df_sender_send(sender_ptr: i64, array_ptr: i64, schema_ptr: i64) -> i64 {
     let mgr = get_rt_manager()?;
     api::sender_send(sender_ptr, array_ptr, schema_ptr, mgr.io_runtime.handle())
-        .map(|_| 0)
+        .map(|outcome| match outcome {
+            crate::partition_stream::SendOutcome::Sent => 0,
+            crate::partition_stream::SendOutcome::ReceiverDropped => SENDER_SEND_RECEIVER_DROPPED,
+        })
         .map_err(|e| e.to_string())
 }
 

sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs

@@ -71,7 +71,7 @@ impl LocalSession {
     pub fn new(runtime_env: &RuntimeEnv) -> Self {
         let runtime_env = Arc::new(runtime_env.clone());
         let mut config = SessionConfig::new();
-        config.options_mut().execution.target_partitions = 4;
+        config.options_mut().execution.target_partitions = crate::api::get_reduce_target_partitions();
         let state = SessionStateBuilder::new()
             .with_config(config)
             .with_runtime_env(runtime_env)
@@ -217,16 +217,15 @@ impl LocalSession {
             ))
         })?;
         let logical_plan = from_substrait_plan(&self.ctx.state(), &plan).await?;
-        log_debug!("DataFusion logical plan (reduce):\n{}", logical_plan.display_indent());
         let dataframe = self.ctx.execute_logical_plan(logical_plan).await?;
         let physical_plan = dataframe.create_physical_plan().await?;
         let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema());
         let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?;
-        log_debug!("DataFusion physical plan (reduce):\n{}", displayable(physical_plan.as_ref()).indent(true));
         let stripped = crate::agg_mode::apply_aggregate_mode(
             physical_plan,
             crate::agg_mode::Mode::Final,
         )?;
+        log_debug!("DataFusion physical plan (reduce):\n{}", displayable(stripped.as_ref()).indent(true));
         self.prepared_plan = Some(stripped);
         Ok(())
     }
@@ -330,9 +329,8 @@ mod tests {
         let handle = Handle::current();
         let producer = std::thread::spawn(move || {
             for chunk in &[vec![1i64, 2, 3], vec![4, 5, 6], vec![7, 8, 9]] {
-                sender
-                    .send_blocking(Ok(i64_batch(&producer_schema, chunk)), &handle)
-                    .expect("send");
+                let outcome = sender.send_blocking(Ok(i64_batch(&producer_schema, chunk)), &handle);
+                assert!(matches!(outcome, crate::partition_stream::SendOutcome::Sent));
             }
             drop(sender); // EOF
         });

sandbox/plugins/analytics-backend-datafusion/rust/src/partition_stream.rs

@@ -60,6 +60,16 @@ pub struct PartitionStreamSender {
     schema: SchemaRef,
 }
 
+/// Outcome of a blocking send. `ReceiverDropped` is the benign terminal case — the
+/// DataFusion consumer finished (e.g. a `LimitExec` satisfied its fetch) and dropped the
+/// receiver. Surfaced as a distinct variant so the FFM layer can signal it without the Java
+/// side substring-matching an error message.
+#[must_use]
+pub enum SendOutcome {
+    Sent,
+    ReceiverDropped,
+}
+
 impl PartitionStreamSender {
     /// Returns the schema this sender was created with.
     pub fn schema(&self) -> &SchemaRef {
@@ -73,16 +83,18 @@ impl PartitionStreamSender {
     /// bridge push without being async itself and without requiring the calling
     /// thread to be a Tokio worker.
     ///
-    /// Blocks while the channel is full (natural backpressure). Returns an
-    /// error only if the receiver has been dropped.
+    /// Blocks while the channel is full (natural backpressure). Returns
+    /// [`SendOutcome::ReceiverDropped`] only if the receiver has been dropped —
+    /// the sole failure mode of `mpsc::Sender::send`.
     pub fn send_blocking(
         &self,
         batch: Result<RecordBatch, DataFusionError>,
         handle: &Handle,
-    ) -> Result<(), DataFusionError> {
-        handle.block_on(self.tx.send(batch)).map_err(|_| {
-            DataFusionError::Execution("partition stream receiver dropped before send".to_string())
-        })
+    ) -> SendOutcome {
+        match handle.block_on(self.tx.send(batch)) {
+            Ok(()) => SendOutcome::Sent,
+            Err(_) => SendOutcome::ReceiverDropped,
+        }
     }
 }
 
@@ -259,9 +271,8 @@ mod tests {
 
         let sender_schema = Arc::clone(&schema);
         let producer = std::thread::spawn(move || {
-            sender
-                .send_blocking(Ok(test_batch(&sender_schema, &[7, 8, 9])), &handle)
-                .unwrap();
+            let outcome = sender.send_blocking(Ok(test_batch(&sender_schema, &[7, 8, 9])), &handle);
+            assert!(matches!(outcome, SendOutcome::Sent));
             drop(sender);
         });
 
@@ -280,14 +291,10 @@ mod tests {
         let (sender, receiver) = channel(Arc::clone(&schema));
         drop(receiver);
 
-        let err = std::thread::spawn(move || {
-            sender
-                .send_blocking(Ok(test_batch(&schema, &[1])), &handle)
-                .unwrap_err()
-        })
-        .join()
-        .unwrap();
-        assert!(err.to_string().contains("receiver dropped"));
+        let outcome = std::thread::spawn(move || sender.send_blocking(Ok(test_batch(&schema, &[1])), &handle))
+            .join()
+            .unwrap();
+        assert!(matches!(outcome, SendOutcome::ReceiverDropped));
     }
 
     #[tokio::test]

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java

@@ -12,10 +12,8 @@
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.types.pojo.Schema;
-import org.opensearch.analytics.spi.CancellableExchangeSink;
 import org.opensearch.analytics.spi.ExchangeSinkContext;
 import org.opensearch.analytics.spi.ReducingExchangeSink;
-import org.opensearch.be.datafusion.nativelib.NativeBridge;
 import org.opensearch.be.datafusion.nativelib.StreamHandle;
 
 import java.util.LinkedHashMap;
@@ -43,7 +41,7 @@
  *
  * @opensearch.internal
  */
-abstract class AbstractDatafusionReduceSink implements ReducingExchangeSink, CancellableExchangeSink {
+abstract class AbstractDatafusionReduceSink implements ReducingExchangeSink {
 
     /** Single-input shortcut for the per-child table id; multi-input uses {@link #inputIdFor(int)}. */
     static final String INPUT_ID = "input-0";
@@ -96,20 +94,6 @@ protected static String inputIdFor(int childStageId) {
         return "input-" + childStageId;
     }
 
-    /**
-     * Fires the Rust CancellationToken for this query so that {@code stream_next}
-     * returns the sentinel {@code 0} immediately and the drain unblocks without
-     * waiting for DataFusion to finish naturally. No-op when {@code taskId} is 0
-     * (no context registered) or when already closed.
-     */
-    @Override
-    public final void cancel() {
-        long taskId = ctx.taskId();
-        if (taskId != 0L) {
-            NativeBridge.cancelQuery(taskId);
-        }
-    }
-
     @Override
     public void feed(VectorSchemaRoot batch) {
         synchronized (feedLock) {
@@ -163,7 +147,8 @@ protected final void drainOutputIntoDownstream(StreamHandle outStream) {
         try (CDataDictionaryProvider dictProvider = new CDataDictionaryProvider()) {
             DatafusionResultStream.BatchIterator it = new DatafusionResultStream.BatchIterator(outStream, alloc, dictProvider);
             while (it.hasNext()) {
-                ctx.downstream().feed(it.next().getArrowRoot());
+                VectorSchemaRoot root = it.next().getArrowRoot();
+                ctx.downstream().feed(root);
             }
         }
     }

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java

@@ -409,6 +409,7 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP
         AggregateFunction.COUNT,
         AggregateFunction.AVG,
         AggregateFunction.APPROX_COUNT_DISTINCT,
+        AggregateFunction.PERCENTILE_APPROX,
         AggregateFunction.TAKE,
         AggregateFunction.FIRST,
         AggregateFunction.LAST,