Aluode/PerceptionLabPortable

1import io2import json3import warnings4from typing import Literal5 6import fsspec7 8from .core import url_to_fs9from .spec import AbstractBufferedFile10from .utils import merge_offset_ranges11 12# Parquet-Specific Utilities for fsspec13#14# Most of the functions defined in this module are NOT15# intended for public consumption. The only exception16# to this is `open_parquet_file`, which should be used17# place of `fs.open()` to open parquet-formatted files18# on remote file systems.19 20 21class AlreadyBufferedFile(AbstractBufferedFile):22    def _fetch_range(self, start, end):23        raise NotImplementedError24 25 26def open_parquet_files(27    path: list[str],28    mode: Literal["rb"] = "rb",29    fs: None | fsspec.AbstractFileSystem = None,30    metadata=None,31    columns: None | list[str] = None,32    row_groups: None | list[int] = None,33    storage_options: None | dict = None,34    engine: str = "auto",35    max_gap: int = 64_000,36    max_block: int = 256_000_000,37    footer_sample_size: int = 1_000_000,38    filters: None | list[list[list[str]]] = None,39    **kwargs,40):41    """42    Return a file-like object for a single Parquet file.43 44    The specified parquet `engine` will be used to parse the45    footer metadata, and determine the required byte ranges46    from the file. The target path will then be opened with47    the "parts" (`KnownPartsOfAFile`) caching strategy.48 49    Note that this method is intended for usage with remote50    file systems, and is unlikely to improve parquet-read51    performance on local file systems.52 53    Parameters54    ----------55    path: str56        Target file path.57    mode: str, optional58        Mode option to be passed through to `fs.open`. Default is "rb".59    metadata: Any, optional60        Parquet metadata object. Object type must be supported61        by the backend parquet engine. For now, only the "fastparquet"62        engine supports an explicit `ParquetFile` metadata object.63        If a metadata object is supplied, the remote footer metadata64        will not need to be transferred into local memory.65    fs: AbstractFileSystem, optional66        Filesystem object to use for opening the file. If nothing is67        specified, an `AbstractFileSystem` object will be inferred.68    engine : str, default "auto"69        Parquet engine to use for metadata parsing. Allowed options70        include "fastparquet", "pyarrow", and "auto". The specified71        engine must be installed in the current environment. If72        "auto" is specified, and both engines are installed,73        "fastparquet" will take precedence over "pyarrow".74    columns: list, optional75        List of all column names that may be read from the file.76    row_groups : list, optional77        List of all row-groups that may be read from the file. This78        may be a list of row-group indices (integers), or it may be79        a list of `RowGroup` metadata objects (if the "fastparquet"80        engine is used).81    storage_options : dict, optional82        Used to generate an `AbstractFileSystem` object if `fs` was83        not specified.84    max_gap : int, optional85        Neighboring byte ranges will only be merged when their86        inter-range gap is <= `max_gap`. Default is 64KB.87    max_block : int, optional88        Neighboring byte ranges will only be merged when the size of89        the aggregated range is <= `max_block`. Default is 256MB.90    footer_sample_size : int, optional91        Number of bytes to read from the end of the path to look92        for the footer metadata. If the sampled bytes do not contain93        the footer, a second read request will be required, and94        performance will suffer. Default is 1MB.95    filters : list[list], optional96        List of filters to apply to prevent reading row groups, of the97        same format as accepted by the loading engines. Ignored if98        ``row_groups`` is specified.99    **kwargs :100        Optional key-word arguments to pass to `fs.open`101    """102 103    # Make sure we have an `AbstractFileSystem` object104    # to work with105    if fs is None:106        path0 = path107        if isinstance(path, (list, tuple)):108            path = path[0]109        fs, path = url_to_fs(path, **(storage_options or {}))110    else:111        path0 = path112 113    # For now, `columns == []` not supported, is the same114    # as all columns115    if columns is not None and len(columns) == 0:116        columns = None117 118    # Set the engine119    engine = _set_engine(engine)120 121    if isinstance(path0, (list, tuple)):122        paths = path0123    elif "*" in path:124        paths = fs.glob(path)125    elif path0.endswith("/"):  # or fs.isdir(path):126        paths = [127            _128            for _ in fs.find(path, withdirs=False, detail=False)129            if _.endswith((".parquet", ".parq"))130        ]131    else:132        paths = [path]133 134    data = _get_parquet_byte_ranges(135        paths,136        fs,137        metadata=metadata,138        columns=columns,139        row_groups=row_groups,140        engine=engine,141        max_gap=max_gap,142        max_block=max_block,143        footer_sample_size=footer_sample_size,144        filters=filters,145    )146 147    # Call self.open with "parts" caching148    options = kwargs.pop("cache_options", {}).copy()149    return [150        AlreadyBufferedFile(151            fs=None,152            path=fn,153            mode=mode,154            cache_type="parts",155            cache_options={156                **options,157                "data": data.get(fn, {}),158            },159            size=max(_[1] for _ in data.get(fn, {})),160            **kwargs,161        )162        for fn in data163    ]164 165 166def open_parquet_file(*args, **kwargs):167    """Create files tailed to reading specific parts of parquet files168 169    Please see ``open_parquet_files`` for details of the arguments. The170    difference is, this function always returns a single ``AleadyBufferedFile``,171    whereas `open_parquet_files`` always returns a list of files, even if172    there are one or zero matching parquet files.173    """174    return open_parquet_files(*args, **kwargs)[0]175 176 177def _get_parquet_byte_ranges(178    paths,179    fs,180    metadata=None,181    columns=None,182    row_groups=None,183    max_gap=64_000,184    max_block=256_000_000,185    footer_sample_size=1_000_000,186    engine="auto",187    filters=None,188):189    """Get a dictionary of the known byte ranges needed190    to read a specific column/row-group selection from a191    Parquet dataset. Each value in the output dictionary192    is intended for use as the `data` argument for the193    `KnownPartsOfAFile` caching strategy of a single path.194    """195 196    # Set engine if necessary197    if isinstance(engine, str):198        engine = _set_engine(engine)199 200    # Pass to specialized function if metadata is defined201    if metadata is not None:202        # Use the provided parquet metadata object203        # to avoid transferring/parsing footer metadata204        return _get_parquet_byte_ranges_from_metadata(205            metadata,206            fs,207            engine,208            columns=columns,209            row_groups=row_groups,210            max_gap=max_gap,211            max_block=max_block,212            filters=filters,213        )214 215    # Get file sizes asynchronously216    file_sizes = fs.sizes(paths)217 218    # Populate global paths, starts, & ends219    result = {}220    data_paths = []221    data_starts = []222    data_ends = []223    add_header_magic = True224    if columns is None and row_groups is None and filters is None:225        # We are NOT selecting specific columns or row-groups.226        #227        # We can avoid sampling the footers, and just transfer228        # all file data with cat_ranges229        for i, path in enumerate(paths):230            result[path] = {}231            data_paths.append(path)232            data_starts.append(0)233            data_ends.append(file_sizes[i])234        add_header_magic = False  # "Magic" should already be included235    else:236        # We ARE selecting specific columns or row-groups.237        #238        # Gather file footers.239        # We just take the last `footer_sample_size` bytes of each240        # file (or the entire file if it is smaller than that)241        footer_starts = []242        footer_ends = []243        for i, path in enumerate(paths):244            footer_ends.append(file_sizes[i])245            sample_size = max(0, file_sizes[i] - footer_sample_size)246            footer_starts.append(sample_size)247        footer_samples = fs.cat_ranges(paths, footer_starts, footer_ends)248 249        # Check our footer samples and re-sample if necessary.250        missing_footer_starts = footer_starts.copy()251        large_footer = 0252        for i, path in enumerate(paths):253            footer_size = int.from_bytes(footer_samples[i][-8:-4], "little")254            real_footer_start = file_sizes[i] - (footer_size + 8)255            if real_footer_start < footer_starts[i]:256                missing_footer_starts[i] = real_footer_start257                large_footer = max(large_footer, (footer_size + 8))258        if large_footer:259            warnings.warn(260                f"Not enough data was used to sample the parquet footer. "261                f"Try setting footer_sample_size >= {large_footer}."262            )263            for i, block in enumerate(264                fs.cat_ranges(265                    paths,266                    missing_footer_starts,267                    footer_starts,268                )269            ):270                footer_samples[i] = block + footer_samples[i]271                footer_starts[i] = missing_footer_starts[i]272 273        # Calculate required byte ranges for each path274        for i, path in enumerate(paths):275            # Use "engine" to collect data byte ranges276            path_data_starts, path_data_ends = engine._parquet_byte_ranges(277                columns,278                row_groups=row_groups,279                footer=footer_samples[i],280                footer_start=footer_starts[i],281                filters=filters,282            )283 284            data_paths += [path] * len(path_data_starts)285            data_starts += path_data_starts286            data_ends += path_data_ends287            result.setdefault(path, {})[(footer_starts[i], file_sizes[i])] = (288                footer_samples[i]289            )290 291        # Merge adjacent offset ranges292        data_paths, data_starts, data_ends = merge_offset_ranges(293            data_paths,294            data_starts,295            data_ends,296            max_gap=max_gap,297            max_block=max_block,298            sort=False,  # Should already be sorted299        )300 301        # Start by populating `result` with footer samples302        for i, path in enumerate(paths):303            result[path] = {(footer_starts[i], footer_ends[i]): footer_samples[i]}304 305    # Transfer the data byte-ranges into local memory306    _transfer_ranges(fs, result, data_paths, data_starts, data_ends)307 308    # Add b"PAR1" to header if necessary309    if add_header_magic:310        _add_header_magic(result)311 312    return result313 314 315def _get_parquet_byte_ranges_from_metadata(316    metadata,317    fs,318    engine,319    columns=None,320    row_groups=None,321    max_gap=64_000,322    max_block=256_000_000,323    filters=None,324):325    """Simplified version of `_get_parquet_byte_ranges` for326    the case that an engine-specific `metadata` object is327    provided, and the remote footer metadata does not need to328    be transferred before calculating the required byte ranges.329    """330 331    # Use "engine" to collect data byte ranges332    data_paths, data_starts, data_ends = engine._parquet_byte_ranges(333        columns, row_groups=row_groups, metadata=metadata, filters=filters334    )335 336    # Merge adjacent offset ranges337    data_paths, data_starts, data_ends = merge_offset_ranges(338        data_paths,339        data_starts,340        data_ends,341        max_gap=max_gap,342        max_block=max_block,343        sort=False,  # Should be sorted344    )345 346    # Transfer the data byte-ranges into local memory347    result = {fn: {} for fn in list(set(data_paths))}348    _transfer_ranges(fs, result, data_paths, data_starts, data_ends)349 350    # Add b"PAR1" to header351    _add_header_magic(result)352 353    return result354 355 356def _transfer_ranges(fs, blocks, paths, starts, ends):357    # Use cat_ranges to gather the data byte_ranges358    ranges = (paths, starts, ends)359    for path, start, stop, data in zip(*ranges, fs.cat_ranges(*ranges)):360        blocks[path][(start, stop)] = data361 362 363def _add_header_magic(data):364    # Add b"PAR1" to file headers365    for path in list(data.keys()):366        add_magic = True367        for k in data[path]:368            if k[0] == 0 and k[1] >= 4:369                add_magic = False370                break371        if add_magic:372            data[path][(0, 4)] = b"PAR1"373 374 375def _set_engine(engine_str):376    # Define a list of parquet engines to try377    if engine_str == "auto":378        try_engines = ("fastparquet", "pyarrow")379    elif not isinstance(engine_str, str):380        raise ValueError(381            "Failed to set parquet engine! "382            "Please pass 'fastparquet', 'pyarrow', or 'auto'"383        )384    elif engine_str not in ("fastparquet", "pyarrow"):385        raise ValueError(f"{engine_str} engine not supported by `fsspec.parquet`")386    else:387        try_engines = [engine_str]388 389    # Try importing the engines in `try_engines`,390    # and choose the first one that succeeds391    for engine in try_engines:392        try:393            if engine == "fastparquet":394                return FastparquetEngine()395            elif engine == "pyarrow":396                return PyarrowEngine()397        except ImportError:398            pass399 400    # Raise an error if a supported parquet engine401    # was not found402    raise ImportError(403        f"The following parquet engines are not installed "404        f"in your python environment: {try_engines}."405        f"Please install 'fastparquert' or 'pyarrow' to "406        f"utilize the `fsspec.parquet` module."407    )408 409 410class FastparquetEngine:411    # The purpose of the FastparquetEngine class is412    # to check if fastparquet can be imported (on initialization)413    # and to define a `_parquet_byte_ranges` method. In the414    # future, this class may also be used to define other415    # methods/logic that are specific to fastparquet.416 417    def __init__(self):418        import fastparquet as fp419 420        self.fp = fp421 422    def _row_group_filename(self, row_group, pf):423        return pf.row_group_filename(row_group)424 425    def _parquet_byte_ranges(426        self,427        columns,428        row_groups=None,429        metadata=None,430        footer=None,431        footer_start=None,432        filters=None,433    ):434        # Initialize offset ranges and define ParqetFile metadata435        pf = metadata436        data_paths, data_starts, data_ends = [], [], []437        if filters and row_groups:438            raise ValueError("filters and row_groups cannot be used together")439        if pf is None:440            pf = self.fp.ParquetFile(io.BytesIO(footer))441 442        # Convert columns to a set and add any index columns443        # specified in the pandas metadata (just in case)444        column_set = None if columns is None else {c.split(".", 1)[0] for c in columns}445        if column_set is not None and hasattr(pf, "pandas_metadata"):446            md_index = [447                ind448                for ind in pf.pandas_metadata.get("index_columns", [])449                # Ignore RangeIndex information450                if not isinstance(ind, dict)451            ]452            column_set |= set(md_index)453 454        # Check if row_groups is a list of integers455        # or a list of row-group metadata456        if filters:457            from fastparquet.api import filter_row_groups458 459            row_group_indices = None460            row_groups = filter_row_groups(pf, filters)461        elif row_groups and not isinstance(row_groups[0], int):462            # Input row_groups contains row-group metadata463            row_group_indices = None464        else:465            # Input row_groups contains row-group indices466            row_group_indices = row_groups467            row_groups = pf.row_groups468 469        # Loop through column chunks to add required byte ranges470        for r, row_group in enumerate(row_groups):471            # Skip this row-group if we are targeting472            # specific row-groups473            if row_group_indices is None or r in row_group_indices:474                # Find the target parquet-file path for `row_group`475                fn = self._row_group_filename(row_group, pf)476 477                for column in row_group.columns:478                    name = column.meta_data.path_in_schema[0]479                    # Skip this column if we are targeting a480                    # specific columns481                    if column_set is None or name in column_set:482                        file_offset0 = column.meta_data.dictionary_page_offset483                        if file_offset0 is None:484                            file_offset0 = column.meta_data.data_page_offset485                        num_bytes = column.meta_data.total_compressed_size486                        if footer_start is None or file_offset0 < footer_start:487                            data_paths.append(fn)488                            data_starts.append(file_offset0)489                            data_ends.append(490                                min(491                                    file_offset0 + num_bytes,492                                    footer_start or (file_offset0 + num_bytes),493                                )494                            )495 496        if metadata:497            # The metadata in this call may map to multiple498            # file paths. Need to include `data_paths`499            return data_paths, data_starts, data_ends500        return data_starts, data_ends501 502 503class PyarrowEngine:504    # The purpose of the PyarrowEngine class is505    # to check if pyarrow can be imported (on initialization)506    # and to define a `_parquet_byte_ranges` method. In the507    # future, this class may also be used to define other508    # methods/logic that are specific to pyarrow.509 510    def __init__(self):511        import pyarrow.parquet as pq512 513        self.pq = pq514 515    def _row_group_filename(self, row_group, metadata):516        raise NotImplementedError517 518    def _parquet_byte_ranges(519        self,520        columns,521        row_groups=None,522        metadata=None,523        footer=None,524        footer_start=None,525        filters=None,526    ):527        if metadata is not None:528            raise ValueError("metadata input not supported for PyarrowEngine")529        if filters:530            raise NotImplementedError531 532        data_starts, data_ends = [], []533        md = self.pq.ParquetFile(io.BytesIO(footer)).metadata534 535        # Convert columns to a set and add any index columns536        # specified in the pandas metadata (just in case)537        column_set = None if columns is None else set(columns)538        if column_set is not None:539            schema = md.schema.to_arrow_schema()540            has_pandas_metadata = (541                schema.metadata is not None and b"pandas" in schema.metadata542            )543            if has_pandas_metadata:544                md_index = [545                    ind546                    for ind in json.loads(547                        schema.metadata[b"pandas"].decode("utf8")548                    ).get("index_columns", [])549                    # Ignore RangeIndex information550                    if not isinstance(ind, dict)551                ]552                column_set |= set(md_index)553 554        # Loop through column chunks to add required byte ranges555        for r in range(md.num_row_groups):556            # Skip this row-group if we are targeting557            # specific row-groups558            if row_groups is None or r in row_groups:559                row_group = md.row_group(r)560                for c in range(row_group.num_columns):561                    column = row_group.column(c)562                    name = column.path_in_schema563                    # Skip this column if we are targeting a564                    # specific columns565                    split_name = name.split(".")[0]566                    if (567                        column_set is None568                        or name in column_set569                        or split_name in column_set570                    ):571                        file_offset0 = column.dictionary_page_offset572                        if file_offset0 is None:573                            file_offset0 = column.data_page_offset574                        num_bytes = column.total_compressed_size575                        if file_offset0 < footer_start:576                            data_starts.append(file_offset0)577                            data_ends.append(578                                min(file_offset0 + num_bytes, footer_start)579                            )580        return data_starts, data_ends581