From bb3c86a9a7392a9aafdd6f8f34534b113687e0b6 Mon Sep 17 00:00:00 2001 From: Nick Crews Date: Tue, 24 Sep 2024 12:42:06 -0800 Subject: [PATCH] fix: handle when structs have different fields in castable() Before, this would error with a KeyError --- ibis/expr/datatypes/cast.py | 3 ++- ibis/expr/datatypes/tests/test_cast.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/ibis/expr/datatypes/cast.py b/ibis/expr/datatypes/cast.py index 93a442786df0..5152326841e9 100644 --- a/ibis/expr/datatypes/cast.py +++ b/ibis/expr/datatypes/cast.py @@ -116,7 +116,8 @@ def castable(source: dt.DataType, target: dt.DataType, value: Any = None) -> boo ) elif target.is_struct(): return source.is_struct() and all( - castable(source[field], target[field]) for field in target.names + field in source and castable(source[field], target[field]) + for field in target.names ) elif target.is_geospatial(): return source.is_geospatial() or source.is_array() diff --git a/ibis/expr/datatypes/tests/test_cast.py b/ibis/expr/datatypes/tests/test_cast.py index 30a37b96ec93..bde76148719a 100644 --- a/ibis/expr/datatypes/tests/test_cast.py +++ b/ibis/expr/datatypes/tests/test_cast.py @@ -101,3 +101,21 @@ def test_implicitly_castable_int_to_bool(value): ) def test_implicitly_uncastable_values(source, target, value): assert not source.castable(target, value=value) + + +def test_struct_different_fields(): + x = dt.Struct({"x": dt.int32}) + x2 = dt.Struct({"x": dt.int64}) + y = dt.Struct({"y": dt.int32}) + xy = dt.Struct({"x": dt.int32, "y": dt.int32}) + + # Can upcast int32 to int64, but not other way + assert x.castable(x2) + assert not x2.castable(x) + # Can remove a field, but not add one + assert xy.castable(x) + assert not x.castable(xy) + + # Missing fields entirely from each other + assert not x.castable(y) + assert not y.castable(x)