diff --git a/cpp/src/arrow/engine/substrait/relation_internal.cc b/cpp/src/arrow/engine/substrait/relation_internal.cc index 55ca899680e9..79abb9fe2356 100644 --- a/cpp/src/arrow/engine/substrait/relation_internal.cc +++ b/cpp/src/arrow/engine/substrait/relation_internal.cc @@ -643,8 +643,7 @@ Result FromProto(const substrait::Rel& rel, const ExtensionSet& ARROW_ASSIGN_OR_RAISE(compute::Expression bound_expr, des_expr.Bind(*input.output_schema)); if (auto* expr_call = bound_expr.call()) { - project_field = field(expr_call->function_name, - expr_call->kernel->signature->out_type().type()); + project_field = field(expr_call->function_name, expr_call->type.GetSharedPtr()); } else if (auto* field_ref = des_expr.field_ref()) { ARROW_ASSIGN_OR_RAISE(FieldPath field_path, field_ref->FindOne(*input.output_schema)); diff --git a/cpp/src/arrow/engine/substrait/serde_test.cc b/cpp/src/arrow/engine/substrait/serde_test.cc index 8c3f5ca69876..1330ea763142 100644 --- a/cpp/src/arrow/engine/substrait/serde_test.cc +++ b/cpp/src/arrow/engine/substrait/serde_test.cc @@ -2905,6 +2905,75 @@ TEST(SubstraitRoundTrip, ProjectRelOnFunctionWithEmit) { /*include_columns=*/{}, conversion_options); } +TEST(SubstraitRoundTrip, ProjectRelWithComputedOutputType) { + auto input_schema = + schema({field("A", boolean()), field("B", boolean()), field("C", int32())}); + auto input_table = TableFromJSON(input_schema, {R"([ + [true, false, 10], + [false, true, 20], + [false, false, 30], + [null, null, 40] + ])"}); + + std::string substrait_json = R"({ + "version": {"major_number": 9999, "minor_number": 9999, "patch_number": 9999}, + "relations": [{"rel": {"project": { + "expressions": [{"if_then": { + "ifs": [ + { + "if": {"selection": { + "directReference": {"structField": {"field": 0}}, "rootReference": {} + }}, + "then": {"selection": { + "directReference": {"structField": {"field": 2}}, "rootReference": {} + }} + }, + { + "if": {"selection": { + "directReference": {"structField": {"field": 1}}, "rootReference": {} + }}, + "then": {"literal": {"i32": 0}} + } + ], + "else": {"literal": {"i32": -1}} + }}], + "input": {"read": { + "base_schema": { + "names": ["A", "B", "C"], + "struct": {"types": [{"bool": {}}, {"bool": {}}, {"i32": {}}]} + }, + "namedTable": {"names": ["A"]} + }} + }}}] + })"; + + ASSERT_OK_AND_ASSIGN(auto buf, + internal::SubstraitFromJSON("Plan", substrait_json, + /*ignore_unknown_fields=*/false)); + ConversionOptions conversion_options; + conversion_options.named_table_provider = AlwaysProvideSameTable(input_table); + + // case_when resolves its output type when the expression is bound. + ASSERT_OK_AND_ASSIGN(auto plan_info, + DeserializePlan(*buf, /*registry=*/nullptr, + /*ext_set_out=*/nullptr, conversion_options)); + const auto& output_schema = plan_info.root.output_schema; + ASSERT_EQ(4, output_schema->num_fields()); + ASSERT_NE(nullptr, output_schema->field(3)->type()); + AssertTypeEqual(*int32(), *output_schema->field(3)->type()); + + auto expected_schema = schema({field("A", boolean()), field("B", boolean()), + field("C", int32()), field("case_when", int32())}); + auto expected_table = TableFromJSON(expected_schema, {R"([ + [true, false, 10, 10], + [false, true, 20, 0], + [false, false, 30, -1], + [null, null, 40, -1] + ])"}); + CheckRoundTripResult(std::move(expected_table), buf, /*include_columns=*/{}, + conversion_options); +} + TEST(SubstraitRoundTrip, ProjectRelOnFunctionWithAllEmit) { compute::ExecContext exec_context; auto dummy_schema = schema({field("A", int32()), field("B", int32())});