This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 80c5ab9593 [python] Align PK input changelogs and exercise Native
batch sharding (#10451)
80c5ab9593 is described below
commit 80c5ab959347bf46bb949d79d908cfdabffc8163
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Oct 8 23:22:09 2026 +0800
[python] Align PK input changelogs and exercise Native batch sharding
(#10451)
---
.../pypaimon/tests/native_changelog_batch_test.py | 285 +++++++++++++++++++++
.../tests/native_cross_partition_write_test.py | 40 +++
.../tests/test_read_builder_nested_projection.py | 35 ++-
.../pypaimon/tests/test_write_merge_buffer.py | 2 +
.../pypaimon/tests/write/row_kind_write_test.py | 193 ++++++++++++++
.../pypaimon/tests/write/table_write_test.py | 11 +-
paimon-python/pypaimon/write/file_store_write.py | 2 +
paimon-python/pypaimon/write/native_write.py | 7 +-
paimon-python/pypaimon/write/row_kind.py | 21 +-
paimon-python/pypaimon/write/table_write.py | 19 +-
paimon-python/pypaimon/write/writer/data_writer.py | 124 ++++-----
.../pypaimon/write/writer/key_value_data_writer.py | 51 +++-
12 files changed, 713 insertions(+), 77 deletions(-)
diff --git a/paimon-python/pypaimon/tests/native_changelog_batch_test.py
b/paimon-python/pypaimon/tests/native_changelog_batch_test.py
new file mode 100644
index 0000000000..87b829cd06
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_changelog_batch_test.py
@@ -0,0 +1,285 @@
+# 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.
+
+"""Java batch changelog packing and sharding through the Native planner."""
+
+import json
+from contextlib import ExitStack
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.read.scan_distribution import java_file_name_shard
+from pypaimon.schema.data_types import AtomicType, PyarrowFieldParser
+from pypaimon.schema.schema_change import SchemaChange
+from pypaimon.table.row.generic_row import GenericRow
+from pypaimon.table.row.internal_row import RowKind
+
+
+pytestmark = pytest.mark.native_plan
+SCHEMA = pa.schema([('key', pa.int64()), ('value', pa.int64())])
+
+
[email protected](params=[False, True], ids=['python', 'native'])
+def native(request):
+ return request.param
+
+
[email protected]
+def catalog(tmp_path):
+ result = CatalogFactory.create({'warehouse': str(tmp_path)})
+ result.create_database('default', True)
+ return result
+
+
+def _table(catalog, options=None, schema=SCHEMA, partitions=None):
+ config = {
+ 'bucket': '4', 'changelog-producer': 'input',
+ 'write.native.enabled': 'false',
+ 'source.split.target-size': '1mb', 'source.split.open-file-cost': '1b',
+ }
+ config.update(options or {})
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(
+ schema, primary_keys=list(partitions or []) + ['key'],
+ partition_keys=partitions, options=config), False)
+ return catalog.get_table('default.t')
+
+
+def _patch_snapshot(table, snapshot_id, **changes):
+ path = table.snapshot_manager().get_snapshot_path(snapshot_id)
+ data = json.loads(table.file_io.read_file_utf8(path))
+ data.update(changes)
+ table.file_io.write_file(path, json.dumps(data), overwrite=True)
+
+
+def _write(table, timestamp, rows, kind=None, overwrite=False):
+ builder = table.new_batch_write_builder()
+ if overwrite:
+ builder.overwrite()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ if kind is None:
+ # Use the declared Arrow types even for NULL nested values.
+ schema = PyarrowFieldParser.from_paimon_schema(table.fields)
+ writer.write_arrow(pa.Table.from_pylist(rows, schema=schema))
+ else:
+ for row in rows:
+ writer.write_row(GenericRow(
+ [row[field.name] for field in table.fields], table.fields,
kind))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ snapshot_id = table.snapshot_manager().get_latest_snapshot().id
+ _patch_snapshot(table, snapshot_id, timeMillis=timestamp)
+ return snapshot_id
+
+
+def _read(table, native, window=(0, 300), shard=None, predicate=None,
+ projection=None, limit=None, with_stats=False):
+ builder = table.copy({
+ 'scan.native-plan.enabled': str(native).lower(),
+ 'read.native.enabled': str(native).lower(),
+ 'scan.mode': 'incremental',
+ 'incremental-between-timestamp': '%s,%s' % window,
+ }).new_read_builder()
+ if predicate is not None:
+ builder.with_filter(predicate)
+ if projection is not None:
+ builder.with_projection(projection)
+ if limit is not None:
+ builder.with_limit(limit)
+ scan = builder.new_scan()
+ if shard is not None:
+ scan.with_shard(*shard)
+ with ExitStack() as stack:
+ if native:
+ for method in ('scan', 'scan_with_stats'):
+ stack.enter_context(patch.object(
+ scan.file_scanner, method,
+ side_effect=AssertionError('batch changelog planning used
Python')))
+ plan = scan.scan_with_stats()[0] if with_stats else scan.plan()
+ read = builder.new_read()
+ read.include_row_kind = True
+ with ExitStack() as stack:
+ if native:
+ assert all(getattr(split, '_native_split', None) is not None for
split in plan.splits())
+ stack.enter_context(patch.object(
+ read, '_create_split_read',
+ side_effect=AssertionError('batch changelog reading used
Python')))
+ rows = read.to_arrow(plan.splits(), parallelism=1).to_pylist()
+ return plan, rows
+
+
[email protected]('bucket', ['1', '4', '-1'])
[email protected]('dv', [False, True])
+def test_changelog_shards_cover_all_events_and_pack_the_whole_window(catalog,
native, bucket, dv):
+ table = _table(catalog, {'bucket': bucket, 'deletion-vectors.enabled':
str(dv).lower()})
+ expected = []
+ for version in range(3):
+ rows = [dict(key=key, value=version) for key in range(32)]
+ _write(table, (version + 1) * 100, rows)
+ expected.extend(('+I', row['key'], row['value']) for row in rows)
+ full, rows = _read(table, native)
+ assert sorted((row['_row_kind'], row['key'], row['value']) for row in
rows) == sorted(expected)
+ # Files from multiple commits in the same bucket must share batch splits.
+ assert len(full.splits()) == len({split.bucket for split in full.splits()})
+ full_files = {file.file_name for split in full.splits() for file in
split.files}
+ for count in (2, 5):
+ seen_files, union = set(), []
+ for index in range(count):
+ plan, rows = _read(table, native, shard=(index, count))
+ assert plan.snapshot_id == 3
+ assert all(split.is_streaming and split.snapshot_id == 3 for split
in plan.splits())
+ names = {file.file_name for split in plan.splits() for file in
split.files}
+ assert seen_files.isdisjoint(names)
+ seen_files.update(names)
+ union.extend((row['_row_kind'], row['key'], row['value']) for row
in rows)
+ for split in plan.splits():
+ if dv:
+ assert all(java_file_name_shard(file.file_name, count) ==
index for file in split.files)
+ else:
+ assert split.bucket % count == index
+ _, limited = _read(table, native, shard=(index, count), limit=1)
+ assert len(limited) == min(1, len(rows))
+ assert all(row in rows for row in limited)
+ assert seen_files == full_files
+ assert sorted(union) == sorted(expected)
+
+
[email protected]('native_write', [False, True])
+def test_changelog_keeps_all_row_kinds_and_snapshot_boundaries(catalog,
native, native_write):
+ table = _table(catalog, {'bucket': '1', 'write.native.enabled':
str(native_write).lower()})
+ kinds = [RowKind.INSERT, RowKind.UPDATE_AFTER, RowKind.UPDATE_BEFORE,
RowKind.DELETE]
+ for index, kind in enumerate(kinds):
+ _write(table, (index + 1) * 100, [dict(key=1, value=index)], kind=kind)
+ plan, rows = _read(table, native, window=(100, 400))
+ assert plan.snapshot_id == 4
+ assert len(plan.splits()) == 1
+ assert sorted((row['_row_kind'], row['key'], row['value']) for row in
rows) == [
+ ('+U', 1, 1), ('-D', 1, 3), ('-U', 1, 2)]
+ for window in [(100, 100), (400, 400), (500, 600)]:
+ plan, rows = _read(table, native, window=window, shard=(1, 3))
+ assert rows == []
+ assert plan.splits() == []
+
+
[email protected]('limit', [0, 1, 2, 5])
+def test_changelog_limit_is_applied_to_the_batch_once(catalog, native, limit):
+ table = _table(catalog, {'bucket': '1', 'source.split.target-size': '1b'})
+ for key in range(4):
+ _write(table, (key + 1) * 100, [dict(key=key, value=key * 10)])
+ plan, rows = _read(table, native, window=(0, 400), limit=limit)
+ assert plan.snapshot_id == 4
+ assert len(rows) == min(limit, 4)
+ if limit:
+ assert len(plan.splits()) == min(limit, 4)
+ assert all(row['_row_kind'] == '+I' for row in rows)
+
+
[email protected]('projection', [['value', 'key'], ['value'], []])
+def
test_sharded_changelog_filters_before_limit_and_preserves_projection(catalog,
native, projection):
+ table = _table(catalog, {'source.split.target-size': '1b'})
+ for version in range(3):
+ _write(table, (version + 1) * 100,
+ [dict(key=key, value=version) for key in range(32)])
+ predicate =
table.new_read_builder().new_predicate_builder().greater_than('value', 0)
+ total = 0
+ for index in range(5):
+ kwargs = dict(shard=(index, 5), predicate=predicate,
projection=projection)
+ _, rows = _read(table, native, **kwargs)
+ total += len(rows)
+ _, limited = _read(table, native, limit=1, **kwargs)
+ assert len(limited) == min(1, len(rows))
+ assert all(row in rows for row in limited)
+ assert all(list(row) == ['_row_kind'] + projection for row in limited)
+ if 'value' in projection:
+ assert all(row['value'] > 0 for row in limited)
+ assert total == 64
+
+
+def
test_changelog_skips_overwrite_but_includes_compact_and_keeps_empty_end_metadata(catalog,
native):
+ table = _table(catalog, {'bucket': '1'})
+ for key in range(4):
+ _write(table, (key + 1) * 100, [dict(key=key, value=key)])
+ _patch_snapshot(table, 2, commitKind='COMPACT')
+ _patch_snapshot(table, 3, commitKind='OVERWRITE')
+ _patch_snapshot(table, 4, changelogManifestList=None)
+ plan, rows = _read(table, native, window=(0, 400), with_stats=True)
+ assert plan.snapshot_id == 4
+ assert sorted((row['key'], row['value']) for row in rows) == [(0, 0), (1,
1)]
+ assert all(split.snapshot_id == 4 for split in plan.splits())
+ plan, rows = _read(table, native, window=(200, 400), shard=(0, 2))
+ assert plan.snapshot_id == 4
+ assert plan.splits() == []
+ assert rows == []
+
+
+def
test_changelog_partition_filter_and_nested_projection_remain_native(catalog,
native):
+ schema = SCHEMA.append(pa.field('part', pa.string())).append(pa.field(
+ 'payload', pa.struct([('score', pa.int64()), ('unused',
pa.string())])))
+ table = _table(catalog, schema=schema, partitions=['part'])
+ for version in range(3):
+ _write(table, (version + 1) * 100, [
+ dict(key=key, value=version, part=part,
+ payload=dict(score=version * 10 + key, unused='large'))
+ for part in ['keep', 'drop'] for key in range(8)])
+ predicate = table.new_read_builder().new_predicate_builder().equal('part',
'keep')
+ union = []
+ for index in range(3):
+ plan, rows = _read(table, native, shard=(index, 3),
predicate=predicate,
+ projection=['payload.score'])
+ assert plan.snapshot_id == 3
+ assert all(split.partition.get_field(0) == 'keep' for split in
plan.splits())
+ union.extend(row['payload_score'] for row in rows)
+ assert all(row['_row_kind'] == '+I' for row in rows)
+ assert sorted(union) == [version * 10 + key for version in range(3) for
key in range(8)]
+
+
[email protected]('engine', ['deduplicate', 'partial-update',
'aggregation'])
+def
test_batch_changelog_preserves_inputs_without_running_merge_engine(catalog,
native, engine):
+ options = {'bucket': '1', 'merge-engine': engine}
+ if engine == 'aggregation':
+ options['fields.value.aggregate-function'] = 'sum'
+ table = _table(catalog, options)
+ for version in range(3):
+ _write(table, (version + 1) * 100, [dict(key=1, value=version + 1)])
+ plan, rows = _read(table, native, shard=(0, 2))
+ assert len(plan.splits()) == 1
+ assert sorted((row['_row_kind'], row['key'], row['value']) for row in
rows) == [
+ ('+I', 1, 1), ('+I', 1, 2), ('+I', 1, 3)]
+
+
+def
test_batch_changelog_packs_evolved_schemas_and_filters_missing_fields(catalog,
native):
+ table = _table(catalog, {'bucket': '1'})
+ _write(table, 100, [dict(key=1, value=10)])
+ catalog.alter_table('default.t', [
+ SchemaChange.rename_column('value', 'renamed'),
+ SchemaChange.add_column('added', AtomicType('BIGINT'))], False)
+ table = catalog.get_table('default.t')
+ _write(table, 200, [dict(key=1, renamed=20, added=2)])
+ _write(table, 300, [dict(key=2, renamed=30, added=3)])
+ plan, rows = _read(table, native, projection=['key', 'renamed', 'added'])
+ assert len(plan.splits()) == 1
+ assert sorted((row['key'], row['renamed'], row['added']) for row in rows)
== [
+ (1, 10, None), (1, 20, 2), (2, 30, 3)]
+ predicate =
table.new_read_builder().new_predicate_builder().greater_than('added', 2)
+ _, rows = _read(table, native, shard=(0, 2), predicate=predicate,
+ projection=['renamed', 'added'], limit=1)
+ assert rows == [dict(_row_kind='+I', renamed=30, added=3)]
diff --git a/paimon-python/pypaimon/tests/native_cross_partition_write_test.py
b/paimon-python/pypaimon/tests/native_cross_partition_write_test.py
index edd9d963c7..de66a715d9 100644
--- a/paimon-python/pypaimon/tests/native_cross_partition_write_test.py
+++ b/paimon-python/pypaimon/tests/native_cross_partition_write_test.py
@@ -23,6 +23,7 @@ import pytest
from pypaimon import CatalogFactory, Schema
from pypaimon.table.row.generic_row import GenericRow
+from pypaimon.table.row.row_kind import RowKind
from pypaimon.write.native_write import NativeTableWrite
pytestmark = pytest.mark.native_plan
@@ -153,6 +154,45 @@ def test_cross_partition_rowkind_and_row_input(tmp_path):
assert _rows(table, native) == [('x', 'a', 1, 30)]
+def test_cross_partition_internal_row_kinds_and_migration_deletes(tmp_path):
+ table = _table(tmp_path, {'changelog-producer': 'input'})
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ for values, kind in [
+ (['x', 'a', 1, 10], RowKind.INSERT),
+ (['x', 'b', 1, 20], RowKind.UPDATE_AFTER),
+ (['x', 'b', 1, 30], RowKind.UPDATE_BEFORE),
+ (['x', 'b', 1, 40], RowKind.DELETE),
+ ]:
+ writer.write_row(GenericRow(values, table.fields, kind))
+ assert writer._python_writer is None
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ for native in (False, True):
+ assert _rows(table, native) == []
+ builder = table.copy({
+ 'scan.native-plan.enabled': str(native).lower(),
+ 'read.native.enabled': str(native).lower(),
+ 'scan.mode': 'incremental',
+ 'incremental-between-timestamp': '0,9223372036854775807',
+ }).new_read_builder()
+ plan = builder.new_scan().plan()
+ read = builder.new_read()
+ read.include_row_kind = True
+ if native:
+ with patch.object(read, '_create_split_read',
side_effect=AssertionError('fallback')):
+ rows = read.to_arrow(plan.splits(), parallelism=1).to_pylist()
+ else:
+ rows = read.to_arrow(plan.splits(), parallelism=1).to_pylist()
+ assert sorted((row['_row_kind'], row['p'], row['v']) for row in rows)
== [
+ ('+I', 'a', 10), ('+U', 'b', 20), ('-D', 'a', 20),
+ ('-D', 'b', 40), ('-U', 'b', 30)]
+
+
@pytest.mark.parametrize('committed', [False, True])
def
test_cross_partition_abort_preserves_prepared_and_committed_files(tmp_path,
committed):
table = _table(tmp_path)
diff --git
a/paimon-python/pypaimon/tests/test_read_builder_nested_projection.py
b/paimon-python/pypaimon/tests/test_read_builder_nested_projection.py
index b5ccc8d108..1d7bd673ac 100644
--- a/paimon-python/pypaimon/tests/test_read_builder_nested_projection.py
+++ b/paimon-python/pypaimon/tests/test_read_builder_nested_projection.py
@@ -405,15 +405,40 @@ class
StreamReadBuilderNestedProjectionTest(_ReadBuilderTestBase):
batch._nested_name_paths(),
stream._nested_name_paths(),
)
- self.assertEqual(
- [field.name for field in batch.new_scan()._read_type],
- [field.name for field in stream.new_streaming_scan()._read_type],
- )
-
table_read = stream.with_include_row_kind().new_read()
self.assertEqual(batch._nested_name_paths(),
table_read.nested_name_paths)
self.assertTrue(table_read.include_row_kind)
+ # Exercise the configured scan through its public iterator. Native
+ # owns its read type in Rust and does not expose Python scanner fields.
+ write_builder = self.table.new_batch_write_builder()
+ writer, commit = write_builder.new_write(), write_builder.new_commit()
+ try:
+ writer.write_arrow(pa.Table.from_pydict({
+ 'pk': [1, 2],
+ 'mv': [{'latest_version': 7, 'latest_value': 'unused'}, None],
+ 'val': ['unused', 'unused'],
+ 'attrs': [[('key.with.dots', 42), ('other', 99)], []],
+ }, schema=self.pa_schema))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+
+ expected = {'mv_latest_version': [7, None],
+ 'attrs_key_with_dots': [42, None], 'pk': [1, 2]}
+ batch_plan = batch.new_scan().plan()
+ self.assertEqual(expected, batch.new_read().to_arrow(
+ batch_plan.splits(), parallelism=1).to_pydict())
+ iterator = stream.new_streaming_scan().stream_sync()
+ try:
+ stream_plan = next(iterator)
+ actual = table_read.to_arrow(stream_plan.splits(),
parallelism=1).to_pydict()
+ finally:
+ iterator.close()
+ self.assertEqual(['+I', '+I'], actual.pop('_row_kind'))
+ self.assertEqual(expected, actual)
+
if __name__ == '__main__':
unittest.main()
diff --git a/paimon-python/pypaimon/tests/test_write_merge_buffer.py
b/paimon-python/pypaimon/tests/test_write_merge_buffer.py
index 42428e6388..d24c65cb73 100644
--- a/paimon-python/pypaimon/tests/test_write_merge_buffer.py
+++ b/paimon-python/pypaimon/tests/test_write_merge_buffer.py
@@ -92,7 +92,9 @@ class _Harness(KeyValueDataWriter):
# 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._input_changelog_buffer = None
self.committed_files = []
+ self.committed_changelog_files = []
self.written_chunks = []
def _write_data_to_file(self, data):
diff --git a/paimon-python/pypaimon/tests/write/row_kind_write_test.py
b/paimon-python/pypaimon/tests/write/row_kind_write_test.py
new file mode 100644
index 0000000000..a8ba0fe6cc
--- /dev/null
+++ b/paimon-python/pypaimon/tests/write/row_kind_write_test.py
@@ -0,0 +1,193 @@
+# 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.
+
+
+"""InternalRow kinds survive the Arrow conversion used by both PK writers."""
+
+import pyarrow as pa
+import pytest
+from unittest.mock import patch
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.table.row.generic_row import GenericRow
+from pypaimon.table.row.row_kind import RowKind
+from pypaimon.write.native_write import NativeTableWrite
+
+
[email protected](params=[False, pytest.param(True,
marks=pytest.mark.native_plan)],
+ ids=['python', 'native'])
+def native_write(request):
+ return request.param
+
+
+def _table(tmp_path, native_write, options=None, with_op=False):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ schema = pa.schema([('id', pa.int64()), ('value', pa.int64())])
+ if with_op:
+ schema = schema.append(pa.field('op', pa.string()))
+ config = {'bucket': '1', 'changelog-producer': 'input',
+ 'write.native.enabled': str(native_write).lower(),
+ 'scan.native-plan.enabled': 'false', 'read.native.enabled':
'false'}
+ config.update(options or {})
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(
+ schema, primary_keys=['id'], options=config), False)
+ return catalog.get_table('default.t')
+
+
+def _write(table, rows, native_write, stream=False):
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ if native_write:
+ assert isinstance(writer, NativeTableWrite)
+ try:
+ for values, kind in rows:
+ writer.write_row(GenericRow(values, table.fields, kind))
+ if native_write:
+ assert writer._python_writer is None
+ if stream:
+ commit.commit(writer.prepare_commit(1), 1)
+ else:
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+
+
+def _rows(table, changelog=False):
+ if changelog:
+ table = table.copy({'scan.mode': 'incremental',
+ 'incremental-between-timestamp':
'0,9223372036854775807'})
+ builder = table.new_read_builder()
+ reader = builder.new_read()
+ reader.include_row_kind = changelog
+ result = reader.to_arrow(builder.new_scan().plan().splits(),
parallelism=1).to_pydict()
+ return sorted(zip(*(result[name] for name in result)))
+
+
[email protected]('stream', [False, True])
[email protected]('bucket', ['1', '-1'])
+def test_internal_row_kinds_reach_data_and_input_changelog(tmp_path,
native_write, stream, bucket):
+ table = _table(tmp_path, native_write, {'bucket': bucket})
+ _write(table, [
+ ([1, 10], RowKind.INSERT), ([1, 20], RowKind.UPDATE_AFTER),
+ ([2, 30], RowKind.DELETE), ([3, 40], RowKind.UPDATE_BEFORE),
+ ([4, 50], RowKind.INSERT)], native_write, stream)
+ assert _rows(table) == [(1, 20), (4, 50)]
+ assert _rows(table, True) == [
+ ('+I', 1, 10), ('+I', 4, 50), ('+U', 1, 20),
+ ('-D', 2, 30), ('-U', 3, 40)]
+
+
[email protected]('options,expected', [
+ ({'ignore-delete': 'true'}, [('+I', 1, 10), ('+U', 3, 40)]),
+ ({'ignore-update-before': 'true'}, [('+I', 1, 10), ('+U', 3, 40), ('-D',
1, 20)]),
+])
+def test_internal_row_kinds_are_filtered_before_dynamic_bucket_assignment(
+ tmp_path, native_write, options, expected):
+ options = dict(options, bucket='-1')
+ table = _table(tmp_path, native_write, options)
+ _write(table, [
+ ([1, 10], RowKind.INSERT), ([1, 20], RowKind.DELETE),
+ ([2, 30], RowKind.UPDATE_BEFORE), ([3, 40], RowKind.UPDATE_AFTER)],
native_write)
+ assert _rows(table, True) == expected
+
+
+def test_rowkind_field_overrides_internal_row_kind(tmp_path, native_write):
+ table = _table(tmp_path, native_write, {'rowkind.field': 'op'},
with_op=True)
+ _write(table, [
+ ([1, 10, '+I'], RowKind.DELETE),
+ ([2, 20, '-D'], RowKind.INSERT)], native_write)
+ assert _rows(table) == [(1, 10, '+I')]
+ assert _rows(table, True) == [('+I', 1, 10, '+I'), ('-D', 2, 20, '-D')]
+
+
[email protected]('row_first', [False, True])
[email protected]('producer', ['none', 'input'])
+def test_mixed_arrow_and_row_writes_preserve_kinds(tmp_path, native_write,
row_first, producer):
+ table = _table(tmp_path, native_write, {'changelog-producer': producer})
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ schema = pa.schema([('id', pa.int64()), ('value', pa.int64())])
+ try:
+ if row_first:
+ writer.write_row(GenericRow([1, 10], table.fields, RowKind.INSERT))
+ writer.write_arrow(pa.Table.from_pydict(
+ {'id': [1], 'value': [20]}, schema=schema))
+ else:
+ writer.write_arrow(pa.Table.from_pydict(
+ {'id': [1], 'value': [10]}, schema=schema))
+ writer.write_row(GenericRow([1, 20], table.fields, RowKind.DELETE))
+ if native_write:
+ assert isinstance(writer, NativeTableWrite)
+ assert writer._python_writer is None
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ assert _rows(table) == ([(1, 20)] if row_first else [])
+ if producer == 'input':
+ assert _rows(table, True) == [('+I', 1, 10), ('+I' if row_first else
'-D', 1, 20)]
+
+
[email protected]('failure', ['later_changelog',
'data_after_changelogs'])
+def
test_input_changelog_retry_preserves_published_chunks_and_all_inputs(tmp_path,
failure):
+ table = _table(tmp_path, False)
+ builder = table.new_stream_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(pa.Table.from_pydict(
+ {'id': [1, 1, 1], 'value': [10, 20, 30]},
+ schema=pa.schema([('id', pa.int64()), ('value', pa.int64())])))
+ # Delay rolling until prepare so a retry can preserve already published
+ # log chunks. These rows fold into one data row, but remain three logs.
+ bucket_writer =
next(iter(writer.file_store_write.data_writers.values()))
+ bucket_writer.target_file_size = 1
+ original = table.file_io.write_parquet
+ state = {'changelogs': 0, 'failed': False}
+
+ def fail_once(path, data, **kwargs):
+ if '/changelog-' in str(path):
+ state['changelogs'] += 1
+ should_fail = failure == 'later_changelog' and
state['changelogs'] == 2
+ else:
+ should_fail = failure == 'data_after_changelogs'
+ if should_fail and not state['failed']:
+ state['failed'] = True
+ raise IOError('transient file write failure')
+ return original(path, data, **kwargs)
+
+ with patch.object(table.file_io, 'write_parquet',
side_effect=fail_once):
+ with pytest.raises(IOError, match='transient file write failure'):
+ writer.prepare_commit(1)
+ messages = writer.prepare_commit(1)
+ logs = [meta for message in messages for meta in
message.changelog_files]
+ assert len(logs) == 3
+ assert sum(meta.row_count for meta in logs) == 3
+ assert [(meta.min_sequence_number, meta.max_sequence_number) for meta
in logs] == [(0, 0), (1, 1), (2, 2)]
+ paths = [meta.physical_path() for meta in logs]
+ assert all(table.file_io.exists(path) for path in paths)
+ commit.commit(messages, 1)
+ assert _rows(table) == [(1, 30)]
+ assert _rows(table, True) == [('+I', 1, 10), ('+I', 1, 20), ('+I', 1,
30)]
+ assert writer.prepare_commit(2) == []
+ finally:
+ writer.close()
+ commit.close()
+ # Successful prepare transfers file ownership to the messages. Closing
+ # the writer must retain the logs even after an earlier failed prepare.
+ assert all(table.file_io.exists(path) for path in paths)
diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py
b/paimon-python/pypaimon/tests/write/table_write_test.py
index 5b48d2ef3b..8aa93a53ad 100644
--- a/paimon-python/pypaimon/tests/write/table_write_test.py
+++ b/paimon-python/pypaimon/tests/write/table_write_test.py
@@ -225,9 +225,10 @@ class TableWriteTest(unittest.TestCase):
return out
@staticmethod
- def _mock_table_write(partitions, buckets):
+ def _mock_table_write(partitions, buckets, schema):
table_write = object.__new__(TableWrite)
table_write.table = Mock(is_primary_key_table=False)
+ table_write.table_pyarrow_schema = schema
table_write._validate_pyarrow_schema = Mock()
table_write.row_key_extractor = Mock()
table_write.file_store_write = Mock()
@@ -242,7 +243,7 @@ class TableWriteTest(unittest.TestCase):
'payload': [b'a', b'b'],
})
table_write = self._mock_table_write(
- [('p1',), ('p1',)], [0, 0])
+ [('p1',), ('p1',)], [0, 0], data.schema)
with patch.object(pa.compute, 'take', wraps=pa.compute.take) as take:
table_write.write_arrow_batch(data)
@@ -258,7 +259,7 @@ class TableWriteTest(unittest.TestCase):
})
table_write = self._mock_table_write(
[('p1',), ('p1',), ('p2',), ('p2',)],
- [0, 0, 1, 1])
+ [0, 0, 1, 1], data.schema)
with patch.object(pa.compute, 'take', wraps=pa.compute.take) as take:
table_write.write_arrow_batch(data)
@@ -280,7 +281,7 @@ class TableWriteTest(unittest.TestCase):
})
table_write = self._mock_table_write(
[('p1',), ('p2',), ('p1',), ('p2',)],
- [0, 1, 0, 1])
+ [0, 1, 0, 1], data.schema)
with patch.object(pa.compute, 'take', wraps=pa.compute.take) as take:
table_write.write_arrow_batch(data)
@@ -303,6 +304,7 @@ class TableWriteTest(unittest.TestCase):
})
table_write = object.__new__(TableWrite)
table_write.table = Mock(is_primary_key_table=False)
+ table_write.table_pyarrow_schema = data.schema
table_write._validate_pyarrow_schema = Mock()
table_write.file_store_write = Mock()
table_write.row_key_extractor = Mock()
@@ -328,6 +330,7 @@ class TableWriteTest(unittest.TestCase):
})
table_write = object.__new__(TableWrite)
table_write.table = Mock(is_primary_key_table=False)
+ table_write.table_pyarrow_schema = data.schema
table_write._validate_pyarrow_schema = Mock()
table_write.file_store_write = Mock()
table_write.row_key_extractor = Mock()
diff --git a/paimon-python/pypaimon/write/file_store_write.py
b/paimon-python/pypaimon/write/file_store_write.py
index 16290ff864..a7bc72db8b 100644
--- a/paimon-python/pypaimon/write/file_store_write.py
+++ b/paimon-python/pypaimon/write/file_store_write.py
@@ -121,6 +121,8 @@ class FileStoreWrite:
self.table.table_schema.fields,
column_names,
)
+ from pypaimon.write.row_kind import with_row_kind
+ data = with_row_kind(self.table, data, row)
writer.write(data.to_batches()[0])
def roll_before_group_if_needed(self, row_count: int):
diff --git a/paimon-python/pypaimon/write/native_write.py
b/paimon-python/pypaimon/write/native_write.py
index bdc7a98518..0a8d0f136f 100644
--- a/paimon-python/pypaimon/write/native_write.py
+++ b/paimon-python/pypaimon/write/native_write.py
@@ -206,6 +206,9 @@ class NativeTableWrite:
if self._python_writer is not None:
return self._python_writer.write_arrow_batch(data, bucket)
data = self._prepare_native_arrow_data(data)
+ self._write_native_batch(data, bucket)
+
+ def _write_native_batch(self, data, bucket=None):
if data.num_rows:
# A failed native write may already have produced files. Never
# retry that batch through Python after this point.
@@ -243,8 +246,10 @@ class NativeTableWrite:
return self._switch_to_python().write_row(row)
values = row_to_named_values(row, self.table.table_schema.fields)
names = self._schema.names
- self.write_arrow_batch(row_values_to_arrow_table(
+ data = self._prepare_native_arrow_data(row_values_to_arrow_table(
values, self.table.table_schema.fields, names).to_batches()[0])
+ from pypaimon.write.row_kind import with_row_kind
+ self._write_native_batch(with_row_kind(self.table, data, row))
def prepare_commit(self, commit_identifier=None):
if self._python_writer is not None:
diff --git a/paimon-python/pypaimon/write/row_kind.py
b/paimon-python/pypaimon/write/row_kind.py
index 9e501e6565..78b5bc5e92 100644
--- a/paimon-python/pypaimon/write/row_kind.py
+++ b/paimon-python/pypaimon/write/row_kind.py
@@ -24,6 +24,8 @@ from pypaimon.table.row.row_kind import RowKind
def row_kinds(options, data):
name = options.options.to_map().get('rowkind.field')
if name is None:
+ if '_VALUE_KIND' in data.schema.names:
+ return data.column('_VALUE_KIND').to_pylist()
return [0] * data.num_rows
if name not in data.schema.names:
raise ValueError('Cannot find rowkind field %s in table schema' % name)
@@ -49,15 +51,26 @@ def filter_write_batch(table, data):
return data.filter(pa.array(keep, type=pa.bool_()))
-def skip_write_row(table, values):
+def with_row_kind(table, data, row):
+ """Carry an InternalRow's kind through the internal Arrow write path.
+
+ A configured rowkind.field takes precedence, as in Java RowKindGenerator.
+ Public Arrow writes still accept only the declared table/write schema.
+ """
+ if table.is_primary_key_table and
table.options.options.to_map().get('rowkind.field') is None:
+ return data.append_column(
+ pa.field('_VALUE_KIND', pa.int8(), nullable=False),
+ pa.array([row.get_row_kind().value] * data.num_rows,
type=pa.int8()))
+ return data
+
+
+def skip_write_row(table, values, row_kind=RowKind.INSERT):
if not table.is_primary_key_table:
return False
options = table.options
raw = options.options.to_map()
name = raw.get('rowkind.field')
- if name is None:
- return False
- kind = RowKind.from_string(values[name]).value
+ kind = row_kind.value if name is None else
RowKind.from_string(values[name]).value
return _is_filtered(kind, options.ignore_delete(),
str(raw.get('ignore-update-before', 'false')).lower()
== 'true')
diff --git a/paimon-python/pypaimon/write/table_write.py
b/paimon-python/pypaimon/write/table_write.py
index 04e7b4df01..44909ea45a 100644
--- a/paimon-python/pypaimon/write/table_write.py
+++ b/paimon-python/pypaimon/write/table_write.py
@@ -159,7 +159,7 @@ class TableWrite:
require_columns(values_by_name, column_names, "write_row")
require_columns(values_by_name, self.table.partition_keys, "write_row")
from pypaimon.write.row_kind import skip_write_row
- if skip_write_row(self.table, values_by_name):
+ if skip_write_row(self.table, values_by_name, row.get_row_kind()):
return
partition, bucket = (
self.row_key_extractor.extract_partition_bucket_row(values_by_name)
@@ -314,9 +314,26 @@ class TableWrite:
table = getattr(self, 'table', None)
if table is None:
return data
+ data = self._normalize_arrow_write_schema(data)
from pypaimon.write.row_kind import filter_write_batch
return filter_write_batch(table, data)
+ def _normalize_arrow_write_schema(self, data):
+ # Compatible input batches may differ in top-level nullability or
+ # metadata. Buffered writes and row-to-Arrow conversion must use the
+ # same declared fields, while keeping normalized Arrow array types.
+ fields = []
+ for field in data.schema:
+ declared = self.table_pyarrow_schema.field(field.name)
+ fields.append(pa.field(field.name, field.type,
nullable=declared.nullable,
+ metadata=declared.metadata))
+ schema = pa.schema(fields, metadata=self.table_pyarrow_schema.metadata)
+ if data.schema.equals(schema, check_metadata=True):
+ return data
+ if isinstance(data, pa.RecordBatch):
+ return pa.RecordBatch.from_arrays(data.columns, schema=schema)
+ return pa.Table.from_arrays(data.columns, schema=schema)
+
def _validate_pyarrow_schema(self, data_schema: pa.Schema):
if self._is_compatible_pyarrow_schema(data_schema,
self.table_pyarrow_schema):
return
diff --git a/paimon-python/pypaimon/write/writer/data_writer.py
b/paimon-python/pypaimon/write/writer/data_writer.py
index 4d78940cdd..d6ded1e0f0 100644
--- a/paimon-python/pypaimon/write/writer/data_writer.py
+++ b/paimon-python/pypaimon/write/writer/data_writer.py
@@ -310,70 +310,26 @@ class DataWriter(ABC):
zstd_level=self.zstd_level)
extra_files.append(row_sidecar_name)
- # min key & max key
-
- selected_table = logical_data.select(self.trimmed_primary_keys)
- key_columns_batch = selected_table.to_batches()[0]
- min_key_row_batch = key_columns_batch.slice(0, 1)
- max_key_row_batch =
key_columns_batch.slice(key_columns_batch.num_rows - 1, 1)
- min_key = [col.to_pylist()[0] for col in min_key_row_batch.columns]
- max_key = [col.to_pylist()[0] for col in max_key_row_batch.columns]
-
- # key stats & value stats
- value_stats_enabled = self.options.metadata_stats_enabled()
- if value_stats_enabled:
- stats_fields = self.table.fields if
self.table.is_primary_key_table \
- else
PyarrowFieldParser.to_paimon_schema(logical_data.schema)
- else:
- stats_fields = self.table.trimmed_primary_keys_fields
- column_stats = {
- field.name: self._get_column_stats(logical_data, field.name)
- for field in stats_fields
- }
- key_fields = self.trimmed_primary_keys_fields
- key_stats = self._collect_value_stats(
- logical_data, key_fields, column_stats)
- if not self.options.primary_key_nullable() and not all(
- count == 0 for count in key_stats.null_counts):
- raise RuntimeError("Primary key should not be null")
-
- value_fields = stats_fields if value_stats_enabled else []
- value_stats = self._collect_value_stats(
- logical_data, value_fields, column_stats)
-
- # Read the range without advancing it: the advance belongs with the
- # append below, so a retried flush derives the same range.
- if self.table.is_primary_key_table and '_SEQUENCE_NUMBER' in
data.schema.names:
- # PK files can contain a deduplicated or split subset of the
- # buffer. Java records the range of the rows actually written.
- sequences = data.column('_SEQUENCE_NUMBER')
- min_seq = pc.min(sequences).as_py()
- max_seq = pc.max(sequences).as_py()
- else:
- min_seq = self.sequence_generator.start
- max_seq = self.sequence_generator.current
- creation_time = Timestamp.now()
+ statistics = self._data_file_statistics(logical_data)
data_meta = self._create_data_file_meta(
file_name=file_name,
file_path=file_path,
row_count=data.num_rows,
- min_key=GenericRow(min_key, self.trimmed_primary_keys_fields),
- max_key=GenericRow(max_key, self.trimmed_primary_keys_fields),
- key_stats=key_stats,
- value_stats=value_stats,
- min_sequence_number=min_seq,
- max_sequence_number=max_seq,
+ min_key=GenericRow(statistics['min_key'],
self.trimmed_primary_keys_fields),
+ max_key=GenericRow(statistics['max_key'],
self.trimmed_primary_keys_fields),
+ key_stats=statistics['key_stats'],
+ value_stats=statistics['value_stats'],
+ min_sequence_number=statistics['min_seq'],
+ max_sequence_number=statistics['max_seq'],
extra_files=extra_files,
- creation_time=creation_time,
+ creation_time=statistics['creation_time'],
delete_row_count=self._count_delete_rows(data),
)
- if self.changelog_producer == ChangelogProducer.INPUT:
+ if (self.changelog_producer == ChangelogProducer.INPUT
+ and not self._writes_changelog_before_merge()):
changelog_meta = self._write_changelog_file(
- data, min_key, max_key, key_stats, value_stats,
- min_seq, max_seq, creation_time,
- value_stats_enabled, external_path_str is not None,
- )
+ data, is_external=external_path_str is not None,
**statistics)
except Exception:
self.file_io.delete_quietly(file_path)
if row_sidecar_path is not None:
@@ -382,6 +338,57 @@ class DataWriter(ABC):
self._finish_data_file(data_meta, changelog_meta,
shared_shredding_stats)
+ def _data_file_statistics(self, data):
+ # min key & max key
+
+ selected_table = data.select(self.trimmed_primary_keys)
+ key_columns_batch = selected_table.to_batches()[0]
+ min_key_row_batch = key_columns_batch.slice(0, 1)
+ max_key_row_batch = key_columns_batch.slice(key_columns_batch.num_rows
- 1, 1)
+ min_key = [col.to_pylist()[0] for col in min_key_row_batch.columns]
+ max_key = [col.to_pylist()[0] for col in max_key_row_batch.columns]
+
+ # key stats & value stats
+ value_stats_enabled = self.options.metadata_stats_enabled()
+ if value_stats_enabled:
+ stats_fields = self.table.fields if
self.table.is_primary_key_table \
+ else PyarrowFieldParser.to_paimon_schema(data.schema)
+ else:
+ stats_fields = self.table.trimmed_primary_keys_fields
+ column_stats = {
+ field.name: self._get_column_stats(data, field.name)
+ for field in stats_fields
+ }
+ key_fields = self.trimmed_primary_keys_fields
+ key_stats = self._collect_value_stats(
+ data, key_fields, column_stats)
+ if not self.options.primary_key_nullable() and not all(
+ count == 0 for count in key_stats.null_counts):
+ raise RuntimeError("Primary key should not be null")
+
+ value_fields = stats_fields if value_stats_enabled else []
+ value_stats = self._collect_value_stats(
+ data, value_fields, column_stats)
+
+ # Read the range without advancing it: the advance belongs with the
+ # append below, so a retried flush derives the same range.
+ if self.table.is_primary_key_table and '_SEQUENCE_NUMBER' in
data.schema.names:
+ # PK files can contain a deduplicated or split subset of the
+ # buffer. Java records the range of the rows actually written.
+ sequences = data.column('_SEQUENCE_NUMBER')
+ min_seq = pc.min(sequences).as_py()
+ max_seq = pc.max(sequences).as_py()
+ else:
+ min_seq = self.sequence_generator.start
+ max_seq = self.sequence_generator.current
+ creation_time = Timestamp.now()
+ return dict(min_key=min_key, max_key=max_key, key_stats=key_stats,
+ value_stats=value_stats, min_seq=min_seq, max_seq=max_seq,
+ creation_time=creation_time,
value_stats_enabled=value_stats_enabled)
+
+ def _writes_changelog_before_merge(self):
+ return False
+
def _finish_data_file(self, data_meta, changelog_meta=None,
shared_shredding_stats=None):
"""Record a fully written data file and its optional changelog."""
@@ -459,11 +466,10 @@ class DataWriter(ABC):
def _write_changelog_file(self, data, min_key, max_key, key_stats,
value_stats,
min_seq, max_seq, creation_time,
value_stats_enabled, is_external) ->
DataFileMeta:
- """Write the changelog file for one data file and return its meta.
+ """Write a complete changelog file and return its metadata.
- The caller appends the returned meta only once the whole data file has
- landed, so a failure here leaves nothing behind: no meta to commit, and
- no file on disk either.
+ The caller publishes this metadata after success. A failure here leaves
+ neither metadata to commit nor a partial changelog file on disk.
"""
cl_fmt = self.changelog_file_format
changelog_file_name = f"changelog-{uuid.uuid4()}-0.{cl_fmt}"
diff --git a/paimon-python/pypaimon/write/writer/key_value_data_writer.py
b/paimon-python/pypaimon/write/writer/key_value_data_writer.py
index 9e9998cd3d..b2e79d50d0 100644
--- a/paimon-python/pypaimon/write/writer/key_value_data_writer.py
+++ b/paimon-python/pypaimon/write/writer/key_value_data_writer.py
@@ -27,6 +27,7 @@ from pypaimon.read.reader.deduplicate_merge_function import \
from pypaimon.common.options.core_options import ChangelogProducer
from pypaimon.table.row.key_value import KeyValue
from pypaimon.write.writer.data_writer import DataWriter
+from pypaimon.write.writer.write_buffer import WriteBuffer
class KeyValueDataWriter(DataWriter):
@@ -50,6 +51,16 @@ class KeyValueDataWriter(DataWriter):
# paths that don't go through FileStoreWrite) don't accidentally
# skip the merge step entirely.
self._merge_function = merge_function or DeduplicateMergeFunction()
+ # Java's INPUT producer records every accepted input before the merge
+ # engine folds equal keys. Changelog files have their own rolling and
+ # metadata, independent of the data files produced by that fold.
+ self._input_changelog_buffer = (
+ WriteBuffer(self._merge_data)
+ if changelog_producer == ChangelogProducer.INPUT
+ and self._writes_changelog_before_merge() else None)
+
+ def _writes_changelog_before_merge(self):
+ return True
def _process_data(self, data: pa.RecordBatch) -> pa.Table:
# No sort here: sorting once at flush is strictly cheaper than
@@ -57,7 +68,10 @@ class KeyValueDataWriter(DataWriter):
# concat of unsorted batches; ``_flush_all`` sorts it exactly
# once before folding.
enhanced_data = self._add_system_fields(data)
- return pa.Table.from_batches([enhanced_data])
+ result = pa.Table.from_batches([enhanced_data])
+ if self._input_changelog_buffer is not None:
+ self._input_changelog_buffer.append(result)
+ return result
def _merge_data(self, existing_data: pa.Table, new_data: pa.Table) ->
pa.Table:
# Plain concat. Sort + fold both run inside ``_flush_all`` so
@@ -65,7 +79,7 @@ class KeyValueDataWriter(DataWriter):
return pa.concat_tables([existing_data, new_data])
def prepare_commit(self) -> List[DataFileMeta]:
- if self._buffer.num_rows > 0:
+ if self._buffer.num_rows > 0 or self._has_pending_input_changelog():
self._flush_all()
# ``_flush_all`` empties the buffer, so super's prepare_commit just
# returns ``committed_files``.
@@ -88,7 +102,7 @@ class KeyValueDataWriter(DataWriter):
# ``_flush_all`` so the contract holds even on the
# close-without-prepare_commit path.
try:
- if self._buffer.num_rows > 0:
+ if self._buffer.num_rows > 0 or
self._has_pending_input_changelog():
self._flush_all()
except Exception as e:
import logging
@@ -109,6 +123,7 @@ class KeyValueDataWriter(DataWriter):
left holding exactly the rows no file has taken yet, so a retried
flush neither loses nor duplicates them.
"""
+ self._flush_input_changelog()
pending = self._buffer.materialize()
if pending is None or pending.num_rows == 0:
self._buffer.reset()
@@ -124,6 +139,34 @@ class KeyValueDataWriter(DataWriter):
self._buffer.reset(folded)
self._roll_write(folded)
+ def _has_pending_input_changelog(self):
+ return self._input_changelog_buffer is not None and
self._input_changelog_buffer.num_rows > 0
+
+ def _flush_input_changelog(self):
+ if not self._has_pending_input_changelog():
+ return
+ data =
self._sort_by_primary_key(self._input_changelog_buffer.materialize())
+ while data.num_rows:
+ split_row = data.num_rows
+ if data.nbytes > self.target_file_size:
+ split_row = max(1, self._find_optimal_split_point(data,
self.target_file_size))
+ logical_data = data.slice(0, split_row)
+ statistics = self._data_file_statistics(logical_data)
+ physical_data = (self._apply_variant_shredding(logical_data)
+ if self._variant_shredding else logical_data)
+ meta = self._write_changelog_file(
+ physical_data, is_external=self.external_path_provider is not
None, **statistics)
+ self.committed_changelog_files.append(meta)
+ # Advance only after file and metadata are complete. A data-file
+ # failure later must neither lose inputs nor write their logs
twice.
+ data = data.slice(split_row)
+ self._input_changelog_buffer.reset(data if data.num_rows else None)
+
+ def abort(self):
+ super().abort()
+ if self._input_changelog_buffer is not None:
+ self._input_changelog_buffer.reset()
+
def _roll_write(self, data: pa.Table) -> None:
"""Write ``data`` as one or more files, each <= target_file_size.
@@ -275,6 +318,8 @@ class KeyValueDataWriter(DataWriter):
new_fields.append(pa.field('_VALUE_KIND', pa.int8(), nullable=False))
for i in range(data.num_columns):
+ if data.schema.field(i).name == '_VALUE_KIND':
+ continue
new_arrays.append(data.column(i))
new_fields.append(data.schema.field(i))
