This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-mosaic.git
The following commit(s) were added to refs/heads/main by this push:
new c082354 fix(python): handle partial stream writes (#94)
c082354 is described below
commit c082354aaa1e7cc4a61154b45aa2066e8b2148cb
Author: jianguotian <[email protected]>
AuthorDate: Wed Sep 30 07:16:25 2026 +0800
fix(python): handle partial stream writes (#94)
---
python/mosaic/mosaic.py | 12 +++++++-
python/tests/test_mosaic.py | 68 +++++++++++++++++++++++++++++++++++++++++++++
2 files changed, 79 insertions(+), 1 deletion(-)
diff --git a/python/mosaic/mosaic.py b/python/mosaic/mosaic.py
index 4796890..16d39ae 100644
--- a/python/mosaic/mosaic.py
+++ b/python/mosaic/mosaic.py
@@ -16,6 +16,7 @@
# under the License.
import ctypes
+import operator
import pyarrow as pa
@@ -221,7 +222,16 @@ class MosaicWriter:
def _on_write(self, ctx, data, length):
try:
buf = (ctypes.c_char * length).from_address(ctypes.cast(data,
ctypes.c_void_p).value)
- self._stream.write(buf)
+ view = memoryview(buf)
+ offset = 0
+ while offset < length:
+ written = self._stream.write(view[offset:])
+ if written is None:
+ return -1
+ written = operator.index(written)
+ if written <= 0 or written > length - offset:
+ return -1
+ offset += written
self._pos += length
return 0
except Exception:
diff --git a/python/tests/test_mosaic.py b/python/tests/test_mosaic.py
index f88720c..0057acc 100644
--- a/python/tests/test_mosaic.py
+++ b/python/tests/test_mosaic.py
@@ -15,9 +15,11 @@
# specific language governing permissions and limitations
# under the License.
+import ctypes
import io
import struct
import traceback
+from types import SimpleNamespace
import pyarrow as pa
import pytest
@@ -46,6 +48,13 @@ def _reader_from_bytes(data):
)
+def _invoke_write_callback(stream, data, position=0):
+ state = SimpleNamespace(_stream=stream, _pos=position)
+ buf = (ctypes.c_uint8 * len(data)).from_buffer_copy(data)
+ result = MosaicWriter._on_write(state, None, buf, len(data))
+ return result, state._pos
+
+
class TestRoundtrip:
def test_basic_roundtrip(self):
pa_schema = pa.schema(
@@ -752,6 +761,65 @@ class TestConvenience:
class TestWriter:
+ def test_write_callback_handles_full_write(self):
+ class FullWriteOutput(io.BytesIO):
+ def __init__(self):
+ super().__init__()
+ self.request_sizes = []
+
+ def write(self, data):
+ self.request_sizes.append(len(data))
+ return super().write(data)
+
+ stream = FullWriteOutput()
+ result, position = _invoke_write_callback(stream, b"abcdef",
position=7)
+
+ assert result == 0
+ assert position == 13
+ assert stream.getvalue() == b"abcdef"
+ assert stream.request_sizes == [6]
+
+ def test_write_callback_retries_short_writes(self):
+ class ShortWriteOutput(io.BytesIO):
+ def __init__(self):
+ super().__init__()
+ self.request_sizes = []
+
+ def write(self, data):
+ self.request_sizes.append(len(data))
+ return super().write(data[:2])
+
+ stream = ShortWriteOutput()
+ result, position = _invoke_write_callback(stream, b"abcdef",
position=7)
+
+ assert result == 0
+ assert position == 13
+ assert stream.getvalue() == b"abcdef"
+ assert stream.request_sizes == [6, 4, 2]
+
+ @pytest.mark.parametrize("stall_result", [0, None], ids=["zero", "none"])
+ def test_write_callback_rejects_stall_without_advancing_position(
+ self, stall_result
+ ):
+ class StallingOutput(io.BytesIO):
+ def __init__(self):
+ super().__init__()
+ self.request_sizes = []
+
+ def write(self, data):
+ self.request_sizes.append(len(data))
+ if len(self.request_sizes) == 1:
+ return super().write(data[:2])
+ return stall_result
+
+ stream = StallingOutput()
+ result, position = _invoke_write_callback(stream, b"abcdef",
position=7)
+
+ assert result == -1
+ assert position == 7
+ assert stream.getvalue() == b"ab"
+ assert stream.request_sizes == [6, 4]
+
def test_estimated_file_size(self):
pa_schema = pa.schema(
[pa.field("x", pa.int32()), pa.field("y", pa.utf8())]
