Skip to content

Commit

Permalink
Correctly resolve None items inside ListPlaceholder
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 597825457
  • Loading branch information
tfx-copybara committed Jun 4, 2024
1 parent 07b193d commit a1a51c4
Show file tree
Hide file tree
Showing 4 changed files with 71 additions and 6 deletions.
7 changes: 4 additions & 3 deletions tfx/dsl/compiler/placeholder_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -433,9 +433,10 @@ def _resolve_list_concat_operator(
"""Evaluates the list concat operator."""
result = []
for sub_expression in op.expressions:
value = self.resolve(sub_expression, pool)
if value is None:
raise NullDereferenceError(sub_expression)
try:
value = self.resolve(sub_expression, pool)
except NullDereferenceError:
value = None
result.append(value)
return result

Expand Down
37 changes: 37 additions & 0 deletions tfx/dsl/compiler/placeholder_utils_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -665,6 +665,43 @@ def testListConcat(self):
placeholder_utils.resolve_placeholder_expression(
pb, self._resolution_context), expected_result)

def testListConcatWithAbsentElement(self):
# When an exec prop has type Union[T, None] and the user passes None, it is
# actually completely absent from the exec_properties dict in
# ExecutionInvocation. See also b/172001324 and the corresponding todo in
# placeholder_utils.py.
placeholder_expression = """
operator {
list_concat_op {
expressions {
value {
string_value: "random_before"
}
}
expressions {
placeholder {
type: EXEC_PROPERTY
key: "doesnotexist"
}
}
expressions {
value {
string_value: "random_after"
}
}
}
}
"""
pb = text_format.Parse(
placeholder_expression, placeholder_pb2.PlaceholderExpression()
)
self.assertEqual(
placeholder_utils.resolve_placeholder_expression(
pb, self._resolution_context
),
["random_before", None, "random_after"],
)

def testListConcatAndSerialize(self):
placeholder_expression = """
operator {
Expand Down
3 changes: 1 addition & 2 deletions tfx/dsl/placeholder/placeholder_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -354,8 +354,7 @@ def serialize_list(
"""Serializes list-value placeholder to JSON or comma-separated string.
Only supports primitive type list element (a.k.a bool, int, float or str) at
the
moment; throws runtime error otherwise.
the moment; throws runtime error otherwise.
Args:
serialization_format: The format of how the proto is serialized.
Expand Down
30 changes: 29 additions & 1 deletion tfx/dsl/placeholder/proto_placeholder_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -220,7 +220,8 @@ def test_NonePlaceholderIntoOptionalField(self):
def test_NoneExecPropIntoOptionalField(self):
# When an exec prop has type Union[T, None] and the user passes None, it is
# actually completely absent from the exec_properties dict in
# ExecutionInvocation.
# ExecutionInvocation. See also b/172001324 and the corresponding todo in
# placeholder_utils.py.
actual = resolve(
_UpdateOptions(reload_policy=ph.exec_property('reload_policy')),
exec_properties={}, # Intentionally empty.
Expand Down Expand Up @@ -385,6 +386,33 @@ def test_RepeatedFieldFalsyItem(self):
parse_text_proto(actual),
)

def test_RepeatedFieldNoneItem(self):
actual = resolve(
ph.make_proto(
execution_invocation_pb2.ExecutionInvocation(
pipeline_node=pipeline_pb2.PipelineNode()
),
pipeline_node=ph.make_proto(
pipeline_pb2.PipelineNode(),
upstream_nodes=[
'foo',
ph.exec_property('reload_policy'), # Will be None.
'bar',
],
),
),
exec_properties={}, # Intentionally empty.
)
self.assertProtoEquals(
"""
pipeline_node {
upstream_nodes: "foo"
upstream_nodes: "bar"
}
""",
parse_text_proto(actual),
)

def test_NoneIntoRepeatedField(self):
actual = resolve(
ph.make_proto(pipeline_pb2.PipelineNode(), upstream_nodes=None)
Expand Down

0 comments on commit a1a51c4

Please sign in to comment.