Simple reproduction case for read local files issue.

EpsilonPrime · May 7, 2024 · dddaf0b · dddaf0b
1 parent e7f5f81
commit dddaf0b
Showing 1 changed file with 79 additions and 0 deletions.
diff --git a/cpp/src/arrow/engine/substrait/serde_test.cc b/cpp/src/arrow/engine/substrait/serde_test.cc
@@ -1064,6 +1064,85 @@ NamedTableProvider AlwaysProvideSameTable(std::shared_ptr<Table> table) {
   };
 }
 
+TEST(Substrait, ReadRelWithRoot) {
+    ASSERT_OK_AND_ASSIGN(std::string dir_string,
+                         arrow::internal::GetEnvVar("PARQUET_TEST_DATA"));
+
+    std::string substrait_json = R"({
+        "relations": [
+        {
+          "root": {
+            "input": {
+              "read": {
+                "common": {
+                  "direct": {}
+                },
+                "baseSchema": {
+                  "names": [
+                    "f32",
+                    "f64"
+                  ],
+                  "struct": {
+                    "types": [
+                      {
+                        "fp32": {
+                          "nullability": "NULLABILITY_REQUIRED"
+                        }
+                      },
+                      {
+                        "fp64": {
+                          "nullability": "NULLABILITY_REQUIRED"
+                        }
+                      }
+                    ],
+                    "nullability": "NULLABILITY_REQUIRED"
+                  }
+                },
+                "localFiles": {
+                  "items": [
+                    {
+                      "uriFile": "file://[DIRECTORY_PLACEHOLDER]/byte_stream_split.zstd.parquet",
+                      "parquet": {}
+                    }
+                  ]
+                }
+              }
+            },
+            "names": [
+              "f32",
+              "f64"
+            ]
+          }
+        }
+        ],
+        "version": {
+        "minorNumber": 42,
+        "producer": "my-producer"
+        }
+    })";
+    const char* placeholder="[DIRECTORY_PLACEHOLDER]";
+    substrait_json.replace(substrait_json.find(placeholder), strlen(placeholder), dir_string);
+
+    ASSERT_OK_AND_ASSIGN(auto buf,
+                         internal::SubstraitFromJSON("Plan", substrait_json,
+                                 /*ignore_unknown_fields=*/false));
+
+    ASSERT_OK_AND_ASSIGN(auto declarations,
+                         DeserializePlans(*buf, acero::NullSinkNodeConsumer::Make));
+    ASSERT_EQ(declarations.size(), 1);
+    acero::Declaration* decl = &declarations[0];
+    ASSERT_EQ(decl->factory_name, "consuming_sink");
+    ASSERT_OK_AND_ASSIGN(auto plan, acero::ExecPlan::Make());
+    ASSERT_OK_AND_ASSIGN(auto sink_node, declarations[0].AddToPlan(plan.get()));
+    ASSERT_STREQ(sink_node->kind_name(), "ConsumingSinkNode");
+    ASSERT_EQ(sink_node->num_inputs(), 1);
+    auto& prev_node = sink_node->inputs()[0];
+    ASSERT_STREQ(prev_node->kind_name(), "SourceNode");
+
+    plan->StartProducing();
+    ASSERT_FINISHES_OK(plan->finished());
+}
+
 TEST(Substrait, RelWithHint) {
   ASSERT_OK_AND_ASSIGN(auto buf,
                        internal::SubstraitFromJSON("Rel", R"({