From e21013f5fbc45a5c5c079fedb15a820fe3063180 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Mon, 17 Aug 2026 13:36:20 -0700 Subject: [PATCH] feat: support the duckdb backend in from_arrow_tables Add build_csr_from_arrow_tables() to the DuckDB backend: registers the node and edge Arrow tables directly in an in-memory DuckDB instance and runs the same relational plan as the Parquet-pair converter (endpoint joins, (source, target) sort, degree histogram), honoring memory_limit via SET memory_limit so the external sort spills to disk. Results are read back as Arrow tables with to_arrow_table() (arrow() returns a RecordBatchReader for empty results in duckdb 1.5). IcebugMemGraph.from_arrow_tables now accepts backend='duckdb' alongside 'pyarrow' and 'datafusion'; the memory_limit auto-switch still prefers datafusion when the default backend is left unchanged. --- icebug_format/_convert_duckdb.py | 156 +++++++++++++++++++++++++++++++ icebug_format/memory.py | 35 +++++-- tests/test_memory.py | 43 +++++---- 3 files changed, 207 insertions(+), 27 deletions(-) diff --git a/icebug_format/_convert_duckdb.py b/icebug_format/_convert_duckdb.py index 67718a1..ed9e024 100644 --- a/icebug_format/_convert_duckdb.py +++ b/icebug_format/_convert_duckdb.py @@ -7,12 +7,18 @@ exceeded, so peak RSS stays bounded for large graphs. Output Parquet files are streamed from DuckDB via ``COPY ... TO ...``. + +The same plan is available from in-memory Arrow tables via +:func:`build_csr_from_arrow_tables`, which backs +:meth:`icebug_format.memory.IcebugMemGraph.from_arrow_tables` when the +``"duckdb"`` backend is selected. """ from __future__ import annotations from pathlib import Path +import pyarrow as pa import pyarrow.parquet as pq from icebug_format.convert_parquet import resolve_rel_column_names @@ -140,3 +146,153 @@ def convert_graph( ) finally: con.close() + + +def build_csr_from_arrow_tables( + node_table: pa.Table, + rel_table: pa.Table, + *, + to_node_table: pa.Table | None = None, + add_reverse_edges: bool = False, + memory_limit: str | None = None, +) -> tuple[pa.Table, pa.Table, pa.Table, pa.Table]: + """ + Build ``(src_out, dst_out, indices, indptr)`` with the DuckDB engine. + + Same contract as :func:`icebug_format._convert_pyarrow.build_csr_from_tables` + (the pure-PyArrow path), but the node sort, endpoint mapping and + ``(source, target)`` sort run inside DuckDB with the vertex/edge tables + registered from Arrow, so *memory_limit* bounds DuckDB's buffer pool and + the external sort spills to disk instead of growing RSS. Mirrors + :func:`convert_graph` for Parquet pairs, with the tables registered from + Arrow rather than Parquet files. Node tables are always ordered by primary + key (there is no ``input_sorted`` shortcut). + """ + from icebug_format._duckdb import require_duckdb + from icebug_format.cli import set_memory_limit + + if add_reverse_edges and to_node_table is not None: + raise ValueError( + "to_node_arrow_table must not be provided when adding reverse edges; " + "from and to node tables must be the same." + ) + if rel_table.num_columns < 2: + raise ValueError( + f"rel_table must have at least 2 columns (source and target), " + f"got {rel_table.num_columns}" + ) + + if to_node_table is None: + to_node_table = node_table + + src_pk = node_table.schema.names[0] + dst_pk = to_node_table.schema.names[0] + src_col, dst_col = resolve_rel_column_names(rel_table.schema.names) + prop_cols = [c for c in rel_table.schema.names if c not in (src_col, dst_col)] + n_src = len(node_table) + + duckdb = require_duckdb("the 'convert' feature (Arrow table conversion)") + con = duckdb.connect() + if memory_limit: + set_memory_limit(con, memory_limit) + try: + con.register("vertices", node_table) + con.register("vertices_dst", to_node_table) + con.register("edges", rel_table) + + # Dense id -> CSR index mapping, ordered by primary key. + con.execute(f""" + CREATE TABLE src_map AS + SELECT row_number() OVER (ORDER BY {_q(src_pk)}) - 1 AS csr_index, + {_q(src_pk)} AS original_node_id + FROM vertices + """) + if to_node_table is node_table: + dst_map = "src_map" + else: + dst_map = "dst_map" + con.execute(f""" + CREATE TABLE dst_map AS + SELECT row_number() OVER (ORDER BY {_q(dst_pk)}) - 1 AS csr_index, + {_q(dst_pk)} AS original_node_id + FROM vertices_dst + """) + + select_fwd = "m1.csr_index AS csr_source, m2.csr_index AS csr_target" + select_rev = "m2.csr_index AS csr_source, m1.csr_index AS csr_target" + if prop_cols: + props = ", " + ", ".join(f"e.{_q(c)}" for c in prop_cols) + select_fwd += props + select_rev += props + join_clause = f""" + FROM edges e + JOIN src_map m1 ON e.{_q(src_col)} = m1.original_node_id + JOIN {dst_map} m2 ON e.{_q(dst_col)} = m2.original_node_id + """ + + if not add_reverse_edges: + rel_query = f"SELECT {select_fwd} {join_clause}" + else: + # Self-loops appear once (forward only); non-self edges get both directions. + rel_query = f""" + SELECT {select_fwd} {join_clause} + UNION ALL + SELECT {select_rev} {join_clause} + WHERE e.{_q(src_col)} != e.{_q(dst_col)} + """ + con.execute(f"CREATE TABLE relations AS {rel_query}") + + # indices: neighbour list sorted by (source, target) + select_props = ", " + ", ".join(_q(c) for c in prop_cols) if prop_cols else "" + con.execute(f""" + CREATE TABLE indices AS + SELECT csr_target::UBIGINT AS target{select_props} + FROM relations + ORDER BY csr_source, csr_target + """) + + # indptr: cumulative degree per source node, zero-filled, N+1 entries. + con.execute(f""" + CREATE TABLE indptr AS + WITH node_range AS ( + SELECT unnest(range(0, {n_src})) AS node_id + ), + degrees AS ( + SELECT csr_source AS src, COUNT(*) AS deg + FROM relations + GROUP BY csr_source + ), + cumulative AS ( + SELECT + node_range.node_id, + COALESCE( + SUM(degrees.deg) OVER ( + ORDER BY node_range.node_id + ROWS UNBOUNDED PRECEDING + ), 0 + ) AS ptr + FROM node_range + LEFT JOIN degrees ON node_range.node_id = degrees.src + ) + SELECT ptr FROM cumulative + ORDER BY node_id + """) + con.execute(""" + CREATE OR REPLACE TABLE indptr AS + SELECT 0::UBIGINT AS ptr + UNION ALL + SELECT ptr::UBIGINT FROM indptr + ORDER BY ptr + """) + + src_out = con.execute( + f"SELECT * FROM vertices ORDER BY {_q(src_pk)}" + ).to_arrow_table() + dst_out = con.execute( + f"SELECT * FROM vertices_dst ORDER BY {_q(dst_pk)}" + ).to_arrow_table() + indices = con.execute("SELECT * FROM indices").to_arrow_table() + indptr = con.execute("SELECT * FROM indptr").to_arrow_table() + return src_out, dst_out, indices, indptr + finally: + con.close() diff --git a/icebug_format/memory.py b/icebug_format/memory.py index 707f664..b5ca378 100644 --- a/icebug_format/memory.py +++ b/icebug_format/memory.py @@ -7,9 +7,10 @@ The conversion is performed by the pure-PyArrow pipeline in :mod:`icebug_format._convert_pyarrow` (the same engine behind the -``pyarrow`` Parquet-pair backend); DuckDB is not required. When a memory -limit is requested (or ``backend="datafusion"``), the DataFusion SQL engine -is used instead so sorts can spill to disk. +``pyarrow`` Parquet-pair backend); DuckDB is not required. The ``"duckdb"`` +and ``"datafusion"`` backends run the same plan inside a SQL engine so a +``memory_limit`` can spill sorts to disk; requesting a memory limit with the +default backend switches to datafusion. """ import warnings @@ -93,8 +94,9 @@ def from_arrow_tables( backend, which always sorts by primary key. Defaults to ``False``. backend: Conversion engine: ``"pyarrow"`` (default; - pure in-memory pipeline) or ``"datafusion"`` - (SQL engine, requires the + pure in-memory pipeline), ``"duckdb"`` + (requires the ``convert`` extra) or + ``"datafusion"`` (requires the ``convert-datafusion`` extra). When *memory_limit* is set and *backend* is left at the default, the backend automatically @@ -104,7 +106,8 @@ def from_arrow_tables( for sorting (size string like ``"4GB"``, percent of RAM like ``"50%"``, or a GB number); sorts spill to disk. Only honored - by the ``"datafusion"`` backend. + by the ``"duckdb"`` and ``"datafusion"`` + backends. Returns: IcebugMemGraph where *src* and *dest* are the node tables in CSR @@ -115,11 +118,13 @@ def from_arrow_tables( ValueError: If *rel_arrow_table* has fewer than 2 columns. ValueError: If ``add_reverse_edges=True`` and *to_node_arrow_table* is provided. - ValueError: If *backend* is not ``"pyarrow"`` or ``"datafusion"``. + ValueError: If *backend* is not ``"pyarrow"``, ``"duckdb"`` or + ``"datafusion"``. """ - if backend not in ("pyarrow", "datafusion"): + if backend not in ("pyarrow", "duckdb", "datafusion"): raise ValueError( - f"Unknown backend {backend!r}; expected 'pyarrow' or 'datafusion'" + f"Unknown backend {backend!r}; expected 'pyarrow', 'duckdb' or " + f"'datafusion'" ) if backend == "pyarrow" and memory_limit: warnings.warn( @@ -141,6 +146,18 @@ def from_arrow_tables( add_reverse_edges=add_reverse_edges, memory_limit=memory_limit, ) + elif backend == "duckdb": + from icebug_format._convert_duckdb import ( + build_csr_from_arrow_tables as _dd_build_csr, + ) + + src, dest, indices, indptr = _dd_build_csr( + from_node_arrow_table, + rel_arrow_table, + to_node_table=to_node_arrow_table, + add_reverse_edges=add_reverse_edges, + memory_limit=memory_limit, + ) else: src, dest, indices, indptr = build_csr_from_tables( from_node_arrow_table, diff --git a/tests/test_memory.py b/tests/test_memory.py index 0fceaf0..49ae228 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -333,7 +333,8 @@ def test_input_sorted_true_agrees_with_default_on_sorted_input(): # backend / memory_limit # --------------------------------------------------------------------------- -BACKENDS = ["pyarrow", "datafusion"] +BACKENDS = ["pyarrow", "duckdb", "datafusion"] +SQL_BACKENDS = ["duckdb", "datafusion"] def _assert_same_csr(g, ref): @@ -346,7 +347,7 @@ def _assert_same_csr(g, ref): @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("add_reverse_edges", [False, True]) def test_backends_agree(backend, add_reverse_edges): - """DataFusion and pyarrow produce identical CSR output.""" + """All backends produce identical CSR output.""" nodes = pa.table({"id": pa.array([30, 10, 20, 40], type=pa.int64())}) rels = _rels([20, 10, 10, 99], [30, 40, 20, 30]) # 99 not in nodes: dropped @@ -371,18 +372,20 @@ def test_memory_limit_switches_to_datafusion_with_warning(): _assert_same_csr(g, ref) -def test_backend_datafusion_with_memory_limit_matches_pyarrow(): +@pytest.mark.parametrize("backend", SQL_BACKENDS) +def test_sql_backends_with_memory_limit_match_pyarrow(backend): nodes = _nodes(30, 10, 20, 40) rels = _rels([20, 10], [30, 40]) g = IcebugMemGraph.from_arrow_tables( - nodes, rels, backend="datafusion", memory_limit="128MB" + nodes, rels, backend=backend, memory_limit="128MB" ) ref = IcebugMemGraph.from_arrow_tables(nodes, rels) _assert_same_csr(g, ref) -def test_backend_datafusion_edge_properties(): +@pytest.mark.parametrize("backend", SQL_BACKENDS) +def test_sql_backends_edge_properties(backend): nodes = _nodes(0, 1, 2) rels = pa.table( { @@ -392,30 +395,32 @@ def test_backend_datafusion_edge_properties(): } ) - g = IcebugMemGraph.from_arrow_tables(nodes, rels, backend="datafusion") + g = IcebugMemGraph.from_arrow_tables(nodes, rels, backend=backend) assert g.indices["weight"].to_pylist() == pytest.approx([0.5, 1.5]) assert g.indptr.schema.names == ["ptr"] -def test_backend_datafusion_empty_edges(): +@pytest.mark.parametrize("backend", SQL_BACKENDS) +def test_sql_backends_empty_edges(backend): nodes = _nodes(0, 1, 2) rels = _rels([], []) - g = IcebugMemGraph.from_arrow_tables(nodes, rels, backend="datafusion") + g = IcebugMemGraph.from_arrow_tables(nodes, rels, backend=backend) assert len(g.indices) == 0 assert g.indices.schema.field("target").type == pa.uint64() assert g.indptr["ptr"].to_pylist() == [0, 0, 0, 0] -def test_backend_datafusion_heterogeneous_nodes(): +@pytest.mark.parametrize("backend", SQL_BACKENDS) +def test_sql_backends_heterogeneous_nodes(backend): src = pa.table({"id": pa.array([20, 10], type=pa.int64())}) dst = pa.table({"id": pa.array([30, 40], type=pa.int64())}) rels = _rels([20], [40]) g = IcebugMemGraph.from_arrow_tables( - src, rels, to_node_arrow_table=dst, backend="datafusion" + src, rels, to_node_arrow_table=dst, backend=backend ) assert g.src["id"].to_pylist() == [10, 20] @@ -423,13 +428,14 @@ def test_backend_datafusion_heterogeneous_nodes(): assert g.indices["target"].to_pylist() == [1] -def test_backend_datafusion_ignores_input_sorted(): - """The datafusion backend always sorts by primary key.""" +@pytest.mark.parametrize("backend", SQL_BACKENDS) +def test_sql_backends_ignore_input_sorted(backend): + """The SQL backends always sort by primary key.""" nodes = pa.table({"id": pa.array([30, 10, 20], type=pa.int64())}) rels = _rels([20], [30]) g = IcebugMemGraph.from_arrow_tables( - nodes, rels, backend="datafusion", input_sorted=True + nodes, rels, backend=backend, input_sorted=True ) ref = IcebugMemGraph.from_arrow_tables(nodes, rels) @@ -441,17 +447,18 @@ def test_invalid_backend_raises(): nodes = _nodes(0, 1) rels = _rels([0], [1]) with pytest.raises(ValueError, match="Unknown backend"): - IcebugMemGraph.from_arrow_tables(nodes, rels, backend="duckdb") + IcebugMemGraph.from_arrow_tables(nodes, rels, backend="bogus") -def test_backend_datafusion_validations_shared(): - """DataFusion path enforces the same input validation as pyarrow.""" +@pytest.mark.parametrize("backend", SQL_BACKENDS) +def test_sql_backends_validations_shared(backend): + """SQL paths enforce the same input validation as pyarrow.""" nodes = _nodes(0, 1) with pytest.raises(ValueError, match="at least 2 columns"): IcebugMemGraph.from_arrow_tables( nodes, pa.table({"source": pa.array([0], type=pa.int64())}), - backend="datafusion", + backend=backend, ) with pytest.raises(ValueError, match="to_node_arrow_table must not be provided"): IcebugMemGraph.from_arrow_tables( @@ -459,5 +466,5 @@ def test_backend_datafusion_validations_shared(): _rels([0], [1]), to_node_arrow_table=nodes, add_reverse_edges=True, - backend="datafusion", + backend=backend, )