1use 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
49pub type GpuResidentInputMap = BTreeMap<String, GpuResultHandle>;
55
56#[derive(Clone)]
62pub struct GpuDataFrame {
63 frame: DataFrame,
64 placement: ExecutionPlacement,
65 resident_owner: Option<ResidentOwnerLease>,
69 retained_source_owners: Vec<RetainedSourceOwner>,
73 retained_normalization_source: Option<ResidentOwnerLease>,
76 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#[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 pub fn new(frame: DataFrame) -> Self {
114 Self::new_with_resident_owner(frame, None)
115 }
116
117 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 Some(source_owner.clone())
143 },
144 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 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 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 pub fn frame(&self) -> &DataFrame {
206 &self.frame
207 }
208
209 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 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 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 pub fn placement(&self) -> ExecutionPlacement {
251 self.placement
252 }
253
254 pub fn device(&self) -> DeviceId {
256 self.placement.device
257 }
258
259 pub fn num_rows(&self) -> usize {
261 self.frame.num_rows()
262 }
263
264 pub fn num_columns(&self) -> usize {
266 self.frame.num_columns()
267 }
268
269 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 #[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#[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
345pub(crate) struct GpuLocalParquetCachedResult {
352 batch: GpuRecordBatch,
353 allocation_lease: Arc<cudf::CacheAllocationLease>,
354}
355
356pub struct GpuHostBatchCursor {
359 cursor: CudfHostBatchCursor,
360 output_schema: SchemaRef,
361 cancellation: ExecutionCancellation,
362 _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#[derive(Debug)]
373pub struct GpuPreparedHostBatch {
374 prepared: CudfPreparedHostBatch,
375}
376
377impl GpuPreparedHostBatch {
378 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 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 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 #[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 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 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 pub fn frame(&self) -> &DataFrame {
539 self.batch.frame()
540 }
541
542 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 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 pub fn schema(&self) -> &SchemaRef {
573 self.batch.schema()
574 }
575
576 pub fn runtime(&self) -> &Arc<cudf::Runtime> {
578 &self.runtime
579 }
580
581 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 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 #[doc(hidden)]
618 pub fn producer_execution_identity(&self) -> Option<Arc<ProducerExecutionIdentity>> {
619 self.producer_execution.clone()
620 }
621
622 #[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 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 pub fn placement(&self) -> ExecutionPlacement {
682 self.batch.placement()
683 }
684
685 pub fn num_rows(&self) -> usize {
687 self.batch.num_rows()
688 }
689
690 pub fn num_columns(&self) -> usize {
692 self.batch.num_columns()
693 }
694
695 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 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 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 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 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 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
889fn 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#[derive(Clone)]
943pub struct GpuHashJoinBuildHandle {
944 hash_join: DataFrameHashJoin,
945 runtime: Arc<cudf::Runtime>,
946 device_allocated_size_bytes: u64,
947}
948
949#[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 pub fn runtime(&self) -> &Arc<cudf::Runtime> {
996 &self.runtime
997 }
998
999 pub fn placement(&self) -> ExecutionPlacement {
1001 self.filtered_join.build().placement()
1002 }
1003
1004 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 pub fn runtime(&self) -> &Arc<cudf::Runtime> {
1046 &self.runtime
1047 }
1048
1049 pub fn placement(&self) -> ExecutionPlacement {
1051 self.hash_join.build().placement()
1052 }
1053
1054 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 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}