add support for new/optional fields to test utils

apache · Oct 29, 2021 · 9c9d62e · 9c9d62e
1 parent cc914e1
commit 9c9d62e
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/cpp/src/arrow/compute/exec/ir_consumer.h b/cpp/src/arrow/compute/exec/ir_consumer.h
@@ -38,11 +38,18 @@ namespace ir = org::apache::arrow::computeir::flatbuf;
 
 class ARROW_EXPORT CatalogSourceNodeOptions : public ExecNodeOptions {
  public:
-  CatalogSourceNodeOptions(std::string name, std::shared_ptr<Schema> schema)
-      : name(std::move(name)), schema(std::move(schema)) {}
+  CatalogSourceNodeOptions(std::string name, std::shared_ptr<Schema> schema,
+                           Expression filter = literal(true),
+                           std::vector<FieldRef> projection = {})
+      : name(std::move(name)),
+        schema(std::move(schema)),
+        filter(std::move(filter)),
+        projection(std::move(projection)) {}
 
   std::string name;
   std::shared_ptr<Schema> schema;
+  Expression filter;
+  std::vector<FieldRef> projection;
 };
 
 ARROW_EXPORT

diff --git a/cpp/src/arrow/compute/exec/test_util.cc b/cpp/src/arrow/compute/exec/test_util.cc
@@ -258,7 +258,13 @@ bool operator==(const Declaration& l, const Declaration& r) {
   if (l.factory_name == "catalog_source") {
     auto l_opts = &OptionsAs<CatalogSourceNodeOptions>(l);
     auto r_opts = &OptionsAs<CatalogSourceNodeOptions>(r);
-    return l_opts->name == r_opts->name && l_opts->schema->Equals(r_opts->schema);
+
+    bool schemas_equal = l_opts->schema == nullptr
+                             ? r_opts->schema == nullptr
+                             : l_opts->schema->Equals(r_opts->schema);
+
+    return l_opts->name == r_opts->name && schemas_equal &&
+           l_opts->filter == r_opts->filter && l_opts->projection == r_opts->projection;
   }
 
   if (l.factory_name == "filter") {
@@ -309,6 +315,16 @@ static inline void PrintToImpl(const std::string& factory_name,
   if (factory_name == "catalog_source") {
     auto o = &OptionsAs<CatalogSourceNodeOptions>(opts);
     *os << o->name << ", schema=" << o->schema->ToString();
+    if (o->filter != literal(true)) {
+      *os << ", filter=" << o->filter.ToString();
+    }
+    if (!o->projection.empty()) {
+      *os << ", projection=[";
+      for (const auto& ref : o->projection) {
+        *os << ref.ToString() << ",";
+      }
+      *os << "]";
+    }
     return;
   }