import io
import json
import warnings
from .core import url_to_fs
from .utils import merge_offset_ranges
# Parquet-Specific Utilities for fsspec
#
# Most of the functions defined in this module are NOT
# intended for public consumption. The only exception
# to this is `open_parquet_file`, which should be used
# place of `fs.open()` to open parquet-formatted files
# on remote file systems.
def open_parquet_file(
path,
mode="rb",
fs=None,
metadata=None,
columns=None,
row_groups=None,
storage_options=None,
strict=False,
engine="auto",
max_gap=64_000,
max_block=256_000_000,
footer_sample_size=1_000_000,
**kwargs,
):
"""
Return a file-like object for a single Parquet file.
The specified parquet `engine` will be used to parse the
footer metadata, and determine the required byte ranges
from the file. The target path will then be opened with
the "parts" (`KnownPartsOfAFile`) caching strategy.
Note that this method is intended for usage with remote
file systems, and is unlikely to improve parquet-read
performance on local file systems.
Parameters
----------
path: str
Target file path.
mode: str, optional
Mode option to be passed through to `fs.open`. Default is "rb".
metadata: Any, optional
Parquet metadata object. Object type must be supported
by the backend parquet engine. For now, only the "fastparquet"
engine supports an explicit `ParquetFile` metadata object.
If a metadata object is supplied, the remote footer metadata
will not need to be transferred into local memory.
fs: AbstractFileSystem, optional
Filesystem object to use for opening the file. If nothing is
specified, an `AbstractFileSystem` object will be inferred.
engine : str, default "auto"
Parquet engine to use for metadata parsing. Allowed options
include "fastparquet", "pyarrow", and "auto". The specified
engine must be installed in the current environment. If
"auto" is specified, and both engines are installed,
"fastparquet" will take precedence over "pyarrow".
columns: list, optional
List of all column names that may be read from the file.
row_groups : list, optional
List of all row-groups that may be read from the file. This
may be a list of row-group indices (integers), or it may be
a list of `RowGroup` metadata objects (if the "fastparquet"
engine is used).
storage_options : dict, optional
Used to generate an `AbstractFileSystem` object if `fs` was
not specified.
strict : bool, optional
Whether the resulting `KnownPartsOfAFile` cache should
fetch reads that go beyond a known byte-range boundary.
If `False` (the default), any read that ends outside a
known part will be zero padded. Note that using
`strict=True` may be useful for debugging.
max_gap : int, optional
Neighboring byte ranges will only be merged when their
inter-range gap is <= `max_gap`. Default is 64KB.
max_block : int, optional
Neighboring byte ranges will only be merged when the size of
the aggregated range is <= `max_block`. Default is 256MB.
footer_sample_size : int, optional
Number of bytes to read from the end of the path to look
for the footer metadata. If the sampled bytes do not contain
the footer, a second read request will be required, and
performance will suffer. Default is 1MB.
**kwargs :
Optional key-word arguments to pass to `fs.open`
"""
# Make sure we have an `AbstractFileSystem` object
# to work with
if fs is None:
fs = url_to_fs(path, **(storage_options or {}))[0]
# For now, `columns == []` not supported. Just use
# default `open` command with `path` input
if columns is not None and len(columns) == 0:
return fs.open(path, mode=mode)
# Set the engine
engine = _set_engine(engine)
# Fetch the known byte ranges needed to read
# `columns` and/or `row_groups`
data = _get_parquet_byte_ranges(
[path],
fs,
metadata=metadata,
columns=columns,
row_groups=row_groups,
engine=engine,
max_gap=max_gap,
max_block=max_block,
footer_sample_size=footer_sample_size,
)
# Extract file name from `data`
fn = next(iter(data)) if data else path
# Call self.open with "parts" caching
options = kwargs.pop("cache_options", {}).copy()
return fs.open(
fn,
mode=mode,
cache_type="parts",
cache_options={
**options,
"data": data.get(fn, {}),
"strict": strict,
},
**kwargs,
)
def _get_parquet_byte_ranges(
paths,
fs,
metadata=None,
columns=None,
row_groups=None,
max_gap=64_000,
max_block=256_000_000,
footer_sample_size=1_000_000,
engine="auto",
):
"""Get a dictionary of the known byte ranges needed
to read a specific column/row-group selection from a
Parquet dataset. Each value in the output dictionary
is intended for use as the `data` argument for the
`KnownPartsOfAFile` caching strategy of a single path.
"""
# Set engine if necessary
if isinstance(engine, str):
engine = _set_engine(engine)
# Pass to specialized function if metadata is defined
if metadata is not None:
# Use the provided parquet metadata object
# to avoid transferring/parsing footer metadata
return _get_parquet_byte_ranges_from_metadata(
metadata,
fs,
engine,
columns=columns,
row_groups=row_groups,
max_gap=max_gap,
max_block=max_block,
)
# Get file sizes asynchronously
file_sizes = fs.sizes(paths)
# Populate global paths, starts, & ends
result = {}
data_paths = []
data_starts = []
data_ends = []
add_header_magic = True
if columns is None and row_groups is None:
# We are NOT selecting specific columns or row-groups.
#
# We can avoid sampling the footers, and just transfer
# all file data with cat_ranges
for i, path in enumerate(paths):
result[path] = {}
for b in range(0, file_sizes[i], max_block):
data_paths.append(path)
data_starts.append(b)
data_ends.append(min(b + max_block, file_sizes[i]))
add_header_magic = False # "Magic" should already be included
else:
# We ARE selecting specific columns or row-groups.
#
# Gather file footers.
# We just take the last `footer_sample_size` bytes of each
# file (or the entire file if it is smaller than that)
footer_starts = []
footer_ends = []
for i, path in enumerate(paths):
footer_ends.append(file_sizes[i])
sample_size = max(0, file_sizes[i] - footer_sample_size)
footer_starts.append(sample_size)
footer_samples = fs.cat_ranges(paths, footer_starts, footer_ends)
# Check our footer samples and re-sample if necessary.
missing_footer_starts = footer_starts.copy()
large_footer = 0
for i, path in enumerate(paths):
footer_size = int.from_bytes(footer_samples[i][-8:-4], "little")
real_footer_start = file_sizes[i] - (footer_size + 8)
if real_footer_start < footer_starts[i]:
missing_footer_starts[i] = real_footer_start
large_footer = max(large_footer, (footer_size + 8))
if large_footer:
warnings.warn(
f"Not enough data was used to sample the parquet footer. "
f"Try setting footer_sample_size >= {large_footer}."
)
for i, block in enumerate(
fs.cat_ranges(
paths,
missing_footer_starts,
footer_starts,
)
):
footer_samples[i] = block + footer_samples[i]
footer_starts[i] = missing_footer_starts[i]
# Calculate required byte ranges for each path
for i, path in enumerate(paths):
# Deal with small-file case.
# Just include all remaining bytes of the file
# in a single range.
if file_sizes[i] < max_block:
if footer_starts[i] > 0:
# Only need to transfer the data if the
# footer sample isn't already the whole file
data_paths.append(path)
data_starts.append(0)
data_ends.append(footer_starts[i])
continue
# Use "engine" to collect data byte ranges
path_data_starts, path_data_ends = engine._parquet_byte_ranges(
columns,
row_groups=row_groups,
footer=footer_samples[i],
footer_start=footer_starts[i],
)
data_paths += [path] * len(path_data_starts)
data_starts += path_data_starts
data_ends += path_data_ends
# Merge adjacent offset ranges
data_paths, data_starts, data_ends = merge_offset_ranges(
data_paths,
data_starts,
data_ends,
max_gap=max_gap,
max_block=max_block,
sort=False, # Should already be sorted
)
# Start by populating `result` with footer samples
for i, path in enumerate(paths):
result[path] = {(footer_starts[i], footer_ends[i]): footer_samples[i]}
# Transfer the data byte-ranges into local memory
_transfer_ranges(fs, result, data_paths, data_starts, data_ends)
# Add b"PAR1" to header if necessary
if add_header_magic:
_add_header_magic(result)
return result
def _get_parquet_byte_ranges_from_metadata(
metadata,
fs,
engine,
columns=None,
row_groups=None,
max_gap=64_000,
max_block=256_000_000,
):
"""Simplified version of `_get_parquet_byte_ranges` for
the case that an engine-specific `metadata` object is
provided, and the remote footer metadata does not need to
be transferred before calculating the required byte ranges.
"""
# Use "engine" to collect data byte ranges
data_paths, data_starts, data_ends = engine._parquet_byte_ranges(
columns,
row_groups=row_groups,
metadata=metadata,
)
# Merge adjacent offset ranges
data_paths, data_starts, data_ends = merge_offset_ranges(
data_paths,
data_starts,
data_ends,
max_gap=max_gap,
max_block=max_block,
sort=False, # Should be sorted
)
# Transfer the data byte-ranges into local memory
result = {fn: {} for fn in list(set(data_paths))}
_transfer_ranges(fs, result, data_paths, data_starts, data_ends)
# Add b"PAR1" to header
_add_header_magic(result)
return result
def _transfer_ranges(fs, blocks, paths, starts, ends):
# Use cat_ranges to gather the data byte_ranges
ranges = (paths, starts, ends)
for path, start, stop, data in zip(*ranges, fs.cat_ranges(*ranges)):
blocks[path][(start, stop)] = data
def _add_header_magic(data):
# Add b"PAR1" to file headers
for path in list(data.keys()):
add_magic = True
for k in data[path].keys():
if k[0] == 0 and k[1] >= 4:
add_magic = False
break
if add_magic:
data[path][(0, 4)] = b"PAR1"
def _set_engine(engine_str):
# Define a list of parquet engines to try
if engine_str == "auto":
try_engines = ("fastparquet", "pyarrow")
elif not isinstance(engine_str, str):
raise ValueError(
"Failed to set parquet engine! "
"Please pass 'fastparquet', 'pyarrow', or 'auto'"
)
elif engine_str not in ("fastparquet", "pyarrow"):
raise ValueError(f"{engine_str} engine not supported by `fsspec.parquet`")
else:
try_engines = [engine_str]
# Try importing the engines in `try_engines`,
# and choose the first one that succeeds
for engine in try_engines:
try:
if engine == "fastparquet":
return FastparquetEngine()
elif engine == "pyarrow":
return PyarrowEngine()
except ImportError:
pass
# Raise an error if a supported parquet engine
# was not found
raise ImportError(
f"The following parquet engines are not installed "
f"in your python environment: {try_engines}."
f"Please install 'fastparquert' or 'pyarrow' to "
f"utilize the `fsspec.parquet` module."
)
class FastparquetEngine:
# The purpose of the FastparquetEngine class is
# to check if fastparquet can be imported (on initialization)
# and to define a `_parquet_byte_ranges` method. In the
# future, this class may also be used to define other
# methods/logic that are specific to fastparquet.
def __init__(self):
import fastparquet as fp
self.fp = fp
def _row_group_filename(self, row_group, pf):
return pf.row_group_filename(row_group)
def _parquet_byte_ranges(
self,
columns,
row_groups=None,
metadata=None,
footer=None,
footer_start=None,
):
# Initialize offset ranges and define ParqetFile metadata
pf = metadata
data_paths, data_starts, data_ends = [], [], []
if pf is None:
pf = self.fp.ParquetFile(io.BytesIO(footer))
# Convert columns to a set and add any index columns
# specified in the pandas metadata (just in case)
column_set = None if columns is None else set(columns)
if column_set is not None and hasattr(pf, "pandas_metadata"):
md_index = [
ind
for ind in pf.pandas_metadata.get("index_columns", [])
# Ignore RangeIndex information
if not isinstance(ind, dict)
]
column_set |= set(md_index)
# Check if row_groups is a list of integers
# or a list of row-group metadata
if row_groups and not isinstance(row_groups[0], int):
# Input row_groups contains row-group metadata
row_group_indices = None
else:
# Input row_groups contains row-group indices
row_group_indices = row_groups
row_groups = pf.row_groups
# Loop through column chunks to add required byte ranges
for r, row_group in enumerate(row_groups):
# Skip this row-group if we are targeting
# specific row-groups
if row_group_indices is None or r in row_group_indices:
# Find the target parquet-file path for `row_group`
fn = self._row_group_filename(row_group, pf)
for column in row_group.columns:
name = column.meta_data.path_in_schema[0]
# Skip this column if we are targeting a
# specific columns
if column_set is None or name in column_set:
file_offset0 = column.meta_data.dictionary_page_offset
if file_offset0 is None:
file_offset0 = column.meta_data.data_page_offset
num_bytes = column.meta_data.total_compressed_size
if footer_start is None or file_offset0 < footer_start:
data_paths.append(fn)
data_starts.append(file_offset0)
data_ends.append(
min(
file_offset0 + num_bytes,
footer_start or (file_offset0 + num_bytes),
)
)
if metadata:
# The metadata in this call may map to multiple
# file paths. Need to include `data_paths`
return data_paths, data_starts, data_ends
return data_starts, data_ends
class PyarrowEngine:
# The purpose of the PyarrowEngine class is
# to check if pyarrow can be imported (on initialization)
# and to define a `_parquet_byte_ranges` method. In the
# future, this class may also be used to define other
# methods/logic that are specific to pyarrow.
def __init__(self):
import pyarrow.parquet as pq
self.pq = pq
def _row_group_filename(self, row_group, metadata):
raise NotImplementedError
def _parquet_byte_ranges(
self,
columns,
row_groups=None,
metadata=None,
footer=None,
footer_start=None,
):
if metadata is not None:
raise ValueError("metadata input not supported for PyarrowEngine")
data_starts, data_ends = [], []
md = self.pq.ParquetFile(io.BytesIO(footer)).metadata
# Convert columns to a set and add any index columns
# specified in the pandas metadata (just in case)
column_set = None if columns is None else set(columns)
if column_set is not None:
schema = md.schema.to_arrow_schema()
has_pandas_metadata = (
schema.metadata is not None and b"pandas" in schema.metadata
)
if has_pandas_metadata:
md_index = [
ind
for ind in json.loads(
schema.metadata[b"pandas"].decode("utf8")
).get("index_columns", [])
# Ignore RangeIndex information
if not isinstance(ind, dict)
]
column_set |= set(md_index)
# Loop through column chunks to add required byte ranges
for r in range(md.num_row_groups):
# Skip this row-group if we are targeting
# specific row-groups
if row_groups is None or r in row_groups:
row_group = md.row_group(r)
for c in range(row_group.num_columns):
column = row_group.column(c)
name = column.path_in_schema
# Skip this column if we are targeting a
# specific columns
split_name = name.split(".")[0]
if (
column_set is None
or name in column_set
or split_name in column_set
):
file_offset0 = column.dictionary_page_offset
if file_offset0 is None:
file_offset0 = column.data_page_offset
num_bytes = column.total_compressed_size
if file_offset0 < footer_start:
data_starts.append(file_offset0)
data_ends.append(
min(file_offset0 + num_bytes, footer_start)
)
return data_starts, data_ends