Skip to main content

query_runtime/runtime/
resident.rs

1//! Frontend-neutral GPU-resident data, result, and reusable build handles.
2//!
3//! Handles in this module are runtime/device-affine. Cloning shares the
4//! underlying native resources; host materialization is explicit.
5//!
6//! Besides the cuDF frame, a handle carries the lifetimes that make the frame
7//! safe to keep on the device: a `ResidentOwnerLease` for the bytes charged to
8//! the producing attempt's grant, transitive source owners for columns that alias
9//! an input frame's memory, a source-cache pin that keeps a cache entry from
10//! being evicted while the frame aliases it, and graph provenance. A consumer
11//! that reads a handle on its own stream wraps these owners with `for_reader`, so
12//! the owner's final release also waits for that stream's work; dropping a handle
13//! therefore never frees buffers a stream may still read. The release rules are
14//! owned by `website/docs/design/admission.mdx#resident-handoff-and-release`.
15
16use std::{
17    collections::{BTreeMap, HashSet},
18    mem,
19    sync::Arc,
20};
21
22use arrow::{
23    datatypes::SchemaRef,
24    record_batch::{RecordBatch, RecordBatchOptions},
25};
26use cudf::{
27    DataFrame, DataFrameFilteredJoin, DataFrameHashJoin, DeviceId, ExecutionPlacement,
28    frame::dataframe::{
29        ColumnMemoryIdentity, GpuRecordBatch, HostBatchCursor as CudfHostBatchCursor, HostBatchCursorConfig,
30        InfoDetail, PreparedHostBatch as CudfPreparedHostBatch,
31    },
32    runtime::CudaStream,
33};
34use query_planning::{
35    error::{QueryError, QueryErrorCode, Result},
36    naming::is_datafusion_order_column_name,
37    plan::{GraphResidencyDescriptor, JoinNullEquality, PlanFingerprint, QueryPlan},
38};
39
40use crate::{
41    native_error::cudf_query_error,
42    runtime::{
43        ExecutionCancellation, ExecutionContext,
44        cache::local_parquet::GpuLocalParquetSourceCachePin,
45        gpu::ledger::{MemoryGrant, ResidentOwnerLease},
46    },
47};
48
49/// Local lookup table for GPU-resident inputs consumed by a query.
50///
51/// Keys are plan-local input names and do not identify global handles. All
52/// handles must be reusable by the same [`ExecutionContext`] that executes the
53/// consuming plan.
54pub type GpuResidentInputMap = BTreeMap<String, GpuResultHandle>;
55
56/// Owned cuDF DataFrame plus the placement observed when it entered the engine.
57///
58/// The wrapper preserves the device placement contract while allowing callers
59/// to move the underlying frame between native operators. Clone semantics are
60/// those of `cudf::DataFrame`.
61#[derive(Clone)]
62pub struct GpuDataFrame {
63    frame: DataFrame,
64    placement: ExecutionPlacement,
65    /// Grant retention for the bytes this frame published. `None` for frames
66    /// whose memory is accounted only by the allocator of the attempt that
67    /// holds them, including cache checkouts.
68    resident_owner: Option<ResidentOwnerLease>,
69    /// Owners of input frames whose column memory this frame still aliases.
70    /// Narrowed on every projection so a derived frame retains only the
71    /// owners of columns it carries.
72    retained_source_owners: Vec<RetainedSourceOwner>,
73    /// Owner of a graph normalization input kept alive until the graph
74    /// consumer takes it; not propagated to derived frames.
75    retained_normalization_source: Option<ResidentOwnerLease>,
76    /// Keeps the local-Parquet cache entry this frame was copied from pinned,
77    /// and therefore non-evictable, for the frame's lifetime.
78    source_cache_pin: Option<GpuLocalParquetSourceCachePin>,
79    graph_residency: Option<Arc<GraphResidencyDescriptor>>,
80    producer_execution_scope: Option<PlanFingerprint>,
81}
82
83struct GpuDataFrameParts {
84    frame: DataFrame,
85    resident_owner: Option<ResidentOwnerLease>,
86    retained_source_owners: Vec<RetainedSourceOwner>,
87    source_cache_pin: Option<GpuLocalParquetSourceCachePin>,
88    graph_residency: Option<Arc<GraphResidencyDescriptor>>,
89    producer_execution_scope: Option<PlanFingerprint>,
90}
91
92// A transitive lease follows only aliases of the columns it protects. Keeping
93// the identity set prevents both premature release and fresh-only over-retention.
94#[derive(Clone)]
95struct RetainedSourceOwner {
96    owner: ResidentOwnerLease,
97    column_identities: HashSet<ColumnMemoryIdentity>,
98}
99
100impl RetainedSourceOwner {
101    fn for_reader(self, stream: &CudaStream) -> Self {
102        Self { owner: self.owner.for_reader(stream.clone()), column_identities: self.column_identities }
103    }
104
105    fn retain_for(&self, output_identities: &HashSet<ColumnMemoryIdentity>) -> Option<Self> {
106        let column_identities: HashSet<_> = self.column_identities.intersection(output_identities).copied().collect();
107        (!column_identities.is_empty()).then(|| Self { owner: self.owner.clone(), column_identities })
108    }
109}
110
111impl GpuDataFrame {
112    /// Wrap a cuDF frame and record its current execution placement.
113    pub fn new(frame: DataFrame) -> Self {
114        Self::new_with_resident_owner(frame, None)
115    }
116
117    /// Wrap a derived cuDF frame with identity-aware live-byte accounting and
118    /// the source's transitive resident and cache eviction-safety lifetimes.
119    pub fn try_new_with_source_residency(frame: DataFrame, source: &Self) -> Result<Self> {
120        if source.resident_owner.is_none()
121            && source.retained_source_owners.is_empty()
122            && source.source_cache_pin.is_none()
123        {
124            return Ok(Self::new(frame));
125        }
126        let source_identities = source_column_memory_identities(&source.frame);
127        let output_identity = frame.memory_identity();
128        let output_identities = output_identity.columns.iter().map(|column| column.identity).collect::<HashSet<_>>();
129        let aliased_output_columns =
130            output_identity.columns.iter().filter(|column| source_identities.contains(&column.identity)).count();
131        let aliases_source = aliased_output_columns > 0;
132        let all_output_columns_alias_source = aliases_source && aliased_output_columns == output_identity.columns.len();
133        let mut retained_source_owners = source
134            .retained_source_owners
135            .iter()
136            .filter_map(|owner| owner.retain_for(&output_identities))
137            .collect::<Vec<_>>();
138        let resident_owner = match &source.resident_owner {
139            Some(source_owner) if all_output_columns_alias_source => {
140                // Pure aliases, including column subsets, keep the source
141                // owner alive without remeasuring a partial owner set.
142                Some(source_owner.clone())
143            },
144            // A frozen grant accepts no further owner charge, so the derived
145            // frame keeps only the alias retention. Its fresh columns are
146            // accounted by the consuming attempt's own allocator.
147            Some(source_owner) if source_owner.is_frozen_resident() => {
148                if aliases_source {
149                    retained_source_owners.push(RetainedSourceOwner {
150                        owner: source_owner.clone(),
151                        column_identities: output_identities.intersection(&source_identities).copied().collect(),
152                    });
153                }
154                None
155            },
156            Some(source_owner) => {
157                // Mixed alias/fresh outputs need one retention handle to cover the
158                // whole escaping frame. This intentionally over-accounts aliased
159                // columns while the input is alive, but avoids under-accounting
160                // alias owners after the input wrapper drops.
161                let observed_bytes = Self::new(frame.clone()).device_allocated_size_bytes()?;
162                source_owner.retain_derived_owner(observed_bytes)?
163            },
164            None => None,
165        };
166        let mut output = Self::new_with_residency(frame, resident_owner, source.source_cache_pin.clone());
167        output.retained_source_owners = retained_source_owners;
168        Ok(output)
169    }
170
171    /// Marks every retained owner as read on `stream`, so each owner's final
172    /// release waits for the work `stream` has enqueued when this frame and
173    /// its derivatives drop.
174    fn read_on_stream(&mut self, stream: &CudaStream) {
175        self.resident_owner = self.resident_owner.take().map(|owner| owner.for_reader(stream.clone()));
176        self.retained_source_owners = mem::take(&mut self.retained_source_owners)
177            .into_iter()
178            .map(|retained| retained.for_reader(stream))
179            .collect();
180    }
181
182    fn new_with_resident_owner(frame: DataFrame, resident_owner: Option<ResidentOwnerLease>) -> Self {
183        Self::new_with_residency(frame, resident_owner, None)
184    }
185
186    fn new_with_residency(
187        frame: DataFrame,
188        resident_owner: Option<ResidentOwnerLease>,
189        source_cache_pin: Option<GpuLocalParquetSourceCachePin>,
190    ) -> Self {
191        let placement = frame.placement();
192        Self {
193            frame,
194            placement,
195            resident_owner,
196            retained_source_owners: Vec::new(),
197            retained_normalization_source: None,
198            source_cache_pin,
199            graph_residency: None,
200            producer_execution_scope: None,
201        }
202    }
203
204    /// Borrow the underlying cuDF frame.
205    pub fn frame(&self) -> &DataFrame {
206        &self.frame
207    }
208
209    /// Consume the wrapper and return the underlying cuDF frame.
210    pub fn into_frame(self) -> DataFrame {
211        let (frame, _) = self.into_frame_and_resident_owner();
212        frame
213    }
214
215    fn into_frame_and_resident_owner(self) -> (DataFrame, Option<ResidentOwnerLease>) {
216        let GpuDataFrameParts { frame, resident_owner, .. } = self.into_frame_and_residency();
217        (frame, resident_owner)
218    }
219
220    fn into_frame_and_residency(self) -> GpuDataFrameParts {
221        debug_assert_eq!(self.placement, self.frame.placement());
222        GpuDataFrameParts {
223            frame: self.frame,
224            resident_owner: self.resident_owner,
225            retained_source_owners: self.retained_source_owners,
226            source_cache_pin: self.source_cache_pin,
227            graph_residency: self.graph_residency,
228            producer_execution_scope: self.producer_execution_scope,
229        }
230    }
231
232    /// Attach graph source provenance to this GPU frame.
233    pub fn with_graph_residency(mut self, graph_residency: GraphResidencyDescriptor) -> Self {
234        self.graph_residency = Some(Arc::new(graph_residency));
235        self
236    }
237
238    /// Graph source provenance carried by this frame, when known.
239    pub fn graph_residency(&self) -> Option<&GraphResidencyDescriptor> {
240        self.graph_residency.as_deref()
241    }
242
243    #[doc(hidden)]
244    pub fn with_producer_execution_scope(mut self, scope: PlanFingerprint) -> Self {
245        self.producer_execution_scope = Some(scope);
246        self
247    }
248
249    /// Execution placement observed when the frame entered the engine.
250    pub fn placement(&self) -> ExecutionPlacement {
251        self.placement
252    }
253
254    /// The CUDA device of the frame.
255    pub fn device(&self) -> DeviceId {
256        self.placement.device
257    }
258
259    /// Current row count reported by the cuDF frame.
260    pub fn num_rows(&self) -> usize {
261        self.frame.num_rows()
262    }
263
264    /// Current column count reported by the cuDF frame.
265    pub fn num_columns(&self) -> usize {
266        self.frame.num_columns()
267    }
268
269    /// Observed device allocation for the frame in bytes.
270    pub fn device_allocated_size_bytes(&self) -> Result<u64> {
271        let info = self.frame.info_with(InfoDetail::DEVICE_ALLOCATED).map_err(|error| {
272            cudf_query_error(
273                "query_engine.gpu_dataframe.device_memory",
274                "failed to observe native GPU dataframe memory",
275                error,
276            )
277        })?;
278        let bytes =
279            info.memory_accounting.and_then(|accounting| accounting.device_allocated_bytes).ok_or_else(|| {
280                QueryError::new(
281                    QueryErrorCode::Internal,
282                    "query_engine.gpu_dataframe.device_memory",
283                    "cudf dataframe memory accounting did not include device-allocated bytes",
284                )
285                .with_fact("rows", self.num_rows())
286                .with_fact("columns", self.num_columns())
287            })?;
288        Ok(bytes as u64)
289    }
290
291    /// Bytes of this frame already charged to `grant`.
292    ///
293    /// This is intentionally distinct from the frame's physical allocation:
294    /// callers use it only to avoid reserving an overlapping live input twice.
295    #[cfg(any(feature = "cugraph", test))]
296    #[cfg(test)]
297    pub(crate) fn resident_owner_bytes_for(&self, grant: Option<&MemoryGrant>) -> u64 {
298        match (self.resident_owner.as_ref(), grant) {
299            (Some(owner), Some(grant)) if owner.belongs_to(grant) => owner.bytes(),
300            _ => 0,
301        }
302    }
303
304    #[doc(hidden)]
305    pub fn take_resident_owner(&mut self) -> Option<ResidentOwnerLease> {
306        self.resident_owner.take()
307    }
308
309    #[doc(hidden)]
310    pub fn with_retained_normalization_source(mut self, owner: Option<ResidentOwnerLease>) -> Self {
311        self.retained_normalization_source = owner;
312        self
313    }
314
315    #[cfg(feature = "cugraph")]
316    #[doc(hidden)]
317    pub fn take_retained_normalization_source(&mut self) -> Option<ResidentOwnerLease> {
318        self.retained_normalization_source.take()
319    }
320}
321
322/// GPU-resident query result bound to an Arrow schema and runtime.
323///
324/// Cloning shares the underlying cuDF batch, runtime, and frozen resident
325/// leases. Same-device consumers may borrow the result as resident input; host
326/// materialization is explicit through [`to_host_record_batch`](Self::to_host_record_batch).
327#[derive(Clone)]
328pub struct GpuResultHandle {
329    batch: GpuRecordBatch,
330    runtime: Arc<cudf::Runtime>,
331    resident_owner: Option<ResidentOwnerLease>,
332    retained_source_owners: Vec<RetainedSourceOwner>,
333    source_cache_pin: Option<GpuLocalParquetSourceCachePin>,
334    graph_residency: Option<Arc<GraphResidencyDescriptor>>,
335    producer_execution: Option<Arc<ProducerExecutionIdentity>>,
336}
337
338#[derive(Debug)]
339#[doc(hidden)]
340pub struct ProducerExecutionIdentity {
341    scope: PlanFingerprint,
342    plan: Option<Arc<QueryPlan>>,
343}
344
345/// Query-neutral GPU allocation retained by the local-Parquet source cache.
346///
347/// The allocation is owned by the cache allocation lease, never by the
348/// attempt whose decoded frame was copied into the cache. The lease keeps the
349/// device-scoped cache allocation domain alive until the concrete frame has
350/// been destroyed and its residency charge can be released.
351pub(crate) struct GpuLocalParquetCachedResult {
352    batch: GpuRecordBatch,
353    allocation_lease: Arc<cudf::CacheAllocationLease>,
354}
355
356/// Bounded native GPU-to-host cursor retaining every owner required by the
357/// resident result until the final chunk is materialized or the cursor drops.
358pub struct GpuHostBatchCursor {
359    cursor: CudfHostBatchCursor,
360    output_schema: SchemaRef,
361    cancellation: ExecutionCancellation,
362    /// The result's runtime, whose stream the cursor copies on. Held so the
363    /// stream outlives the last copy and the owners below can fence on it.
364    _runtime: Arc<cudf::Runtime>,
365    _resident_owner: Option<ResidentOwnerLease>,
366    _retained_source_owners: Vec<RetainedSourceOwner>,
367    _source_cache_pin: Option<GpuLocalParquetSourceCachePin>,
368    _graph_residency: Option<Arc<GraphResidencyDescriptor>>,
369}
370
371/// Metadata-sized GPU output chunk awaiting a reserved device-to-host copy.
372#[derive(Debug)]
373pub struct GpuPreparedHostBatch {
374    prepared: CudfPreparedHostBatch,
375}
376
377impl GpuPreparedHostBatch {
378    /// Conservative final Arrow buffer capacity that must be reserved before
379    /// materializing this chunk.
380    pub fn estimated_materialization_bytes(&self) -> usize {
381        self.prepared.estimated_materialization_bytes()
382    }
383
384    pub fn num_rows(&self) -> usize {
385        self.prepared.num_rows()
386    }
387}
388
389fn host_batch_prepare_error(error: cudf::Error) -> QueryError {
390    const HARD_LIMIT_REASON: &str = "host_materialization_hard_limit_exceeded";
391    let hard_limit_details = (error.detail_value("reason") == Some(HARD_LIMIT_REASON)).then(|| {
392        [
393            ("reason", HARD_LIMIT_REASON.to_string()),
394            ("row_offset", error.detail_value("row_offset").unwrap_or("unknown").to_string()),
395            ("candidate_rows", error.detail_value("candidate_rows").unwrap_or("unknown").to_string()),
396            (
397                "estimated_materialization_bytes",
398                error.detail_value("estimated_materialization_bytes").unwrap_or("unknown").to_string(),
399            ),
400            (
401                "max_materialization_bytes",
402                error.detail_value("max_materialization_bytes").unwrap_or("unknown").to_string(),
403            ),
404        ]
405    });
406    let mut error = cudf_query_error(
407        "query_engine.host_batch_cursor.prepare",
408        "failed to size bounded native output chunk before host materialization",
409        error,
410    );
411    if let Some(details) = hard_limit_details {
412        for (key, value) in details {
413            error = error.with_fact(key, value);
414        }
415    }
416    error
417}
418
419impl GpuHostBatchCursor {
420    /// Size the next bounded host chunk without copying GPU values to host.
421    pub fn prepare_next_batch(&self) -> Result<Option<GpuPreparedHostBatch>> {
422        if !self.cursor.has_next_batch() {
423            return Ok(None);
424        }
425        self.cancellation.check("query_engine.host_batch_cursor.prepare")?;
426        self.cursor
427            .prepare_next_batch()
428            .map(|prepared| prepared.map(|prepared| GpuPreparedHostBatch { prepared }))
429            .map_err(host_batch_prepare_error)
430    }
431
432    /// Materialize one prepared chunk. Cancellation is checked after the host
433    /// reservation is acquired and immediately before the device-to-host copy.
434    pub fn materialize_prepared_batch(&mut self, prepared: GpuPreparedHostBatch) -> Result<RecordBatch> {
435        self.cancellation.check("query_engine.host_batch_cursor.materialize")?;
436        let batch = self.cursor.materialize_prepared_batch(prepared.prepared).map_err(|error| {
437            cudf_query_error(
438                "query_engine.host_batch_cursor.materialize",
439                "failed to materialize bounded native output chunk to host Arrow",
440                error,
441            )
442        })?;
443        bind_host_output_schema(batch, &self.output_schema)
444    }
445
446    /// Materialize the next bounded host chunk. Cancellation is checked before
447    /// any new cuDF slice or device-to-host copy starts.
448    #[cfg(test)]
449    pub(crate) fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
450        let Some(prepared) = self.prepare_next_batch()? else {
451            return Ok(None);
452        };
453        self.materialize_prepared_batch(prepared).map(Some)
454    }
455
456    #[cfg(test)]
457    pub(crate) fn remaining_rows(&self) -> usize {
458        self.cursor.remaining_rows()
459    }
460}
461
462impl GpuResultHandle {
463    /// Bind a GPU frame to a schema and the runtime from `context`.
464    ///
465    /// The frame must already reside on the context's device. Hidden
466    /// DataFusion order columns may be stripped before schema binding.
467    pub fn new(context: &ExecutionContext, frame: GpuDataFrame, schema: SchemaRef) -> Result<Self> {
468        Self::new_with_resident_owner(context, frame, schema, None)
469    }
470
471    #[doc(hidden)]
472    pub fn new_with_resident_owner(
473        context: &ExecutionContext,
474        frame: GpuDataFrame,
475        schema: SchemaRef,
476        resident_owner: Option<ResidentOwnerLease>,
477    ) -> Result<Self> {
478        if frame.device() != context.runtime()?.device() {
479            return Err(QueryError::new(
480                QueryErrorCode::InvalidPlan,
481                "query_engine.gpu_result_handle",
482                "native GPU result frame device does not match the execution context",
483            )
484            .with_fact("frame_device", frame.device().ordinal())
485            .with_fact("runtime_device", context.runtime()?.device().ordinal()));
486        }
487        let frame = strip_hidden_datafusion_order_columns(frame, &schema)?;
488        let GpuDataFrameParts {
489            frame,
490            resident_owner: frame_resident_owner,
491            retained_source_owners,
492            source_cache_pin,
493            graph_residency,
494            producer_execution_scope,
495        } = frame.into_frame_and_residency();
496        let resident_owner = resident_owner.or(frame_resident_owner);
497        // Binding is where the producer's stream is known. Its fence makes
498        // the owner's final release wait for the work that produced the
499        // output, and the backing pointers let the report commit attest that
500        // the frozen charge matches the registry's live allocations.
501        if let Some(owner) = &resident_owner {
502            owner.record_completion(context.runtime()?.stream())?;
503            owner.set_backing_allocation_pointers(frame.backing_allocation_pointers().map_err(|error| {
504                cudf_query_error(
505                    "query_engine.gpu_result_handle",
506                    "failed to identify the resident output's backing allocations",
507                    error,
508                )
509            })?);
510        }
511        let batch = GpuRecordBatch::new(frame, schema).map_err(|error| {
512            cudf_query_error(
513                "query_engine.gpu_result_handle",
514                "failed to bind native GPU result to its Arrow schema contract",
515                error,
516            )
517        })?;
518        Ok(Self {
519            batch,
520            runtime: context.shared_runtime()?,
521            resident_owner,
522            retained_source_owners,
523            source_cache_pin,
524            graph_residency,
525            producer_execution: producer_execution_scope
526                .map(|scope| Arc::new(ProducerExecutionIdentity { scope, plan: None })),
527        })
528    }
529
530    pub(crate) fn retained_owner_identity_for(&self, grant: &MemoryGrant) -> Option<(usize, u64)> {
531        self.resident_owner
532            .as_ref()
533            .filter(|owner| owner.belongs_to(grant))
534            .map(|owner| (owner.identity_key(), owner.bytes()))
535    }
536
537    /// Borrow the underlying GPU cuDF frame.
538    pub fn frame(&self) -> &DataFrame {
539        self.batch.frame()
540    }
541
542    /// Borrows this result as a frame read on `context`'s stream. Once that
543    /// frame and every frame derived from it have dropped, the resident
544    /// owner's final release also waits for the work the stream had enqueued
545    /// by then.
546    pub fn read_on(&self, context: &ExecutionContext) -> Result<GpuDataFrame> {
547        let stream = context.runtime()?.stream().clone();
548        let mut frame = self.clone().into_gpu_frame();
549        frame.read_on_stream(&stream);
550        Ok(frame)
551    }
552
553    /// Consume the handle and return its GPU frame wrapper.
554    pub fn into_gpu_frame(self) -> GpuDataFrame {
555        let Self {
556            batch,
557            runtime: _,
558            resident_owner,
559            retained_source_owners,
560            source_cache_pin,
561            graph_residency,
562            producer_execution,
563        } = self;
564        let mut frame = GpuDataFrame::new_with_residency(batch.into_frame(), resident_owner, source_cache_pin);
565        frame.retained_source_owners = retained_source_owners;
566        frame.graph_residency = graph_residency;
567        frame.producer_execution_scope = producer_execution.map(|identity| identity.scope);
568        frame
569    }
570
571    /// Arrow schema bound to this result.
572    pub fn schema(&self) -> &SchemaRef {
573        self.batch.schema()
574    }
575
576    /// Shared cuDF runtime that created the handle.
577    pub fn runtime(&self) -> &Arc<cudf::Runtime> {
578        &self.runtime
579    }
580
581    /// Observed physical device bytes for this result handle.
582    ///
583    /// Source-cache promotion uses this as its pre-reservation upper bound
584    /// before creating any cache-domain allocation.
585    pub(crate) fn device_allocated_size_bytes(&self) -> Result<u64> {
586        let info = self.batch.frame().info_with(InfoDetail::DEVICE_ALLOCATED).map_err(|error| {
587            cudf_query_error(
588                "query_engine.gpu_result_handle.device_memory",
589                "failed to observe native GPU result device memory",
590                error,
591            )
592        })?;
593        let bytes =
594            info.memory_accounting.and_then(|accounting| accounting.device_allocated_bytes).ok_or_else(|| {
595                QueryError::new(
596                    QueryErrorCode::Internal,
597                    "query_engine.gpu_result_handle.device_memory",
598                    "GPU result memory accounting did not include device-allocated bytes",
599                )
600            })?;
601        Ok(bytes as u64)
602    }
603
604    /// Graph source provenance carried by this handle, when known.
605    pub fn graph_residency(&self) -> Option<&GraphResidencyDescriptor> {
606        self.graph_residency.as_deref()
607    }
608
609    #[doc(hidden)]
610    pub fn with_producer_execution_plan(mut self, plan: QueryPlan) -> Self {
611        let scope = plan.execution_fingerprint();
612        self.producer_execution = Some(Arc::new(ProducerExecutionIdentity { scope, plan: Some(Arc::new(plan)) }));
613        self
614    }
615
616    /// Returns the opaque typed producer identity for an ownership transfer.
617    #[doc(hidden)]
618    pub fn producer_execution_identity(&self) -> Option<Arc<ProducerExecutionIdentity>> {
619        self.producer_execution.clone()
620    }
621
622    /// Attaches an opaque typed producer identity retained before a source
623    /// handle was consumed by a host transfer.
624    #[doc(hidden)]
625    pub fn with_producer_execution_identity(mut self, identity: Option<Arc<ProducerExecutionIdentity>>) -> Self {
626        self.producer_execution = identity;
627        self
628    }
629
630    pub(crate) fn producer_execution_scope(&self) -> Option<PlanFingerprint> {
631        self.producer_execution.as_ref().map(|identity| identity.scope)
632    }
633
634    pub(crate) fn producer_execution_plan(&self) -> Option<&QueryPlan> {
635        self.producer_execution.as_ref().and_then(|identity| identity.plan.as_deref())
636    }
637
638    /// Copies this attempt-owned source frame into the supplied cache-domain
639    /// lease. The source handle is borrowed and remains owned by the attempt
640    /// on every success, cancellation, or allocation-failure path.
641    pub(crate) fn copy_for_local_parquet_cache(
642        &self,
643        allocation_lease: Arc<cudf::CacheAllocationLease>,
644    ) -> Result<GpuLocalParquetCachedResult> {
645        const OPERATION: &str = "query_engine.gpu_local_parquet_source_cache.copy_for_cache";
646        if self.source_cache_pin.is_some() || self.graph_residency.is_some() {
647            return Err(QueryError::new(
648                QueryErrorCode::InvalidPlan,
649                OPERATION,
650                "local Parquet cache input already carries non-attempt cache residency",
651            )
652            .with_fact("reason", "non_attempt_source_residency_in_cache_input"));
653        }
654        if self.runtime.device() != allocation_lease.domain().device() {
655            return Err(QueryError::new(
656                QueryErrorCode::InvalidPlan,
657                OPERATION,
658                "local Parquet cache allocation domain is on another CUDA device",
659            )
660            .with_fact("source_device", self.runtime.device().ordinal())
661            .with_fact("cache_domain_device", allocation_lease.domain().device().ordinal()));
662        }
663        if self.runtime.binding_id() == allocation_lease.runtime().binding_id() {
664            return Err(QueryError::new(
665                QueryErrorCode::InvalidPlan,
666                OPERATION,
667                "local Parquet cache allocation lease must use an independent execution binding",
668            )
669            .with_fact("reason", "local_parquet_cache_runtime_binding_not_independent"));
670        }
671        let frame = self.batch.frame().copy_into_runtime(allocation_lease.runtime()).map_err(|error| {
672            cudf_query_error(OPERATION, "failed to copy attempt-owned local Parquet source into cache ownership", error)
673        })?;
674        let batch = GpuRecordBatch::new(frame, Arc::clone(self.schema())).map_err(|error| {
675            cudf_query_error(OPERATION, "failed to bind query-neutral local Parquet cache value", error)
676        })?;
677        Ok(GpuLocalParquetCachedResult { batch, allocation_lease })
678    }
679
680    /// Current cuDF execution placement for the underlying batch.
681    pub fn placement(&self) -> ExecutionPlacement {
682        self.batch.placement()
683    }
684
685    /// Row count in the GPU result.
686    pub fn num_rows(&self) -> usize {
687        self.batch.num_rows()
688    }
689
690    /// Column count in the GPU result.
691    pub fn num_columns(&self) -> usize {
692        self.batch.num_columns()
693    }
694
695    /// Materialize this GPU result to host Arrow.
696    pub fn to_host_record_batch(&self) -> Result<RecordBatch> {
697        self.batch.to_host_record_batch().map_err(|error| {
698            cudf_query_error(
699                "query_engine.host_arrow.materialize",
700                "failed to materialize algeon native output to host Arrow",
701                error,
702            )
703        })
704    }
705
706    /// Consume this result into a bounded incremental host materializer.
707    pub fn into_host_batch_cursor(
708        self,
709        cancellation: ExecutionCancellation,
710        config: HostBatchCursorConfig,
711    ) -> Result<GpuHostBatchCursor> {
712        let output_schema = Arc::clone(self.schema());
713        self.into_host_batch_cursor_with_bound_schema(output_schema, cancellation, config)
714    }
715
716    /// Consume this result into a bounded host materializer using an external
717    /// schema whose field names, nullability, and metadata are bound by ordinal.
718    pub fn into_host_batch_cursor_with_host_schema(
719        self,
720        output_schema: SchemaRef,
721        cancellation: ExecutionCancellation,
722        config: HostBatchCursorConfig,
723    ) -> Result<GpuHostBatchCursor> {
724        validate_host_output_schema_binding(self.schema(), &output_schema)?;
725        self.into_host_batch_cursor_with_bound_schema(output_schema, cancellation, config)
726    }
727
728    fn into_host_batch_cursor_with_bound_schema(
729        self,
730        output_schema: SchemaRef,
731        cancellation: ExecutionCancellation,
732        config: HostBatchCursorConfig,
733    ) -> Result<GpuHostBatchCursor> {
734        let Self {
735            batch,
736            runtime,
737            resident_owner,
738            retained_source_owners,
739            source_cache_pin,
740            graph_residency,
741            producer_execution: _,
742        } = self;
743        // The cursor gathers its batches on the result's stream after the
744        // owner bound it.
745        let stream = runtime.stream().clone();
746        let resident_owner = resident_owner.map(|owner| owner.for_reader(stream.clone()));
747        let retained_source_owners =
748            retained_source_owners.into_iter().map(|retained| retained.for_reader(&stream)).collect();
749        let cursor = batch.into_host_batch_cursor(config).map_err(|error| {
750            cudf_query_error(
751                "query_engine.host_batch_cursor.create",
752                "failed to create bounded native host-output cursor",
753                error,
754            )
755        })?;
756        Ok(GpuHostBatchCursor {
757            cursor,
758            output_schema,
759            cancellation,
760            _runtime: runtime,
761            _resident_owner: resident_owner,
762            _retained_source_owners: retained_source_owners,
763            _source_cache_pin: source_cache_pin,
764            _graph_residency: graph_residency,
765        })
766    }
767}
768
769fn validate_host_output_schema_binding(internal_schema: &SchemaRef, output_schema: &SchemaRef) -> Result<()> {
770    const OPERATION: &str = "query_engine.host_output_schema_binding";
771    if internal_schema.fields().len() != output_schema.fields().len() {
772        return Err(QueryError::new(
773            QueryErrorCode::InvalidPlan,
774            OPERATION,
775            "internal executable schema and external host schema have different field counts",
776        )
777        .with_fact("reason", "host_output_field_count_mismatch")
778        .with_fact("internal_field_count", internal_schema.fields().len())
779        .with_fact("external_field_count", output_schema.fields().len()));
780    }
781
782    for (ordinal, (internal, external)) in
783        internal_schema.fields().iter().zip(output_schema.fields().iter()).enumerate()
784    {
785        if internal.data_type() != external.data_type() {
786            return Err(QueryError::new(
787                QueryErrorCode::InvalidPlan,
788                OPERATION,
789                "internal executable field and external host field have different Arrow types",
790            )
791            .with_fact("reason", "host_output_field_type_mismatch")
792            .with_fact("ordinal", ordinal)
793            .with_fact("internal_field", internal.name())
794            .with_fact("external_field", external.name())
795            .with_fact("internal_arrow_type", format!("{:?}", internal.data_type()))
796            .with_fact("external_arrow_type", format!("{:?}", external.data_type())));
797        }
798    }
799    Ok(())
800}
801
802fn bind_host_output_schema(batch: RecordBatch, output_schema: &SchemaRef) -> Result<RecordBatch> {
803    let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
804    RecordBatch::try_new_with_options(Arc::clone(output_schema), batch.columns().to_vec(), &options).map_err(|error| {
805        QueryError::new(
806            QueryErrorCode::Internal,
807            "query_engine.host_output_schema_binding",
808            "failed to apply the validated external host schema to a materialized native output batch",
809        )
810        .with_fact("reason", "host_output_schema_projection_failed")
811        .with_source(error)
812    })
813}
814
815impl GpuLocalParquetCachedResult {
816    pub(crate) fn schema(&self) -> &SchemaRef {
817        self.batch.schema()
818    }
819
820    /// Completes work submitted on the cache lease before the cached value is
821    /// destroyed. Cache code must call this before releasing its ledger charge.
822    pub(crate) fn synchronize_for_destruction(&self) -> Result<()> {
823        self.allocation_lease.runtime().synchronize().map_err(|error| {
824            cudf_query_error(
825                "query_engine.gpu_local_parquet_source_cache.destroy",
826                "failed to synchronize the cache allocation lease before source-cache destruction",
827                error,
828            )
829        })
830    }
831
832    pub(crate) fn device_allocated_size_bytes(&self) -> Result<u64> {
833        let info = self.batch.frame().info_with(InfoDetail::DEVICE_ALLOCATED).map_err(|error| {
834            cudf_query_error(
835                "query_engine.gpu_local_parquet_source_cache.device_memory",
836                "failed to inspect cached local Parquet device memory",
837                error,
838            )
839        })?;
840        let bytes =
841            info.memory_accounting.and_then(|accounting| accounting.device_allocated_bytes).ok_or_else(|| {
842                QueryError::new(
843                    QueryErrorCode::Internal,
844                    "query_engine.gpu_local_parquet_source_cache.device_memory",
845                    "cached local Parquet memory accounting did not include device-allocated bytes",
846                )
847            })?;
848        Ok(bytes as u64)
849    }
850
851    /// Copies the cached value into `context`'s runtime, so the attempt
852    /// computes on memory its own allocation domain owns and the cache can
853    /// account its bytes independently of any attempt. The returned handle
854    /// holds `pin` for as long as the copied frame exists.
855    pub(crate) fn checkout(
856        &self,
857        context: &ExecutionContext,
858        pin: GpuLocalParquetSourceCachePin,
859    ) -> Result<GpuResultHandle> {
860        const OPERATION: &str = "query_engine.gpu_local_parquet_source_cache.checkout";
861        let query_runtime = context.shared_runtime()?;
862        if query_runtime.device() != self.allocation_lease.domain().device() {
863            return Err(QueryError::new(
864                QueryErrorCode::InvalidPlan,
865                OPERATION,
866                "local Parquet cache entry does not match the current query device placement",
867            )
868            .with_fact("cache_domain_device", self.allocation_lease.domain().device().ordinal())
869            .with_fact("query_device", query_runtime.device().ordinal()));
870        }
871        let frame = self.batch.frame().copy_into_runtime(&query_runtime).map_err(|error| {
872            cudf_query_error(OPERATION, "failed to copy cached local Parquet source into current attempt", error)
873        })?;
874        let batch = GpuRecordBatch::new(frame, Arc::clone(self.schema())).map_err(|error| {
875            cudf_query_error(OPERATION, "failed to bind current-query local Parquet cache checkout", error)
876        })?;
877        Ok(GpuResultHandle {
878            batch,
879            runtime: query_runtime,
880            resident_owner: None,
881            retained_source_owners: Vec::new(),
882            source_cache_pin: Some(pin),
883            graph_residency: None,
884            producer_execution: None,
885        })
886    }
887}
888
889/// Drops trailing DataFusion order columns the schema does not declare. The
890/// resident owner is kept whole because the selection aliases the same
891/// buffers; the transitive source owners are narrowed to the columns that
892/// remain.
893fn strip_hidden_datafusion_order_columns(frame: GpuDataFrame, schema: &SchemaRef) -> Result<GpuDataFrame> {
894    if frame.num_columns() == schema.fields().len() {
895        return Ok(frame);
896    }
897
898    let GpuDataFrameParts {
899        frame: raw,
900        resident_owner,
901        retained_source_owners,
902        source_cache_pin,
903        graph_residency,
904        producer_execution_scope,
905    } = frame.into_frame_and_residency();
906    let frame_names = raw.column_names();
907    let expected_names = schema.fields().iter().map(|field| field.name().as_str()).collect::<Vec<_>>();
908    let visible_prefix_matches =
909        frame_names.iter().take(expected_names.len()).copied().eq(expected_names.iter().copied());
910    let trailing_are_order_columns =
911        frame_names.iter().skip(expected_names.len()).all(|name| is_datafusion_order_column_name(name));
912    if visible_prefix_matches && trailing_are_order_columns {
913        let selected = raw.select(expected_names).map_err(|error| {
914            cudf_query_error(
915                "query_engine.gpu_result_handle",
916                "failed to strip hidden native DataFusion order columns before binding result schema",
917                error,
918            )
919        })?;
920        let selected_identities = source_column_memory_identities(&selected);
921        let retained_source_owners =
922            retained_source_owners.iter().filter_map(|owner| owner.retain_for(&selected_identities)).collect();
923        let mut frame = GpuDataFrame::new_with_residency(selected, resident_owner, source_cache_pin);
924        frame.retained_source_owners = retained_source_owners;
925        frame.graph_residency = graph_residency;
926        frame.producer_execution_scope = producer_execution_scope;
927        return Ok(frame);
928    }
929
930    let mut frame = GpuDataFrame::new_with_residency(raw, resident_owner, source_cache_pin);
931    frame.retained_source_owners = retained_source_owners;
932    frame.graph_residency = graph_residency;
933    frame.producer_execution_scope = producer_execution_scope;
934    Ok(frame)
935}
936
937/// Reusable cuDF hash-join build artifact.
938///
939/// The handle is runtime/device-affine and carries the retained hash-table byte
940/// contract needed to reject unsafe reuse. Cloning shares the underlying cuDF
941/// artifact.
942#[derive(Clone)]
943pub struct GpuHashJoinBuildHandle {
944    hash_join: DataFrameHashJoin,
945    runtime: Arc<cudf::Runtime>,
946    device_allocated_size_bytes: u64,
947}
948
949/// Reusable cuDF filtered-join build artifact.
950///
951/// Like [`GpuHashJoinBuildHandle`], this is runtime/device-affine and reports
952/// observed retained device bytes for cache accounting.
953#[derive(Clone)]
954pub struct GpuFilteredJoinBuildHandle {
955    filtered_join: DataFrameFilteredJoin,
956    runtime: Arc<cudf::Runtime>,
957    device_allocated_size_bytes: u64,
958}
959
960impl GpuFilteredJoinBuildHandle {
961    #[doc(hidden)]
962    pub fn new(
963        context: &ExecutionContext,
964        filtered_join: DataFrameFilteredJoin,
965        null_equality: JoinNullEquality,
966    ) -> Result<Self> {
967        if filtered_join.build().device() != context.runtime()?.device() {
968            return Err(QueryError::new(
969                QueryErrorCode::InvalidPlan,
970                "query_engine.gpu_filtered_join_build_handle",
971                "native GPU filtered-join build artifact device does not match the execution context",
972            )
973            .with_fact("build_device", filtered_join.build().device().ordinal())
974            .with_fact("runtime_device", context.runtime()?.device().ordinal()));
975        }
976        let device_allocated_size_bytes =
977            filtered_join.device_allocated_size_bytes().map(|bytes| bytes as u64).map_err(|error| {
978                cudf_query_error(
979                    "query_engine.gpu_filtered_join_build_handle",
980                    "failed to observe algeon native reusable filtered join memory",
981                    error,
982                )
983                .with_fact("key_count", filtered_join.build_keys().len())
984                .with_fact("null_equality", null_equality)
985            })?;
986        Ok(Self { filtered_join, runtime: context.shared_runtime()?, device_allocated_size_bytes })
987    }
988
989    #[doc(hidden)]
990    pub fn filtered_join(&self) -> &DataFrameFilteredJoin {
991        &self.filtered_join
992    }
993
994    /// Shared cuDF runtime that created this reusable build artifact.
995    pub fn runtime(&self) -> &Arc<cudf::Runtime> {
996        &self.runtime
997    }
998
999    /// Current cuDF execution placement for the retained build side.
1000    pub fn placement(&self) -> ExecutionPlacement {
1001        self.filtered_join.build().placement()
1002    }
1003
1004    /// Observed retained device bytes for this artifact.
1005    pub fn device_allocated_size_bytes(&self) -> u64 {
1006        self.device_allocated_size_bytes
1007    }
1008}
1009
1010impl GpuHashJoinBuildHandle {
1011    #[doc(hidden)]
1012    pub fn new(
1013        context: &ExecutionContext,
1014        hash_join: DataFrameHashJoin,
1015        null_equality: JoinNullEquality,
1016    ) -> Result<Self> {
1017        if hash_join.build().device() != context.runtime()?.device() {
1018            return Err(QueryError::new(
1019                QueryErrorCode::InvalidPlan,
1020                "query_engine.gpu_hash_join_build_handle",
1021                "native GPU hash-join build artifact device does not match the execution context",
1022            )
1023            .with_fact("build_device", hash_join.build().device().ordinal())
1024            .with_fact("runtime_device", context.runtime()?.device().ordinal()));
1025        }
1026        let device_allocated_size_bytes =
1027            hash_join.device_allocated_size_bytes().map(|bytes| bytes as u64).map_err(|error| {
1028                cudf_query_error(
1029                    "query_engine.gpu_hash_join_build_handle",
1030                    "failed to observe algeon native reusable hash join memory",
1031                    error,
1032                )
1033                .with_fact("key_count", hash_join.build_keys().len())
1034                .with_fact("null_equality", null_equality)
1035            })?;
1036        Ok(Self { hash_join, runtime: context.shared_runtime()?, device_allocated_size_bytes })
1037    }
1038
1039    #[doc(hidden)]
1040    pub fn hash_join(&self) -> &DataFrameHashJoin {
1041        &self.hash_join
1042    }
1043
1044    /// Shared cuDF runtime that created this reusable build artifact.
1045    pub fn runtime(&self) -> &Arc<cudf::Runtime> {
1046        &self.runtime
1047    }
1048
1049    /// Current cuDF execution placement for the retained build side.
1050    pub fn placement(&self) -> ExecutionPlacement {
1051        self.hash_join.build().placement()
1052    }
1053
1054    /// Observed retained device bytes for this artifact.
1055    pub fn device_allocated_size_bytes(&self) -> u64 {
1056        self.device_allocated_size_bytes
1057    }
1058}
1059
1060fn source_column_memory_identities(frame: &DataFrame) -> HashSet<ColumnMemoryIdentity> {
1061    frame.memory_identity().columns.into_iter().map(|column| column.identity).collect()
1062}
1063
1064#[cfg(test)]
1065mod tests {
1066    use std::{collections::HashMap, sync::Arc};
1067
1068    use arrow::{
1069        array::{Array, Int32Array, StringArray, UInt64Array},
1070        datatypes::{DataType, Field, Schema},
1071        record_batch::RecordBatch,
1072    };
1073    use cudf::{DataFrame, Exec, ops::expr::col, testing::run_named_cases};
1074    use query_planning::error::{QueryErrorCode, QueryErrorKind, QueryErrorStatus};
1075
1076    use super::*;
1077    use crate::{
1078        observability::AdmissionRequestId,
1079        runtime::{
1080            StandaloneConfig,
1081            gpu::ledger::{DeviceLedger, MemoryGrant},
1082        },
1083    };
1084
1085    fn single_i32_frame(runtime: &Arc<cudf::Runtime>, name: &str, values: Vec<i32>) -> DataFrame {
1086        let schema = Arc::new(Schema::new(vec![Field::new(name, DataType::Int32, false)]));
1087        let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values))]).unwrap();
1088        DataFrame::from_record_batch(runtime, &batch).unwrap()
1089    }
1090
1091    fn frame_with_resident_owner(
1092        runtime: &Arc<cudf::Runtime>,
1093        grant: &MemoryGrant,
1094        name: &str,
1095        values: Vec<i32>,
1096    ) -> (GpuDataFrame, u64) {
1097        let frame = single_i32_frame(runtime, name, values);
1098        let observed_bytes = GpuDataFrame::new(frame.clone()).device_allocated_size_bytes().unwrap();
1099        let resident_owner = grant.retain_resident_owner(observed_bytes).expect("test frame owner is recorded");
1100        (GpuDataFrame::new_with_resident_owner(frame, resident_owner), observed_bytes)
1101    }
1102
1103    fn gpu_frame_counts_resident_bytes_only_for_the_owning_grant() {
1104        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1105        let governor = DeviceLedger::for_test(runtime.device().ordinal(), 2048, 4);
1106        let owner = governor
1107            .admit_with_configured_floor_and_soft_working_set_hint_for_test(AdmissionRequestId::new(1), 1, 1024)
1108            .unwrap()
1109            .expect("owner grant must fit");
1110        let foreign = governor
1111            .admit_with_configured_floor_and_soft_working_set_hint_for_test(AdmissionRequestId::new(2), 1, 1024)
1112            .unwrap()
1113            .expect("foreign grant must fit");
1114        let (frame, bytes) = frame_with_resident_owner(&runtime, &owner, "id", vec![1, 2, 3]);
1115
1116        assert_eq!(frame.resident_owner_bytes_for(Some(&owner)), bytes);
1117        assert_eq!(frame.resident_owner_bytes_for(Some(&foreign)), 0);
1118        assert_eq!(frame.resident_owner_bytes_for(None), 0);
1119    }
1120
1121    #[test]
1122    fn resident_cases() {
1123        run_named_cases(&[
1124            (
1125                "runtime::resident::tests::gpu_frame_counts_resident_bytes_only_for_the_owning_grant",
1126                gpu_frame_counts_resident_bytes_only_for_the_owning_grant,
1127            ),
1128            (
1129                "runtime::resident::tests::gpu_result_handle_retained_bytes_require_the_current_grant",
1130                gpu_result_handle_retained_bytes_require_the_current_grant,
1131            ),
1132            (
1133                "runtime::resident::tests::host_batch_cursor_observes_cancellation_before_starting_the_next_device_to_host_copy",
1134                host_batch_cursor_observes_cancellation_before_starting_the_next_device_to_host_copy,
1135            ),
1136            (
1137                "runtime::resident::tests::host_output_schema_binding_preserves_duplicate_names_metadata_and_values",
1138                host_output_schema_binding_preserves_duplicate_names_metadata_and_values,
1139            ),
1140            (
1141                "runtime::resident::tests::host_batch_cursor_preserves_structured_hard_limit_facts",
1142                host_batch_cursor_preserves_structured_hard_limit_facts,
1143            ),
1144            (
1145                "runtime::resident::tests::host_batch_cursor_reports_clean_eof_when_cancelled_after_final_chunk",
1146                host_batch_cursor_reports_clean_eof_when_cancelled_after_final_chunk,
1147            ),
1148            (
1149                "runtime::resident::tests::host_batch_cursor_prepares_each_chunk_without_early_success",
1150                host_batch_cursor_prepares_each_chunk_without_early_success,
1151            ),
1152            (
1153                "runtime::resident::tests::host_batch_cursor_prepare_reports_cancellation_before_any_copy",
1154                host_batch_cursor_prepare_reports_cancellation_before_any_copy,
1155            ),
1156            (
1157                "runtime::resident::tests::source_residency_clones_owner_for_pure_alias_output",
1158                source_residency_clones_owner_for_pure_alias_output,
1159            ),
1160            (
1161                "runtime::resident::tests::source_residency_records_derived_owner_for_fresh_output",
1162                source_residency_records_derived_owner_for_fresh_output,
1163            ),
1164            (
1165                "runtime::resident::tests::source_residency_retains_full_output_for_mixed_alias_output",
1166                source_residency_retains_full_output_for_mixed_alias_output,
1167            ),
1168            (
1169                "runtime::resident::tests::source_residency_preserves_transitive_frozen_owner_across_alias_chain",
1170                source_residency_preserves_transitive_frozen_owner_across_alias_chain,
1171            ),
1172            (
1173                "runtime::resident::tests::source_residency_releases_transitive_owner_when_only_fresh_columns_survive",
1174                source_residency_releases_transitive_owner_when_only_fresh_columns_survive,
1175            ),
1176        ]);
1177    }
1178
1179    fn retained_guard_keeps_gpu_result_handle_send_safe() {
1180        fn assert_send_sync<T: Send + Sync>() {}
1181
1182        assert_send_sync::<GpuResultHandle>();
1183    }
1184
1185    const _: fn() = retained_guard_keeps_gpu_result_handle_send_safe;
1186
1187    fn gpu_result_handle_retained_bytes_require_the_current_grant() {
1188        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1189        let governor = DeviceLedger::for_test(runtime.device().ordinal(), 2048, 4);
1190        let owner = governor
1191            .admit_with_configured_floor_and_soft_working_set_hint_for_test(AdmissionRequestId::new(11), 1, 1024)
1192            .unwrap()
1193            .expect("owner grant must fit");
1194        let foreign = governor
1195            .admit_with_configured_floor_and_soft_working_set_hint_for_test(AdmissionRequestId::new(12), 1, 1024)
1196            .unwrap()
1197            .expect("foreign grant must fit");
1198        let (frame, bytes) = frame_with_resident_owner(&runtime, &owner, "id", vec![1, 2, 3]);
1199        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1200        let context = ExecutionContext::standalone(StandaloneConfig {
1201            runtime: Some(Arc::clone(&runtime)),
1202            ..StandaloneConfig::default()
1203        })
1204        .unwrap();
1205        let handle = GpuResultHandle::new(&context, frame, Arc::clone(&schema)).unwrap();
1206
1207        assert_eq!(handle.retained_owner_identity_for(&owner).map(|(_, value)| value), Some(bytes));
1208        assert_eq!(handle.retained_owner_identity_for(&foreign), None);
1209
1210        let fresh =
1211            GpuResultHandle::new(&context, GpuDataFrame::new(single_i32_frame(&runtime, "id", vec![4, 5, 6])), schema)
1212                .unwrap();
1213        assert_eq!(fresh.retained_owner_identity_for(&owner), None);
1214        drop(fresh);
1215        assert_eq!(owner.resident_owner_bytes(), bytes);
1216        drop(handle);
1217        assert_eq!(owner.resident_owner_bytes(), 0);
1218    }
1219
1220    fn host_batch_cursor_observes_cancellation_before_starting_the_next_device_to_host_copy() {
1221        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1222        let cancellation = ExecutionCancellation::new();
1223        let context = ExecutionContext::standalone(StandaloneConfig {
1224            runtime: Some(Arc::clone(&runtime)),
1225            cancellation: cancellation.clone(),
1226            ..StandaloneConfig::default()
1227        })
1228        .unwrap();
1229        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1230        let frame = GpuDataFrame::new(single_i32_frame(&runtime, "id", vec![1, 2, 3]));
1231        let handle = GpuResultHandle::new(&context, frame, schema).unwrap();
1232        let mut cursor = handle
1233            .into_host_batch_cursor(
1234                cancellation.clone(),
1235                HostBatchCursorConfig { target_materialization_bytes: 1024, max_rows: 1, ..Default::default() },
1236            )
1237            .unwrap();
1238
1239        assert_eq!(cursor.next_batch().unwrap().unwrap().num_rows(), 1);
1240        cancellation.cancel();
1241        let error = cursor.next_batch().expect_err("cancellation must stop the next host chunk");
1242        assert_eq!(error.code(), QueryErrorCode::Cancelled);
1243        assert_eq!(cursor.remaining_rows(), 2);
1244    }
1245
1246    fn host_output_schema_binding_preserves_duplicate_names_metadata_and_values() {
1247        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1248        let context = ExecutionContext::standalone(StandaloneConfig {
1249            runtime: Some(Arc::clone(&runtime)),
1250            ..StandaloneConfig::default()
1251        })
1252        .unwrap();
1253        let internal_schema = Arc::new(Schema::new(vec![
1254            Field::new("left_first", DataType::Utf8, true),
1255            Field::new("right_first", DataType::Utf8, true),
1256            Field::new("shared_targets_internal", DataType::UInt64, false),
1257        ]));
1258        let input = RecordBatch::try_new(
1259            Arc::clone(&internal_schema),
1260            vec![
1261                Arc::new(StringArray::from(vec![Some("a"), None])) as Arc<dyn Array>,
1262                Arc::new(StringArray::from(vec![Some("b"), Some("c")])) as Arc<dyn Array>,
1263                Arc::new(UInt64Array::from(vec![2, 1])) as Arc<dyn Array>,
1264            ],
1265        )
1266        .unwrap();
1267        let field_metadata = HashMap::from([("source".to_string(), "datafusion".to_string())]);
1268        let output_schema = Arc::new(
1269            Schema::new(vec![
1270                Field::new("first", DataType::Utf8, true).with_metadata(field_metadata),
1271                Field::new("first", DataType::Utf8, false),
1272                Field::new("shared_targets", DataType::UInt64, true),
1273            ])
1274            .with_metadata(HashMap::from([("query".to_string(), "self_join".to_string())])),
1275        );
1276        let frame = GpuDataFrame::new(DataFrame::from_record_batch(&runtime, &input).unwrap());
1277        let handle = GpuResultHandle::new(&context, frame, internal_schema).unwrap();
1278        let mut cursor = handle
1279            .into_host_batch_cursor_with_host_schema(
1280                Arc::clone(&output_schema),
1281                ExecutionCancellation::new(),
1282                HostBatchCursorConfig::default(),
1283            )
1284            .unwrap();
1285
1286        let output = cursor.next_batch().unwrap().unwrap();
1287        assert_eq!(output.schema(), output_schema);
1288        assert_eq!(output.column(0).as_any().downcast_ref::<StringArray>().unwrap().value(0), "a");
1289        assert!(output.column(0).as_any().downcast_ref::<StringArray>().unwrap().is_null(1));
1290        assert_eq!(output.column(1).as_any().downcast_ref::<StringArray>().unwrap().value(1), "c");
1291        assert_eq!(output.column(2).as_any().downcast_ref::<UInt64Array>().unwrap().values(), &[2, 1]);
1292    }
1293
1294    #[test]
1295    fn host_output_schema_binding_rejects_count_and_type_mismatches_with_ordinal_facts() {
1296        let internal = Arc::new(Schema::new(vec![
1297            Field::new("left_first", DataType::Utf8, true),
1298            Field::new("shared_targets_internal", DataType::UInt64, false),
1299        ]));
1300        let short = Arc::new(Schema::new(vec![Field::new("first", DataType::Utf8, true)]));
1301        let count_error = validate_host_output_schema_binding(&internal, &short).unwrap_err();
1302        assert_eq!(count_error.code(), QueryErrorCode::InvalidPlan);
1303        assert_eq!(count_error.fact_value("reason"), Some("host_output_field_count_mismatch"));
1304        assert_eq!(count_error.fact_value("internal_field_count"), Some("2"));
1305        assert_eq!(count_error.fact_value("external_field_count"), Some("1"));
1306
1307        let reordered = Arc::new(Schema::new(vec![
1308            Field::new("shared_targets", DataType::UInt64, false),
1309            Field::new("first", DataType::Utf8, true),
1310        ]));
1311        let type_error = validate_host_output_schema_binding(&internal, &reordered).unwrap_err();
1312        assert_eq!(type_error.code(), QueryErrorCode::InvalidPlan);
1313        assert_eq!(type_error.fact_value("reason"), Some("host_output_field_type_mismatch"));
1314        assert_eq!(type_error.fact_value("ordinal"), Some("0"));
1315        assert_eq!(type_error.fact_value("internal_field"), Some("left_first"));
1316        assert_eq!(type_error.fact_value("external_field"), Some("shared_targets"));
1317        assert_eq!(type_error.fact_value("internal_arrow_type"), Some("Utf8"));
1318        assert_eq!(type_error.fact_value("external_arrow_type"), Some("UInt64"));
1319    }
1320
1321    fn host_batch_cursor_preserves_structured_hard_limit_facts() {
1322        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1323        let cancellation = ExecutionCancellation::new();
1324        let context = ExecutionContext::standalone(StandaloneConfig {
1325            runtime: Some(Arc::clone(&runtime)),
1326            ..StandaloneConfig::default()
1327        })
1328        .unwrap();
1329        let value = "x".repeat(16 * 1024 * 1024);
1330        let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, false)]));
1331        let batch = RecordBatch::try_new(
1332            Arc::clone(&schema),
1333            vec![Arc::new(StringArray::from(vec![value.as_str()])) as Arc<dyn Array>],
1334        )
1335        .unwrap();
1336        let frame = GpuDataFrame::new(DataFrame::from_record_batch(&runtime, &batch).unwrap());
1337        let handle = GpuResultHandle::new(&context, frame, schema).unwrap();
1338        let cursor = handle.into_host_batch_cursor(cancellation, HostBatchCursorConfig::default()).unwrap();
1339
1340        let error = cursor.prepare_next_batch().expect_err("oversized scalar must reject before D2H");
1341        assert_eq!(error.code(), QueryErrorCode::HostMaterializationLimitExceeded);
1342        assert_eq!(error.kind(), QueryErrorKind::ResourceExhausted);
1343        assert_eq!(error.status(), QueryErrorStatus::Permanent);
1344        assert_eq!(error.fact_value("reason"), Some("host_materialization_hard_limit_exceeded"));
1345        assert_eq!(error.fact_value("candidate_rows"), Some("1"));
1346        assert_eq!(
1347            error.fact_value("max_materialization_bytes"),
1348            Some(HostBatchCursorConfig::DEFAULT_MAX_MATERIALIZATION_BYTES.to_string().as_str())
1349        );
1350        assert!(
1351            error.fact_value("estimated_materialization_bytes").unwrap().parse::<usize>().unwrap()
1352                > HostBatchCursorConfig::DEFAULT_MAX_MATERIALIZATION_BYTES
1353        );
1354    }
1355
1356    fn host_batch_cursor_reports_clean_eof_when_cancelled_after_final_chunk() {
1357        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1358        let cancellation = ExecutionCancellation::new();
1359        let context = ExecutionContext::standalone(StandaloneConfig {
1360            runtime: Some(Arc::clone(&runtime)),
1361            cancellation: cancellation.clone(),
1362            ..StandaloneConfig::default()
1363        })
1364        .unwrap();
1365        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1366        let frame = GpuDataFrame::new(single_i32_frame(&runtime, "id", vec![1]));
1367        let handle = GpuResultHandle::new(&context, frame, schema).unwrap();
1368        let mut cursor = handle
1369            .into_host_batch_cursor(
1370                cancellation.clone(),
1371                HostBatchCursorConfig { target_materialization_bytes: 1024, max_rows: 1, ..Default::default() },
1372            )
1373            .unwrap();
1374
1375        assert_eq!(cursor.next_batch().unwrap().unwrap().num_rows(), 1);
1376        assert_eq!(cursor.remaining_rows(), 0);
1377        cancellation.cancel();
1378        assert!(cursor.next_batch().unwrap().is_none());
1379    }
1380
1381    /// The deferred host cursor prepares each bounded chunk inside the existing sink
1382    /// contract without declaring success early. Every prepared chunk materializes exactly
1383    /// once, the cursor reports more work until the final chunk drains, and dropping the
1384    /// cursor after a partial drain releases its retained owners and does not leak a grant.
1385    fn host_batch_cursor_prepares_each_chunk_without_early_success() {
1386        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1387        let cancellation = ExecutionCancellation::new();
1388        let context = ExecutionContext::standalone(StandaloneConfig {
1389            runtime: Some(Arc::clone(&runtime)),
1390            cancellation: cancellation.clone(),
1391            ..StandaloneConfig::default()
1392        })
1393        .unwrap();
1394        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1395        let frame = GpuDataFrame::new(single_i32_frame(&runtime, "id", vec![1, 2, 3]));
1396        let handle = GpuResultHandle::new(&context, frame, schema).unwrap();
1397        let mut cursor = handle
1398            .into_host_batch_cursor(
1399                cancellation,
1400                HostBatchCursorConfig { target_materialization_bytes: 1024, max_rows: 1, ..Default::default() },
1401            )
1402            .unwrap();
1403
1404        let mut rows = Vec::new();
1405        for expected in [1, 2, 3] {
1406            assert_eq!(cursor.remaining_rows(), 4 - expected, "unprepared chunks are not success");
1407            let prepared = cursor.prepare_next_batch().expect("chunk prepares").expect("chunk remains");
1408            assert_eq!(prepared.num_rows(), 1);
1409            let batch = cursor.materialize_prepared_batch(prepared).expect("prepared chunk materializes");
1410            assert_eq!(batch.num_rows(), 1);
1411            rows.push(batch.column(0).as_any().downcast_ref::<Int32Array>().unwrap().value(0));
1412        }
1413        assert!(cursor.prepare_next_batch().expect("final prepare observes").is_none());
1414        assert_eq!(cursor.remaining_rows(), 0);
1415        rows.sort();
1416        assert_eq!(rows, vec![1, 2, 3], "every chunk materializes exactly once");
1417    }
1418
1419    fn host_batch_cursor_prepare_reports_cancellation_before_any_copy() {
1420        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1421        let cancellation = ExecutionCancellation::new();
1422        let context = ExecutionContext::standalone(StandaloneConfig {
1423            runtime: Some(Arc::clone(&runtime)),
1424            cancellation: cancellation.clone(),
1425            ..StandaloneConfig::default()
1426        })
1427        .unwrap();
1428        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
1429        let frame = GpuDataFrame::new(single_i32_frame(&runtime, "id", vec![1, 2, 3]));
1430        let handle = GpuResultHandle::new(&context, frame, schema).unwrap();
1431        let cursor = handle
1432            .into_host_batch_cursor(
1433                cancellation.clone(),
1434                HostBatchCursorConfig { target_materialization_bytes: 1024, max_rows: 1, ..Default::default() },
1435            )
1436            .unwrap();
1437        cancellation.cancel();
1438        let error = cursor.prepare_next_batch().expect_err("cancelled prepare must not size a chunk");
1439        assert_eq!(error.code(), QueryErrorCode::Cancelled);
1440    }
1441
1442    fn source_residency_clones_owner_for_pure_alias_output() {
1443        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1444        let governor = DeviceLedger::for_test(runtime.device().ordinal(), 128 * 1024 * 1024, 4);
1445        let grant = governor
1446            .admit_with_configured_floor_and_soft_working_set_hint_for_test(
1447                AdmissionRequestId::new(1),
1448                64 * 1024 * 1024,
1449                64 * 1024 * 1024,
1450            )
1451            .unwrap()
1452            .unwrap();
1453        let (source, source_bytes) = frame_with_resident_owner(&runtime, &grant, "id", vec![1, 2, 3]);
1454        assert!(source_bytes > 0);
1455        assert_eq!(grant.resident_owner_bytes(), source_bytes);
1456
1457        let selected = source.frame().select(["id"]).unwrap();
1458        let derived = GpuDataFrame::try_new_with_source_residency(selected, &source).unwrap();
1459        assert_eq!(grant.resident_owner_bytes(), source_bytes);
1460
1461        drop(source);
1462        assert_eq!(grant.resident_owner_bytes(), source_bytes);
1463        drop(derived);
1464        assert_eq!(grant.resident_owner_bytes(), 0);
1465    }
1466
1467    fn source_residency_records_derived_owner_for_fresh_output() {
1468        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1469        let governor = DeviceLedger::for_test(runtime.device().ordinal(), 128 * 1024 * 1024, 4);
1470        let grant = governor
1471            .admit_with_configured_floor_and_soft_working_set_hint_for_test(
1472                AdmissionRequestId::new(2),
1473                64 * 1024 * 1024,
1474                64 * 1024 * 1024,
1475            )
1476            .unwrap()
1477            .unwrap();
1478        let (source, source_bytes) = frame_with_resident_owner(&runtime, &grant, "id", vec![1, 2, 3]);
1479        let fresh = single_i32_frame(&runtime, "fresh", vec![10, 20, 30]);
1480        let fresh_bytes = GpuDataFrame::new(fresh.clone()).device_allocated_size_bytes().unwrap();
1481        assert!(fresh_bytes > 0);
1482
1483        let derived = GpuDataFrame::try_new_with_source_residency(fresh, &source).unwrap();
1484        assert_eq!(grant.resident_owner_bytes(), source_bytes.saturating_add(fresh_bytes));
1485
1486        drop(source);
1487        assert_eq!(grant.resident_owner_bytes(), fresh_bytes);
1488        drop(derived);
1489        assert_eq!(grant.resident_owner_bytes(), 0);
1490    }
1491
1492    fn source_residency_retains_full_output_for_mixed_alias_output() {
1493        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1494        let governor = DeviceLedger::for_test(runtime.device().ordinal(), 128 * 1024 * 1024, 4);
1495        let grant = governor
1496            .admit_with_configured_floor_and_soft_working_set_hint_for_test(
1497                AdmissionRequestId::new(3),
1498                64 * 1024 * 1024,
1499                64 * 1024 * 1024,
1500            )
1501            .unwrap()
1502            .unwrap();
1503        let (source, source_bytes) = frame_with_resident_owner(&runtime, &grant, "id", vec![1, 2, 3]);
1504        let mixed = source.frame().on(Exec::RETAINED).with_columns([("fresh", col("id") + 1i32)]).unwrap();
1505        let fresh_only = mixed.select(["fresh"]).unwrap();
1506        let fresh_bytes = GpuDataFrame::new(fresh_only).device_allocated_size_bytes().unwrap();
1507        let mixed_bytes = GpuDataFrame::new(mixed.clone()).device_allocated_size_bytes().unwrap();
1508        assert!(fresh_bytes > 0);
1509        assert!(
1510            mixed_bytes > fresh_bytes,
1511            "mixed output should include both the aliased source column and the fresh derived column"
1512        );
1513
1514        let derived = GpuDataFrame::try_new_with_source_residency(mixed, &source).unwrap();
1515        assert_eq!(grant.resident_owner_bytes(), source_bytes.saturating_add(mixed_bytes));
1516
1517        drop(source);
1518        assert_eq!(grant.resident_owner_bytes(), mixed_bytes);
1519        drop(derived);
1520        assert_eq!(grant.resident_owner_bytes(), 0);
1521    }
1522
1523    fn source_residency_preserves_transitive_frozen_owner_across_alias_chain() {
1524        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1525        let governor = DeviceLedger::for_test(runtime.device().ordinal(), 128 * 1024 * 1024, 4);
1526        let grant = governor
1527            .admit_with_configured_floor_and_soft_working_set_hint_for_test(
1528                AdmissionRequestId::new(4),
1529                64 * 1024 * 1024,
1530                64 * 1024 * 1024,
1531            )
1532            .unwrap()
1533            .unwrap();
1534        let (source, source_bytes) = frame_with_resident_owner(&runtime, &grant, "id", vec![1, 2, 3]);
1535        source.resident_owner.as_ref().unwrap().record_completion(runtime.stream()).unwrap();
1536        grant.freeze_resident(source_bytes).unwrap();
1537
1538        let mixed = source.frame().on(Exec::RETAINED).with_columns([("fresh", col("id") + 1i32)]).unwrap();
1539        let mixed = GpuDataFrame::try_new_with_source_residency(mixed, &source).unwrap();
1540        let selected = mixed.frame().select(["id"]).unwrap();
1541        let selected = GpuDataFrame::try_new_with_source_residency(selected, &mixed).unwrap();
1542
1543        drop(source);
1544        drop(mixed);
1545        assert_eq!(grant.resident_owner_bytes(), source_bytes);
1546
1547        drop(selected);
1548        assert_eq!(grant.resident_owner_bytes(), 0);
1549    }
1550
1551    fn source_residency_releases_transitive_owner_when_only_fresh_columns_survive() {
1552        let runtime = Arc::new(cudf::Runtime::new().unwrap());
1553        let governor = DeviceLedger::for_test(runtime.device().ordinal(), 128 * 1024 * 1024, 4);
1554        let grant = governor
1555            .admit_with_configured_floor_and_soft_working_set_hint_for_test(
1556                AdmissionRequestId::new(5),
1557                64 * 1024 * 1024,
1558                64 * 1024 * 1024,
1559            )
1560            .unwrap()
1561            .unwrap();
1562        let (source, source_bytes) = frame_with_resident_owner(&runtime, &grant, "id", vec![1, 2, 3]);
1563        source.resident_owner.as_ref().unwrap().record_completion(runtime.stream()).unwrap();
1564        grant.freeze_resident(source_bytes).unwrap();
1565
1566        let mixed = source.frame().on(Exec::RETAINED).with_columns([("fresh", col("id") + 1i32)]).unwrap();
1567        let mixed = GpuDataFrame::try_new_with_source_residency(mixed, &source).unwrap();
1568        let selected = mixed.frame().select(["fresh"]).unwrap();
1569        let selected = GpuDataFrame::try_new_with_source_residency(selected, &mixed).unwrap();
1570
1571        drop(source);
1572        drop(mixed);
1573        assert_eq!(grant.resident_owner_bytes(), 0);
1574
1575        drop(selected);
1576    }
1577}