paimon-python/pypaimon/tests/test_write_merge_buffer.py

################################################################################
#  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.
################################################################################

"""Unit tests for ``KeyValueDataWriter`` buffer behaviour.

Covers the fold algorithm (`_merge_pending_by_pk`), the flush lifecycle
(`_flush_all` drains the buffer), and the roll-write helper
(`_roll_write` splits oversized buffers across multiple files).
Drives a thin harness that bypasses
``DataWriter.__init__`` so tests can exercise these paths without
spinning up the real catalog/write stack.
"""

import unittest
from unittest.mock import Mock

import pyarrow as pa

from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.options.options import Options
from pypaimon.read.reader.deduplicate_merge_function import \
    DeduplicateMergeFunction
from pypaimon.read.reader.partial_update_merge_function import \
    PartialUpdateMergeFunction
from pypaimon.write.writer.key_value_data_writer import KeyValueDataWriter
from pypaimon.write.writer.write_buffer import WriteBuffer


# Layout matches what ``KeyValueDataWriter._add_system_fields`` emits:
# ``[_KEY_id, _SEQUENCE_NUMBER, _VALUE_KIND, id, a, b]``. ``id`` is
# duplicated on the value side because the value layout in Paimon's
# row tuple includes every original column.
_SCHEMA = pa.schema([
    pa.field('_KEY_id', pa.int64(), nullable=False),
    pa.field('_SEQUENCE_NUMBER', pa.int64(), nullable=False),
    pa.field('_VALUE_KIND', pa.int8(), nullable=False),
    pa.field('id', pa.int64(), nullable=False),
    pa.field('a', pa.string()),
    pa.field('b', pa.string()),
])

_NULLABLE_PK_SCHEMA = pa.schema([
    pa.field('_KEY_id', pa.int64()),
    pa.field('_SEQUENCE_NUMBER', pa.int64(), nullable=False),
    pa.field('_VALUE_KIND', pa.int8(), nullable=False),
    pa.field('id', pa.int64()),
    pa.field('a', pa.string()),
    pa.field('b', pa.string()),
])


def _row(pk, seq, a, b):
    return {
        '_KEY_id': pk,
        '_SEQUENCE_NUMBER': seq,
        '_VALUE_KIND': 0,
        'id': pk,
        'a': a,
        'b': b,
    }


class _Harness(KeyValueDataWriter):
    """Bypass ``DataWriter.__init__`` to keep tests focused.

    Provides just the attributes ``_merge_pending_by_pk`` / ``_flush_all``
    / ``_roll_write`` read, plus a recording stub for
    ``_write_data_to_file`` so the roll-write path can be exercised
    without touching the filesystem.
    """

    def __init__(self, merge_function, target_file_size: int = 10 ** 12):
        self.options = CoreOptions(Options({}))
        self.trimmed_primary_keys = ['id']
        self._merge_function = merge_function
        # Large enough that ``_check_and_roll_if_needed`` does not
        # trigger on its own in tests that don't care about rolling.
        self.target_file_size = target_file_size
        self._buffer = WriteBuffer(self._merge_data)
        self.committed_files = []
        self.written_chunks = []

    def _write_data_to_file(self, data):
        # Record each chunk instead of writing to disk; mirrors the
        # base writer's contract of appending to ``committed_files``.
        self.written_chunks.append(data)


class WriteMergeBufferTest(unittest.TestCase):

    # -- deduplicate ------------------------------------------------------

    def test_dedupe_collapses_same_pk_run_to_latest(self):
        writer = _Harness(DeduplicateMergeFunction())
        data = pa.Table.from_pylist(
            [_row(1, 1, 'old', None), _row(1, 2, 'new', None)],
            schema=_SCHEMA,
        )
        out = writer._merge_pending_by_pk(data)
        self.assertEqual(out.num_rows, 1)
        self.assertEqual(
            out.to_pylist(),
            [_row(1, 2, 'new', None)],
        )

    def test_dedupe_keeps_disjoint_keys(self):
        writer = _Harness(DeduplicateMergeFunction())
        data = pa.Table.from_pylist(
            [_row(1, 1, 'A', None),
             _row(2, 2, 'B', None),
             _row(3, 3, 'C', None)],
            schema=_SCHEMA,
        )
        out = writer._merge_pending_by_pk(data)
        self.assertEqual(out.num_rows, 3)
        self.assertEqual(
            sorted(out.to_pylist(), key=lambda r: r['id']),
            [_row(1, 1, 'A', None),
             _row(2, 2, 'B', None),
             _row(3, 3, 'C', None)],
        )

    def test_dedupe_nullable_pk_uses_null_safe_equality_and_nulls_first(self):
        writer = _Harness(DeduplicateMergeFunction())
        writer._buffer.append(pa.Table.from_pylist(
            [_row(2, 1, 'two', None),
             _row(None, 2, 'null-old', None),
             _row(1, 3, 'one', None),
             _row(None, 4, 'null-new', None)],
            schema=_NULLABLE_PK_SCHEMA,
        ))

        writer._flush_all()

        self.assertEqual(len(writer.written_chunks), 1)
        self.assertEqual(
            writer.written_chunks[0].to_pylist(),
            [_row(None, 4, 'null-new', None),
             _row(1, 3, 'one', None),
             _row(2, 1, 'two', None)],
        )

    # -- partial-update ---------------------------------------------------

    def _partial_update(self):
        # Value-side carries 3 columns (id, a, b). The PK column ``id``
        # is duplicated into the value side so partial-update can apply
        # last-non-null semantics uniformly across every original
        # user column.
        return PartialUpdateMergeFunction(
            key_arity=1, value_arity=3, nullables=[True, True, True])

    def test_partial_update_merges_non_null_per_field(self):
        writer = _Harness(self._partial_update())
        data = pa.Table.from_pylist(
            [_row(1, 1, 'A', None), _row(1, 2, None, 'B')],
            schema=_SCHEMA,
        )
        out = writer._merge_pending_by_pk(data)
        self.assertEqual(out.num_rows, 1)
        self.assertEqual(out.to_pylist(), [_row(1, 2, 'A', 'B')])

    def test_partial_update_three_writes_compose(self):
        writer = _Harness(self._partial_update())
        data = pa.Table.from_pylist(
            [_row(1, 1, 'A', None),
             _row(1, 2, None, 'B'),
             _row(1, 3, None, None)],
            schema=_SCHEMA,
        )
        out = writer._merge_pending_by_pk(data)
        self.assertEqual(out.to_pylist(), [_row(1, 3, 'A', 'B')])

    def test_partial_update_later_null_does_not_clobber_earlier_value(self):
        writer = _Harness(self._partial_update())
        data = pa.Table.from_pylist(
            [_row(1, 1, 'KEEP', 'B'), _row(1, 2, None, None)],
            schema=_SCHEMA,
        )
        out = writer._merge_pending_by_pk(data)
        self.assertEqual(out.to_pylist(), [_row(1, 2, 'KEEP', 'B')])

    # -- edge cases -------------------------------------------------------

    def test_empty_buffer_returns_empty(self):
        writer = _Harness(DeduplicateMergeFunction())
        empty = pa.Table.from_pylist([], schema=_SCHEMA)
        out = writer._merge_pending_by_pk(empty)
        self.assertEqual(out.num_rows, 0)

    def test_single_row_buffer_skips_merge(self):
        # Mock to confirm the merge function isn't invoked: a single
        # row cannot have duplicates, so we sidestep the to_pylist
        # round-trip.
        mock_mf = Mock()
        writer = _Harness(mock_mf)
        data = pa.Table.from_pylist(
            [_row(1, 1, 'X', None)], schema=_SCHEMA)
        out = writer._merge_pending_by_pk(data)
        self.assertEqual(out.num_rows, 1)
        mock_mf.reset.assert_not_called()
        mock_mf.add.assert_not_called()
        mock_mf.get_result.assert_not_called()

    def test_get_result_none_drops_pk_run(self):
        # Future-proof: contract says ``get_result`` returning ``None``
        # means the entire PK group should be dropped.
        class DropAll:
            def reset(self):
                pass

            def add(self, _):
                pass

            def get_result(self):
                return None

        writer = _Harness(DropAll())
        data = pa.Table.from_pylist(
            [_row(1, 1, 'A', None), _row(1, 2, 'B', None)],
            schema=_SCHEMA,
        )
        out = writer._merge_pending_by_pk(data)
        self.assertEqual(out.num_rows, 0)

    # -- KeyValue pooling -------------------------------------------------

    def test_keyvalue_pool_does_not_alias_results_across_runs(self):
        # Pooling reuses one KeyValue across the whole fold. If the
        # PartialUpdateMergeFunction's get_result snapshotting were
        # broken, run 1's result would mutate when run 2's data is
        # written into the pooled instance. This test would catch that
        # regression: build a buffer with two distinct PK runs and
        # verify both results stand on their own.
        writer = _Harness(self._partial_update())
        data = pa.Table.from_pylist(
            [_row(1, 1, 'A1', None),
             _row(1, 2, None, 'B1'),
             _row(2, 3, 'A2', None),
             _row(2, 4, None, 'B2')],
            schema=_SCHEMA,
        )
        out = writer._merge_pending_by_pk(data)
        self.assertEqual(
            sorted(out.to_pylist(), key=lambda r: r['id']),
            [_row(1, 2, 'A1', 'B1'), _row(2, 4, 'A2', 'B2')],
        )

    # -- _process_data (no longer sorts) ----------------------------------

    def test_process_data_adds_system_fields_without_sorting(self):
        # Deferred-sort design: ``_process_data`` must not pre-sort the
        # incoming batch. The global sort happens once inside
        # ``_flush_all`` over the concatenated buffer.
        writer = _Harness(DeduplicateMergeFunction())
        writer.sequence_generator = _StubSeqGen()
        # Intentionally pass rows in descending PK order; if _process_data
        # were still sorting, the output would come back ascending.
        batch = pa.RecordBatch.from_pylist(
            [{'id': 3, 'a': 'C', 'b': None},
             {'id': 1, 'a': 'A', 'b': None},
             {'id': 2, 'a': 'B', 'b': None}],
            schema=pa.schema([
                pa.field('id', pa.int64(), nullable=False),
                pa.field('a', pa.string()),
                pa.field('b', pa.string()),
            ]),
        )
        out = writer._process_data(batch)
        self.assertEqual(
            [r['id'] for r in out.to_pylist()],
            [3, 1, 2],
        )

    # -- _flush_all -------------------------------------------------------

    def test_flush_all_sorts_folds_and_writes_one_file(self):
        # Buffer with duplicate PKs in arbitrary order. _flush_all is
        # responsible for sorting before folding, so unsorted input is
        # the right stress case.
        writer = _Harness(DeduplicateMergeFunction())
        writer._buffer.append(pa.Table.from_pylist(
            [_row(2, 5, 'B2-new', None),
             _row(1, 2, 'A1-mid', None),
             _row(1, 1, 'A1-old', None),
             _row(2, 4, 'B2-old', None),
             _row(1, 3, 'A1-new', None)],
            schema=_SCHEMA,
        ))
        writer._flush_all()

        # Buffer cleared.
        self.assertTrue(writer._buffer.is_empty)
        # Exactly one file written (size well under target).
        self.assertEqual(len(writer.written_chunks), 1)
        flushed = writer.written_chunks[0]
        result = sorted(flushed.to_pylist(), key=lambda r: r['id'])
        # Dedup -> 1 row per PK, with the highest seq value retained.
        self.assertEqual(result, [
            _row(1, 3, 'A1-new', None),
            _row(2, 5, 'B2-new', None),
        ])

    def test_flush_all_on_empty_buffer_is_noop(self):
        writer = _Harness(DeduplicateMergeFunction())
        writer._flush_all()
        self.assertTrue(writer._buffer.is_empty)
        self.assertEqual(writer.written_chunks, [])

    def test_flush_all_clears_buffer_even_when_fold_drops_everything(self):
        # MergeFunction that returns None for every group; verifies
        # ``_flush_all`` still drains the buffer so a subsequent write
        # starts from a clean slate.
        class DropAll:
            def reset(self):
                pass

            def add(self, _):
                pass

            def get_result(self):
                return None

        writer = _Harness(DropAll())
        writer._buffer.append(pa.Table.from_pylist(
            [_row(1, 1, 'A', None), _row(1, 2, 'B', None)],
            schema=_SCHEMA,
        ))
        writer._flush_all()
        self.assertTrue(writer._buffer.is_empty)
        self.assertEqual(writer.written_chunks, [])

    # -- _roll_write ------------------------------------------------------

    def test_roll_write_single_chunk_when_under_target(self):
        writer = _Harness(DeduplicateMergeFunction(),
                          target_file_size=10 ** 9)
        data = pa.Table.from_pylist(
            [_row(1, 1, 'A', None), _row(2, 2, 'B', None)],
            schema=_SCHEMA,
        )
        writer._roll_write(data)
        self.assertEqual(len(writer.written_chunks), 1)
        self.assertEqual(writer.written_chunks[0].num_rows, 2)

    def test_roll_write_splits_oversized_buffer_into_multiple_files(self):
        # Build a buffer whose nbytes comfortably exceeds the chosen
        # target. With a small target_file_size the writer should hand
        # back at least two files. Use long strings so nbytes scales
        # predictably with row count.
        rows = [
            _row(i, i, 'x' * 64, 'y' * 64) for i in range(1, 401)
        ]
        data = pa.Table.from_pylist(rows, schema=_SCHEMA)
        # Target small enough that 400 rows will not fit in one file.
        target = data.nbytes // 4
        writer = _Harness(DeduplicateMergeFunction(),
                          target_file_size=target)
        writer._roll_write(data)

        self.assertGreaterEqual(len(writer.written_chunks), 2)
        total_rows = sum(c.num_rows for c in writer.written_chunks)
        self.assertEqual(total_rows, data.num_rows)
        # Each chunk except possibly the last should respect the target.
        for chunk in writer.written_chunks[:-1]:
            self.assertLessEqual(chunk.nbytes, target)

    def test_roll_write_leaves_only_unwritten_rows_when_a_file_fails(self):
        # A flush spans several files, so it cannot be all-or-nothing: keeping
        # every row would make the retry rewrite what the first files already
        # took. The buffer has to hold the remainder and nothing else.
        rows = [_row(i, i, 'x' * 64, 'y' * 64) for i in range(1, 401)]
        data = pa.Table.from_pylist(rows, schema=_SCHEMA)

        class _FailOnSecondFile(_Harness):
            """Fails once, on the second file, then writes normally."""

            failed = False

            def _write_data_to_file(self, chunk):
                if not self.failed and len(self.written_chunks) == 1:
                    self.failed = True
                    raise IOError('transient storage failure')
                super()._write_data_to_file(chunk)

        writer = _FailOnSecondFile(DeduplicateMergeFunction(),
                                   target_file_size=data.nbytes // 4)
        writer._buffer.append(data)
        with self.assertRaises(IOError):
            writer._flush_all()

        written = sum(c.num_rows for c in writer.written_chunks)
        self.assertEqual(len(writer.written_chunks), 1)
        self.assertEqual(written + writer._buffer.num_rows, data.num_rows)

        # The retry picks up exactly where the failure left off.
        writer._flush_all()
        self.assertTrue(writer._buffer.is_empty)
        self.assertEqual(sum(c.num_rows for c in writer.written_chunks),
                         data.num_rows)
        flushed_ids = [r['id'] for c in writer.written_chunks
                       for r in c.to_pylist()]
        self.assertEqual(sorted(flushed_ids), [r['id'] for r in rows])


class KeyValueWriteBufferTest(unittest.TestCase):
    """The write path must not fold the buffer on every write.

    ``_check_and_roll_if_needed`` flushes on size alone, and it reads that size
    from the buffer's running total, so a write that stays under
    ``target_file_size`` never needs the buffered batches as one table. Folding
    there anyway would make writing N batches O(N^2): each fold leaves one more
    chunk per column for the next ``pa.Table.nbytes`` walk to visit.
    """

    def test_writes_under_target_size_do_not_fold(self):
        writer = _Harness(DeduplicateMergeFunction())
        writer.sequence_generator = _StubSeqGen()
        batch = pa.RecordBatch.from_pylist(
            [{'id': 1, 'a': 'A', 'b': None}],
            schema=pa.schema([
                pa.field('id', pa.int64(), nullable=False),
                pa.field('a', pa.string()),
                pa.field('b', pa.string()),
            ]),
        )

        folds = []
        original = pa.concat_tables

        def counting(*args, **kwargs):
            folds.append(1)
            return original(*args, **kwargs)

        pa.concat_tables = counting
        try:
            for _ in range(100):
                writer.write(batch)
        finally:
            pa.concat_tables = original

        self.assertEqual(folds, [])
        self.assertEqual(writer.written_chunks, [])
        self.assertEqual(writer.pending_row_count, 100)

        # The rows are all still there; flushing folds them exactly once.
        writer._flush_all()
        self.assertEqual(len(writer.written_chunks), 1)
        # Dedup collapses the 100 same-PK rows to one.
        self.assertEqual(writer.written_chunks[0].num_rows, 1)


class _StubSeqGen:
    """Stand-in for ``SequenceGenerator`` so the harness can call
    ``_process_data`` without going through the real ``DataWriter.__init__``.
    """

    def __init__(self):
        self._n = 0

    def next(self) -> int:
        self._n += 1
        return self._n


if __name__ == '__main__':
    unittest.main()