[python] Fix OSS directory deletion with PyArrow 22+

.github/workflows/ci-python.yml

@@ -145,6 +145,32 @@ jobs:
           chmod +x paimon-python/dev/lint-python.sh
           ./paimon-python/dev/lint-python.sh -e pytest_torch
 
+  s3-delete-compatibility:
+    name: S3 deletion / PyArrow ${{ matrix.pyarrow }}
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    strategy:
+      fail-fast: false
+      matrix:
+        pyarrow: ['22.0.0', '23.0.0']
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - name: Install test dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e ./paimon-python --group './paimon-python/pyproject.toml:test'
+          # Exercise the compatibility path before widening the package's Arrow range.
+          python -m pip install --no-deps 'pyarrow==${{ matrix.pyarrow }}'
+      - name: Test S3 compatibility
+        working-directory: paimon-python
+        run: |
+          python -m pytest -q pypaimon/tests/s3_delete_compatibility_test.py \
+            pypaimon/tests/s3_client_compatibility_test.py pypaimon/tests/oss_legacy_mode_test.py \
+            pypaimon/tests/file_io_test.py pypaimon/tests/filesystem_catalog_test.py
+
   rust-plan:
     name: Native CI
     timeout-minutes: 90

paimon-python/pypaimon/filesystem/pyarrow_file_io.py

@@ -15,14 +15,20 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import base64
+import hashlib
+import json
 import logging
 import os
 import re
 import subprocess
+import tempfile
 import threading
+import time
 from datetime import datetime, timezone
+from itertools import islice
 from pathlib import PurePosixPath
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Iterable, List, Optional
 from urllib.parse import splitport, urlparse
 
 import pyarrow
@@ -46,6 +52,9 @@ def _pyarrow_lt_7():
     return parse(pyarrow.__version__) < parse("7.0.0")
 
 
+_S3_DELETE_TIMEOUT_SECONDS = 3600
+
+
 class LegacyOssDirectoryListingError(RuntimeError):
     """Raised when legacy PyArrow OSS cannot enumerate a directory."""
 
@@ -63,6 +72,7 @@ def __init__(self, path: str, catalog_options: Options):
             self._get_s3_property("endpoint", S3Options.S3_ENDPOINT.key())
             if self._is_s3 else None
         )
+        self._s3_delete_client = None
         self._oss_bucket = None
         _oss_impl = self.properties.get(OssOptions.OSS_IMPL)
         self._use_jindo = False
@@ -96,6 +106,7 @@ def __init__(self, path: str, catalog_options: Options):
             raise ValueError(f"Unrecognized filesystem type in URI: {scheme}")
 
     def _set_pyarrow_version(self):
+        self._pyarrow_gte_22 = parse(pyarrow.__version__) >= parse("22.0.0")
         self._pyarrow_gte_8 = parse(pyarrow.__version__) >= parse("8.0.0")
         # force_virtual_addressing landed in PyArrow 16; below it the OSS bucket
         # goes into endpoint_override, so keys must omit it (init + path share
@@ -107,6 +118,9 @@ def _uses_s3_compatibility(self) -> bool:
         return (not self._use_jindo
                 and (self._is_oss or bool(self._s3_endpoint)))
 
+    def _uses_s3_delete_fallback(self) -> bool:
+        return self._uses_s3_compatibility() and self._pyarrow_gte_22
+
     def _configure_s3_checksums(self):
         if parse(pyarrow.__version__) < parse("22.0.0"):
             return
@@ -119,13 +133,15 @@ def __getstate__(self):
         # threading.Lock cannot be pickled; recreated in __setstate__.
         state.pop("_legacy_bucket_lock", None)
         state.pop("logger", None)
+        state.pop("_s3_delete_client", None)
         # Recreate S3-compatible clients with the worker's AWS SDK settings.
         if self._uses_s3_compatibility():
             state.pop("filesystem", None)
         return state
 
     def __setstate__(self, state):
         self.__dict__.update(state)
+        self._s3_delete_client = None
         self.logger = logging.getLogger(__name__)
         self._set_pyarrow_version()
         if "_is_s3" not in state:
@@ -300,6 +316,25 @@ def _initialize_s3_fs(self) -> FileSystem:
 
         return pafs.S3FileSystem(**client_kwargs)
 
+    def _s3_connection_options(self):
+        return {
+            "access_key": self._get_property(
+                S3Options.S3_ACCESS_KEY_ID.key(),
+                *self._s3_key_variants("access-key", "access.key")),
+            "secret_key": self._get_property(
+                S3Options.S3_ACCESS_KEY_SECRET.key(),
+                *self._s3_key_variants("secret-key", "secret.key")),
+            "session_token": self._get_property(
+                S3Options.S3_SECURITY_TOKEN.key(),
+                *self._s3_key_variants(
+                    "session-token", "session.token",
+                    "security-token", "security.token")),
+            "region": self._get_s3_property("region", S3Options.S3_REGION.key()),
+            "path_style": (self._pyarrow_gte_16 and (
+                self._get_s3_boolean_property("path-style-access") or
+                self._get_s3_boolean_property("path.style.access"))),
+        }
+
     def _initialize_hdfs_fs(self, scheme: str, netloc: Optional[str]) -> FileSystem:
         if 'HADOOP_HOME' not in os.environ:
             raise RuntimeError("HADOOP_HOME environment variable is not set.")
@@ -486,13 +521,38 @@ def exists_batch(self, paths: List[str]) -> Dict[str, bool]:
         }
 
     def delete(self, path: str, recursive: bool = False) -> bool:
+        use_fallback = self._uses_s3_delete_fallback()
+        if use_fallback and self._is_oss:
+            parsed = urlparse(path)
+            if parsed.scheme == "oss" and parsed.netloc:
+                bucket = self._extract_oss_bucket(path)
+                key = parsed.path.lstrip("/")
+                if "@" in parsed.netloc:
+                    _, _, key = key.partition("/")
+                path = "oss://" + bucket + "/" + key
         path_str = self.to_filesystem_path(path)
+        if use_fallback:
+            bucket, key = self._split_s3_path(path_str)
+            if not bucket or not key.strip("/"):
+                raise OSError(f"Refusing to delete bucket root: {path}")
         file_info = self._get_file_info(path_str)
 
         if file_info.type == pafs.FileType.NotFound:
             return False
 
         if file_info.type == pafs.FileType.Directory:
+            if use_fallback:
+                if recursive:
+                    return self._delete_s3_compatible_directory(path_str)
+                selector = pafs.FileSelector(
+                    path_str, recursive=False, allow_not_found=True)
+                if self.filesystem.get_file_info(selector):
+                    raise OSError(f"Directory {path} is not empty")
+                bucket, key = self._split_s3_path(path_str)
+                if key:
+                    self._get_s3_delete_client().delete_object(
+                        Bucket=bucket, Key=key.rstrip("/") + "/")
+                return True
             if not recursive:
                 selector = pafs.FileSelector(path_str, recursive=False, allow_not_found=True)
                 dir_contents = self.filesystem.get_file_info(selector)
@@ -507,6 +567,163 @@ def delete(self, path: str, recursive: bool = False) -> bool:
             self.filesystem.delete_file(path_str)
         return True
 
+    def _delete_s3_compatible_directory(self, path_str: str) -> bool:
+        client = self._get_s3_delete_client()
+        bucket, key = self._split_s3_path(path_str)
+        prefix = key.rstrip("/")
+        if prefix:
+            prefix += "/"
+        deadline = time.monotonic() + _S3_DELETE_TIMEOUT_SECONDS
+        # Finish listing before deleting: mutating a paginated listing may skip keys.
+        with tempfile.TemporaryFile(mode="w+t", encoding="utf-8") as listed:
+            for name in self._list_s3_keys(client, bucket, prefix, deadline, path_str):
+                listed.write(json.dumps(name) + "\n")
+            listed.seek(0)
+            self._delete_s3_objects(
+                client, bucket, (json.loads(line) for line in listed),
+                deadline, path_str)
+        if any(self._list_s3_keys(client, bucket, prefix, deadline, path_str)):
+            raise OSError(f"S3 directory {path_str} changed during deletion")
+        return True
+
+    def _list_s3_keys(self, client, bucket: str, prefix: str,
+                      deadline: float, path_str: str):
+        token = None
+        while True:
+            self._check_s3_delete_deadline(deadline, path_str)
+            params = {"Bucket": bucket, "Prefix": prefix, "MaxKeys": 1000}
+            if token is not None:
+                params["ContinuationToken"] = token
+            response = client.list_objects_v2(**params)
+            yield from self._listed_s3_keys(response, bucket, prefix)
+            if not response.get("IsTruncated"):
+                return
+            next_token = response.get("NextContinuationToken")
+            if not next_token or next_token == token:
+                raise OSError("S3 listing did not advance")
+            token = next_token
+
+    @staticmethod
+    def _listed_s3_keys(response, bucket: str, prefix: str):
+        if (response.get("Name", bucket) != bucket
+                or response.get("Prefix", prefix) != prefix):
+            raise OSError("S3 listing returned a different bucket or prefix")
+        keys = [item["Key"] for item in response.get("Contents", ())]
+        if any(not key.startswith(prefix) for key in keys):
+            raise OSError(f"S3 listing returned a key outside prefix {prefix}")
+        return keys
+
+    @staticmethod
+    def _check_s3_delete_deadline(deadline: float, path_str: str):
+        if time.monotonic() >= deadline:
+            raise TimeoutError(f"Timed out deleting S3 directory {path_str}")
+
+    @staticmethod
+    def _delete_s3_objects(
+            client, bucket: str, keys: Iterable[str],
+            deadline: float, path_str: str):
+        keys = iter(keys)
+        while True:
+            batch = list(islice(keys, 1000))
+            if not batch:
+                return
+            PyArrowFileIO._check_s3_delete_deadline(deadline, path_str)
+            ordinary = []
+            for key in batch:
+                if key.isprintable():
+                    ordinary.append(key)
+                else:
+                    # Non-printable keys may not round-trip through DeleteObjects XML.
+                    PyArrowFileIO._check_s3_delete_deadline(deadline, path_str)
+                    client.delete_object(Bucket=bucket, Key=key)
+            if not ordinary:
+                continue
+            PyArrowFileIO._check_s3_delete_deadline(deadline, path_str)
+            response = client.delete_objects(
+                Bucket=bucket,
+                Delete={"Objects": [{"Key": key} for key in ordinary],
+                        "Quiet": False})
+            deleted = [item["Key"] for item in response.get("Deleted", ())]
+            errors = response.get("Errors", ())
+            if errors:
+                examples = ", ".join(
+                    "{}: {!r}".format(item.get("Code"), item.get("Key"))
+                    for item in errors[:3])
+                raise OSError(
+                    f"S3 batch delete incomplete for {path_str}: {examples}")
+            if set(deleted) != set(ordinary) or len(deleted) != len(ordinary):
+                raise OSError(f"S3 batch delete incomplete for {path_str}")
+
+    @staticmethod
+    def _split_s3_path(path_str: str):
+        bucket, _, key = path_str.partition("/")
+        return bucket, key
+
+    def _get_s3_delete_client(self):
+        if self._s3_delete_client is not None:
+            return self._s3_delete_client
+
+        import boto3
+        from botocore.config import Config
+
+        if self._is_oss:
+            endpoint = self.properties.get(OssOptions.OSS_ENDPOINT)
+            access_key = self.properties.get(OssOptions.OSS_ACCESS_KEY_ID)
+            secret_key = self.properties.get(OssOptions.OSS_ACCESS_KEY_SECRET)
+            session_token = self.properties.get(OssOptions.OSS_SECURITY_TOKEN)
+            region = self.properties.get(OssOptions.OSS_REGION)
+            addressing_style = "virtual"
+        else:
+            endpoint = self._s3_endpoint
+            connection = self._s3_connection_options()
+            access_key = connection["access_key"]
+            secret_key = connection["secret_key"]
+            session_token = connection["session_token"]
+            region = connection["region"]
+            addressing_style = "path" if connection["path_style"] else "virtual"
+
+        region = region or self.filesystem.region
+        if endpoint and "://" not in endpoint:
+            endpoint = "https://" + endpoint
+        config_args = {
+            "retries": {"max_attempts": 10, "mode": "standard"},
+            "s3": {"addressing_style": addressing_style},
+            "connect_timeout": 60,
+            "read_timeout": 60,
+        }
+        try:
+            config = Config(request_checksum_calculation="when_required",
+                            **config_args)
+            uses_new_checksums = True
+        except TypeError:
+            config = Config(**config_args)
+            uses_new_checksums = False
+
+        client = boto3.session.Session().client(
+            "s3",
+            endpoint_url=endpoint,
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key,
+            aws_session_token=session_token,
+            region_name=region,
+            config=config,
+        )
+        if uses_new_checksums:
+            # OSS requires Content-MD5; newer Botocore defaults to CRC32.
+            def use_content_md5(request, **kwargs):
+                request.headers["Content-MD5"] = base64.b64encode(
+                    hashlib.md5(request.body, usedforsecurity=False).digest()
+                ).decode("ascii")
+                for name in list(request.headers):
+                    if (name.lower().startswith("x-amz-checksum-") or
+                            name.lower() == "x-amz-sdk-checksum-algorithm"):
+                        del request.headers[name]
+
+            client.meta.events.register(
+                "before-sign.s3.DeleteObjects", use_content_md5)
+        self._s3_delete_client = client
+        return client
+
     def mkdirs(self, path: str) -> bool:
         path_str = self.to_filesystem_path(path)
         file_info = self._get_file_info(path_str)