feat: support file-level parquet row selections (#22940)

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

- Closes #22939 

## Rationale for this change

- see issue #22939 

## What changes are included in this PR?

- Add public `ParquetRowSelection`.
- Add `ParquetAccessPlan::try_new_from_overall_row_selection`.
- Allow Parquet opener setup to read either `ParquetAccessPlan` or
`ParquetRowSelection`.
- Reject using both extension types on the same file.
- Validate that the selection row count matches the file row count.
- Document the new extension path in `ParquetSource`.


## Are these changes tested?

Yes. This PR adds tests for:

- converting a file-level selection into row-group access
- rejecting invalid selection row counts
- creating an initial plan from `ParquetRowSelection`
- rejecting both `ParquetAccessPlan` and `ParquetRowSelection` on the
same file

## Are there any user-facing changes?

Yes. This adds a new public `ParquetRowSelection` type for callers that
want to attach a file-level Parquet `RowSelection` to a
`PartitionedFile`.

Author: haohuaijin

datafusion/core/tests/parquet/external_access_plan.rs

@@ -29,8 +29,10 @@ use datafusion::common::Result;
 use datafusion::datasource::listing::PartitionedFile;
 use datafusion::datasource::physical_plan::ParquetSource;
 use datafusion::prelude::SessionContext;
-use datafusion_common::{DFSchema, assert_contains};
-use datafusion_datasource_parquet::{ParquetAccessPlan, RowGroupAccess};
+use datafusion_common::{DFSchema, assert_batches_eq, assert_contains};
+use datafusion_datasource_parquet::{
+    ParquetAccessPlan, ParquetRowSelection, RowGroupAccess,
+};
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_expr::{Expr, col, lit};
 use datafusion_physical_plan::ExecutionPlan;
@@ -152,6 +154,94 @@ async fn skip_scan() {
     }
 }
 
+#[tokio::test]
+async fn row_selection_extension() {
+    // The file has 2 row groups of 5 rows each (10 rows total). Attach a
+    // file-level `ParquetRowSelection` to the `PartitionedFile` and verify it
+    // survives the path from `PartitionedFile` into the parquet opener/reader.
+
+    // select a single row in the first row group
+    let parquet_metrics = TestFull {
+        access_plan: None,
+        row_selection: Some(ParquetRowSelection::new(RowSelection::from(vec![
+            RowSelector::skip(2),
+            RowSelector::select(1),
+            RowSelector::skip(7),
+        ]))),
+        expected_rows: 1,
+        expected_output: Some(&[
+            "+------+------------+",
+            "| utf8 | large_utf8 |",
+            "+------+------------+",
+            "| c    | c          |",
+            "+------+------------+",
+        ]),
+        predicate: None,
+    }
+    .run()
+    .await
+    .unwrap();
+
+    // only the first row group is read, so some bytes are scanned
+    let bytes_scanned = metric_value(&parquet_metrics, "bytes_scanned").unwrap();
+    assert_ne!(bytes_scanned, 0, "metrics : {parquet_metrics:#?}",);
+}
+
+#[tokio::test]
+async fn row_selection_extension_spanning_row_groups() {
+    // A selection whose selectors straddle the row group boundary (row 4 is the
+    // last row of group 0, rows 5-6 are the first rows of group 1).
+    let parquet_metrics = TestFull {
+        access_plan: None,
+        row_selection: Some(ParquetRowSelection::new(RowSelection::from(vec![
+            RowSelector::skip(4),
+            RowSelector::select(3),
+            RowSelector::skip(3),
+        ]))),
+        expected_rows: 3,
+        expected_output: Some(&[
+            "+------+------------+",
+            "| utf8 | large_utf8 |",
+            "+------+------------+",
+            "|      |            |",
+            "| e    | e          |",
+            "| f    | f          |",
+            "+------+------------+",
+        ]),
+        predicate: None,
+    }
+    .run()
+    .await
+    .unwrap();
+
+    let bytes_scanned = metric_value(&parquet_metrics, "bytes_scanned").unwrap();
+    assert_ne!(bytes_scanned, 0, "metrics : {parquet_metrics:#?}",);
+}
+
+#[tokio::test]
+async fn bad_row_selection_extension() {
+    // selection specifies fewer rows than the file actually contains
+    let err = TestFull {
+        access_plan: None,
+        row_selection: Some(ParquetRowSelection::new(RowSelection::from(vec![
+            RowSelector::skip(2),
+            RowSelector::select(1),
+        ]))),
+        expected_rows: 10000,
+        expected_output: None,
+        predicate: None,
+    }
+    .run()
+    .await
+    .unwrap_err();
+    let err_string = err.to_string();
+    assert_contains!(&err_string, "Invalid Parquet RowSelection");
+    assert_contains!(
+        &err_string,
+        "File has 10 rows, but selection specifies 3 rows."
+    );
+}
+
 #[tokio::test]
 async fn plan_and_filter() {
     // show that row group pruning is applied even when an initial plan is supplied
@@ -170,7 +260,9 @@ async fn plan_and_filter() {
     // initial
     let parquet_metrics = TestFull {
         access_plan,
+        row_selection: None,
         expected_rows: 0,
+        expected_output: None,
         predicate: Some(predicate),
     }
     .run()
@@ -227,7 +319,9 @@ async fn bad_row_groups() {
             RowGroupAccess::Skip,
             RowGroupAccess::Scan,
         ])),
+        row_selection: None,
         expected_rows: 0,
+        expected_output: None,
         predicate: None,
     }
     .run()
@@ -249,8 +343,10 @@ async fn bad_selection() {
             ])),
             RowGroupAccess::Skip,
         ])),
+        row_selection: None,
         // expects that we hit an error, this should not be run
         expected_rows: 10000,
+        expected_output: None,
         predicate: None,
     }
     .run()
@@ -300,7 +396,9 @@ impl Test {
         } = self;
         TestFull {
             access_plan,
+            row_selection: None,
             expected_rows,
+            expected_output: None,
             predicate: None,
         }
         .run()
@@ -317,7 +415,9 @@ impl Test {
 /// 4. Returns the statistics from running the plan
 struct TestFull {
     access_plan: Option<ParquetAccessPlan>,
+    row_selection: Option<ParquetRowSelection>,
     expected_rows: usize,
+    expected_output: Option<&'static [&'static str]>,
     predicate: Option<Expr>,
 }
 
@@ -327,7 +427,9 @@ impl TestFull {
 
         let Self {
             access_plan,
+            row_selection,
             expected_rows,
+            expected_output,
             predicate,
         } = self;
 
@@ -352,6 +454,11 @@ impl TestFull {
             partitioned_file = partitioned_file.with_extension(access_plan);
         }
 
+        // add the file-level row selection, if any, as an extension
+        if let Some(row_selection) = row_selection {
+            partitioned_file = partitioned_file.with_extension(row_selection);
+        }
+
         // Create a DataSourceExec to read the file
         let object_store_url = ObjectStoreUrl::local_filesystem();
         // add the predicate, if requested
@@ -380,6 +487,9 @@ impl TestFull {
             "results: \n{}",
             pretty_format_batches(&results).unwrap()
         );
+        if let Some(expected_output) = expected_output {
+            assert_batches_eq!(expected_output, &results);
+        }
 
         std::fs::remove_file(file_name).unwrap();