Skip to content

feat(conversion): add LTable and KTable decoders for list and dict bindings #767

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 74 additions & 4 deletions python/cocoindex/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,10 +92,14 @@ def make_engine_value_decoder(
if src_type_kind == "Struct":
return _make_engine_struct_to_dict_decoder(field_path, src_type["fields"])
if src_type_kind in TABLE_TYPES:
raise ValueError(
f"Missing type annotation for `{''.join(field_path)}`."
f"It's required for {src_type_kind} type."
)
if src_type_kind == "LTable":
return _make_engine_ltable_to_list_dict_decoder(
field_path, src_type["row"]["fields"]
)
elif src_type_kind == "KTable":
return _make_engine_ktable_to_dict_dict_decoder(
field_path, src_type["row"]["fields"]
)
return lambda value: value

# Handle struct -> dict binding for explicit dict annotations
Expand Down Expand Up @@ -340,6 +344,72 @@ def decode_to_dict(values: list[Any] | None) -> dict[str, Any] | None:
return decode_to_dict


def _make_engine_ltable_to_list_dict_decoder(
field_path: list[str],
src_fields: list[dict[str, Any]],
) -> Callable[[list[Any] | None], list[dict[str, Any]] | None]:
"""Make a decoder from engine LTable values to a list of dicts."""

# Create a decoder for each row (struct) to dict
row_decoder = _make_engine_struct_to_dict_decoder(field_path, src_fields)

def decode_to_list_dict(values: list[Any] | None) -> list[dict[str, Any]] | None:
if values is None:
return None
result = []
for row_values in values:
decoded_row = row_decoder(row_values)
if decoded_row is not None:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we may raise an exception if a row is None

result.append(decoded_row)
return result

return decode_to_list_dict


def _make_engine_ktable_to_dict_dict_decoder(
field_path: list[str],
src_fields: list[dict[str, Any]],
) -> Callable[[list[Any] | None], dict[Any, dict[str, Any]] | None]:
"""Make a decoder from engine KTable values to a dict of dicts."""

if not src_fields:
raise ValueError("KTable must have at least one field for the key")

# First field is the key, remaining fields are the value
key_field_schema = src_fields[0]
value_fields_schema = src_fields[1:]

# Create decoders
field_path.append(f".{key_field_schema.get('name', KEY_FIELD_NAME)}")
key_decoder = make_engine_value_decoder(field_path, key_field_schema["type"], Any)
field_path.pop()

value_decoder = _make_engine_struct_to_dict_decoder(field_path, value_fields_schema)

def decode_to_dict_dict(
values: list[Any] | None,
) -> dict[Any, dict[str, Any]] | None:
if values is None:
return None
result = {}
for row_values in values:
if len(row_values) < 2:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually a key-only KTable is valid. In this case the value can be an empty dict.

raise ValueError(
f"KTable row must have at least 2 values (key + value), got {len(row_values)}"
)
key = key_decoder(row_values[0])
value = value_decoder(row_values[1:])
# Handle case where key is a dict (from struct key) - convert to tuple
if isinstance(key, dict):
key = tuple(key.values())
# Skip None values to maintain type consistency
if value is not None:
result[key] = value
return result

return decode_to_dict_dict


def dump_engine_object(v: Any) -> Any:
"""Recursively dump an object for engine. Engine side uses `Pythonized` to catch."""
if v is None:
Expand Down
127 changes: 127 additions & 0 deletions python/cocoindex/tests/test_convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -1341,3 +1341,130 @@ class Point(NamedTuple):
validate_full_roundtrip(
instance, Point, (expected_dict, dict), (expected_dict, Any)
)


def test_roundtrip_ltable_to_list_dict_binding() -> None:
"""Test LTable -> list[dict] binding with Any annotation."""

@dataclass
class User:
id: str
name: str
age: int

users = [User("u1", "Alice", 25), User("u2", "Bob", 30), User("u3", "Charlie", 35)]
expected_list_dict = [
{"id": "u1", "name": "Alice", "age": 25},
{"id": "u2", "name": "Bob", "age": 30},
{"id": "u3", "name": "Charlie", "age": 35},
]

# Test Any annotation
validate_full_roundtrip(users, list[User], (expected_list_dict, Any))


def test_roundtrip_ktable_to_dict_dict_binding() -> None:
"""Test KTable -> dict[K, dict] binding with Any annotation."""

@dataclass
class Product:
name: str
price: float
active: bool

products = {
"p1": Product("Widget", 29.99, True),
"p2": Product("Gadget", 49.99, False),
"p3": Product("Tool", 19.99, True),
}
expected_dict_dict = {
"p1": {"name": "Widget", "price": 29.99, "active": True},
"p2": {"name": "Gadget", "price": 49.99, "active": False},
"p3": {"name": "Tool", "price": 19.99, "active": True},
}

# Test Any annotation
validate_full_roundtrip(products, dict[str, Product], (expected_dict_dict, Any))


def test_roundtrip_ktable_with_complex_key() -> None:
"""Test KTable with complex key types -> dict binding."""

@dataclass(frozen=True)
class OrderKey:
shop_id: str
version: int

@dataclass
class Order:
customer: str
total: float

orders = {
OrderKey("shop1", 1): Order("Alice", 100.0),
OrderKey("shop2", 2): Order("Bob", 200.0),
}
expected_dict_dict = {
("shop1", 1): {"customer": "Alice", "total": 100.0},
("shop2", 2): {"customer": "Bob", "total": 200.0},
}

# Test Any annotation
validate_full_roundtrip(orders, dict[OrderKey, Order], (expected_dict_dict, Any))


def test_roundtrip_ltable_with_nested_structs() -> None:
"""Test LTable with nested structs -> list[dict] binding."""

@dataclass
class Address:
street: str
city: str

@dataclass
class Person:
name: str
age: int
address: Address

people = [
Person("John", 30, Address("123 Main St", "Anytown")),
Person("Jane", 25, Address("456 Oak Ave", "Somewhere")),
]
expected_list_dict = [
{
"name": "John",
"age": 30,
"address": {"street": "123 Main St", "city": "Anytown"},
},
{
"name": "Jane",
"age": 25,
"address": {"street": "456 Oak Ave", "city": "Somewhere"},
},
]

# Test Any annotation
validate_full_roundtrip(people, list[Person], (expected_list_dict, Any))


def test_roundtrip_ktable_with_list_fields() -> None:
"""Test KTable with list fields -> dict binding."""

@dataclass
class Team:
name: str
members: list[str]
active: bool

teams = {
"team1": Team("Dev Team", ["Alice", "Bob"], True),
"team2": Team("QA Team", ["Charlie", "David"], False),
}
expected_dict_dict = {
"team1": {"name": "Dev Team", "members": ["Alice", "Bob"], "active": True},
"team2": {"name": "QA Team", "members": ["Charlie", "David"], "active": False},
}

# Test Any annotation
validate_full_roundtrip(teams, dict[str, Team], (expected_dict_dict, Any))
Loading