| # 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. | |
| """Tests for ``WriteBuffer`` and the deferred-fold write path. | |
| The regression these guard: ``DataWriter.write`` used to fold every incoming | |
| batch into its buffer and then measure the result, which made writing N batches | |
| O(N^2) -- ``pa.concat_tables`` leaves N chunks per column and the | |
| ``pa.Table.nbytes`` walk behind each rolling decision re-visits all of them. | |
| Folding is now deferred, so the tests below assert both halves of that: the | |
| fold count no longer scales with the number of writes, and every rolling | |
| trigger the eager path had still fires. | |
| """ | |
| import contextlib | |
| import unittest | |
| import pyarrow as pa | |
| from pypaimon.manifest.schema.data_file_meta import DataFileMeta | |
| from pypaimon.execution.python.write.writer.append_only_data_writer import AppendOnlyDataWriter | |
| from pypaimon.execution.python.write.writer.data_vector_writer import DataVectorWriter | |
| from pypaimon.execution.python.write.writer.data_writer import DataWriter | |
| from pypaimon.execution.python.write.writer.dedicated_format_writer import DedicatedFormatWriter | |
| from pypaimon.execution.python.write.writer.write_buffer import WriteBuffer | |
| _SCHEMA = pa.schema([ | |
| pa.field('id', pa.int64(), nullable=False), | |
| pa.field('name', pa.string()), | |
| ]) | |
| # Anything past this is effectively "no rolling"; matches how | |
| # ``target_file_row_num`` defaults to the max long when the option is unset. | |
| _NO_LIMIT = 2 ** 63 - 1 | |
| def _batch(start: int, num_rows: int) -> pa.RecordBatch: | |
| return pa.RecordBatch.from_pydict( | |
| { | |
| 'id': list(range(start, start + num_rows)), | |
| 'name': ['n%d' % i for i in range(start, start + num_rows)], | |
| }, | |
| schema=_SCHEMA, | |
| ) | |
| def _table(start: int, num_rows: int) -> pa.Table: | |
| return pa.Table.from_batches([_batch(start, num_rows)]) | |
| # Same columns and types as ``_SCHEMA``, but ``id`` is nullable. This is one of | |
| # the differences ``TableWrite._validate_pyarrow_schema`` lets through (it only | |
| # compares field types) while ``pa.concat_tables`` rejects it. | |
| _NULLABLE_SCHEMA = pa.schema([ | |
| pa.field('id', pa.int64()), | |
| pa.field('name', pa.string()), | |
| ]) | |
| # ``_SCHEMA`` carrying metadata, which neither ``Schema.equals`` nor | |
| # ``concat_tables`` looks at. | |
| _ANNOTATED_SCHEMA = pa.schema( | |
| [ | |
| pa.field('id', pa.int64(), nullable=False, metadata={b'k': b'v'}), | |
| pa.field('name', pa.string()), | |
| ], | |
| metadata={b'origin': b'test'}, | |
| ) | |
| def _batch_with(schema: pa.Schema, start: int, num_rows: int) -> pa.RecordBatch: | |
| return pa.RecordBatch.from_pydict( | |
| { | |
| 'id': list(range(start, start + num_rows)), | |
| 'name': ['n%d' % i for i in range(start, start + num_rows)], | |
| }, | |
| schema=schema, | |
| ) | |
| def _table_with(schema: pa.Schema, start: int, num_rows: int) -> pa.Table: | |
| return pa.Table.from_batches([_batch_with(schema, start, num_rows)]) | |
| @contextlib.contextmanager | |
| def _count_folds(): | |
| """Count the concat calls a fold performs, as a list of one entry each. | |
| ``pa.concat_tables`` is the only way the buffer or a writer's | |
| ``_merge_data`` collapses tables, so its call count is the fold count. The | |
| eager path called it once per write; the deferred path calls it once per | |
| ``materialize`` that has something to fold. | |
| """ | |
| calls = [] | |
| original = pa.concat_tables | |
| def counting(*args, **kwargs): | |
| calls.append(1) | |
| return original(*args, **kwargs) | |
| pa.concat_tables = counting | |
| try: | |
| yield calls | |
| finally: | |
| pa.concat_tables = original | |
| class _Harness(AppendOnlyDataWriter): | |
| """Append-only writer with the file layer stubbed out. | |
| Bypasses ``DataWriter.__init__`` -- which needs a real table, catalog and | |
| file IO -- and sets up only what the write/roll path reads. | |
| """ | |
| def __init__(self, target_file_size: int = _NO_LIMIT, | |
| target_file_row_num: int = _NO_LIMIT): | |
| self.target_file_size = target_file_size | |
| self.target_file_row_num = target_file_row_num | |
| self._buffer = WriteBuffer(self._merge_data) | |
| self.committed_files = [] | |
| self.written_chunks = [] | |
| self.aborted = False | |
| def _write_data_to_file(self, data: pa.Table): | |
| self.written_chunks.append(data) | |
| def abort(self): | |
| # The real ``abort`` deletes the committed files through file IO this | |
| # harness has none of; record that it ran and do the rest. | |
| self.aborted = True | |
| self._buffer.reset() | |
| self.committed_files.clear() | |
| class WriteBufferTest(unittest.TestCase): | |
| def _buffer(self): | |
| return WriteBuffer(lambda a, b: pa.concat_tables([a, b])) | |
| def test_append_tracks_counts_without_folding(self): | |
| buffer = self._buffer() | |
| with _count_folds() as folds: | |
| for i in range(50): | |
| buffer.append(_table(i * 10, 10)) | |
| self.assertEqual(folds, []) | |
| self.assertEqual(buffer.num_rows, 500) | |
| self.assertFalse(buffer.is_empty) | |
| def test_running_nbytes_matches_the_folded_table(self): | |
| # The gate on the write path trusts the running sum, so it has to agree | |
| # with what the folded table reports -- separate tables share no Arrow | |
| # buffers, so the per-table sizes add up exactly. | |
| buffer = self._buffer() | |
| for i in range(10): | |
| buffer.append(_table(i * 10, 10)) | |
| running = buffer.nbytes | |
| self.assertEqual(running, buffer.materialize().nbytes) | |
| self.assertEqual(running, buffer.nbytes) | |
| def test_materialize_folds_once_and_is_idempotent(self): | |
| buffer = self._buffer() | |
| for i in range(20): | |
| buffer.append(_table(i, 1)) | |
| with _count_folds() as folds: | |
| first = buffer.materialize() | |
| second = buffer.materialize() | |
| # One concat for the 20 appended tables, and nothing on re-read. | |
| self.assertEqual(len(folds), 1) | |
| self.assertIs(first, second) | |
| self.assertEqual(first.num_rows, 20) | |
| self.assertEqual(first.column('id').to_pylist(), list(range(20))) | |
| def test_single_append_skips_the_concat_entirely(self): | |
| buffer = self._buffer() | |
| buffer.append(_table(0, 5)) | |
| with _count_folds() as folds: | |
| self.assertEqual(buffer.materialize().num_rows, 5) | |
| self.assertEqual(folds, []) | |
| def test_append_after_materialize_goes_through_merge(self): | |
| merged = [] | |
| def merge(existing, new): | |
| merged.append((existing.num_rows, new.num_rows)) | |
| return pa.concat_tables([existing, new]) | |
| buffer = WriteBuffer(merge) | |
| buffer.append(_table(0, 3)) | |
| buffer.materialize() | |
| buffer.append(_table(3, 4)) | |
| self.assertEqual(buffer.materialize().num_rows, 7) | |
| self.assertEqual(merged, [(3, 4)]) | |
| def test_empty_buffer_materializes_to_none(self): | |
| buffer = self._buffer() | |
| self.assertTrue(buffer.is_empty) | |
| self.assertIsNone(buffer.materialize()) | |
| self.assertEqual(buffer.nbytes, 0) | |
| self.assertEqual(buffer.num_rows, 0) | |
| def test_reset_replaces_contents_and_remeasures(self): | |
| buffer = self._buffer() | |
| buffer.append(_table(0, 100)) | |
| replacement = _table(0, 7) | |
| buffer.reset(replacement) | |
| self.assertIs(buffer.materialize(), replacement) | |
| self.assertEqual(buffer.num_rows, 7) | |
| self.assertEqual(buffer.nbytes, replacement.nbytes) | |
| def test_reset_to_none_empties_the_buffer(self): | |
| buffer = self._buffer() | |
| buffer.append(_table(0, 100)) | |
| buffer.reset() | |
| self.assertTrue(buffer.is_empty) | |
| self.assertIsNone(buffer.materialize()) | |
| self.assertEqual(buffer.num_rows, 0) | |
| self.assertEqual(buffer.nbytes, 0) | |
| def test_zero_row_table_is_not_reported_empty(self): | |
| # ``is_empty`` has to distinguish "nothing set" from "an empty table was | |
| # set", because the writers use it to decide whether there is anything | |
| # to roll at all. | |
| buffer = self._buffer() | |
| buffer.reset(_table(0, 0)) | |
| self.assertFalse(buffer.is_empty) | |
| self.assertEqual(buffer.num_rows, 0) | |
| self.assertIsNotNone(buffer.materialize()) | |
| class SchemaGuardTest(unittest.TestCase): | |
| """The guard keeps a mismatch failing where the eager fold failed. | |
| Folding on append meant a batch ``concat_tables`` could not accept raised | |
| inside ``DataWriter.write``, which aborts and cleans up. Deferring the fold | |
| would otherwise push that failure out to ``prepare_commit``, which has no | |
| such handler, so ``append`` rejects up front exactly what concat rejects. | |
| """ | |
| def _buffer(self): | |
| return WriteBuffer(lambda a, b: pa.concat_tables([a, b])) | |
| def test_append_rejects_a_schema_concat_would_reject(self): | |
| buffer = self._buffer() | |
| buffer.append(_table(0, 3)) | |
| with self.assertRaises(ValueError) as caught: | |
| buffer.append(_table_with(_NULLABLE_SCHEMA, 3, 3)) | |
| self.assertIn('schema differs', str(caught.exception)) | |
| # And a rejected batch leaves the running counts describing the rows | |
| # that are actually buffered. | |
| self.assertEqual(buffer.num_rows, 3) | |
| self.assertEqual(buffer.materialize().num_rows, 3) | |
| def test_append_rejects_a_mismatch_after_a_materialize(self): | |
| # ``materialize`` rebases the schema onto the folded table, so the | |
| # second half of a buffer's life is guarded too. | |
| buffer = self._buffer() | |
| buffer.append(_table(0, 3)) | |
| buffer.materialize() | |
| with self.assertRaises(ValueError): | |
| buffer.append(_table_with(_NULLABLE_SCHEMA, 3, 3)) | |
| def test_append_accepts_a_metadata_only_difference(self): | |
| # ``concat_tables`` ignores schema and field metadata, so the guard has | |
| # to as well or it would reject batches the old path folded fine. | |
| buffer = self._buffer() | |
| buffer.append(_table(0, 3)) | |
| buffer.append(_table_with(_ANNOTATED_SCHEMA, 3, 4)) | |
| self.assertEqual(buffer.materialize().num_rows, 7) | |
| def test_reset_rebases_the_schema(self): | |
| # A writer that resets to a new table starts a new file; batches | |
| # matching that table's schema have to be accepted afterwards. | |
| buffer = self._buffer() | |
| buffer.append(_table(0, 3)) | |
| buffer.reset(_table_with(_NULLABLE_SCHEMA, 0, 2)) | |
| buffer.append(_table_with(_NULLABLE_SCHEMA, 2, 2)) | |
| self.assertEqual(buffer.materialize().num_rows, 4) | |
| with self.assertRaises(ValueError): | |
| buffer.append(_table(4, 1)) | |
| def test_first_append_after_an_empty_reset_sets_the_schema(self): | |
| buffer = self._buffer() | |
| buffer.append(_table(0, 3)) | |
| buffer.reset() | |
| buffer.append(_table_with(_NULLABLE_SCHEMA, 0, 3)) | |
| self.assertEqual(buffer.materialize().num_rows, 3) | |
| class DeferredFoldWritePathTest(unittest.TestCase): | |
| def test_many_small_writes_fold_a_constant_number_of_times(self): | |
| # The regression: with an eager fold this was 199 concats for 200 | |
| # writes, each walking a buffer one chunk longer than the last. | |
| writer = _Harness() | |
| with _count_folds() as folds: | |
| for i in range(200): | |
| writer.write(_batch(i * 5, 5)) | |
| self.assertEqual(folds, []) | |
| # Reading the buffer folds it -- once, not once per write. | |
| with _count_folds() as folds: | |
| pending = writer._buffer.materialize() | |
| self.assertEqual(len(folds), 1) | |
| self.assertEqual(pending.num_rows, 1000) | |
| self.assertEqual(pending.column('id').to_pylist(), list(range(1000))) | |
| def test_row_num_rolling_still_fires(self): | |
| # ``target_file_row_num`` postdates the eager fold, so a byte-only gate | |
| # would leave row-based rolling silently dead for tables that set it | |
| # while staying well under target_file_size. | |
| writer = _Harness(target_file_row_num=10) | |
| for i in range(10): | |
| writer.write(_batch(i * 3, 3)) | |
| self.assertEqual([c.num_rows for c in writer.written_chunks], | |
| [10, 10]) | |
| self.assertEqual(writer.pending_row_count, 10) | |
| self.assertEqual( | |
| [row for c in writer.written_chunks | |
| for row in c.column('id').to_pylist()], | |
| list(range(20)), | |
| ) | |
| def test_row_num_rolling_splits_a_single_oversized_write(self): | |
| writer = _Harness(target_file_row_num=4) | |
| writer.write(_batch(0, 14)) | |
| self.assertEqual([c.num_rows for c in writer.written_chunks], | |
| [4, 4, 4]) | |
| self.assertEqual(writer.pending_row_count, 2) | |
| def test_size_rolling_still_bounds_written_files(self): | |
| target = _table(0, 100).nbytes // 4 | |
| writer = _Harness(target_file_size=target) | |
| for i in range(20): | |
| writer.write(_batch(i * 5, 5)) | |
| self.assertGreaterEqual(len(writer.written_chunks), 3) | |
| for chunk in writer.written_chunks: | |
| self.assertLessEqual(chunk.nbytes, target) | |
| written = sum(c.num_rows for c in writer.written_chunks) | |
| self.assertEqual(written + writer.pending_row_count, 100) | |
| def test_row_that_alone_exceeds_target_size_is_rolled_by_itself(self): | |
| writer = _Harness(target_file_size=1) | |
| writer.write(_batch(0, 3)) | |
| # Rolls the first two rows one at a time; the last stays buffered | |
| # because the loop stops once the buffer is down to a single row. | |
| self.assertEqual([c.num_rows for c in writer.written_chunks], [1, 1]) | |
| self.assertEqual(writer.pending_row_count, 1) | |
| def test_rolling_does_not_refold_on_every_subsequent_write(self): | |
| # After a roll the remainder goes back into the buffer, which | |
| # re-measures it. If that left the roll condition stuck open, every | |
| # later write would fold again and the quadratic would be back. | |
| writer = _Harness(target_file_row_num=8) | |
| writer.write(_batch(0, 12)) | |
| # Rolled 8 rows, 4 left buffered. | |
| self.assertEqual(len(writer.written_chunks), 1) | |
| with _count_folds() as folds: | |
| for i in range(20): | |
| writer.write(_batch(100 + i, 1)) | |
| # 4 buffered + 20 single-row writes crosses 8 rows exactly twice, so | |
| # only two of those 20 writes fold anything. Asserting a bound rather | |
| # than an exact count: how many concats one fold takes is up to | |
| # ``WriteBuffer``, but it must not be one per write. | |
| self.assertLess(len(folds), 20) | |
| self.assertEqual(len(writer.written_chunks), 3) | |
| def test_pending_row_count_does_not_fold(self): | |
| # The composite writers read this on every write to size the next | |
| # slice; folding here would bring the quadratic back through the side | |
| # door. | |
| writer = _Harness() | |
| with _count_folds() as folds: | |
| for i in range(30): | |
| writer.write(_batch(i, 1)) | |
| self.assertEqual(writer.pending_row_count, i + 1) | |
| self.assertEqual(folds, []) | |
| def test_prepare_commit_flushes_the_deferred_buffer(self): | |
| writer = _Harness() | |
| for i in range(10): | |
| writer.write(_batch(i * 2, 2)) | |
| writer.prepare_commit() | |
| self.assertEqual([c.num_rows for c in writer.written_chunks], [20]) | |
| self.assertTrue(writer._buffer.is_empty) | |
| def test_close_flushes_the_deferred_buffer(self): | |
| writer = _Harness() | |
| for i in range(10): | |
| writer.write(_batch(i * 2, 2)) | |
| writer.close() | |
| self.assertEqual([c.num_rows for c in writer.written_chunks], [20]) | |
| self.assertTrue(writer._buffer.is_empty) | |
| def test_mismatched_schema_fails_the_write_that_carries_it(self): | |
| # Not ``prepare_commit``: only ``write`` aborts, so a failure deferred | |
| # to commit time would leave the files rolled so far orphaned. | |
| writer = _Harness() | |
| writer.write(_batch(0, 3)) | |
| with self.assertRaises(ValueError): | |
| writer.write(_batch_with(_NULLABLE_SCHEMA, 3, 3)) | |
| self.assertTrue(writer.aborted) | |
| class FlushFailureTest(unittest.TestCase): | |
| """A failed flush leaves the rows buffered for the retry. | |
| ``StreamTableWrite`` is reusable, so a transient storage error followed by | |
| another ``prepare_commit`` on the same writer has to write the same rows, | |
| not silently skip them. Draining the buffer before the write would lose | |
| them. | |
| """ | |
| class _FailOnceHarness(_Harness): | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.fail_next = True | |
| def _write_data_to_file(self, data: pa.Table): | |
| if self.fail_next: | |
| self.fail_next = False | |
| raise IOError('transient storage failure') | |
| super()._write_data_to_file(data) | |
| def test_failed_prepare_commit_keeps_the_rows_for_the_retry(self): | |
| writer = self._FailOnceHarness() | |
| for i in range(3): | |
| writer.write(_batch(i, 1)) | |
| with self.assertRaises(IOError): | |
| writer.prepare_commit() | |
| self.assertEqual(writer.pending_row_count, 3) | |
| writer.prepare_commit() | |
| self.assertEqual([c.num_rows for c in writer.written_chunks], [3]) | |
| self.assertEqual(writer.pending_row_count, 0) | |
| class _FailOnceVectorHarness(DataVectorWriter): | |
| """``_close_current_writers`` with the file layer stubbed out. | |
| No ``vector_writer``, so this covers the normal half on its own; the | |
| point is only where the buffer is cleared relative to the write. | |
| """ | |
| def __init__(self): | |
| self.target_file_size = _NO_LIMIT | |
| self.target_file_row_num = _NO_LIMIT | |
| self.record_count = 0 | |
| self.vector_writer = None | |
| self._normal_buffer = WriteBuffer(self._merge_data) | |
| self.committed_files = [] | |
| self._committed_files_to_delete_on_abort = [] | |
| self._pending_normal_meta = None | |
| self.written = [] | |
| self.fail_next = True | |
| def _write_normal_data_to_file(self, data: pa.Table): | |
| if self.fail_next: | |
| self.fail_next = False | |
| raise IOError('transient storage failure') | |
| self.written.append(data) | |
| return object() | |
| def test_failed_normal_flush_keeps_the_rows_for_the_retry(self): | |
| # Otherwise the retry finds no normal_meta, flushes the sidecars alone | |
| # and skips the row-count check, committing sidecar-only metadata. | |
| writer = self._FailOnceVectorHarness() | |
| writer._normal_buffer.append(_table(0, 3)) | |
| with self.assertRaises(IOError): | |
| writer._close_current_writers() | |
| self.assertEqual(writer._normal_buffer.num_rows, 3) | |
| writer._close_current_writers() | |
| self.assertEqual([t.num_rows for t in writer.written], [3]) | |
| self.assertEqual(writer._normal_buffer.num_rows, 0) | |
| def _file_meta(row_count: int, file_name: str) -> DataFileMeta: | |
| """Use real metadata so flush and abort exercise the shared path resolver.""" | |
| return DataFileMeta( | |
| file_name=file_name, file_size=1, row_count=row_count, | |
| min_key=None, max_key=None, key_stats=None, value_stats=None, | |
| min_sequence_number=0, max_sequence_number=row_count - 1, | |
| schema_id=0, level=0, extra_files=[], | |
| file_path='/warehouse/%s' % file_name) | |
| class _RecordingFileIO: | |
| def __init__(self): | |
| self.deleted = [] | |
| def delete_quietly(self, path): | |
| self.deleted.append(path) | |
| class _StubSidecarWriter: | |
| """A blob/vector writer that fails its first ``prepare_commit``. | |
| Models the real ones in the way that matters here: a failure produces no | |
| metadata, and because the sub-writer drains its own buffer as it writes, a | |
| later call returns whatever has landed so far -- so the parent must be able | |
| to harvest the same metas twice without double-counting them. | |
| """ | |
| def __init__(self, row_count: int, file_name: str, fail_times: int = 0, | |
| delete_on_abort: bool = True): | |
| self.committed_files = [] | |
| self.committed_changelog_files = [] | |
| self.pending_row_count = 0 | |
| self.prepare_commit_calls = 0 | |
| self.aborted = False | |
| self._row_count = row_count | |
| self._file_name = file_name | |
| self._fail_times = fail_times | |
| self._delete_on_abort = delete_on_abort | |
| def prepare_commit(self): | |
| self.prepare_commit_calls += 1 | |
| if self._fail_times > 0: | |
| self._fail_times -= 1 | |
| raise IOError('transient sidecar failure') | |
| if not self.committed_files: | |
| self.committed_files.append( | |
| _file_meta(self._row_count, self._file_name)) | |
| return self.committed_files.copy() | |
| def _release_prepared_files(self): | |
| return DataWriter._release_prepared_files(self) | |
| def delete_file_upon_abort(self): | |
| return self._delete_on_abort | |
| def abort(self): | |
| self.aborted = True | |
| self.committed_files.clear() | |
| class CompositeFlushResumeTest(unittest.TestCase): | |
| """A composite flush publishes all of its files or none of them. | |
| One flush writes the normal data file and then the blob/vector sidecars. The | |
| sidecar writers drain their own buffers as they go, so a failure part way | |
| through cannot be rolled back -- deleting the sidecars that already landed | |
| would lose rows nothing can replay. So the flush resumes instead: the normal | |
| rows stay buffered until their file lands, the landed file is remembered so | |
| the retry skips it, and no metadata is published until every phase is done. | |
| Publishing the normal meta before the sidecars ran, as the code used to, | |
| left the retry writing a second copy of rows the first meta already covered | |
| -- and ``_validate_consistency`` then checked the sidecars against that | |
| second copy only. | |
| """ | |
| class _VectorHarness(DataVectorWriter): | |
| def __init__(self, vector_writer, fail_normal_times: int = 0): | |
| self.target_file_size = _NO_LIMIT | |
| self.target_file_row_num = _NO_LIMIT | |
| self.record_count = 0 | |
| self.vector_writer = vector_writer | |
| self.normal_column_names = ['id', 'name'] | |
| self.vector_write_columns = [] | |
| self._normal_buffer = WriteBuffer(self._merge_data) | |
| self._buffer = WriteBuffer(self._merge_data) | |
| self.committed_files = [] | |
| self._committed_files_to_delete_on_abort = [] | |
| self._pending_normal_meta = None | |
| self.committed_changelog_files = [] | |
| self.file_io = _RecordingFileIO() | |
| self.written = [] | |
| self._fail_normal_times = fail_normal_times | |
| def _write_normal_data_to_file(self, data: pa.Table): | |
| if self._fail_normal_times > 0: | |
| self._fail_normal_times -= 1 | |
| raise IOError('transient storage failure') | |
| self.written.append(data) | |
| return _file_meta(data.num_rows, 'data-%d' % len(self.written)) | |
| class _DedicatedHarness(DedicatedFormatWriter): | |
| def __init__(self, blob_writers, vector_writer=None): | |
| self.target_file_size = _NO_LIMIT | |
| self.target_file_row_num = _NO_LIMIT | |
| self.record_count = 0 | |
| self.blob_writers = blob_writers | |
| self.blob_file_column_names = list(blob_writers) | |
| self.vector_writer = vector_writer | |
| self._normal_buffer = WriteBuffer(self._merge_normal_data) | |
| self._buffer = WriteBuffer(self._merge_normal_data) | |
| self.committed_files = [] | |
| self.committed_changelog_files = [] | |
| self._committed_files_to_delete_on_abort = [] | |
| self._pending_normal_meta = None | |
| self.file_io = _RecordingFileIO() | |
| self.written = [] | |
| self._video_group_policy = None | |
| def _write_normal_data_to_file(self, data: pa.Table): | |
| self.written.append(data) | |
| return _file_meta(data.num_rows, 'data-%d' % len(self.written)) | |
| def test_failed_sidecar_publishes_nothing_and_the_retry_resumes(self): | |
| vector = _StubSidecarWriter(3, 'vector-0', fail_times=1) | |
| writer = self._VectorHarness(vector) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| with self.assertRaises(IOError): | |
| writer._close_current_writers() | |
| # The normal file landed, so the rows are no longer buffered -- but | |
| # nothing is committed and the file is remembered for the retry. | |
| self.assertEqual([t.num_rows for t in writer.written], [3]) | |
| self.assertEqual(writer.committed_files, []) | |
| self.assertEqual(writer._normal_buffer.num_rows, 0) | |
| self.assertIsNotNone(writer._pending_normal_meta) | |
| writer._close_current_writers() | |
| # Still one normal file: the retry resumed at the sidecar phase rather | |
| # than writing the same 3 rows again. | |
| self.assertEqual([t.num_rows for t in writer.written], [3]) | |
| self.assertEqual([m.file_name for m in writer.committed_files], | |
| ['data-1', 'vector-0']) | |
| self.assertIsNone(writer._pending_normal_meta) | |
| # Harvested once the flush completed, and cleared only then. | |
| self.assertEqual(vector.committed_files, []) | |
| def test_successful_flush_publishes_normal_then_sidecars(self): | |
| vector = _StubSidecarWriter(3, 'vector-0') | |
| writer = self._VectorHarness(vector) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| writer._close_current_writers() | |
| self.assertEqual([m.file_name for m in writer.committed_files], | |
| ['data-1', 'vector-0']) | |
| self.assertEqual(vector.committed_files, []) | |
| self.assertIsNone(writer._pending_normal_meta) | |
| self.assertEqual(writer.record_count, 0) | |
| def test_failed_normal_write_keeps_the_rows_and_publishes_nothing(self): | |
| # The other half of the same rule: the sidecars are never reached, so | |
| # the rows have to stay where a retry can find them. | |
| vector = _StubSidecarWriter(3, 'vector-0') | |
| writer = self._VectorHarness(vector, fail_normal_times=1) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| with self.assertRaises(IOError): | |
| writer._close_current_writers() | |
| self.assertEqual(writer._normal_buffer.num_rows, 3) | |
| self.assertEqual(writer.committed_files, []) | |
| self.assertIsNone(writer._pending_normal_meta) | |
| self.assertEqual(vector.prepare_commit_calls, 0) | |
| writer._close_current_writers() | |
| self.assertEqual([t.num_rows for t in writer.written], [3]) | |
| self.assertEqual([m.file_name for m in writer.committed_files], | |
| ['data-1', 'vector-0']) | |
| def test_write_is_rejected_while_a_flush_is_unfinished(self): | |
| # Rows appended between a failed flush and its retry would belong to no | |
| # file: the resumed flush skips the normal write, while the sidecar | |
| # writer would drain them -- breaking the row-count check. | |
| vector = _StubSidecarWriter(3, 'vector-0', fail_times=1) | |
| writer = self._VectorHarness(vector) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| with self.assertRaises(IOError): | |
| writer._close_current_writers() | |
| with self.assertRaises(RuntimeError) as caught: | |
| writer.write(_batch(3, 1)) | |
| self.assertIn('Cannot write', str(caught.exception)) | |
| # Rejecting does not abort, so the flush is still resumable. | |
| self.assertIsNotNone(writer._pending_normal_meta) | |
| writer._close_current_writers() | |
| self.assertEqual([t.num_rows for t in writer.written], [3]) | |
| def test_abort_deletes_the_unpublished_normal_file(self): | |
| # It is in no committed list, so abort has to know about it separately | |
| # or it leaks a data file no snapshot references. | |
| vector = _StubSidecarWriter(3, 'vector-0', fail_times=1) | |
| writer = self._VectorHarness(vector) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| with self.assertRaises(IOError): | |
| writer._close_current_writers() | |
| writer.abort() | |
| self.assertEqual(writer.file_io.deleted, ['/warehouse/data-1']) | |
| self.assertIsNone(writer._pending_normal_meta) | |
| self.assertTrue(vector.aborted) | |
| def test_abort_deletes_external_unpublished_file_and_sidecars(self): | |
| for kind in ('vector', 'blob'): | |
| with self.subTest(kind=kind): | |
| sidecar = _StubSidecarWriter(3, 'sidecar-0', fail_times=1) | |
| writer = (self._VectorHarness(sidecar) if kind == 'vector' | |
| else self._DedicatedHarness({'payload': sidecar})) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| with self.assertRaises(IOError): | |
| writer._close_current_writers() | |
| meta = writer._pending_normal_meta | |
| external_path = 'file:/external/pt=a%2Fb%25/data-1' | |
| meta.external_path = external_path | |
| meta.extra_files = ['data-1.index'] | |
| writer.abort() | |
| self.assertEqual(writer.file_io.deleted, [ | |
| '/external/pt=a%2Fb%25/data-1', | |
| '/external/pt=a%2Fb%25/data-1.index', | |
| ]) | |
| self.assertEqual(meta.external_path, external_path) | |
| self.assertEqual(writer.committed_files, []) | |
| self.assertIsNone(writer._pending_normal_meta) | |
| self.assertTrue(sidecar.aborted) | |
| def test_dedicated_writer_failed_blob_phase_publishes_nothing(self): | |
| blob = _StubSidecarWriter(3, 'blob-0', fail_times=1) | |
| writer = self._DedicatedHarness({'payload': blob}) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| with self.assertRaises(IOError): | |
| writer._close_current_writers() | |
| self.assertEqual([t.num_rows for t in writer.written], [3]) | |
| self.assertEqual(writer.committed_files, []) | |
| # Already tracked for abort, since no committed list holds it yet. | |
| self.assertEqual( | |
| [m.file_name for m in writer._committed_files_to_delete_on_abort], | |
| ['data-1']) | |
| writer._close_current_writers() | |
| self.assertEqual([t.num_rows for t in writer.written], [3]) | |
| self.assertEqual([m.file_name for m in writer.committed_files], | |
| ['data-1', 'blob-0']) | |
| self.assertEqual(blob.committed_files, []) | |
| self.assertIsNone(writer._pending_normal_meta) | |
| def test_dedicated_writer_keeps_the_documented_meta_order(self): | |
| blob = _StubSidecarWriter(3, 'blob-0') | |
| vector = _StubSidecarWriter(3, 'vector-0', fail_times=1) | |
| writer = self._DedicatedHarness({'payload': blob}, vector) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| with self.assertRaises(IOError): | |
| writer._close_current_writers() | |
| self.assertEqual(writer.committed_files, []) | |
| self.assertEqual([m.file_name for m in blob.committed_files], ['blob-0']) | |
| writer._close_current_writers() | |
| self.assertEqual([m.file_name for m in writer.committed_files], | |
| ['data-1', 'blob-0', 'vector-0']) | |
| self.assertEqual( | |
| [m.file_name for m in writer._committed_files_to_delete_on_abort], | |
| ['data-1', 'blob-0', 'vector-0']) | |
| def test_dedicated_writer_respects_the_blob_delete_policy(self): | |
| # Externally managed blob files are not the writer's to delete, so they | |
| # must stay out of the abort list even now that it is filled at publish | |
| # time rather than as each sidecar lands. | |
| blob = _StubSidecarWriter(3, 'blob-0', delete_on_abort=False) | |
| writer = self._DedicatedHarness({'payload': blob}) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| writer._close_current_writers() | |
| self.assertEqual([m.file_name for m in writer.committed_files], | |
| ['data-1', 'blob-0']) | |
| self.assertEqual( | |
| [m.file_name for m in writer._committed_files_to_delete_on_abort], | |
| ['data-1']) | |
| class VectorNormalBufferTest(unittest.TestCase): | |
| """``DataVectorWriter`` keeps its own buffer for the normal columns.""" | |
| class _VectorHarness(DataVectorWriter): | |
| """Only the fields ``_should_roll_normal`` reads. | |
| ``CHECK_ROLLING_RECORD_CNT`` is 1 rather than the real 1000 so the size | |
| branch is reached on every write instead of being short-circuited by the | |
| periodic check. | |
| """ | |
| CHECK_ROLLING_RECORD_CNT = 1 | |
| def __init__(self, target_file_size: int = _NO_LIMIT): | |
| self.target_file_size = target_file_size | |
| self.target_file_row_num = _NO_LIMIT | |
| self.record_count = 0 | |
| self.vector_writer = None | |
| self._normal_buffer = WriteBuffer(self._merge_data) | |
| self._buffer = WriteBuffer(self._merge_data) | |
| def test_should_roll_normal_does_not_fold(self): | |
| writer = self._VectorHarness() | |
| with _count_folds() as folds: | |
| for i in range(50): | |
| writer._normal_buffer.append(_table(i * 5, 5)) | |
| writer.record_count += 5 | |
| self.assertFalse(writer._should_roll_normal()) | |
| self.assertEqual(folds, []) | |
| self.assertEqual(writer._normal_buffer.num_rows, 250) | |
| def test_should_roll_normal_fires_off_the_running_size(self): | |
| writer = self._VectorHarness(target_file_size=_table(0, 20).nbytes) | |
| writer.record_count = 1 | |
| for i in range(10): | |
| writer._normal_buffer.append(_table(i * 5, 5)) | |
| if writer._should_roll_normal(): | |
| break | |
| self.assertTrue(writer._should_roll_normal()) | |
| # Fires before the buffer grows far past the target, i.e. off the | |
| # accumulated size rather than at some arbitrary later point. | |
| self.assertLessEqual(writer._normal_buffer.num_rows, 30) | |
| def test_pending_row_count_reports_the_normal_buffer(self): | |
| # The inherited property reads the base ``_buffer``, which this writer | |
| # never fills, so it has to be overridden or it always answers zero. | |
| writer = self._VectorHarness() | |
| self.assertEqual(writer.pending_row_count, 0) | |
| writer._normal_buffer.append(_table(0, 7)) | |
| self.assertEqual(writer.pending_row_count, 7) | |
| self.assertTrue(writer._buffer.is_empty) | |
| def test_pending_row_count_falls_back_to_the_vector_writer(self): | |
| # A table whose columns are all vectors buffers nothing normal, so the | |
| # count has to come from the sidecar instead of reading as zero. | |
| writer = self._VectorHarness() | |
| writer.vector_writer = _StubSidecarWriter(0, 'vector-0') | |
| writer.vector_writer.pending_row_count = 4 | |
| self.assertEqual(writer.pending_row_count, 4) | |
| writer._normal_buffer.append(_table(0, 7)) | |
| self.assertEqual(writer.pending_row_count, 7) | |
| def test_dedicated_writer_pending_row_count_reports_the_normal_buffer(self): | |
| writer = CompositeFlushResumeTest._DedicatedHarness({}) | |
| self.assertEqual(writer.pending_row_count, 0) | |
| writer._normal_buffer.append(_table(0, 3)) | |
| self.assertEqual(writer.pending_row_count, 3) | |
| self.assertTrue(writer._buffer.is_empty) | |
| if __name__ == '__main__': | |
| unittest.main() |