Enable skip_rows in the chunked parquet reader.

[mhaseeb123]

Description

Closes #16186
Closes #16249

This PR adds multiple bug fixes required to enable the skip_rows option for the Chunked Parquet reader.

CC: @etseidl

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[Matt711]

I think we can support skipping row groups in cudf.polars (xref #16257) once this PR goes through.

[mhaseeb123]

Not 100% sure about this fix yet but seems like do end up needing more pages than found with start_row and end_row for list columns in some cases. I am guessing the page_row_index doesn't contain exact row indices for list columns? Maybe @etseidl or @nvdbaranec can elaborate if possible.

Question: There shouldn't be any performance penalties for this change as we do a trim pass later on with accurate row counts for lists pages, right?

The following case fails (not all pages for the list col c selected) without this change.

import cudf
from cudf.testing import assert_eq
import pyarrow as pa
from io import BytesIO

data = [
    {
        "a": "g",
        "b": {
            "b_a": 10,
            "b_b": {"b_b_b": None, "b_b_a": 2},
        },
        "c": None,
    },
    {"a": None, "b": {"b_a": None, "b_b": None}, "c": [15, 16]},
    {"a": "j", "b": None, "c": [8, 10]},
    {"a": None, "b": {"b_a": None, "b_b": None}, "c": None},
    None,
    {
        "a": None,
        "b": {"b_a": None, "b_b": {"b_b_b": 1}},
        "c": [18, 19],
    },
    {"a": None, "b": None, "c": None},
] * 1000

pa_struct = pa.Table.from_pydict({"struct": data})
df = cudf.DataFrame.from_arrow(pa_struct)

buffer = BytesIO()
df.to_parquet(buffer, row_group_size_rows=7000, max_page_size_rows=100)

# Number of rows to read
with cudf.option_context("io.parquet.low_memory", True):
    actual = cudf.read_parquet(
        [buffer],
        _chunk_read_limit=0,
        _pass_read_limit=1024000000,
        nrows=999,
        skip_rows=6001,
    ).reset_index(drop=True)
expected = cudf.read_parquet(buffer, nrows=999, skip_rows=6001).reset_index(
    drop=True
)
assert_eq(expected, actual)

[nvdbaranec]

I'm not sure this is the right way to do this. If I'm understanding what's going on here, this will just arbitrarily return the entire row group. Which I'm pretty sure will cause us to read all pages for all columns in a row group, which reverts us back to just being a row group reader in the case where lists exist.

I think the only real option here, if you have lists, is for each eligible row group

• march through the entire row group, using the row estimates as we do now.

• for each subpass that we read in, after we compute the real row counts (roughly here

  cudf/cpp/src/io/parquet/reader_impl_preprocess.cu

  Line 1485 in ea2bca3

  size_t first_page_index = 0;

  ) determine re-apply skip_rows/num_rows. If we have managed to read a subpass that contains none of the rows we care about, early out and move on to the next subpass.

• stop early if we cross the last row index we care about.

[mhaseeb123]

Catch the infinite for loop here if occurs.

[nvdbaranec]

I'm not sure this is the right way to do this. If I'm understanding what's going on here, this will just arbitrarily return the entire row group. Which I'm pretty sure will cause us to read all pages for all columns in a row group, which reverts us back to just being a row group reader in the case where lists exist.

I think the only real option here, if you have lists, is for each eligible row group

• march through the entire row group, using the row estimates as we do now.

• for each subpass that we read in, after we compute the real row counts (roughly here

  cudf/cpp/src/io/parquet/reader_impl_preprocess.cu

  Line 1485 in ea2bca3

  size_t first_page_index = 0;

  ) determine re-apply skip_rows/num_rows. If we have managed to read a subpass that contains none of the rows we care about, early out and move on to the next subpass.

• stop early if we cross the last row index we care about.

[mhaseeb123]

Accidentally clicked Ready for review instead of Update branch. Will mark this PR ready once it is

[mhaseeb123]

/ok to test

[nvdbaranec]

I vaguely don't like the idea of putting an accurate computation for row counts inside an function named estimates. This is probably the right place for it though, so maybe the function name should be more specific.

[mhaseeb123]

Updated the name of the function with an input argument indicating if we are writing estimates or not