| # -*- encoding: utf-8 -*- |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| import threading |
| import unittest |
| |
| from pyspark.messages.zero_copy_byte_stream import ZeroCopyByteStream |
| |
| |
| class ZeroCopyByteStreamTests(unittest.TestCase): |
| """Tests for ZeroCopyByteStream.""" |
| |
| # ---- Basic single-chunk reads (zero-copy fast path) ---- |
| |
| def test_read_exact_chunk(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"hello")) |
| result = stream.read(5) |
| self.assertEqual(bytes(result), b"hello") |
| |
| def test_read_partial_chunk(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"hello world")) |
| stream.finish() |
| r1 = stream.read(5) |
| r2 = stream.read(6) |
| self.assertEqual(bytes(r1), b"hello") |
| self.assertEqual(bytes(r2), b" world") |
| # Check EOF read |
| with self.assertRaises(EOFError): |
| stream.read(1) |
| |
| def test_read_multiple_chunks_sequentially(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"aaa")) |
| stream.add_next_chunk(memoryview(b"bbb")) |
| stream.add_next_chunk(memoryview(b"ccc")) |
| |
| self.assertEqual(bytes(stream.read(3)), b"aaa") |
| self.assertEqual(bytes(stream.read(3)), b"bbb") |
| self.assertEqual(bytes(stream.read(3)), b"ccc") |
| |
| # ---- Cross-boundary reads (slow path with copy) ---- |
| |
| def test_read_across_two_chunks(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"aaa")) |
| stream.add_next_chunk(memoryview(b"bbb")) |
| |
| result = stream.read(6) |
| self.assertEqual(bytes(result), b"aaabbb") |
| |
| def test_read_across_three_chunks(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"ab")) |
| stream.add_next_chunk(memoryview(b"cd")) |
| stream.add_next_chunk(memoryview(b"ef")) |
| |
| result = stream.read(6) |
| self.assertEqual(bytes(result), b"abcdef") |
| |
| def test_read_partial_then_cross_boundary(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"aabb")) |
| stream.add_next_chunk(memoryview(b"ccdd")) |
| |
| # Read first 2 bytes from chunk 1 (zero-copy) |
| r1 = stream.read(2) |
| self.assertEqual(bytes(r1), b"aa") |
| |
| # Read 4 bytes crossing chunk boundary |
| r2 = stream.read(4) |
| self.assertEqual(bytes(r2), b"bbcc") |
| |
| # Read remaining 2 bytes from chunk 2 (zero-copy) |
| r3 = stream.read(2) |
| self.assertEqual(bytes(r3), b"dd") |
| |
| def test_cross_boundary_read_consumes_full_middle_chunk(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"aa")) |
| stream.add_next_chunk(memoryview(b"bb")) |
| stream.add_next_chunk(memoryview(b"cc")) |
| |
| # Read 1 byte to offset into first chunk |
| stream.read(1) |
| |
| # Read 5 bytes: 1 from chunk1 + 2 from chunk2 + 2 from chunk3 |
| result = stream.read(5) |
| self.assertEqual(bytes(result), b"abbcc") |
| |
| # ---- EOF handling ---- |
| |
| def test_eof_throws_eof_error(self): |
| stream = ZeroCopyByteStream() |
| stream.finish() |
| |
| self.assertTrue(stream.finished) |
| with self.assertRaises(EOFError): |
| stream.read(1) |
| |
| def test_eof_after_consuming_all_data(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"data")) |
| stream.finish() |
| |
| self.assertFalse(stream.finished) |
| |
| result = stream.read(4) |
| self.assertEqual(bytes(result), b"data") |
| |
| self.assertTrue(stream.finished) |
| with self.assertRaises(EOFError): |
| stream.read(1) |
| |
| def test_eof_with_out_of_bounds_read(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"data")) |
| stream.finish() |
| |
| result = stream.read(3) |
| self.assertEqual(bytes(result), b"dat") |
| |
| self.assertFalse(stream.finished) |
| with self.assertRaises(EOFError): |
| stream.read(2) |
| |
| def test_eof_during_cross_boundary_read(self): |
| """EOF mid-cross-boundary read returns None.""" |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"ab")) |
| stream.add_next_chunk(memoryview(b"cd")) |
| stream.finish() |
| |
| # Request more bytes than available; after consuming "abcd", |
| # _try_read_bytes hits EOF and throws |
| with self.assertRaises(EOFError): |
| stream.read(5) |
| |
| def test_finished_property(self): |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"x")) |
| self.assertFalse(stream.finished) |
| |
| stream.read(1) |
| self.assertFalse(stream.finished) # not yet marked EOF |
| |
| stream.finish() |
| self.assertTrue(stream.finished) |
| |
| # ---- Threading / blocking behavior ---- |
| |
| def test_read_blocks_until_chunk_available(self): |
| stream = ZeroCopyByteStream() |
| result = None |
| |
| def reader(): |
| nonlocal result |
| result = stream.read(3) |
| |
| t = threading.Thread(target=reader) |
| t.start() |
| |
| # Give reader time to block |
| t.join(timeout=2) |
| self.assertTrue(t.is_alive(), "Reader should be blocked waiting for data") |
| |
| stream.add_next_chunk(memoryview(b"abc")) |
| t.join(timeout=2) |
| self.assertFalse(t.is_alive(), "Reader should have unblocked") |
| self.assertEqual(bytes(result), b"abc") |
| |
| def test_read_blocks_until_eof(self): |
| stream = ZeroCopyByteStream() |
| read_raised_eof = False |
| read_called = threading.Event() |
| |
| def reader(): |
| nonlocal read_raised_eof |
| read_called.set() |
| try: |
| stream.read(1) |
| except EOFError: |
| read_raised_eof = True |
| |
| t = threading.Thread(target=reader) |
| t.start() |
| read_called.wait() |
| |
| # Give reader time to block |
| t.join(timeout=2) |
| self.assertTrue(t.is_alive()) |
| |
| stream.finish() |
| t.join(timeout=2) |
| self.assertFalse(t.is_alive()) |
| self.assertTrue(read_raised_eof) |
| |
| def test_cross_boundary_read_blocks_for_next_chunk(self): |
| """Cross-boundary read blocks when second chunk isn't available yet.""" |
| stream = ZeroCopyByteStream(initial_view=memoryview(b"aa")) |
| result = None |
| |
| def reader(): |
| nonlocal result |
| result = stream.read(4) |
| |
| t = threading.Thread(target=reader) |
| t.start() |
| |
| # Reader consumed "aa" but needs 2 more bytes |
| t.join(timeout=2) |
| self.assertTrue(t.is_alive(), "Reader should block waiting for more data") |
| |
| stream.add_next_chunk(memoryview(b"bb")) |
| t.join(timeout=2) |
| self.assertFalse(t.is_alive()) |
| self.assertEqual(bytes(result), b"aabb") |
| |
| def test_read_of_zero_bytes_succeeds_without_data(self): |
| """Reading zero bytes should immediately return, even if no data is present""" |
| stream = ZeroCopyByteStream(initial_view=None) |
| |
| res = stream.read(0) |
| |
| self.assertEqual(res, memoryview(b"")) |
| |
| def test_negative_read_throws_value_error(self): |
| stream = ZeroCopyByteStream(initial_view=None) |
| |
| with self.assertRaises(ValueError): |
| stream.read(-1) |
| |
| # ---- add_next_chunk assertions ---- |
| |
| def test_add_none_chunk_raises(self): |
| stream = ZeroCopyByteStream() |
| with self.assertRaises(TypeError): |
| stream.add_next_chunk(None) |
| |
| def test_add_chunk_after_finish_raises(self): |
| stream = ZeroCopyByteStream() |
| stream.finish() |
| with self.assertRaises(ValueError): |
| stream.add_next_chunk(memoryview(b"data")) |
| |
| # ---- Initial view ---- |
| |
| def test_no_initial_view(self): |
| stream = ZeroCopyByteStream() |
| stream.add_next_chunk(memoryview(b"hello")) |
| result = stream.read(5) |
| self.assertEqual(bytes(result), b"hello") |
| |
| def test_initial_view_none(self): |
| stream = ZeroCopyByteStream(initial_view=None) |
| stream.add_next_chunk(memoryview(b"test")) |
| result = stream.read(4) |
| self.assertEqual(bytes(result), b"test") |
| |
| def test_invalid_initial_view(self): |
| with self.assertRaises(TypeError): |
| ZeroCopyByteStream(initial_view=5) |
| |
| |
| if __name__ == "__main__": |
| from pyspark.testing import main |
| |
| main() |
