Skip to content

Commit 0a536e4

Browse files
PyArrow: Derive the ADLS account from the location when it is not configured (#3936)
* PyArrow: Derive the ADLS account from the location when it is not configured _initialize_azure_fs took no netloc, unlike the S3 and HDFS initialisers next to it, so the account could only come from adls.account-name and the account in an abfss location was dropped. Pass the netloc through and fall back to the account in its host part. An explicit adls.account-name still wins, same precedence as FsspecFileIO. The fallback is gated on the userinfo part being present, since a container only netloc carries no account. * Remove unnecessary comment
1 parent 068aae5 commit 0a536e4

2 files changed

Lines changed: 48 additions & 4 deletions

File tree

‎pyiceberg/io/pyarrow.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -447,7 +447,7 @@ def _initialize_fs(self, scheme: str, netloc: str | None = None) -> FileSystem:
447447
return self._initialize_gcs_fs()
448448

449449
elif scheme in ADLS_SCHEMES:
450-
return self._initialize_azure_fs()
450+
return self._initialize_azure_fs(netloc)
451451

452452
elif scheme in {"file"}:
453453
return self._initialize_local_fs()
@@ -555,7 +555,7 @@ def _initialize_s3_fs(self, netloc: str | None) -> FileSystem:
555555

556556
return S3FileSystem(**client_kwargs)
557557

558-
def _initialize_azure_fs(self) -> FileSystem:
558+
def _initialize_azure_fs(self, netloc: str | None = None) -> FileSystem:
559559
# https://arrow.apache.org/docs/python/generated/pyarrow.fs.AzureFileSystem.html
560560
from packaging import version
561561

@@ -570,7 +570,11 @@ def _initialize_azure_fs(self) -> FileSystem:
570570

571571
client_kwargs: dict[str, str] = {}
572572

573-
if account_name := self.properties.get(ADLS_ACCOUNT_NAME):
573+
account_name = self.properties.get(ADLS_ACCOUNT_NAME)
574+
if account_name is None and netloc and "@" in netloc:
575+
account_name = netloc.rpartition("@")[2].split(".")[0] or None
576+
577+
if account_name:
574578
client_kwargs["account_name"] = account_name
575579

576580
if account_key := self.properties.get(ADLS_ACCOUNT_KEY):

‎tests/io/test_pyarrow.py‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@
6363
Or,
6464
)
6565
from pyiceberg.expressions.literals import literal
66-
from pyiceberg.io import S3_RETRY_STRATEGY_IMPL, S3_SSL_CA_CERT, InputStream, OutputStream, load_file_io
66+
from pyiceberg.io import ADLS_ACCOUNT_NAME, S3_RETRY_STRATEGY_IMPL, S3_SSL_CA_CERT, InputStream, OutputStream, load_file_io
6767
from pyiceberg.io.pyarrow import (
6868
ICEBERG_SCHEMA,
6969
PYARROW_PARQUET_FIELD_ID_KEY,
@@ -3272,6 +3272,46 @@ def test__to_requested_schema_float_promotion(
32723272
assert result.column(0).to_pylist() == [1.5, 2.25, 3.0, None]
32733273

32743274

3275+
@skip_if_pyarrow_too_old
3276+
@pytest.mark.parametrize("scheme", ["abfs", "abfss", "wasb", "wasbs"])
3277+
def test_adls_account_name_from_location(scheme: str) -> None:
3278+
"""The account in an account qualified location is used when the property is not set."""
3279+
with patch("pyarrow.fs.AzureFileSystem") as mock_azure_fs:
3280+
PyArrowFileIO().fs_by_scheme(scheme, "mycontainer@myaccount.dfs.core.windows.net")
3281+
3282+
assert mock_azure_fs.call_args.kwargs["account_name"] == "myaccount"
3283+
3284+
3285+
@skip_if_pyarrow_too_old
3286+
def test_adls_account_name_property_wins_over_location() -> None:
3287+
"""An explicit adls.account-name is not overridden by the account in the location."""
3288+
with patch("pyarrow.fs.AzureFileSystem") as mock_azure_fs:
3289+
PyArrowFileIO({ADLS_ACCOUNT_NAME: "configured"}).fs_by_scheme("abfss", "mycontainer@myaccount.dfs.core.windows.net")
3290+
3291+
assert mock_azure_fs.call_args.kwargs["account_name"] == "configured"
3292+
3293+
3294+
@skip_if_pyarrow_too_old
3295+
def test_adls_no_account_name_from_container_only_location() -> None:
3296+
"""A container only netloc carries no account, so nothing should be inferred from it."""
3297+
with patch("pyarrow.fs.AzureFileSystem") as mock_azure_fs:
3298+
PyArrowFileIO().fs_by_scheme("abfss", "warehouse")
3299+
3300+
assert "account_name" not in mock_azure_fs.call_args.kwargs
3301+
3302+
3303+
@skip_if_pyarrow_too_old
3304+
def test_adls_account_name_per_location() -> None:
3305+
"""Two accounts served by one FileIO must each get their own filesystem."""
3306+
file_io = PyArrowFileIO()
3307+
3308+
with patch("pyarrow.fs.AzureFileSystem") as mock_azure_fs:
3309+
file_io.fs_by_scheme("abfss", "data@accountone.dfs.core.windows.net")
3310+
file_io.fs_by_scheme("abfss", "data@accounttwo.dfs.core.windows.net")
3311+
3312+
assert [call.kwargs["account_name"] for call in mock_azure_fs.call_args_list] == ["accountone", "accounttwo"]
3313+
3314+
32753315
def test_pyarrow_file_io_fs_by_scheme_cache() -> None:
32763316
# It's better to set up multi-region minio servers for an integration test once `endpoint_url` argument
32773317
# becomes available for `resolve_s3_region`

0 commit comments

Comments
 (0)