| # 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. |
| |
| import asyncio |
| import traceback |
| from datetime import date, datetime |
| from datetime import time as dt_time |
| from decimal import Decimal |
| |
| import pandas as pd |
| import pyarrow as pa |
| |
| import fluss |
| |
| |
| async def main(): |
| # Create connection configuration |
| config_spec = { |
| "bootstrap.servers": "127.0.0.1:9123", |
| # Add other configuration options as needed |
| "writer.request-max-size": "10485760", # 10 MB |
| "writer.acks": "all", # Wait for all replicas to acknowledge |
| "writer.retries": "3", # Retry up to 3 times on failure |
| "writer.batch-size": "1000", # Batch size for writes |
| } |
| config = fluss.Config(config_spec) |
| |
| # Create connection using the static create method |
| conn = await fluss.FlussConnection.create(config) |
| |
| # Define fields for PyArrow |
| fields = [ |
| pa.field("id", pa.int32()), |
| pa.field("name", pa.string()), |
| pa.field("score", pa.float32()), |
| pa.field("age", pa.int32()), |
| pa.field("birth_date", pa.date32()), |
| pa.field("check_in_time", pa.time32("ms")), |
| pa.field("created_at", pa.timestamp("us")), # TIMESTAMP (NTZ) |
| pa.field("updated_at", pa.timestamp("us", tz="UTC")), # TIMESTAMP_LTZ |
| pa.field("salary", pa.decimal128(10, 2)), |
| ] |
| |
| # Create a PyArrow schema |
| schema = pa.schema(fields) |
| |
| # Create a Fluss Schema first (this is what TableDescriptor expects) |
| fluss_schema = fluss.Schema(schema) |
| |
| # Create a Fluss TableDescriptor |
| table_descriptor = fluss.TableDescriptor(fluss_schema) |
| |
| # Get the admin for Fluss |
| admin = await conn.get_admin() |
| |
| # Create a Fluss table |
| table_path = fluss.TablePath("fluss", "sample_table_types") |
| |
| try: |
| await admin.create_table(table_path, table_descriptor, True) |
| print(f"Created table: {table_path}") |
| except Exception as e: |
| print(f"Table creation failed: {e}") |
| |
| # Get table information via admin |
| try: |
| table_info = await admin.get_table_info(table_path) |
| print(f"Table info: {table_info}") |
| print(f"Table ID: {table_info.table_id}") |
| print(f"Schema ID: {table_info.schema_id}") |
| print(f"Created time: {table_info.created_time}") |
| print(f"Primary keys: {table_info.get_primary_keys()}") |
| except Exception as e: |
| print(f"Failed to get table info: {e}") |
| |
| # Demo: List offsets |
| print("\n--- Testing list_offsets() ---") |
| try: |
| # Query latest offsets using OffsetSpec factory method |
| offsets = await admin.list_offsets( |
| table_path, |
| bucket_ids=[0], |
| offset_spec=fluss.OffsetSpec.latest() |
| ) |
| print(f"Latest offsets for table (before writes): {offsets}") |
| except Exception as e: |
| print(f"Failed to list offsets: {e}") |
| |
| # Get the table instance |
| table = await conn.get_table(table_path) |
| print(f"Got table: {table}") |
| |
| # Create a writer for the table |
| append_writer = table.new_append().create_writer() |
| print(f"Created append writer: {append_writer}") |
| |
| try: |
| # Demo: Write PyArrow Table |
| print("\n--- Testing PyArrow Table write ---") |
| pa_table = pa.Table.from_arrays( |
| [ |
| pa.array([1, 2, 3], type=pa.int32()), |
| pa.array(["Alice", "Bob", "Charlie"], type=pa.string()), |
| pa.array([95.2, 87.2, 92.1], type=pa.float32()), |
| pa.array([25, 30, 35], type=pa.int32()), |
| pa.array( |
| [date(1999, 5, 15), date(1994, 3, 20), date(1989, 11, 8)], |
| type=pa.date32(), |
| ), |
| pa.array( |
| [dt_time(9, 0, 0), dt_time(9, 30, 0), dt_time(10, 0, 0)], |
| type=pa.time32("ms"), |
| ), |
| pa.array( |
| [ |
| datetime(2024, 1, 15, 10, 30), |
| datetime(2024, 1, 15, 11, 0), |
| datetime(2024, 1, 15, 11, 30), |
| ], |
| type=pa.timestamp("us"), |
| ), |
| pa.array( |
| [ |
| datetime(2024, 1, 15, 10, 30), |
| datetime(2024, 1, 15, 11, 0), |
| datetime(2024, 1, 15, 11, 30), |
| ], |
| type=pa.timestamp("us", tz="UTC"), |
| ), |
| pa.array( |
| [Decimal("75000.00"), Decimal("82000.50"), Decimal("95000.75")], |
| type=pa.decimal128(10, 2), |
| ), |
| ], |
| schema=schema, |
| ) |
| |
| append_writer.write_arrow(pa_table) |
| print("Successfully wrote PyArrow Table") |
| |
| # Demo: Write PyArrow RecordBatch |
| print("\n--- Testing PyArrow RecordBatch write ---") |
| pa_record_batch = pa.RecordBatch.from_arrays( |
| [ |
| pa.array([4, 5], type=pa.int32()), |
| pa.array(["David", "Eve"], type=pa.string()), |
| pa.array([88.5, 91.0], type=pa.float32()), |
| pa.array([28, 32], type=pa.int32()), |
| pa.array([date(1996, 7, 22), date(1992, 12, 1)], type=pa.date32()), |
| pa.array([dt_time(14, 15, 0), dt_time(8, 45, 0)], type=pa.time32("ms")), |
| pa.array( |
| [datetime(2024, 1, 16, 9, 0), datetime(2024, 1, 16, 9, 30)], |
| type=pa.timestamp("us"), |
| ), |
| pa.array( |
| [datetime(2024, 1, 16, 9, 0), datetime(2024, 1, 16, 9, 30)], |
| type=pa.timestamp("us", tz="UTC"), |
| ), |
| pa.array( |
| [Decimal("68000.00"), Decimal("72500.25")], |
| type=pa.decimal128(10, 2), |
| ), |
| ], |
| schema=schema, |
| ) |
| |
| append_writer.write_arrow_batch(pa_record_batch) |
| print("Successfully wrote PyArrow RecordBatch") |
| |
| # Test 3: Append single rows with Date, Time, Timestamp, Decimal |
| print("\n--- Testing single row append with temporal/decimal types ---") |
| # Dict input with all types including Date, Time, Timestamp, Decimal |
| append_writer.append( |
| { |
| "id": 8, |
| "name": "Helen", |
| "score": 93.5, |
| "age": 26, |
| "birth_date": date(1998, 4, 10), |
| "check_in_time": dt_time(11, 30, 45), |
| "created_at": datetime(2024, 1, 17, 14, 0, 0), |
| "updated_at": datetime(2024, 1, 17, 14, 0, 0), |
| "salary": Decimal("88000.00"), |
| } |
| ) |
| print("Successfully appended row (dict with Date, Time, Timestamp, Decimal)") |
| |
| # List input with all types |
| append_writer.append( |
| [ |
| 9, |
| "Ivan", |
| 90.0, |
| 31, |
| date(1993, 8, 25), |
| dt_time(16, 45, 0), |
| datetime(2024, 1, 17, 15, 30, 0), |
| datetime(2024, 1, 17, 15, 30, 0), |
| Decimal("91500.50"), |
| ] |
| ) |
| print("Successfully appended row (list with Date, Time, Timestamp, Decimal)") |
| |
| # Demo: Write Pandas DataFrame |
| print("\n--- Testing Pandas DataFrame write ---") |
| df = pd.DataFrame( |
| { |
| "id": [10, 11], |
| "name": ["Frank", "Grace"], |
| "score": [89.3, 94.7], |
| "age": [29, 27], |
| "birth_date": [date(1995, 2, 14), date(1997, 9, 30)], |
| "check_in_time": [dt_time(10, 0, 0), dt_time(10, 30, 0)], |
| "created_at": [ |
| datetime(2024, 1, 18, 8, 0), |
| datetime(2024, 1, 18, 8, 30), |
| ], |
| "updated_at": [ |
| datetime(2024, 1, 18, 8, 0), |
| datetime(2024, 1, 18, 8, 30), |
| ], |
| "salary": [Decimal("79000.00"), Decimal("85500.75")], |
| } |
| ) |
| |
| append_writer.write_pandas(df) |
| print("Successfully wrote Pandas DataFrame") |
| |
| # Flush all pending data |
| print("\n--- Flushing data ---") |
| await append_writer.flush() |
| print("Successfully flushed data") |
| |
| # Demo: Check offsets after writes |
| print("\n--- Checking offsets after writes ---") |
| try: |
| offsets = await admin.list_offsets( |
| table_path, |
| bucket_ids=[0], |
| offset_spec=fluss.OffsetSpec.latest() |
| ) |
| print(f"Latest offsets after writing 7 records: {offsets}") |
| except Exception as e: |
| print(f"Failed to list offsets: {e}") |
| |
| except Exception as e: |
| print(f"Error during writing: {e}") |
| |
| # Now scan the table to verify data was written |
| print("\n--- Scanning table (batch scanner) ---") |
| try: |
| # Use new_scan().create_record_batch_log_scanner() for batch-based operations |
| batch_scanner = await table.new_scan().create_record_batch_log_scanner() |
| print(f"Created batch scanner: {batch_scanner}") |
| |
| # Subscribe to buckets (required before to_arrow/to_pandas) |
| # Use subscribe_buckets to subscribe all buckets from EARLIEST_OFFSET |
| 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)}) |
| print(f"Subscribed to {num_buckets} buckets from EARLIEST_OFFSET") |
| |
| # Read all data using to_arrow() |
| print("Scanning results using to_arrow():") |
| |
| # Try to get as PyArrow Table |
| try: |
| pa_table_result = batch_scanner.to_arrow() |
| print(f"\nAs PyArrow Table: {pa_table_result}") |
| except Exception as e: |
| print(f"Could not convert to PyArrow: {e}") |
| |
| # Create a new batch scanner for to_pandas() test |
| batch_scanner2 = await table.new_scan().create_record_batch_log_scanner() |
| batch_scanner2.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| # Try to get as Pandas DataFrame |
| try: |
| df_result = batch_scanner2.to_pandas() |
| print(f"\nAs Pandas DataFrame:\n{df_result}") |
| except Exception as e: |
| print(f"Could not convert to Pandas: {e}") |
| |
| # TODO: support to_arrow_batch_reader() |
| # which is reserved for streaming use cases |
| |
| # TODO: support to_duckdb() |
| |
| # Test poll_arrow() method for incremental reading as Arrow Table |
| print("\n--- Testing poll_arrow() method ---") |
| batch_scanner3 = await table.new_scan().create_record_batch_log_scanner() |
| batch_scanner3.subscribe(bucket_id=0, start_offset=fluss.EARLIEST_OFFSET) |
| print(f"Subscribed to bucket 0 at EARLIEST_OFFSET ({fluss.EARLIEST_OFFSET})") |
| |
| # Poll with a timeout of 5000ms (5 seconds) |
| # Note: poll_arrow() returns an empty table (not an error) on timeout |
| try: |
| poll_result = batch_scanner3.poll_arrow(5000) |
| print(f"Number of rows: {poll_result.num_rows}") |
| |
| if poll_result.num_rows > 0: |
| poll_df = poll_result.to_pandas() |
| print(f"Polled data:\n{poll_df}") |
| else: |
| print("Empty result (no records available)") |
| # Empty table still has schema - this is useful! |
| print(f"Schema: {poll_result.schema}") |
| |
| except Exception as e: |
| print(f"Error during poll_arrow: {e}") |
| |
| # Test poll_record_batch() method for batches with metadata |
| print("\n--- Testing poll_record_batch() method ---") |
| batch_scanner4 = await table.new_scan().create_record_batch_log_scanner() |
| batch_scanner4.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| try: |
| batches = batch_scanner4.poll_record_batch(5000) |
| print(f"Number of batches: {len(batches)}") |
| |
| for i, batch in enumerate(batches): |
| print(f" Batch {i}: bucket={batch.bucket}, " |
| f"offsets={batch.base_offset}-{batch.last_offset}, " |
| f"rows={batch.batch.num_rows}") |
| |
| except Exception as e: |
| print(f"Error during poll_record_batch: {e}") |
| |
| except Exception as e: |
| print(f"Error during batch scanning: {e}") |
| |
| # Test record-based scanning with poll() |
| print("\n--- Scanning table (record scanner) ---") |
| try: |
| # Use new_scan().create_log_scanner() for record-based operations |
| record_scanner = await table.new_scan().create_log_scanner() |
| print(f"Created record scanner: {record_scanner}") |
| |
| record_scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| |
| # Poll returns ScanRecords — records grouped by bucket |
| print("\n--- Testing poll() method (record-by-record) ---") |
| try: |
| scan_records = record_scanner.poll(5000) |
| print(f"Total records: {scan_records.count()}, buckets: {len(scan_records.buckets())}") |
| |
| # Flat iteration over all records (regardless of bucket) |
| print(f" Flat iteration: {scan_records.count()} records") |
| for record in scan_records: |
| print(f" offset={record.offset}, timestamp={record.timestamp}") |
| |
| # Per-bucket access |
| for bucket in scan_records.buckets(): |
| bucket_recs = scan_records.records(bucket) |
| print(f" Bucket {bucket}: {len(bucket_recs)} records") |
| for record in bucket_recs[:3]: |
| print(f" offset={record.offset}, " |
| f"timestamp={record.timestamp}, " |
| f"change_type={record.change_type}, " |
| f"row={record.row}") |
| |
| except Exception as e: |
| print(f"Error during poll: {e}") |
| |
| except Exception as e: |
| print(f"Error during record scanning: {e}") |
| |
| # Demo: unsubscribe — unsubscribe from a bucket (non-partitioned tables) |
| print("\n--- Testing unsubscribe ---") |
| try: |
| unsub_scanner = await table.new_scan().create_record_batch_log_scanner() |
| unsub_scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| print(f"Subscribed to {num_buckets} buckets") |
| # Unsubscribe from bucket 0 — future polls will skip this bucket |
| unsub_scanner.unsubscribe(bucket_id=0) |
| print("Unsubscribed from bucket 0") |
| remaining = unsub_scanner.poll_arrow(5000) |
| print(f"After unsubscribe, got {remaining.num_rows} records (from remaining buckets)") |
| except Exception as e: |
| print(f"Error during unsubscribe test: {e}") |
| |
| # ===================================================== |
| # Demo: Primary Key Table with Lookup and Upsert |
| # ===================================================== |
| print("\n" + "=" * 60) |
| print("--- Testing Primary Key Table (Lookup & Upsert) ---") |
| print("=" * 60) |
| |
| # Create a primary key table for lookup/upsert tests |
| # Include temporal and decimal types to test full conversion |
| pk_table_fields = [ |
| pa.field("user_id", pa.int32()), |
| pa.field("name", pa.string()), |
| pa.field("email", pa.string()), |
| pa.field("age", pa.int32()), |
| pa.field("birth_date", pa.date32()), |
| pa.field("login_time", pa.time32("ms")), |
| pa.field("created_at", pa.timestamp("us")), # TIMESTAMP (NTZ) |
| pa.field("updated_at", pa.timestamp("us", tz="UTC")), # TIMESTAMP_LTZ |
| pa.field("balance", pa.decimal128(10, 2)), |
| ] |
| pk_schema = pa.schema(pk_table_fields) |
| fluss_pk_schema = fluss.Schema(pk_schema, primary_keys=["user_id"]) |
| |
| # Create table descriptor |
| pk_table_descriptor = fluss.TableDescriptor( |
| fluss_pk_schema, |
| bucket_count=3, |
| ) |
| |
| pk_table_path = fluss.TablePath("fluss", "users_pk_table_v3") |
| |
| try: |
| await admin.create_table(pk_table_path, pk_table_descriptor, True) |
| print(f"Created PK table: {pk_table_path}") |
| except Exception as e: |
| print(f"PK Table creation failed (may already exist): {e}") |
| |
| # Get the PK table |
| pk_table = await conn.get_table(pk_table_path) |
| print(f"Got PK table: {pk_table}") |
| print(f"Has primary key: {pk_table.has_primary_key()}") |
| |
| # --- Test Upsert --- |
| print("\n--- Testing Upsert (fire-and-forget) ---") |
| try: |
| upsert_writer = pk_table.new_upsert().create_writer() |
| print(f"Created upsert writer: {upsert_writer}") |
| |
| # Fire-and-forget: queue writes synchronously, flush at end. |
| # Records are batched internally for efficiency. |
| upsert_writer.upsert( |
| { |
| "user_id": 1, |
| "name": "Alice", |
| "email": "alice@example.com", |
| "age": 25, |
| "birth_date": date(1999, 5, 15), |
| "login_time": dt_time(9, 30, 45, 123000), # 09:30:45.123 |
| "created_at": datetime( |
| 2024, 1, 15, 10, 30, 45, 123456 |
| ), # with microseconds |
| "updated_at": datetime(2024, 1, 15, 10, 30, 45, 123456), |
| "balance": Decimal("1234.56"), |
| } |
| ) |
| print("Queued user_id=1 (Alice)") |
| |
| upsert_writer.upsert( |
| { |
| "user_id": 2, |
| "name": "Bob", |
| "email": "bob@example.com", |
| "age": 30, |
| "birth_date": date(1994, 3, 20), |
| "login_time": dt_time(14, 15, 30, 500000), # 14:15:30.500 |
| "created_at": datetime(2024, 1, 16, 11, 22, 33, 444555), |
| "updated_at": datetime(2024, 1, 16, 11, 22, 33, 444555), |
| "balance": Decimal("5678.91"), |
| } |
| ) |
| print("Queued user_id=2 (Bob)") |
| |
| upsert_writer.upsert( |
| { |
| "user_id": 3, |
| "name": "Charlie", |
| "email": "charlie@example.com", |
| "age": 35, |
| "birth_date": date(1989, 11, 8), |
| "login_time": dt_time(16, 45, 59, 999000), # 16:45:59.999 |
| "created_at": datetime(2024, 1, 17, 23, 59, 59, 999999), |
| "updated_at": datetime(2024, 1, 17, 23, 59, 59, 999999), |
| "balance": Decimal("9876.54"), |
| } |
| ) |
| print("Queued user_id=3 (Charlie)") |
| |
| # flush() waits for all queued writes to be acknowledged by the server |
| await upsert_writer.flush() |
| print("Flushed — all 3 rows acknowledged by server") |
| |
| # Per-record acknowledgment: await the returned handle to block until |
| # the server confirms this specific write, useful when you need to |
| # read-after-write or verify critical updates. |
| print("\n--- Testing Upsert (per-record acknowledgment) ---") |
| handle = upsert_writer.upsert( |
| { |
| "user_id": 1, |
| "name": "Alice Updated", |
| "email": "alice.new@example.com", |
| "age": 26, |
| "birth_date": date(1999, 5, 15), |
| "login_time": dt_time(10, 11, 12, 345000), # 10:11:12.345 |
| "created_at": datetime(2024, 1, 15, 10, 30, 45, 123456), # unchanged |
| "updated_at": datetime( |
| 2024, 1, 20, 15, 45, 30, 678901 |
| ), # new update time |
| "balance": Decimal("2345.67"), |
| } |
| ) |
| await handle.wait() # wait for server acknowledgment |
| print("Updated user_id=1 (Alice -> Alice Updated) — server acknowledged") |
| |
| except Exception as e: |
| print(f"Error during upsert: {e}") |
| traceback.print_exc() |
| |
| # --- Test Lookup --- |
| print("\n--- Testing Lookup ---") |
| try: |
| lookuper = pk_table.new_lookup().create_lookuper() |
| print(f"Created lookuper: {lookuper}") |
| |
| result = await lookuper.lookup({"user_id": 1}) |
| if result: |
| print("Lookup user_id=1: Found!") |
| print(f" name: {result['name']}") |
| print(f" email: {result['email']}") |
| print(f" age: {result['age']}") |
| print( |
| f" birth_date: {result['birth_date']} (type: {type(result['birth_date']).__name__})" |
| ) |
| print( |
| f" login_time: {result['login_time']} (type: {type(result['login_time']).__name__})" |
| ) |
| print( |
| f" created_at: {result['created_at']} (type: {type(result['created_at']).__name__})" |
| ) |
| print( |
| f" updated_at: {result['updated_at']} (type: {type(result['updated_at']).__name__})" |
| ) |
| print( |
| f" balance: {result['balance']} (type: {type(result['balance']).__name__})" |
| ) |
| else: |
| print("Lookup user_id=1: Not found") |
| |
| # Lookup another row |
| result = await lookuper.lookup({"user_id": 2}) |
| if result: |
| print(f"Lookup user_id=2: Found! -> {result}") |
| else: |
| print("Lookup user_id=2: Not found") |
| |
| # Lookup non-existent row |
| result = await lookuper.lookup({"user_id": 999}) |
| if result: |
| print(f"Lookup user_id=999: Found! -> {result}") |
| else: |
| print("Lookup user_id=999: Not found (as expected)") |
| |
| except Exception as e: |
| print(f"Error during lookup: {e}") |
| traceback.print_exc() |
| |
| # --- Test Delete --- |
| print("\n--- Testing Delete ---") |
| try: |
| upsert_writer = pk_table.new_upsert().create_writer() |
| |
| handle = upsert_writer.delete({"user_id": 3}) |
| await handle.wait() |
| print("Deleted user_id=3 — server acknowledged") |
| |
| lookuper = pk_table.new_lookup().create_lookuper() |
| result = await lookuper.lookup({"user_id": 3}) |
| if result: |
| print(f"Lookup user_id=3 after delete: Still found! -> {result}") |
| else: |
| print("Lookup user_id=3 after delete: Not found (deletion confirmed)") |
| |
| except Exception as e: |
| print(f"Error during delete: {e}") |
| traceback.print_exc() |
| |
| # --- Test Partial Update by column names --- |
| print("\n--- Testing Partial Update (by column names) ---") |
| try: |
| partial_writer = pk_table.new_upsert().partial_update_by_name(["user_id", "balance"]).create_writer() |
| handle = partial_writer.upsert({"user_id": 1, "balance": Decimal("9999.99")}) |
| await handle.wait() |
| print("Partial update: set balance=9999.99 for user_id=1") |
| |
| lookuper = pk_table.new_lookup().create_lookuper() |
| result = await lookuper.lookup({"user_id": 1}) |
| if result: |
| print(f"Partial update verified:" |
| f"\n name={result['name']} (unchanged)" |
| f"\n balance={result['balance']} (updated)") |
| else: |
| print("ERROR: Expected to find user_id=1") |
| |
| except Exception as e: |
| print(f"Error during partial update by names: {e}") |
| traceback.print_exc() |
| |
| # --- Test Partial Update by column indices --- |
| print("\n--- Testing Partial Update (by column indices) ---") |
| try: |
| # Columns: 0=user_id (PK), 1=name — update name only |
| partial_writer_idx = pk_table.new_upsert().partial_update_by_index([0, 1]).create_writer() |
| handle = partial_writer_idx.upsert([1, "Alice Renamed"]) |
| await handle.wait() |
| print("Partial update by indices: set name='Alice Renamed' for user_id=1") |
| |
| lookuper = pk_table.new_lookup().create_lookuper() |
| result = await lookuper.lookup({"user_id": 1}) |
| if result: |
| print(f"Partial update by indices verified:" |
| f"\n name={result['name']} (updated)" |
| f"\n balance={result['balance']} (unchanged)") |
| else: |
| print("ERROR: Expected to find user_id=1") |
| |
| except Exception as e: |
| print(f"Error during partial update by indices: {e}") |
| traceback.print_exc() |
| |
| # Demo: Column projection using builder pattern |
| print("\n--- Testing Column Projection ---") |
| try: |
| # Get bucket count for subscriptions |
| num_buckets = (await admin.get_table_info(table_path)).num_buckets |
| |
| # Project specific columns by index (using batch scanner for to_pandas) |
| print("\n1. Projection by index [0, 1] (id, name):") |
| scanner_index = await table.new_scan().project([0, 1]).create_record_batch_log_scanner() |
| scanner_index.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| df_projected = scanner_index.to_pandas() |
| print(df_projected.head()) |
| print( |
| f" Projected {df_projected.shape[1]} columns: {list(df_projected.columns)}" |
| ) |
| |
| # Project specific columns by name (Pythonic!) |
| print("\n2. Projection by name ['name', 'score'] (Pythonic):") |
| scanner_names = await table.new_scan() \ |
| .project_by_name(["name", "score"]) \ |
| .create_record_batch_log_scanner() |
| scanner_names.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| df_named = scanner_names.to_pandas() |
| print(df_named.head()) |
| print(f" Projected {df_named.shape[1]} columns: {list(df_named.columns)}") |
| |
| # Test empty result schema with projection |
| print("\n3. Testing empty result schema with projection:") |
| scanner_proj = await table.new_scan().project([0, 2]).create_record_batch_log_scanner() |
| scanner_proj.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}) |
| # Quick poll that may return empty |
| result = scanner_proj.poll_arrow(100) |
| print(f" Schema columns: {result.schema.names}") |
| |
| except Exception as e: |
| print(f"Error during projection: {e}") |
| |
| # Demo: Drop tables |
| print("\n--- Testing drop_table() ---") |
| try: |
| # Drop the log table |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| print(f"Successfully dropped table: {table_path}") |
| # Drop the PK table |
| await admin.drop_table(pk_table_path, ignore_if_not_exists=True) |
| print(f"Successfully dropped table: {pk_table_path}") |
| except Exception as e: |
| print(f"Failed to drop table: {e}") |
| |
| # ===================================================== |
| # Demo: Partitioned Table with list_partition_offsets |
| # ===================================================== |
| print("\n" + "=" * 60) |
| print("--- Testing Partitioned Table ---") |
| print("=" * 60) |
| |
| # Create a partitioned log table |
| partitioned_fields = [ |
| pa.field("id", pa.int32()), |
| pa.field("region", pa.string()), # partition key |
| pa.field("value", pa.int64()), |
| ] |
| partitioned_schema = pa.schema(partitioned_fields) |
| fluss_partitioned_schema = fluss.Schema(partitioned_schema) |
| |
| partitioned_table_descriptor = fluss.TableDescriptor( |
| fluss_partitioned_schema, |
| partition_keys=["region"], # Partition by region |
| bucket_count=1, |
| ) |
| |
| partitioned_table_path = fluss.TablePath("fluss", "partitioned_log_table_py") |
| |
| try: |
| # Drop if exists first |
| await admin.drop_table(partitioned_table_path, ignore_if_not_exists=True) |
| print(f"Dropped existing table: {partitioned_table_path}") |
| |
| # Create the partitioned table |
| await admin.create_table(partitioned_table_path, partitioned_table_descriptor, False) |
| print(f"Created partitioned table: {partitioned_table_path}") |
| |
| # Create partitions for US and EU regions |
| print("\n--- Creating partitions ---") |
| await admin.create_partition(partitioned_table_path, {"region": "US"}, ignore_if_exists=True) |
| print("Created partition: region=US") |
| await admin.create_partition(partitioned_table_path, {"region": "EU"}, ignore_if_exists=True) |
| print("Created partition: region=EU") |
| |
| # List partitions |
| print("\n--- Listing partitions ---") |
| partition_infos = await admin.list_partition_infos(partitioned_table_path) |
| for p in partition_infos: |
| print(f" {p}") # PartitionInfo(partition_id=..., partition_name='region=...') |
| |
| # Get the table and write some data |
| partitioned_table = await conn.get_table(partitioned_table_path) |
| partitioned_writer = partitioned_table.new_append().create_writer() |
| |
| # Append data to US partition |
| partitioned_writer.append({"id": 1, "region": "US", "value": 100}) |
| partitioned_writer.append({"id": 2, "region": "US", "value": 200}) |
| # Append data to EU partition |
| partitioned_writer.append({"id": 3, "region": "EU", "value": 300}) |
| partitioned_writer.append({"id": 4, "region": "EU", "value": 400}) |
| await partitioned_writer.flush() |
| print("\nWrote 4 records (2 to US, 2 to EU)") |
| |
| # Demo: list_partition_infos with partial spec filter |
| print("\n--- Testing list_partition_infos with spec ---") |
| us_partitions = await admin.list_partition_infos( |
| partitioned_table_path, partition_spec={"region": "US"} |
| ) |
| print(f"Filtered partitions (region=US): {us_partitions}") |
| |
| # Demo: list_partition_offsets |
| print("\n--- Testing list_partition_offsets ---") |
| |
| # Query offsets for US partition |
| # Note: partition_name is just the value (e.g., "US"), not "region=US" |
| us_offsets = await admin.list_partition_offsets( |
| partitioned_table_path, |
| partition_name="US", |
| bucket_ids=[0], |
| offset_spec=fluss.OffsetSpec.latest() |
| ) |
| print(f"US partition latest offsets: {us_offsets}") |
| |
| # Query offsets for EU partition |
| eu_offsets = await admin.list_partition_offsets( |
| partitioned_table_path, |
| partition_name="EU", |
| bucket_ids=[0], |
| offset_spec=fluss.OffsetSpec.latest() |
| ) |
| print(f"EU partition latest offsets: {eu_offsets}") |
| |
| # Demo: subscribe_partition for reading partitioned data |
| print("\n--- Testing subscribe_partition + to_arrow() ---") |
| partitioned_scanner = await partitioned_table.new_scan().create_record_batch_log_scanner() |
| |
| # Subscribe to each partition using partition_id |
| for p in partition_infos: |
| partitioned_scanner.subscribe_partition( |
| partition_id=p.partition_id, |
| bucket_id=0, |
| start_offset=fluss.EARLIEST_OFFSET |
| ) |
| print(f"Subscribed to partition {p.partition_name} (id={p.partition_id})") |
| |
| # Use to_arrow() - now works for partitioned tables! |
| partitioned_arrow = partitioned_scanner.to_arrow() |
| print(f"\nto_arrow() returned {partitioned_arrow.num_rows} records from partitioned table:") |
| print(partitioned_arrow.to_pandas()) |
| |
| # Demo: subscribe_partition_buckets for batch subscribing to multiple partitions at once |
| print("\n--- Testing subscribe_partition_buckets + to_arrow() ---") |
| partitioned_scanner_batch = await partitioned_table.new_scan().create_record_batch_log_scanner() |
| partition_bucket_offsets = { |
| (p.partition_id, 0): fluss.EARLIEST_OFFSET for p in partition_infos |
| } |
| partitioned_scanner_batch.subscribe_partition_buckets(partition_bucket_offsets) |
| print(f"Batch subscribed to {len(partition_bucket_offsets)} partition+bucket combinations") |
| partitioned_batch_arrow = partitioned_scanner_batch.to_arrow() |
| print(f"to_arrow() returned {partitioned_batch_arrow.num_rows} records:") |
| print(partitioned_batch_arrow.to_pandas()) |
| |
| # Demo: unsubscribe_partition - unsubscribe from one partition, read remaining |
| print("\n--- Testing unsubscribe_partition ---") |
| partitioned_scanner3 = await partitioned_table.new_scan().create_record_batch_log_scanner() |
| for p in partition_infos: |
| partitioned_scanner3.subscribe_partition(p.partition_id, 0, fluss.EARLIEST_OFFSET) |
| # Unsubscribe from the first partition |
| first_partition = partition_infos[0] |
| partitioned_scanner3.unsubscribe_partition(first_partition.partition_id, 0) |
| print(f"Unsubscribed from partition {first_partition.partition_name} (id={first_partition.partition_id})") |
| remaining_arrow = partitioned_scanner3.to_arrow() |
| print(f"After unsubscribe, to_arrow() returned {remaining_arrow.num_rows} records (from remaining partitions):") |
| print(remaining_arrow.to_pandas()) |
| |
| # Demo: to_pandas() also works for partitioned tables |
| print("\n--- Testing to_pandas() on partitioned table ---") |
| partitioned_scanner2 = await partitioned_table.new_scan().create_record_batch_log_scanner() |
| for p in partition_infos: |
| partitioned_scanner2.subscribe_partition(p.partition_id, 0, fluss.EARLIEST_OFFSET) |
| partitioned_df = partitioned_scanner2.to_pandas() |
| print(f"to_pandas() returned {len(partitioned_df)} records:") |
| print(partitioned_df) |
| |
| # Cleanup |
| await admin.drop_table(partitioned_table_path, ignore_if_not_exists=True) |
| print(f"\nDropped partitioned table: {partitioned_table_path}") |
| |
| except Exception as e: |
| print(f"Error with partitioned table: {e}") |
| traceback.print_exc() |
| |
| # ===================================================== |
| # Demo: Partitioned KV Table (Upsert, Lookup, Delete) |
| # ===================================================== |
| print("\n" + "=" * 60) |
| print("--- Testing Partitioned KV Table ---") |
| print("=" * 60) |
| |
| partitioned_kv_fields = [ |
| pa.field("region", pa.string()), # partition key + part of PK |
| pa.field("user_id", pa.int32()), # part of PK |
| pa.field("name", pa.string()), |
| pa.field("score", pa.int64()), |
| ] |
| partitioned_kv_schema = pa.schema(partitioned_kv_fields) |
| fluss_partitioned_kv_schema = fluss.Schema( |
| partitioned_kv_schema, primary_keys=["region", "user_id"] |
| ) |
| |
| partitioned_kv_descriptor = fluss.TableDescriptor( |
| fluss_partitioned_kv_schema, |
| partition_keys=["region"], |
| ) |
| |
| partitioned_kv_path = fluss.TablePath("fluss", "partitioned_kv_table_py") |
| |
| try: |
| await admin.drop_table(partitioned_kv_path, ignore_if_not_exists=True) |
| await admin.create_table(partitioned_kv_path, partitioned_kv_descriptor, False) |
| print(f"Created partitioned KV table: {partitioned_kv_path}") |
| |
| # Create partitions |
| await admin.create_partition(partitioned_kv_path, {"region": "US"}) |
| await admin.create_partition(partitioned_kv_path, {"region": "EU"}) |
| await admin.create_partition(partitioned_kv_path, {"region": "APAC"}) |
| print("Created partitions: US, EU, APAC") |
| |
| partitioned_kv_table = await conn.get_table(partitioned_kv_path) |
| upsert_writer = partitioned_kv_table.new_upsert().create_writer() |
| |
| # Upsert rows across partitions |
| test_data = [ |
| ("US", 1, "Gustave", 100), |
| ("US", 2, "Lune", 200), |
| ("EU", 1, "Sciel", 150), |
| ("EU", 2, "Maelle", 250), |
| ("APAC", 1, "Noco", 300), |
| ] |
| for region, user_id, name, score in test_data: |
| upsert_writer.upsert({ |
| "region": region, "user_id": user_id, |
| "name": name, "score": score, |
| }) |
| await upsert_writer.flush() |
| print(f"Upserted {len(test_data)} rows across 3 partitions") |
| |
| # Lookup all rows across partitions |
| print("\n--- Lookup across partitions ---") |
| lookuper = partitioned_kv_table.new_lookup().create_lookuper() |
| for region, user_id, name, score in test_data: |
| result = await lookuper.lookup({"region": region, "user_id": user_id}) |
| assert result is not None, f"Expected to find region={region} user_id={user_id}" |
| assert result["name"] == name, f"Name mismatch: {result['name']} != {name}" |
| assert result["score"] == score, f"Score mismatch: {result['score']} != {score}" |
| print(f"All {len(test_data)} rows verified across partitions") |
| |
| # Update within a partition |
| print("\n--- Update within partition ---") |
| handle = upsert_writer.upsert({ |
| "region": "US", "user_id": 1, |
| "name": "Gustave Updated", "score": 999, |
| }) |
| await handle.wait() |
| result = await lookuper.lookup({"region": "US", "user_id": 1}) |
| assert result is not None, "Expected to find region=US user_id=1 after update" |
| assert result["name"] == "Gustave Updated" |
| assert result["score"] == 999 |
| print(f"Update verified: US/1 name={result['name']} score={result['score']}") |
| |
| # Lookup in non-existent partition |
| print("\n--- Lookup in non-existent partition ---") |
| result = await lookuper.lookup({"region": "UNKNOWN", "user_id": 1}) |
| assert result is None, "Expected UNKNOWN partition lookup to return None" |
| print("UNKNOWN partition lookup: not found (expected)") |
| |
| # Delete within a partition |
| print("\n--- Delete within partition ---") |
| handle = upsert_writer.delete({"region": "EU", "user_id": 1}) |
| await handle.wait() |
| result = await lookuper.lookup({"region": "EU", "user_id": 1}) |
| assert result is None, "Expected EU/1 to be deleted" |
| print("Delete verified: EU/1 not found") |
| |
| # Verify sibling record still exists |
| result = await lookuper.lookup({"region": "EU", "user_id": 2}) |
| assert result is not None, "Expected EU/2 to still exist" |
| assert result["name"] == "Maelle" |
| print(f"EU/2 still exists: name={result['name']}") |
| |
| # Cleanup |
| await admin.drop_table(partitioned_kv_path, ignore_if_not_exists=True) |
| print(f"\nDropped partitioned KV table: {partitioned_kv_path}") |
| |
| except Exception as e: |
| print(f"Error with partitioned KV table: {e}") |
| traceback.print_exc() |
| |
| # Close connection |
| conn.close() |
| print("\nConnection closed") |
| |
| |
| if __name__ == "__main__": |
| # Run the async main function |
| asyncio.run(main()) |
