PyArrow: Derive the ADLS account from the location when it is not configured

pyiceberg/io/pyarrow.py

@@ -209,6 +209,7 @@
 MAP_VALUE_NAME = "value"
 DOC = "doc"
 UTC_ALIASES = {"UTC", "+00:00", "Etc/UTC", "Z"}
+ADLS_SCHEMES = frozenset({"abfs", "abfss", "wasb", "wasbs"})
 
 T = TypeVar("T")
 
@@ -421,6 +422,11 @@ def parse_location(location: str, properties: Properties = EMPTY_DICT) -> tuple[
             return default_scheme, default_netloc, os.path.abspath(location)
         elif uri.scheme in ("hdfs", "viewfs"):
             return uri.scheme, uri.netloc, uri.path
+        elif uri.scheme in ADLS_SCHEMES and uri.username:
+            # Azure locations are <container>@<account>.<host>/<path>. The account is
+            # configured on the AzureFileSystem itself, which expects paths of the form
+            # <container>/<path>, so only the container belongs in the path here.
+            return uri.scheme, uri.netloc, f"{uri.username}{uri.path}"
         else:
             return uri.scheme, uri.netloc, f"{uri.netloc}{uri.path}"
 
@@ -438,8 +444,8 @@ def _initialize_fs(self, scheme: str, netloc: str | None = None) -> FileSystem:
         elif scheme in {"gs", "gcs"}:
             return self._initialize_gcs_fs()
 
-        elif scheme in {"abfs", "abfss", "wasb", "wasbs"}:
-            return self._initialize_azure_fs()
+        elif scheme in ADLS_SCHEMES:
+            return self._initialize_azure_fs(netloc)
 
         elif scheme in {"file"}:
             return self._initialize_local_fs()
@@ -533,7 +539,7 @@ def _initialize_s3_fs(self, netloc: str | None) -> FileSystem:
 
         return S3FileSystem(**client_kwargs)
 
-    def _initialize_azure_fs(self) -> FileSystem:
+    def _initialize_azure_fs(self, netloc: str | None = None) -> FileSystem:
         # https://arrow.apache.org/docs/python/generated/pyarrow.fs.AzureFileSystem.html
         from packaging import version
 
@@ -548,7 +554,15 @@ def _initialize_azure_fs(self) -> FileSystem:
 
         client_kwargs: dict[str, str] = {}
 
-        if account_name := self.properties.get(ADLS_ACCOUNT_NAME):
+        account_name = self.properties.get(ADLS_ACCOUNT_NAME)
+        if account_name is None and netloc and "@" in netloc:
+            # An account qualified location carries the account in its host part,
+            # abfss://<container>@<account>.<host>/<path>. Only that form has a userinfo
+            # part, without it the netloc is just the container and there is no account
+            # to read. An explicit adls.account-name always wins, same as FsspecFileIO.
+            account_name = netloc.rpartition("@")[2].split(".")[0] or None
+
+        if account_name:
             client_kwargs["account_name"] = account_name
 
         if account_key := self.properties.get(ADLS_ACCOUNT_KEY):

tests/io/test_pyarrow.py

@@ -63,7 +63,7 @@
     Or,
 )
 from pyiceberg.expressions.literals import literal
-from pyiceberg.io import S3_RETRY_STRATEGY_IMPL, InputStream, OutputStream, load_file_io
+from pyiceberg.io import ADLS_ACCOUNT_NAME, S3_RETRY_STRATEGY_IMPL, InputStream, OutputStream, load_file_io
 from pyiceberg.io.pyarrow import (
     ICEBERG_SCHEMA,
     PYARROW_PARQUET_FIELD_ID_KEY,
@@ -2326,6 +2326,30 @@ def check_results(location: str, expected_schema: str, expected_netloc: str, exp
     check_results("/root/foo.txt", "file", "", os.path.abspath("/root/foo.txt"))
     check_results("/root/tmp/foo.txt", "file", "", os.path.abspath("/root/tmp/foo.txt"))
 
+    check_results("s3://bucket/root/foo.txt", "s3", "bucket", "bucket/root/foo.txt")
+
+
+@pytest.mark.parametrize("scheme", ["abfs", "abfss", "wasb", "wasbs"])
+def test_parse_location_adls_account_qualified(scheme: str) -> None:
+    """The account must not leak into the path, PyArrow takes it on the filesystem instead."""
+    scheme_, netloc, path = PyArrowFileIO.parse_location(
+        f"{scheme}://mycontainer@myaccount.dfs.core.windows.net/wh/db/tbl/data.parquet"
+    )
+
+    assert scheme_ == scheme
+    assert netloc == "mycontainer@myaccount.dfs.core.windows.net"
+    assert path == "mycontainer/wh/db/tbl/data.parquet"
+
+
+@pytest.mark.parametrize("scheme", ["abfs", "abfss", "wasb", "wasbs"])
+def test_parse_location_adls_container_only(scheme: str) -> None:
+    """Locations without an account keep the netloc as the container."""
+    scheme_, netloc, path = PyArrowFileIO.parse_location(f"{scheme}://mycontainer/wh/db/tbl/data.parquet")
+
+    assert scheme_ == scheme
+    assert netloc == "mycontainer"
+    assert path == "mycontainer/wh/db/tbl/data.parquet"
+
 
 @pytest.mark.skipif(sys.platform != "win32", reason="Windows-only behavior")
 def test_parse_location_windows_drive_letter() -> None:
@@ -3210,6 +3234,46 @@ def test__to_requested_schema_float_promotion(
     assert result.column(0).to_pylist() == [1.5, 2.25, 3.0, None]
 
 
+@skip_if_pyarrow_too_old
+@pytest.mark.parametrize("scheme", ["abfs", "abfss", "wasb", "wasbs"])
+def test_adls_account_name_from_location(scheme: str) -> None:
+    """The account in an account qualified location is used when the property is not set."""
+    with patch("pyarrow.fs.AzureFileSystem") as mock_azure_fs:
+        PyArrowFileIO().fs_by_scheme(scheme, "mycontainer@myaccount.dfs.core.windows.net")
+
+    assert mock_azure_fs.call_args.kwargs["account_name"] == "myaccount"
+
+
+@skip_if_pyarrow_too_old
+def test_adls_account_name_property_wins_over_location() -> None:
+    """An explicit adls.account-name is not overridden by the account in the location."""
+    with patch("pyarrow.fs.AzureFileSystem") as mock_azure_fs:
+        PyArrowFileIO({ADLS_ACCOUNT_NAME: "configured"}).fs_by_scheme("abfss", "mycontainer@myaccount.dfs.core.windows.net")
+
+    assert mock_azure_fs.call_args.kwargs["account_name"] == "configured"
+
+
+@skip_if_pyarrow_too_old
+def test_adls_no_account_name_from_container_only_location() -> None:
+    """A container only netloc carries no account, so nothing should be inferred from it."""
+    with patch("pyarrow.fs.AzureFileSystem") as mock_azure_fs:
+        PyArrowFileIO().fs_by_scheme("abfss", "warehouse")
+
+    assert "account_name" not in mock_azure_fs.call_args.kwargs
+
+
+@skip_if_pyarrow_too_old
+def test_adls_account_name_per_location() -> None:
+    """Two accounts served by one FileIO must each get their own filesystem."""
+    file_io = PyArrowFileIO()
+
+    with patch("pyarrow.fs.AzureFileSystem") as mock_azure_fs:
+        file_io.fs_by_scheme("abfss", "data@accountone.dfs.core.windows.net")
+        file_io.fs_by_scheme("abfss", "data@accounttwo.dfs.core.windows.net")
+
+    assert [call.kwargs["account_name"] for call in mock_azure_fs.call_args_list] == ["accountone", "accounttwo"]
+
+
 def test_pyarrow_file_io_fs_by_scheme_cache() -> None:
     # It's better to set up multi-region minio servers for an integration test once `endpoint_url` argument
     # becomes available for `resolve_s3_region`