1use std::sync::Arc;
19
20use futures::stream::BoxStream;
21use serde::{Deserialize, Serialize};
22use typed_builder::TypedBuilder;
23
24use crate::expr::BoundPredicate;
25use crate::spec::{
26 DataContentType, DataFileFormat, ManifestEntryRef, NameMapping, PartitionSpec, Schema,
27 SchemaRef, Struct, StructType,
28};
29use crate::{Error, ErrorKind, Result};
30
31pub type FileScanTaskStream = BoxStream<'static, Result<FileScanTask>>;
33
34#[derive(Debug, Clone, Deserialize, PartialEq, TypedBuilder)]
36#[serde(try_from = "crate::scan::task::_serde::FileScanTaskSerde")]
37#[builder(
38 field_defaults(setter(prefix = "with_")),
39 build_method(into = Result<FileScanTask>)
40)]
41pub struct FileScanTask {
42 file_size_in_bytes: u64,
45 start: u64,
47 length: u64,
49 #[builder(default)]
54 record_count: Option<u64>,
55
56 #[builder(default)]
61 first_row_id: Option<i64>,
62
63 #[builder(default)]
70 data_sequence_number: Option<i64>,
71
72 data_file_path: String,
74
75 data_file_format: DataFileFormat,
77
78 schema: SchemaRef,
80 project_field_ids: Vec<i32>,
82 #[builder(default)]
84 predicate: Option<BoundPredicate>,
85
86 #[builder(default)]
88 deletes: Vec<FileScanTaskDeleteFile>,
89
90 #[builder(default)]
94 partition: Option<Struct>,
95
96 #[builder(default)]
100 partition_spec: Option<Arc<PartitionSpec>>,
101
102 #[builder(default)]
106 name_mapping: Option<Arc<NameMapping>>,
107
108 #[builder(default)]
117 unified_partition_type: Option<Arc<StructType>>,
118
119 case_sensitive: bool,
121
122 #[builder(default)]
131 key_metadata: Option<Box<[u8]>>,
132}
133
134impl FileScanTask {
135 pub fn file_size_in_bytes(&self) -> u64 {
137 self.file_size_in_bytes
138 }
139
140 pub fn start(&self) -> u64 {
142 self.start
143 }
144
145 pub fn length(&self) -> u64 {
147 self.length
148 }
149
150 pub fn record_count(&self) -> Option<u64> {
152 self.record_count
153 }
154
155 pub fn first_row_id(&self) -> Option<i64> {
157 self.first_row_id
158 }
159
160 pub fn data_sequence_number(&self) -> Option<i64> {
162 self.data_sequence_number
163 }
164
165 pub fn data_file_path(&self) -> &str {
167 &self.data_file_path
168 }
169
170 pub fn data_file_format(&self) -> DataFileFormat {
172 self.data_file_format
173 }
174
175 pub fn schema(&self) -> &Schema {
177 &self.schema
178 }
179
180 pub fn schema_ref(&self) -> SchemaRef {
182 self.schema.clone()
183 }
184
185 pub fn project_field_ids(&self) -> &[i32] {
187 &self.project_field_ids
188 }
189
190 pub fn predicate(&self) -> Option<&BoundPredicate> {
192 self.predicate.as_ref()
193 }
194
195 pub fn deletes(&self) -> &[FileScanTaskDeleteFile] {
197 &self.deletes
198 }
199
200 pub fn partition(&self) -> Option<&Struct> {
202 self.partition.as_ref()
203 }
204
205 pub fn partition_spec(&self) -> Option<&Arc<PartitionSpec>> {
207 self.partition_spec.as_ref()
208 }
209
210 pub fn name_mapping(&self) -> Option<&Arc<NameMapping>> {
212 self.name_mapping.as_ref()
213 }
214
215 pub fn unified_partition_type(&self) -> Option<&Arc<StructType>> {
217 self.unified_partition_type.as_ref()
218 }
219
220 pub fn case_sensitive(&self) -> bool {
222 self.case_sensitive
223 }
224
225 pub fn key_metadata(&self) -> Option<&[u8]> {
227 self.key_metadata.as_deref()
228 }
229
230 fn validate(&self) -> Result<()> {
231 match (self.partition.as_ref(), self.partition_spec.as_deref()) {
232 (None, None) => Ok(()),
233 (None, Some(partition_spec)) if partition_spec.is_unpartitioned() => Ok(()),
234 (None, Some(_)) => Err(Error::new(
235 ErrorKind::DataInvalid,
236 "FileScanTask with a partitioned spec requires partition values",
237 )),
238 (Some(partition), None) if partition.fields().is_empty() => Ok(()),
239 (Some(_), None) => Err(Error::new(
240 ErrorKind::DataInvalid,
241 "Non-empty FileScanTask partition requires a partition spec",
242 )),
243 (Some(partition), Some(partition_spec))
244 if partition.fields().len() != partition_spec.fields().len() =>
245 {
246 Err(Error::new(
247 ErrorKind::DataInvalid,
248 format!(
249 "FileScanTask partition has {} fields but partition spec has {} fields",
250 partition.fields().len(),
251 partition_spec.fields().len()
252 ),
253 ))
254 }
255 (Some(_), Some(partition_spec)) => {
256 partition_spec.partition_type(&self.schema)?;
257 Ok(())
258 }
259 }
260 }
261}
262
263impl From<FileScanTask> for Result<FileScanTask> {
264 fn from(task: FileScanTask) -> Self {
265 task.validate()?;
266 Ok(task)
267 }
268}
269
270#[derive(Debug)]
271pub(crate) struct DeleteFileContext {
272 pub(crate) manifest_entry: ManifestEntryRef,
273 pub(crate) partition_spec_id: i32,
274}
275
276impl From<&DeleteFileContext> for FileScanTaskDeleteFile {
277 fn from(ctx: &DeleteFileContext) -> Self {
278 FileScanTaskDeleteFile::builder()
279 .with_file_path(ctx.manifest_entry.file_path().to_string())
280 .with_file_size_in_bytes(ctx.manifest_entry.file_size_in_bytes())
281 .with_file_type(ctx.manifest_entry.content_type())
282 .with_file_format(ctx.manifest_entry.data_file().file_format())
283 .with_partition_spec_id(ctx.partition_spec_id)
284 .with_equality_ids(ctx.manifest_entry.data_file.equality_ids.clone())
285 .with_referenced_data_file(ctx.manifest_entry.data_file.referenced_data_file.clone())
286 .with_content_offset(ctx.manifest_entry.data_file.content_offset)
287 .with_content_size_in_bytes(ctx.manifest_entry.data_file.content_size_in_bytes)
288 .with_record_count(Some(ctx.manifest_entry.record_count()))
289 .with_key_metadata(
290 ctx.manifest_entry
291 .data_file
292 .key_metadata
293 .as_deref()
294 .map(Box::from),
295 )
296 .build()
297 }
298}
299
300#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TypedBuilder)]
302#[builder(field_defaults(setter(prefix = "with_")))]
303pub struct FileScanTaskDeleteFile {
304 pub file_path: String,
306
307 pub file_size_in_bytes: u64,
309
310 pub file_type: DataContentType,
312
313 pub file_format: DataFileFormat,
316
317 pub partition_spec_id: i32,
319
320 #[builder(default)]
322 pub equality_ids: Option<Vec<i32>>,
323
324 #[serde(default)]
327 #[serde(skip_serializing_if = "Option::is_none")]
328 #[builder(default)]
329 pub referenced_data_file: Option<String>,
330
331 #[serde(default)]
334 #[serde(skip_serializing_if = "Option::is_none")]
335 #[builder(default)]
336 pub content_offset: Option<i64>,
337
338 #[serde(default)]
341 #[serde(skip_serializing_if = "Option::is_none")]
342 #[builder(default)]
343 pub content_size_in_bytes: Option<i64>,
344
345 #[serde(default)]
348 #[serde(skip_serializing_if = "Option::is_none")]
349 #[builder(default)]
350 pub record_count: Option<u64>,
351
352 #[serde(default)]
362 #[serde(skip_serializing_if = "Option::is_none")]
363 #[builder(default)]
364 pub key_metadata: Option<Box<[u8]>>,
365}
366
367mod _serde {
368 use std::sync::Arc;
369
370 use serde::{Deserialize, Serialize};
371
372 use super::{FileScanTask, FileScanTaskDeleteFile};
373 use crate::expr::BoundPredicate;
374 use crate::spec::{
375 DataFileFormat, Literal, NameMapping, PartitionSpec, RawLiteral, SchemaRef, StructType,
376 Type,
377 };
378 use crate::{Error, ErrorKind, Result};
379
380 #[derive(Deserialize)]
381 pub(super) struct FileScanTaskSerde {
382 file_size_in_bytes: u64,
383 start: u64,
384 length: u64,
385 record_count: Option<u64>,
386 first_row_id: Option<i64>,
387 data_sequence_number: Option<i64>,
388 data_file_path: String,
389 data_file_format: DataFileFormat,
390 schema: SchemaRef,
391 project_field_ids: Vec<i32>,
392 predicate: Option<BoundPredicate>,
393 deletes: Vec<FileScanTaskDeleteFile>,
394 #[serde(default)]
395 partition: Option<RawLiteral>,
396 #[serde(default)]
397 partition_spec: Option<Arc<PartitionSpec>>,
398 #[serde(default)]
399 name_mapping: Option<Arc<NameMapping>>,
400 #[serde(default)]
401 unified_partition_type: Option<Arc<StructType>>,
402 case_sensitive: bool,
403 #[serde(default)]
404 key_metadata: Option<Box<[u8]>>,
405 }
406
407 #[derive(Serialize)]
408 struct FileScanTaskRefSerde<'a> {
409 file_size_in_bytes: u64,
410 start: u64,
411 length: u64,
412 #[serde(skip_serializing_if = "Option::is_none")]
413 record_count: Option<u64>,
414 #[serde(skip_serializing_if = "Option::is_none")]
415 first_row_id: Option<i64>,
416 #[serde(skip_serializing_if = "Option::is_none")]
417 data_sequence_number: Option<i64>,
418 data_file_path: &'a str,
419 data_file_format: DataFileFormat,
420 schema: &'a SchemaRef,
421 project_field_ids: &'a [i32],
422 #[serde(skip_serializing_if = "Option::is_none")]
423 predicate: Option<&'a BoundPredicate>,
424 deletes: &'a [FileScanTaskDeleteFile],
425 #[serde(skip_serializing_if = "Option::is_none")]
426 partition: Option<RawLiteral>,
427 #[serde(skip_serializing_if = "Option::is_none")]
428 partition_spec: Option<&'a Arc<PartitionSpec>>,
429 #[serde(skip_serializing_if = "Option::is_none")]
430 name_mapping: Option<&'a Arc<NameMapping>>,
431 #[serde(skip_serializing_if = "Option::is_none")]
432 unified_partition_type: Option<&'a Arc<StructType>>,
433 case_sensitive: bool,
434 #[serde(skip_serializing_if = "Option::is_none")]
435 key_metadata: Option<&'a [u8]>,
436 }
437
438 fn partition_type(
439 partition_spec: Option<&PartitionSpec>,
440 schema: &crate::spec::Schema,
441 ) -> Result<Type> {
442 let partition_type = match partition_spec {
443 Some(partition_spec) => partition_spec.partition_type(schema)?,
444 None => PartitionSpec::unpartition_spec().partition_type(schema)?,
445 };
446 Ok(Type::Struct(partition_type))
447 }
448
449 impl<'a> TryFrom<&'a FileScanTask> for FileScanTaskRefSerde<'a> {
450 type Error = Error;
451
452 fn try_from(value: &'a FileScanTask) -> Result<Self> {
453 let partition = value
454 .partition
455 .as_ref()
456 .map(|partition| {
457 let partition_type =
458 partition_type(value.partition_spec.as_deref(), &value.schema)?;
459 RawLiteral::try_from(Literal::Struct(partition.clone()), &partition_type)
460 })
461 .transpose()?;
462
463 Ok(Self {
464 file_size_in_bytes: value.file_size_in_bytes,
465 start: value.start,
466 length: value.length,
467 record_count: value.record_count,
468 first_row_id: value.first_row_id,
469 data_sequence_number: value.data_sequence_number,
470 data_file_path: &value.data_file_path,
471 data_file_format: value.data_file_format,
472 schema: &value.schema,
473 project_field_ids: &value.project_field_ids,
474 predicate: value.predicate.as_ref(),
475 deletes: &value.deletes,
476 partition,
477 partition_spec: value.partition_spec.as_ref(),
478 name_mapping: value.name_mapping.as_ref(),
479 unified_partition_type: value.unified_partition_type.as_ref(),
480 case_sensitive: value.case_sensitive,
481 key_metadata: value.key_metadata.as_deref(),
482 })
483 }
484 }
485
486 impl Serialize for FileScanTask {
487 fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
488 where S: serde::Serializer {
489 FileScanTaskRefSerde::try_from(self)
490 .map_err(serde::ser::Error::custom)?
491 .serialize(serializer)
492 }
493 }
494
495 impl TryFrom<FileScanTaskSerde> for FileScanTask {
496 type Error = Error;
497
498 fn try_from(value: FileScanTaskSerde) -> Result<Self> {
499 let partition = value
500 .partition
501 .map(|partition| {
502 let partition_type =
503 partition_type(value.partition_spec.as_deref(), &value.schema)?;
504 match partition.try_into(&partition_type)? {
505 Some(Literal::Struct(partition)) => Ok(partition),
506 _ => Err(Error::new(
507 ErrorKind::DataInvalid,
508 "FileScanTask partition must be a struct",
509 )),
510 }
511 })
512 .transpose()?;
513
514 Self::builder()
515 .with_file_size_in_bytes(value.file_size_in_bytes)
516 .with_start(value.start)
517 .with_length(value.length)
518 .with_record_count(value.record_count)
519 .with_first_row_id(value.first_row_id)
520 .with_data_sequence_number(value.data_sequence_number)
521 .with_data_file_path(value.data_file_path)
522 .with_data_file_format(value.data_file_format)
523 .with_schema(value.schema)
524 .with_project_field_ids(value.project_field_ids)
525 .with_predicate(value.predicate)
526 .with_deletes(value.deletes)
527 .with_partition(partition)
528 .with_partition_spec(value.partition_spec)
529 .with_name_mapping(value.name_mapping)
530 .with_unified_partition_type(value.unified_partition_type)
531 .with_case_sensitive(value.case_sensitive)
532 .with_key_metadata(value.key_metadata)
533 .build()
534 }
535 }
536}
537
538#[cfg(test)]
539mod tests {
540 use super::*;
541 use crate::ErrorKind;
542 use crate::spec::{Literal, NestedField, PrimitiveType, Transform, Type};
543
544 fn build_file_scan_task(
545 schema: SchemaRef,
546 partition: Option<Struct>,
547 partition_spec: Option<Arc<PartitionSpec>>,
548 ) -> Result<FileScanTask> {
549 FileScanTask::builder()
550 .with_file_size_in_bytes(100)
551 .with_start(0)
552 .with_length(100)
553 .with_data_file_path("data_file_path".to_string())
554 .with_data_file_format(DataFileFormat::Parquet)
555 .with_schema(schema)
556 .with_project_field_ids(vec![])
557 .with_partition(partition)
558 .with_partition_spec(partition_spec)
559 .with_case_sensitive(false)
560 .build()
561 }
562
563 fn schema_and_spec(
564 primitive_type: PrimitiveType,
565 transform: Transform,
566 ) -> (SchemaRef, Arc<PartitionSpec>) {
567 let schema = Arc::new(
568 Schema::builder()
569 .with_fields(vec![Arc::new(NestedField::required(
570 1,
571 "x",
572 Type::Primitive(primitive_type),
573 ))])
574 .build()
575 .unwrap(),
576 );
577 let partition_spec = Arc::new(
578 PartitionSpec::builder(schema.clone())
579 .add_partition_field("x", "x_partition", transform)
580 .unwrap()
581 .build()
582 .unwrap(),
583 );
584 (schema, partition_spec)
585 }
586
587 #[test]
588 fn test_file_scan_task_builder_rejects_non_empty_partition_without_spec() {
589 let err = build_file_scan_task(
591 Arc::new(Schema::builder().build().unwrap()),
592 Some(Struct::from_iter([Some(Literal::long(42))])),
593 None,
594 )
595 .unwrap_err();
596
597 assert_eq!(err.kind(), ErrorKind::DataInvalid);
598 assert_eq!(
599 err.message(),
600 "Non-empty FileScanTask partition requires a partition spec"
601 );
602 }
603
604 #[test]
605 fn test_file_scan_task_builder_accepts_empty_partition_without_spec() {
606 build_file_scan_task(
607 Arc::new(Schema::builder().build().unwrap()),
608 Some(Struct::empty()),
609 None,
610 )
611 .unwrap();
612 }
613
614 #[test]
615 fn test_file_scan_task_builder_rejects_partitioned_spec_without_partition() {
616 let (schema, partition_spec) = schema_and_spec(PrimitiveType::Long, Transform::Identity);
617
618 let err = build_file_scan_task(schema, None, Some(partition_spec)).unwrap_err();
619
620 assert_eq!(err.kind(), ErrorKind::DataInvalid);
621 assert_eq!(
622 err.message(),
623 "FileScanTask with a partitioned spec requires partition values"
624 );
625 }
626
627 #[test]
628 fn test_file_scan_task_builder_accepts_unpartitioned_spec_without_partition() {
629 build_file_scan_task(
630 Arc::new(Schema::builder().build().unwrap()),
631 None,
632 Some(Arc::new(PartitionSpec::unpartition_spec())),
633 )
634 .unwrap();
635 }
636
637 #[test]
638 fn test_file_scan_task_builder_rejects_partition_arity_mismatch() {
639 let (schema, partition_spec) = schema_and_spec(PrimitiveType::Long, Transform::Identity);
640
641 let err =
642 build_file_scan_task(schema, Some(Struct::empty()), Some(partition_spec)).unwrap_err();
643
644 assert_eq!(err.kind(), ErrorKind::DataInvalid);
645 assert!(err.message().contains("partition has 0 fields"));
646 assert!(err.message().contains("partition spec has 1 fields"));
647 }
648
649 #[test]
650 fn test_file_scan_task_builder_rejects_dropped_partition_source_column() {
651 let (_historical_schema, partition_spec) =
652 schema_and_spec(PrimitiveType::Long, Transform::Identity);
653 let current_schema = Arc::new(
654 Schema::builder()
655 .with_fields(vec![Arc::new(NestedField::required(
656 2,
657 "y",
658 Type::Primitive(PrimitiveType::String),
659 ))])
660 .build()
661 .unwrap(),
662 );
663
664 let err = build_file_scan_task(
665 current_schema,
666 Some(Struct::from_iter([Some(Literal::long(42))])),
667 Some(partition_spec),
668 )
669 .unwrap_err();
670
671 assert_eq!(err.kind(), ErrorKind::Unexpected);
672 assert!(err.message().contains("No column with source column id 1"));
673 }
674
675 #[test]
676 fn test_file_scan_task_builder_rejects_partition_spec_incompatible_with_schema() {
677 let (_historical_schema, partition_spec) =
678 schema_and_spec(PrimitiveType::Timestamp, Transform::Day);
679 let current_schema = Arc::new(
680 Schema::builder()
681 .with_fields(vec![Arc::new(NestedField::required(
682 1,
683 "x",
684 Type::Primitive(PrimitiveType::String),
685 ))])
686 .build()
687 .unwrap(),
688 );
689
690 let err = build_file_scan_task(
691 current_schema,
692 Some(Struct::from_iter([Some(Literal::date(20_000))])),
693 Some(partition_spec),
694 )
695 .unwrap_err();
696
697 assert_eq!(err.kind(), ErrorKind::DataInvalid);
698 }
699}