vibhatha commented on code in PR #14118: URL: https://github.com/apache/arrow/pull/14118#discussion_r970682427
########## cpp/src/arrow/engine/substrait/serde_test.cc: ########## @@ -97,6 +97,78 @@ Result<std::shared_ptr<Table>> GetTableFromPlan( return arrow::Table::FromRecordBatchReader(sink_reader.get()); } +void AssertScanRelation(const compute::Declaration& output_scan, + const std::shared_ptr<dataset::Dataset>& expected_dataset, + const std::shared_ptr<Schema> schema) { + const auto& dataset_opts = + checked_cast<const dataset::ScanNodeOptions&>(*(output_scan.options)); + const auto& output_ds = dataset_opts.dataset; + ASSERT_TRUE(output_ds->schema()->Equals(*schema)); + ASSERT_OK_AND_ASSIGN(auto output_frgs, output_ds->GetFragments()); + ASSERT_OK_AND_ASSIGN(auto expected_frgs, expected_dataset->GetFragments()); + + auto output_frg_vec = IteratorToVector(std::move(output_frgs)); + auto expected_frg_vec = IteratorToVector(std::move(expected_frgs)); + ASSERT_EQ(expected_frg_vec.size(), output_frg_vec.size()); + int64_t idx = 0; + for (auto fragment : expected_frg_vec) { + const auto* l_frag = checked_cast<const dataset::FileFragment*>(fragment.get()); + const auto* r_frag = + checked_cast<const dataset::FileFragment*>(output_frg_vec[idx++].get()); + ASSERT_TRUE(l_frag->Equals(*r_frag)); + } +} + +void AssertExpressionCall(const std::shared_ptr<Schema> schema, + const compute::Expression output_expr, + const compute::Expression& expected_expr) { + if (auto* out_call = output_expr.call()) { + if (auto* exp_call = expected_expr.call()) { + ASSERT_EQ(out_call->function_name, exp_call->function_name); + auto out_args = out_call->arguments; + auto exp_args = exp_call->arguments; + ASSERT_EQ(out_args.size(), exp_args.size()); + int exp_id = 0; + for (const auto& arg : exp_args) { + auto lhs = out_args[exp_id++].field_ref()->field_path()->indices()[0]; + ASSERT_EQ(schema->field_names()[lhs], *(arg.field_ref()->name())); + } + } + } +} + +void AssertFilterRelation(const compute::Declaration& output_filter, + const std::string& filter_func_name, + const compute::Expression& exp_filter_expr, + const std::shared_ptr<Schema>& schema) { + const auto& filter_opts = + checked_cast<const compute::FilterNodeOptions&>(*(output_filter.options)); + auto out_filter_expr = filter_opts.filter_expression; + AssertExpressionCall(schema, out_filter_expr, exp_filter_expr); +} + +void AssertProjectRelation(const compute::Declaration& output_projection, + const std::vector<compute::Expression>& exp_expressions, + const std::shared_ptr<Schema>& schema) { + const auto& project_opts = + checked_cast<const compute::ProjectNodeOptions&>(*(output_projection.options)); + auto out_expressions = project_opts.expressions; + int expr_id = 0; + ASSERT_EQ(out_expressions.size(), exp_expressions.size()); + for (const auto& out_expr : out_expressions) { + AssertExpressionCall(schema, out_expr, exp_expressions[expr_id++]); + } +} + +void AssertPlanExecutionResult(const std::shared_ptr<Table> expected_table, Review Comment: Should we keep this method or just copy and paste the following lines each test case? Note: I created this one in case of adding additional options for asserting results. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org