| # 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. |
| |
| """Integration tests for log (append-only) table operations. |
| |
| Mirrors the Rust integration tests in crates/fluss/tests/integration/log_table.rs. |
| """ |
| |
| import asyncio |
| import time |
| |
| import pyarrow as pa |
| import pytest |
| from conftest import ( |
| assert_complex_edge, |
| assert_complex_full, |
| complex_column_names, |
| complex_edge_row, |
| complex_full_row, |
| complex_null_row, |
| complex_schema, |
| ) |
| |
| import fluss |
| |
| |
| async def test_append_and_scan(connection, admin): |
| """Test appending record batches and scanning with a record-based scanner.""" |
| table_path = fluss.TablePath("fluss", "py_test_append_and_scan") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("c1", pa.int32()), pa.field("c2", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema, bucket_count=3, bucket_keys=["c1"]) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| |
| batch1 = pa.RecordBatch.from_arrays( |
| [pa.array([1, 2, 3], type=pa.int32()), pa.array(["a1", "a2", "a3"])], |
| schema=pa.schema([pa.field("c1", pa.int32()), pa.field("c2", pa.string())]), |
| ) |
| append_writer.write_arrow_batch(batch1) |
| |
| batch2 = pa.RecordBatch.from_arrays( |
| [pa.array([4, 5, 6], type=pa.int32()), pa.array(["a4", "a5", "a6"])], |
| schema=pa.schema([pa.field("c1", pa.int32()), pa.field("c2", pa.string())]), |
| ) |
| append_writer.write_arrow_batch(batch2) |
| |
| await append_writer.flush() |
| |
| # Scan with record-based scanner |
| scanner = await table.new_scan().create_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| records = await _poll_records(scanner, expected_count=6) |
| |
| assert len(records) == 6, f"Expected 6 records, got {len(records)}" |
| |
| records.sort(key=lambda r: r.row["c1"]) |
| |
| expected_c1 = [1, 2, 3, 4, 5, 6] |
| expected_c2 = ["a1", "a2", "a3", "a4", "a5", "a6"] |
| for i, record in enumerate(records): |
| assert record.row["c1"] == expected_c1[i], f"c1 mismatch at row {i}" |
| assert record.row["c2"] == expected_c2[i], f"c2 mismatch at row {i}" |
| |
| # Test unsubscribe |
| scanner.unsubscribe(bucket_id=0) |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_append_dict_rows(connection, admin): |
| """Test appending rows as dicts and scanning.""" |
| table_path = fluss.TablePath("fluss", "py_test_append_dict_rows") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| |
| # Append using dicts |
| append_writer.append({"id": 1, "name": "Alice"}) |
| append_writer.append({"id": 2, "name": "Bob"}) |
| # Append using lists |
| append_writer.append([3, "Charlie"]) |
| await append_writer.flush() |
| |
| scanner = await table.new_scan().create_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| records = await _poll_records(scanner, expected_count=3) |
| assert len(records) == 3 |
| |
| rows = sorted([r.row for r in records], key=lambda r: r["id"]) |
| assert rows[0] == {"id": 1, "name": "Alice"} |
| assert rows[1] == {"id": 2, "name": "Bob"} |
| assert rows[2] == {"id": 3, "name": "Charlie"} |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_list_offsets(connection, admin, wait_for_table_ready): |
| """Test listing earliest, latest, and timestamp-based offsets.""" |
| table_path = fluss.TablePath("fluss", "py_test_list_offsets") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| await wait_for_table_ready(table_path) |
| |
| # Earliest offset should be 0 for empty table |
| earliest = await admin.list_offsets( |
| table_path, bucket_ids=[0], offset_spec=fluss.OffsetSpec.earliest() |
| ) |
| assert earliest[0] == 0 |
| |
| # Latest offset should be 0 for empty table |
| latest = await admin.list_offsets( |
| table_path, bucket_ids=[0], offset_spec=fluss.OffsetSpec.latest() |
| ) |
| assert latest[0] == 0 |
| |
| before_append_ms = int(time.time() * 1000) |
| |
| # Append some records |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| batch = pa.RecordBatch.from_arrays( |
| [ |
| pa.array([1, 2, 3], type=pa.int32()), |
| pa.array(["alice", "bob", "charlie"]), |
| ], |
| schema=pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]), |
| ) |
| append_writer.write_arrow_batch(batch) |
| await append_writer.flush() |
| |
| await asyncio.sleep(1) |
| |
| after_append_ms = int(time.time() * 1000) |
| |
| # Latest offset should be 3 after appending 3 records |
| latest_after = await admin.list_offsets( |
| table_path, bucket_ids=[0], offset_spec=fluss.OffsetSpec.latest() |
| ) |
| assert latest_after[0] == 3 |
| |
| # Earliest offset should still be 0 |
| earliest_after = await admin.list_offsets( |
| table_path, bucket_ids=[0], offset_spec=fluss.OffsetSpec.earliest() |
| ) |
| assert earliest_after[0] == 0 |
| |
| # Timestamp before append should resolve to offset 0 |
| ts_before = await admin.list_offsets( |
| table_path, |
| bucket_ids=[0], |
| offset_spec=fluss.OffsetSpec.timestamp(before_append_ms), |
| ) |
| assert ts_before[0] == 0 |
| |
| # Intentional sleep to avoid race condition FlussError(code=38) The timestamp is invalid |
| await asyncio.sleep(1) |
| |
| # Timestamp after append should resolve to offset 3 |
| ts_after = await admin.list_offsets( |
| table_path, |
| bucket_ids=[0], |
| offset_spec=fluss.OffsetSpec.timestamp(after_append_ms), |
| ) |
| assert ts_after[0] == 3 |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_project(connection, admin): |
| """Test column projection by name and by index.""" |
| table_path = fluss.TablePath("fluss", "py_test_project") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("col_a", pa.int32()), |
| pa.field("col_b", pa.string()), |
| pa.field("col_c", pa.int32()), |
| ] |
| ) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| |
| batch = pa.RecordBatch.from_arrays( |
| [ |
| pa.array([1, 2, 3], type=pa.int32()), |
| pa.array(["x", "y", "z"]), |
| pa.array([10, 20, 30], type=pa.int32()), |
| ], |
| schema=pa.schema( |
| [ |
| pa.field("col_a", pa.int32()), |
| pa.field("col_b", pa.string()), |
| pa.field("col_c", pa.int32()), |
| ] |
| ), |
| ) |
| append_writer.write_arrow_batch(batch) |
| await append_writer.flush() |
| |
| # Test project_by_name: select col_b and col_c only |
| scan = table.new_scan().project_by_name(["col_b", "col_c"]) |
| scanner = await scan.create_log_scanner() |
| scanner.subscribe_buckets({0: 0}) |
| |
| records = await _poll_records(scanner, expected_count=3) |
| assert len(records) == 3 |
| |
| records.sort(key=lambda r: r.row["col_c"]) |
| expected_col_b = ["x", "y", "z"] |
| expected_col_c = [10, 20, 30] |
| for i, record in enumerate(records): |
| assert record.row["col_b"] == expected_col_b[i] |
| assert record.row["col_c"] == expected_col_c[i] |
| # col_a should not be present in projected results |
| assert "col_a" not in record.row |
| |
| # Test project by indices [1, 0] -> (col_b, col_a) |
| scanner2 = await table.new_scan().project([1, 0]).create_log_scanner() |
| scanner2.subscribe_buckets({0: 0}) |
| |
| records2 = await _poll_records(scanner2, expected_count=3) |
| assert len(records2) == 3 |
| |
| records2.sort(key=lambda r: r.row["col_a"]) |
| for i, record in enumerate(records2): |
| assert record.row["col_b"] == expected_col_b[i] |
| assert record.row["col_a"] == [1, 2, 3][i] |
| assert "col_c" not in record.row |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_project_compound_types(connection, admin): |
| """Projection selects and reorders ROW/MAP/ARRAY columns and drops others |
| (mirrors the Rust projection_with_compound_types IT).""" |
| table_path = fluss.TablePath("fluss", "py_test_project_compound") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field( |
| "nested", pa.struct([("seq", pa.int32()), ("label", pa.string())]) |
| ), |
| pa.field("attrs", pa.map_(pa.string(), pa.int32())), |
| pa.field("tags", pa.list_(pa.string())), |
| pa.field("extra", pa.string()), |
| ] |
| ) |
| ) |
| await admin.create_table( |
| table_path, fluss.TableDescriptor(schema), ignore_if_exists=False |
| ) |
| table = await connection.get_table(table_path) |
| |
| aw = table.new_append().create_writer() |
| aw.append( |
| { |
| "id": 7, |
| "nested": {"seq": 42, "label": "hello"}, |
| "attrs": {"x": 1, "y": 2}, |
| "tags": ["alpha", "beta"], |
| "extra": "ignore-me", |
| } |
| ) |
| await aw.flush() |
| |
| # Reorder and drop `extra`. |
| scan = table.new_scan().project_by_name(["nested", "attrs", "tags", "id"]) |
| scanner = await scan.create_log_scanner() |
| scanner.subscribe_buckets({0: fluss.EARLIEST_OFFSET}) |
| records = await _poll_records(scanner, expected_count=1) |
| assert len(records) == 1 |
| row = records[0].row |
| assert row["nested"] == {"seq": 42, "label": "hello"} |
| assert dict(row["attrs"]) == {"x": 1, "y": 2} |
| assert row["tags"] == ["alpha", "beta"] |
| assert row["id"] == 7 |
| assert "extra" not in row |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_poll_batches(connection, admin, wait_for_table_ready): |
| """Test batch-based scanning with poll_arrow and poll_record_batch.""" |
| table_path = fluss.TablePath("fluss", "py_test_poll_batches") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| await wait_for_table_ready(table_path) |
| |
| table = await connection.get_table(table_path) |
| scanner = await table.new_scan().create_record_batch_log_scanner() |
| scanner.subscribe(bucket_id=0, start_offset=0) |
| |
| # Empty table should return empty result |
| result = await scanner.poll_arrow(500) |
| assert result.num_rows == 0 |
| |
| writer = table.new_append().create_writer() |
| pa_schema = pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [pa.array([1, 2], type=pa.int32()), pa.array(["a", "b"])], |
| schema=pa_schema, |
| ) |
| ) |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [pa.array([3, 4], type=pa.int32()), pa.array(["c", "d"])], |
| schema=pa_schema, |
| ) |
| ) |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [pa.array([5, 6], type=pa.int32()), pa.array(["e", "f"])], |
| schema=pa_schema, |
| ) |
| ) |
| await writer.flush() |
| |
| # Poll until we get all 6 records |
| all_ids = await _poll_arrow_ids(scanner, expected_count=6) |
| assert all_ids == [1, 2, 3, 4, 5, 6] |
| |
| # Append more and verify offset continuation (no duplicates) |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [pa.array([7, 8], type=pa.int32()), pa.array(["g", "h"])], |
| schema=pa_schema, |
| ) |
| ) |
| await writer.flush() |
| |
| new_ids = await _poll_arrow_ids(scanner, expected_count=2) |
| assert new_ids == [7, 8] |
| |
| # Subscribe from mid-offset should truncate (skip earlier records) |
| trunc_scanner = await table.new_scan().create_record_batch_log_scanner() |
| trunc_scanner.subscribe(bucket_id=0, start_offset=3) |
| |
| trunc_ids = await _poll_arrow_ids(trunc_scanner, expected_count=5) |
| assert trunc_ids == [4, 5, 6, 7, 8] |
| |
| # Projection with batch scanner |
| proj_scanner = ( |
| await table.new_scan().project_by_name(["id"]).create_record_batch_log_scanner() |
| ) |
| proj_scanner.subscribe(bucket_id=0, start_offset=0) |
| batches = await proj_scanner.poll_record_batch(10000) |
| assert len(batches) > 0 |
| assert batches[0].batch.num_columns == 1 |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_to_arrow_and_to_pandas(connection, admin): |
| """Test to_arrow() and to_pandas() convenience methods.""" |
| table_path = fluss.TablePath("fluss", "py_test_to_arrow_pandas") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| |
| pa_schema = pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [pa.array([1, 2, 3], type=pa.int32()), pa.array(["a", "b", "c"])], |
| schema=pa_schema, |
| ) |
| ) |
| await writer.flush() |
| |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| |
| # to_arrow() |
| scanner = await table.new_scan().create_record_batch_log_scanner() |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| arrow_table = await scanner.to_arrow() |
| assert arrow_table.num_rows == 3 |
| assert arrow_table.schema.names == ["id", "name"] |
| |
| # to_pandas() |
| scanner2 = await table.new_scan().create_record_batch_log_scanner() |
| scanner2.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| df = await scanner2.to_pandas() |
| assert len(df) == 3 |
| assert list(df.columns) == ["id", "name"] |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_to_arrow_batch_reader(connection, admin): |
| """Test to_arrow_batch_reader() returns a lazy PyArrow RecordBatchReader.""" |
| table_path = fluss.TablePath("fluss", "py_test_to_arrow_batch_reader") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| |
| pa_schema = pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [pa.array([10, 20, 30], type=pa.int32()), pa.array(["x", "y", "z"])], |
| schema=pa_schema, |
| ) |
| ) |
| await writer.flush() |
| |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| |
| scanner = await table.new_scan().create_record_batch_log_scanner() |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| # to_arrow_batch_reader() is a blocking/sync API; run in a thread to |
| # avoid starving the asyncio event loop (see docstring warning). |
| def _read_all(): |
| reader = scanner.to_arrow_batch_reader() |
| assert isinstance(reader, pa.RecordBatchReader) |
| assert reader.schema == pa_schema |
| |
| batches = list(reader) |
| total_rows = sum(b.num_rows for b in batches) |
| assert total_rows == 3 |
| |
| result_table = pa.Table.from_batches(batches, schema=pa_schema) |
| assert result_table.column("id").to_pylist() == [10, 20, 30] |
| assert result_table.column("name").to_pylist() == ["x", "y", "z"] |
| |
| await asyncio.to_thread(_read_all) |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_to_arrow_batch_reader_drop_and_guard(connection, admin): |
| """Test reader-active guard and Drop cleanup on mid-iteration drop.""" |
| table_path = fluss.TablePath("fluss", "py_test_batch_reader_drop_guard") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| |
| pa_schema = pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| # Write multiple separate flushes so the server stores multiple log |
| # batches per bucket. This makes it likely that the reader's first poll |
| # only drains a subset, leaving real work for the Drop cleanup loop. |
| num_flushes = 10 |
| rows_per_flush = 200 |
| total_rows = num_flushes * rows_per_flush |
| for f in range(num_flushes): |
| start = f * rows_per_flush |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [ |
| pa.array( |
| list(range(start, start + rows_per_flush)), type=pa.int32() |
| ), |
| pa.array( |
| [f"row_{i}" for i in range(start, start + rows_per_flush)] |
| ), |
| ], |
| schema=pa_schema, |
| ) |
| ) |
| await writer.flush() |
| |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| |
| scanner = await table.new_scan().create_record_batch_log_scanner() |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| # to_arrow_batch_reader() is a blocking/sync API; run all blocking |
| # interactions in a thread to avoid starving the asyncio event loop. |
| def _test_guard_and_drop(): |
| # --- Guard blocks subscribe / unsubscribe while reader is active --- |
| reader = scanner.to_arrow_batch_reader() |
| with pytest.raises(fluss.FlussError, match="RecordBatchLogReader is active"): |
| scanner.subscribe_buckets({0: fluss.EARLIEST_OFFSET}) |
| with pytest.raises(fluss.FlussError, match="RecordBatchLogReader is active"): |
| scanner.unsubscribe(0) |
| |
| # --- Drop mid-iteration: read one batch, then discard --- |
| first_batch = next(reader) |
| assert first_batch.num_rows > 0 |
| del reader |
| |
| # --- Drop unsubscribed leftover buckets: creating a reader without |
| # re-subscribing must fail with "No buckets subscribed" --- |
| with pytest.raises(fluss.FlussError, match="No buckets subscribed"): |
| scanner.to_arrow_batch_reader() |
| |
| # --- Guard cleared after drop: scanner is reusable from a fresh subscribe --- |
| scanner.subscribe_buckets( |
| {i: fluss.EARLIEST_OFFSET for i in range(num_buckets)} |
| ) |
| reader2 = scanner.to_arrow_batch_reader() |
| batches = list(reader2) |
| assert sum(b.num_rows for b in batches) == total_rows |
| |
| await asyncio.to_thread(_test_guard_and_drop) |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_partitioned_table_append_scan(connection, admin, wait_for_table_ready): |
| """Test append and scan on a partitioned log table.""" |
| table_path = fluss.TablePath("fluss", "py_test_partitioned_log_append") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("region", pa.string()), |
| pa.field("value", pa.int64()), |
| ] |
| ) |
| ) |
| table_descriptor = fluss.TableDescriptor( |
| schema, |
| partition_keys=["region"], |
| ) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| # Create partitions |
| for region in ["US", "EU"]: |
| await admin.create_partition( |
| table_path, {"region": region}, ignore_if_exists=True |
| ) |
| await wait_for_table_ready(table_path, partition_name=region) |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| |
| # Append rows |
| test_data = [ |
| (1, "US", 100), |
| (2, "US", 200), |
| (3, "EU", 300), |
| (4, "EU", 400), |
| ] |
| for id_, region, value in test_data: |
| append_writer.append({"id": id_, "region": region, "value": value}) |
| await append_writer.flush() |
| |
| # Append arrow batches per partition |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("region", pa.string()), |
| pa.field("value", pa.int64()), |
| ] |
| ) |
| us_batch = pa.RecordBatch.from_arrays( |
| [ |
| pa.array([5, 6], type=pa.int32()), |
| pa.array(["US", "US"]), |
| pa.array([500, 600], type=pa.int64()), |
| ], |
| schema=pa_schema, |
| ) |
| append_writer.write_arrow_batch(us_batch) |
| |
| eu_batch = pa.RecordBatch.from_arrays( |
| [ |
| pa.array([7, 8], type=pa.int32()), |
| pa.array(["EU", "EU"]), |
| pa.array([700, 800], type=pa.int64()), |
| ], |
| schema=pa_schema, |
| ) |
| append_writer.write_arrow_batch(eu_batch) |
| await append_writer.flush() |
| |
| # Verify partition offsets |
| us_offsets = await admin.list_partition_offsets( |
| table_path, |
| partition_name="US", |
| bucket_ids=[0], |
| offset_spec=fluss.OffsetSpec.latest(), |
| ) |
| assert us_offsets[0] == 4, "US partition should have 4 records" |
| |
| eu_offsets = await admin.list_partition_offsets( |
| table_path, |
| partition_name="EU", |
| bucket_ids=[0], |
| offset_spec=fluss.OffsetSpec.latest(), |
| ) |
| assert eu_offsets[0] == 4, "EU partition should have 4 records" |
| |
| # Scan all partitions |
| scanner = await table.new_scan().create_log_scanner() |
| partition_infos = await admin.list_partition_infos(table_path) |
| for p in partition_infos: |
| scanner.subscribe_partition( |
| partition_id=p.partition_id, bucket_id=0, start_offset=0 |
| ) |
| |
| expected = [ |
| (1, "US", 100), |
| (2, "US", 200), |
| (3, "EU", 300), |
| (4, "EU", 400), |
| (5, "US", 500), |
| (6, "US", 600), |
| (7, "EU", 700), |
| (8, "EU", 800), |
| ] |
| |
| # Poll and verify per-bucket grouping |
| all_records = [] |
| deadline = time.monotonic() + 10 |
| while len(all_records) < 8 and time.monotonic() < deadline: |
| scan_records = await scanner.poll(5000) |
| for bucket, bucket_records in scan_records.items(): |
| assert bucket.partition_id is not None, ( |
| "Partitioned table should have partition_id" |
| ) |
| # All records in a bucket should belong to the same partition |
| regions = {r.row["region"] for r in bucket_records} |
| assert len(regions) == 1, f"Bucket has mixed regions: {regions}" |
| all_records.extend(bucket_records) |
| |
| assert len(all_records) == 8 |
| |
| collected = sorted( |
| [(r.row["id"], r.row["region"], r.row["value"]) for r in all_records], |
| key=lambda x: x[0], |
| ) |
| assert collected == expected |
| |
| # Test unsubscribe_partition: unsubscribe from EU, only US data should remain |
| unsub_scanner = await table.new_scan().create_log_scanner() |
| eu_partition_id = next( |
| p.partition_id for p in partition_infos if p.partition_name == "EU" |
| ) |
| for p in partition_infos: |
| unsub_scanner.subscribe_partition(p.partition_id, 0, 0) |
| unsub_scanner.unsubscribe_partition(eu_partition_id, 0) |
| |
| remaining = await _poll_records(unsub_scanner, expected_count=4, timeout_s=5) |
| assert len(remaining) == 4 |
| assert all(r.row["region"] == "US" for r in remaining) |
| |
| # Test subscribe_partition_buckets (batch subscribe) |
| batch_scanner = await table.new_scan().create_log_scanner() |
| partition_bucket_offsets = { |
| (p.partition_id, 0): fluss.EARLIEST_OFFSET for p in partition_infos |
| } |
| batch_scanner.subscribe_partition_buckets(partition_bucket_offsets) |
| |
| batch_records = await _poll_records(batch_scanner, expected_count=8) |
| assert len(batch_records) == 8 |
| batch_collected = sorted( |
| [(r.row["id"], r.row["region"], r.row["value"]) for r in batch_records], |
| key=lambda x: x[0], |
| ) |
| assert batch_collected == expected |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_write_arrow(connection, admin): |
| """Test writing a full PyArrow Table via write_arrow().""" |
| table_path = fluss.TablePath("fluss", "py_test_write_arrow") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| |
| pa_schema = pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| arrow_table = pa.table( |
| { |
| "id": pa.array([1, 2, 3, 4, 5], type=pa.int32()), |
| "name": pa.array(["alice", "bob", "charlie", "dave", "eve"]), |
| }, |
| schema=pa_schema, |
| ) |
| writer.write_arrow(arrow_table) |
| await writer.flush() |
| |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner = await table.new_scan().create_record_batch_log_scanner() |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| result = await scanner.to_arrow() |
| assert result.num_rows == 5 |
| |
| ids = sorted(result.column("id").to_pylist()) |
| names = [ |
| n |
| for _, n in sorted( |
| zip(result.column("id").to_pylist(), result.column("name").to_pylist()) |
| ) |
| ] |
| assert ids == [1, 2, 3, 4, 5] |
| assert names == ["alice", "bob", "charlie", "dave", "eve"] |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_write_pandas(connection, admin): |
| """Test writing a Pandas DataFrame via write_pandas().""" |
| import pandas as pd |
| |
| table_path = fluss.TablePath("fluss", "py_test_write_pandas") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| |
| df = pd.DataFrame({"id": [10, 20, 30], "name": ["x", "y", "z"]}) |
| writer.write_pandas(df) |
| await writer.flush() |
| |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner = await table.new_scan().create_record_batch_log_scanner() |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| result = await scanner.to_pandas() |
| assert len(result) == 3 |
| |
| result_sorted = result.sort_values("id").reset_index(drop=True) |
| assert result_sorted["id"].tolist() == [10, 20, 30] |
| assert result_sorted["name"].tolist() == ["x", "y", "z"] |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_partitioned_table_to_arrow(connection, admin, wait_for_table_ready): |
| """Test to_arrow() on partitioned tables.""" |
| table_path = fluss.TablePath("fluss", "py_test_partitioned_to_arrow") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("region", pa.string()), |
| pa.field("value", pa.int64()), |
| ] |
| ) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema, partition_keys=["region"]) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| for region in ["US", "EU"]: |
| await admin.create_partition( |
| table_path, {"region": region}, ignore_if_exists=True |
| ) |
| await wait_for_table_ready(table_path, partition_name=region) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| writer.append({"id": 1, "region": "US", "value": 100}) |
| writer.append({"id": 2, "region": "EU", "value": 200}) |
| await writer.flush() |
| |
| scanner = await table.new_scan().create_record_batch_log_scanner() |
| partition_infos = await admin.list_partition_infos(table_path) |
| for p in partition_infos: |
| scanner.subscribe_partition(p.partition_id, 0, fluss.EARLIEST_OFFSET) |
| |
| arrow_table = await scanner.to_arrow() |
| assert arrow_table.num_rows == 2 |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_scan_records_indexing_and_slicing(connection, admin): |
| """Test ScanRecords indexing, slicing (incl. negative steps), and iteration consistency.""" |
| table_path = fluss.TablePath("fluss", "py_test_scan_records_indexing") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("val", pa.string())]) |
| ) |
| await admin.create_table(table_path, fluss.TableDescriptor(schema)) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [ |
| pa.array(list(range(1, 9)), type=pa.int32()), |
| pa.array([f"v{i}" for i in range(1, 9)]), |
| ], |
| schema=pa.schema( |
| [pa.field("id", pa.int32()), pa.field("val", pa.string())] |
| ), |
| ) |
| ) |
| await writer.flush() |
| |
| scanner = await table.new_scan().create_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| # Poll until we get a non-empty ScanRecords (need ≥2 records for slice tests) |
| sr = None |
| deadline = time.monotonic() + 10 |
| while time.monotonic() < deadline: |
| sr = await scanner.poll(5000) |
| if len(sr) >= 2: |
| break |
| assert sr is not None and len(sr) >= 2, "Expected at least 2 records" |
| n = len(sr) |
| offsets = [sr[i].offset for i in range(n)] |
| |
| # Iteration and indexing must produce the same order |
| assert [r.offset for r in sr] == offsets |
| |
| # Negative indexing |
| assert sr[-1].offset == offsets[-1] |
| assert sr[-n].offset == offsets[0] |
| |
| # Verify slices match the same operation on the offsets reference list |
| test_slices = [ |
| slice(1, n - 1), # forward subrange |
| slice(None, None, -1), # [::-1] full reverse |
| slice(n - 2, 0, -1), # reverse with bounds |
| slice(n - 1, 0, -2), # reverse with step |
| slice(None, None, 2), # [::2] |
| slice(1, None, 3), # [1::3] |
| slice(2, 2), # empty |
| ] |
| for s in test_slices: |
| result = [r.offset for r in sr[s]] |
| assert result == offsets[s], f"slice {s}: got {result}, expected {offsets[s]}" |
| |
| # Bucket-based indexing |
| for bucket in sr.buckets(): |
| assert len(sr[bucket]) > 0 |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_async_iterator(connection, admin): |
| """Test the Python asynchronous iterator loop (`async for`) on LogScanner.""" |
| table_path = fluss.TablePath("fluss", "py_test_async_iterator") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("val", pa.string())]) |
| ) |
| await admin.create_table(table_path, fluss.TableDescriptor(schema)) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| |
| # Write 5 records |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [ |
| pa.array(list(range(1, 6)), type=pa.int32()), |
| pa.array([f"async{i}" for i in range(1, 6)]), |
| ], |
| schema=pa.schema( |
| [pa.field("id", pa.int32()), pa.field("val", pa.string())] |
| ), |
| ) |
| ) |
| await writer.flush() |
| |
| scanner = await table.new_scan().create_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| collected = [] |
| |
| # Here is the magical Issue #424 async iterator logic at work: |
| async def consume_scanner(): |
| async for record in scanner: |
| collected.append(record) |
| if len(collected) == 5: |
| break |
| |
| await consume_scanner() |
| |
| assert len(collected) == 5, f"Expected 5 records, got {len(collected)}" |
| |
| collected.sort(key=lambda r: r.row["id"]) |
| for i, record in enumerate(collected): |
| assert record.row["id"] == i + 1 |
| assert record.row["val"] == f"async{i + 1}" |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_async_iterator_break_no_leak(connection, admin): |
| """Verify that breaking out of `async for` does not leak resources. |
| |
| After breaking, the scanner must still be usable for synchronous |
| `poll()` calls. If the old implementation's tokio::spawn'd task |
| were still alive, it would hold the Mutex and cause `poll()` to |
| deadlock or error. |
| """ |
| table_path = fluss.TablePath("fluss", "py_test_async_break_leak") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("val", pa.string())]) |
| ) |
| await admin.create_table(table_path, fluss.TableDescriptor(schema)) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [ |
| pa.array(list(range(1, 11)), type=pa.int32()), |
| pa.array([f"v{i}" for i in range(1, 11)]), |
| ], |
| schema=pa.schema( |
| [pa.field("id", pa.int32()), pa.field("val", pa.string())] |
| ), |
| ) |
| ) |
| await writer.flush() |
| |
| scanner = await table.new_scan().create_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| # Phase 1: async for with early break (collect only 3 of 10) |
| collected_async = [] |
| |
| async def consume_and_break(): |
| async for record in scanner: |
| collected_async.append(record) |
| if len(collected_async) >= 3: |
| break |
| |
| await consume_and_break() |
| assert len(collected_async) == 3, ( |
| f"Expected 3 records from async for, got {len(collected_async)}" |
| ) |
| |
| # Phase 2: sync poll() must still work — proves no leaked task / lock. |
| # With small data and few buckets, _async_poll may have fetched all |
| # records in one batch. After break, the un-yielded records from that |
| # batch are lost. So sync poll may return 0 records — the key assertion |
| # is that poll() completes without deadlock (returns within timeout). |
| remaining = await scanner.poll(2000) |
| assert remaining is not None, "poll() should return (not deadlock)" |
| |
| # If we got records, verify no duplicates |
| async_ids = {r.row["id"] for r in collected_async} |
| sync_ids = {r.row["id"] for r in remaining} |
| assert async_ids.isdisjoint(sync_ids), ( |
| f"Duplicate IDs between async and sync: {async_ids & sync_ids}" |
| ) |
| |
| # All IDs must be from the original 1-10 range |
| all_ids = async_ids | sync_ids |
| assert all_ids.issubset(set(range(1, 11))), ( |
| f"Unexpected IDs: {all_ids - set(range(1, 11))}" |
| ) |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_async_iterator_multiple_batches(connection, admin): |
| """Verify async iteration works across multiple network poll cycles. |
| |
| _async_poll does a single bounded poll per call. Writing 20 records |
| to multiple buckets ensures the Python generator must loop through |
| several _async_poll calls to collect them all. |
| """ |
| table_path = fluss.TablePath("fluss", "py_test_async_multi_batch") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("val", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema, bucket_count=3, bucket_keys=["id"]) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| |
| num_records = 20 |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [ |
| pa.array(list(range(1, num_records + 1)), type=pa.int32()), |
| pa.array([f"multi{i}" for i in range(1, num_records + 1)]), |
| ], |
| schema=pa.schema( |
| [pa.field("id", pa.int32()), pa.field("val", pa.string())] |
| ), |
| ) |
| ) |
| await writer.flush() |
| |
| scanner = await table.new_scan().create_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| collected = [] |
| |
| async def consume_all(): |
| async for record in scanner: |
| collected.append(record) |
| if len(collected) >= num_records: |
| break |
| |
| await consume_all() |
| assert len(collected) == num_records, ( |
| f"Expected {num_records} records, got {len(collected)}" |
| ) |
| |
| # Verify all IDs are present (order may vary due to bucketing) |
| ids = sorted(r.row["id"] for r in collected) |
| assert ids == list(range(1, num_records + 1)) |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_batch_async_iterator(connection, admin): |
| """Test the Python asynchronous iterator loop (`async for`) on a batch LogScanner. |
| |
| With our __aiter__ dispatch, a batch-based scanner should yield RecordBatch |
| objects (not ScanRecord). Each yielded item has .batch (PyArrow RecordBatch), |
| .bucket, .base_offset, .last_offset. |
| """ |
| table_path = fluss.TablePath("fluss", "py_test_batch_async_iter") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("val", pa.string())]) |
| ) |
| await admin.create_table(table_path, fluss.TableDescriptor(schema)) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [ |
| pa.array(list(range(1, 7)), type=pa.int32()), |
| pa.array([f"bv{i}" for i in range(1, 7)]), |
| ], |
| schema=pa.schema( |
| [pa.field("id", pa.int32()), pa.field("val", pa.string())] |
| ), |
| ) |
| ) |
| await writer.flush() |
| |
| batch_scanner = await table.new_scan().create_record_batch_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| batch_scanner.subscribe_buckets( |
| {i: fluss.EARLIEST_OFFSET for i in range(num_buckets)} |
| ) |
| |
| collected_batches = [] |
| total_rows = 0 |
| |
| async def consume_batches(): |
| nonlocal total_rows |
| async for rb in batch_scanner: |
| collected_batches.append(rb) |
| total_rows += rb.batch.num_rows |
| if total_rows >= 6: |
| break |
| |
| await consume_batches() |
| |
| assert total_rows >= 6, f"Expected >=6 total rows, got {total_rows}" |
| assert len(collected_batches) > 0 |
| |
| # Verify each yielded item is a RecordBatch with expected attributes |
| for rb in collected_batches: |
| assert hasattr(rb, "batch"), "RecordBatch should have .batch" |
| assert hasattr(rb, "bucket"), "RecordBatch should have .bucket" |
| assert hasattr(rb, "base_offset"), "RecordBatch should have .base_offset" |
| assert hasattr(rb, "last_offset"), "RecordBatch should have .last_offset" |
| # .batch should be a PyArrow RecordBatch |
| arrow_batch = rb.batch |
| assert isinstance(arrow_batch, pa.RecordBatch), ( |
| f"Expected PyArrow RecordBatch, got {type(arrow_batch).__name__}" |
| ) |
| assert arrow_batch.num_columns == 2 |
| assert set(arrow_batch.schema.names) == {"id", "val"} |
| |
| # Verify all 6 IDs are present |
| all_ids = [] |
| for rb in collected_batches: |
| all_ids.extend(rb.batch.column("id").to_pylist()) |
| assert sorted(all_ids[:6]) == [1, 2, 3, 4, 5, 6] |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_batch_async_iterator_break_no_leak(connection, admin): |
| """Verify that breaking out of batch `async for` does not leak resources. |
| |
| After breaking, the scanner must still be usable for synchronous |
| poll_record_batch() calls, proving no leaked task or lock. |
| """ |
| table_path = fluss.TablePath("fluss", "py_test_batch_async_break") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("val", pa.string())]) |
| ) |
| await admin.create_table(table_path, fluss.TableDescriptor(schema)) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [ |
| pa.array(list(range(1, 11)), type=pa.int32()), |
| pa.array([f"bl{i}" for i in range(1, 11)]), |
| ], |
| schema=pa.schema( |
| [pa.field("id", pa.int32()), pa.field("val", pa.string())] |
| ), |
| ) |
| ) |
| await writer.flush() |
| |
| batch_scanner = await table.new_scan().create_record_batch_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| batch_scanner.subscribe_buckets( |
| {i: fluss.EARLIEST_OFFSET for i in range(num_buckets)} |
| ) |
| |
| # Phase 1: async for with early break (collect just 1 batch) |
| first_batch = None |
| |
| async def consume_and_break(): |
| nonlocal first_batch |
| async for rb in batch_scanner: |
| first_batch = rb |
| break |
| |
| await consume_and_break() |
| assert first_batch is not None, "Should have received at least 1 batch" |
| assert first_batch.batch.num_rows > 0 |
| |
| # Phase 2: sync poll_record_batch() must still work — proves no leak |
| remaining = await batch_scanner.poll_record_batch(2000) |
| assert remaining is not None, "poll_record_batch() should return (not deadlock)" |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_batch_async_iterator_multiple_batches(connection, admin): |
| """Verify batch async iteration works across multiple network poll cycles. |
| |
| Writing 20 records to 3 buckets ensures the generator must loop through |
| several _async_poll_batches calls to collect them all. |
| """ |
| table_path = fluss.TablePath("fluss", "py_test_batch_async_multi") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema([pa.field("id", pa.int32()), pa.field("val", pa.string())]) |
| ) |
| table_descriptor = fluss.TableDescriptor(schema, bucket_count=3, bucket_keys=["id"]) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_append().create_writer() |
| |
| num_records = 20 |
| writer.write_arrow_batch( |
| pa.RecordBatch.from_arrays( |
| [ |
| pa.array(list(range(1, num_records + 1)), type=pa.int32()), |
| pa.array([f"bm{i}" for i in range(1, num_records + 1)]), |
| ], |
| schema=pa.schema( |
| [pa.field("id", pa.int32()), pa.field("val", pa.string())] |
| ), |
| ) |
| ) |
| await writer.flush() |
| |
| batch_scanner = await table.new_scan().create_record_batch_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| batch_scanner.subscribe_buckets( |
| {i: fluss.EARLIEST_OFFSET for i in range(num_buckets)} |
| ) |
| |
| all_ids = [] |
| |
| async def consume_all(): |
| async for rb in batch_scanner: |
| all_ids.extend(rb.batch.column("id").to_pylist()) |
| if len(all_ids) >= num_records: |
| break |
| |
| await consume_all() |
| assert len(all_ids) >= num_records, ( |
| f"Expected >={num_records} IDs, got {len(all_ids)}" |
| ) |
| assert sorted(all_ids[:num_records]) == list(range(1, num_records + 1)) |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Helpers |
| # --------------------------------------------------------------------------- |
| |
| |
| async def _poll_records(scanner, expected_count, timeout_s=10): |
| """Poll a record-based scanner until expected_count records are collected.""" |
| collected = [] |
| deadline = time.monotonic() + timeout_s |
| while len(collected) < expected_count and time.monotonic() < deadline: |
| records = await scanner.poll(5000) |
| collected.extend(records) |
| return collected |
| |
| |
| async def _poll_arrow_ids(scanner, expected_count, timeout_s=10): |
| """Poll a batch scanner and extract 'id' column values.""" |
| all_ids = [] |
| deadline = time.monotonic() + timeout_s |
| while len(all_ids) < expected_count and time.monotonic() < deadline: |
| arrow_table = await scanner.poll_arrow(5000) |
| if arrow_table.num_rows > 0: |
| all_ids.extend(arrow_table.column("id").to_pylist()) |
| return all_ids |
| |
| |
| async def test_append_and_scan_with_array(connection, admin): |
| """Test appending and scanning with array columns.""" |
| table_path = fluss.TablePath("fluss", "py_test_append_and_scan_with_array") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("tags", pa.list_(pa.string())), |
| pa.field("scores", pa.list_(pa.int32())), |
| ] |
| ) |
| schema = fluss.Schema(pa_schema) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| |
| # Batch 1: Testing standard lists |
| batch1 = pa.RecordBatch.from_arrays( |
| [ |
| pa.array([1, 2], type=pa.int32()), |
| pa.array([["a", "b"], ["c"]], type=pa.list_(pa.string())), |
| pa.array([[10, 20], [30]], type=pa.list_(pa.int32())), |
| ], |
| schema=pa_schema, |
| ) |
| append_writer.write_arrow_batch(batch1) |
| |
| # Batch 2: Testing null values inside arrays and null arrays |
| batch2 = pa.RecordBatch.from_arrays( |
| [ |
| pa.array([3, 4, 5, 6], type=pa.int32()), |
| pa.array([["d", None], None, [], [None]], type=pa.list_(pa.string())), |
| pa.array([[40, 50], [60], None, []], type=pa.list_(pa.int32())), |
| ], |
| schema=pa_schema, |
| ) |
| append_writer.write_arrow_batch(batch2) |
| await append_writer.flush() |
| |
| # Verify via LogScanner (record-by-record) |
| scanner = await table.new_scan().create_log_scanner() |
| scanner.subscribe_buckets({0: fluss.EARLIEST_OFFSET}) |
| records = await _poll_records(scanner, expected_count=6) |
| |
| assert len(records) == 6 |
| records.sort(key=lambda r: r.row["id"]) |
| |
| # Verify Batch 1 |
| assert records[0].row["tags"] == ["a", "b"] |
| assert records[0].row["scores"] == [10, 20] |
| assert records[1].row["tags"] == ["c"] |
| assert records[1].row["scores"] == [30] |
| |
| # Verify Batch 2 |
| assert records[2].row["tags"] == ["d", None] |
| assert records[2].row["scores"] == [40, 50] |
| assert records[3].row["tags"] is None |
| assert records[3].row["scores"] == [60] |
| assert records[4].row["tags"] == [] |
| assert records[4].row["scores"] is None |
| assert records[5].row["tags"] == [None] |
| assert records[5].row["scores"] == [] |
| |
| # Verify via to_arrow (batch-based) |
| scanner2 = await table.new_scan().create_record_batch_log_scanner() |
| scanner2.subscribe_buckets({0: fluss.EARLIEST_OFFSET}) |
| result_table = await scanner2.to_arrow() |
| |
| assert result_table.num_rows == 6 |
| assert result_table.column("tags").to_pylist() == [ |
| ["a", "b"], |
| ["c"], |
| ["d", None], |
| None, |
| [], |
| [None], |
| ] |
| assert result_table.column("scores").to_pylist() == [ |
| [10, 20], |
| [30], |
| [40, 50], |
| [60], |
| None, |
| [], |
| ] |
| |
| |
| async def test_append_rows_with_array(connection, admin): |
| """Test appending rows with array data as Python lists and scanning.""" |
| table_path = fluss.TablePath("fluss", "py_test_append_rows_with_array") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("tags", pa.list_(pa.string())), |
| pa.field("scores", pa.list_(pa.int32())), |
| ] |
| ) |
| schema = fluss.Schema(pa_schema) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| |
| # Append rows using dicts with lists |
| append_writer.append({"id": 1, "tags": ["a", "b"], "scores": [10, 20]}) |
| append_writer.append({"id": 2, "tags": ["c"], "scores": [30]}) |
| # Append row using list with nested list (null handling) |
| append_writer.append([3, None, [40, None, 60]]) |
| |
| await append_writer.flush() |
| |
| scanner = await table.new_scan().create_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| records = await _poll_records(scanner, expected_count=3) |
| assert len(records) == 3 |
| |
| rows = sorted([r.row for r in records], key=lambda r: r["id"]) |
| assert rows[0] == {"id": 1, "tags": ["a", "b"], "scores": [10, 20]} |
| assert rows[1] == {"id": 2, "tags": ["c"], "scores": [30]} |
| # Note: records[2].row["tags"] will be None, records[2].row["scores"] will be [40, None, 60] |
| assert rows[2]["id"] == 3 |
| assert rows[2]["tags"] is None |
| assert rows[2]["scores"] == [40, None, 60] |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_append_rows_with_nested_array(connection, admin): |
| """Test appending rows with nested array data (ARRAY<ARRAY<INT>>) and scanning.""" |
| table_path = fluss.TablePath("fluss", "py_test_append_rows_with_nested_array") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("matrix", pa.list_(pa.list_(pa.int32()))), |
| ] |
| ) |
| schema = fluss.Schema(pa_schema) |
| await admin.create_table( |
| table_path, fluss.TableDescriptor(schema), ignore_if_exists=False |
| ) |
| |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| |
| # Append nested lists |
| append_writer.append({"id": 1, "matrix": [[1, 2], [3, 4]]}) |
| append_writer.append({"id": 2, "matrix": [[], [5], [6, 7, 8]]}) |
| append_writer.append({"id": 3, "matrix": None}) |
| append_writer.append({"id": 4, "matrix": [[1, None], None, []]}) |
| append_writer.append({"id": 5, "matrix": [[None, None]]}) |
| |
| await append_writer.flush() |
| |
| scanner = await table.new_scan().create_log_scanner() |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| records = await _poll_records(scanner, expected_count=5) |
| assert len(records) == 5 |
| |
| rows = sorted([r.row for r in records], key=lambda r: r["id"]) |
| assert rows[0] == {"id": 1, "matrix": [[1, 2], [3, 4]]} |
| assert rows[1] == {"id": 2, "matrix": [[], [5], [6, 7, 8]]} |
| assert rows[2] == {"id": 3, "matrix": None} |
| assert rows[3] == {"id": 4, "matrix": [[1, None], None, []]} |
| assert rows[4] == {"id": 5, "matrix": [[None, None]]} |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_append_rows_with_invalid_array(connection, admin): |
| """Test that appending invalid data to an array column raises an error.""" |
| table_path = fluss.TablePath("fluss", "py_test_append_rows_with_invalid_array") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("tags", pa.list_(pa.string())), |
| ] |
| ) |
| schema = fluss.Schema(pa_schema) |
| await admin.create_table( |
| table_path, fluss.TableDescriptor(schema), ignore_if_exists=False |
| ) |
| |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| |
| # Appending a string instead of a list should raise an error |
| with pytest.raises(Exception, match="Expected sequence for Array column"): |
| append_writer.append({"id": 4, "tags": "not_a_list"}) |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_all_complex_datatypes(connection, admin): |
| """Comprehensive ARRAY/MAP/ROW coverage on a log table. |
| |
| Appends fully-populated, edge-case, and all-null rows (see complex_types.py) |
| and verifies them through both the record (dict) scan path and the Arrow |
| scan path -- the two paths must agree. Mirrors the section-based |
| ``all_supported_datatypes`` structure of the Rust integration tests. |
| """ |
| table_path = fluss.TablePath("fluss", "py_test_log_all_complex") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = complex_schema() |
| await admin.create_table( |
| table_path, fluss.TableDescriptor(schema), ignore_if_exists=False |
| ) |
| |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| append_writer.append(complex_full_row(1)) |
| append_writer.append(complex_edge_row(2)) |
| append_writer.append(complex_null_row(3)) |
| await append_writer.flush() |
| |
| # Record (dict) scan path. |
| scanner = await table.new_scan().create_log_scanner() |
| scanner.subscribe_buckets({0: fluss.EARLIEST_OFFSET}) |
| records = await _poll_records(scanner, expected_count=3) |
| assert len(records) == 3 |
| poll_rows = sorted([r.row for r in records], key=lambda r: r["id"]) |
| assert_complex_full(poll_rows[0]) |
| assert_complex_edge(poll_rows[1]) |
| for col in complex_column_names(): |
| assert poll_rows[2][col] is None |
| |
| # Arrow scan path: to_pylist() agrees with the dict path column-for-column, |
| # except NaN-bearing float columns (NaN != NaN), which are checked elsewhere. |
| scanner2 = await table.new_scan().create_record_batch_log_scanner() |
| scanner2.subscribe_buckets({0: fluss.EARLIEST_OFFSET}) |
| result_table = await scanner2.to_arrow() |
| assert result_table.num_rows == 3 |
| arrow_rows = result_table.sort_by("id").to_pylist() |
| float_cols = {"arr_float", "arr_double", "map_float", "map_double"} |
| for i in range(3): |
| for col in complex_column_names(): |
| if col in float_cols: |
| continue |
| assert arrow_rows[i][col] == poll_rows[i][col], ( |
| f"scan-path mismatch at row {i}, column {col}" |
| ) |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_append_arrow_batch_complex_types(connection, admin): |
| """Arrow write path: write MAP and ROW columns via write_arrow_batch and |
| verify through both the record and Arrow scan paths.""" |
| table_path = fluss.TablePath("fluss", "py_test_arrow_batch_complex") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| row_type = pa.struct([("seq", pa.int32()), ("label", pa.string())]) |
| map_type = pa.map_(pa.string(), pa.int32()) |
| pa_schema = pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("attrs", map_type), |
| pa.field("nested", row_type), |
| ] |
| ) |
| schema = fluss.Schema(pa_schema) |
| await admin.create_table( |
| table_path, fluss.TableDescriptor(schema), ignore_if_exists=False |
| ) |
| |
| table = await connection.get_table(table_path) |
| append_writer = table.new_append().create_writer() |
| batch = pa.RecordBatch.from_arrays( |
| [ |
| pa.array([1, 2], type=pa.int32()), |
| pa.array([[("a", 1), ("b", 2)], []], type=map_type), |
| pa.array( |
| [{"seq": 10, "label": "open"}, {"seq": 20, "label": None}], |
| type=row_type, |
| ), |
| ], |
| schema=pa_schema, |
| ) |
| append_writer.write_arrow_batch(batch) |
| await append_writer.flush() |
| |
| # Record scan path. |
| scanner = await table.new_scan().create_log_scanner() |
| scanner.subscribe_buckets({0: fluss.EARLIEST_OFFSET}) |
| records = await _poll_records(scanner, expected_count=2) |
| assert len(records) == 2 |
| rows = sorted([r.row for r in records], key=lambda r: r["id"]) |
| assert dict(rows[0]["attrs"]) == {"a": 1, "b": 2} |
| assert rows[0]["nested"] == {"seq": 10, "label": "open"} |
| assert rows[1]["attrs"] == [] |
| assert rows[1]["nested"] == {"seq": 20, "label": None} |
| |
| # Arrow scan path. |
| scanner2 = await table.new_scan().create_record_batch_log_scanner() |
| scanner2.subscribe_buckets({0: fluss.EARLIEST_OFFSET}) |
| result_table = await scanner2.to_arrow() |
| assert result_table.column("id").to_pylist() == [1, 2] |
| assert result_table.column("attrs").to_pylist() == [[("a", 1), ("b", 2)], []] |
| assert result_table.column("nested").to_pylist() == [ |
| {"seq": 10, "label": "open"}, |
| {"seq": 20, "label": None}, |
| ] |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
