[Parquet][Python]: Reading subset by feeding data asynchronously to parquet parser #45352

class MyAsyncReader(Protocol):
    async def parquet_size(self) -> int: ...  # Size stored separately when writing elsewhere
    async def parquet_meta(self) -> bytes: ...  # Metadata stored separately when writing elsewhere
    async def parquet_data(self, start_offset: int, end_offset: int) -> bytes: ...


async def query(reader: MyAsyncReader, filter: Expression) -> Table:
    size = await reader.parquet_size()
    meta = await reader.parquet_meta()

    anon_mmap = mmap.mmap(-1, size, flags=mmap.MAP_ANONYMOUS | mmap.MAP_PRIVATE)
    try:
        anon_mmap.seek(size - len(meta))  # Meta to tail
        anon_mmap.write(meta)

        frag = ParquetFileFormat().make_fragment(anon_mmap).subset(filter)

        first_row_col = frag.row_groups[0].metadata.column(0)
        last_row_col = frag.row_groups[-1].metadata.column(frag.metadata.num_columns - 1)
        start_offset = offset(first_row_col)
        end_offset = offset(last_row_col) + last_row_col.total_compressed_size

        anon_mmap.seek(start_offset)
        anon_mmap.write(await reader.parquet_data(start_offset, end_offset))  # Feed needed data for parser

        return frag.to_table()  # Parse the subset of row groups
    finally:
        anon_mmap.close()


def offset(meta: ColumnChunkMetaData) -> int:
    return (
        min(meta.dictionary_page_offset, meta.data_page_offset)  # Is there a better way to get this?
        if meta.dictionary_page_offset is not None
        else meta.data_page_offset
    )