| # 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 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 = _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 = _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): |
| """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 asyncio.sleep(2) # Wait for table initialization |
| |
| # 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 = _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 = _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_poll_batches(connection, admin): |
| """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 asyncio.sleep(1) |
| |
| 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 = 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 = _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 = _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 = _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 = 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 = 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 = 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_partitioned_table_append_scan(connection, admin): |
| """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 asyncio.sleep(2) # Wait for partitions to be available |
| |
| 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 = 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 = _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 = _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 = 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 = 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): |
| """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 asyncio.sleep(2) |
| |
| 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 = 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 = 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) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Helpers |
| # --------------------------------------------------------------------------- |
| |
| |
| 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 = scanner.poll(5000) |
| collected.extend(records) |
| return collected |
| |
| |
| |
| |
| 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 = scanner.poll_arrow(5000) |
| if arrow_table.num_rows > 0: |
| all_ids.extend(arrow_table.column("id").to_pylist()) |
| return all_ids |
