Add s3.allow-delayed-open to PyArrowFileIO

mkdocs/docs/configuration.md

@@ -170,6 +170,7 @@ For the FileIO there are several configuration options available:
 | s3.connect-timeout          | 60.0                       | Configure socket connection timeout, in seconds.                                                                                                                                                                                                                        |
 | s3.request-timeout          | 60.0                       | Configure socket read timeouts on Windows and macOS, in seconds.                                                                                                                                                                                                        |
 | s3.force-virtual-addressing | False                      | Whether to use virtual addressing of buckets. If true, then virtual addressing is always enabled. If false, then virtual addressing is only enabled if endpoint_override is empty. This can be used for non-AWS backends that only support virtual hosted-style access. |
+| s3.allow-delayed-open       | True                       | Only supported by `PyArrowFileIO`, requires pyarrow >= 21. If False (default), every write is a multipart upload. If True, an object smaller than 10 MiB is written with a single `PutObject` request instead of a three-request multipart upload. Errors such as a missing bucket are then raised when the file is closed rather than when it is opened. |
 | s3.retry-strategy-impl      | None                       | Ability to set a custom S3 retry strategy. A full path to a class needs to be given that extends the [S3RetryStrategy](https://github.com/apache/arrow/blob/639201bfa412db26ce45e73851432018af6c945e/python/pyarrow/_s3fs.pyx#L110) base class.                         |
 | s3.anonymous                | True                       | Configure whether to use anonymous connection. If False (default), uses key/secret if configured or boto's credential resolver.                                                                                                                                         |
 | s3.server-side-encryption   | aws:kms                    | Configure server-side encryption (e.g. `AES256` or `aws:kms`). Only supported by `FsspecFileIO`.                                                                                                                                     |

pyiceberg/io/__init__.py

@@ -80,6 +80,7 @@ def _is_local_path(path: str) -> bool:
 S3_ROLE_ARN = "s3.role-arn"
 S3_ROLE_SESSION_NAME = "s3.role-session-name"
 S3_FORCE_VIRTUAL_ADDRESSING = "s3.force-virtual-addressing"
+S3_ALLOW_DELAYED_OPEN = "s3.allow-delayed-open"
 S3_RETRY_STRATEGY_IMPL = "s3.retry-strategy-impl"
 S3_SERVER_SIDE_ENCRYPTION = "s3.server-side-encryption"
 S3_SSE_KMS_KEY_ID = "s3.sse-kms-key-id"

pyiceberg/io/pyarrow.py

@@ -104,6 +104,7 @@
     HDFS_PORT,
     HDFS_USER,
     S3_ACCESS_KEY_ID,
+    S3_ALLOW_DELAYED_OPEN,
     S3_ANONYMOUS,
     S3_CONNECT_TIMEOUT,
     S3_ENDPOINT,
@@ -465,6 +466,16 @@ def _set_tls_ca_file_path(self, client_kwargs: dict[str, Any]) -> None:
                 )
             client_kwargs["tls_ca_file_path"] = tls_ca_file_path
 
+    def _set_allow_delayed_open(self, client_kwargs: dict[str, Any]) -> None:
+        if self.properties.get(S3_ALLOW_DELAYED_OPEN) is not None:
+            min_pyarrow_version_supporting_allow_delayed_open = (21, 0)
+            if tuple(map(int, pa.__version__.split(".")[:2])) < min_pyarrow_version_supporting_allow_delayed_open:
+                raise ImportError(
+                    "pyarrow version >= 21.0.0 required for S3FileSystem allow_delayed_open support, "
+                    f"but found version {pa.__version__}."
+                )
+            client_kwargs["allow_delayed_open"] = property_as_bool(self.properties, S3_ALLOW_DELAYED_OPEN, False)
+
     def _initialize_oss_fs(self) -> FileSystem:
         from pyarrow.fs import S3FileSystem
 
@@ -496,6 +507,7 @@ def _initialize_oss_fs(self) -> FileSystem:
             client_kwargs["anonymous"] = strtobool(s3_anonymous)
 
         self._set_tls_ca_file_path(client_kwargs)
+        self._set_allow_delayed_open(client_kwargs)
 
         return S3FileSystem(**client_kwargs)
 
@@ -552,6 +564,7 @@ def _initialize_s3_fs(self, netloc: str | None) -> FileSystem:
             client_kwargs["anonymous"] = strtobool(s3_anonymous)
 
         self._set_tls_ca_file_path(client_kwargs)
+        self._set_allow_delayed_open(client_kwargs)
 
         return S3FileSystem(**client_kwargs)
 

tests/io/test_pyarrow.py

@@ -63,7 +63,15 @@
     Or,
 )
 from pyiceberg.expressions.literals import literal
-from pyiceberg.io import ADLS_ACCOUNT_NAME, S3_RETRY_STRATEGY_IMPL, S3_SSL_CA_CERT, InputStream, OutputStream, load_file_io
+from pyiceberg.io import (
+    ADLS_ACCOUNT_NAME,
+    S3_ALLOW_DELAYED_OPEN,
+    S3_RETRY_STRATEGY_IMPL,
+    S3_SSL_CA_CERT,
+    InputStream,
+    OutputStream,
+    load_file_io,
+)
 from pyiceberg.io.pyarrow import (
     ICEBERG_SCHEMA,
     PYARROW_PARQUET_FIELD_ID_KEY,
@@ -420,6 +428,51 @@ def test_pyarrow_s3_ssl_ca_cert_requires_supported_pyarrow_version() -> None:
             PyArrowFileIO(properties=session_properties).new_input(location=f"s3://warehouse/{uuid.uuid4()}")
 
 
+def test_pyarrow_s3_session_properties_with_allow_delayed_open() -> None:
+    session_properties: Properties = {
+        S3_ALLOW_DELAYED_OPEN: "true",
+        "s3.endpoint": "http://localhost:9000",
+        "s3.access-key-id": "admin",
+        "s3.secret-access-key": "password",
+        "s3.region": "us-east-1",
+        "s3.session-token": "s3.session-token",
+        **UNIFIED_AWS_SESSION_PROPERTIES,
+    }
+
+    with patch("pyarrow.fs.S3FileSystem") as mock_s3fs, patch("pyarrow.fs.resolve_s3_region") as mock_s3_region_resolver:
+        s3_fileio = PyArrowFileIO(properties=session_properties)
+        filename = str(uuid.uuid4())
+
+        # Mock `resolve_s3_region` to prevent from the location used resolving to a different s3 region
+        mock_s3_region_resolver.side_effect = OSError("S3 bucket is not found")
+        s3_fileio.new_input(location=f"s3://warehouse/{filename}")
+
+        mock_s3fs.assert_called_with(
+            allow_delayed_open=True,
+            endpoint_override="http://localhost:9000",
+            access_key="admin",
+            secret_key="password",
+            region="us-east-1",
+            session_token="s3.session-token",
+        )
+
+
+def test_pyarrow_s3_allow_delayed_open_requires_supported_pyarrow_version() -> None:
+    session_properties: Properties = {
+        S3_ALLOW_DELAYED_OPEN: "true",
+    }
+
+    with (
+        patch("pyiceberg.io.pyarrow.pa.__version__", "20.0.0"),
+        patch("pyarrow.fs.S3FileSystem"),
+        patch("pyarrow.fs.resolve_s3_region") as mock_s3_region_resolver,
+    ):
+        mock_s3_region_resolver.side_effect = OSError("S3 bucket is not found")
+
+        with pytest.raises(ImportError, match="pyarrow version >= 21.0.0 required"):
+            PyArrowFileIO(properties=session_properties).new_input(location=f"s3://warehouse/{uuid.uuid4()}")
+
+
 def test_pyarrow_s3_session_properties_with_anonymous() -> None:
     session_properties: Properties = {
         "s3.anonymous": "true",