GH-40502: [C++][Python] Expose NativeFile.abort()

cpp/src/arrow/filesystem/s3fs.cc

@@ -1750,8 +1750,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));
@@ -1767,10 +1772,6 @@ class ObjectOutputStream final : public io::OutputStream {
       }
     }
 
-    current_part_.reset();
-    holder_ = nullptr;
-    closed_ = true;
-
     return Status::OK();
   }
 

python/pyarrow/includes/libarrow.pxd

@@ -1634,6 +1634,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()

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 abruptly.
+
+        The stream is always closed, but whether written data is flushed or
+        discarded depends on the implementation.
+        """
+        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 = <shared_ptr[CInputStream]> handle
         self.random_access = handle

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
+
+    # Only S3 is guaranteed to discard the written data on abort
+    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)

python/pyarrow/tests/test_io.py

@@ -102,6 +102,16 @@ def test_python_file_write():
         pa.PythonFile(StringIO())
 
 
+def test_python_file_abort():
+    buf = BytesIO()
+    f = pa.PythonFile(buf, mode='w')
+    f.write(b'data')
+    f.abort()
+    assert f.closed
+    assert not buf.closed
+    assert buf.getvalue() == b'data'
+
+
 def test_python_file_read():
     data = b'some sample data'
 
@@ -989,6 +999,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))