Skip to content

Commit e6d6da2

Browse files
haohuaijinhaohuaijin
committed
update
1 parent 185798b commit e6d6da2

2 files changed

Lines changed: 233 additions & 10 deletions

File tree

‎datafusion/datasource-parquet/src/access_plan.rs‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -397,6 +397,151 @@ impl ParquetAccessPlan {
397397
}
398398
}
399399

400+
/// Return an overall `RowSelection`, if needed
401+
///
402+
/// This is used to compute the row selection for the parquet reader. See
403+
/// [`ArrowReaderBuilder::with_row_selection`] for more details.
404+
///
405+
/// Returns
406+
/// * `None` if there are no [`RowGroupAccess::Selection`]
407+
/// * `Some(selection)` if there are [`RowGroupAccess::Selection`]s
408+
///
409+
/// The returned selection represents which rows to scan across any row
410+
/// groups which are not skipped.
411+
///
412+
/// # Deprecated
413+
///
414+
/// DataFusion scans now keep selections local to each row group. For custom
415+
/// readers, use [`Self::into_inner`] to obtain the [`RowGroupAccess`] entries
416+
/// and adapt them to the reader's selection API. This compatibility method
417+
/// retains its existing conversion and length validation behavior.
418+
///
419+
/// # Notes
420+
///
421+
/// If there are no [`RowGroupAccess::Selection`]s, the overall row
422+
/// selection is `None` because each row group is either entirely skipped or
423+
/// scanned, which is covered by [`Self::row_group_indexes`].
424+
///
425+
/// If there are any [`RowGroupAccess::Selection`], an overall row selection
426+
/// is returned for *all* the rows in the row groups that are not skipped.
427+
/// Thus it includes a `Select` selection for any [`RowGroupAccess::Scan`].
428+
///
429+
/// # Errors
430+
///
431+
/// Returns an error if any specified row selection does not specify
432+
/// the same number of rows as in its corresponding `row_group_metadata`.
433+
///
434+
/// # Example: No Selections
435+
///
436+
/// Given an access plan like this
437+
///
438+
/// ```text
439+
/// RowGroupAccess::Scan (scan all row group 0)
440+
/// RowGroupAccess::Skip (skip row group 1)
441+
/// RowGroupAccess::Scan (scan all row group 2)
442+
/// RowGroupAccess::Scan (scan all row group 3)
443+
/// ```
444+
///
445+
/// The overall row selection would be `None` because there are no
446+
/// [`RowGroupAccess::Selection`]s. The row group indexes
447+
/// returned by [`Self::row_group_indexes`] would be `0, 2, 3` .
448+
///
449+
/// # Example: With Selections
450+
///
451+
/// Given an access plan like this:
452+
///
453+
/// ```text
454+
/// RowGroupAccess::Scan (scan all row group 0)
455+
/// RowGroupAccess::Skip (skip row group 1)
456+
/// RowGroupAccess::Selection (skip 50, scan 50, skip 900) (scan rows 50-100 in row group 2)
457+
/// RowGroupAccess::Scan (scan all row group 3)
458+
/// ```
459+
///
460+
/// Assuming each row group has 1000 rows, the resulting row selection would
461+
/// be the rows to scan in row group 0, 2 and 3:
462+
///
463+
/// ```text
464+
/// RowSelection::Select(1000) (scan all rows in row group 0)
465+
/// RowSelection::Skip(50) (skip first 50 rows in row group 2)
466+
/// RowSelection::Select(50) (scan rows 50-100 in row group 2)
467+
/// RowSelection::Skip(900) (skip last 900 rows in row group 2)
468+
/// RowSelection::Select(1000) (scan all rows in row group 3)
469+
/// ```
470+
///
471+
/// Note there is no entry for the (entirely) skipped row group 1.
472+
///
473+
/// The row group indexes returned by [`Self::row_group_indexes`] would
474+
/// still be `0, 2, 3` .
475+
///
476+
/// [`ArrowReaderBuilder::with_row_selection`]: parquet::arrow::arrow_reader::ArrowReaderBuilder::with_row_selection
477+
#[deprecated(
478+
since = "56.0.0",
479+
note = "Use into_inner() to obtain row-group-local access entries"
480+
)]
481+
pub fn into_overall_row_selection(
482+
self,
483+
row_group_meta_data: &[RowGroupMetaData],
484+
) -> Result<Option<RowSelection>> {
485+
assert_eq!(row_group_meta_data.len(), self.row_groups.len());
486+
// Intuition: entire row groups are filtered out using
487+
// `row_group_indexes` which come from Skip and Scan. An overall
488+
// RowSelection is only useful if there is any parts *within* a row group
489+
// which can be filtered out, that is a `Selection`.
490+
if !self
491+
.row_groups
492+
.iter()
493+
.any(|rg| matches!(rg, RowGroupAccess::Selection(_)))
494+
{
495+
return Ok(None);
496+
}
497+
498+
// validate all Selections
499+
for (idx, (rg, rg_meta)) in self
500+
.row_groups
501+
.iter()
502+
.zip(row_group_meta_data.iter())
503+
.enumerate()
504+
{
505+
let RowGroupAccess::Selection(selection) = rg else {
506+
continue;
507+
};
508+
let rows_in_selection = selection
509+
.iter()
510+
.map(|selection| selection.row_count)
511+
.sum::<usize>();
512+
513+
let row_group_row_count = rg_meta.num_rows();
514+
assert_eq_or_internal_err!(
515+
rows_in_selection as i64,
516+
row_group_row_count,
517+
"Invalid ParquetAccessPlan Selection. Row group {idx} has {row_group_row_count} rows \
518+
but selection only specifies {rows_in_selection} rows. \
519+
Selection: {selection:?}"
520+
);
521+
}
522+
523+
let total_selection: RowSelection = self
524+
.row_groups
525+
.into_iter()
526+
.zip(row_group_meta_data.iter())
527+
.flat_map(|(rg, rg_meta)| {
528+
match rg {
529+
RowGroupAccess::Skip => vec![],
530+
RowGroupAccess::Scan => {
531+
// need a row group access to scan the entire row group (need row group counts)
532+
vec![RowSelector::select(rg_meta.num_rows() as usize)]
533+
}
534+
RowGroupAccess::Selection(selection) => {
535+
let selection: Vec<RowSelector> = selection.into();
536+
selection
537+
}
538+
}
539+
})
540+
.collect();
541+
542+
Ok(Some(total_selection))
543+
}
544+
400545
/// Return an iterator over the row group indexes that should be scanned
401546
pub fn row_group_index_iter(&self) -> impl Iterator<Item = usize> + '_ {
402547
self.row_groups
@@ -692,6 +837,81 @@ mod test {
692837
use parquet::schema::types::{SchemaDescPtr, SchemaDescriptor};
693838
use std::sync::{Arc, LazyLock};
694839

840+
#[test]
841+
#[allow(deprecated)]
842+
fn test_deprecated_overall_row_selection() {
843+
for plan in [
844+
ParquetAccessPlan::new_all(4),
845+
ParquetAccessPlan::new_none(4),
846+
] {
847+
assert_eq!(
848+
plan.into_overall_row_selection(&ROW_GROUP_METADATA)
849+
.unwrap(),
850+
None
851+
);
852+
}
853+
854+
// Skipped groups do not occupy coordinates in the combined selection.
855+
// Check both input representations retain the same conversion behavior.
856+
let selectors =
857+
RowSelection::from(vec![RowSelector::skip(10), RowSelector::select(20)]);
858+
let bitmap = RowSelection::from(arrow::buffer::BooleanBuffer::from(
859+
(0..30).map(|i| i >= 10).collect::<Vec<_>>(),
860+
));
861+
for selection in [selectors, bitmap] {
862+
let plan = ParquetAccessPlan::new(vec![
863+
RowGroupAccess::Scan,
864+
RowGroupAccess::Skip,
865+
RowGroupAccess::Selection(selection),
866+
RowGroupAccess::Scan,
867+
]);
868+
assert_eq!(
869+
plan.into_overall_row_selection(&ROW_GROUP_METADATA)
870+
.unwrap(),
871+
Some(RowSelection::from(vec![
872+
RowSelector::select(10),
873+
RowSelector::skip(10),
874+
RowSelector::select(60),
875+
]))
876+
);
877+
}
878+
879+
let plan = ParquetAccessPlan::new(vec![
880+
RowGroupAccess::Selection(vec![RowSelector::skip(10)].into()),
881+
RowGroupAccess::Skip,
882+
RowGroupAccess::Skip,
883+
RowGroupAccess::Skip,
884+
]);
885+
assert_eq!(
886+
plan.into_overall_row_selection(&ROW_GROUP_METADATA)
887+
.unwrap(),
888+
Some(RowSelection::from(vec![RowSelector::skip(10)]))
889+
);
890+
}
891+
892+
#[test]
893+
#[allow(deprecated)]
894+
fn test_deprecated_overall_row_selection_validates_length() {
895+
for rows in [19, 21] {
896+
let plan = ParquetAccessPlan::new(vec![
897+
RowGroupAccess::Scan,
898+
RowGroupAccess::Selection(vec![RowSelector::select(rows)].into()),
899+
RowGroupAccess::Skip,
900+
RowGroupAccess::Skip,
901+
]);
902+
let err = plan
903+
.into_overall_row_selection(&ROW_GROUP_METADATA)
904+
.unwrap_err()
905+
.to_string();
906+
assert_contains!(
907+
err,
908+
format!(
909+
"Row group 1 has 20 rows but selection only specifies {rows} rows"
910+
)
911+
);
912+
}
913+
}
914+
695915
#[test]
696916
fn test_only_scans() {
697917
let plan = ParquetAccessPlan::new_all(4)

‎docs/source/library-user-guide/upgrading/56.0.0.md‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -65,22 +65,25 @@ change. Use `FieldMetadata::to_hashmap` where a `HashMap` is still required.
6565

6666
[`metadata`]: https://docs.rs/arrow/60.0.0/arrow/datatypes/struct.Metadata.html
6767

68-
### `ParquetAccessPlan::into_overall_row_selection` was removed
68+
### `ParquetAccessPlan::into_overall_row_selection` is deprecated
6969

70-
`ParquetAccessPlan::into_overall_row_selection` has been removed. DataFusion now
71-
keeps selections local to each row group during Parquet scanning.
70+
`ParquetAccessPlan::into_overall_row_selection` is deprecated starting in 56.0.0.
71+
DataFusion now keeps selections local to each row group during Parquet scanning.
72+
The method remains available with its existing conversion and selection-length
73+
validation behavior, so direct callers continue to work but receive a deprecation
74+
warning.
7275

7376
Users who attach a `ParquetAccessPlan` to a `PartitionedFile` for DataFusion to
74-
read need no changes. Only code that directly calls the removed conversion
75-
method needs to migrate; there is no replacement public conversion method.
77+
read need no changes.
7678

7779
For custom readers, use `ParquetAccessPlan::into_inner()` to obtain the
7880
`RowGroupAccess` entries. Their positions are the original file row-group
79-
indexes, and each selection uses coordinates within its own row group. Callers
80-
must now handle conversion and selection-length validation themselves. Parquet's
81-
push decoder and async reader accept `RowGroupSelection` values through
82-
`with_row_group_selections`; the synchronous reader still requires
83-
`with_row_groups` and an overall `with_row_selection`.
81+
indexes, and each selection uses coordinates within its own row group. When
82+
migrating away from the deprecated method, callers must handle conversion and
83+
selection-length validation themselves. Parquet's push decoder and async reader
84+
accept `RowGroupSelection` values through `with_row_group_selections`; the
85+
synchronous reader still requires `with_row_groups` and an overall
86+
`with_row_selection`.
8487

8588
### Missing Parquet null counts are treated as unknown
8689

0 commit comments

Comments
 (0)