Skip to content

Commit 787130b

Browse files
PyArrow: Keep the storage account out of ADLS paths in parse_location
For Azure the netloc is <container>@<account>.<host>, so building the path as netloc + path put the account inside the path and PyArrow then read the whole first segment as the container name. Return only the container instead, which matches what PyArrow's own from_uri produces for the same location.
1 parent db5bd04 commit 787130b

2 files changed

Lines changed: 31 additions & 1 deletion

File tree

‎pyiceberg/io/pyarrow.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,7 @@
209209
MAP_VALUE_NAME = "value"
210210
DOC = "doc"
211211
UTC_ALIASES = {"UTC", "+00:00", "Etc/UTC", "Z"}
212+
ADLS_SCHEMES = frozenset({"abfs", "abfss", "wasb", "wasbs"})
212213

213214
T = TypeVar("T")
214215

@@ -421,6 +422,11 @@ def parse_location(location: str, properties: Properties = EMPTY_DICT) -> tuple[
421422
return default_scheme, default_netloc, os.path.abspath(location)
422423
elif uri.scheme in ("hdfs", "viewfs"):
423424
return uri.scheme, uri.netloc, uri.path
425+
elif uri.scheme in ADLS_SCHEMES and uri.username:
426+
# Azure locations are <container>@<account>.<host>/<path>. The account is
427+
# configured on the AzureFileSystem itself, which expects paths of the form
428+
# <container>/<path>, so only the container belongs in the path here.
429+
return uri.scheme, uri.netloc, f"{uri.username}{uri.path}"
424430
else:
425431
return uri.scheme, uri.netloc, f"{uri.netloc}{uri.path}"
426432

@@ -438,7 +444,7 @@ def _initialize_fs(self, scheme: str, netloc: str | None = None) -> FileSystem:
438444
elif scheme in {"gs", "gcs"}:
439445
return self._initialize_gcs_fs()
440446

441-
elif scheme in {"abfs", "abfss", "wasb", "wasbs"}:
447+
elif scheme in ADLS_SCHEMES:
442448
return self._initialize_azure_fs()
443449

444450
elif scheme in {"file"}:

‎tests/io/test_pyarrow.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2326,6 +2326,30 @@ def check_results(location: str, expected_schema: str, expected_netloc: str, exp
23262326
check_results("/root/foo.txt", "file", "", os.path.abspath("/root/foo.txt"))
23272327
check_results("/root/tmp/foo.txt", "file", "", os.path.abspath("/root/tmp/foo.txt"))
23282328

2329+
check_results("s3://bucket/root/foo.txt", "s3", "bucket", "bucket/root/foo.txt")
2330+
2331+
2332+
@pytest.mark.parametrize("scheme", ["abfs", "abfss", "wasb", "wasbs"])
2333+
def test_parse_location_adls_account_qualified(scheme: str) -> None:
2334+
"""The account must not leak into the path, PyArrow takes it on the filesystem instead."""
2335+
scheme_, netloc, path = PyArrowFileIO.parse_location(
2336+
f"{scheme}://mycontainer@myaccount.dfs.core.windows.net/wh/db/tbl/data.parquet"
2337+
)
2338+
2339+
assert scheme_ == scheme
2340+
assert netloc == "mycontainer@myaccount.dfs.core.windows.net"
2341+
assert path == "mycontainer/wh/db/tbl/data.parquet"
2342+
2343+
2344+
@pytest.mark.parametrize("scheme", ["abfs", "abfss", "wasb", "wasbs"])
2345+
def test_parse_location_adls_container_only(scheme: str) -> None:
2346+
"""Locations without an account keep the netloc as the container."""
2347+
scheme_, netloc, path = PyArrowFileIO.parse_location(f"{scheme}://mycontainer/wh/db/tbl/data.parquet")
2348+
2349+
assert scheme_ == scheme
2350+
assert netloc == "mycontainer"
2351+
assert path == "mycontainer/wh/db/tbl/data.parquet"
2352+
23292353

23302354
@pytest.mark.skipif(sys.platform != "win32", reason="Windows-only behavior")
23312355
def test_parse_location_windows_drive_letter() -> None:

0 commit comments

Comments
 (0)