def _resolve_paths_and_filesystem(
paths: Union[str, List[str]],
filesystem: "pyarrow.fs.FileSystem" = None,
) -> Tuple[List[str], "pyarrow.fs.FileSystem"]:
"""
Resolves and normalizes all provided paths, infers a filesystem from the
paths and ensures that all paths use the same filesystem.
Args:
paths: A single file/directory path or a list of file/directory paths.
A list of paths can contain both files and directories.
filesystem: The filesystem implementation that should be used for
reading these files. If None, a filesystem will be inferred. If not
None, the provided filesystem will still be validated against all
filesystems inferred from the provided paths to ensure
compatibility.
"""
import pyarrow as pa
from pyarrow.fs import (FileSystem, PyFileSystem, FSSpecHandler,
_resolve_filesystem_and_path)
import fsspec
if isinstance(paths, str):
paths = [paths]
elif (not isinstance(paths, list)
or any(not isinstance(p, str) for p in paths)):
raise ValueError(
"paths must be a path string or a list of path strings.")
elif len(paths) == 0:
raise ValueError("Must provide at least one path.")
if filesystem and not isinstance(filesystem, FileSystem):
if not isinstance(filesystem, fsspec.spec.AbstractFileSystem):
raise TypeError(f"The filesystem passed must either conform to "
f"pyarrow.fs.FileSystem, or "
f"fsspec.spec.AbstractFileSystem. The provided "
f"filesystem was: {filesystem}")
filesystem = PyFileSystem(FSSpecHandler(filesystem))
resolved_paths = []
for path in paths:
try:
resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
path, filesystem)
except pa.lib.ArrowInvalid as e:
if "Cannot parse URI" in str(e):
resolved_filesystem, resolved_path = (
_resolve_filesystem_and_path(_encode_url(path),
filesystem))
resolved_path = _decode_url(resolved_path)
else:
raise
if filesystem is None:
filesystem = resolved_filesystem
else:
resolved_path = _unwrap_protocol(resolved_path)
resolved_path = filesystem.normalize_path(resolved_path)
resolved_paths.append(resolved_path)
return resolved_paths, filesystem
def read_table(source, columns=None, filesystem=None):
"""
Read a Table from an ORC file.
Parameters
----------
source : str, pyarrow.NativeFile, or file-like object
If a string passed, can be a single file name. For file-like objects,
only read a single file. Use pyarrow.BufferReader to read a file
contained in a bytes or buffer-like object.
columns : list
If not None, only these columns will be read from the file. A column
name may be a prefix of a nested field, e.g. 'a' will select 'a.b',
'a.c', and 'a.d.e'. If empty, no columns will be read. Note
that the table will still have the correct num_rows set despite having
no columns.
filesystem : FileSystem, default None
If nothing passed, paths assumed to be found in the local on-disk
filesystem.
"""
filesystem, path = _resolve_filesystem_and_path(source, filesystem)
if filesystem is not None:
source = filesystem.open_input_file(path)
if columns is not None and len(columns) == 0:
result = ORCFile(source).read().select(columns)
else:
result = ORCFile(source).read(columns=columns)
return result
def _resolve_paths_and_filesystem(
paths: Union[str, List[str]],
filesystem: "pyarrow.fs.FileSystem" = None
) -> Tuple[List[str], "pyarrow.fs.FileSystem"]:
"""
Resolves and normalizes all provided paths, infers a filesystem from the
paths and ensures that all paths use the same filesystem, and expands all
directory paths to the underlying file paths.
Args:
paths: A single file/directory path or a list of file/directory paths.
A list of paths can contain both files and directories.
filesystem: The filesystem implementation that should be used for
reading these files. If None, a filesystem will be inferred. If not
None, the provided filesystem will still be validated against all
filesystems inferred from the provided paths to ensure
compatibility.
"""
from pyarrow.fs import FileType, _resolve_filesystem_and_path
if isinstance(paths, str):
paths = [paths]
elif (not isinstance(paths, list)
or any(not isinstance(p, str) for p in paths)):
raise ValueError(
"paths must be a path string or a list of path strings.")
elif len(paths) == 0:
raise ValueError("Must provide at least one path.")
resolved_paths = []
for path in paths:
if filesystem is not None:
# If we provide a filesystem, _resolve_filesystem_and_path will not
# slice off the protocol from the provided URI/path when resolved.
path = _unwrap_protocol(path)
resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
path, filesystem)
if filesystem is None:
filesystem = resolved_filesystem
elif type(resolved_filesystem) != type(filesystem):
raise ValueError("All paths must use same filesystem.")
resolved_path = filesystem.normalize_path(resolved_path)
resolved_paths.append(resolved_path)
expanded_paths = []
file_infos = []
for path in resolved_paths:
file_info = filesystem.get_file_info(path)
if file_info.type == FileType.Directory:
paths, file_infos_ = _expand_directory(path, filesystem)
expanded_paths.extend(paths)
file_infos.extend(file_infos_)
elif file_info.type == FileType.File:
expanded_paths.append(path)
file_infos.append(file_info)
else:
raise FileNotFoundError(path)
return expanded_paths, file_infos, filesystem
def read_table(source, columns=None, filesystem=None):
filesystem, path = _resolve_filesystem_and_path(source, filesystem)
if filesystem is not None:
source = filesystem.open_input_file(path)
if columns is not None and len(columns) == 0:
result = ORCFile(source).read().select(columns)
else:
result = ORCFile(source).read(columns=columns)
return result
def _ensure_single_source(path, filesystem=None):
"""
Treat path as either a recursively traversable directory or a single file.
Parameters
----------
path : path-like
filesystem : FileSystem or str, optional
If an URI is passed, then its path component will act as a prefix for
the file paths.
Returns
-------
(FileSystem, list of str or fs.Selector)
File system object and either a single item list pointing to a file or
an fs.Selector object pointing to a directory.
Raises
------
TypeError
If the passed filesystem has wrong type.
FileNotFoundError
If the referenced file or directory doesn't exist.
"""
from pyarrow.fs import FileType, FileSelector, _resolve_filesystem_and_path
# at this point we already checked that `path` is a path-like
filesystem, path = _resolve_filesystem_and_path(path, filesystem)
# ensure that the path is normalized before passing to dataset discovery
path = filesystem.normalize_path(path)
# retrieve the file descriptor
file_info = filesystem.get_file_info(path)
# depending on the path type either return with a recursive
# directory selector or as a list containing a single file
if file_info.type == FileType.Directory:
paths_or_selector = FileSelector(path, recursive=True)
elif file_info.type == FileType.File:
paths_or_selector = [path]
else:
raise FileNotFoundError(path)
return filesystem, paths_or_selector
def write_dataset(data,
base_dir,
basename_template=None,
format=None,
partitioning=None,
schema=None,
filesystem=None,
file_options=None,
use_threads=True,
use_async=False,
max_partitions=None,
file_visitor=None):
"""
Write a dataset to a given format and partitioning.
Parameters
----------
data : Dataset, Table/RecordBatch, RecordBatchReader, list of
Table/RecordBatch, or iterable of RecordBatch
The data to write. This can be a Dataset instance or
in-memory Arrow data. If an iterable is given, the schema must
also be given.
base_dir : str
The root directory where to write the dataset.
basename_template : str, optional
A template string used to generate basenames of written data files.
The token '{i}' will be replaced with an automatically incremented
integer. If not specified, it defaults to
"part-{i}." + format.default_extname
format : FileFormat or str
The format in which to write the dataset. Currently supported:
"parquet", "ipc"/"feather". If a FileSystemDataset is being written
and `format` is not specified, it defaults to the same format as the
specified FileSystemDataset. When writing a Table or RecordBatch, this
keyword is required.
partitioning : Partitioning, optional
The partitioning scheme specified with the ``partitioning()``
function.
schema : Schema, optional
filesystem : FileSystem, optional
file_options : FileWriteOptions, optional
FileFormat specific write options, created using the
``FileFormat.make_write_options()`` function.
use_threads : bool, default True
Write files in parallel. If enabled, then maximum parallelism will be
used determined by the number of available CPU cores.
use_async : bool, default False
If enabled, an async scanner will be used that should offer
better performance with high-latency/highly-parallel filesystems
(e.g. S3)
max_partitions : int, default 1024
Maximum number of partitions any batch may be written into.
file_visitor : Function
If set, this function will be called with a WrittenFile instance
for each file created during the call. This object will have both
a path attribute and a metadata attribute.
The path attribute will be a string containing the path to
the created file.
The metadata attribute will be the parquet metadata of the file.
This metadata will have the file path attribute set and can be used
to build a _metadata file. The metadata attribute will be None if
the format is not parquet.
Example visitor which simple collects the filenames created::
visited_paths = []
def file_visitor(written_file):
visited_paths.append(written_file.path)
"""
from pyarrow.fs import _resolve_filesystem_and_path
if isinstance(data, (list, tuple)):
schema = schema or data[0].schema
data = InMemoryDataset(data, schema=schema)
elif isinstance(data, (pa.RecordBatch, pa.Table)):
schema = schema or data.schema
data = InMemoryDataset(data, schema=schema)
elif isinstance(data, pa.ipc.RecordBatchReader) or _is_iterable(data):
data = Scanner.from_batches(data, schema=schema)
schema = None
elif not isinstance(data, (Dataset, Scanner)):
raise ValueError(
"Only Dataset, Scanner, Table/RecordBatch, RecordBatchReader, "
"a list of Tables/RecordBatches, or iterable of batches are "
"supported.")
if format is None and isinstance(data, FileSystemDataset):
format = data.format
else:
format = _ensure_format(format)
if file_options is None:
file_options = format.make_write_options()
if format != file_options.format:
raise TypeError("Supplied FileWriteOptions have format {}, "
"which doesn't match supplied FileFormat {}".format(
format, file_options))
if basename_template is None:
basename_template = "part-{i}." + format.default_extname
if max_partitions is None:
max_partitions = 1024
partitioning = _ensure_write_partitioning(partitioning)
filesystem, base_dir = _resolve_filesystem_and_path(base_dir, filesystem)
if isinstance(data, Dataset):
scanner = data.scanner(use_threads=use_threads, use_async=use_async)
else:
# scanner was passed directly by the user, in which case a schema
# cannot be passed
if schema is not None:
raise ValueError("Cannot specify a schema when writing a Scanner")
scanner = data
_filesystemdataset_write(scanner, base_dir, basename_template, filesystem,
partitioning, file_options, max_partitions,
file_visitor)
def write_dataset(data,
base_dir,
basename_template=None,
format=None,
partitioning=None,
schema=None,
filesystem=None,
file_options=None,
use_threads=True,
max_partitions=None):
"""
Write a dataset to a given format and partitioning.
Parameters
----------
data : Dataset, Table/RecordBatch, RecordBatchReader, list of
Table/RecordBatch, or iterable of RecordBatch
The data to write. This can be a Dataset instance or
in-memory Arrow data. If an iterable is given, the schema must
also be given.
base_dir : str
The root directory where to write the dataset.
basename_template : str, optional
A template string used to generate basenames of written data files.
The token '{i}' will be replaced with an automatically incremented
integer. If not specified, it defaults to
"part-{i}." + format.default_extname
format : FileFormat or str
The format in which to write the dataset. Currently supported:
"parquet", "ipc"/"feather". If a FileSystemDataset is being written
and `format` is not specified, it defaults to the same format as the
specified FileSystemDataset. When writing a Table or RecordBatch, this
keyword is required.
partitioning : Partitioning, optional
The partitioning scheme specified with the ``partitioning()``
function.
schema : Schema, optional
filesystem : FileSystem, optional
file_options : FileWriteOptions, optional
FileFormat specific write options, created using the
``FileFormat.make_write_options()`` function.
use_threads : bool, default True
Write files in parallel. If enabled, then maximum parallelism will be
used determined by the number of available CPU cores.
max_partitions : int, default 1024
Maximum number of partitions any batch may be written into.
"""
from pyarrow.fs import _resolve_filesystem_and_path
if isinstance(data, Dataset):
schema = schema or data.schema
elif isinstance(data, (list, tuple)):
schema = schema or data[0].schema
data = InMemoryDataset(data, schema=schema)
elif isinstance(data, (pa.RecordBatch, pa.ipc.RecordBatchReader,
pa.Table)) or _is_iterable(data):
data = InMemoryDataset(data, schema=schema)
schema = schema or data.schema
else:
raise ValueError(
"Only Dataset, Table/RecordBatch, RecordBatchReader, a list "
"of Tables/RecordBatches, or iterable of batches are supported.")
if format is None and isinstance(data, FileSystemDataset):
format = data.format
else:
format = _ensure_format(format)
if file_options is None:
file_options = format.make_write_options()
if format != file_options.format:
raise TypeError("Supplied FileWriteOptions have format {}, "
"which doesn't match supplied FileFormat {}".format(
format, file_options))
if basename_template is None:
basename_template = "part-{i}." + format.default_extname
if max_partitions is None:
max_partitions = 1024
partitioning = _ensure_write_partitioning(partitioning)
filesystem, base_dir = _resolve_filesystem_and_path(base_dir, filesystem)
_filesystemdataset_write(data, base_dir, basename_template, schema,
filesystem, partitioning, file_options,
use_threads, max_partitions)
def write_dataset(data, base_dir, basename_template=None, format=None,
partitioning=None, partitioning_flavor=None, schema=None,
filesystem=None, file_options=None, use_threads=True,
max_partitions=None, max_open_files=None,
max_rows_per_file=None, min_rows_per_group=None,
max_rows_per_group=None, file_visitor=None,
existing_data_behavior='error', create_dir=True):
"""
Write a dataset to a given format and partitioning.
Parameters
----------
data : Dataset, Table/RecordBatch, RecordBatchReader, list of \
Table/RecordBatch, or iterable of RecordBatch
The data to write. This can be a Dataset instance or
in-memory Arrow data. If an iterable is given, the schema must
also be given.
base_dir : str
The root directory where to write the dataset.
basename_template : str, optional
A template string used to generate basenames of written data files.
The token '{i}' will be replaced with an automatically incremented
integer. If not specified, it defaults to
"part-{i}." + format.default_extname
format : FileFormat or str
The format in which to write the dataset. Currently supported:
"parquet", "ipc"/"arrow"/"feather", and "csv". If a FileSystemDataset
is being written and `format` is not specified, it defaults to the
same format as the specified FileSystemDataset. When writing a
Table or RecordBatch, this keyword is required.
partitioning : Partitioning or list[str], optional
The partitioning scheme specified with the ``partitioning()``
function or a list of field names. When providing a list of
field names, you can use ``partitioning_flavor`` to drive which
partitioning type should be used.
partitioning_flavor : str, optional
One of the partitioning flavors supported by
``pyarrow.dataset.partitioning``. If omitted will use the
default of ``partitioning()`` which is directory partitioning.
schema : Schema, optional
filesystem : FileSystem, optional
file_options : pyarrow.dataset.FileWriteOptions, optional
FileFormat specific write options, created using the
``FileFormat.make_write_options()`` function.
use_threads : bool, default True
Write files in parallel. If enabled, then maximum parallelism will be
used determined by the number of available CPU cores.
max_partitions : int, default 1024
Maximum number of partitions any batch may be written into.
max_open_files : int, default 1024
If greater than 0 then this will limit the maximum number of
files that can be left open. If an attempt is made to open
too many files then the least recently used file will be closed.
If this setting is set too low you may end up fragmenting your
data into many small files.
max_rows_per_file : int, default 0
Maximum number of rows per file. If greater than 0 then this will
limit how many rows are placed in any single file. Otherwise there
will be no limit and one file will be created in each output
directory unless files need to be closed to respect max_open_files
min_rows_per_group : int, default 0
Minimum number of rows per group. When the value is greater than 0,
the dataset writer will batch incoming data and only write the row
groups to the disk when sufficient rows have accumulated.
max_rows_per_group : int, default 1024 * 1024
Maximum number of rows per group. If the value is greater than 0,
then the dataset writer may split up large incoming batches into
multiple row groups. If this value is set, then min_rows_per_group
should also be set. Otherwise it could end up with very small row
groups.
file_visitor : function
If set, this function will be called with a WrittenFile instance
for each file created during the call. This object will have both
a path attribute and a metadata attribute.
The path attribute will be a string containing the path to
the created file.
The metadata attribute will be the parquet metadata of the file.
This metadata will have the file path attribute set and can be used
to build a _metadata file. The metadata attribute will be None if
the format is not parquet.
Example visitor which simple collects the filenames created::
visited_paths = []
def file_visitor(written_file):
visited_paths.append(written_file.path)
existing_data_behavior : 'error' | 'overwrite_or_ignore' | \
'delete_matching'
Controls how the dataset will handle data that already exists in
the destination. The default behavior ('error') is to raise an error
if any data exists in the destination.
'overwrite_or_ignore' will ignore any existing data and will
overwrite files with the same name as an output file. Other
existing files will be ignored. This behavior, in combination
with a unique basename_template for each write, will allow for
an append workflow.
'delete_matching' is useful when you are writing a partitioned
dataset. The first time each partition directory is encountered
the entire directory will be deleted. This allows you to overwrite
old partitions completely.
create_dir : bool, default True
If False, directories will not be created. This can be useful for
filesystems that do not require directories.
"""
from pyarrow.fs import _resolve_filesystem_and_path
if isinstance(data, (list, tuple)):
schema = schema or data[0].schema
data = InMemoryDataset(data, schema=schema)
elif isinstance(data, (pa.RecordBatch, pa.Table)):
schema = schema or data.schema
data = InMemoryDataset(data, schema=schema)
elif isinstance(data, pa.ipc.RecordBatchReader) or _is_iterable(data):
data = Scanner.from_batches(data, schema=schema)
schema = None
elif not isinstance(data, (Dataset, Scanner)):
raise ValueError(
"Only Dataset, Scanner, Table/RecordBatch, RecordBatchReader, "
"a list of Tables/RecordBatches, or iterable of batches are "
"supported."
)
if format is None and isinstance(data, FileSystemDataset):
format = data.format
else:
format = _ensure_format(format)
if file_options is None:
file_options = format.make_write_options()
if format != file_options.format:
raise TypeError("Supplied FileWriteOptions have format {}, "
"which doesn't match supplied FileFormat {}".format(
format, file_options))
if basename_template is None:
basename_template = "part-{i}." + format.default_extname
if max_partitions is None:
max_partitions = 1024
if max_open_files is None:
max_open_files = 1024
if max_rows_per_file is None:
max_rows_per_file = 0
if max_rows_per_group is None:
max_rows_per_group = 1 << 20
if min_rows_per_group is None:
min_rows_per_group = 0
# at this point data is a Scanner or a Dataset, anything else
# was converted to one of those two. So we can grab the schema
# to build the partitioning object from Dataset.
if isinstance(data, Scanner):
partitioning_schema = data.dataset_schema
else:
partitioning_schema = data.schema
partitioning = _ensure_write_partitioning(partitioning,
schema=partitioning_schema,
flavor=partitioning_flavor)
filesystem, base_dir = _resolve_filesystem_and_path(base_dir, filesystem)
if isinstance(data, Dataset):
scanner = data.scanner(use_threads=use_threads)
else:
# scanner was passed directly by the user, in which case a schema
# cannot be passed
if schema is not None:
raise ValueError("Cannot specify a schema when writing a Scanner")
scanner = data
_filesystemdataset_write(
scanner, base_dir, basename_template, filesystem, partitioning,
file_options, max_partitions, file_visitor, existing_data_behavior,
max_open_files, max_rows_per_file,
min_rows_per_group, max_rows_per_group, create_dir
)
def _resolve_paths_and_filesystem(
paths: Union[str, List[str]],
filesystem: "pyarrow.fs.FileSystem" = None,
) -> Tuple[List[str], "pyarrow.fs.FileSystem"]:
"""
Resolves and normalizes all provided paths, infers a filesystem from the
paths and ensures that all paths use the same filesystem.
Args:
paths: A single file/directory path or a list of file/directory paths.
A list of paths can contain both files and directories.
filesystem: The filesystem implementation that should be used for
reading these files. If None, a filesystem will be inferred. If not
None, the provided filesystem will still be validated against all
filesystems inferred from the provided paths to ensure
compatibility.
"""
import pyarrow as pa
from pyarrow.fs import (
FileSystem,
FSSpecHandler,
PyFileSystem,
_resolve_filesystem_and_path,
)
if isinstance(paths, str):
paths = [paths]
elif not isinstance(paths, list) or any(not isinstance(p, str)
for p in paths):
raise ValueError(
"paths must be a path string or a list of path strings.")
elif len(paths) == 0:
raise ValueError("Must provide at least one path.")
need_unwrap_path_protocol = True
if filesystem and not isinstance(filesystem, FileSystem):
err_msg = (f"The filesystem passed must either conform to "
f"pyarrow.fs.FileSystem, or "
f"fsspec.spec.AbstractFileSystem. The provided "
f"filesystem was: {filesystem}")
try:
import fsspec
from fsspec.implementations.http import HTTPFileSystem
except ModuleNotFoundError:
# If filesystem is not a pyarrow filesystem and fsspec isn't
# installed, then filesystem is neither a pyarrow filesystem nor
# an fsspec filesystem, so we raise a TypeError.
raise TypeError(err_msg) from None
if not isinstance(filesystem, fsspec.spec.AbstractFileSystem):
raise TypeError(err_msg) from None
if isinstance(filesystem, HTTPFileSystem):
# If filesystem is fsspec HTTPFileSystem, the protocol/scheme of paths
# should not be unwrapped/removed, because HTTPFileSystem expects full file
# paths including protocol/scheme. This is different behavior compared to
# file systems implementation in pyarrow.fs.FileSystem.
need_unwrap_path_protocol = False
filesystem = PyFileSystem(FSSpecHandler(filesystem))
resolved_paths = []
for path in paths:
path = _resolve_example_path(path)
try:
resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
path, filesystem)
except pa.lib.ArrowInvalid as e:
if "Cannot parse URI" in str(e):
resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
_encode_url(path), filesystem)
resolved_path = _decode_url(resolved_path)
elif "Unrecognized filesystem type in URI" in str(e):
scheme = urllib.parse.urlparse(path,
allow_fragments=False).scheme
if scheme in ["http", "https"]:
# If scheme of path is HTTP and filesystem is not resolved,
# try to use fsspec HTTPFileSystem. This expects fsspec is
# installed.
try:
from fsspec.implementations.http import HTTPFileSystem
except ModuleNotFoundError:
raise ImportError(
"Please install fsspec to read files from HTTP."
) from None
resolved_filesystem = PyFileSystem(
FSSpecHandler(HTTPFileSystem()))
resolved_path = path
need_unwrap_path_protocol = False
else:
raise
else:
raise
if filesystem is None:
filesystem = resolved_filesystem
elif need_unwrap_path_protocol:
resolved_path = _unwrap_protocol(resolved_path)
resolved_path = filesystem.normalize_path(resolved_path)
resolved_paths.append(resolved_path)
return resolved_paths, filesystem
