wehub-resource-sync
164 lines
6.8 KiB
Python
164 lines
6.8 KiB
Python
from unittest import mock
|
|
|
|
import pyarrow
|
|
import pyarrow.fs
|
|
import pytest
|
|
from fsspec.implementations.http import HTTPFileSystem
|
|
from pyarrow.fs import FSSpecHandler, PyFileSystem
|
|
|
|
from ray.data._internal.util import RetryingPyFileSystem
|
|
from ray.data.datasource.path_util import (
|
|
_has_file_extension,
|
|
_is_filesystem_compatible_with_scheme,
|
|
_is_local_windows_path,
|
|
_resolve_paths_and_filesystem,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"path, extensions, has_extension",
|
|
[
|
|
("foo.csv", ["csv"], True),
|
|
("foo.csv", ["json", "csv"], True),
|
|
("foo.csv", ["json", "jsonl"], False),
|
|
("foo.csv", [".csv"], True),
|
|
("foo.parquet.crc", ["parquet"], False),
|
|
("foo.parquet.crc", ["crc"], True),
|
|
("s3://bucket/foo.parquet?versionId=abc123", ["parquet"], True),
|
|
("bucket/foo.parquet?versionId=abc123", ["parquet"], True),
|
|
("s3://bucket/foo.parquet?versionId=abc123", ["csv"], False),
|
|
("s3://bucket/data#v2/file.parquet", ["parquet"], True),
|
|
("C:\\data\\test.parquet", ["parquet"], True),
|
|
("C:\\data\\parquet", ["parquet"], False),
|
|
("foo.csv", None, True),
|
|
],
|
|
)
|
|
def test_has_file_extension(path, extensions, has_extension):
|
|
assert _has_file_extension(path, extensions) == has_extension
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"filesystem", [None, PyFileSystem(FSSpecHandler(HTTPFileSystem()))]
|
|
)
|
|
def test_resolve_http_paths(filesystem):
|
|
resolved_paths, resolve_filesystem = _resolve_paths_and_filesystem(
|
|
"https://google.com", filesystem
|
|
)
|
|
# `_resolve_paths_and_filesystem` shouldn't remove the protocol/scheme from the
|
|
# path for HTTP paths.
|
|
assert resolved_paths == ["https://google.com"]
|
|
assert isinstance(resolve_filesystem, pyarrow.fs.PyFileSystem)
|
|
assert isinstance(resolve_filesystem.handler, pyarrow.fs.FSSpecHandler)
|
|
assert isinstance(resolve_filesystem.handler.fs, HTTPFileSystem)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"path",
|
|
[
|
|
"c:/some/where",
|
|
"c:\\some\\where",
|
|
"c:\\some\\where/mixed",
|
|
],
|
|
)
|
|
def test_windows_path(path):
|
|
with mock.patch("sys.platform", "win32"):
|
|
assert _is_local_windows_path(path)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"path",
|
|
[
|
|
"some/file",
|
|
"some/file;semicolon",
|
|
"some/file?questionmark",
|
|
"some/file#hash",
|
|
"some/file;all?of the#above",
|
|
],
|
|
)
|
|
def test_weird_local_paths(path):
|
|
resolved_paths, _ = _resolve_paths_and_filesystem(path)
|
|
assert resolved_paths[0] == path
|
|
|
|
|
|
class TestIsFilesystemCompatibleWithScheme:
|
|
"""Tests for _is_filesystem_compatible_with_scheme with real filesystem implementations."""
|
|
|
|
def test_native_local_filesystem(self):
|
|
"""Native PyArrow LocalFileSystem should be compatible with empty scheme."""
|
|
fs = pyarrow.fs.LocalFileSystem()
|
|
assert _is_filesystem_compatible_with_scheme(fs, "") is True
|
|
assert _is_filesystem_compatible_with_scheme(fs, "s3") is False
|
|
|
|
def test_native_s3_filesystem(self):
|
|
"""Native PyArrow S3FileSystem should be compatible with s3 scheme."""
|
|
fs = pyarrow.fs.S3FileSystem(anonymous=True, region="us-east-1")
|
|
assert _is_filesystem_compatible_with_scheme(fs, "s3") is True
|
|
assert _is_filesystem_compatible_with_scheme(fs, "") is False
|
|
|
|
def test_fsspec_s3_filesystem_with_s3_scheme(self):
|
|
"""fsspec S3FileSystem wrapped in PyFileSystem should be compatible with s3 scheme."""
|
|
s3fs = pytest.importorskip("s3fs")
|
|
wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "s3") is True
|
|
|
|
def test_fsspec_s3_filesystem_with_bare_paths(self):
|
|
"""fsspec S3FileSystem wrapped in PyFileSystem should be compatible with bare paths."""
|
|
s3fs = pytest.importorskip("s3fs")
|
|
wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "") is True
|
|
|
|
def test_fsspec_s3_filesystem_with_s3a_scheme(self):
|
|
"""Real s3fs has protocol=('s3', 's3a'), so it should match s3a scheme."""
|
|
s3fs = pytest.importorskip("s3fs")
|
|
wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "s3a") is True
|
|
|
|
def test_fsspec_gcs_filesystem_with_gs_scheme(self):
|
|
"""fsspec GCS filesystem should be compatible with gs scheme."""
|
|
gcsfs = pytest.importorskip("gcsfs")
|
|
wrapped_fs = PyFileSystem(FSSpecHandler(gcsfs.GCSFileSystem(token="anon")))
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "gs") is True
|
|
|
|
def test_fsspec_gcs_filesystem_with_gcs_scheme(self):
|
|
"""Real gcsfs has protocol=('gcs', 'gs'), so it should match gcs scheme."""
|
|
gcsfs = pytest.importorskip("gcsfs")
|
|
wrapped_fs = PyFileSystem(FSSpecHandler(gcsfs.GCSFileSystem(token="anon")))
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "gcs") is True
|
|
|
|
def test_fsspec_http_filesystem(self):
|
|
"""fsspec HTTPFileSystem wrapped in PyFileSystem should be compatible with http."""
|
|
wrapped_fs = PyFileSystem(FSSpecHandler(HTTPFileSystem()))
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "http") is True
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "https") is True
|
|
|
|
def test_fsspec_s3_not_compatible_with_gs(self):
|
|
"""fsspec S3FileSystem should NOT be compatible with gs scheme."""
|
|
s3fs = pytest.importorskip("s3fs")
|
|
wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "gs") is False
|
|
|
|
def test_fsspec_s3_not_compatible_with_http(self):
|
|
"""fsspec S3FileSystem should NOT be compatible with http/https schemes."""
|
|
s3fs = pytest.importorskip("s3fs")
|
|
wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "http") is False
|
|
assert _is_filesystem_compatible_with_scheme(wrapped_fs, "https") is False
|
|
|
|
def test_retrying_wrapper_around_native_s3(self):
|
|
"""RetryingPyFileSystem wrapping a native S3FileSystem should be compatible with s3."""
|
|
s3_fs = pyarrow.fs.S3FileSystem(anonymous=True)
|
|
retrying_fs = RetryingPyFileSystem.wrap(s3_fs, retryable_errors=["AWS Error"])
|
|
assert _is_filesystem_compatible_with_scheme(retrying_fs, "s3") is True
|
|
assert _is_filesystem_compatible_with_scheme(retrying_fs, "gs") is False
|
|
|
|
def test_unknown_scheme_trusts_filesystem(self):
|
|
"""Unknown schemes should always return True (trust user's filesystem)."""
|
|
fs = pyarrow.fs.LocalFileSystem()
|
|
assert _is_filesystem_compatible_with_scheme(fs, "custom") is True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|