diff --git a/cpp/src/arrow/filesystem/s3fs.cc b/cpp/src/arrow/filesystem/s3fs.cc index 1c6763a4aee9..e72695c096d8 100644 --- a/cpp/src/arrow/filesystem/s3fs.cc +++ b/cpp/src/arrow/filesystem/s3fs.cc @@ -1748,8 +1748,13 @@ class ObjectOutputStream final : public io::OutputStream { return Status::OK(); } + // Close even if the abort request fails + auto holder = std::move(holder_); + current_part_.reset(); + closed_ = true; + if (IsMultipartCreated()) { - ARROW_ASSIGN_OR_RAISE(auto client_lock, holder_->Lock()); + ARROW_ASSIGN_OR_RAISE(auto client_lock, holder->Lock()); S3Model::AbortMultipartUploadRequest req; req.SetBucket(ToAwsString(path_.bucket)); @@ -1765,10 +1770,6 @@ class ObjectOutputStream final : public io::OutputStream { } } - current_part_.reset(); - holder_ = nullptr; - closed_ = true; - return Status::OK(); } diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 17dcb87a80e9..6f54eb646a76 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -1633,6 +1633,7 @@ cdef extern from "arrow/io/api.h" namespace "arrow::io" nogil: cdef cppclass FileInterface: CStatus Close() + CStatus Abort() CResult[int64_t] Tell() FileMode mode() c_bool closed() diff --git a/python/pyarrow/io.pxi b/python/pyarrow/io.pxi index dbecd80909c6..aad002afce8b 100644 --- a/python/pyarrow/io.pxi +++ b/python/pyarrow/io.pxi @@ -193,6 +193,20 @@ cdef class NativeFile(_Weakrefable): else: check_status(self.output_stream.get().Close()) + def abort(self): + """ + Close the stream, discarding written data if the stream supports it. + + For example, an S3 output stream aborts its multipart upload, so no + object is written. Other streams are simply closed. + """ + if not self.closed: + with nogil: + if self.is_readable: + check_status(self.input_stream.get().Abort()) + else: + check_status(self.output_stream.get().Abort()) + cdef set_random_access_file(self, shared_ptr[CRandomAccessFile] handle): self.input_stream = handle self.random_access = handle diff --git a/python/pyarrow/tests/test_fs.py b/python/pyarrow/tests/test_fs.py index 5bf1950c0654..4db78a04ad95 100644 --- a/python/pyarrow/tests/test_fs.py +++ b/python/pyarrow/tests/test_fs.py @@ -1143,6 +1143,68 @@ def test_open_output_stream_metadata(fs, pathfn): assert got_metadata == {} +@pytest.mark.gzip +@pytest.mark.parametrize( + ('compression', 'buffer_size'), + [ + (None, None), + (None, 64), + ('gzip', None), + ('gzip', 256), + ] +) +def test_open_output_stream_abort(fs, pathfn, compression, buffer_size): + p = pathfn('open-output-stream-abort') + with fs.open_output_stream(p, compression, buffer_size) as f: + f.write(b'some data') + f.abort() + assert f.closed + + # Note that only S3 discards the written data on abort for now + if fs.type_name == 's3': + assert fs.get_file_info(p).type == FileType.NotFound + elif 'mock' in fs.type_name: + with fs.open_input_stream(p) as f: + assert f.read().startswith(b'MockFSOutputStream aborted') + + +@pytest.mark.s3 +def test_s3_output_stream_abort_after_part_upload(s3fs): + fs, pathfn = s3fs['fs'], s3fs['pathfn'] + p = pathfn('abort-after-part-upload') + with fs.open_output_stream(p) as f: + # Flushing a full 10 MiB part waits for its upload to complete + f.write(b'x' * 10 * 1024 * 1024) + f.flush() + f.abort() + + assert fs.get_file_info(p).type == FileType.NotFound + + +@pytest.mark.s3 +def test_s3_output_stream_failed_abort(s3_server): + from pyarrow.fs import S3FileSystem + # The limited user isn't allowed to abort multipart uploads + _configure_s3_limited_user(s3_server, _minio_limited_policy, + 'test_fs_abort_user', 'abort123') + host, port, _, _ = s3_server['connection'] + fs = S3FileSystem( + access_key='test_fs_abort_user', + secret_key='abort123', + endpoint_override=f'{host}:{port}', + scheme='http' + ) + p = 'existing-bucket/failed-abort' + with pytest.raises(OSError, match="AbortMultipartUpload"): + with fs.open_output_stream(p) as f: + f.write(b'some data') + f.abort() + assert f.closed + + del f + assert fs.get_file_info(p).type == FileType.NotFound + + def test_localfs_options(): # LocalFileSystem instantiation LocalFileSystem(use_mmap=False) diff --git a/python/pyarrow/tests/test_io.py b/python/pyarrow/tests/test_io.py index 1cfebb4936a1..bafa89ff99dc 100644 --- a/python/pyarrow/tests/test_io.py +++ b/python/pyarrow/tests/test_io.py @@ -989,6 +989,25 @@ def test_inmemory_write_after_closed(): f.write(b'not ok') +def test_inmemory_write_after_abort(): + f = pa.BufferOutputStream() + f.write(b'ok') + f.abort() + assert f.closed + + with pytest.raises(ValueError): + f.write(b'not ok') + + +def test_inmemory_read_after_abort(): + f = pa.BufferReader(b'data') + f.abort() + assert f.closed + + with pytest.raises(ValueError): + f.read() + + def test_buffer_protocol_ref_counting(): def make_buffer(bytes_obj): return bytearray(pa.py_buffer(bytes_obj))