From 7db6ccec99a314880a5e88f340be3443129bc397 Mon Sep 17 00:00:00 2001 From: Mohit Nirmal Date: Sun, 11 Oct 2026 01:10:24 +0530 Subject: [PATCH] Fix mixed FileInfo and path dataset sources --- python/pyarrow/_dataset.pyx | 7 +- python/pyarrow/dataset.py | 42 +++++++---- python/pyarrow/tests/test_dataset.py | 108 +++++++++++++++++++++++++++ 3 files changed, 142 insertions(+), 15 deletions(-) diff --git a/python/pyarrow/_dataset.pyx b/python/pyarrow/_dataset.pyx index d40614a61fc0..8c3a75b0e3f0 100644 --- a/python/pyarrow/_dataset.pyx +++ b/python/pyarrow/_dataset.pyx @@ -3355,6 +3355,8 @@ cdef class FileSystemFactoryOptions(_Weakrefable): cdef vector[CFileInfo] unwrap_finfos(finfos): cdef vector[CFileInfo] o_vect for fi in finfos: + if not isinstance(fi, FileInfo): + raise TypeError("Expected all items to be FileInfo objects") o_vect.push_back(( fi).unwrap()) return o_vect @@ -3367,8 +3369,9 @@ cdef class FileSystemDatasetFactory(DatasetFactory): ---------- filesystem : pyarrow.fs.FileSystem Filesystem to discover. - paths_or_selector : pyarrow.fs.FileSelector or list of path-likes - Either a Selector object or a list of path-like objects. + paths_or_selector : pyarrow.fs.FileSelector, list of path-likes or list of FileInfo + Either a Selector object, a list of path-like objects, or a list of + FileInfo objects. Lists must not mix path-likes and FileInfo objects. format : FileFormat Currently only ParquetFileFormat and IpcFileFormat are supported. options : FileSystemFactoryOptions, optional diff --git a/python/pyarrow/dataset.py b/python/pyarrow/dataset.py index f696d5f03ba5..c55c544c4aa8 100644 --- a/python/pyarrow/dataset.py +++ b/python/pyarrow/dataset.py @@ -319,7 +319,7 @@ def _ensure_format(obj): def _ensure_multiple_sources(paths, filesystem=None): """ - Treat a list of paths as files belonging to a single file system + Treat a list of paths or FileInfo objects as files on one file system. If the file system is local then also validates that all paths are referencing existing *files* otherwise any non-file paths will be @@ -327,7 +327,7 @@ def _ensure_multiple_sources(paths, filesystem=None): Parameters ---------- - paths : list of path-like + paths : list of path-like or FileInfo Note that URIs are not allowed. filesystem : FileSystem or str, optional If an URI is passed, then its path component will act as a prefix for @@ -335,8 +335,9 @@ def _ensure_multiple_sources(paths, filesystem=None): Returns ------- - (FileSystem, list of str) - File system object and a list of normalized paths. + (FileSystem, list of str or FileInfo) + File system object and a list of normalized paths, or FileInfo + objects when both kinds of source are present. Raises ------ @@ -347,7 +348,7 @@ def _ensure_multiple_sources(paths, filesystem=None): not a file. """ from pyarrow.fs import ( - LocalFileSystem, SubTreeFileSystem, _MockFileSystem, FileType, + LocalFileSystem, SubTreeFileSystem, _MockFileSystem, FileInfo, FileType, _ensure_filesystem ) @@ -364,14 +365,24 @@ def _ensure_multiple_sources(paths, filesystem=None): isinstance(filesystem.base_fs, LocalFileSystem)) ) - # allow normalizing irregular paths such as Windows local paths - paths = [filesystem.normalize_path(_stringify_path(p)) for p in paths] + # Preserve supplied FileInfo objects (including their size and mtime). + # Normalize path-like entries as for a list containing only paths. + has_file_infos = any(isinstance(p, FileInfo) for p in paths) + if has_file_infos: + paths = [p if isinstance(p, FileInfo) else + filesystem.normalize_path(_stringify_path(p)) for p in paths] + path_entries = [p for p in paths if not isinstance(p, FileInfo)] + else: + paths = [filesystem.normalize_path(_stringify_path(p)) for p in paths] + path_entries = paths # validate that all of the paths are pointing to existing *files* # possible improvement is to group the file_infos by type and raise for # multiple paths per error category + if is_local or has_file_infos: + path_infos = filesystem.get_file_info(path_entries) if is_local: - for info in filesystem.get_file_info(paths): + for info in path_infos: file_type = info.type if file_type == FileType.File: continue @@ -389,6 +400,11 @@ def _ensure_multiple_sources(paths, filesystem=None): 'special file such as a Unix socket or character device, ' 'or Windows NUL / CON / ...)') + if has_file_infos: + path_infos = iter(path_infos) + paths = [p if isinstance(p, FileInfo) else next(path_infos) + for p in paths] + return filesystem, paths @@ -458,7 +474,7 @@ def _filesystem_dataset(source, schema=None, filesystem=None, partitioning = _ensure_partitioning(partitioning) if isinstance(source, (list, tuple)): - if source and isinstance(source[0], FileInfo): + if source and all(isinstance(item, FileInfo) for item in source): if filesystem is None: # fall back to local file system as the default fs = LocalFileSystem() @@ -598,17 +614,17 @@ def dataset(source, schema=None, format=None, filesystem=None, Parameters ---------- - source : path, list of paths, dataset, list of datasets, (list of) \ -RecordBatch or Table, iterable of RecordBatch, RecordBatchReader, or URI + source : path, list of paths or FileInfo, dataset, list of datasets, \ +(list of) RecordBatch or Table, iterable of RecordBatch, RecordBatchReader, or URI Path pointing to a single file: Open a FileSystemDataset from a single file. Path pointing to a directory: The directory gets discovered recursively according to a partitioning scheme if given. - List of file paths: + List of file paths or FileInfo objects: Create a FileSystemDataset from explicitly given files. The files must be located on the same filesystem given by the filesystem - parameter. + parameter. Paths and FileInfo objects may be mixed. Note that in contrary of construction from a single file, passing URIs as paths is not allowed. List of datasets: diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index 9ddfdfdc9c49..7d7dae54e76f 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -2562,6 +2562,114 @@ def test_construct_from_list_of_files(tempdir, dataset_reader): assert t1 == t2 == t3 == t4 +@pytest.mark.parquet +@pytest.mark.parametrize('container', [list, tuple]) +@pytest.mark.parametrize('source_kinds', [ + ('path', 'path'), + ('info', 'info'), + ('path', 'info'), + ('info', 'path'), +]) +@pytest.mark.parametrize('duplicate', [False, True]) +def test_construct_from_paths_and_file_infos(mockfs, container, + source_kinds, duplicate): + paths = ['subdir/1/xxx/file0.parquet', + 'subdir/2/yyy/file1.parquet'] + if duplicate: + paths[1] = paths[0] + infos = mockfs.get_file_info(paths) + sources = container(path if kind == 'path' else info + for path, info, kind in zip(paths, infos, source_kinds)) + + dataset = ds.dataset(sources, filesystem=mockfs) + + assert dataset.files == paths + expected = [0] * 5 + [0 if duplicate else 1] * 5 + assert dataset.to_table()['const'].to_pylist() == expected + + +@pytest.mark.parquet +def test_construct_from_interleaved_paths_and_file_infos(mockfs): + path0 = 'subdir/1/xxx/file0.parquet' + path1 = 'subdir/2/yyy/file1.parquet' + info1 = mockfs.get_file_info(path1) + + dataset = ds.dataset([path0, info1, path0], filesystem=mockfs) + + assert dataset.files == [path0, path1, path0] + assert dataset.to_table()['const'].to_pylist() == [0] * 5 + [1] * 5 + [0] * 5 + + +@pytest.mark.parquet +@pytest.mark.parametrize('source_kinds', [('path', 'info'), + ('info', 'path')]) +def test_construct_from_mixed_sources_with_subtree_filesystem(mockfs, + source_kinds): + filesystem = fs.SubTreeFileSystem('subdir', mockfs) + paths = ['1/xxx/file0.parquet', '2/yyy/file1.parquet'] + infos = filesystem.get_file_info(paths) + sources = [path if kind == 'path' else info + for path, info, kind in zip(paths, infos, source_kinds)] + + dataset = ds.dataset(sources, filesystem=filesystem) + + assert dataset.files == paths + assert dataset.to_table()['const'].to_pylist() == [0] * 5 + [1] * 5 + + +@pytest.mark.parquet +def test_mixed_sources_preserve_file_info(mockfs): + path = 'subdir/1/xxx/file0.parquet' + info = mockfs.get_file_info(path) + + _, sources = ds._ensure_multiple_sources([info, path], mockfs) + + assert sources[0] is info + assert sources[1].path == path + + +@pytest.mark.parquet +@pytest.mark.parametrize('source_kinds', [('path', 'info'), + ('info', 'path')]) +def test_construct_from_mixed_sources_missing_with_subtree(mockfs, + source_kinds): + filesystem = fs.SubTreeFileSystem('subdir', mockfs) + info = filesystem.get_file_info('1/xxx/file0.parquet') + sources = ['missing.parquet' if kind == 'path' else info + for kind in source_kinds] + + with pytest.raises(FileNotFoundError): + ds.dataset(sources, filesystem=filesystem).to_table() + + +@pytest.mark.parquet +@pytest.mark.parametrize('source_kinds', [('path', 'info'), + ('info', 'path')]) +@pytest.mark.parametrize('invalid, error', [ + ('subdir/1/xxx/doesnt-exist.parquet', FileNotFoundError), + ('subdir/1/xxx', IsADirectoryError), +]) +def test_construct_from_mixed_sources_invalid_path(mockfs, source_kinds, + invalid, error): + path = 'subdir/1/xxx/file0.parquet' + info = mockfs.get_file_info(path) + sources = [invalid if kind == 'path' else info + for kind in source_kinds] + + with pytest.raises(error): + ds.dataset(sources, filesystem=mockfs) + + +@pytest.mark.parquet +def test_filesystem_factory_rejects_mixed_file_infos(mockfs): + path = 'subdir/1/xxx/file0.parquet' + info = mockfs.get_file_info(path) + + with pytest.raises(TypeError, match='FileInfo'): + ds.FileSystemDatasetFactory( + mockfs, [info, path], ds.ParquetFileFormat()) + + @pytest.mark.parquet def test_construct_from_list_of_mixed_paths_fails(mockfs): # instantiate from a list of mixed paths