From ac8afa5ff6b80bf09da29e4b7a35a231b7d261be Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Thu, 8 Oct 2026 12:12:33 +0300 Subject: [PATCH 1/4] fix(type-inference)!: follow nested struct field references --- src/substrait/type_inference.py | 25 ++++++++--- tests/test_type_inference.py | 74 +++++++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+), 6 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 22ca03a..ea3cf05 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -645,12 +645,25 @@ def infer_expression_type( if reference_type == "direct_reference": segment = expression.selection.direct_reference - segment_reference_type = segment.WhichOneof("reference_type") - - if segment_reference_type == "struct_field": - return schema.types[segment.struct_field.field] - else: - raise Exception(f"Unknown reference_type {reference_type}") + while segment.WhichOneof("reference_type") == "struct_field": + field = segment.struct_field + if not 0 <= field.field < len(schema.types): + raise IndexError( + f"Struct field index {field.field} is out of range" + ) + result = schema.types[field.field] + if not field.HasField("child"): + return result + segment = field.child + if ( + segment.WhichOneof("reference_type") == "struct_field" + and result.WhichOneof("kind") != "struct" + ): + raise ValueError("Struct field child requires a struct") + schema = result.struct + raise Exception( + f"Unknown reference_type {segment.WhichOneof('reference_type')}" + ) else: raise Exception(f"Unknown reference_type {reference_type}") diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 696e988..16e3bc2 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -983,6 +983,80 @@ def test_infer_expression_type_selection(): assert result == expected +@pytest.mark.parametrize("root", ["row", "outer_steps", "outer_anchor", "lambda"]) +@pytest.mark.parametrize("depth", [1, 2]) +@pytest.mark.parametrize("nullable", [False, True]) +def test_infer_expression_type_nested_struct_selection(root, depth, nullable): + from substrait.type_inference import ( + _outer_anchor_binding, + lambda_scope, + outer_schemas, + ) + + expected = stt.Type(string=stt.Type.String(nullability=_NULL if nullable else _REQ)) + selected = expected + segment = stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField(field=1) + ) + for level in range(depth): + selected = stt.Type( + struct=stt.Type.Struct(types=[struct.types[0], selected], nullability=_REQ) + ) + segment = stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField( + field=0 if level == depth - 1 else 1, child=segment + ) + ) + row = stt.Type.Struct(types=[selected], nullability=_REQ) + ref = stalg.Expression.FieldReference(direct_reference=segment) + if root == "row": + ref.root_reference.SetInParent() + actual = infer_expression_type(stalg.Expression(selection=ref), row) + elif root == "lambda": + ref.lambda_parameter_reference.steps_out = 0 + with lambda_scope(row): + actual = infer_expression_type(stalg.Expression(selection=ref), struct) + elif root == "outer_anchor": + ref.outer_reference.rel_reference = 5 + with _outer_anchor_binding(5, row): + actual = infer_expression_type(stalg.Expression(selection=ref), struct) + else: + ref.outer_reference.steps_out = 1 + token = outer_schemas.set((stt.NamedStruct(struct=row),)) + try: + actual = infer_expression_type(stalg.Expression(selection=ref), struct) + finally: + outer_schemas.reset(token) + assert actual == expected + + +@pytest.mark.parametrize("index", [-1, 3]) +def test_nested_struct_selection_rejects_out_of_range_field(index): + row = stt.Type.Struct(types=[stt.Type(struct=struct)], nullability=_REQ) + expr = _field_reference(0) + expr.selection.direct_reference.struct_field.child.struct_field.field = index + with pytest.raises(IndexError, match="Struct field index .* is out of range"): + infer_expression_type(expr, row) + + +def test_nested_struct_selection_rejects_child_on_scalar(): + expr = _field_reference(0) + expr.selection.direct_reference.struct_field.child.struct_field.field = 0 + with pytest.raises(ValueError, match="Struct field child requires a struct"): + infer_expression_type(expr, struct) + + +def test_nested_struct_selection_rejects_unsupported_child_segment(): + expr = _field_reference(0) + expr.selection.direct_reference.struct_field.child.list_element.offset = 0 + row = stt.Type.Struct( + types=[stt.Type(list=stt.Type.List(type=struct.types[0], nullability=_REQ))], + nullability=_REQ, + ) + with pytest.raises(Exception, match="Unknown reference_type list_element"): + infer_expression_type(expr, row) + + def test_infer_expression_type_window_function(): """Test infer_expression_type with a window function expression.""" expr = stalg.Expression( From 19475ec3d74d6582654d0aa3ec6ca3d13a62662d Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Thu, 8 Oct 2026 16:49:42 +0300 Subject: [PATCH 2/4] fix(type-inference): resolve list and map reference children --- src/substrait/type_inference.py | 44 +++++++++++++++++------------ tests/dataframe/test_frame.py | 21 ++++++++++++++ tests/test_type_inference.py | 50 +++++++++++++++++++++++++++++---- 3 files changed, 92 insertions(+), 23 deletions(-) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index ea3cf05..4fd979b 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -645,25 +645,33 @@ def infer_expression_type( if reference_type == "direct_reference": segment = expression.selection.direct_reference - while segment.WhichOneof("reference_type") == "struct_field": - field = segment.struct_field - if not 0 <= field.field < len(schema.types): - raise IndexError( - f"Struct field index {field.field} is out of range" - ) - result = schema.types[field.field] - if not field.HasField("child"): + result = stt.Type(struct=schema) + while True: + kind = segment.WhichOneof("reference_type") + if kind == "struct_field": + if result.WhichOneof("kind") != "struct": + raise ValueError("Struct field child requires a struct") + child = segment.struct_field + if not 0 <= child.field < len(result.struct.types): + raise IndexError( + f"Struct field index {child.field} is out of range" + ) + result = result.struct.types[child.field] + elif kind == "list_element": + if result.WhichOneof("kind") != "list": + raise ValueError("List element reference requires a list") + child = segment.list_element + result = result.list.type + elif kind == "map_key": + if result.WhichOneof("kind") != "map": + raise ValueError("Map key reference requires a map") + child = segment.map_key + result = result.map.value + else: + raise Exception(f"Unknown reference_type {kind}") + if not child.HasField("child"): return result - segment = field.child - if ( - segment.WhichOneof("reference_type") == "struct_field" - and result.WhichOneof("kind") != "struct" - ): - raise ValueError("Struct field child requires a struct") - schema = result.struct - raise Exception( - f"Unknown reference_type {segment.WhichOneof('reference_type')}" - ) + segment = child.child else: raise Exception(f"Unknown reference_type {reference_type}") diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index 8502944..b9b5c03 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -79,6 +79,27 @@ def test_filter_select_matches_builder(): assert fluent.SerializeToString() == raw.SerializeToString() +@pytest.mark.parametrize("kind", ["list", "map"]) +def test_filter_collection_access_is_null(kind): + from substrait.builders.type import list as list_type + from substrait.builders.type import map as map_type + from substrait.type_inference import infer_plan_schema + + value_type = string() + collection = ( + list_type(value_type) if kind == "list" else map_type(string(), value_type) + ) + ns = named_struct( + names=["values"], struct=struct(types=[collection], nullable=False) + ) + df = sub.read_named_table("t", ns) + access = ( + sub.col("values")[0] if kind == "list" else sub.col("values").map_key("key") + ) + plan = df.filter(access.is_null()).select(access.alias("value")).to_plan() + assert infer_plan_schema(plan).struct.types[0] == value_type + + def test_with_columns_named_appends_projection(): fluent = people_df().with_columns(bonus=sub.col("age") + 1).to_plan() # ProjectRel appends: output has original columns + the new one. diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 16e3bc2..1683840 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -1046,15 +1046,55 @@ def test_nested_struct_selection_rejects_child_on_scalar(): infer_expression_type(expr, struct) -def test_nested_struct_selection_rejects_unsupported_child_segment(): +@pytest.mark.parametrize("kind", ["list_element", "map_key"]) +@pytest.mark.parametrize("nullable", [False, True]) +def test_nested_collection_selection(kind, nullable): expr = _field_reference(0) - expr.selection.direct_reference.struct_field.child.list_element.offset = 0 + child = expr.selection.direct_reference.struct_field.child + expected = stt.Type(string=stt.Type.String(nullability=_NULL if nullable else _REQ)) + if kind == "list_element": + child.list_element.offset = -1 + selected = stt.Type(list=stt.Type.List(type=expected, nullability=_REQ)) + else: + child.map_key.map_key.string = "key" + selected = stt.Type( + map=stt.Type.Map(key=struct.types[0], value=expected, nullability=_REQ) + ) + row = stt.Type.Struct(types=[selected], nullability=_REQ) + assert infer_expression_type(expr, row) == expected + + +def test_nested_collection_selection_continues_through_struct(): + expected = stt.Type(string=stt.Type.String(nullability=_NULL)) + value = stt.Type(struct=stt.Type.Struct(types=[expected], nullability=_REQ)) + mapping = stt.Type( + map=stt.Type.Map(key=struct.types[0], value=value, nullability=_REQ) + ) row = stt.Type.Struct( - types=[stt.Type(list=stt.Type.List(type=struct.types[0], nullability=_REQ))], + types=[stt.Type(list=stt.Type.List(type=mapping, nullability=_REQ))], nullability=_REQ, ) - with pytest.raises(Exception, match="Unknown reference_type list_element"): - infer_expression_type(expr, row) + expr = _field_reference(0) + element = expr.selection.direct_reference.struct_field.child.list_element + element.offset = 0 + key = element.child.map_key + key.map_key.string = "key" + key.child.struct_field.field = 0 + assert infer_expression_type(expr, row) == expected + + +@pytest.mark.parametrize( + "kind, message", + [ + ("list_element", "List element reference requires a list"), + ("map_key", "Map key reference requires a map"), + ], +) +def test_nested_collection_selection_rejects_wrong_container(kind, message): + expr = _field_reference(0) + getattr(expr.selection.direct_reference.struct_field.child, kind).SetInParent() + with pytest.raises(ValueError, match=message): + infer_expression_type(expr, struct) def test_infer_expression_type_window_function(): From b154ec09837aa80a243087a9c49c037c5379117a Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Thu, 8 Oct 2026 16:57:20 +0300 Subject: [PATCH 3/4] test(type-inference): use matching map key types --- tests/test_type_inference.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 1683840..8e826f1 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -1058,7 +1058,11 @@ def test_nested_collection_selection(kind, nullable): else: child.map_key.map_key.string = "key" selected = stt.Type( - map=stt.Type.Map(key=struct.types[0], value=expected, nullability=_REQ) + map=stt.Type.Map( + key=stt.Type(string=stt.Type.String(nullability=_REQ)), + value=expected, + nullability=_REQ, + ) ) row = stt.Type.Struct(types=[selected], nullability=_REQ) assert infer_expression_type(expr, row) == expected @@ -1068,7 +1072,11 @@ def test_nested_collection_selection_continues_through_struct(): expected = stt.Type(string=stt.Type.String(nullability=_NULL)) value = stt.Type(struct=stt.Type.Struct(types=[expected], nullability=_REQ)) mapping = stt.Type( - map=stt.Type.Map(key=struct.types[0], value=value, nullability=_REQ) + map=stt.Type.Map( + key=stt.Type(string=stt.Type.String(nullability=_REQ)), + value=value, + nullability=_REQ, + ) ) row = stt.Type.Struct( types=[stt.Type(list=stt.Type.List(type=mapping, nullability=_REQ))], From 77b2a02c9d0768fb4f143e02ea15bc8b1973413d Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Thu, 8 Oct 2026 19:16:03 +0300 Subject: [PATCH 4/4] fix(type-inference): validate map key literal types --- src/substrait/type_inference.py | 16 ++++++ tests/test_type_inference.py | 89 +++++++++++++++++++++++++++++++++ 2 files changed, 105 insertions(+) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 4fd979b..f58f649 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -666,6 +666,22 @@ def infer_expression_type( if result.WhichOneof("kind") != "map": raise ValueError("Map key reference requires a map") child = segment.map_key + key_type = stt.Type() + key_type.CopyFrom(infer_literal_type(child.map_key)) + key_kind = key_type.WhichOneof("kind") + if key_kind == result.map.key.WhichOneof("kind"): + detail = getattr(key_type, key_kind) + detail.nullability = getattr( + result.map.key, key_kind + ).nullability + if child.map_key.WhichOneof("literal_type") != "null": + detail.type_variation_reference = ( + child.map_key.type_variation_reference + ) + if key_type != result.map.key: + raise ValueError( + "Map key literal type does not match map key type" + ) result = result.map.value else: raise Exception(f"Unknown reference_type {kind}") diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 8e826f1..2573044 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -1068,6 +1068,95 @@ def test_nested_collection_selection(kind, nullable): assert infer_expression_type(expr, row) == expected +@pytest.mark.parametrize("literal_nullable", [False, True]) +@pytest.mark.parametrize("compatible", [False, True]) +def test_map_key_reference_validates_literal_type(literal_nullable, compatible): + expr = _field_reference(0) + key = expr.selection.direct_reference.struct_field.child.map_key.map_key + key.nullable = literal_nullable + if compatible: + key.i32 = 7 + else: + key.string = "wrong" + expected = stt.Type(string=stt.Type.String(nullability=_NULL)) + row = stt.Type.Struct( + types=[ + stt.Type( + map=stt.Type.Map( + key=stt.Type(i32=stt.Type.I32(nullability=_REQ)), + value=expected, + nullability=_REQ, + ) + ) + ], + nullability=_REQ, + ) + before = row.SerializeToString(), expr.SerializeToString() + if compatible: + assert infer_expression_type(expr, row) == expected + else: + with pytest.raises(ValueError, match="Map key literal type"): + infer_expression_type(expr, row) + assert before == (row.SerializeToString(), expr.SerializeToString()) + + +@pytest.mark.parametrize("scale", [2, 3]) +def test_map_key_reference_validates_decimal_parameters(scale): + expr = _field_reference(0) + key = expr.selection.direct_reference.struct_field.child.map_key.map_key + key.decimal.precision = 10 + key.decimal.scale = scale + key.decimal.value = (1000).to_bytes(16, "little", signed=True) + row = stt.Type.Struct( + types=[ + stt.Type( + map=stt.Type.Map( + key=stt.Type( + decimal=stt.Type.Decimal( + precision=10, scale=2, nullability=_REQ + ) + ), + value=stt.Type(i32=stt.Type.I32(nullability=_REQ)), + nullability=_REQ, + ) + ) + ], + nullability=_REQ, + ) + if scale == 2: + assert infer_expression_type(expr, row) == row.types[0].map.value + else: + with pytest.raises(ValueError, match="Map key literal type"): + infer_expression_type(expr, row) + + +@pytest.mark.parametrize("variation", [0, 1]) +def test_map_key_reference_validates_type_variation(variation): + expr = _field_reference(0) + literal = expr.selection.direct_reference.struct_field.child.map_key.map_key + literal.i32 = 7 + literal.type_variation_reference = variation + row = stt.Type.Struct( + types=[ + stt.Type( + map=stt.Type.Map( + key=stt.Type( + i32=stt.Type.I32(nullability=_REQ, type_variation_reference=1) + ), + value=stt.Type(i32=stt.Type.I32(nullability=_REQ)), + nullability=_REQ, + ) + ) + ], + nullability=_REQ, + ) + if variation == 1: + assert infer_expression_type(expr, row) == row.types[0].map.value + else: + with pytest.raises(ValueError, match="Map key literal type"): + infer_expression_type(expr, row) + + def test_nested_collection_selection_continues_through_struct(): expected = stt.Type(string=stt.Type.String(nullability=_NULL)) value = stt.Type(struct=stt.Type.Struct(types=[expected], nullability=_REQ))