Merge pull request #86 from cwida/85-with-queries-lead-to-segmentatio…

…n-fault

WITH statements and subqueries now seem to work

duckpgq/src/duckpgq_extension.cpp

@@ -17,6 +17,7 @@
 #include "duckdb/parser/query_node/select_node.hpp"
 #include "duckdb/parser/statement/copy_statement.hpp"
 #include "duckdb/parser/parsed_data/create_table_info.hpp"
+#include "duckdb/parser/tableref/joinref.hpp"
 
 #include "duckdb/parser/statement/extension_statement.hpp"
 
@@ -109,32 +110,40 @@ BoundStatement duckpgq_bind(ClientContext &context, Binder &binder,
   throw BinderException("Unable to find DuckPGQ Parse Data");
 }
 
-ParserExtensionPlanResult duckpgq_handle_statement(SQLStatement *statement, DuckPGQState &duckpgq_state) {
-  if (statement->type == StatementType::SELECT_STATEMENT) {
-    auto select_statement = dynamic_cast<SelectStatement *>(statement);
-    auto select_node = dynamic_cast<SelectNode *>(select_statement->node.get());
-    auto from_table_function =
-        dynamic_cast<TableFunctionRef *>(select_node->from_table.get());
-    auto function =
-        dynamic_cast<FunctionExpression *>(from_table_function->function.get());
+void duckpgq_find_match_function(TableRef* table_ref, DuckPGQState &duckpgq_state) {
+  if (auto table_function_ref = dynamic_cast<TableFunctionRef *>(table_ref)) {
+    // Handle TableFunctionRef case
+    auto function = dynamic_cast<FunctionExpression *>(table_function_ref->function.get());
     if (function->function_name == "duckpgq_match") {
       duckpgq_state.transform_expression =
           std::move(std::move(function->children[0]));
       function->children.pop_back();
     }
+  } else if (auto join_ref = dynamic_cast<JoinRef *>(table_ref)) {
+    // Handle JoinRef case
+    duckpgq_find_match_function(join_ref->left.get(), duckpgq_state);
+    duckpgq_find_match_function(join_ref->right.get(), duckpgq_state);
+  }
+}
+
+ParserExtensionPlanResult duckpgq_handle_statement(SQLStatement *statement, DuckPGQState &duckpgq_state) {
+  if (statement->type == StatementType::SELECT_STATEMENT) {
+    auto select_statement = dynamic_cast<SelectStatement *>(statement);
+    auto select_node = dynamic_cast<SelectNode *>(select_statement->node.get());
+    duckpgq_find_match_function(select_node->from_table.get(), duckpgq_state);
     throw Exception("use duckpgq_bind instead");
   }
   if (statement->type == StatementType::CREATE_STATEMENT) {
-    auto &create_statement = statement->Cast<CreateStatement>();
-    auto create_property_graph = dynamic_cast<CreatePropertyGraphInfo*>(create_statement.info.get());
+    const auto &create_statement = statement->Cast<CreateStatement>();
+    const auto create_property_graph = dynamic_cast<CreatePropertyGraphInfo*>(create_statement.info.get());
     if (create_property_graph) {
       ParserExtensionPlanResult result;
       result.function = CreatePropertyGraphFunction();
       result.requires_valid_transaction = true;
       result.return_type = StatementReturnType::QUERY_RESULT;
       return result;
     }
-    auto create_table = reinterpret_cast<CreateTableInfo*>(create_statement.info.get());
+    const auto create_table = reinterpret_cast<CreateTableInfo*>(create_statement.info.get());
     duckpgq_handle_statement(create_table->query.get(), duckpgq_state);
   }
   if (statement->type == StatementType::DROP_STATEMENT) {
@@ -152,15 +161,7 @@ ParserExtensionPlanResult duckpgq_handle_statement(SQLStatement *statement, Duck
   if (statement->type == StatementType::COPY_STATEMENT) {
     auto &copy_statement = statement->Cast<CopyStatement>();
     auto select_node = dynamic_cast<SelectNode *>(copy_statement.select_statement.get());
-    auto from_table_function =
-        dynamic_cast<TableFunctionRef *>(select_node->from_table.get());
-    auto function =
-        dynamic_cast<FunctionExpression *>(from_table_function->function.get());
-    if (function->function_name == "duckpgq_match") {
-      duckpgq_state.transform_expression =
-          std::move(std::move(function->children[0]));
-      function->children.pop_back();
-    }
+    duckpgq_find_match_function(select_node->from_table.get(), duckpgq_state);
     throw Exception("use duckpgq_bind instead");
   }
   if (statement->type == StatementType::INSERT_STATEMENT) {

test/sql/with_statement_duckpgq.test

@@ -0,0 +1,73 @@
+# name: test/sql/sqlpgq/snb.test
+# group: [duckpgq]
+
+require duckpgq
+
+statement ok
+import database 'duckdb-pgq/data/SNB0.003';
+
+statement ok
+-CREATE PROPERTY GRAPH snb_projected
+VERTEX TABLES (Message);
+
+query IIIIIII
+-WITH message_count AS (
+    SELECT count(*) as m_count
+    FROM Message m
+    WHERE m.creationDate < '2010-05-27 11:16:36.013'
+)
+SELECT year, isComment,
+    CASE    WHEN m_length < 40 THEN 0
+            WHEN m_length < 80 THEN 1
+            WHEN m_length < 160 THEN 2
+            ELSE 3 END as lengthCategory,
+    count(*) as messageCount,
+    avg(m_length) as averageMessageLength,
+    sum(m_length) as sumMessageLength,
+    count(*) / mc.m_count as percentageOfMessages
+FROM GRAPH_TABLE(snb_projected
+    MATCH (message:Message where message.creationDate < '2010-05-27 11:16:36.013')
+    COLUMNS (date_part('year', message.creationDate::TIMESTAMP) as year, message.ImageFile is NULL as isComment, message.length as m_length, message.id)
+    ) tmp, message_count mc
+GROUP BY year, isComment, lengthCategory, m_count
+ORDER BY year DESC, isComment ASC, lengthCategory ASC;
+----
+2010	false	0	63	0.0	0	0.9692307692307692
+2010	true	2	2	109.0	218	0.03076923076923077
+
+
+query II
+-FROM GRAPH_TABLE (snb_projected
+    MATCH (m:message)
+    COLUMNS (m.id)
+    ) tmp, (SELECT id from message limit 1)
+LIMIT 10;
+----
+618475290624	618475290624
+343597383683	618475290624
+343597383684	618475290624
+962072674309	618475290624
+962072674310	618475290624
+962072674311	618475290624
+962072674312	618475290624
+962072674313	618475290624
+962072674314	618475290624
+962072674315	618475290624
+
+query II
+-FROM (SELECT id from message limit 1), GRAPH_TABLE (snb_projected
+    MATCH (m:message)
+    COLUMNS (m.id)
+    ) tmp
+LIMIT 10;
+----
+618475290624	618475290624
+618475290624	343597383683
+618475290624	343597383684
+618475290624	962072674309
+618475290624	962072674310
+618475290624	962072674311
+618475290624	962072674312
+618475290624	962072674313
+618475290624	962072674314
+618475290624	962072674315