Skip to content

Commit 10e0c64

Browse files
authored
feat: optimize export_graph for polardb (#1337)
* feat:optimize export_graph * feat:optimize export_graph
1 parent 813ede3 commit 10e0c64

1 file changed

Lines changed: 26 additions & 36 deletions

File tree

src/memos/graph_dbs/polardb.py

Lines changed: 26 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1800,7 +1800,6 @@ def search_by_fulltext(
18001800
FROM "{self.db_name}_graph"."Memory" m
18011801
CROSS JOIN q
18021802
{where_clause_cte}
1803-
ORDER BY rank DESC
18041803
LIMIT {top_k};
18051804
"""
18061805
params = [tsquery_string]
@@ -2411,56 +2410,47 @@ def _extract_special_filter_values(filter_obj):
24112410
order_clause = """
24122411
ORDER BY ag_catalog.agtype_access_operator(properties, '"created_at"'::agtype) DESC NULLS LAST,id DESC
24132412
"""
2413+
count_query = f"""
2414+
SELECT COUNT(*) AS total_count
2415+
FROM "{self.db_name}_graph"."Memory"
2416+
{where_clause}
2417+
"""
24142418
if include_embedding:
2415-
node_query = f"""
2416-
WITH filtered AS (
2417-
SELECT id, properties, embedding
2418-
FROM "{self.db_name}_graph"."Memory"
2419-
{where_clause}
2420-
)
2421-
SELECT p.id, p.properties, p.embedding, c.total_count
2422-
FROM (SELECT COUNT(*) AS total_count FROM filtered) c
2423-
LEFT JOIN LATERAL (
2424-
SELECT id, properties, embedding
2425-
FROM filtered
2426-
{order_clause}
2427-
{pagination_clause}
2428-
) p ON TRUE
2419+
data_query = f"""
2420+
SELECT id, properties, embedding
2421+
FROM "{self.db_name}_graph"."Memory"
2422+
{where_clause}
2423+
{order_clause}
2424+
{pagination_clause}
24292425
"""
24302426
else:
2431-
node_query = f"""
2432-
WITH filtered AS (
2433-
SELECT id, properties
2434-
FROM "{self.db_name}_graph"."Memory"
2435-
{where_clause}
2436-
)
2437-
SELECT p.id, p.properties, c.total_count
2438-
FROM (SELECT COUNT(*) AS total_count FROM filtered) c
2439-
LEFT JOIN LATERAL (
2440-
SELECT id, properties
2441-
FROM filtered
2442-
{order_clause}
2443-
{pagination_clause}
2444-
) p ON TRUE
2427+
data_query = f"""
2428+
SELECT id, properties
2429+
FROM "{self.db_name}_graph"."Memory"
2430+
{where_clause}
2431+
{order_clause}
2432+
{pagination_clause}
24452433
"""
2446-
logger.info(f"[export_graph nodes] Query: {node_query}")
2434+
logger.info(f"[export_graph nodes] count_query: {count_query}")
2435+
logger.info(f"[export_graph nodes] data_query: {data_query}")
24472436

24482437
try:
24492438
with self._get_connection() as conn, conn.cursor() as cursor:
2450-
cursor.execute(node_query)
2439+
cursor.execute(count_query)
2440+
count_row = cursor.fetchone()
2441+
total_nodes = int(count_row[0]) if count_row and count_row[0] is not None else 0
2442+
2443+
cursor.execute(data_query)
24512444
node_results = cursor.fetchall()
24522445
nodes = []
24532446

24542447
for row in node_results:
24552448
if include_embedding:
2456-
row_id, properties_json, embedding_json, row_total_count = row
2449+
row_id, properties_json, embedding_json = row
24572450
else:
2458-
row_id, properties_json, row_total_count = row
2451+
row_id, properties_json = row
24592452
embedding_json = None
24602453

2461-
if row_total_count is not None:
2462-
total_nodes = int(row_total_count)
2463-
24642454
if row_id is None:
24652455
continue
24662456

0 commit comments

Comments
 (0)