feat: support the duckdb backend in from_arrow_tables

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()

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,

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,44 +395,47 @@ 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]
     assert g.dest["id"].to_pylist() == [30, 40]
     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,23 +447,24 @@ 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(
             nodes,
             _rels([0], [1]),
             to_node_arrow_table=nodes,
             add_reverse_edges=True,
-            backend="datafusion",
+            backend=backend,
         )