refactor(exec): own path hydration and query resources

crates/graphforge-exec/src/demand.rs

@@ -188,6 +188,8 @@ pub struct DemandSnapshot {
     pub cancellations: u64,
     /// Maximum simultaneous filtered-read calls.
     pub max_in_flight_reads: u64,
+    /// Query-owned path hydration work, using physical operator metric names.
+    pub hydration: BTreeMap<String, u64>,
     /// Query memory-pool reservation before physical execution.
     pub memory_reserved_before: u64,
     /// Query memory-pool reservation after every operator stream was dropped.
@@ -430,6 +432,9 @@ pub(crate) fn record_plan_completion(
             operator.peak_bytes = operator.peak_bytes.max(completion_rss);
         }
     }
+    let mut hydration = BTreeMap::new();
+    hydration_metrics(plan, &mut hydration);
+    state.snapshot.hydration = hydration;
     state.snapshot.sorts = sorts;
     state.snapshot.memory_reserved_before = memory_reserved_before as u64;
     state.snapshot.memory_reserved_after = memory_reserved_after as u64;
@@ -1593,6 +1598,22 @@ impl RecordBatchStream for DemandGuardStream {
     }
 }
 
+fn hydration_metrics(plan: &Arc<dyn ExecutionPlan>, result: &mut BTreeMap<String, u64>) {
+    if plan
+        .downcast_ref::<crate::path_hydration::HydrationExec>()
+        .is_some()
+        && let Some(metrics) = plan.metrics()
+    {
+        for metric in metrics.iter() {
+            let value = metric.value();
+            *result.entry(value.name().to_owned()).or_default() += value.as_usize() as u64;
+        }
+    }
+    for child in plan.children() {
+        hydration_metrics(child, result);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

crates/graphforge-exec/src/lib.rs

@@ -62,6 +62,7 @@ pub(crate) mod algorithm_embedding_hashgnn;
 mod algorithm_embedding_invocation;
 pub(crate) mod algorithm_embedding_options;
 pub mod mutation;
+mod path_hydration;
 pub mod read_resource;
 pub mod write_resource;
 pub use algorithm_embedding_options::validate_embedding_options;
@@ -4901,13 +4902,14 @@ impl QueryPlanner for GraphForgeQueryPlanner {
         logical_plan: &LogicalPlan,
         session_state: &SessionState,
     ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
-        let logical_plan = read_resource::bind(logical_plan, session_state)?;
+        let (logical_plan, hydration) = read_resource::bind(logical_plan, session_state)?;
         let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
             GraphForgeExtensionPlanner,
         )]);
-        planner
+        let physical = planner
             .create_physical_plan(&logical_plan, session_state)
-            .await
+            .await?;
+        Ok(path_hydration::wrap(physical, hydration))
     }
 }
 
@@ -4963,6 +4965,7 @@ impl QueryEvidenceStream {
             return;
         }
         self.finalized = true;
+        path_hydration::cancel_plan(&self.physical);
         drop(self.inner.take());
         demand::record_plan_completion(
             &self.physical,
@@ -5362,7 +5365,7 @@ impl ExecutionSession {
             .await
             .map_err(GfError::from_plan_error)?;
 
-        let batches = collect(Arc::clone(&physical), self.ctx.task_ctx())
+        let batches = path_hydration::collect_guarded(Arc::clone(&physical), self.ctx.task_ctx())
             .await
             .map_err(GfError::from_execution_error)?;
 
@@ -5547,7 +5550,7 @@ impl ExecutionSession {
             .create_physical_plan(&logical)
             .await
             .map_err(GfError::from_plan_error)?;
-        let batches = collect(physical, self.ctx.task_ctx())
+        let batches = path_hydration::collect_guarded(physical, self.ctx.task_ctx())
             .await
             .map_err(GfError::from_execution_error)?;
         let mut frontier = write_driver::Frontier { df_schema, batches };
@@ -5560,6 +5563,10 @@ impl ExecutionSession {
             mode: resource.mode,
             params,
             type_map: lowerer.entity_name_map(),
+            hydration: path_hydration::HydrationResource::new(
+                read_resource::required(&self.ctx.state()).map_err(GfError::from_plan_error)?,
+                Arc::clone(&self.ctx.runtime_env().memory_pool),
+            ),
         };
         let mut wctx = write_driver::StatementWriteContext::new(&resource.dir, resource.mode)?
             .with_semantic_composition_fingerprint(self.semantic_composition_fingerprint.clone());
@@ -5813,7 +5820,9 @@ impl ExecutionSession {
 
         let task_ctx = self.ctx.task_ctx();
         let memory_reserved_before = task_ctx.memory_pool().reserved();
-        let collected = collect(Arc::clone(&physical), Arc::clone(&task_ctx)).await;
+        let collected =
+            path_hydration::collect_guarded(Arc::clone(&physical), Arc::clone(&task_ctx)).await;
+        path_hydration::cancel_plan(&physical);
         let memory_reserved_after = task_ctx.memory_pool().reserved();
         let returned_batch_bytes = collected.as_ref().map_or(0, |batches| {
             batches
@@ -5897,7 +5906,7 @@ impl ExecutionSession {
                 .create_physical_plan(&logical)
                 .await
                 .map_err(GfError::from_plan_error)?;
-            let batches = collect(physical, self.ctx.task_ctx())
+            let batches = path_hydration::collect_guarded(physical, self.ctx.task_ctx())
                 .await
                 .map_err(GfError::from_execution_error)?;
             let batch = batches.first().ok_or_else(|| {