GH-53033: [Python][Dataset] Support mixed FileInfo and path sources without crashing

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((<FileInfo> 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

python/pyarrow/dataset.py

@@ -319,24 +319,25 @@ 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
     silently skipped (for example on a remote filesystem).
 
     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
         the file paths.
 
     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:

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