ocient-graph 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,3211 @@
1
+ """
2
+ The oc_graph module provides the same operations as the Spark GraphX library,
3
+ but for an Ocient database. Graphs are represented as tables in the database,
4
+ with vertices and edges stored in separate tables.
5
+
6
+ A vertices table must contain a column named "id" that is not nullable and is
7
+ of type BIGINT. An edges table must contain two columns named "srcid" and
8
+ "destid" that are not nullable and are of type BIGINT. Besides these columns,
9
+ the vertices and edges tables can contain any number of additional columns
10
+ with any data types.
11
+ """
12
+
13
+ import uuid
14
+ from enum import Enum
15
+ from typing import List, Optional, Set
16
+
17
+ from pyocient import Connection, Cursor, SQLException
18
+
19
+ # -----------------------------------------------------------------------------
20
+ # Private Helper Functions
21
+ # -----------------------------------------------------------------------------
22
+
23
+
24
+ def _check_not_null(obj: object, message: str) -> None:
25
+ """Raises SQLException if the object is None."""
26
+ if obj is None:
27
+ raise SQLException(f"Argument Error: {message}")
28
+
29
+
30
+ def _check_not_null_or_empty(s: Optional[str], message: str) -> None:
31
+ """Raises SQLException if the string is None, empty, or only whitespace."""
32
+ if not s or not s.strip():
33
+ raise SQLException(f"Argument Error: {message}")
34
+
35
+
36
+ def _check_positive(value: int, message: str) -> None:
37
+ """Raises SQLException if the value is not positive."""
38
+ if value <= 0:
39
+ raise SQLException(f"Argument Error: {message}")
40
+
41
+
42
+ def _table_name(schema: str, table: str) -> str:
43
+ """Returns a fully qualified and quoted table name."""
44
+ return f'"{schema}"."{table}"'
45
+
46
+
47
+ def _check_table_conflict(result_schema: str, result_table: str, input_qualified_table_names: Set[str]) -> None:
48
+ """Raises SQLException if the result table name conflicts with any input table name."""
49
+ qualified_result_name = _table_name(result_schema, result_table)
50
+ if qualified_result_name in input_qualified_table_names:
51
+ raise SQLException(
52
+ f"Naming Conflict: Result table/view '{qualified_result_name}' cannot have the same name as an input table."
53
+ )
54
+
55
+
56
+ def _safe_drop_table(cursor: Cursor, qualified_table_name: str) -> None:
57
+ """Safely drops a table, ignoring errors if it doesn't exist."""
58
+ if cursor is None or not qualified_table_name:
59
+ return
60
+ try:
61
+ cursor.execute(f"DROP TABLE IF EXISTS {qualified_table_name}")
62
+ except SQLException:
63
+ # Ignore potential errors during cleanup
64
+ pass
65
+
66
+
67
+ def _safe_drop_view(cursor: Cursor, qualified_view_name: str) -> None:
68
+ """Safely drops a view, ignoring errors if it doesn't exist."""
69
+ if cursor is None or not qualified_view_name:
70
+ return
71
+ try:
72
+ cursor.execute(f"DROP VIEW IF EXISTS {qualified_view_name}")
73
+ except SQLException:
74
+ # Ignore potential errors during cleanup
75
+ pass
76
+
77
+
78
+ def _create_table_from_template(
79
+ template_schema: str,
80
+ template_table: str,
81
+ result_schema: str,
82
+ result_table: str,
83
+ result_indexes: List[str],
84
+ cursor: Cursor,
85
+ ) -> None:
86
+ """
87
+ Creates a new table based on an existing template table and optionally
88
+ creates indexes on specified columns.
89
+ """
90
+ template_table_name = _table_name(template_schema, template_table)
91
+ result_table_name = _table_name(result_schema, result_table)
92
+
93
+ # Create an empty copy of the template table
94
+ cursor.execute(f"CREATE TABLE {result_table_name} AS SELECT * FROM {template_table_name} WHERE 1=0")
95
+
96
+ # Create indexes on specified columns
97
+ for index_col in result_indexes:
98
+ _check_not_null_or_empty(index_col, "Index column name cannot be null or empty.")
99
+ clean_col_name = "".join(filter(str.isalnum, index_col))
100
+ index_name = f"idx_{result_table}_{clean_col_name}_{uuid.uuid4().hex[:8]}"
101
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_table_name} ("{index_col}")')
102
+
103
+
104
+ # -----------------------------------------------------------------------------
105
+ # Public API
106
+ # -----------------------------------------------------------------------------
107
+
108
+
109
+ class EdgeDirection(Enum):
110
+ """
111
+ Specifies the direction of edges to follow relative to a source vertex
112
+ when traversing a graph or collecting neighbors. Used to determine whether
113
+ to consider incoming edges, outgoing edges, or both.
114
+ """
115
+
116
+ IN = "IN"
117
+ """
118
+ Follow incoming edges only. Selects neighbors connected by an edge pointing
119
+ *towards* the source vertex (i.e., where the source vertex is the
120
+ destination of the edge).
121
+ """
122
+
123
+ OUT = "OUT"
124
+ """
125
+ Follow outgoing edges only. Selects neighbors connected by an edge
126
+ originating *from* the source vertex (i.e., where the source vertex is the
127
+ source of the edge).
128
+ """
129
+
130
+ BOTH = "BOTH"
131
+ """
132
+ Follow edges in both directions. Selects neighbors connected by either
133
+ incoming or outgoing edges relative to the source vertex.
134
+ """
135
+
136
+
137
+ def subgraph(
138
+ connection: Connection,
139
+ input_schema: str,
140
+ input_vertices_table: str,
141
+ input_edges_table: str,
142
+ result_schema: str,
143
+ result_vertices_table: str,
144
+ result_edges_table: str,
145
+ vertex_filter: str,
146
+ edge_filter: str,
147
+ result_vertices_indexes: List[str],
148
+ result_edges_indexes: List[str],
149
+ ) -> None:
150
+ """
151
+ Creates a subgraph consisting of new vertex and edge tables based on
152
+ filtering criteria applied to existing input graph tables.
153
+
154
+ This method performs the following steps:
155
+ 1. Validates input parameters.
156
+ 2. Checks for potential conflicts where result tables might overwrite inputs.
157
+ 3. Creates the structure for the result vertex and edge tables, copying the
158
+ schema from the corresponding input tables and applying specified indexes.
159
+ 4. Populates the result vertex table by selecting vertices from the input
160
+ vertex table that satisfy the `vertex_filter`.
161
+ 5. Populates the result edge table by selecting edges from the input edge
162
+ table that satisfy the `edge_filter` AND whose source and destination
163
+ vertices BOTH exist in the newly created (filtered) result vertex table.
164
+
165
+ Important Filter Logic:
166
+ - The `vertex_filter` is a standard SQL WHERE clause applied to the input
167
+ vertex table.
168
+ - The `edge_filter` is a standard SQL WHERE clause condition. Within this
169
+ filter string:
170
+ - Reference columns from the *source vertex* using the alias `a`.
171
+ - Reference columns from the *edge* itself using the alias `b`.
172
+ - Reference columns from the *destination vertex* using the alias `c`.
173
+ An edge is only included if it passes this filter AND its corresponding
174
+ source and destination vertices are present in the *result* vertex table.
175
+
176
+ Atomicity and Cleanup:
177
+ The operation attempts to be atomic. If any step fails, it will attempt to
178
+ drop the newly created result tables to avoid leaving partial results.
179
+
180
+ :param connection: The pyocient Connection object.
181
+ :param input_schema: The schema name for the original graph tables.
182
+ :param input_vertices_table: The table name for the original vertices.
183
+ :param input_edges_table: The table name for the original edges.
184
+ :param result_schema: The schema name for the resulting subgraph tables.
185
+ :param result_vertices_table: Name for the new table for filtered vertices.
186
+ :param result_edges_table: Name for the new table for filtered edges.
187
+ :param vertex_filter: An SQL WHERE clause (without "WHERE") to filter
188
+ vertices. Use "1=1" for no filter.
189
+ :param edge_filter: An SQL WHERE clause (without "WHERE") to filter edges.
190
+ Use aliases a, b, and c for source vertex, edge, and destination
191
+ vertex. Use "1=1" for no filter.
192
+ :param result_vertices_indexes: A list of column names for indexes on the
193
+ result vertices table.
194
+ :param result_edges_indexes: A list of column names for indexes on the
195
+ result edges table.
196
+ :raises SQLException: If any database error occurs or filters are invalid.
197
+ """
198
+ _check_not_null(connection, "Connection cannot be null.")
199
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
200
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
201
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
202
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
203
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
204
+ _check_not_null_or_empty(result_edges_table, "Result edges table name cannot be null or empty.")
205
+ _check_not_null(vertex_filter, "Vertex filter cannot be null (use '1=1').")
206
+ _check_not_null(edge_filter, "Edge filter cannot be null (use '1=1').")
207
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
208
+ _check_not_null(result_edges_indexes, "Result edges index list cannot be null.")
209
+
210
+ input_tables = {
211
+ _table_name(input_schema, input_vertices_table),
212
+ _table_name(input_schema, input_edges_table),
213
+ }
214
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
215
+ _check_table_conflict(result_schema, result_edges_table, input_tables)
216
+
217
+ qualified_result_vertices = _table_name(result_schema, result_vertices_table)
218
+ qualified_result_edges = _table_name(result_schema, result_edges_table)
219
+ operation_successful = False
220
+
221
+ with connection.cursor() as cursor:
222
+ try:
223
+ _create_table_from_template(
224
+ input_schema,
225
+ input_vertices_table,
226
+ result_schema,
227
+ result_vertices_table,
228
+ result_vertices_indexes,
229
+ cursor,
230
+ )
231
+ _create_table_from_template(
232
+ input_schema,
233
+ input_edges_table,
234
+ result_schema,
235
+ result_edges_table,
236
+ result_edges_indexes,
237
+ cursor,
238
+ )
239
+
240
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
241
+ v_filter_clause = (
242
+ f" WHERE {vertex_filter}" if vertex_filter.strip() and vertex_filter.strip() != "1=1" else ""
243
+ )
244
+ cursor.execute(
245
+ f"INSERT INTO {qualified_result_vertices} SELECT * FROM {input_vertices_name}{v_filter_clause}"
246
+ )
247
+
248
+ input_edges_name = _table_name(input_schema, input_edges_table)
249
+ e_filter_clause = f" AND ({edge_filter})" if edge_filter.strip() and edge_filter.strip() != "1=1" else ""
250
+
251
+ # This complex join ensures edges are included only if they pass the
252
+ # edge_filter AND their endpoints exist in the filtered result vertices.
253
+ insert_edges_sql = f"""
254
+ INSERT INTO {qualified_result_edges}
255
+ SELECT b.* FROM
256
+ {input_vertices_name} a,
257
+ {input_edges_name} b,
258
+ {input_vertices_name} c,
259
+ {qualified_result_vertices} d,
260
+ {qualified_result_vertices} e
261
+ WHERE a.id = b.srcid AND c.id = b.destid
262
+ AND d.id = b.srcid AND e.id = b.destid
263
+ {e_filter_clause}
264
+ """
265
+ cursor.execute(insert_edges_sql)
266
+
267
+ operation_successful = True
268
+ finally:
269
+ if not operation_successful:
270
+ _safe_drop_table(cursor, qualified_result_vertices)
271
+ _safe_drop_table(cursor, qualified_result_edges)
272
+
273
+
274
+ def join_vertices(
275
+ connection: Connection,
276
+ input_schema: str,
277
+ input_vertices_table: str,
278
+ modification_schema: str,
279
+ modification_vertices_table: str,
280
+ result_schema: str,
281
+ result_vertices_table: str,
282
+ result_attribute_expressions: List[str],
283
+ result_vertices_indexes: List[str],
284
+ ) -> None:
285
+ """
286
+ Joins two sets of vertices (input and modification) into a result table.
287
+
288
+ This method performs a join operation based on the vertex `id`.
289
+ 1. It inserts all vertices from `input_vertices_table` that are NOT
290
+ present (by `id`) in the `modification_vertices_table`.
291
+ 2. It inserts vertices that ARE present (by `id`) in BOTH tables. For these,
292
+ the `id` is from the input table, and other attributes are calculated
293
+ based on the provided `result_attribute_expressions`.
294
+
295
+ Usage of `result_attribute_expressions`:
296
+ - Each string is a SQL expression to calculate a column value.
297
+ - Use alias `a` to refer to columns from `input_vertices_table`.
298
+ - Use alias `b` to refer to columns from `modification_vertices_table`.
299
+ - It's highly recommended that each expression ends with `AS alias_name`.
300
+ - The order of expressions must match the column order of the result table
301
+ (after the 'id' column).
302
+
303
+ :param connection: The pyocient Connection object.
304
+ :param input_schema: Schema of the primary input vertices table.
305
+ :param input_vertices_table: Name of the primary input vertices table (alias `a`).
306
+ :param modification_schema: Schema of the modification vertices table.
307
+ :param modification_vertices_table: Name of the modification vertices table (alias `b`).
308
+ :param result_schema: Schema for the result table.
309
+ :param result_vertices_table: Name of the table for the joined vertices.
310
+ :param result_attribute_expressions: Ordered list of SQL expressions for
311
+ merged vertex attributes.
312
+ :param result_vertices_indexes: A list of column names for indexes on the
313
+ result table.
314
+ :raises SQLException: If any database error occurs.
315
+ """
316
+ _check_not_null(connection, "Connection cannot be null.")
317
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
318
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
319
+ _check_not_null_or_empty(modification_schema, "Modification schema cannot be null or empty.")
320
+ _check_not_null_or_empty(
321
+ modification_vertices_table,
322
+ "Modification vertices table name cannot be null or empty.",
323
+ )
324
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
325
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
326
+ _check_not_null(result_attribute_expressions, "Result attribute expressions list cannot be null.")
327
+ if not result_attribute_expressions:
328
+ raise SQLException("Argument Error: Result attribute expressions list cannot be empty.")
329
+
330
+ # ADDED VALIDATION
331
+ for expr in result_attribute_expressions:
332
+ _check_not_null_or_empty(expr, "Result attribute expression cannot be null or empty.")
333
+
334
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
335
+
336
+ input_tables = {
337
+ _table_name(input_schema, input_vertices_table),
338
+ _table_name(modification_schema, modification_vertices_table),
339
+ }
340
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
341
+
342
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
343
+ operation_successful = False
344
+
345
+ with connection.cursor() as cursor:
346
+ try:
347
+ _create_table_from_template(
348
+ input_schema,
349
+ input_vertices_table,
350
+ result_schema,
351
+ result_vertices_table,
352
+ result_vertices_indexes,
353
+ cursor,
354
+ )
355
+
356
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
357
+ mod_vertices_name = _table_name(modification_schema, modification_vertices_table)
358
+
359
+ # Insert vertices from input not in modification table
360
+ cursor.execute(
361
+ f"INSERT INTO {result_vertices_name} "
362
+ f"SELECT * FROM {input_vertices_name} "
363
+ f"WHERE id NOT IN (SELECT id FROM {mod_vertices_name})"
364
+ )
365
+
366
+ # Insert vertices present in both, applying expressions
367
+ expressions = ", ".join(result_attribute_expressions)
368
+ cursor.execute(
369
+ f"INSERT INTO {result_vertices_name} "
370
+ f"SELECT a.id, {expressions} "
371
+ f"FROM {input_vertices_name} a, {mod_vertices_name} b "
372
+ f"WHERE a.id = b.id"
373
+ )
374
+
375
+ operation_successful = True
376
+ finally:
377
+ if not operation_successful:
378
+ _safe_drop_table(cursor, result_vertices_name)
379
+
380
+
381
+ def map_vertices(
382
+ connection: Connection,
383
+ input_schema: str,
384
+ input_vertices_table: str,
385
+ result_schema: str,
386
+ result_vertices_table: str,
387
+ result_column_expressions: List[str],
388
+ result_vertices_indexes: List[str],
389
+ ) -> None:
390
+ """
391
+ Transforms vertex attributes from an input table to a new result table.
392
+
393
+ This function creates a new table containing vertices with transformed
394
+ attributes based on the data in the input table. The transformation is
395
+ defined by the SQL expressions provided in `result_column_expressions`.
396
+
397
+ The mandatory `id` column from the input table is always included as the
398
+ first column in the result.
399
+
400
+ Using `result_column_expressions`:
401
+ - Each string defines a column in the `SELECT` list.
402
+ - Expressions can refer to any column in the input vertices table by name.
403
+ - It is *highly recommended* to use an alias for each expression using
404
+ `AS alias_name` to ensure predictable column names in the result table.
405
+ - Do *not* include the `id` column in this list; it is handled automatically.
406
+
407
+ :param connection: The pyocient Connection object.
408
+ :param input_schema: Schema of the source vertices table.
409
+ :param input_vertices_table: Name of the source vertices table.
410
+ :param result_schema: Schema for the target vertices table.
411
+ :param result_vertices_table: Name for the target vertices table to create.
412
+ :param result_column_expressions: Ordered list of SQL expressions defining
413
+ the result columns (beyond `id`).
414
+ :param result_vertices_indexes: A list of column names for indexes on the
415
+ result table.
416
+ :raises SQLException: If any database error occurs.
417
+ """
418
+ _check_not_null(connection, "Connection cannot be null.")
419
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
420
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
421
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
422
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
423
+ _check_not_null(result_column_expressions, "Result column expressions list cannot be null.")
424
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
425
+
426
+ input_tables = {_table_name(input_schema, input_vertices_table)}
427
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
428
+
429
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
430
+ operation_successful = False
431
+
432
+ with connection.cursor() as cursor:
433
+ try:
434
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
435
+
436
+ all_columns = ["id"]
437
+ for expr in result_column_expressions:
438
+ _check_not_null_or_empty(expr, "Result column expression cannot be null or empty.")
439
+ all_columns.append(expr)
440
+
441
+ columns_list_str = ", ".join(all_columns)
442
+
443
+ create_query = (
444
+ f"CREATE TABLE {result_vertices_name} AS SELECT {columns_list_str} FROM {input_vertices_name} WHERE 1=0"
445
+ )
446
+ cursor.execute(create_query)
447
+
448
+ # Create indexes
449
+ for col in result_vertices_indexes:
450
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
451
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
452
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_vertices_name} ("{col}")')
453
+
454
+ # Populate table
455
+ insert_query = f"INSERT INTO {result_vertices_name} SELECT {columns_list_str} FROM {input_vertices_name}"
456
+ cursor.execute(insert_query)
457
+
458
+ operation_successful = True
459
+ finally:
460
+ if not operation_successful:
461
+ _safe_drop_table(cursor, result_vertices_name)
462
+
463
+
464
+ def map_edges(
465
+ connection: Connection,
466
+ input_schema: str,
467
+ input_edges_table: str,
468
+ result_schema: str,
469
+ result_edges_table: str,
470
+ result_column_expressions: List[str],
471
+ result_edges_indexes: List[str],
472
+ ) -> None:
473
+ """
474
+ Transforms edge attributes by creating a new edges table based on expressions.
475
+
476
+ This method creates a new table with `srcid`, `destid`, and additional
477
+ columns defined by `result_column_expressions`.
478
+
479
+ Using `result_column_expressions`:
480
+ - Each string is a SQL expression evaluated against the `input_edges_table`.
481
+ - Expressions can refer to any column in the input edges table by name.
482
+ - Each expression *must* define a result column name using `AS alias_name`.
483
+ - Do *not* include expressions for `srcid` or `destid`; they are handled
484
+ automatically.
485
+
486
+ :param connection: The pyocient Connection object.
487
+ :param input_schema: Schema of the source edges table.
488
+ :param input_edges_table: Name of the source edges table.
489
+ :param result_schema: Schema for the target edges table.
490
+ :param result_edges_table: Name for the target edges table to create.
491
+ :param result_column_expressions: List of SQL expressions defining the new
492
+ columns.
493
+ :param result_edges_indexes: A list of column names for indexes on the
494
+ result table.
495
+ :raises SQLException: If any database error occurs.
496
+ """
497
+ _check_not_null(connection, "Connection cannot be null.")
498
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
499
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
500
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
501
+ _check_not_null_or_empty(result_edges_table, "Result edges table name cannot be null or empty.")
502
+ _check_not_null(result_column_expressions, "Result column expressions list cannot be null.")
503
+ _check_not_null(result_edges_indexes, "Result edges index list cannot be null.")
504
+
505
+ input_tables = {_table_name(input_schema, input_edges_table)}
506
+ _check_table_conflict(result_schema, result_edges_table, input_tables)
507
+
508
+ result_edges_name = _table_name(result_schema, result_edges_table)
509
+ operation_successful = False
510
+
511
+ with connection.cursor() as cursor:
512
+ try:
513
+ input_edges_name = _table_name(input_schema, input_edges_table)
514
+
515
+ all_columns = ["srcid", "destid"]
516
+ for expr in result_column_expressions:
517
+ _check_not_null_or_empty(expr, "Result column expression cannot be null or empty.")
518
+ all_columns.append(expr)
519
+
520
+ columns_list_str = ", ".join(all_columns)
521
+
522
+ create_query = (
523
+ f"CREATE TABLE {result_edges_name} AS SELECT {columns_list_str} FROM {input_edges_name} WHERE 1=0"
524
+ )
525
+ cursor.execute(create_query)
526
+
527
+ # Create indexes
528
+ for col in result_edges_indexes:
529
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
530
+ index_name = f"{result_edges_table}_{col}_idx".replace('"', "")
531
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_edges_name} ("{col}")')
532
+
533
+ # Populate table
534
+ insert_query = f"INSERT INTO {result_edges_name} SELECT {columns_list_str} FROM {input_edges_name}"
535
+ cursor.execute(insert_query)
536
+
537
+ operation_successful = True
538
+ finally:
539
+ if not operation_successful:
540
+ _safe_drop_table(cursor, result_edges_name)
541
+
542
+
543
+ def map_triplets(
544
+ connection: Connection,
545
+ input_schema: str,
546
+ input_vertices_table: str,
547
+ input_edges_table: str,
548
+ result_schema: str,
549
+ result_edges_table: str,
550
+ result_column_expressions: List[str],
551
+ result_edges_indexes: List[str],
552
+ ) -> None:
553
+ """
554
+ Transforms edge attributes by incorporating information from the source and
555
+ destination vertices (the "triplet").
556
+
557
+ This method creates a new edge table based on joining the input edge table
558
+ with the input vertex table for both source and destination vertices.
559
+
560
+ Using `result_column_expressions`:
561
+ - Each string is a SQL expression defining a new column.
562
+ - Expressions *must* reference columns using specific aliases:
563
+ - `a`: for the source vertex.
564
+ - `b`: for the edge.
565
+ - `c`: for the destination vertex.
566
+ - It is *highly recommended* that each expression use `AS alias_name` to
567
+ define the result column name.
568
+ - Do *not* include `srcid` or `destid`; they are handled automatically.
569
+
570
+ :param connection: The pyocient Connection object.
571
+ :param input_schema: Schema of the source graph tables.
572
+ :param input_vertices_table: Name of the source vertices table.
573
+ :param input_edges_table: Name of the source edges table.
574
+ :param result_schema: Schema for the target edges table.
575
+ :param result_edges_table: Name for the target edges table to create.
576
+ :param result_column_expressions: List of SQL expressions defining the new
577
+ columns.
578
+ :param result_edges_indexes: A list of column names for indexes on the
579
+ result table.
580
+ :raises SQLException: If any database error occurs.
581
+ """
582
+ _check_not_null(connection, "Connection cannot be null.")
583
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
584
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
585
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
586
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
587
+ _check_not_null_or_empty(result_edges_table, "Result edges table name cannot be null or empty.")
588
+ _check_not_null(result_column_expressions, "Result column expressions list cannot be null.")
589
+ _check_not_null(result_edges_indexes, "Result edges index list cannot be null.")
590
+
591
+ input_tables = {
592
+ _table_name(input_schema, input_vertices_table),
593
+ _table_name(input_schema, input_edges_table),
594
+ }
595
+ _check_table_conflict(result_schema, result_edges_table, input_tables)
596
+
597
+ result_edges_name = _table_name(result_schema, result_edges_table)
598
+ operation_successful = False
599
+
600
+ with connection.cursor() as cursor:
601
+ try:
602
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
603
+ input_edges_name = _table_name(input_schema, input_edges_table)
604
+
605
+ select_expressions = ["b.srcid", "b.destid"]
606
+ for expr in result_column_expressions:
607
+ _check_not_null_or_empty(expr, "Result column expression cannot be null or empty.")
608
+ select_expressions.append(expr)
609
+
610
+ select_list_str = ", ".join(select_expressions)
611
+
612
+ from_join_clause = (
613
+ f"FROM {input_vertices_name} a "
614
+ f"JOIN {input_edges_name} b ON a.id = b.srcid "
615
+ f"JOIN {input_vertices_name} c ON c.id = b.destid"
616
+ )
617
+
618
+ create_query = f"CREATE TABLE {result_edges_name} AS SELECT {select_list_str} {from_join_clause} WHERE 1=0"
619
+ cursor.execute(create_query)
620
+
621
+ # Create indexes
622
+ for col in result_edges_indexes:
623
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
624
+ index_name = f"{result_edges_table}_{col}_idx".replace('"', "")
625
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_edges_name} ("{col}")')
626
+
627
+ # Populate table
628
+ insert_query = f"INSERT INTO {result_edges_name} SELECT {select_list_str} {from_join_clause}"
629
+ cursor.execute(insert_query)
630
+
631
+ operation_successful = True
632
+ finally:
633
+ if not operation_successful:
634
+ _safe_drop_table(cursor, result_edges_name)
635
+
636
+
637
+ def create_triplets_view(
638
+ connection: Connection,
639
+ input_schema: str,
640
+ input_vertices_table: str,
641
+ input_edges_table: str,
642
+ result_schema: str,
643
+ result_triplets_view: str,
644
+ ) -> None:
645
+ """
646
+ Creates a database view representing graph triplets (source vertex, edge,
647
+ destination vertex).
648
+
649
+ The view joins the edges table with the vertices table twice (for source
650
+ and destination). It includes all columns from the original edges table,
651
+ plus all columns from the vertices table (except `id`), prefixed with
652
+ `src_` and `dest_` respectively.
653
+
654
+ :param connection: The pyocient Connection object.
655
+ :param input_schema: Schema of the source graph tables.
656
+ :param input_vertices_table: Name of the source vertices table.
657
+ :param input_edges_table: Name of the source edges table.
658
+ :param result_schema: Schema for the new triplets view.
659
+ :param result_triplets_view: Name for the new view to be created.
660
+ :raises SQLException: If any database error occurs.
661
+ """
662
+ _check_not_null(connection, "Connection cannot be null.")
663
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
664
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
665
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
666
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
667
+ _check_not_null_or_empty(result_triplets_view, "Result triplets view name cannot be null or empty.")
668
+
669
+ input_tables = {
670
+ _table_name(input_schema, input_vertices_table),
671
+ _table_name(input_schema, input_edges_table),
672
+ }
673
+ _check_table_conflict(result_schema, result_triplets_view, input_tables)
674
+
675
+ result_triplets_name = _table_name(result_schema, result_triplets_view)
676
+ operation_successful = False
677
+
678
+ with connection.cursor() as cursor:
679
+ try:
680
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
681
+ input_edges_name = _table_name(input_schema, input_edges_table)
682
+
683
+ # Get vertex columns (excluding id) to build the select list
684
+ cursor.execute(f"SELECT * FROM {input_vertices_name} WHERE 1=0")
685
+ vertex_columns = [col_desc[0] for col_desc in (cursor.description or []) if col_desc[0].lower() != "id"]
686
+
687
+ src_columns = [f'a."{col}" AS src_{col}' for col in vertex_columns]
688
+ dest_columns = [f'c."{col}" AS dest_{col}' for col in vertex_columns]
689
+
690
+ all_select_list = ["b.*"] + src_columns + dest_columns
691
+ all_columns_str = ", ".join(all_select_list)
692
+
693
+ create_view_sql = f"""
694
+ CREATE VIEW {result_triplets_name} AS
695
+ SELECT {all_columns_str}
696
+ FROM {input_edges_name} b
697
+ JOIN {input_vertices_name} a ON b.srcid = a.id
698
+ JOIN {input_vertices_name} c ON b.destid = c.id
699
+ """
700
+ cursor.execute(create_view_sql)
701
+ operation_successful = True
702
+ finally:
703
+ if not operation_successful:
704
+ _safe_drop_view(cursor, result_triplets_name)
705
+
706
+
707
+ def create_triplets_table(
708
+ connection: Connection,
709
+ input_schema: str,
710
+ input_vertices_table: str,
711
+ input_edges_table: str,
712
+ result_schema: str,
713
+ result_triplets_table: str,
714
+ result_triplets_indexes: List[str],
715
+ ) -> None:
716
+ """
717
+ Creates a materialized triplet table by joining vertex and edge data.
718
+
719
+ This method is similar to `create_triplets_view` but creates a physical
720
+ table instead of a view. It includes all columns from the edges table, plus
721
+ prefixed columns (`src_`, `dest_`) for attributes from the source and
722
+ destination vertices.
723
+
724
+ :param connection: The pyocient Connection object.
725
+ :param input_schema: Schema of the source graph tables.
726
+ :param input_vertices_table: Name of the source vertices table.
727
+ :param input_edges_table: Name of the source edges table.
728
+ :param result_schema: Schema for the new triplets table.
729
+ :param result_triplets_table: Name for the new table to be created.
730
+ :param result_triplets_indexes: List of column names to index in the result.
731
+ :raises SQLException: If any database error occurs.
732
+ """
733
+ _check_not_null(connection, "Connection cannot be null.")
734
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
735
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
736
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
737
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
738
+ _check_not_null_or_empty(result_triplets_table, "Result triplets table name cannot be null or empty.")
739
+ _check_not_null(result_triplets_indexes, "Result triplets index list cannot be null.")
740
+
741
+ input_tables = {
742
+ _table_name(input_schema, input_vertices_table),
743
+ _table_name(input_schema, input_edges_table),
744
+ }
745
+ _check_table_conflict(result_schema, result_triplets_table, input_tables)
746
+
747
+ result_triplets_name = _table_name(result_schema, result_triplets_table)
748
+ operation_successful = False
749
+
750
+ with connection.cursor() as cursor:
751
+ try:
752
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
753
+ input_edges_name = _table_name(input_schema, input_edges_table)
754
+
755
+ cursor.execute(f"SELECT * FROM {input_vertices_name} WHERE 1=0")
756
+ vertex_columns = [col_desc[0] for col_desc in (cursor.description or []) if col_desc[0].lower() != "id"]
757
+
758
+ src_columns = [f'a."{col}" AS src_{col}' for col in vertex_columns]
759
+ dest_columns = [f'c."{col}" AS dest_{col}' for col in vertex_columns]
760
+ all_select_list = ["b.*"] + src_columns + dest_columns
761
+ select_list_str = ", ".join(all_select_list)
762
+
763
+ from_join_clause = (
764
+ f"FROM {input_edges_name} b "
765
+ f"JOIN {input_vertices_name} a ON b.srcid = a.id "
766
+ f"JOIN {input_vertices_name} c ON b.destid = c.id"
767
+ )
768
+
769
+ create_table_sql = (
770
+ f"CREATE TABLE {result_triplets_name} AS SELECT {select_list_str} {from_join_clause} WHERE 1=0"
771
+ )
772
+ cursor.execute(create_table_sql)
773
+
774
+ for col in result_triplets_indexes:
775
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
776
+ index_name = f"{result_triplets_table}_{col}_idx".replace('"', "")
777
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_triplets_name} ("{col}")')
778
+
779
+ insert_sql = f"INSERT INTO {result_triplets_name} SELECT {select_list_str} {from_join_clause}"
780
+ cursor.execute(insert_sql)
781
+
782
+ operation_successful = True
783
+ finally:
784
+ if not operation_successful:
785
+ _safe_drop_table(cursor, result_triplets_name)
786
+
787
+
788
+ def reverse_edges(
789
+ connection: Connection,
790
+ input_schema: str,
791
+ input_edges_table: str,
792
+ result_schema: str,
793
+ result_edges_table: str,
794
+ result_edges_indexes: List[str],
795
+ ) -> None:
796
+ """
797
+ Reverses the direction of edges stored in a database table.
798
+
799
+ Creates a new result table with the same columns as the input table, but
800
+ the values in `srcid` and `destid` are swapped. All other columns are
801
+ copied verbatim.
802
+
803
+ :param connection: The pyocient Connection object.
804
+ :param input_schema: Schema of the source edges table.
805
+ :param input_edges_table: Name of the source edges table.
806
+ :param result_schema: Schema for the target table.
807
+ :param result_edges_table: Name for the table with reversed edges.
808
+ :param result_edges_indexes: List of columns to index in the result table.
809
+ :raises SQLException: If any database error occurs.
810
+ """
811
+ _check_not_null(connection, "Connection cannot be null.")
812
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
813
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
814
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
815
+ _check_not_null_or_empty(result_edges_table, "Result edges table name cannot be null or empty.")
816
+ _check_not_null(result_edges_indexes, "Result edges index list cannot be null.")
817
+
818
+ input_tables = {_table_name(input_schema, input_edges_table)}
819
+ _check_table_conflict(result_schema, result_edges_table, input_tables)
820
+
821
+ result_edges_name = _table_name(result_schema, result_edges_table)
822
+ operation_successful = False
823
+
824
+ with connection.cursor() as cursor:
825
+ try:
826
+ _create_table_from_template(
827
+ input_schema,
828
+ input_edges_table,
829
+ result_schema,
830
+ result_edges_table,
831
+ result_edges_indexes,
832
+ cursor,
833
+ )
834
+
835
+ input_edges_name = _table_name(input_schema, input_edges_table)
836
+ cursor.execute(f"SELECT * FROM {input_edges_name} WHERE 1=0")
837
+ other_columns = [
838
+ desc[0] for desc in (cursor.description or []) if desc[0].lower() not in ("srcid", "destid")
839
+ ]
840
+
841
+ other_cols_quoted_str = ", ".join([f'"{c}"' for c in other_columns])
842
+ select_clause = f"destid AS srcid, srcid AS destid"
843
+ if other_cols_quoted_str:
844
+ select_clause += f", {other_cols_quoted_str}"
845
+
846
+ target_cols_str = "srcid, destid"
847
+ if other_cols_quoted_str:
848
+ target_cols_str += f", {other_cols_quoted_str}"
849
+
850
+ insert_sql = (
851
+ f"INSERT INTO {result_edges_name} ({target_cols_str}) SELECT {select_clause} FROM {input_edges_name}"
852
+ )
853
+ cursor.execute(insert_sql)
854
+
855
+ operation_successful = True
856
+ finally:
857
+ if not operation_successful:
858
+ _safe_drop_table(cursor, result_edges_name)
859
+
860
+
861
+ def filter_edges(
862
+ connection: Connection,
863
+ input_schema: str,
864
+ input_edges_table: str,
865
+ result_schema: str,
866
+ result_edges_table: str,
867
+ edge_filter: str,
868
+ result_edges_indexes: List[str],
869
+ ) -> None:
870
+ """
871
+ Filters edges from a source table into a new result table.
872
+
873
+ The `edge_filter` is a standard SQL conditional expression (without the
874
+ `WHERE` keyword) that references columns in the `input_edges_table`.
875
+ Use "1=1" for no filtering.
876
+
877
+ :param connection: The pyocient Connection object.
878
+ :param input_schema: Schema of the source edges table.
879
+ :param input_edges_table: Name of the source edges table.
880
+ :param result_schema: Schema for the target table.
881
+ :param result_edges_table: Name for the table with filtered edges.
882
+ :param edge_filter: SQL predicate string to filter edges.
883
+ :param result_edges_indexes: List of columns to index in the result table.
884
+ :raises SQLException: If any database error occurs.
885
+ """
886
+ _check_not_null(connection, "Connection cannot be null.")
887
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
888
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
889
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
890
+ _check_not_null_or_empty(result_edges_table, "Result edges table name cannot be null or empty.")
891
+ _check_not_null(edge_filter, "Edge filter cannot be null (use '1=1').")
892
+ _check_not_null(result_edges_indexes, "Result edges index list cannot be null.")
893
+
894
+ input_tables = {_table_name(input_schema, input_edges_table)}
895
+ _check_table_conflict(result_schema, result_edges_table, input_tables)
896
+
897
+ result_edges_name = _table_name(result_schema, result_edges_table)
898
+ operation_successful = False
899
+
900
+ with connection.cursor() as cursor:
901
+ try:
902
+ _create_table_from_template(
903
+ input_schema,
904
+ input_edges_table,
905
+ result_schema,
906
+ result_edges_table,
907
+ result_edges_indexes,
908
+ cursor,
909
+ )
910
+
911
+ input_edges_name = _table_name(input_schema, input_edges_table)
912
+ filter_clause = ""
913
+ if edge_filter.strip() and edge_filter.strip() != "1=1":
914
+ filter_clause = f" WHERE {edge_filter}"
915
+
916
+ cursor.execute(f"INSERT INTO {result_edges_name} SELECT * FROM {input_edges_name}{filter_clause}")
917
+
918
+ operation_successful = True
919
+ finally:
920
+ if not operation_successful:
921
+ _safe_drop_table(cursor, result_edges_name)
922
+
923
+
924
+ def filter_vertices(
925
+ connection: Connection,
926
+ input_schema: str,
927
+ input_vertices_table: str,
928
+ input_edges_table: str,
929
+ result_schema: str,
930
+ result_vertices_table: str,
931
+ result_edges_table: str,
932
+ vertex_filter: str,
933
+ result_vertices_indexes: List[str],
934
+ result_edges_indexes: List[str],
935
+ ) -> None:
936
+ """
937
+ Creates a new subgraph by filtering vertices and retaining their edges.
938
+
939
+ Populates `result_vertices_table` with vertices that satisfy `vertex_filter`.
940
+ Then populates `result_edges_table` with edges from `input_edges_table`
941
+ where *both* endpoints are present in the new `result_vertices_table`.
942
+
943
+ :param connection: The pyocient Connection object.
944
+ :param input_schema: Schema of the source graph tables.
945
+ :param input_vertices_table: Name of the source vertices table.
946
+ :param input_edges_table: Name of the source edges table.
947
+ :param result_schema: Schema for the target tables.
948
+ :param result_vertices_table: Name for the table with filtered vertices.
949
+ :param result_edges_table: Name for the table with corresponding edges.
950
+ :param vertex_filter: SQL predicate string to filter vertices.
951
+ :param result_vertices_indexes: List of columns to index in the result vertices table.
952
+ :param result_edges_indexes: List of columns to index in the result edges table.
953
+ :raises SQLException: If any database error occurs.
954
+ """
955
+ _check_not_null(connection, "Connection cannot be null.")
956
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
957
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
958
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
959
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
960
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
961
+ _check_not_null_or_empty(result_edges_table, "Result edges table name cannot be null or empty.")
962
+ _check_not_null(vertex_filter, "Vertex filter cannot be null (use '1=1').")
963
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
964
+ _check_not_null(result_edges_indexes, "Result edges index list cannot be null.")
965
+
966
+ input_tables = {_table_name(input_schema, input_vertices_table), _table_name(input_schema, input_edges_table)}
967
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
968
+ _check_table_conflict(result_schema, result_edges_table, input_tables)
969
+
970
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
971
+ result_edges_name = _table_name(result_schema, result_edges_table)
972
+ operation_successful = False
973
+
974
+ with connection.cursor() as cursor:
975
+ try:
976
+ _create_table_from_template(
977
+ input_schema,
978
+ input_vertices_table,
979
+ result_schema,
980
+ result_vertices_table,
981
+ result_vertices_indexes,
982
+ cursor,
983
+ )
984
+ _create_table_from_template(
985
+ input_schema,
986
+ input_edges_table,
987
+ result_schema,
988
+ result_edges_table,
989
+ result_edges_indexes,
990
+ cursor,
991
+ )
992
+
993
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
994
+ filter_clause = ""
995
+ if vertex_filter.strip() and vertex_filter.strip() != "1=1":
996
+ filter_clause = f" WHERE {vertex_filter}"
997
+
998
+ cursor.execute(f"INSERT INTO {result_vertices_name} SELECT * FROM {input_vertices_name}{filter_clause}")
999
+
1000
+ input_edges_name = _table_name(input_schema, input_edges_table)
1001
+ cursor.execute(
1002
+ f"INSERT INTO {result_edges_name} "
1003
+ f"SELECT * FROM {input_edges_name} "
1004
+ f"WHERE srcid IN (SELECT id FROM {result_vertices_name}) "
1005
+ f"AND destid IN (SELECT id FROM {result_vertices_name})"
1006
+ )
1007
+
1008
+ operation_successful = True
1009
+ finally:
1010
+ if not operation_successful:
1011
+ _safe_drop_table(cursor, result_vertices_name)
1012
+ _safe_drop_table(cursor, result_edges_name)
1013
+
1014
+
1015
+ def in_degrees(
1016
+ connection: Connection,
1017
+ input_schema: str,
1018
+ input_edges_table: str,
1019
+ result_schema: str,
1020
+ result_in_degrees_table: str,
1021
+ result_indexes: List[str],
1022
+ ) -> None:
1023
+ """
1024
+ Calculates the in-degree for each destination node in an edge table.
1025
+
1026
+ The result table will contain two columns: `id` (the vertex ID) and
1027
+ `in_degree` (the count of incoming edges).
1028
+
1029
+ :param connection: The pyocient Connection object.
1030
+ :param input_schema: Schema of the source edges table.
1031
+ :param input_edges_table: Name of the source edges table.
1032
+ :param result_schema: Schema for the target table.
1033
+ :param result_in_degrees_table: Name for the table with in-degree counts.
1034
+ :param result_indexes: List of columns to index ('id' or 'in_degree').
1035
+ :raises SQLException: If any database error occurs.
1036
+ """
1037
+ _check_not_null(connection, "Connection cannot be null.")
1038
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1039
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
1040
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1041
+ _check_not_null_or_empty(result_in_degrees_table, "Result in-degrees table name cannot be null or empty.")
1042
+ _check_not_null(result_indexes, "Result index list cannot be null.")
1043
+
1044
+ input_tables = {_table_name(input_schema, input_edges_table)}
1045
+ _check_table_conflict(result_schema, result_in_degrees_table, input_tables)
1046
+
1047
+ result_table_name = _table_name(result_schema, result_in_degrees_table)
1048
+ operation_successful = False
1049
+
1050
+ with connection.cursor() as cursor:
1051
+ try:
1052
+ create_sql = f"CREATE TABLE {result_table_name} (id BIGINT NOT NULL, in_degree BIGINT NOT NULL)"
1053
+ cursor.execute(create_sql)
1054
+
1055
+ for col in result_indexes:
1056
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1057
+ if col.lower() not in ("id", "in_degree"):
1058
+ raise SQLException(f"Argument Error: Invalid index column '{col}' for inDegrees.")
1059
+ index_name = f"{result_in_degrees_table}_{col}_idx".replace('"', "")
1060
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_table_name} ("{col}")')
1061
+
1062
+ input_edges_name = _table_name(input_schema, input_edges_table)
1063
+ insert_sql = (
1064
+ f"INSERT INTO {result_table_name} (id, in_degree) "
1065
+ f"SELECT destid AS id, COUNT(*) AS in_degree "
1066
+ f"FROM {input_edges_name} "
1067
+ f"GROUP BY destid"
1068
+ )
1069
+ cursor.execute(insert_sql)
1070
+
1071
+ operation_successful = True
1072
+ finally:
1073
+ if not operation_successful:
1074
+ _safe_drop_table(cursor, result_table_name)
1075
+
1076
+
1077
+ def out_degrees(
1078
+ connection: Connection,
1079
+ input_schema: str,
1080
+ input_edges_table: str,
1081
+ result_schema: str,
1082
+ result_out_degrees_table: str,
1083
+ result_indexes: List[str],
1084
+ ) -> None:
1085
+ """
1086
+ Calculates the out-degree for each source node in an edge table.
1087
+
1088
+ The result table will contain two columns: `id` (the vertex ID) and
1089
+ `out_degree` (the count of outgoing edges).
1090
+
1091
+ :param connection: The pyocient Connection object.
1092
+ :param input_schema: Schema of the source edges table.
1093
+ :param input_edges_table: Name of the source edges table.
1094
+ :param result_schema: Schema for the target table.
1095
+ :param result_out_degrees_table: Name for the table with out-degree counts.
1096
+ :param result_indexes: List of columns to index ('id' or 'out_degree').
1097
+ :raises SQLException: If any database error occurs.
1098
+ """
1099
+ _check_not_null(connection, "Connection cannot be null.")
1100
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1101
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
1102
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1103
+ _check_not_null_or_empty(result_out_degrees_table, "Result out-degrees table name cannot be null or empty.")
1104
+ _check_not_null(result_indexes, "Result index list cannot be null.")
1105
+
1106
+ input_tables = {_table_name(input_schema, input_edges_table)}
1107
+ _check_table_conflict(result_schema, result_out_degrees_table, input_tables)
1108
+
1109
+ result_table_name = _table_name(result_schema, result_out_degrees_table)
1110
+ operation_successful = False
1111
+
1112
+ with connection.cursor() as cursor:
1113
+ try:
1114
+ create_sql = f"CREATE TABLE {result_table_name} (id BIGINT NOT NULL, out_degree BIGINT NOT NULL)"
1115
+ cursor.execute(create_sql)
1116
+
1117
+ for col in result_indexes:
1118
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1119
+ if col.lower() not in ("id", "out_degree"):
1120
+ raise SQLException(f"Argument Error: Invalid index column '{col}' for outDegrees.")
1121
+ index_name = f"{result_out_degrees_table}_{col}_idx".replace('"', "")
1122
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_table_name} ("{col}")')
1123
+
1124
+ input_edges_name = _table_name(input_schema, input_edges_table)
1125
+ insert_sql = (
1126
+ f"INSERT INTO {result_table_name} (id, out_degree) "
1127
+ f"SELECT srcid AS id, COUNT(*) AS out_degree "
1128
+ f"FROM {input_edges_name} "
1129
+ f"GROUP BY srcid"
1130
+ )
1131
+ cursor.execute(insert_sql)
1132
+
1133
+ operation_successful = True
1134
+ finally:
1135
+ if not operation_successful:
1136
+ _safe_drop_table(cursor, result_table_name)
1137
+
1138
+
1139
+ def degrees(
1140
+ connection: Connection,
1141
+ input_schema: str,
1142
+ input_edges_table: str,
1143
+ result_schema: str,
1144
+ result_degrees_table: str,
1145
+ result_indexes: List[str],
1146
+ ) -> None:
1147
+ """
1148
+ Calculates the total degree (in + out) for each node in an edge table.
1149
+
1150
+ The result table will contain two columns: `id` (the vertex ID) and
1151
+ `degree` (the total count of incident edges).
1152
+
1153
+ :param connection: The pyocient Connection object.
1154
+ :param input_schema: Schema of the source edges table.
1155
+ :param input_edges_table: Name of the source edges table.
1156
+ :param result_schema: Schema for the target table.
1157
+ :param result_degrees_table: Name for the table with total degree counts.
1158
+ :param result_indexes: List of columns to index ('id' or 'degree').
1159
+ :raises SQLException: If any database error occurs.
1160
+ """
1161
+ _check_not_null(connection, "Connection cannot be null.")
1162
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1163
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
1164
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1165
+ _check_not_null_or_empty(result_degrees_table, "Result degrees table name cannot be null or empty.")
1166
+ _check_not_null(result_indexes, "Result index list cannot be null.")
1167
+
1168
+ input_tables = {_table_name(input_schema, input_edges_table)}
1169
+ _check_table_conflict(result_schema, result_degrees_table, input_tables)
1170
+
1171
+ result_table_name = _table_name(result_schema, result_degrees_table)
1172
+ operation_successful = False
1173
+
1174
+ with connection.cursor() as cursor:
1175
+ try:
1176
+ create_sql = f"CREATE TABLE {result_table_name} (id BIGINT NOT NULL, degree BIGINT NOT NULL)"
1177
+ cursor.execute(create_sql)
1178
+
1179
+ for col in result_indexes:
1180
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1181
+ if col.lower() not in ("id", "degree"):
1182
+ raise SQLException(f"Argument Error: Invalid index column '{col}' for degrees.")
1183
+ index_name = f"{result_degrees_table}_{col}_idx".replace('"', "")
1184
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_table_name} ("{col}")')
1185
+
1186
+ input_edges_name = _table_name(input_schema, input_edges_table)
1187
+ insert_sql = f"""
1188
+ INSERT INTO {result_table_name} (id, degree)
1189
+ SELECT id, COUNT(*) AS degree FROM (
1190
+ SELECT srcid AS id FROM {input_edges_name}
1191
+ UNION ALL
1192
+ SELECT destid AS id FROM {input_edges_name}
1193
+ ) combined_ids GROUP BY id
1194
+ """
1195
+ cursor.execute(insert_sql)
1196
+
1197
+ operation_successful = True
1198
+ finally:
1199
+ if not operation_successful:
1200
+ _safe_drop_table(cursor, result_table_name)
1201
+
1202
+
1203
+ def from_edges(
1204
+ connection: Connection,
1205
+ input_schema: str,
1206
+ input_edges_table: str,
1207
+ result_schema: str,
1208
+ result_vertices_table: str,
1209
+ result_column_expressions: List[str],
1210
+ result_vertices_indexes: List[str],
1211
+ ) -> None:
1212
+ """
1213
+ Creates a vertices table from an existing edges table by extracting
1214
+ unique source and destination IDs.
1215
+
1216
+ Using `result_column_expressions`:
1217
+ - Expressions can define additional columns in the result table.
1218
+ - The unique vertex ID is available within expressions via the alias `ids.id`.
1219
+ - Each expression *must* use `AS alias_name` to name the result column.
1220
+ - Example: `["ids.id % 100 AS bucket_id"]`
1221
+
1222
+ :param connection: The pyocient Connection object.
1223
+ :param input_schema: Schema of the source edges table.
1224
+ :param input_edges_table: Name of the source edges table.
1225
+ :param result_schema: Schema for the target vertices table.
1226
+ :param result_vertices_table: Name for the target vertices table.
1227
+ :param result_column_expressions: List of SQL expressions for additional columns.
1228
+ :param result_vertices_indexes: List of columns to index in the result table.
1229
+ :raises SQLException: If any database error occurs.
1230
+ """
1231
+ _check_not_null(connection, "Connection cannot be null.")
1232
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1233
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
1234
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1235
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
1236
+ _check_not_null(result_column_expressions, "Result column expressions list cannot be null.")
1237
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
1238
+
1239
+ input_tables = {_table_name(input_schema, input_edges_table)}
1240
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
1241
+
1242
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
1243
+ operation_successful = False
1244
+
1245
+ with connection.cursor() as cursor:
1246
+ try:
1247
+ input_edges_name = _table_name(input_schema, input_edges_table)
1248
+ id_sub_query = (
1249
+ f"(SELECT srcid AS id FROM {input_edges_name} "
1250
+ f"UNION DISTINCT SELECT destid AS id FROM {input_edges_name}) ids"
1251
+ )
1252
+
1253
+ columns_builder = ["ids.id"]
1254
+ for expr in result_column_expressions:
1255
+ _check_not_null_or_empty(expr, "Result column expression cannot be null or empty.")
1256
+ columns_builder.append(expr)
1257
+ columns_list_str = ", ".join(columns_builder)
1258
+
1259
+ create_sql = (
1260
+ f"CREATE TABLE {result_vertices_name} AS SELECT {columns_list_str} FROM {id_sub_query} WHERE 1=0"
1261
+ )
1262
+ cursor.execute(create_sql)
1263
+
1264
+ for col in result_vertices_indexes:
1265
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1266
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
1267
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_vertices_name} ("{col}")')
1268
+
1269
+ insert_sql = f"INSERT INTO {result_vertices_name} SELECT {columns_list_str} FROM {id_sub_query}"
1270
+ cursor.execute(insert_sql)
1271
+
1272
+ operation_successful = True
1273
+ finally:
1274
+ if not operation_successful:
1275
+ _safe_drop_table(cursor, result_vertices_name)
1276
+
1277
+
1278
+ def inner_join_vertices(
1279
+ connection: Connection,
1280
+ input_schema: str,
1281
+ input_vertices_table: str,
1282
+ other_schema: str,
1283
+ other_vertices_table: str,
1284
+ result_schema: str,
1285
+ result_vertices_table: str,
1286
+ result_column_expressions: List[str],
1287
+ result_vertices_indexes: List[str],
1288
+ ) -> None:
1289
+ """
1290
+ Performs an SQL INNER JOIN between two vertex tables on their 'id' columns.
1291
+
1292
+ Using `result_column_expressions`:
1293
+ - The `id` column from the first table is always included.
1294
+ - Each string defines an additional column in the `SELECT` list.
1295
+ - Use alias `a` for columns from `input_vertices_table`.
1296
+ - Use alias `b` for columns from `other_vertices_table`.
1297
+ - It is highly recommended to use `AS alias_name` for each expression.
1298
+
1299
+ :param connection: The pyocient Connection object.
1300
+ :param input_schema: Schema of the first (left) vertices table.
1301
+ :param input_vertices_table: Name of the first vertices table (alias `a`).
1302
+ :param other_schema: Schema of the second (right) vertices table.
1303
+ :param other_vertices_table: Name of the second vertices table (alias `b`).
1304
+ :param result_schema: Schema for the result table.
1305
+ :param result_vertices_table: Name for the table with join results.
1306
+ :param result_column_expressions: List of SQL expressions for result columns.
1307
+ :param result_vertices_indexes: List of columns to index in the result table.
1308
+ :raises SQLException: If any database error occurs.
1309
+ """
1310
+ _check_not_null(connection, "Connection cannot be null.")
1311
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1312
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
1313
+ _check_not_null_or_empty(other_schema, "Other schema cannot be null or empty.")
1314
+ _check_not_null_or_empty(other_vertices_table, "Other vertices table name cannot be null or empty.")
1315
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1316
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
1317
+ _check_not_null(result_column_expressions, "Result column expressions list cannot be null.")
1318
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
1319
+
1320
+ input_tables = {
1321
+ _table_name(input_schema, input_vertices_table),
1322
+ _table_name(other_schema, other_vertices_table),
1323
+ }
1324
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
1325
+
1326
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
1327
+ operation_successful = False
1328
+
1329
+ with connection.cursor() as cursor:
1330
+ try:
1331
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
1332
+ other_vertices_name = _table_name(other_schema, other_vertices_table)
1333
+
1334
+ all_columns = ["a.id"]
1335
+ for expr in result_column_expressions:
1336
+ _check_not_null_or_empty(expr, "Result column expression cannot be null or empty.")
1337
+ all_columns.append(expr)
1338
+
1339
+ columns_list_str = ", ".join(all_columns)
1340
+
1341
+ from_join_clause = f"FROM {input_vertices_name} a INNER JOIN {other_vertices_name} b ON a.id = b.id"
1342
+
1343
+ create_sql = (
1344
+ f"CREATE TABLE {result_vertices_name} AS SELECT {columns_list_str} {from_join_clause} WHERE 1=0"
1345
+ )
1346
+ cursor.execute(create_sql)
1347
+
1348
+ for col in result_vertices_indexes:
1349
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1350
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
1351
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_vertices_name} ("{col}")')
1352
+
1353
+ insert_sql = f"INSERT INTO {result_vertices_name} SELECT {columns_list_str} {from_join_clause}"
1354
+ cursor.execute(insert_sql)
1355
+
1356
+ operation_successful = True
1357
+ finally:
1358
+ if not operation_successful:
1359
+ _safe_drop_table(cursor, result_vertices_name)
1360
+
1361
+
1362
+ def outer_join_vertices(
1363
+ connection: Connection,
1364
+ input_schema: str,
1365
+ input_vertices_table: str,
1366
+ other_schema: str,
1367
+ other_vertices_table: str,
1368
+ result_schema: str,
1369
+ result_vertices_table: str,
1370
+ result_column_expressions: List[str],
1371
+ result_vertices_indexes: List[str],
1372
+ ) -> None:
1373
+ """
1374
+ Performs a LEFT OUTER JOIN between two vertex tables on their 'id' columns.
1375
+
1376
+ All rows from the left table (`input_vertices_table`) are included.
1377
+ If a row has no match in the right table, columns from the right table
1378
+ will be NULL.
1379
+
1380
+ Using `result_column_expressions`:
1381
+ - The `id` column from the first table is always included.
1382
+ - Use alias `a` for columns from `input_vertices_table` (left).
1383
+ - Use alias `b` for columns from `other_vertices_table` (right).
1384
+ - It is highly recommended to use `AS alias_name` for each expression.
1385
+
1386
+ :param connection: The pyocient Connection object.
1387
+ :param input_schema: Schema of the first (left) vertices table.
1388
+ :param input_vertices_table: Name of the first vertices table (alias `a`).
1389
+ :param other_schema: Schema of the second (right) vertices table.
1390
+ :param other_vertices_table: Name of the second vertices table (alias `b`).
1391
+ :param result_schema: Schema for the result table.
1392
+ :param result_vertices_table: Name for the table with join results.
1393
+ :param result_column_expressions: List of SQL expressions for result columns.
1394
+ :param result_vertices_indexes: List of columns to index in the result table.
1395
+ :raises SQLException: If any database error occurs.
1396
+ """
1397
+ _check_not_null(connection, "Connection cannot be null.")
1398
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1399
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
1400
+ _check_not_null_or_empty(other_schema, "Other schema cannot be null or empty.")
1401
+ _check_not_null_or_empty(other_vertices_table, "Other vertices table name cannot be null or empty.")
1402
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1403
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
1404
+ _check_not_null(result_column_expressions, "Result column expressions list cannot be null.")
1405
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
1406
+
1407
+ input_tables = {
1408
+ _table_name(input_schema, input_vertices_table),
1409
+ _table_name(other_schema, other_vertices_table),
1410
+ }
1411
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
1412
+
1413
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
1414
+ operation_successful = False
1415
+
1416
+ with connection.cursor() as cursor:
1417
+ try:
1418
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
1419
+ other_vertices_name = _table_name(other_schema, other_vertices_table)
1420
+
1421
+ all_columns = ["a.id"]
1422
+ for expr in result_column_expressions:
1423
+ _check_not_null_or_empty(expr, "Result column expression cannot be null or empty.")
1424
+ all_columns.append(expr)
1425
+
1426
+ columns_list_str = ", ".join(all_columns)
1427
+
1428
+ from_join_clause = f"FROM {input_vertices_name} a LEFT JOIN {other_vertices_name} b ON a.id = b.id"
1429
+
1430
+ create_sql = (
1431
+ f"CREATE TABLE {result_vertices_name} AS SELECT {columns_list_str} {from_join_clause} WHERE 1=0"
1432
+ )
1433
+ cursor.execute(create_sql)
1434
+
1435
+ for col in result_vertices_indexes:
1436
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1437
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
1438
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_vertices_name} ("{col}")')
1439
+
1440
+ insert_sql = f"INSERT INTO {result_vertices_name} SELECT {columns_list_str} {from_join_clause}"
1441
+ cursor.execute(insert_sql)
1442
+
1443
+ operation_successful = True
1444
+ finally:
1445
+ if not operation_successful:
1446
+ _safe_drop_table(cursor, result_vertices_name)
1447
+
1448
+
1449
+ def mask(
1450
+ connection: Connection,
1451
+ input_schema: str,
1452
+ input_vertices_table: str,
1453
+ input_edges_table: str,
1454
+ other_schema: str,
1455
+ other_vertices_table: str,
1456
+ other_edges_table: str,
1457
+ result_schema: str,
1458
+ result_vertices_table: str,
1459
+ result_edges_table: str,
1460
+ result_vertices_indexes: List[str],
1461
+ result_edges_indexes: List[str],
1462
+ ) -> None:
1463
+ """
1464
+ Creates a subgraph by retaining elements from an input graph that are
1465
+ also present in an 'other' graph.
1466
+
1467
+ - Vertices from the input graph are kept if an ID match is found in the
1468
+ 'other' vertex table.
1469
+ - Edges are kept if a (srcid, destid) match is found in the 'other' edge
1470
+ table AND both endpoints exist in the new result vertex table.
1471
+
1472
+ :param connection: The pyocient Connection object.
1473
+ :param input_schema: Schema of the primary input graph.
1474
+ :param input_vertices_table: Vertices table of the primary graph.
1475
+ :param input_edges_table: Edges table of the primary graph.
1476
+ :param other_schema: Schema of the masking ('other') graph.
1477
+ :param other_vertices_table: Vertices table of the masking graph.
1478
+ :param other_edges_table: Edges table of the masking graph.
1479
+ :param result_schema: Schema for the result subgraph tables.
1480
+ :param result_vertices_table: Name for the result vertices table.
1481
+ :param result_edges_table: Name for the result edges table.
1482
+ :param result_vertices_indexes: List of columns to index in the result vertices.
1483
+ :param result_edges_indexes: List of columns to index in the result edges.
1484
+ :raises SQLException: If any database error occurs.
1485
+ """
1486
+ _check_not_null(connection, "Connection cannot be null.")
1487
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1488
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table cannot be null or empty.")
1489
+ _check_not_null_or_empty(input_edges_table, "Input edges table cannot be null or empty.")
1490
+ _check_not_null_or_empty(other_schema, "Other schema cannot be null or empty.")
1491
+ _check_not_null_or_empty(other_vertices_table, "Other vertices table cannot be null or empty.")
1492
+ _check_not_null_or_empty(other_edges_table, "Other edges table cannot be null or empty.")
1493
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1494
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table cannot be null or empty.")
1495
+ _check_not_null_or_empty(result_edges_table, "Result edges table cannot be null or empty.")
1496
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
1497
+ _check_not_null(result_edges_indexes, "Result edges index list cannot be null.")
1498
+
1499
+ input_v = _table_name(input_schema, input_vertices_table)
1500
+ input_e = _table_name(input_schema, input_edges_table)
1501
+ other_v = _table_name(other_schema, other_vertices_table)
1502
+ other_e = _table_name(other_schema, other_edges_table)
1503
+ result_v = _table_name(result_schema, result_vertices_table)
1504
+ result_e = _table_name(result_schema, result_edges_table)
1505
+
1506
+ source_tables = {input_v, input_e, other_v, other_e}
1507
+ _check_table_conflict(result_schema, result_vertices_table, source_tables)
1508
+ _check_table_conflict(result_schema, result_edges_table, source_tables)
1509
+
1510
+ operation_successful = False
1511
+
1512
+ with connection.cursor() as cursor:
1513
+ try:
1514
+ _create_table_from_template(
1515
+ input_schema,
1516
+ input_vertices_table,
1517
+ result_schema,
1518
+ result_vertices_table,
1519
+ result_vertices_indexes,
1520
+ cursor,
1521
+ )
1522
+ _create_table_from_template(
1523
+ input_schema, input_edges_table, result_schema, result_edges_table, result_edges_indexes, cursor
1524
+ )
1525
+
1526
+ populate_v_sql = f"INSERT INTO {result_v} SELECT v.* FROM {input_v} v JOIN {other_v} o ON v.id = o.id"
1527
+ cursor.execute(populate_v_sql)
1528
+
1529
+ populate_e_sql = f"""
1530
+ INSERT INTO {result_e}
1531
+ SELECT e.* FROM {input_e} e
1532
+ INNER JOIN {other_e} oe ON e.srcid = oe.srcid AND e.destid = oe.destid
1533
+ WHERE EXISTS (SELECT 1 FROM {result_v} rv_src WHERE rv_src.id = e.srcid)
1534
+ AND EXISTS (SELECT 1 FROM {result_v} rv_dst WHERE rv_dst.id = e.destid)
1535
+ """
1536
+ cursor.execute(populate_e_sql)
1537
+
1538
+ operation_successful = True
1539
+ finally:
1540
+ if not operation_successful:
1541
+ _safe_drop_table(cursor, result_v)
1542
+ _safe_drop_table(cursor, result_e)
1543
+
1544
+
1545
+ def group_edges(
1546
+ connection: Connection,
1547
+ input_schema: str,
1548
+ input_edges_table: str,
1549
+ result_schema: str,
1550
+ result_edges_table: str,
1551
+ result_column_expressions: List[str],
1552
+ result_edges_indexes: List[str],
1553
+ ) -> None:
1554
+ """
1555
+ Groups duplicate edges from an input table into a new result table
1556
+ using SQL aggregation.
1557
+
1558
+ The result table contains unique (srcid, destid) pairs.
1559
+
1560
+ Using `result_column_expressions`:
1561
+ - Each expression computes an additional column for the result table.
1562
+ - Each expression *must* use an SQL aggregate function (e.g., COUNT, SUM,
1563
+ AVG, MIN, MAX) on columns from the input table.
1564
+ - Each expression *must* include an `AS alias_name` clause to name the
1565
+ resulting column.
1566
+
1567
+ :param connection: The pyocient Connection object.
1568
+ :param input_schema: Schema of the source edges table.
1569
+ :param input_edges_table: Name of the source edges table.
1570
+ :param result_schema: Schema for the target table.
1571
+ :param result_edges_table: Name for the table with grouped edges.
1572
+ :param result_column_expressions: List of SQL aggregate expressions.
1573
+ :param result_edges_indexes: List of columns to index in the result table.
1574
+ :raises SQLException: If any database error occurs.
1575
+ """
1576
+ _check_not_null(connection, "Connection cannot be null.")
1577
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1578
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
1579
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1580
+ _check_not_null_or_empty(result_edges_table, "Result edges table name cannot be null or empty.")
1581
+ _check_not_null(result_column_expressions, "Result column expressions list cannot be null.")
1582
+ _check_not_null(result_edges_indexes, "Result edges index list cannot be null.")
1583
+
1584
+ input_tables = {_table_name(input_schema, input_edges_table)}
1585
+ _check_table_conflict(result_schema, result_edges_table, input_tables)
1586
+
1587
+ result_edges_name = _table_name(result_schema, result_edges_table)
1588
+ operation_successful = False
1589
+
1590
+ with connection.cursor() as cursor:
1591
+ try:
1592
+ input_edges_name = _table_name(input_schema, input_edges_table)
1593
+
1594
+ columns_builder = ["srcid", "destid"]
1595
+ for expr in result_column_expressions:
1596
+ _check_not_null_or_empty(expr, "Result column expression cannot be null or empty.")
1597
+ columns_builder.append(expr)
1598
+ columns_list_str = ", ".join(columns_builder)
1599
+
1600
+ create_query = (
1601
+ f"CREATE TABLE {result_edges_name} AS "
1602
+ f"SELECT {columns_list_str} FROM {input_edges_name} "
1603
+ f"GROUP BY srcid, destid HAVING 1=0"
1604
+ )
1605
+ cursor.execute(create_query)
1606
+
1607
+ for col in result_edges_indexes:
1608
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1609
+ index_name = f"{result_edges_table}_{col}_idx".replace('"', "")
1610
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_edges_name} ("{col}")')
1611
+
1612
+ insert_query = (
1613
+ f"INSERT INTO {result_edges_name} "
1614
+ f"SELECT {columns_list_str} FROM {input_edges_name} "
1615
+ f"GROUP BY srcid, destid"
1616
+ )
1617
+ cursor.execute(insert_query)
1618
+
1619
+ operation_successful = True
1620
+ finally:
1621
+ if not operation_successful:
1622
+ _safe_drop_table(cursor, result_edges_name)
1623
+
1624
+
1625
+ def aggregate_messages(
1626
+ connection: Connection,
1627
+ input_schema: str,
1628
+ input_vertices_table: str,
1629
+ input_edges_table: str,
1630
+ result_schema: str,
1631
+ result_vertices_table: str,
1632
+ send_to_source_expr: Optional[str],
1633
+ send_to_dest_expr: Optional[str],
1634
+ aggregate_expr: str,
1635
+ result_vertices_indexes: List[str],
1636
+ ) -> None:
1637
+ """
1638
+ Aggregates messages passed between connected vertices in a graph.
1639
+
1640
+ This method simulates a message-passing process. For each edge, messages
1641
+ are generated via `send_to_source_expr` and `send_to_dest_expr`. These
1642
+ messages are then grouped by target vertex and combined using `aggregate_expr`.
1643
+
1644
+ Expression Aliases:
1645
+ - In `send_..._expr`, use `a` for source vertex, `b` for edge, `c` for dest vertex.
1646
+ - In `aggregate_expr`, operate on the generated message column, which is
1647
+ internally named `msg` (e.g., `SUM(msg)`).
1648
+ - It is recommended to alias the result of `aggregate_expr` (e.g.,
1649
+ `SUM(msg) AS total_value`). If no alias is given, the result column
1650
+ is named `aggregated_message`.
1651
+
1652
+ :param connection: The pyocient Connection object.
1653
+ :param input_schema: Schema of the source graph tables.
1654
+ :param input_vertices_table: Name of the source vertices table.
1655
+ :param input_edges_table: Name of the source edges table.
1656
+ :param result_schema: Schema for the result table.
1657
+ :param result_vertices_table: Name for the table with aggregated messages.
1658
+ :param send_to_source_expr: SQL expression for the message sent to the source vertex.
1659
+ :param send_to_dest_expr: SQL expression for the message sent to the destination vertex.
1660
+ :param aggregate_expr: SQL aggregate expression to combine messages (e.g., 'SUM(msg)').
1661
+ :param result_vertices_indexes: List of columns to index in the result table.
1662
+ :raises SQLException: If any database error occurs.
1663
+ """
1664
+ _check_not_null(connection, "Connection cannot be null.")
1665
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1666
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
1667
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
1668
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1669
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
1670
+ _check_not_null_or_empty(aggregate_expr, "Aggregate expression cannot be null or empty.")
1671
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
1672
+
1673
+ send_source = send_to_source_expr and send_to_source_expr.strip()
1674
+ send_dest = send_to_dest_expr and send_to_dest_expr.strip()
1675
+ if not send_source and not send_dest:
1676
+ raise SQLException("Argument Error: At least one of send_to_source_expr or send_to_dest_expr must be provided.")
1677
+
1678
+ aggregate_alias = "aggregated_message"
1679
+ agg_parts = aggregate_expr.lower().split(" as ")
1680
+ if len(agg_parts) > 1:
1681
+ aggregate_alias = agg_parts[-1].strip()
1682
+
1683
+ current_agg_expr = aggregate_expr
1684
+ if f" as {aggregate_alias}" not in aggregate_expr.lower():
1685
+ current_agg_expr = f"{aggregate_expr} AS {aggregate_alias}"
1686
+
1687
+ input_tables = {
1688
+ _table_name(input_schema, input_vertices_table),
1689
+ _table_name(input_schema, input_edges_table),
1690
+ }
1691
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
1692
+
1693
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
1694
+ operation_successful = False
1695
+
1696
+ with connection.cursor() as cursor:
1697
+ try:
1698
+ input_v = _table_name(input_schema, input_vertices_table)
1699
+ input_e = _table_name(input_schema, input_edges_table)
1700
+
1701
+ messages_sql_parts = []
1702
+ from_join_clause = (
1703
+ f"FROM {input_v} a JOIN {input_e} b ON a.id = b.srcid JOIN {input_v} c ON b.destid = c.id"
1704
+ )
1705
+
1706
+ if send_source:
1707
+ messages_sql_parts.append(
1708
+ f"SELECT b.srcid AS id, ({send_to_source_expr}) AS msg {from_join_clause} WHERE ({send_to_source_expr}) IS NOT NULL"
1709
+ )
1710
+ if send_dest:
1711
+ messages_sql_parts.append(
1712
+ f"SELECT b.destid AS id, ({send_to_dest_expr}) AS msg {from_join_clause} WHERE ({send_to_dest_expr}) IS NOT NULL"
1713
+ )
1714
+
1715
+ messages_sql = " UNION ALL ".join(messages_sql_parts)
1716
+
1717
+ create_sql = (
1718
+ f"CREATE TABLE {result_vertices_name} AS "
1719
+ f"SELECT id, {current_agg_expr} FROM ({messages_sql}) msgs "
1720
+ "GROUP BY id HAVING 1=0"
1721
+ )
1722
+ cursor.execute(create_sql)
1723
+
1724
+ for col in result_vertices_indexes:
1725
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1726
+ if col.lower() not in ("id", aggregate_alias.lower()):
1727
+ raise SQLException(
1728
+ f"Argument Error: Invalid index column '{col}'. Must be 'id' or '{aggregate_alias}'."
1729
+ )
1730
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
1731
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_vertices_name} ("{col}")')
1732
+
1733
+ insert_sql = (
1734
+ f"INSERT INTO {result_vertices_name} (id, {aggregate_alias}) "
1735
+ f"SELECT id, {current_agg_expr} FROM ({messages_sql}) msgs "
1736
+ "GROUP BY id"
1737
+ )
1738
+ cursor.execute(insert_sql)
1739
+
1740
+ operation_successful = True
1741
+ finally:
1742
+ if not operation_successful:
1743
+ _safe_drop_table(cursor, result_vertices_name)
1744
+
1745
+
1746
+ def collect_neighbors(
1747
+ connection: Connection,
1748
+ input_schema: str,
1749
+ input_vertices_table: str,
1750
+ input_edges_table: str,
1751
+ result_schema: str,
1752
+ result_vertices_table: str,
1753
+ direction: EdgeDirection,
1754
+ result_vertices_indexes: List[str],
1755
+ ) -> None:
1756
+ """
1757
+ Collects neighbor information for each vertex and stores the results
1758
+ (as an array of tuples) in a new table.
1759
+
1760
+ The result table will contain `id` (the source vertex) and `neighbors`
1761
+ (an array of tuples, where each tuple represents a neighbor vertex and
1762
+ contains all of its attributes).
1763
+
1764
+ :param connection: The pyocient Connection object.
1765
+ :param input_schema: Schema of the source graph tables.
1766
+ :param input_vertices_table: Name of the source vertices table.
1767
+ :param input_edges_table: Name of the source edges table.
1768
+ :param result_schema: Schema for the result table.
1769
+ :param result_vertices_table: Name for the table with neighbor arrays.
1770
+ :param direction: The direction of edges to follow (IN, OUT, or BOTH).
1771
+ :param result_vertices_indexes: List of columns to index ('id' or 'neighbors').
1772
+ :raises SQLException: If any database error occurs.
1773
+ """
1774
+ _check_not_null(connection, "Connection cannot be null.")
1775
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1776
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
1777
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
1778
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1779
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
1780
+ _check_not_null(direction, "Edge direction cannot be null.")
1781
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
1782
+
1783
+ input_tables = {_table_name(input_schema, input_vertices_table), _table_name(input_schema, input_edges_table)}
1784
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
1785
+
1786
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
1787
+ operation_successful = False
1788
+
1789
+ with connection.cursor() as cursor:
1790
+ try:
1791
+ input_v = _table_name(input_schema, input_vertices_table)
1792
+ input_e = _table_name(input_schema, input_edges_table)
1793
+
1794
+ cursor.execute(f"SELECT * FROM {input_v} WHERE 1=0")
1795
+ neighbor_vertex_cols = [f'n."{c[0]}"' for c in (cursor.description or []) if c[0].lower() != "id"]
1796
+
1797
+ tuple_cols = ["n.id"] + neighbor_vertex_cols
1798
+ tuple_expr = f"TUPLE({', '.join(tuple_cols)})"
1799
+ tuple_alias = "neighbor_tuple"
1800
+ result_col_name = "neighbors"
1801
+ group_by_col = "vertex_id"
1802
+
1803
+ if direction == EdgeDirection.IN:
1804
+ base_query = (
1805
+ f"SELECT e.destid AS {group_by_col}, {tuple_expr} AS {tuple_alias} "
1806
+ f"FROM {input_v} v JOIN {input_e} e ON e.destid = v.id JOIN {input_v} n ON e.srcid = n.id"
1807
+ )
1808
+ elif direction == EdgeDirection.OUT:
1809
+ base_query = (
1810
+ f"SELECT e.srcid AS {group_by_col}, {tuple_expr} AS {tuple_alias} "
1811
+ f"FROM {input_v} v JOIN {input_e} e ON e.srcid = v.id JOIN {input_v} n ON e.destid = n.id"
1812
+ )
1813
+ elif direction == EdgeDirection.BOTH:
1814
+ out_part = f"(SELECT e.srcid AS {group_by_col}, {tuple_expr} AS {tuple_alias} FROM {input_v} v JOIN {input_e} e ON e.srcid = v.id JOIN {input_v} n ON e.destid = n.id)"
1815
+ in_part = f"(SELECT e.destid AS {group_by_col}, {tuple_expr} AS {tuple_alias} FROM {input_v} v JOIN {input_e} e ON e.destid = v.id JOIN {input_v} n ON e.srcid = n.id)"
1816
+ # FIXED: Use UNION ALL to preserve duplicates
1817
+ base_query = f"{out_part} UNION ALL {in_part}"
1818
+ else:
1819
+ raise SQLException("Internal Error: Invalid edge direction.")
1820
+
1821
+ create_sql = (
1822
+ f"CREATE TABLE {result_vertices_name} AS "
1823
+ f"SELECT {group_by_col} AS id, ARRAY_AGG({tuple_alias}) AS {result_col_name} "
1824
+ f"FROM ({base_query}) subquery GROUP BY {group_by_col} HAVING 1=0"
1825
+ )
1826
+ cursor.execute(create_sql)
1827
+
1828
+ for col in result_vertices_indexes:
1829
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1830
+ if col.lower() not in ("id", result_col_name.lower()):
1831
+ raise SQLException(
1832
+ f"Argument Error: Invalid index column '{col}'. Must be 'id' or '{result_col_name}'."
1833
+ )
1834
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
1835
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_vertices_name} ("{col}")')
1836
+
1837
+ insert_sql = (
1838
+ f"INSERT INTO {result_vertices_name} (id, {result_col_name}) "
1839
+ f"SELECT {group_by_col} AS id, ARRAY_AGG({tuple_alias}) "
1840
+ f"FROM ({base_query}) subquery GROUP BY {group_by_col}"
1841
+ )
1842
+ cursor.execute(insert_sql)
1843
+
1844
+ operation_successful = True
1845
+ finally:
1846
+ if not operation_successful:
1847
+ _safe_drop_table(cursor, result_vertices_name)
1848
+
1849
+
1850
+ def collect_edges(
1851
+ connection: Connection,
1852
+ input_schema: str,
1853
+ input_vertices_table: str,
1854
+ input_edges_table: str,
1855
+ result_schema: str,
1856
+ result_vertices_table: str,
1857
+ direction: EdgeDirection,
1858
+ result_vertices_indexes: List[str],
1859
+ ) -> None:
1860
+ """
1861
+ Collects edge information for each vertex and stores the results
1862
+ (as an array of tuples) in a new table.
1863
+
1864
+ The result table will contain `id` (the vertex) and `edges` (an array of
1865
+ tuples, where each tuple represents a connected edge and contains all of
1866
+ its attributes).
1867
+
1868
+ :param connection: The pyocient Connection object.
1869
+ :param input_schema: Schema of the source graph tables.
1870
+ :param input_vertices_table: Name of the source vertices table.
1871
+ :param input_edges_table: Name of the source edges table.
1872
+ :param result_schema: Schema for the result table.
1873
+ :param result_vertices_table: Name for the table with edge arrays.
1874
+ :param direction: The direction of edges to follow (IN, OUT, or BOTH).
1875
+ :param result_vertices_indexes: List of columns to index ('id' or 'edges').
1876
+ :raises SQLException: If any database error occurs.
1877
+ """
1878
+ _check_not_null(connection, "Connection cannot be null.")
1879
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1880
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
1881
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
1882
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1883
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
1884
+ _check_not_null(direction, "Edge direction cannot be null.")
1885
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
1886
+
1887
+ input_tables = {_table_name(input_schema, input_vertices_table), _table_name(input_schema, input_edges_table)}
1888
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
1889
+
1890
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
1891
+ operation_successful = False
1892
+
1893
+ with connection.cursor() as cursor:
1894
+ try:
1895
+ input_v = _table_name(input_schema, input_vertices_table)
1896
+ input_e = _table_name(input_schema, input_edges_table)
1897
+
1898
+ cursor.execute(f"SELECT * FROM {input_e} WHERE 1=0")
1899
+ edge_cols_list = [f'e."{c[0]}"' for c in (cursor.description or [])]
1900
+ if not edge_cols_list:
1901
+ raise SQLException(f"Input edges table '{input_e}' has no columns.")
1902
+
1903
+ edge_tuple_expr = f"TUPLE({', '.join(edge_cols_list)})"
1904
+ edge_tuple_alias = "edge_tuple"
1905
+ result_col_name = "edges"
1906
+ group_by_col = "vertex_id"
1907
+
1908
+ if direction == EdgeDirection.IN:
1909
+ base_query = f"SELECT e.destid AS {group_by_col}, {edge_tuple_expr} AS {edge_tuple_alias} FROM {input_v} v JOIN {input_e} e ON e.destid = v.id"
1910
+ elif direction == EdgeDirection.OUT:
1911
+ base_query = f"SELECT e.srcid AS {group_by_col}, {edge_tuple_expr} AS {edge_tuple_alias} FROM {input_v} v JOIN {input_e} e ON e.srcid = v.id"
1912
+ elif direction == EdgeDirection.BOTH:
1913
+ out_part = f"(SELECT e.srcid AS {group_by_col}, {edge_tuple_expr} AS {edge_tuple_alias} FROM {input_v} v JOIN {input_e} e ON e.srcid = v.id)"
1914
+ in_part = f"(SELECT e.destid AS {group_by_col}, {edge_tuple_expr} AS {edge_tuple_alias} FROM {input_v} v JOIN {input_e} e ON e.destid = v.id)"
1915
+ # FIXED: Use UNION ALL to preserve duplicates
1916
+ base_query = f"{out_part} UNION ALL {in_part}"
1917
+ else:
1918
+ raise SQLException("Internal Error: Invalid edge direction.")
1919
+
1920
+ create_sql = (
1921
+ f"CREATE TABLE {result_vertices_name} AS "
1922
+ f"SELECT {group_by_col} AS id, ARRAY_AGG({edge_tuple_alias}) AS {result_col_name} "
1923
+ f"FROM ({base_query}) subquery GROUP BY {group_by_col} HAVING 1=0"
1924
+ )
1925
+ cursor.execute(create_sql)
1926
+
1927
+ for col in result_vertices_indexes:
1928
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
1929
+ if col.lower() not in ("id", result_col_name.lower()):
1930
+ raise SQLException(
1931
+ f"Argument Error: Invalid index column '{col}'. Must be 'id' or '{result_col_name}'."
1932
+ )
1933
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
1934
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_vertices_name} ("{col}")')
1935
+
1936
+ insert_sql = (
1937
+ f"INSERT INTO {result_vertices_name} (id, {result_col_name}) "
1938
+ f"SELECT {group_by_col} AS id, ARRAY_AGG({edge_tuple_alias}) "
1939
+ f"FROM ({base_query}) subquery GROUP BY {group_by_col}"
1940
+ )
1941
+ cursor.execute(insert_sql)
1942
+
1943
+ operation_successful = True
1944
+ finally:
1945
+ if not operation_successful:
1946
+ _safe_drop_table(cursor, result_vertices_name)
1947
+
1948
+
1949
+ def label_propagation(
1950
+ connection: Connection,
1951
+ input_schema: str,
1952
+ input_vertices_table: str,
1953
+ input_edges_table: str,
1954
+ result_schema: str,
1955
+ result_vertices_table: str,
1956
+ max_iterations: int,
1957
+ result_vertices_indexes: List[str],
1958
+ ) -> None:
1959
+ """
1960
+ Performs the Label Propagation Algorithm (LPA) iteratively to find communities.
1961
+
1962
+ The algorithm initializes each vertex's label to its own ID. Then, for a
1963
+ specified number of iterations, each vertex updates its label to the most
1964
+ frequent label among its neighbors. Ties are broken by choosing the
1965
+ numerically smaller label ID.
1966
+
1967
+ :param connection: The pyocient Connection object.
1968
+ :param input_schema: Schema of the source graph tables.
1969
+ :param input_vertices_table: Name of the source vertices table.
1970
+ :param input_edges_table: Name of the source edges table.
1971
+ :param result_schema: Schema for the result and temporary tables.
1972
+ :param result_vertices_table: Name for the table with final vertex labels.
1973
+ :param max_iterations: The maximum number of iterations to perform.
1974
+ :param result_vertices_indexes: List of columns to index ('id' or 'label').
1975
+ :raises SQLException: If any database error occurs.
1976
+ """
1977
+ _check_not_null(connection, "Connection cannot be null.")
1978
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
1979
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
1980
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
1981
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
1982
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
1983
+ _check_positive(max_iterations, "Maximum iterations must be >= 1.")
1984
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
1985
+
1986
+ input_tables = {_table_name(input_schema, input_vertices_table), _table_name(input_schema, input_edges_table)}
1987
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
1988
+
1989
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
1990
+ temp_base_name = f"lpa_temp_{uuid.uuid4().hex[:8]}"
1991
+ current_table = _table_name(result_schema, f"{temp_base_name}_iter0")
1992
+
1993
+ temp_tables_created = []
1994
+ operation_successful = False
1995
+
1996
+ with connection.cursor() as cursor:
1997
+ try:
1998
+ # Initial step (Iteration 0)
1999
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
2000
+ initial_sql = f"CREATE TABLE {current_table} AS SELECT id, id AS label FROM {input_vertices_name}"
2001
+ cursor.execute(initial_sql)
2002
+ temp_tables_created.append(current_table)
2003
+
2004
+ input_edges_name = _table_name(input_schema, input_edges_table)
2005
+
2006
+ for i in range(max_iterations):
2007
+ is_last_iteration = i == max_iterations - 1
2008
+ next_table = (
2009
+ result_vertices_name
2010
+ if is_last_iteration
2011
+ else _table_name(result_schema, f"{temp_base_name}_iter{i + 1}")
2012
+ )
2013
+
2014
+ # This CTE-based query finds the most frequent neighbor label for each vertex
2015
+ new_labels_sql = f"""
2016
+ WITH neighbors AS (
2017
+ SELECT e.srcid AS vertex, e.destid AS neighbor FROM {input_edges_name} e
2018
+ UNION ALL
2019
+ SELECT e.destid AS vertex, e.srcid AS neighbor FROM {input_edges_name} e
2020
+ ), neighbor_labels AS (
2021
+ SELECT n.vertex, l.label FROM neighbors n JOIN {current_table} l ON n.neighbor = l.id
2022
+ ), label_counts AS (
2023
+ SELECT vertex, label, COUNT(*) AS cnt FROM neighbor_labels GROUP BY vertex, label
2024
+ ), ranked_labels AS (
2025
+ SELECT vertex, label, ROW_NUMBER() OVER (PARTITION BY vertex ORDER BY cnt DESC, label ASC) AS rn FROM label_counts
2026
+ ), preferred_new_labels AS (
2027
+ SELECT vertex AS id, label FROM ranked_labels WHERE rn = 1
2028
+ )
2029
+ SELECT ct.id, COALESCE(pnl.label, ct.label) AS label
2030
+ FROM {current_table} ct LEFT JOIN preferred_new_labels pnl ON ct.id = pnl.id
2031
+ """
2032
+
2033
+ if is_last_iteration:
2034
+ cursor.execute(f"CREATE TABLE {next_table} (id BIGINT NOT NULL, label BIGINT)")
2035
+
2036
+ for col in result_vertices_indexes:
2037
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
2038
+ if col.lower() not in ("id", "label"):
2039
+ raise SQLException(
2040
+ f"Argument Error: Invalid index column '{col}'. Must be 'id' or 'label'."
2041
+ )
2042
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
2043
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {next_table} ("{col}")')
2044
+
2045
+ cursor.execute(f"INSERT INTO {next_table} (id, label) {new_labels_sql}")
2046
+ else:
2047
+ cursor.execute(f"CREATE TABLE {next_table} AS {new_labels_sql}")
2048
+ temp_tables_created.append(next_table)
2049
+
2050
+ current_table = next_table
2051
+
2052
+ operation_successful = True
2053
+ finally:
2054
+ for temp_table in temp_tables_created:
2055
+ if temp_table != result_vertices_name or not operation_successful:
2056
+ _safe_drop_table(cursor, temp_table)
2057
+ if not operation_successful:
2058
+ _safe_drop_table(cursor, result_vertices_name)
2059
+
2060
+
2061
+ class TriangleCount:
2062
+ """
2063
+ Provides utility methods for performing Triangle Counting operations.
2064
+ """
2065
+
2066
+ @staticmethod
2067
+ def run(
2068
+ connection: Connection,
2069
+ input_schema: str,
2070
+ input_vertices_table: str,
2071
+ input_edges_table: str,
2072
+ result_schema: str,
2073
+ result_vertices_table: str,
2074
+ result_vertices_indexes: List[str],
2075
+ ) -> None:
2076
+ """
2077
+ Runs the full triangle counting process.
2078
+
2079
+ This handles creating a temporary table for canonical edges (where srcid <
2080
+ destid, and duplicates/self-loops are removed), executing the core
2081
+ counting logic, and cleaning up.
2082
+
2083
+ :param connection: The pyocient Connection object.
2084
+ :param input_schema: Schema of the source graph tables.
2085
+ :param input_vertices_table: Name of the source vertices table.
2086
+ :param input_edges_table: Name of the source edges table.
2087
+ :param result_schema: Schema for the result table.
2088
+ :param result_vertices_table: Name for the table with triangle counts.
2089
+ :param result_vertices_indexes: List of columns to index ('id' or 'triangle_count').
2090
+ :raises SQLException: If any database error occurs.
2091
+ """
2092
+ _check_not_null(connection, "Connection cannot be null.")
2093
+
2094
+ canonical_edges_temp_name = f"tc_canon_{uuid.uuid4().hex[:8]}"
2095
+ # Create temp table in the result schema for permissions reasons
2096
+ canonical_edges_qualified_name = _table_name(result_schema, canonical_edges_temp_name)
2097
+
2098
+ with connection.cursor() as cursor:
2099
+ try:
2100
+ # Store edges with srcid < destid, remove self-loops and duplicates
2101
+ qualified_input_edges = _table_name(input_schema, input_edges_table)
2102
+ cursor.execute(f"""
2103
+ CREATE TABLE {canonical_edges_qualified_name} AS
2104
+ SELECT
2105
+ CASE WHEN srcid < destid THEN srcid ELSE destid END AS srcid,
2106
+ CASE WHEN srcid < destid THEN destid ELSE srcid END AS destid
2107
+ FROM {qualified_input_edges}
2108
+ WHERE srcid != destid
2109
+ GROUP BY 1, 2
2110
+ """)
2111
+
2112
+ # Run the core logic on the canonicalized edges
2113
+ TriangleCount.run_pre_canonicalized(
2114
+ connection,
2115
+ input_schema,
2116
+ input_vertices_table,
2117
+ result_schema, # Use result_schema for temp table
2118
+ canonical_edges_temp_name,
2119
+ result_schema,
2120
+ result_vertices_table,
2121
+ result_vertices_indexes,
2122
+ )
2123
+ finally:
2124
+ _safe_drop_table(cursor, canonical_edges_qualified_name)
2125
+
2126
+ @staticmethod
2127
+ def run_pre_canonicalized(
2128
+ connection: Connection,
2129
+ input_schema: str,
2130
+ input_vertices_table: str,
2131
+ canonical_edges_schema: str,
2132
+ canonical_edges_table: str,
2133
+ result_schema: str,
2134
+ result_vertices_table: str,
2135
+ result_vertices_indexes: List[str],
2136
+ ) -> None:
2137
+ """
2138
+ Runs triangle count on a graph with pre-canonicalized edges.
2139
+
2140
+ Assumes edges are already directed (srcid < destid) with no duplicates
2141
+ or self-loops.
2142
+ """
2143
+ _check_not_null_or_empty(canonical_edges_table, "Canonical edges table cannot be null or empty.")
2144
+
2145
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
2146
+ operation_successful = False
2147
+
2148
+ with connection.cursor() as cursor:
2149
+ try:
2150
+ cursor.execute(
2151
+ f"CREATE TABLE {result_vertices_name} (id BIGINT NOT NULL, triangle_count BIGINT NOT NULL)"
2152
+ )
2153
+
2154
+ for col in result_vertices_indexes:
2155
+ if col.lower() not in ("id", "triangle_count"):
2156
+ raise SQLException(
2157
+ f"Argument Error: Invalid index column '{col}'. Must be 'id' or 'triangle_count'."
2158
+ )
2159
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
2160
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_vertices_name} ("{col}")')
2161
+
2162
+ qualified_input_vertices = _table_name(input_schema, input_vertices_table)
2163
+ qualified_canonical_edges = _table_name(canonical_edges_schema, canonical_edges_table)
2164
+
2165
+ insert_sql = f"""
2166
+ INSERT INTO {result_vertices_name} (id, triangle_count)
2167
+ WITH triangles AS (
2168
+ SELECT e1.srcid AS a, e1.destid AS b, e2.destid AS c
2169
+ FROM {qualified_canonical_edges} e1
2170
+ JOIN {qualified_canonical_edges} e2 ON e1.destid = e2.srcid
2171
+ ), closed_triangles AS (
2172
+ SELECT t.a, t.b, t.c FROM triangles t
2173
+ JOIN {qualified_canonical_edges} e3 ON t.a = e3.srcid AND t.c = e3.destid
2174
+ ), vertex_counts AS (
2175
+ SELECT a AS vertex, COUNT(*) AS cnt FROM closed_triangles GROUP BY a
2176
+ UNION ALL
2177
+ SELECT b AS vertex, COUNT(*) AS cnt FROM closed_triangles GROUP BY b
2178
+ UNION ALL
2179
+ SELECT c AS vertex, COUNT(*) AS cnt FROM closed_triangles GROUP BY c
2180
+ ), aggregated_counts AS (
2181
+ SELECT vertex, SUM(cnt) as total_count FROM vertex_counts GROUP BY vertex
2182
+ )
2183
+ SELECT v.id, COALESCE(agg.total_count, 0) AS triangle_count
2184
+ FROM {qualified_input_vertices} v
2185
+ LEFT JOIN aggregated_counts agg ON v.id = agg.vertex
2186
+ """
2187
+ cursor.execute(insert_sql)
2188
+ operation_successful = True
2189
+ finally:
2190
+ if not operation_successful:
2191
+ _safe_drop_table(cursor, result_vertices_name)
2192
+
2193
+
2194
+ def pregel(
2195
+ connection: Connection,
2196
+ input_schema: str,
2197
+ input_vertices_table: str,
2198
+ input_edges_table: str,
2199
+ result_schema: str,
2200
+ result_vertices_table: str,
2201
+ initializer_expr: str,
2202
+ send_to_source_expr: Optional[str],
2203
+ send_to_dest_expr: Optional[str],
2204
+ aggregate_expr: str,
2205
+ updater_expr: str,
2206
+ max_iterations: int,
2207
+ result_vertices_indexes: List[str],
2208
+ ) -> None:
2209
+ """
2210
+ Performs a Pregel-like graph computation using SQL.
2211
+
2212
+ This method implements a vertex-centric graph processing algorithm. It
2213
+ iteratively computes vertex states based on messages passed along edges.
2214
+
2215
+ Important Note on Expressions and Aliases:
2216
+ - initializer_expr: Computes the starting 'state' for each vertex. Can
2217
+ reference any column from `input_vertices_table`.
2218
+ - send_to_..._expr: Computes a message. MUST use aliases:
2219
+ - `a`: for the source vertex's current state row (`id`, `state`).
2220
+ - `b`: for the edge table row (`srcid`, `destid`, ...).
2221
+ - `c`: for the destination vertex's current state row (`id`, `state`).
2222
+ Return NULL to not send a message. Computation terminates early if all
2223
+ send expressions for an iteration return NULL.
2224
+ - aggregate_expr: Combines incoming messages. MUST operate on the `msg`
2225
+ column (e.g., `SUM(msg)`). The result is aliased as `aggregated_message`.
2226
+ - updater_expr: Computes the next state. MUST use aliases:
2227
+ - `s`: for the vertex's current state row (`id`, `state`).
2228
+ - `m`: for the aggregated messages row (`id`, `aggregated_message`).
2229
+ `m.aggregated_message` will be NULL if no messages were received.
2230
+
2231
+ :param connection: The pyocient Connection object.
2232
+ :param input_schema: Schema of the source graph tables.
2233
+ :param input_vertices_table: Name of the source vertices table.
2234
+ :param input_edges_table: Name of the source edges table.
2235
+ :param result_schema: Schema for the result and temporary tables.
2236
+ :param result_vertices_table: Name for the table with final vertex states.
2237
+ :param initializer_expr: SQL expression for the initial vertex state.
2238
+ :param send_to_source_expr: SQL expression for messages sent to the source.
2239
+ :param send_to_dest_expr: SQL expression for messages sent to the destination.
2240
+ :param aggregate_expr: SQL aggregate expression to combine messages.
2241
+ :param updater_expr: SQL expression to compute the next vertex state.
2242
+ :param max_iterations: The maximum number of iterations to perform.
2243
+ :param result_vertices_indexes: List of columns to index ('id' or 'result').
2244
+ :raises SQLException: If any database error occurs.
2245
+ """
2246
+ _check_not_null(connection, "Connection cannot be null.")
2247
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
2248
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
2249
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
2250
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
2251
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
2252
+ _check_not_null_or_empty(initializer_expr, "Initializer expression cannot be null or empty.")
2253
+ _check_not_null_or_empty(aggregate_expr, "Aggregate expression cannot be null or empty.")
2254
+ _check_not_null_or_empty(updater_expr, "Updater expression cannot be null or empty.")
2255
+ _check_positive(max_iterations, "Maximum iterations must be >= 1.")
2256
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
2257
+
2258
+ send_source = send_to_source_expr and send_to_source_expr.strip()
2259
+ send_dest = send_to_dest_expr and send_to_dest_expr.strip()
2260
+
2261
+ current_initializer_expr = initializer_expr
2262
+ if " as state" not in initializer_expr.lower().strip():
2263
+ current_initializer_expr += " AS state"
2264
+
2265
+ aggregate_alias = "aggregated_message"
2266
+ current_aggregate_expr = aggregate_expr
2267
+ if f" as {aggregate_alias}" not in aggregate_expr.lower().strip():
2268
+ current_aggregate_expr += f" AS {aggregate_alias}"
2269
+
2270
+ input_tables = {_table_name(input_schema, input_vertices_table), _table_name(input_schema, input_edges_table)}
2271
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
2272
+
2273
+ result_vertices_name = _table_name(result_schema, result_vertices_table)
2274
+ temp_base_name = f"pregel_temp_{uuid.uuid4().hex[:8]}"
2275
+ current_state_table = _table_name(result_schema, f"{temp_base_name}_iter0")
2276
+
2277
+ operation_successful = False
2278
+ temp_tables_created = []
2279
+
2280
+ with connection.cursor() as cursor:
2281
+ try:
2282
+ input_vertices_name = _table_name(input_schema, input_vertices_table)
2283
+
2284
+ # Create initial state table
2285
+ initial_sql = (
2286
+ f"CREATE TABLE {current_state_table} AS "
2287
+ f"SELECT id, {current_initializer_expr} FROM {input_vertices_name}"
2288
+ )
2289
+ cursor.execute(initial_sql)
2290
+ temp_tables_created.append(current_state_table)
2291
+
2292
+ input_edges_name = _table_name(input_schema, input_edges_table)
2293
+
2294
+ for i in range(max_iterations):
2295
+ next_state_table = _table_name(result_schema, f"{temp_base_name}_iter{i + 1}")
2296
+ message_table = _table_name(result_schema, f"{temp_base_name}_msgs{i}")
2297
+
2298
+ # 1. Send messages
2299
+ messages_sent = 0
2300
+ if send_source or send_dest:
2301
+ messages_sql_parts = []
2302
+ from_join_clause = (
2303
+ f"FROM {current_state_table} a "
2304
+ f"JOIN {input_edges_name} b ON a.id = b.srcid "
2305
+ f"JOIN {current_state_table} c ON b.destid = c.id"
2306
+ )
2307
+
2308
+ if send_source:
2309
+ messages_sql_parts.append(
2310
+ f"SELECT b.srcid AS id, ({send_to_source_expr}) AS msg {from_join_clause} "
2311
+ f"WHERE ({send_to_source_expr}) IS NOT NULL"
2312
+ )
2313
+ if send_dest:
2314
+ messages_sql_parts.append(
2315
+ f"SELECT b.destid AS id, ({send_to_dest_expr}) AS msg {from_join_clause} "
2316
+ f"WHERE ({send_to_dest_expr}) IS NOT NULL"
2317
+ )
2318
+
2319
+ full_messages_sql = " UNION ALL ".join(messages_sql_parts)
2320
+ cursor.execute(f"CREATE TABLE {message_table} AS {full_messages_sql}")
2321
+ temp_tables_created.append(message_table)
2322
+
2323
+ cursor.execute(f"SELECT COUNT(*) FROM {message_table}")
2324
+ count_result = cursor.fetchone()
2325
+ if count_result:
2326
+ messages_sent = count_result[0]
2327
+
2328
+ if messages_sent == 0:
2329
+ break # Early termination if no messages were sent
2330
+
2331
+ # 2. Aggregate messages and update state
2332
+ agg_msg_subquery = f"(SELECT id, {current_aggregate_expr} FROM {message_table} GROUP BY id) m"
2333
+
2334
+ new_state_sql = (
2335
+ f"CREATE TABLE {next_state_table} AS "
2336
+ f"SELECT s.id, ({updater_expr}) AS state "
2337
+ f"FROM {current_state_table} s LEFT JOIN {agg_msg_subquery} ON s.id = m.id"
2338
+ )
2339
+ cursor.execute(new_state_sql)
2340
+ temp_tables_created.append(next_state_table)
2341
+
2342
+ # 3. Check for convergence
2343
+ diff_check_sql = (
2344
+ f"SELECT COUNT(*) FROM ("
2345
+ f"SELECT id, state FROM {next_state_table} "
2346
+ "EXCEPT "
2347
+ f"SELECT id, state FROM {current_state_table}"
2348
+ ") AS diff_query"
2349
+ )
2350
+ cursor.execute(diff_check_sql)
2351
+ row = cursor.fetchone()
2352
+ diff_count = row[0] if row is not None else 0
2353
+
2354
+ current_state_table = next_state_table
2355
+ if diff_count == 0:
2356
+ break # Converged
2357
+
2358
+ # Create and populate final result table
2359
+ final_table_sql = (
2360
+ f"CREATE TABLE {result_vertices_name} AS SELECT id, state AS result FROM {current_state_table}"
2361
+ )
2362
+ cursor.execute(final_table_sql)
2363
+
2364
+ for col in result_vertices_indexes:
2365
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
2366
+ if col.lower() not in ("id", "result"):
2367
+ raise SQLException(f"Argument Error: Invalid index column '{col}'. Must be 'id' or 'result'.")
2368
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
2369
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result_vertices_name} ("{col}")')
2370
+
2371
+ operation_successful = True
2372
+ finally:
2373
+ for temp_table in temp_tables_created:
2374
+ _safe_drop_table(cursor, temp_table)
2375
+ if not operation_successful:
2376
+ _safe_drop_table(cursor, result_vertices_name)
2377
+
2378
+
2379
+ def connected_components(
2380
+ connection: Connection,
2381
+ input_schema: str,
2382
+ input_vertices_table: str,
2383
+ input_edges_table: str,
2384
+ result_schema: str,
2385
+ result_vertices_table: str,
2386
+ max_iterations: int,
2387
+ result_vertices_indexes: List[str],
2388
+ ) -> None:
2389
+ """
2390
+ Finds the connected components of a graph using a Pregel-based algorithm.
2391
+
2392
+ Each vertex is assigned a component ID, which is the minimum vertex ID
2393
+ within its connected component.
2394
+
2395
+ :param connection: The pyocient Connection object.
2396
+ :param input_schema: Schema of the source graph tables.
2397
+ :param input_vertices_table: Name of the source vertices table.
2398
+ :param input_edges_table: Name of the source edges table.
2399
+ :param result_schema: Schema for the result table.
2400
+ :param result_vertices_table: Name for the table with component results.
2401
+ :param max_iterations: The maximum number of Pregel iterations.
2402
+ :param result_vertices_indexes: List of columns to index ('id' or 'result').
2403
+ :raises SQLException: If any database error occurs.
2404
+ """
2405
+ _check_not_null(connection, "Connection cannot be null.")
2406
+ _check_positive(max_iterations, "Maximum iterations must be >= 1.")
2407
+
2408
+ # Define Pregel parameters for Connected Components
2409
+ initializer_expr = "id" # Initial component ID = vertex ID
2410
+ send_to_source_expr = "CASE WHEN c.state < a.state THEN c.state ELSE NULL END"
2411
+ send_to_dest_expr = "CASE WHEN a.state < c.state THEN a.state ELSE NULL END"
2412
+ aggregate_expr = "MIN(msg)" # Aggregate by taking the minimum received ID
2413
+ # Update to smallest of current or received ID. COALESCE handles nodes with no messages.
2414
+ updater_expr = "LEAST(s.state, COALESCE(m.aggregated_message, s.state))"
2415
+
2416
+ pregel(
2417
+ connection=connection,
2418
+ input_schema=input_schema,
2419
+ input_vertices_table=input_vertices_table,
2420
+ input_edges_table=input_edges_table,
2421
+ result_schema=result_schema,
2422
+ result_vertices_table=result_vertices_table,
2423
+ initializer_expr=initializer_expr,
2424
+ send_to_source_expr=send_to_source_expr,
2425
+ send_to_dest_expr=send_to_dest_expr,
2426
+ aggregate_expr=aggregate_expr,
2427
+ updater_expr=updater_expr,
2428
+ max_iterations=max_iterations,
2429
+ result_vertices_indexes=result_vertices_indexes,
2430
+ )
2431
+
2432
+
2433
+ def strongly_connected_components(
2434
+ connection: Connection,
2435
+ input_schema: str,
2436
+ input_vertices_table: str,
2437
+ input_edges_table: str,
2438
+ result_schema: str,
2439
+ result_vertices_table: str,
2440
+ result_vertices_indexes: List[str],
2441
+ ) -> None:
2442
+ """
2443
+ Computes the Strongly Connected Components (SCCs) of a directed graph.
2444
+
2445
+ This method implements a recursive algorithm to partition the graph and
2446
+ identify sets of vertices where every vertex is reachable from every other
2447
+ vertex within the same set.
2448
+
2449
+ The final result is populated into the `result_vertices_table`, which maps
2450
+ each vertex `id` to a representative `component` ID (the minimum vertex
2451
+ ID within that component).
2452
+
2453
+ :param connection: The pyocient Connection object.
2454
+ :param input_schema: Schema of the source graph tables.
2455
+ :param input_vertices_table: Name of the source vertices table.
2456
+ :param input_edges_table: Name of the source edges table.
2457
+ :param result_schema: Schema for the result and temporary tables.
2458
+ :param result_vertices_table: Name for the table with SCC results.
2459
+ :param result_vertices_indexes: List of columns to index ('id' or 'component').
2460
+ :raises SQLException: If any database error occurs.
2461
+ """
2462
+ _check_not_null(connection, "Connection cannot be null.")
2463
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
2464
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
2465
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
2466
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
2467
+ _check_not_null_or_empty(result_vertices_table, "Result vertices table name cannot be null or empty.")
2468
+ _check_not_null(result_vertices_indexes, "Result vertices index list cannot be null.")
2469
+
2470
+ input_tables = {
2471
+ _table_name(input_schema, input_vertices_table),
2472
+ _table_name(input_schema, input_edges_table),
2473
+ }
2474
+ _check_table_conflict(result_schema, result_vertices_table, input_tables)
2475
+ final_result_table = _table_name(result_schema, result_vertices_table)
2476
+
2477
+ temp_tables = [] # Central tracker for ALL temp tables
2478
+ unique_suffix = f"scc_{uuid.uuid4().hex[:10]}"
2479
+ initial_vertices_base = f"v_start_{unique_suffix}"
2480
+ initial_edges_base = f"e_start_{unique_suffix}"
2481
+
2482
+ with connection.cursor() as cursor:
2483
+ try:
2484
+ # Create the final result table structure up front
2485
+ cursor.execute(f"CREATE TABLE {final_result_table} (id BIGINT NOT NULL, component BIGINT NOT NULL)")
2486
+ for col in result_vertices_indexes:
2487
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
2488
+ if col.lower() not in ("id", "component"):
2489
+ raise SQLException(f"Argument Error: Invalid index column '{col}' for SCC result.")
2490
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
2491
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {final_result_table} ("{col}")')
2492
+
2493
+ # Create initial temporary copies of the graph
2494
+ initial_vertices_table = _table_name(result_schema, initial_vertices_base)
2495
+ initial_edges_table = _table_name(result_schema, initial_edges_base)
2496
+
2497
+ cursor.execute(
2498
+ f"CREATE TABLE {initial_vertices_table} AS SELECT * FROM {_table_name(input_schema, input_vertices_table)}"
2499
+ )
2500
+ temp_tables.append(initial_vertices_table)
2501
+
2502
+ cursor.execute(
2503
+ f"CREATE TABLE {initial_edges_table} AS SELECT * FROM {_table_name(input_schema, input_edges_table)}"
2504
+ )
2505
+ temp_tables.append(initial_edges_table)
2506
+
2507
+ # Start the recursive process
2508
+ _process_component(
2509
+ cursor,
2510
+ result_schema,
2511
+ initial_vertices_base,
2512
+ initial_edges_base,
2513
+ result_schema,
2514
+ result_vertices_table,
2515
+ temp_tables,
2516
+ f"{unique_suffix}_r0",
2517
+ )
2518
+ except SQLException as e:
2519
+ # If anything fails, drop the final table too
2520
+ _safe_drop_table(cursor, final_result_table)
2521
+ raise e
2522
+ finally:
2523
+ # Clean up all temporary tables
2524
+ for temp_table in reversed(temp_tables):
2525
+ _safe_drop_table(cursor, temp_table)
2526
+
2527
+
2528
+ def _process_component(
2529
+ cursor: Cursor,
2530
+ schema: str,
2531
+ vertices_base: str,
2532
+ edges_base: str,
2533
+ result_schema: str,
2534
+ result_table: str,
2535
+ temp_tables: List[str],
2536
+ recur_suffix: str,
2537
+ ) -> None:
2538
+ """Recursive core of the SCC algorithm."""
2539
+ current_vertices = _table_name(schema, vertices_base)
2540
+
2541
+ # Base case: if the current graph partition is empty, stop.
2542
+ cursor.execute(f"SELECT COUNT(*) FROM {current_vertices}")
2543
+ row = cursor.fetchone()
2544
+ if row is None or row[0] == 0:
2545
+ return
2546
+
2547
+ # 1. Select a pivot vertex
2548
+ cursor.execute(f"SELECT MIN(id) FROM {current_vertices}")
2549
+ pivot_result = cursor.fetchone()
2550
+ if not pivot_result or pivot_result[0] is None:
2551
+ return
2552
+ pivot = pivot_result[0]
2553
+
2554
+ # 2. Compute predecessors (Pred) and descendants (Desc) of the pivot
2555
+ pred_base = f"pred_{recur_suffix}"
2556
+ desc_base = f"desc_{recur_suffix}"
2557
+
2558
+ _compute_descendants(cursor, schema, edges_base, pivot, desc_base, temp_tables)
2559
+ _compute_predecessors(cursor, schema, edges_base, pivot, pred_base, temp_tables, recur_suffix)
2560
+
2561
+ # 3. The current SCC is the intersection of Pred and Desc
2562
+ scc_base = f"scc_{recur_suffix}"
2563
+ scc_table = _table_name(schema, scc_base)
2564
+ pred_table = _table_name(schema, pred_base)
2565
+ desc_table = _table_name(schema, desc_base)
2566
+
2567
+ cursor.execute(f"CREATE TABLE {scc_table} AS SELECT p.id FROM {pred_table} p JOIN {desc_table} d ON p.id = d.id")
2568
+ temp_tables.append(scc_table)
2569
+
2570
+ # 4. Insert the found SCC into the final result table
2571
+ final_result_table = _table_name(result_schema, result_table)
2572
+ cursor.execute(
2573
+ f"INSERT INTO {final_result_table} (id, component) SELECT id, MIN(id) OVER () AS component FROM {scc_table}"
2574
+ )
2575
+
2576
+ # 5. Recurse on the three remaining partitions of the graph
2577
+ # Partition 1: Desc \ SCC
2578
+ _process_partition(
2579
+ cursor,
2580
+ schema,
2581
+ vertices_base,
2582
+ edges_base,
2583
+ result_schema,
2584
+ result_table,
2585
+ desc_base,
2586
+ scc_base,
2587
+ "desc",
2588
+ temp_tables,
2589
+ f"{recur_suffix}_d",
2590
+ )
2591
+
2592
+ # Partition 2: Pred \ SCC
2593
+ _process_partition(
2594
+ cursor,
2595
+ schema,
2596
+ vertices_base,
2597
+ edges_base,
2598
+ result_schema,
2599
+ result_table,
2600
+ pred_base,
2601
+ scc_base,
2602
+ "pred",
2603
+ temp_tables,
2604
+ f"{recur_suffix}_p",
2605
+ )
2606
+
2607
+ # Partition 3: Remainder = V \ (Pred U Desc)
2608
+ _process_remaining(
2609
+ cursor,
2610
+ schema,
2611
+ vertices_base,
2612
+ edges_base,
2613
+ result_schema,
2614
+ result_table,
2615
+ pred_base,
2616
+ desc_base,
2617
+ temp_tables,
2618
+ f"{recur_suffix}_r",
2619
+ )
2620
+
2621
+
2622
+ def _compute_descendants(
2623
+ cursor: Cursor, schema: str, edges_base: str, pivot: int, result_base: str, temp_tables: List[str]
2624
+ ) -> None:
2625
+ """Iteratively finds all vertices reachable from the pivot."""
2626
+ result_table = _table_name(schema, result_base)
2627
+ current_edges = _table_name(schema, edges_base)
2628
+
2629
+ cursor.execute(f"CREATE TABLE {result_table} (id BIGINT NOT NULL)")
2630
+ temp_tables.append(result_table)
2631
+ cursor.execute(f"INSERT INTO {result_table} (id) VALUES ({pivot})")
2632
+
2633
+ while True:
2634
+ new_nodes_sql = f"""
2635
+ INSERT INTO {result_table} (id)
2636
+ SELECT DISTINCT e.destid FROM {result_table} d JOIN {current_edges} e ON d.id = e.srcid
2637
+ WHERE NOT EXISTS (SELECT 1 FROM {result_table} existed WHERE existed.id = e.destid)
2638
+ """
2639
+ cursor.execute(new_nodes_sql)
2640
+ if cursor.rowcount == 0:
2641
+ break
2642
+
2643
+
2644
+ def _compute_predecessors(
2645
+ cursor: Cursor,
2646
+ schema: str,
2647
+ edges_base: str,
2648
+ pivot: int,
2649
+ result_base: str,
2650
+ temp_tables: List[str],
2651
+ recur_suffix: str,
2652
+ ) -> None:
2653
+ """Finds all vertices that can reach the pivot by reversing edges and finding descendants."""
2654
+ temp_rev_edges_base = f"rev_{recur_suffix}"
2655
+ temp_rev_edges_table = _table_name(schema, temp_rev_edges_base)
2656
+
2657
+ try:
2658
+ # Create a temporary reversed graph
2659
+ reverse_edges(
2660
+ connection=cursor.connection,
2661
+ input_schema=schema,
2662
+ input_edges_table=edges_base,
2663
+ result_schema=schema,
2664
+ result_edges_table=temp_rev_edges_base,
2665
+ result_edges_indexes=[],
2666
+ )
2667
+ temp_tables.append(temp_rev_edges_table)
2668
+
2669
+ # Find "descendants" in the reversed graph, which are predecessors in the original
2670
+ _compute_descendants(cursor, schema, temp_rev_edges_base, pivot, result_base, temp_tables)
2671
+ finally:
2672
+ _safe_drop_table(cursor, temp_rev_edges_table)
2673
+
2674
+
2675
+ def _process_partition(
2676
+ cursor: Cursor,
2677
+ schema: str,
2678
+ vertices_base: str,
2679
+ edges_base: str,
2680
+ result_schema: str,
2681
+ result_table: str,
2682
+ partition_set_base: str,
2683
+ scc_base: str,
2684
+ type_prefix: str,
2685
+ temp_tables: List[str],
2686
+ recur_suffix: str,
2687
+ ) -> None:
2688
+ r"""Creates a subgraph for a partition (e.g., Desc \ SCC) and recurses."""
2689
+ filtered_v_base = f"{type_prefix}_v_{recur_suffix}"
2690
+ filtered_e_base = f"{type_prefix}_e_{recur_suffix}"
2691
+
2692
+ current_vertices = _table_name(schema, vertices_base)
2693
+ current_edges = _table_name(schema, edges_base)
2694
+ partition_set_table = _table_name(schema, partition_set_base)
2695
+ scc_table = _table_name(schema, scc_base)
2696
+
2697
+ # Create the new vertex set for the partition
2698
+ filtered_v_table = _table_name(schema, filtered_v_base)
2699
+ cursor.execute(f"""
2700
+ CREATE TABLE {filtered_v_table} AS
2701
+ SELECT v.* FROM {current_vertices} v JOIN {partition_set_table} p ON v.id = p.id
2702
+ WHERE NOT EXISTS (SELECT 1 FROM {scc_table} scc WHERE scc.id = v.id)
2703
+ """)
2704
+ temp_tables.append(filtered_v_table)
2705
+
2706
+ cursor.execute(f"SELECT COUNT(*) FROM {filtered_v_table}")
2707
+ row = cursor.fetchone()
2708
+ if row is not None and row[0] > 0:
2709
+ # Create the corresponding edge set
2710
+ filtered_e_table = _table_name(schema, filtered_e_base)
2711
+ cursor.execute(f"""
2712
+ CREATE TABLE {filtered_e_table} AS
2713
+ SELECT e.* FROM {current_edges} e
2714
+ JOIN {filtered_v_table} src_v ON e.srcid = src_v.id
2715
+ JOIN {filtered_v_table} dest_v ON e.destid = dest_v.id
2716
+ """)
2717
+ temp_tables.append(filtered_e_table)
2718
+
2719
+ # Recurse on this new subgraph
2720
+ _process_component(
2721
+ cursor, schema, filtered_v_base, filtered_e_base, result_schema, result_table, temp_tables, recur_suffix
2722
+ )
2723
+
2724
+
2725
+ def _process_remaining(
2726
+ cursor: Cursor,
2727
+ schema: str,
2728
+ vertices_base: str,
2729
+ edges_base: str,
2730
+ result_schema: str,
2731
+ result_table: str,
2732
+ pred_base: str,
2733
+ desc_base: str,
2734
+ temp_tables: List[str],
2735
+ recur_suffix: str,
2736
+ ) -> None:
2737
+ r"""Creates the subgraph for the remainder partition V \ (Pred U Desc) and recurses."""
2738
+ remaining_v_base = f"rem_v_{recur_suffix}"
2739
+ remaining_e_base = f"rem_e_{recur_suffix}"
2740
+
2741
+ current_vertices = _table_name(schema, vertices_base)
2742
+ current_edges = _table_name(schema, edges_base)
2743
+ pred_table = _table_name(schema, pred_base)
2744
+ desc_table = _table_name(schema, desc_base)
2745
+
2746
+ # Create the vertex set for the remainder
2747
+ remaining_v_table = _table_name(schema, remaining_v_base)
2748
+ cursor.execute(f"""
2749
+ CREATE TABLE {remaining_v_table} AS
2750
+ SELECT v.* FROM {current_vertices} v
2751
+ WHERE NOT EXISTS (SELECT 1 FROM {pred_table} pred WHERE pred.id = v.id)
2752
+ AND NOT EXISTS (SELECT 1 FROM {desc_table} desc WHERE desc.id = v.id)
2753
+ """)
2754
+ temp_tables.append(remaining_v_table)
2755
+
2756
+ cursor.execute(f"SELECT COUNT(*) FROM {remaining_v_table}")
2757
+ row = cursor.fetchone()
2758
+ if row is not None and row[0] > 0:
2759
+ # Create the corresponding edge set
2760
+ remaining_e_table = _table_name(schema, remaining_e_base)
2761
+ cursor.execute(f"""
2762
+ CREATE TABLE {remaining_e_table} AS
2763
+ SELECT e.* FROM {current_edges} e
2764
+ JOIN {remaining_v_table} src_v ON e.srcid = src_v.id
2765
+ JOIN {remaining_v_table} dest_v ON e.destid = dest_v.id
2766
+ """)
2767
+ temp_tables.append(remaining_e_table)
2768
+
2769
+ # Recurse on this new subgraph
2770
+ _process_component(
2771
+ cursor, schema, remaining_v_base, remaining_e_base, result_schema, result_table, temp_tables, recur_suffix
2772
+ )
2773
+
2774
+
2775
+ def shortest_paths(
2776
+ connection: Connection,
2777
+ input_schema: str,
2778
+ input_vertices_table: str,
2779
+ input_edges_table: str,
2780
+ result_schema: str,
2781
+ result_table: str,
2782
+ landmarks: List[int],
2783
+ edge_weight_column: Optional[str],
2784
+ max_iterations: int,
2785
+ result_indexes: List[str],
2786
+ ) -> None:
2787
+ """
2788
+ Computes the shortest paths from all vertices to a set of landmark vertices.
2789
+ This method uses an iterative relaxation algorithm (similar to Bellman-Ford)
2790
+ to find the shortest distance from every vertex to each landmark. It assumes
2791
+ all edge weights are non-negative. If no weight column is specified,
2792
+ edges have a uniform weight of 1.0.
2793
+ The result table will contain `srcid`, `destid` (the landmark), and the
2794
+ computed `distance`. Only paths with finite distances are stored.
2795
+ :param connection: The pyocient Connection object.
2796
+ :param input_schema: Schema of the source graph tables.
2797
+ :param input_vertices_table: Name of the source vertices table.
2798
+ :param input_edges_table: Name of the source edges table.
2799
+ :param result_schema: Schema for the result and temporary tables.
2800
+ :param result_table: Name for the table with shortest path results.
2801
+ :param landmarks: A non-empty list of vertex IDs to use as landmarks.
2802
+ :param edge_weight_column: The name of the column containing edge weights.
2803
+ If None or empty, weight is 1.0.
2804
+ :param max_iterations: Maximum number of iterations for the algorithm.
2805
+ :param result_indexes: List of columns to index ('srcid', 'destid', 'distance').
2806
+ :raises SQLException: If any database error occurs.
2807
+ """
2808
+ _check_not_null(connection, "Connection cannot be null.")
2809
+ _check_not_null_or_empty(input_schema, "Input schema cannot be null or empty.")
2810
+ _check_not_null_or_empty(input_vertices_table, "Input vertices table name cannot be null or empty.")
2811
+ _check_not_null_or_empty(input_edges_table, "Input edges table name cannot be null or empty.")
2812
+ _check_not_null_or_empty(result_schema, "Result schema cannot be null or empty.")
2813
+ _check_not_null_or_empty(result_table, "Result table name cannot be null or empty.")
2814
+ _check_not_null(landmarks, "Landmarks list cannot be null.")
2815
+ if not landmarks:
2816
+ raise SQLException("Argument Error: Landmarks list cannot be empty.")
2817
+ if any(lm is None for lm in landmarks):
2818
+ raise SQLException("Argument Error: Landmarks list cannot contain null values.")
2819
+ _check_positive(max_iterations, "Maximum iterations must be >= 1.")
2820
+ _check_not_null(result_indexes, "Result index list cannot be null.")
2821
+
2822
+ input_tables = {
2823
+ _table_name(input_schema, input_vertices_table),
2824
+ _table_name(input_schema, input_edges_table),
2825
+ }
2826
+ _check_table_conflict(result_schema, result_table, input_tables)
2827
+
2828
+ weight_expr = f'e."{edge_weight_column.strip()}"' if edge_weight_column and edge_weight_column.strip() else "1.0"
2829
+
2830
+ vertices = _table_name(input_schema, input_vertices_table)
2831
+ edges = _table_name(input_schema, input_edges_table)
2832
+ result = _table_name(result_schema, result_table)
2833
+
2834
+ temp_base = f"sp_temp_{uuid.uuid4().hex[:8]}"
2835
+ current_dists_base = f"{temp_base}_curr"
2836
+ next_dists_base = f"{temp_base}_next"
2837
+
2838
+ qualified_current_dists = _table_name(result_schema, current_dists_base)
2839
+ qualified_next_dists = _table_name(result_schema, next_dists_base)
2840
+
2841
+ temp_tables_created = []
2842
+
2843
+ with connection.cursor() as cursor:
2844
+ try:
2845
+ landmarks_str = f"ARRAY[{','.join(map(str, landmarks))}]"
2846
+
2847
+ init_sql = f"""
2848
+ CREATE TABLE {qualified_current_dists} AS
2849
+ SELECT v.id AS srcid, l.landmark AS destid,
2850
+ CASE WHEN v.id = l.landmark THEN 0.0 ELSE CAST('1.0e308' AS DOUBLE PRECISION) END AS distance
2851
+ FROM {vertices} v
2852
+ CROSS JOIN (SELECT UNNEST({landmarks_str}) AS landmark) l
2853
+ """
2854
+ cursor.execute(init_sql)
2855
+ temp_tables_created.append(qualified_current_dists) # CORRECTED
2856
+
2857
+ changed = True
2858
+ for i in range(max_iterations):
2859
+ if not changed:
2860
+ break
2861
+
2862
+ _safe_drop_table(cursor, qualified_next_dists)
2863
+
2864
+ merge_sql = f"""
2865
+ CREATE TABLE {qualified_next_dists} AS
2866
+ WITH new_paths AS (
2867
+ SELECT e.srcid, d.destid, MIN(d.distance + {weight_expr}) AS new_distance
2868
+ FROM {qualified_current_dists} d JOIN {edges} e ON d.srcid = e.destid
2869
+ WHERE d.distance != CAST('1.0e308' AS DOUBLE PRECISION)
2870
+ GROUP BY e.srcid, d.destid
2871
+ )
2872
+ SELECT
2873
+ COALESCE(d.srcid, n.srcid) AS srcid,
2874
+ COALESCE(d.destid, n.destid) AS destid,
2875
+ LEAST(
2876
+ COALESCE(d.distance, CAST('1.0e308' AS DOUBLE PRECISION)),
2877
+ COALESCE(n.new_distance, CAST('1.0e308' AS DOUBLE PRECISION))
2878
+ ) AS distance
2879
+ FROM {qualified_current_dists} d FULL OUTER JOIN new_paths n
2880
+ ON d.srcid = n.srcid AND d.destid = n.destid
2881
+ """
2882
+ cursor.execute(merge_sql)
2883
+ temp_tables_created.append(qualified_next_dists) # CORRECTED
2884
+
2885
+ check_change_sql = f"""
2886
+ SELECT COUNT(*) FROM {qualified_next_dists} next JOIN {qualified_current_dists} curr
2887
+ ON next.srcid = curr.srcid AND next.destid = curr.destid
2888
+ WHERE next.distance < curr.distance
2889
+ """
2890
+ cursor.execute(check_change_sql)
2891
+ row = cursor.fetchone()
2892
+ num_changes = row[0] if row is not None else 0
2893
+ changed = num_changes > 0
2894
+
2895
+ qualified_current_dists, qualified_next_dists = qualified_next_dists, qualified_current_dists
2896
+
2897
+ cursor.execute(
2898
+ f"CREATE TABLE {result} (srcid BIGINT NOT NULL, destid BIGINT NOT NULL, distance DOUBLE PRECISION)"
2899
+ )
2900
+ for col in result_indexes:
2901
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
2902
+ if col.lower() not in ("srcid", "destid", "distance"):
2903
+ raise SQLException(f"Argument Error: Invalid index column '{col}'.")
2904
+ index_name = f"{result_table}_{col}_idx".replace('"', "")
2905
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {result} ("{col}")')
2906
+
2907
+ cursor.execute(
2908
+ f"INSERT INTO {result} (srcid, destid, distance) "
2909
+ f"SELECT srcid, destid, distance FROM {qualified_current_dists} "
2910
+ f"WHERE distance != CAST('1.0e308' AS DOUBLE PRECISION)"
2911
+ )
2912
+ finally:
2913
+ _safe_drop_table(cursor, qualified_current_dists)
2914
+ _safe_drop_table(cursor, qualified_next_dists)
2915
+
2916
+
2917
+ def _get_vertex_count(cursor: Cursor, schema: str, table: str) -> int:
2918
+ """Gets the total number of rows from a table."""
2919
+ qualified_table_name = _table_name(schema, table)
2920
+ cursor.execute(f"SELECT COUNT(*) FROM {qualified_table_name}")
2921
+ result = cursor.fetchone()
2922
+ if result and result[0] is not None:
2923
+ return int(result[0])
2924
+ raise SQLException(f"Database Error: Could not retrieve vertex count from {qualified_table_name}.")
2925
+
2926
+
2927
+ def static_page_rank(
2928
+ connection: Connection,
2929
+ input_schema: str,
2930
+ input_vertices_table: str,
2931
+ input_edges_table: str,
2932
+ result_schema: str,
2933
+ result_vertices_table: str,
2934
+ num_iterations: int,
2935
+ reset_prob: float,
2936
+ result_vertices_indexes: List[str],
2937
+ personalization_src_id: Optional[int],
2938
+ ) -> None:
2939
+ """
2940
+ Computes static PageRank for a fixed number of iterations.
2941
+
2942
+ This implementation is optimized to avoid copying tables between iterations
2943
+ by swapping the roles of two temporary tables.
2944
+
2945
+ :param connection: The pyocient Connection object.
2946
+ :param input_schema: Schema of the source graph tables.
2947
+ :param input_vertices_table: Name of the source vertices table.
2948
+ :param input_edges_table: Name of the source edges table.
2949
+ :param result_schema: Schema for the result table.
2950
+ :param result_vertices_table: Name for the table with PageRank scores.
2951
+ :param num_iterations: The number of PageRank iterations to perform.
2952
+ :param reset_prob: The reset probability (damping factor), typically 0.85.
2953
+ :param result_vertices_indexes: List of columns to index in the result.
2954
+ :param personalization_src_id: If not None, computes Personalized PageRank
2955
+ biased towards this vertex ID.
2956
+ :raises SQLException: If any database error occurs.
2957
+ """
2958
+ _check_not_null(connection, "Connection cannot be null.")
2959
+ _check_positive(num_iterations, "Number of iterations must be positive.")
2960
+ if not (0.0 <= reset_prob <= 1.0):
2961
+ raise SQLException("Reset probability must be between 0.0 and 1.0.")
2962
+
2963
+ qualified_input_v = _table_name(input_schema, input_vertices_table)
2964
+ qualified_input_e = _table_name(input_schema, input_edges_table)
2965
+ qualified_result_v = _table_name(result_schema, result_vertices_table)
2966
+
2967
+ _check_table_conflict(result_schema, result_vertices_table, {qualified_input_v, qualified_input_e})
2968
+
2969
+ unique_suffix = uuid.uuid4().hex
2970
+ current_ranks_base = f"pagerank_curr_{unique_suffix}"
2971
+ previous_ranks_base = f"pagerank_prev_{unique_suffix}"
2972
+ qualified_current_ranks = _table_name(result_schema, current_ranks_base)
2973
+ qualified_previous_ranks = _table_name(result_schema, previous_ranks_base)
2974
+
2975
+ created_tables = set()
2976
+
2977
+ with connection.cursor() as cursor:
2978
+ try:
2979
+ vertex_count = _get_vertex_count(cursor, input_schema, input_vertices_table)
2980
+ if vertex_count == 0:
2981
+ cursor.execute(
2982
+ f"CREATE TABLE {qualified_result_v} AS "
2983
+ f"SELECT v.*, CAST(NULL AS DOUBLE PRECISION) AS pagerank "
2984
+ f"FROM {qualified_input_v} v WHERE 1=0"
2985
+ )
2986
+ return
2987
+
2988
+ # Initialize the first "current" ranks table
2989
+ if personalization_src_id is not None:
2990
+ init_sql = (
2991
+ f"SELECT id, CASE WHEN id = {personalization_src_id} "
2992
+ f"THEN 1.0 ELSE 0.0 END AS rank FROM {qualified_input_v}"
2993
+ )
2994
+ else:
2995
+ init_rank = 1.0 / vertex_count
2996
+ init_sql = f"SELECT id, CAST({init_rank} AS DOUBLE PRECISION) AS rank FROM {qualified_input_v}"
2997
+ cursor.execute(f"CREATE TABLE {qualified_current_ranks} AS {init_sql}")
2998
+ created_tables.add(qualified_current_ranks)
2999
+
3000
+ # --- Iterations ---
3001
+ for _ in range(num_iterations):
3002
+ # OPTIMIZATION: Swap variables, not table data.
3003
+ # This is an instantaneous metadata change.
3004
+ qualified_current_ranks, qualified_previous_ranks = (
3005
+ qualified_previous_ranks,
3006
+ qualified_current_ranks,
3007
+ )
3008
+
3009
+ _safe_drop_table(cursor, qualified_current_ranks)
3010
+ created_tables.discard(qualified_current_ranks)
3011
+
3012
+ # The `qualified_previous_ranks` table now holds the ranks from the last iteration.
3013
+ # The iteration SQL will create a new `qualified_current_ranks` table.
3014
+
3015
+ teleport_sql = f"(1.0 - {reset_prob}) / {vertex_count}" if personalization_src_id is None else "0.0"
3016
+ personalization_sql = (
3017
+ ""
3018
+ if personalization_src_id is None
3019
+ else (f" + CASE WHEN v.id = {personalization_src_id} THEN {1.0 - reset_prob} ELSE 0.0 END")
3020
+ )
3021
+
3022
+ iteration_sql = f"""
3023
+ CREATE TABLE {qualified_current_ranks} AS
3024
+ WITH EdgeContributions AS (
3025
+ SELECT e.destid, SUM(pr.rank / out_degree.out_cnt) as contrib
3026
+ FROM {qualified_input_e} e
3027
+ JOIN {qualified_previous_ranks} pr ON e.srcid = pr.id
3028
+ JOIN (SELECT srcid, COUNT(*) AS out_cnt FROM {qualified_input_e} GROUP BY srcid HAVING COUNT(*) > 0) out_degree ON e.srcid = out_degree.srcid
3029
+ GROUP BY e.destid
3030
+ ), DanglingRank AS (
3031
+ SELECT COALESCE(SUM(pr.rank), 0.0) as total_rank
3032
+ FROM {qualified_previous_ranks} pr
3033
+ WHERE pr.id NOT IN (SELECT DISTINCT srcid FROM {qualified_input_e})
3034
+ )
3035
+ SELECT
3036
+ v.id,
3037
+ CAST( ({teleport_sql})
3038
+ + ({reset_prob} * COALESCE(ec.contrib, 0.0))
3039
+ + ({reset_prob} * dr.total_rank / {vertex_count})
3040
+ {personalization_sql}
3041
+ AS DOUBLE PRECISION) AS rank
3042
+ FROM {qualified_input_v} v
3043
+ LEFT JOIN EdgeContributions ec ON v.id = ec.destid
3044
+ CROSS JOIN DanglingRank dr
3045
+ """
3046
+ cursor.execute(iteration_sql)
3047
+ created_tables.add(qualified_current_ranks)
3048
+
3049
+ # Final result is in the last table pointed to by `qualified_current_ranks`
3050
+ cursor.execute(
3051
+ f"CREATE TABLE {qualified_result_v} AS "
3052
+ f"SELECT v.*, tr.rank AS pagerank "
3053
+ f"FROM {qualified_input_v} v JOIN {qualified_current_ranks} tr ON v.id = tr.id"
3054
+ )
3055
+ created_tables.add(qualified_result_v)
3056
+
3057
+ for col in result_vertices_indexes:
3058
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
3059
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
3060
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {qualified_result_v} ("{col}")')
3061
+
3062
+ finally:
3063
+ _safe_drop_table(cursor, qualified_current_ranks)
3064
+ _safe_drop_table(cursor, qualified_previous_ranks)
3065
+
3066
+
3067
+ def dynamic_page_rank(
3068
+ connection: Connection,
3069
+ input_schema: str,
3070
+ input_vertices_table: str,
3071
+ input_edges_table: str,
3072
+ result_schema: str,
3073
+ result_vertices_table: str,
3074
+ tolerance: float,
3075
+ reset_prob: float,
3076
+ result_vertices_indexes: List[str],
3077
+ personalization_src_id: Optional[int],
3078
+ ) -> None:
3079
+ """
3080
+ Computes PageRank scores, iterating until the scores converge.
3081
+
3082
+ This implementation is optimized to avoid copying tables between iterations
3083
+ by swapping the roles of two temporary tables.
3084
+
3085
+ :param connection: The pyocient Connection object.
3086
+ :param input_schema: Schema of the source graph tables.
3087
+ :param input_vertices_table: Name of the source vertices table.
3088
+ :param input_edges_table: Name of the source edges table.
3089
+ :param result_schema: Schema for the result table.
3090
+ :param result_vertices_table: Name for the table with PageRank scores.
3091
+ :param tolerance: The convergence tolerance.
3092
+ :param reset_prob: The reset probability (damping factor), typically 0.85.
3093
+ :param result_vertices_indexes: List of columns to index in the result.
3094
+ :param personalization_src_id: If not None, computes Personalized PageRank
3095
+ biased towards this vertex ID.
3096
+ :raises SQLException: If any database error occurs.
3097
+ """
3098
+ _check_not_null(connection, "Connection cannot be null.")
3099
+ if tolerance < 0.0:
3100
+ raise SQLException("Tolerance must be non-negative.")
3101
+ if not (0.0 <= reset_prob <= 1.0):
3102
+ raise SQLException("Reset probability must be between 0.0 and 1.0.")
3103
+
3104
+ qualified_input_v = _table_name(input_schema, input_vertices_table)
3105
+ qualified_input_e = _table_name(input_schema, input_edges_table)
3106
+ qualified_result_v = _table_name(result_schema, result_vertices_table)
3107
+
3108
+ _check_table_conflict(result_schema, result_vertices_table, {qualified_input_v, qualified_input_e})
3109
+
3110
+ unique_suffix = uuid.uuid4().hex
3111
+ current_ranks_base = f"pagerank_curr_{unique_suffix}"
3112
+ previous_ranks_base = f"pagerank_prev_{unique_suffix}"
3113
+ qualified_current_ranks = _table_name(result_schema, current_ranks_base)
3114
+ qualified_previous_ranks = _table_name(result_schema, previous_ranks_base)
3115
+
3116
+ created_tables = set()
3117
+
3118
+ with connection.cursor() as cursor:
3119
+ try:
3120
+ vertex_count = _get_vertex_count(cursor, input_schema, input_vertices_table)
3121
+ if vertex_count == 0:
3122
+ cursor.execute(
3123
+ f"CREATE TABLE {qualified_result_v} AS "
3124
+ f"SELECT v.*, CAST(NULL AS DOUBLE PRECISION) AS pagerank "
3125
+ f"FROM {qualified_input_v} v WHERE 1=0"
3126
+ )
3127
+ return
3128
+
3129
+ # Initialize the first "current" ranks table
3130
+ if personalization_src_id is not None:
3131
+ init_sql = (
3132
+ f"SELECT id, CASE WHEN id = {personalization_src_id} "
3133
+ f"THEN 1.0 ELSE 0.0 END AS rank FROM {qualified_input_v}"
3134
+ )
3135
+ else:
3136
+ init_rank = 1.0 / vertex_count
3137
+ init_sql = f"SELECT id, CAST({init_rank} AS DOUBLE PRECISION) AS rank FROM {qualified_input_v}"
3138
+ cursor.execute(f"CREATE TABLE {qualified_current_ranks} AS {init_sql}")
3139
+ created_tables.add(qualified_current_ranks)
3140
+
3141
+ delta = float("inf")
3142
+ while delta > tolerance:
3143
+ # OPTIMIZATION: Swap variables, not table data.
3144
+ qualified_current_ranks, qualified_previous_ranks = (
3145
+ qualified_previous_ranks,
3146
+ qualified_current_ranks,
3147
+ )
3148
+
3149
+ _safe_drop_table(cursor, qualified_current_ranks)
3150
+ created_tables.discard(qualified_current_ranks)
3151
+
3152
+ teleport_sql = f"(1.0 - {reset_prob}) / {vertex_count}" if personalization_src_id is None else "0.0"
3153
+ personalization_sql = (
3154
+ ""
3155
+ if personalization_src_id is None
3156
+ else (f" + CASE WHEN v.id = {personalization_src_id} THEN {1.0 - reset_prob} ELSE 0.0 END")
3157
+ )
3158
+
3159
+ iteration_sql = f"""
3160
+ CREATE TABLE {qualified_current_ranks} AS
3161
+ WITH EdgeContributions AS (
3162
+ SELECT e.destid, SUM(pr.rank / out_degree.out_cnt) as contrib
3163
+ FROM {qualified_input_e} e
3164
+ JOIN {qualified_previous_ranks} pr ON e.srcid = pr.id
3165
+ JOIN (SELECT srcid, COUNT(*) AS out_cnt FROM {qualified_input_e} GROUP BY srcid HAVING COUNT(*) > 0) out_degree ON e.srcid = out_degree.srcid
3166
+ GROUP BY e.destid
3167
+ ), DanglingRank AS (
3168
+ SELECT COALESCE(SUM(pr.rank), 0.0) as total_rank
3169
+ FROM {qualified_previous_ranks} pr
3170
+ WHERE pr.id NOT IN (SELECT DISTINCT srcid FROM {qualified_input_e})
3171
+ )
3172
+ SELECT
3173
+ v.id,
3174
+ CAST( ({teleport_sql})
3175
+ + ({reset_prob} * COALESCE(ec.contrib, 0.0))
3176
+ + ({reset_prob} * dr.total_rank / {vertex_count})
3177
+ {personalization_sql}
3178
+ AS DOUBLE PRECISION) AS rank
3179
+ FROM {qualified_input_v} v
3180
+ LEFT JOIN EdgeContributions ec ON v.id = ec.destid
3181
+ CROSS JOIN DanglingRank dr
3182
+ """
3183
+ cursor.execute(iteration_sql)
3184
+ created_tables.add(qualified_current_ranks)
3185
+
3186
+ # Check for convergence by comparing the new and old tables
3187
+ delta_sql = (
3188
+ f"SELECT SUM(ABS(curr.rank - prev.rank)) "
3189
+ f"FROM {qualified_current_ranks} curr "
3190
+ f"JOIN {qualified_previous_ranks} prev ON curr.id = prev.id"
3191
+ )
3192
+ cursor.execute(delta_sql)
3193
+ delta_result = cursor.fetchone()
3194
+ delta = delta_result[0] if delta_result and delta_result[0] is not None else 0.0
3195
+
3196
+ # Final result
3197
+ cursor.execute(
3198
+ f"CREATE TABLE {qualified_result_v} AS "
3199
+ f"SELECT v.*, tr.rank AS pagerank "
3200
+ f"FROM {qualified_input_v} v JOIN {qualified_current_ranks} tr ON v.id = tr.id"
3201
+ )
3202
+ created_tables.add(qualified_result_v)
3203
+
3204
+ for col in result_vertices_indexes:
3205
+ _check_not_null_or_empty(col, "Index column name cannot be null or empty.")
3206
+ index_name = f"{result_vertices_table}_{col}_idx".replace('"', "")
3207
+ cursor.execute(f'CREATE INDEX "{index_name}" ON {qualified_result_v} ("{col}")')
3208
+
3209
+ finally:
3210
+ _safe_drop_table(cursor, qualified_current_ranks)
3211
+ _safe_drop_table(cursor, qualified_previous_ranks)