[python] Support newer PyArrow versions - #9821
XiaoHongbo-Hope wants to merge 3 commits into
Conversation
a8e0f85 to
04df066
Compare
04df066 to
a06b25c
Compare
|
@TheR1sing3un could you please help take a loot at? |
Thank you for your invitation. I'll come and have a look |
TheR1sing3un
left a comment
There was a problem hiding this comment.
I support the direction of this change. Unlike #7736, this PR addresses the S3-compatible upload behavior that prevented us from relaxing the upper bound, and retaining <24 is reasonable given apache/arrow#50188.
I reran the 54 filesystem tests with PyArrow 23.0.0 and exercised real PyArrow clients against a local HTTP S3 test server. The checksum workaround allowed a 12 MiB multipart upload to complete without checksum trailers. Uploading from a spawned worker after deserialization also passed, including when that worker had already initialized an S3 client. Static directory deletion passed as well.
However, I reproduced two edge cases in recursive deletion that the mocked tests do not cover: an object disappearing after listing aborts cleanup, and an object arriving during cleanup can send execution back through the incompatible DeleteObjects API. Please address these and add request-level regression tests before merging. I also left a separate scalability concern about applying this fallback to every custom S3 endpoint.
These checks used a local S3 test server and official PyArrow wheels; I did not independently rerun the real DLF canary.
|
|
||
| def _delete_s3_directory_marker(self, path_str: str): | ||
| try: | ||
| self.filesystem.delete_dir(path_str.rstrip("/")) |
There was a problem hiding this comment.
Avoid re-entering DeleteObjects during directory-marker cleanup.
S3FileSystem.delete_dir() does more than remove the marker: it recursively lists the directory and deletes any remaining children via DeleteObjects before deleting the marker. See the PyArrow 23 implementation.
I reproduced this with a real PyArrow 23 client against a local endpoint requiring Content-MD5 for batch deletion: list table/data/{a,b}, delete those files, and insert table/data/late before this call. The call sends POST ?delete without Content-MD5 and fails with MissingContentMD5, leaving the new file and directory marker behind.
The current workaround therefore avoids batch deletion only while the directory remains unchanged. Could we ensure marker cleanup cannot re-enter that incompatible API, and add a request-level regression test with an object appearing between listing and marker cleanup? Mocking delete_dir() hides this behavior.
There was a problem hiding this comment.
Avoid re-entering
DeleteObjectsduring directory-marker cleanup.
S3FileSystem.delete_dir()does more than remove the marker: it recursively lists the directory and deletes any remaining children viaDeleteObjectsbefore deleting the marker. See the PyArrow 23 implementation.I reproduced this with a real PyArrow 23 client against a local endpoint requiring
Content-MD5for batch deletion: listtable/data/{a,b}, delete those files, and inserttable/data/latebefore this call. The call sendsPOST ?deletewithoutContent-MD5and fails withMissingContentMD5, leaving the new file and directory marker behind.The current workaround therefore avoids batch deletion only while the directory remains unchanged. Could we ensure marker cleanup cannot re-enter that incompatible API, and add a request-level regression test with an object appearing between listing and marker cleanup? Mocking
delete_dir()hides this behavior.
Thanks, fixed
| ] | ||
| if files: | ||
| with ThreadPoolExecutor(max_workers=min(16, len(files))) as executor: | ||
| list(executor.map(self.filesystem.delete_file, files)) |
There was a problem hiding this comment.
Tolerate objects that disappear after the directory listing.
PyArrow's delete_file() performs a HEAD request first and raises FileNotFoundError if the object is already gone. Here that exception propagates through executor.map() and skips the directory-marker cleanup.
I reproduced this by removing one listed object before its HEAD request, simulating another cleanup task deleting it. The PyArrow 23 fallback raised FileNotFoundError and left a marker behind; the original batch-deletion path with PyArrow 21 completed successfully for the same scenario.
Please wrap individual deletions so that FileNotFoundError is treated as successful deletion, while permission, transport, and other failures still propagate. A regression test should remove an object after listing and verify that cleanup completes.
There was a problem hiding this comment.
Tolerate objects that disappear after the directory listing.
PyArrow's
delete_file()performs a HEAD request first and raisesFileNotFoundErrorif the object is already gone. Here that exception propagates throughexecutor.map()and skips the directory-marker cleanup.I reproduced this by removing one listed object before its HEAD request, simulating another cleanup task deleting it. The PyArrow 23 fallback raised
FileNotFoundErrorand left a marker behind; the original batch-deletion path with PyArrow 21 completed successfully for the same scenario.Please wrap individual deletions so that
FileNotFoundErroris treated as successful deletion, while permission, transport, and other failures still propagate. A regression test should remove an object after listing and verify that cleanup completes.
Thanks, fixed
| def _delete_s3_compatible_directory(self, path_str: str) -> bool: | ||
| selector = pafs.FileSelector( | ||
| path_str, recursive=True, allow_not_found=True) | ||
| file_infos = self.filesystem.get_file_info(selector) |
There was a problem hiding this comment.
Scalability concern: this fallback is substantially more expensive and applies to every custom S3 endpoint.
With PyArrow 23, each delete_file() performs HEAD + DELETE + a PUT to preserve the parent directory. In a local HTTP request-count comparison for 1,000 files under one directory, the PyArrow 21 native path issued 10 requests, while this path issued 3,010: 1,003 HEADs, 4 GETs, 1,002 DELETEs, and 1,001 PUTs. These are request counts from a local test server, not production throughput measurements.
Also, get_file_info(selector) materializes the entire tree, and executor.map() eagerly submits the files on the tested Python version. Limiting workers to 16 does not bound the number of queued futures.
Could we provide a way to retain native batch deletion for endpoints that support it, and bound listing/submission memory for the fallback? _uses_s3_compatibility() currently includes all explicit S3 endpoints, even services that support the newer requests. This is separate from the two correctness issues above, but the cost should be considered before enabling it broadly.
There was a problem hiding this comment.
Scalability concern: this fallback is substantially more expensive and applies to every custom S3 endpoint.
With PyArrow 23, each
delete_file()performs HEAD + DELETE + a PUT to preserve the parent directory. In a local HTTP request-count comparison for 1,000 files under one directory, the PyArrow 21 native path issued 10 requests, while this path issued 3,010: 1,003 HEADs, 4 GETs, 1,002 DELETEs, and 1,001 PUTs. These are request counts from a local test server, not production throughput measurements.Also,
get_file_info(selector)materializes the entire tree, andexecutor.map()eagerly submits the files on the tested Python version. Limiting workers to 16 does not bound the number of queued futures.Could we provide a way to retain native batch deletion for endpoints that support it, and bound listing/submission memory for the fallback?
_uses_s3_compatibility()currently includes all explicit S3 endpoints, even services that support the newer requests. This is separate from the two correctness issues above, but the cost should be considered before enabling it broadly.
Thanks, fixed
JingsongLi
left a comment
There was a problem hiding this comment.
I found four correctness and compatibility issues in the current revision.
| if self._delete_s3_objects(client, bucket, keys): | ||
| continue | ||
| if prefix: | ||
| client.delete_object(Bucket=bucket, Key=prefix) |
There was a problem hiding this comment.
[P2] Preserve implicit parent directory markers
Both compatibility paths delete the target directory marker and return without recreating its immediate parent. PyArrow's S3FileSystem::DeleteDir explicitly calls EnsureParentExists after deleting a directory because an implicit parent can otherwise disappear. This is reachable in the filesystem catalog: list_databases() and get_database() accept an implicit S3 prefix, but dropping its last table removes the prefix's only child and makes the database become NotFound. The request-level cross-bucket test already models parent/child/ without a parent/ marker, but currently expects the bucket to become empty. Please recreate the immediate parent marker after successful recursive and non-recursive directory deletion, except when the parent is the bucket root.
| self.__dict__.update(state) | ||
| self._legacy_bucket_lock = threading.Lock() | ||
| self._s3_delete_client = None | ||
| if self._uses_s3_compatibility(): |
There was a problem hiding this comment.
[P2] Recompute compatibility state in the worker
This rebuilds a worker-local filesystem while retaining the producer's serialized PyArrow gates (_pyarrow_gte_8, _pyarrow_gte_16, _pyarrow_gte_22, and _oss_bucket_in_endpoint). With a supported PyArrow 21 driver and PyArrow 23 worker, _pyarrow_gte_22 remains false, so recursive deletion takes the incompatible native batch path that this PR is intended to avoid. A pre-PR non-OSS pickle also has no _s3_endpoint, causing __setstate__ to fail with AttributeError. Please migrate missing fields and recompute all version-dependent state from the worker's installed PyArrow before rebuilding the client, or reconstruct from stable path/options state.
|
|
||
| @staticmethod | ||
| def _configure_s3_compatibility(): | ||
| os.environ.setdefault( |
There was a problem hiding this comment.
[P2] Do not implement endpoint-local behavior through conditional process-global state
setdefault has two conflicting failure modes. If the host already sets the valid value AWS_REQUEST_CHECKSUM_CALCULATION=WHEN_SUPPORTED, the OSS/custom-endpoint workaround is silently skipped and uploads can still use the incompatible optional checksum trailers. If the variable is initially absent, this permanently changes the checksum policy of native AWS S3 clients created later in the same process. The tests clear os.environ before each construction, so they cover neither case. Please scope and restore this setting around compatible client construction where possible, or detect an incompatible pre-existing value and fail with an actionable message rather than silently leaving the workaround disabled.
| prefix = key.rstrip("/") | ||
| if prefix: | ||
| prefix += "/" | ||
| while True: |
There was a problem hiding this comment.
[P2] Bound the relist loop under concurrent writers
This loop exits only after a verification listing contains no non-marker object. A writer that continues adding objects under the prefix—or an eventually consistent compatible endpoint that repeatedly returns a recently deleted key—can therefore keep drop_table/drop_database issuing GET and DELETE requests forever. The regression test injects exactly one late object and does not cover sustained churn. Please add a pass, deadline, or request budget and raise a clear concurrent-modification error when the prefix does not quiesce.
025ef2b to
2e2636c
Compare
JingsongLi
left a comment
There was a problem hiding this comment.
The compatibility direction looks reasonable, but the implementation in PyArrowFileIO has grown into a second S3 deletion engine. I left a simplification proposal and three concrete correctness/compatibility comments inline.
| self.filesystem.delete_file(path_str) | ||
| return True | ||
|
|
||
| def _delete_s3_compatible_directory(self, path_str: str) -> bool: |
There was a problem hiding this comment.
Simplify this fallback around the boto3 client
This method now reimplements pagination, disk spooling, concurrency, deadlines, schema ordering, marker cleanup, and scope validation inside PyArrowFileIO. Since _get_s3_delete_client already sets request_checksum_calculation="when_required", could we first validate boto3 delete_objects against the real OSS/DLF endpoint and delete batches of up to 1,000 keys? OSS's S3 compatibility explicitly supports DeleteObjects: https://www.alibabacloud.com/help/en/oss/developer-reference/compatibility-with-amazon-s3. A smaller flow would snapshot/spool ordinary keys once, batch-delete them, relist while retaining schema-0, then delete schema-0 and the marker. That removes the thread pool, per-object requests, three staging streams, and much of this branching. If a specific endpoint truly cannot batch-delete, please isolate the individual-object strategy in a small helper rather than embedding the entire state machine here.
| raise OSError("S3 listing did not advance") | ||
| token = next_token | ||
|
|
||
| for staged in (listed, schemas, schema_zero): |
There was a problem hiding this comment.
[P2] Validate concurrent changes before deleting schema-0
This loop deletes schema_zero before the late-object check below. If a writer adds a key after the initial snapshot, the method raises changed during deletion, but schema/schema-0 has already been removed, so the failed drop leaves the table undiscoverable and cannot be retried as an existing table. I reproduced this by combining the existing late-object scenario with a staged schema-0: the exception is raised after both the data key and schema-0 appear in delete_object calls. Please delete ordinary objects first, relist while allowing only the marker and known schema keys, and delete schema-0 only after that validation succeeds.
| if not batch: | ||
| return 0 | ||
| deleted = 0 | ||
| with ThreadPoolExecutor(max_workers=16) as executor: |
There was a problem hiding this comment.
[P2] Enforce the deadline inside the deletion loop
_delete_s3_objects does not receive or check the deadline, so a slow endpoint can continue submitting every later 16-key batch long after the one-hour budget has expired. With 60-second connect/read timeouts and ten retry attempts per request, the final check may run much later, potentially only after schemas have already been deleted. Pass the deadline into this helper, check it before every new batch/stage, and avoid entering schema deletion once the budget is exhausted.
|
|
||
| if (self._is_oss and (self._use_jindo or self._oss_bucket_in_endpoint) | ||
| and (parsed.scheme or parsed.netloc)): | ||
| if (not parsed.netloc or "@" in parsed.netloc |
There was a problem hiding this comment.
[P2] Preserve credential-bearing OSS URIs
Rejecting every authority containing @ breaks the existing supported form oss://access_id:secret_key@Endpoint/bucket/.... _extract_oss_bucket explicitly parses this form, and ao_simple_test.py verifies its construction, but in legacy bucket-in-endpoint and Jindo modes every subsequent read/write/delete now fails here with OSS path is outside current bucket. I reproduced the failure in both modes through to_filesystem_path. Please validate the bucket returned by _extract_oss_bucket instead of rejecting @ unconditionally, and add an I/O-path regression test for this URI form.
163d215 to
bf2baa1
Compare
bf2baa1 to
42ca891
Compare
wangzhigang1999
left a comment
There was a problem hiding this comment.
Real OSS testing found a deletion regression: keys containing \r fail the returned-key comparison, while partition values containing \x01 cause drop_table to exhaust HTTP 500 retries after deleting data, leaving schema-0. The baseline succeeds.
Deletion also slowed from ~3s to ~11s for 8,060 objects. This is non-blocking given infrequent table drops.
Minor suggestions: retain per-object error details and share connection-option parsing between Arrow and Boto3.
| if response.get("Errors") or set(deleted) != set(batch) \ | ||
| or len(deleted) != len(batch): | ||
| raise OSError(f"S3 batch delete incomplete for {path_str}") |
There was a problem hiding this comment.
Could we include the error codes and a few failed keys from Errors? The current message makes it difficult to distinguish per-object failures from a mismatch in the returned keys.
| endpoint = self._s3_endpoint | ||
| 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._get_s3_boolean_property("path-style-access") or | ||
| self._get_s3_boolean_property("path.style.access")) | ||
| addressing_style = "path" if path_style else "virtual" |
There was a problem hiding this comment.
Could we share the connection-option parsing with _initialize_s3_fs? Credentials, region, and path-style settings are parsed in both places and could drift as options change.
| if not batch: | ||
| return | ||
| PyArrowFileIO._check_s3_delete_deadline(deadline, path_str) | ||
| response = client.delete_objects( |
There was a problem hiding this comment.
I reproduced a regression on real OSS: keys containing \r are deleted, but XML normalizes CR to LF, causing the returned-key comparison to fail. For partition values containing \x01, append and PK tables write/read successfully, but drop_table exhausts HTTP 500 retries after deleting the data, leaving schema-0. The baseline succeeds in both cases. Could we investigate the request/response handling and add regression tests, including partial failures?
JingsongLi
left a comment
There was a problem hiding this comment.
Reviewed the latest revision, including the PyArrow and S3 compatibility handling raised in earlier feedback. The updated implementation keeps the supported dependency paths usable and the coverage looks good to me.
LeRobot 0.4's dependency stack needs PyArrow >=21. Support PyArrow 20–23 while keeping
<24until Arrow #50188 is fixed; the LeRobot extra pin will be updated separately.For OSS and custom S3 endpoints, configure checksums during PyArrow client creation and rebuild clients in workers. On PyArrow 22+, use Boto3 batch deletion for ordinary keys and individual deletion for keys that cannot safely round-trip through XML. Validate bucket/prefix scope and batch results, and retain schema metadata until the directory is otherwise clear.
Validation: PyArrow 23 filesystem/catalog tests 108 passed; PyArrow 19 101 passed, 7 skipped. Real OSS recursive deletion passed with PyArrow 23 and Boto3 1.36 for keys containing
\rand\x01; earlier batch-delete canaries passed with Boto3 1.34, 1.36, and 1.43.