Allow multiple graph tables inside single query

src/duckpgq_extension.cpp

@@ -12,6 +12,8 @@
 
 #include "duckdb/parser/tableref/table_function_ref.hpp"
 #include "duckdb/parser/expression/function_expression.hpp"
+#include "duckdb/parser/expression/constant_expression.hpp"
+
 #include "duckdb/parser/parsed_data/create_scalar_function_info.hpp"
 #include "duckdb/parser/parsed_data/create_table_function_info.hpp"
 #include "duckdb/parser/query_node/select_node.hpp"
@@ -83,16 +85,17 @@
 
 ParserExtensionParseResult duckpgq_parse(ParserExtensionInfo *info,
                                          const std::string &query) {
   auto parse_info = (DuckPGQParserExtensionInfo &)(info);
   Parser parser;
   parser.ParseQuery((query[0] == '-') ? query.substr(1, query.length())
                                       : query);
   if (parser.statements.size() != 1) {
     throw Exception(ExceptionType::PARSER,
-        "More than 1 statement detected, please only give one.");
+                    "More than one statement detected, please only give one.");
   }
-  return ParserExtensionParseResult(make_uniq_base<ParserExtensionParseData, DuckPGQParseData>(
-      std::move(parser.statements[0])));
+  return ParserExtensionParseResult(
+      make_uniq_base<ParserExtensionParseData, DuckPGQParseData>(
+          std::move(parser.statements[0])));
 }
 
 BoundStatement duckpgq_bind(ClientContext &context, Binder &binder,
@@ -119,15 +122,16 @@
     // Handle TableFunctionRef case
     auto function =
         dynamic_cast<FunctionExpression *>(table_function_ref->function.get());
-    if (function->function_name == "duckpgq_match") {
-      if (duckpgq_state.transform_expression != nullptr) {
-        throw Exception(ExceptionType::INVALID,
-                        "Multiple graph tables in a single query are not supported yet");
-      }
-      duckpgq_state.transform_expression =
-          std::move(std::move(function->children[0]));
-      function->children.pop_back();
+    if (function->function_name != "duckpgq_match") {
+      return;
     }
+    int32_t match_index = duckpgq_state.match_index++;
+    duckpgq_state.transform_expression[match_index] =
+        std::move(function->children[0]);
+    function->children.pop_back();
+    auto function_identifier =
+        make_uniq<ConstantExpression>(Value::CreateValue(match_index));
+    function->children.push_back(std::move(function_identifier));
   } else if (auto join_ref = dynamic_cast<JoinRef *>(table_ref)) {
     // Handle JoinRef case
     duckpgq_find_match_function(join_ref->left.get(), duckpgq_state);
@@ -141,7 +145,8 @@
     const auto select_statement = dynamic_cast<SelectStatement *>(statement);
     const auto select_node =
         dynamic_cast<SelectNode *>(select_statement->node.get());
-    const auto describe_node = dynamic_cast<ShowRef *>(select_node->from_table.get());
+    const auto describe_node =
+        dynamic_cast<ShowRef *>(select_node->from_table.get());
     if (describe_node) {
       ParserExtensionPlanResult result;
       result.function = DescribePropertyGraphFunction();
@@ -192,7 +197,8 @@
   }
 
   throw Exception(ExceptionType::NOT_IMPLEMENTED,
-    StatementTypeToString(statement->type) + "has not been implemented yet for DuckPGQ queries");
+                  StatementTypeToString(statement->type) +
+                      "has not been implemented yet for DuckPGQ queries");
 }
 
 ParserExtensionPlanResult

src/functions/scalar/csr_creation.cpp

@@ -36,7 +36,8 @@ static void CsrInitializeVertex(DuckPGQState &context, int32_t id,
     csr->initialized_v = true;
     context.csr_list[id] = std::move(csr);
   } catch (std::bad_alloc const &) {
-    throw Exception(ExceptionType::INTERNAL, "Unable to initialize vector of size for csr vertex table "
+    throw Exception(ExceptionType::INTERNAL,
+                    "Unable to initialize vector of size for csr vertex table "
                     "representation");
   }
 
@@ -55,7 +56,8 @@ static void CsrInitializeEdge(DuckPGQState &context, int32_t id, int64_t v_size,
     csr_entry->second->e.resize(e_size, 0);
     csr_entry->second->edge_ids.resize(e_size, 0);
   } catch (std::bad_alloc const &) {
-    throw Exception(ExceptionType::INTERNAL, "Unable to initialize vector of size for csr edge table "
+    throw Exception(ExceptionType::INTERNAL,
+                    "Unable to initialize vector of size for csr edge table "
                     "representation");
   }
   for (auto i = 1; i < v_size + 2; i++) {
@@ -82,7 +84,8 @@ static void CsrInitializeWeight(DuckPGQState &context, int32_t id,
       throw NotImplementedException("Unrecognized weight type detected.");
     }
   } catch (std::bad_alloc const &) {
-    throw Exception(ExceptionType::INTERNAL, "Unable to initialize vector of size for csr weight table "
+    throw Exception(ExceptionType::INTERNAL,
+                    "Unable to initialize vector of size for csr weight table "
                     "representation");
   }
 

src/functions/scalar/reachability.cpp

@@ -245,7 +245,8 @@ static void ReachabilityFunction(DataChunk &args, ExpressionState &state,
           break;
         }
         default:
-          throw Exception(ExceptionType::INTERNAL, "Unknown reachability mode encountered");
+          throw Exception(ExceptionType::INTERNAL,
+                          "Unknown reachability mode encountered");
         }
       } else {
         exit_early = BfsWithoutArray(exit_early, csr, input_size, seen, visit,

src/functions/tablefunctions/create_property_graph.cpp

@@ -9,15 +9,17 @@ void CreatePropertyGraphFunction::CheckPropertyGraphTableLabels(
     const shared_ptr<PropertyGraphTable> &pg_table, TableCatalogEntry &table) {
   if (!pg_table->discriminator.empty()) {
     if (!table.ColumnExists(pg_table->discriminator)) {
-      throw Exception(ExceptionType::INVALID, "Column " + pg_table->discriminator +
-        " not found in table " + pg_table->table_name);
+      throw Exception(ExceptionType::INVALID,
+                      "Column " + pg_table->discriminator +
+                          " not found in table " + pg_table->table_name);
     }
     auto &column = table.GetColumn(pg_table->discriminator);
     if (!(column.GetType() == LogicalType::BIGINT ||
           column.GetType() == LogicalType::INTEGER)) {
-      throw Exception(ExceptionType::INVALID, "The discriminator column " +
-        pg_table->discriminator + " of table " +
-        pg_table->table_name + " should be of type BIGINT or INTEGER");
+      throw Exception(ExceptionType::INVALID,
+                      "The discriminator column " + pg_table->discriminator +
+                          " of table " + pg_table->table_name +
+                          " should be of type BIGINT or INTEGER");
     }
   }
 }
@@ -31,8 +33,9 @@ void CreatePropertyGraphFunction::CheckPropertyGraphTableColumns(
   if (pg_table->all_columns) {
     for (auto &except_column : pg_table->except_columns) {
       if (!table.ColumnExists(except_column)) {
-        throw Exception(ExceptionType::INVALID, "Except column " + except_column +
-          " not found in table " + pg_table->table_name);
+        throw Exception(ExceptionType::INVALID,
+                        "Except column " + except_column +
+                            " not found in table " + pg_table->table_name);
       }
     }
 
@@ -52,13 +55,13 @@ void CreatePropertyGraphFunction::CheckPropertyGraphTableColumns(
   for (auto &column : pg_table->column_names) {
     if (!table.ColumnExists(column)) {
       throw Exception(ExceptionType::INVALID, "Column " + column +
-        " not found in table " + pg_table->table_name);
+                                                  " not found in table " +
+                                                  pg_table->table_name);
     }
   }
 }
 
-unique_ptr<FunctionData>
-CreatePropertyGraphFunction::CreatePropertyGraphBind(
+unique_ptr<FunctionData> CreatePropertyGraphFunction::CreatePropertyGraphBind(
     ClientContext &context, TableFunctionBindInput &input,
     vector<LogicalType> &return_types, vector<string> &names) {
   names.emplace_back("Success");
@@ -83,21 +86,24 @@ CreatePropertyGraphFunction::CreatePropertyGraphBind(
 
   if (pg_table != duckpgq_state->registered_property_graphs.end() &&
       info->on_conflict == OnCreateConflict::ERROR_ON_CONFLICT) {
-    throw Exception(ExceptionType::INVALID, "Property graph table with name " + info->property_graph_name + " already exists");
+    throw Exception(ExceptionType::INVALID, "Property graph table with name " +
+                                                info->property_graph_name +
+                                                " already exists");
   }
 
   auto &catalog = Catalog::GetCatalog(context, info->catalog);
   case_insensitive_set_t v_table_names;
   for (auto &vertex_table : info->vertex_tables) {
     try {
-        auto &table = catalog.GetEntry<TableCatalogEntry>(context, info->schema,
-                                                        vertex_table->table_name);
-
-        CheckPropertyGraphTableColumns(vertex_table, table);
-        CheckPropertyGraphTableLabels(vertex_table, table);
-      } catch (Exception &) {
-        throw Exception(ExceptionType::INVALID, vertex_table->table_name + " does not exist");
-      }
+      auto &table = catalog.GetEntry<TableCatalogEntry>(
+          context, info->schema, vertex_table->table_name);
+
+      CheckPropertyGraphTableColumns(vertex_table, table);
+      CheckPropertyGraphTableLabels(vertex_table, table);
+    } catch (Exception &) {
+      throw Exception(ExceptionType::INVALID,
+                      vertex_table->table_name + " does not exist");
+    }
 
     v_table_names.insert(vertex_table->table_name);
     if (vertex_table->hasTableNameAlias()) {
@@ -108,54 +114,65 @@ CreatePropertyGraphFunction::CreatePropertyGraphBind(
   for (auto &edge_table : info->edge_tables) {
     try {
       auto &table = catalog.GetEntry<TableCatalogEntry>(context, info->schema,
-                                                      edge_table->table_name);
+                                                        edge_table->table_name);
 
       CheckPropertyGraphTableColumns(edge_table, table);
       CheckPropertyGraphTableLabels(edge_table, table);
 
       for (auto &fk : edge_table->source_fk) {
         if (!table.ColumnExists(fk)) {
-          throw Exception(ExceptionType::INVALID,"Foreign key " + fk + " does not exist in table " + edge_table->table_name);
+          throw Exception(ExceptionType::INVALID,
+                          "Foreign key " + fk + " does not exist in table " +
+                              edge_table->table_name);
         }
       }
 
       for (auto &fk : edge_table->destination_fk) {
         if (!table.ColumnExists(fk)) {
-          throw Exception(ExceptionType::INVALID,"Foreign key " + fk + " does not exist in table " + edge_table->table_name);
+          throw Exception(ExceptionType::INVALID,
+                          "Foreign key " + fk + " does not exist in table " +
+                              edge_table->table_name);
         }
       }
-    } catch(const Exception &) {
-      throw Exception(ExceptionType::INVALID, edge_table->table_name + " does not exist");
+    } catch (const Exception &) {
+      throw Exception(ExceptionType::INVALID,
+                      edge_table->table_name + " does not exist");
     }
 
     if (v_table_names.find(edge_table->source_reference) ==
         v_table_names.end()) {
-      throw Exception(ExceptionType::INVALID, "Referenced vertex table " + edge_table->source_reference + " does not exist.");
+      throw Exception(ExceptionType::INVALID, "Referenced vertex table " +
+                                                  edge_table->source_reference +
+                                                  " does not exist.");
     }
 
     auto &pk_source_table = catalog.GetEntry<TableCatalogEntry>(
         context, info->schema, edge_table->source_reference);
     for (auto &pk : edge_table->source_pk) {
       if (!pk_source_table.ColumnExists(pk)) {
-        throw Exception(ExceptionType::INVALID, "Primary key " + pk + " does not exist in table " + edge_table->source_reference);
+        throw Exception(ExceptionType::INVALID,
+                        "Primary key " + pk + " does not exist in table " +
+                            edge_table->source_reference);
       }
     }
 
     if (v_table_names.find(edge_table->source_reference) ==
         v_table_names.end()) {
-      throw Exception(ExceptionType::INVALID, "Referenced vertex table " + edge_table->source_reference + " does not exist");
+      throw Exception(ExceptionType::INVALID, "Referenced vertex table " +
+                                                  edge_table->source_reference +
+                                                  " does not exist");
     }
 
     auto &pk_destination_table = catalog.GetEntry<TableCatalogEntry>(
         context, info->schema, edge_table->destination_reference);
 
     for (auto &pk : edge_table->destination_pk) {
       if (!pk_destination_table.ColumnExists(pk)) {
-        throw Exception(ExceptionType::INVALID,"Primary key " + pk + " does not exist in table " + edge_table->destination_reference);
+        throw Exception(ExceptionType::INVALID,
+                        "Primary key " + pk + " does not exist in table " +
+                            edge_table->destination_reference);
       }
     }
-
-
   }
   return make_uniq<CreatePropertyGraphBindData>(info);
 }
@@ -173,10 +190,11 @@ void CreatePropertyGraphFunction::CreatePropertyGraphFunc(
   auto pg_info = bind_data.create_pg_info;
   auto lookup = context.registered_state.find("duckpgq");
   if (lookup == context.registered_state.end()) {
-    throw Exception(ExceptionType::INVALID,"Registered DuckPGQ state not found");
+    throw Exception(ExceptionType::INVALID,
+                    "Registered DuckPGQ state not found");
   }
   auto duckpgq_state = (DuckPGQState *)lookup->second.get();
   duckpgq_state->registered_property_graphs[pg_info->property_graph_name] =
-    pg_info->Copy();
+      pg_info->Copy();
 }
 }; // namespace duckdb

src/functions/tablefunctions/describe_property_graph.cpp

@@ -29,10 +29,13 @@ DescribePropertyGraphFunction::DescribePropertyGraphBind(
 
   auto pg_table =
       duckpgq_state->registered_property_graphs.find(show_ref->table_name);
-  if (pg_table == duckpgq_state->registered_property_graphs.end() ) {
-    throw Exception(ExceptionType::INVALID, "Property graph " + show_ref->table_name + " does not exist.");
+  if (pg_table == duckpgq_state->registered_property_graphs.end()) {
+    throw Exception(ExceptionType::INVALID, "Property graph " +
+                                                show_ref->table_name +
+                                                " does not exist.");
   }
-  auto property_graph = dynamic_cast<CreatePropertyGraphInfo *>(pg_table->second.get());
+  auto property_graph =
+      dynamic_cast<CreatePropertyGraphInfo *>(pg_table->second.get());
 
   names.emplace_back("table_name");
   return_types.emplace_back(LogicalType::VARCHAR);
@@ -57,7 +60,6 @@ DescribePropertyGraphFunction::DescribePropertyGraphBind(
   names.emplace_back("sub_labels");
   return_types.emplace_back(LogicalType::LIST(LogicalType::VARCHAR));
 
-
   return make_uniq<DescribePropertyGraphBindData>(property_graph);
 }
 
@@ -76,7 +78,7 @@ void DescribePropertyGraphFunction::DescribePropertyGraphFunc(
   }
   auto pg_info = bind_data.describe_pg_info;
   idx_t vector_idx = 0;
-  for (const auto& vertex_table : pg_info->vertex_tables) {
+  for (const auto &vertex_table : pg_info->vertex_tables) {
     output.SetValue(0, vector_idx, Value(vertex_table->table_name));
     output.SetValue(1, vector_idx, Value(vertex_table->main_label));
     output.SetValue(2, vector_idx, Value(vertex_table->is_vertex_table));
@@ -89,49 +91,55 @@ void DescribePropertyGraphFunction::DescribePropertyGraphFunc(
     if (!vertex_table->discriminator.empty()) {
       output.SetValue(9, vector_idx, Value(vertex_table->discriminator));
       vector<Value> sub_labels;
-      for (const auto& label : vertex_table->sub_labels) {
+      for (const auto &label : vertex_table->sub_labels) {
         sub_labels.push_back(Value(label));
       }
-      output.SetValue(10, vector_idx, Value::LIST(LogicalType::VARCHAR, sub_labels));
+      output.SetValue(10, vector_idx,
+                      Value::LIST(LogicalType::VARCHAR, sub_labels));
     } else {
       output.SetValue(9, vector_idx, Value());
       output.SetValue(10, vector_idx, Value());
     }
     vector_idx++;
   }
-  for (const auto& edge_table : pg_info->edge_tables) {
+  for (const auto &edge_table : pg_info->edge_tables) {
     output.SetValue(0, vector_idx, Value(edge_table->table_name));
     output.SetValue(1, vector_idx, Value(edge_table->main_label));
     output.SetValue(2, vector_idx, Value(edge_table->is_vertex_table));
     output.SetValue(3, vector_idx, Value(edge_table->source_reference));
     vector<Value> source_pk_list;
-    for (const auto& col : edge_table->source_pk) {
+    for (const auto &col : edge_table->source_pk) {
       source_pk_list.push_back(Value(col));
     }
-    output.SetValue(4, vector_idx, Value::LIST(LogicalType::VARCHAR,source_pk_list));
+    output.SetValue(4, vector_idx,
+                    Value::LIST(LogicalType::VARCHAR, source_pk_list));
     vector<Value> source_fk_list;
-    for (const auto& col : edge_table->source_fk) {
+    for (const auto &col : edge_table->source_fk) {
       source_fk_list.push_back(Value(col));
     }
-    output.SetValue(5, vector_idx, Value::LIST(LogicalType::VARCHAR,source_fk_list));
+    output.SetValue(5, vector_idx,
+                    Value::LIST(LogicalType::VARCHAR, source_fk_list));
     output.SetValue(6, vector_idx, Value(edge_table->destination_reference));
     vector<Value> destination_pk_list;
-    for (const auto& col : edge_table->destination_pk) {
+    for (const auto &col : edge_table->destination_pk) {
       destination_pk_list.push_back(Value(col));
     }
-    output.SetValue(7, vector_idx, Value::LIST(LogicalType::VARCHAR,destination_pk_list));
+    output.SetValue(7, vector_idx,
+                    Value::LIST(LogicalType::VARCHAR, destination_pk_list));
     vector<Value> destination_fk_list;
-    for (const auto& col : edge_table->destination_fk) {
+    for (const auto &col : edge_table->destination_fk) {
       destination_fk_list.push_back(Value(col));
     }
-    output.SetValue(8, vector_idx, Value::LIST(LogicalType::VARCHAR,destination_fk_list));
+    output.SetValue(8, vector_idx,
+                    Value::LIST(LogicalType::VARCHAR, destination_fk_list));
     if (!edge_table->discriminator.empty()) {
       output.SetValue(9, vector_idx, Value(edge_table->discriminator));
       vector<Value> sub_labels;
-      for (const auto& label : edge_table->sub_labels) {
+      for (const auto &label : edge_table->sub_labels) {
         sub_labels.push_back(Value(label));
       }
-      output.SetValue(10, vector_idx, Value::LIST(LogicalType::VARCHAR, sub_labels));
+      output.SetValue(10, vector_idx,
+                      Value::LIST(LogicalType::VARCHAR, sub_labels));
     } else {
       output.SetValue(9, vector_idx, Value());
       output.SetValue(10, vector_idx, Value());