| # 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 KV (primary key) table operations. |
| |
| Mirrors the Rust integration tests in crates/fluss/tests/integration/kv_table.rs. |
| """ |
| |
| import math |
| from datetime import date, datetime, timezone |
| from datetime import time as dt_time |
| from decimal import Decimal |
| |
| 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 _upsert_and_wait(writer, row): |
| handle = writer.upsert(row) |
| await handle.wait() |
| |
| |
| def _assert_float_specials(values): |
| assert math.isnan(values[0]) |
| assert math.isinf(values[1]) and values[1] > 0 |
| assert math.isinf(values[2]) and values[2] < 0 |
| |
| |
| async def test_upsert_delete_and_lookup(connection, admin): |
| """Test upsert, lookup, update, delete, and non-existent key lookup.""" |
| table_path = fluss.TablePath("fluss", "py_test_upsert_and_lookup") |
| 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()), |
| pa.field("age", pa.int64()), |
| ] |
| ), |
| primary_keys=["id"], |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| upsert_writer = table.new_upsert().create_writer() |
| |
| test_data = [(1, "Verso", 32), (2, "Noco", 25), (3, "Esquie", 35)] |
| |
| # Upsert rows (fire-and-forget, then flush) |
| for id_, name, age in test_data: |
| upsert_writer.upsert({"id": id_, "name": name, "age": age}) |
| await upsert_writer.flush() |
| |
| # Lookup and verify |
| lookuper = table.new_lookup().create_lookuper() |
| |
| for id_, expected_name, expected_age in test_data: |
| result = await lookuper.lookup({"id": id_}) |
| assert result is not None, f"Row with id={id_} should exist" |
| assert result["id"] == id_ |
| assert result["name"] == expected_name |
| assert result["age"] == expected_age |
| |
| # Update record with id=1 (await acknowledgment) |
| handle = upsert_writer.upsert({"id": 1, "name": "Verso", "age": 33}) |
| await handle.wait() |
| |
| result = await lookuper.lookup({"id": 1}) |
| assert result is not None |
| assert result["age"] == 33 |
| assert result["name"] == "Verso" |
| |
| # Delete record with id=1 (await acknowledgment) |
| handle = upsert_writer.delete({"id": 1}) |
| await handle.wait() |
| |
| result = await lookuper.lookup({"id": 1}) |
| assert result is None, "Record 1 should not exist after delete" |
| |
| # Verify other records still exist |
| for id_ in [2, 3]: |
| result = await lookuper.lookup({"id": id_}) |
| assert result is not None, f"Record {id_} should still exist" |
| |
| # Lookup non-existent key |
| result = await lookuper.lookup({"id": 999}) |
| assert result is None, "Non-existent key should return None" |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_composite_primary_keys(connection, admin): |
| """Test upsert/lookup with composite PKs, including prefix lookup.""" |
| table_path = fluss.TablePath("fluss", "py_test_composite_pk") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| # PK columns intentionally interleaved with non-PK column to verify |
| # that lookup correctly handles non-contiguous primary key indices. |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("region", pa.string()), |
| pa.field("score", pa.int64()), |
| pa.field("user_id", pa.int32()), |
| pa.field("event_id", pa.int64()), |
| ] |
| ), |
| primary_keys=["region", "user_id", "event_id"], |
| ) |
| table_descriptor = fluss.TableDescriptor( |
| schema, bucket_count=3, bucket_keys=["region", "user_id"] |
| ) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| upsert_writer = table.new_upsert().create_writer() |
| |
| test_data = [ |
| ("US", 1, 1, 100), |
| ("US", 1, 2, 200), |
| ("US", 2, 1, 300), |
| ("EU", 1, 1, 150), |
| ("EU", 2, 1, 250), |
| ] |
| |
| for region, user_id, event_id, score in test_data: |
| upsert_writer.upsert( |
| { |
| "region": region, |
| "user_id": user_id, |
| "event_id": event_id, |
| "score": score, |
| } |
| ) |
| await upsert_writer.flush() |
| |
| lookuper = table.new_lookup().create_lookuper() |
| |
| # Lookup (US, 1, 1) -> score 100 |
| result = await lookuper.lookup({"region": "US", "user_id": 1, "event_id": 1}) |
| assert result is not None |
| assert result["score"] == 100 |
| |
| # Lookup (EU, 2, 1) -> score 250 |
| result = await lookuper.lookup({"region": "EU", "user_id": 2, "event_id": 1}) |
| assert result is not None |
| assert result["score"] == 250 |
| |
| # Update (US, 1, 1) score (await acknowledgment) |
| handle = upsert_writer.upsert( |
| {"region": "US", "user_id": 1, "event_id": 1, "score": 500} |
| ) |
| await handle.wait() |
| |
| result = await lookuper.lookup({"region": "US", "user_id": 1, "event_id": 1}) |
| assert result is not None |
| assert result["score"] == 500 |
| |
| prefix_lookuper = ( |
| table.new_lookup().lookup_by(["region", "user_id"]).create_lookuper() |
| ) |
| |
| # Prefix (US, 1) should match 2 rows (event_id 1 and 2) |
| rows = await prefix_lookuper.lookup({"region": "US", "user_id": 1}) |
| assert len(rows) == 2 |
| event_ids = sorted(row["event_id"] for row in rows) |
| assert event_ids == [1, 2] |
| |
| # Also validate list/tuple prefix input |
| rows = await prefix_lookuper.lookup(["US", 1]) |
| assert len(rows) == 2 |
| rows = await prefix_lookuper.lookup(("EU", 2)) |
| assert len(rows) == 1 |
| assert rows[0]["event_id"] == 1 |
| |
| # Validate empty-result case: valid prefix shape but no matching rows. |
| rows = await prefix_lookuper.lookup({"region": "APAC", "user_id": 999}) |
| assert rows == [] |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_partial_update(connection, admin): |
| """Test partial column update via partial_update_by_name.""" |
| table_path = fluss.TablePath("fluss", "py_test_partial_update") |
| 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()), |
| pa.field("age", pa.int64()), |
| pa.field("score", pa.int64()), |
| ] |
| ), |
| primary_keys=["id"], |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| |
| # Insert initial record |
| upsert_writer = table.new_upsert().create_writer() |
| handle = upsert_writer.upsert({"id": 1, "name": "Verso", "age": 32, "score": 6942}) |
| await handle.wait() |
| |
| lookuper = table.new_lookup().create_lookuper() |
| result = await lookuper.lookup({"id": 1}) |
| assert result is not None |
| assert result["id"] == 1 |
| assert result["name"] == "Verso" |
| assert result["age"] == 32 |
| assert result["score"] == 6942 |
| |
| # Partial update: only update score column |
| partial_writer = ( |
| table.new_upsert().partial_update_by_name(["id", "score"]).create_writer() |
| ) |
| handle = partial_writer.upsert({"id": 1, "score": 420}) |
| await handle.wait() |
| |
| result = await lookuper.lookup({"id": 1}) |
| assert result is not None |
| assert result["id"] == 1 |
| assert result["name"] == "Verso", "name should remain unchanged" |
| assert result["age"] == 32, "age should remain unchanged" |
| assert result["score"] == 420, "score should be updated to 420" |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_partial_update_by_index(connection, admin): |
| """Test partial column update via partial_update_by_index.""" |
| table_path = fluss.TablePath("fluss", "py_test_partial_update_by_index") |
| 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()), |
| pa.field("age", pa.int64()), |
| pa.field("score", pa.int64()), |
| ] |
| ), |
| primary_keys=["id"], |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| |
| upsert_writer = table.new_upsert().create_writer() |
| handle = upsert_writer.upsert({"id": 1, "name": "Verso", "age": 32, "score": 6942}) |
| await handle.wait() |
| |
| # Partial update by indices: columns 0=id (PK), 1=name |
| partial_writer = table.new_upsert().partial_update_by_index([0, 1]).create_writer() |
| handle = partial_writer.upsert([1, "Verso Renamed"]) |
| await handle.wait() |
| |
| lookuper = table.new_lookup().create_lookuper() |
| result = await lookuper.lookup({"id": 1}) |
| assert result is not None |
| assert result["name"] == "Verso Renamed", "name should be updated" |
| assert result["score"] == 6942, "score should remain unchanged" |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_partial_update_complex(connection, admin): |
| """Partial updates preserve non-targeted ROW/MAP/ARRAY columns and update the |
| targeted ones (mirrors the Rust partial_update IT).""" |
| table_path = fluss.TablePath("fluss", "py_test_partial_update_complex") |
| 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()), |
| 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())), |
| ] |
| ), |
| primary_keys=["id"], |
| ) |
| await admin.create_table( |
| table_path, fluss.TableDescriptor(schema), ignore_if_exists=False |
| ) |
| table = await connection.get_table(table_path) |
| |
| await _upsert_and_wait( |
| table.new_upsert().create_writer(), |
| { |
| "id": 1, |
| "name": "Verso", |
| "nested": {"seq": 10, "label": "alpha"}, |
| "attrs": {"a": 1, "b": 2}, |
| "tags": ["alpha-tag", "beta-tag"], |
| }, |
| ) |
| |
| lookuper = table.new_lookup().create_lookuper() |
| |
| # Update only `name`: ROW/MAP/ARRAY columns must be preserved. |
| await _upsert_and_wait( |
| table.new_upsert().partial_update_by_name(["id", "name"]).create_writer(), |
| {"id": 1, "name": "Renamed"}, |
| ) |
| r = await lookuper.lookup({"id": 1}) |
| assert r["name"] == "Renamed" |
| assert r["nested"] == {"seq": 10, "label": "alpha"} |
| assert dict(r["attrs"]) == {"a": 1, "b": 2} |
| assert r["tags"] == ["alpha-tag", "beta-tag"] |
| |
| # Update only `nested`: scalar + other complex columns preserved. |
| await _upsert_and_wait( |
| table.new_upsert().partial_update_by_name(["id", "nested"]).create_writer(), |
| {"id": 1, "nested": {"seq": 99, "label": "omega"}}, |
| ) |
| r = await lookuper.lookup({"id": 1}) |
| assert r["nested"] == {"seq": 99, "label": "omega"} |
| assert r["name"] == "Renamed" |
| assert dict(r["attrs"]) == {"a": 1, "b": 2} |
| assert r["tags"] == ["alpha-tag", "beta-tag"] |
| |
| # Update `attrs` and `tags`. |
| await _upsert_and_wait( |
| table.new_upsert() |
| .partial_update_by_name(["id", "attrs", "tags"]) |
| .create_writer(), |
| {"id": 1, "attrs": {"z": 99}, "tags": ["gamma"]}, |
| ) |
| r = await lookuper.lookup({"id": 1}) |
| assert dict(r["attrs"]) == {"z": 99} |
| assert r["tags"] == ["gamma"] |
| assert r["nested"] == {"seq": 99, "label": "omega"} |
| assert r["name"] == "Renamed" |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_partitioned_table_upsert_and_lookup(connection, admin): |
| """Test upsert/lookup/delete on a partitioned KV table.""" |
| table_path = fluss.TablePath("fluss", "py_test_partitioned_kv_table") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("region", pa.string()), |
| pa.field("user_id", pa.int32()), |
| pa.field("name", pa.string()), |
| pa.field("score", pa.int64()), |
| ] |
| ), |
| primary_keys=["region", "user_id"], |
| ) |
| 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", "APAC"]: |
| await admin.create_partition( |
| table_path, {"region": region}, ignore_if_exists=True |
| ) |
| |
| table = await connection.get_table(table_path) |
| upsert_writer = table.new_upsert().create_writer() |
| |
| 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() |
| |
| lookuper = table.new_lookup().create_lookuper() |
| |
| # Verify all rows across partitions |
| for region, user_id, expected_name, expected_score in test_data: |
| result = await lookuper.lookup({"region": region, "user_id": user_id}) |
| assert result is not None, f"Row ({region}, {user_id}) should exist" |
| assert result["region"] == region |
| assert result["user_id"] == user_id |
| assert result["name"] == expected_name |
| assert result["score"] == expected_score |
| |
| # Update within a partition (await acknowledgment) |
| 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 |
| assert result["name"] == "Gustave Updated" |
| assert result["score"] == 999 |
| |
| # Lookup in non-existent partition should return None |
| result = await lookuper.lookup({"region": "UNKNOWN_REGION", "user_id": 1}) |
| assert result is None, "Lookup in non-existent partition should return None" |
| |
| # Delete within a partition (await acknowledgment) |
| 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, "Deleted record should not exist" |
| |
| # Verify sibling record still exists |
| result = await lookuper.lookup({"region": "EU", "user_id": 2}) |
| assert result is not None |
| assert result["name"] == "Maelle" |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_omit_nullable_vs_required_fields(connection, admin): |
| """Omitting a nullable field defaults it to null; omitting a non-nullable (or |
| PK) field is a clear client-side error -- at both top level and inside a ROW.""" |
| table_path = fluss.TablePath("fluss", "py_test_omit_fields") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| nested_type = pa.struct( |
| [ |
| pa.field("a", pa.int32()), # nullable |
| pa.field("b", pa.int32(), nullable=False), # NOT NULL |
| ] |
| ) |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("id", pa.int32()), |
| pa.field("opt", pa.int32()), # nullable |
| pa.field("req", pa.int32(), nullable=False), # NOT NULL |
| pa.field("nested", nested_type), |
| ] |
| ), |
| primary_keys=["id"], |
| ) |
| await admin.create_table( |
| table_path, fluss.TableDescriptor(schema), ignore_if_exists=False |
| ) |
| table = await connection.get_table(table_path) |
| writer = table.new_upsert().create_writer() |
| lookuper = table.new_lookup().create_lookuper() |
| |
| # Omit nullable top-level field `opt` -> Null. |
| await _upsert_and_wait(writer, {"id": 1, "req": 5, "nested": {"a": 1, "b": 2}}) |
| r = await lookuper.lookup({"id": 1}) |
| assert r["opt"] is None |
| assert r["req"] == 5 |
| assert r["nested"] == {"a": 1, "b": 2} |
| |
| # Omit nullable nested field `a` -> Null. |
| await _upsert_and_wait(writer, {"id": 2, "req": 9, "nested": {"b": 3}}) |
| r = await lookuper.lookup({"id": 2}) |
| assert r["nested"] == {"a": None, "b": 3} |
| |
| # Omit non-nullable top-level field `req` -> error. |
| with pytest.raises(Exception, match="non-nullable field 'req'"): |
| writer.upsert({"id": 3, "opt": 1, "nested": {"a": 1, "b": 2}}) |
| |
| # Omit non-nullable nested field `b` -> error. |
| with pytest.raises(Exception, match="non-nullable row field 'b'"): |
| writer.upsert({"id": 4, "req": 5, "nested": {"a": 1}}) |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_partitioned_complex(connection, admin): |
| """ROW/MAP/ARRAY columns round-trip and update correctly across partitions |
| (mirrors the complex columns in the Rust partitioned_table_upsert_and_lookup IT).""" |
| table_path = fluss.TablePath("fluss", "py_test_partitioned_complex") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("region", pa.string()), |
| pa.field("user_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())), |
| ] |
| ), |
| primary_keys=["region", "user_id"], |
| ) |
| await admin.create_table( |
| table_path, |
| fluss.TableDescriptor(schema, partition_keys=["region"]), |
| ignore_if_exists=False, |
| ) |
| for region in ["US", "EU"]: |
| await admin.create_partition( |
| table_path, {"region": region}, ignore_if_exists=True |
| ) |
| |
| table = await connection.get_table(table_path) |
| writer = table.new_upsert().create_writer() |
| await _upsert_and_wait( |
| writer, |
| { |
| "region": "US", |
| "user_id": 1, |
| "nested": {"seq": 11, "label": "a"}, |
| "attrs": {"x": 1}, |
| "tags": ["alpha"], |
| }, |
| ) |
| await _upsert_and_wait( |
| writer, |
| { |
| "region": "EU", |
| "user_id": 1, |
| "nested": {"seq": 22, "label": "b"}, |
| "attrs": {"y": 2}, |
| "tags": ["beta"], |
| }, |
| ) |
| |
| lookuper = table.new_lookup().create_lookuper() |
| r = await lookuper.lookup({"region": "US", "user_id": 1}) |
| assert r["nested"] == {"seq": 11, "label": "a"} |
| assert dict(r["attrs"]) == {"x": 1} |
| assert r["tags"] == ["alpha"] |
| r = await lookuper.lookup({"region": "EU", "user_id": 1}) |
| assert r["nested"] == {"seq": 22, "label": "b"} |
| assert dict(r["attrs"]) == {"y": 2} |
| assert r["tags"] == ["beta"] |
| |
| # Update complex columns within one partition; sibling partition unaffected. |
| await _upsert_and_wait( |
| writer, |
| { |
| "region": "US", |
| "user_id": 1, |
| "nested": {"seq": 99, "label": "updated"}, |
| "attrs": {"z": 99}, |
| "tags": ["renamed"], |
| }, |
| ) |
| r = await lookuper.lookup({"region": "US", "user_id": 1}) |
| assert r["nested"] == {"seq": 99, "label": "updated"} |
| assert dict(r["attrs"]) == {"z": 99} |
| assert r["tags"] == ["renamed"] |
| r = await lookuper.lookup({"region": "EU", "user_id": 1}) |
| assert r["nested"] == {"seq": 22, "label": "b"} |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_all_complex_datatypes(connection, admin): |
| """Comprehensive ARRAY/MAP/ROW coverage via upsert+lookup: full, edge, null rows. |
| |
| Mirrors the section-based ``all_supported_datatypes`` structure of the Rust |
| integration tests: the schema is composed from the shared array/row/map |
| sections, and three rows exercise fully-populated, edge-case, and all-null |
| values (see complex_types.py). |
| """ |
| table_path = fluss.TablePath("fluss", "py_test_kv_all_complex") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = complex_schema(primary_keys=["id"]) |
| await admin.create_table( |
| table_path, fluss.TableDescriptor(schema), ignore_if_exists=False |
| ) |
| |
| table = await connection.get_table(table_path) |
| upsert_writer = table.new_upsert().create_writer() |
| await _upsert_and_wait(upsert_writer, complex_full_row(1)) |
| await _upsert_and_wait(upsert_writer, complex_edge_row(2)) |
| await _upsert_and_wait(upsert_writer, complex_null_row(3)) |
| |
| lookuper = table.new_lookup().create_lookuper() |
| |
| r1 = await lookuper.lookup({"id": 1}) |
| assert r1 is not None |
| assert_complex_full(r1) |
| |
| r2 = await lookuper.lookup({"id": 2}) |
| assert r2 is not None |
| assert_complex_edge(r2) |
| |
| r3 = await lookuper.lookup({"id": 3}) |
| assert r3 is not None |
| for col in complex_column_names(): |
| assert r3[col] is None, f"{col} should be null" |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_all_supported_datatypes(connection, admin): |
| """Test upsert/lookup for all supported data types, including nulls.""" |
| table_path = fluss.TablePath("fluss", "py_test_kv_all_datatypes") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("pk_int", pa.int32()), |
| pa.field("col_boolean", pa.bool_()), |
| pa.field("col_tinyint", pa.int8()), |
| pa.field("col_smallint", pa.int16()), |
| pa.field("col_int", pa.int32()), |
| pa.field("col_bigint", pa.int64()), |
| pa.field("col_float", pa.float32()), |
| pa.field("col_double", pa.float64()), |
| pa.field("col_string", pa.string()), |
| pa.field("col_decimal", pa.decimal128(10, 2)), |
| pa.field("col_date", pa.date32()), |
| pa.field("col_time", pa.time32("ms")), |
| pa.field("col_timestamp_ntz", pa.timestamp("us")), |
| pa.field("col_timestamp_ltz", pa.timestamp("us", tz="UTC")), |
| pa.field("col_bytes", pa.binary()), |
| pa.field("col_array", pa.list_(pa.string())), |
| pa.field("col_binary", pa.binary(16)), |
| ] |
| ), |
| primary_keys=["pk_int"], |
| ) |
| table_descriptor = fluss.TableDescriptor(schema) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| upsert_writer = table.new_upsert().create_writer() |
| |
| # Test data for all types |
| row_data = { |
| "pk_int": 1, |
| "col_boolean": True, |
| "col_tinyint": 127, |
| "col_smallint": 32767, |
| "col_int": 2147483647, |
| "col_bigint": 9223372036854775807, |
| "col_float": 3.14, |
| "col_double": 2.718281828459045, |
| "col_string": "world of fluss python client", |
| "col_decimal": Decimal("123.45"), |
| "col_date": date(2026, 1, 23), |
| "col_time": dt_time(10, 13, 47, 123000), # millisecond precision |
| "col_timestamp_ntz": datetime(2026, 1, 23, 10, 13, 47, 123000), |
| "col_timestamp_ltz": datetime(2026, 1, 23, 10, 13, 47, 123000), |
| "col_bytes": b"binary data", |
| "col_array": ["fluss", "python"], |
| "col_binary": b"binary_data_0123", |
| } |
| |
| await _upsert_and_wait(upsert_writer, row_data) |
| |
| lookuper = table.new_lookup().create_lookuper() |
| result = await lookuper.lookup({"pk_int": 1}) |
| assert result is not None, "Row should exist" |
| |
| assert result["pk_int"] == 1 |
| assert result["col_boolean"] is True |
| assert result["col_tinyint"] == 127 |
| assert result["col_smallint"] == 32767 |
| assert result["col_int"] == 2147483647 |
| assert result["col_bigint"] == 9223372036854775807 |
| assert math.isclose(result["col_float"], 3.14, rel_tol=1e-6) |
| assert math.isclose(result["col_double"], 2.718281828459045, rel_tol=1e-15) |
| assert result["col_string"] == "world of fluss python client" |
| assert result["col_decimal"] == Decimal("123.45") |
| assert result["col_date"] == date(2026, 1, 23) |
| assert result["col_time"] == dt_time(10, 13, 47, 123000) |
| assert result["col_timestamp_ntz"] == datetime(2026, 1, 23, 10, 13, 47, 123000) |
| assert result["col_timestamp_ltz"] == datetime( |
| 2026, 1, 23, 10, 13, 47, 123000, tzinfo=timezone.utc |
| ) |
| assert result["col_bytes"] == b"binary data" |
| assert result["col_array"] == ["fluss", "python"] |
| assert result["col_binary"] == b"binary_data_0123" |
| |
| # Test with null values for all nullable columns |
| null_row = {"pk_int": 2} |
| for col in row_data: |
| if col != "pk_int": |
| null_row[col] = None |
| await _upsert_and_wait(upsert_writer, null_row) |
| |
| result = await lookuper.lookup({"pk_int": 2}) |
| assert result is not None, "Row with nulls should exist" |
| assert result["pk_int"] == 2 |
| for col in row_data: |
| if col != "pk_int": |
| assert result[col] is None, f"{col} should be null" |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| |
| async def test_prefix_lookup_validation_errors(connection, admin): |
| """Test that prefix lookup raises errors for invalid column configurations.""" |
| table_path = fluss.TablePath("fluss", "py_test_prefix_lookup_validation") |
| await admin.drop_table(table_path, ignore_if_not_exists=True) |
| |
| schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("a", pa.int32()), |
| pa.field("b", pa.string()), |
| pa.field("c", pa.int64()), |
| ] |
| ), |
| primary_keys=["a", "b", "c"], |
| ) |
| table_descriptor = fluss.TableDescriptor( |
| schema, bucket_count=3, bucket_keys=["a", "b"] |
| ) |
| await admin.create_table(table_path, table_descriptor, ignore_if_exists=False) |
| |
| table = await connection.get_table(table_path) |
| |
| # lookup_by with columns equal to full PK should error |
| with pytest.raises(fluss.FlussError, match="prefix lookup"): |
| table.new_lookup().lookup_by(["a", "b", "c"]).create_lookuper() |
| |
| # lookup_by with wrong column names should error |
| with pytest.raises(fluss.FlussError, match="bucket keys"): |
| table.new_lookup().lookup_by(["a", "c"]).create_lookuper() |
| |
| # lookup_by with unknown column should error |
| with pytest.raises(fluss.FlussError, match="Unknown column name"): |
| table.new_lookup().lookup_by(["a", "missing_col"]).create_lookuper() |
| |
| await admin.drop_table(table_path, ignore_if_not_exists=False) |
| |
| # Partitioned table: lookup columns must include partition keys first, |
| # followed by bucket keys. |
| partitioned_table_path = fluss.TablePath( |
| "fluss", "py_test_prefix_lookup_validation_pt" |
| ) |
| await admin.drop_table(partitioned_table_path, ignore_if_not_exists=True) |
| |
| partitioned_schema = fluss.Schema( |
| pa.schema( |
| [ |
| pa.field("region", pa.string()), |
| pa.field("user_id", pa.int32()), |
| pa.field("event_id", pa.int64()), |
| ] |
| ), |
| primary_keys=["region", "user_id", "event_id"], |
| ) |
| partitioned_table_descriptor = fluss.TableDescriptor( |
| partitioned_schema, |
| partition_keys=["region"], |
| bucket_count=3, |
| bucket_keys=["user_id"], |
| ) |
| await admin.create_table( |
| partitioned_table_path, partitioned_table_descriptor, ignore_if_exists=False |
| ) |
| |
| partitioned_table = await connection.get_table(partitioned_table_path) |
| |
| # Missing partition key in lookup columns. |
| with pytest.raises(fluss.FlussError, match="partition fields"): |
| partitioned_table.new_lookup().lookup_by(["user_id"]).create_lookuper() |
| |
| # A non-existent partition returns empty list. |
| partitioned_prefix_lookuper = ( |
| partitioned_table.new_lookup() |
| .lookup_by(["region", "user_id"]) |
| .create_lookuper() |
| ) |
| rows = await partitioned_prefix_lookuper.lookup( |
| {"region": "UNKNOWN_REGION", "user_id": 1} |
| ) |
| assert rows == [] |
| |
| # After partition keys, remaining columns must equal bucket keys. |
| with pytest.raises(fluss.FlussError, match="bucket keys"): |
| partitioned_table.new_lookup().lookup_by( |
| ["region", "event_id"] |
| ).create_lookuper() |
| |
| await admin.drop_table(partitioned_table_path, ignore_if_not_exists=False) |
