vibhatha commented on code in PR #12672:
URL: https://github.com/apache/arrow/pull/12672#discussion_r862282740
##########
cpp/src/arrow/engine/substrait/serde_test.cc:
##########
@@ -724,5 +728,103 @@ TEST(Substrait, ExtensionSetFromPlan) {
EXPECT_EQ(decoded_add_func.name, "add");
}
+TEST(Substrait, GetRecordBatchIterator) {
+ const auto parquet_root = std::getenv("PARQUET_TEST_DATA");
+ std::string dir_string(parquet_root);
+ std::stringstream ss;
+ ss << dir_string << "/binary.parquet";
+ auto file_path = ss.str();
+
+ std::string substrait_json = R"({
+ "relations": [
+ {"rel": {
+ "read": {
+ "base_schema": {
+ "struct": {
+ "types": [
+ {"binary": {}}
+ ]
+ },
+ "names": [
+ "foo"
+ ]
+ },
+ "local_files": {
+ "items": [
+ {
+ "uri_file": "file://FILENAME_PLACEHOLDER",
+ "format": "FILE_FORMAT_PARQUET"
+ }
+ ]
+ }
+ }
+ }}
+ ]
+ })";
+
+ std::string filename_placeholder = "FILENAME_PLACEHOLDER";
+ substrait_json.replace(substrait_json.find(filename_placeholder),
+ filename_placeholder.size(), file_path);
+ auto in_schema = schema({field("foo", binary())});
+ AsyncGenerator<util::optional<cp::ExecBatch>> sink_gen;
+ cp::ExecContext exec_context(default_memory_pool(),
+ arrow::internal::GetCpuThreadPool());
+ ASSERT_OK_AND_ASSIGN(auto plan, cp::ExecPlan::Make());
+ engine::SubstraitExecutor executor(substrait_json, &sink_gen, plan,
in_schema,
+ exec_context);
+ auto status = executor.MakePlan();
+ ASSERT_OK(status);
+ ASSERT_OK_AND_ASSIGN(auto reader, executor.Execute());
+ auto finish = executor.Finalize();
+ ASSERT_OK(finish);
+ ASSERT_OK_AND_ASSIGN(auto table, Table::FromRecordBatchReader(reader.get()));
+ EXPECT_GT(table->num_rows(), 0);
+}
+
+TEST(Substrait, GetRecordBatchIteratorUtil) {
+ const auto parquet_root = std::getenv("PARQUET_TEST_DATA");
+ std::string dir_string(parquet_root);
+ std::stringstream ss;
+ ss << dir_string << "/binary.parquet";
+ auto file_path = ss.str();
+
+ std::string substrait_json = R"({
+ "relations": [
+ {"rel": {
+ "read": {
+ "base_schema": {
+ "struct": {
+ "types": [
+ {"binary": {}}
+ ]
+ },
+ "names": [
+ "foo"
+ ]
+ },
+ "local_files": {
+ "items": [
+ {
+ "uri_file": "file://FILENAME_PLACEHOLDER",
+ "format": "FILE_FORMAT_PARQUET"
+ }
+ ]
+ }
+ }
+ }}
+ ]
+ })";
+
+ std::string filename_placeholder = "FILENAME_PLACEHOLDER";
+ substrait_json.replace(substrait_json.find(filename_placeholder),
+ filename_placeholder.size(), file_path);
+ auto in_schema = schema({field("foo", binary())});
+
+ ASSERT_OK_AND_ASSIGN(auto reader,
engine::SubstraitExecutor::GetRecordBatchReader(
+ substrait_json, in_schema));
+ ASSERT_OK_AND_ASSIGN(auto table, Table::FromRecordBatchReader(reader.get()));
+ EXPECT_GT(table->num_rows(), 0);
Review Comment:
cc @lidavidm I updated the test case and also added note in case of a
modification to the test file could cause a test failure. I think being exact
is better than checking a non-zero value and thank you for pointing this out in
the first place. Plus in Python we verify this to the letter by using the
parquet reader. Using parquet reader here seemed like an overkill. So one way
or the other, the objective is covered :)
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
