Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 31 additions & 4 deletions graphbench/engines/issundb_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,29 @@
from ..schema import Schema


def _verify_import_counts(payload: dict, expected: dict[str, tuple[str, int]]) -> None:
"""Check IMPORT DATABASE's per-file report against the prepared row counts.

Since IssunDB 0.1.0a16 the import returns one `(target, kind, count)` row
per COPY statement; a mismatch (for example an edge file misclassified as
nodes) previously surfaced only as silently empty query results. Older
versions return `{"imported": true}` and cannot be verified, so the check
is skipped for them.
"""
if payload.get("columns") != ["target", "kind", "count"]:
return
imported = {
rec["values"][0]: (rec["values"][1], rec["values"][2])
for rec in payload["records"]
}
for target, want in expected.items():
got = imported.get(target)
if got != want:
raise RuntimeError(
f"IMPORT DATABASE ingested {got} for '{target}', expected {want}"
)


class IssunDBEngine(Engine):
name = "issundb"
kind = "embedded"
Expand Down Expand Up @@ -68,11 +91,13 @@ def build(self, data_dir: Path) -> BuildResult:

# 3. Process and write node Parquet files with _id column
n_node_rows = 0
expected_counts: dict[str, tuple[str, int]] = {}
for label in self.schema.nodes:
parquet_path = data_dir / "nodes" / f"{label.name}.parquet"
dst_parquet = import_dir / f"{label.name}.parquet"
df = pl.read_parquet(parquet_path)
n_node_rows += df.height
expected_counts[label.name] = ("nodes", df.height)
df = df.with_columns(
(
pl.col(self.schema.id_column).cast(pl.Int64) + offsets[label.name]
Expand All @@ -89,14 +114,15 @@ def build(self, data_dir: Path) -> BuildResult:
jsonl_path = import_dir / f"{rel.name}.jsonl"
df = pl.read_parquet(parquet_path)
n_edge_rows += df.height
expected_counts[rel.name] = ("relationships", df.height)
df = df.with_columns(
[
(
pl.col(self.schema.src_column).cast(pl.Int64) + offsets[rel.src]
).alias("from"),
).alias("_from"),
(
pl.col(self.schema.dst_column).cast(pl.Int64) + offsets[rel.dst]
).alias("to"),
).alias("_to"),
]
)
# Drop original src/dst columns to avoid importing them as edge properties
Expand All @@ -118,10 +144,11 @@ def build(self, data_dir: Path) -> BuildResult:

(import_dir / "copy.cypher").write_text("\n".join(copy_lines))

# 6. Execute IMPORT DATABASE
# 6. Execute IMPORT DATABASE and verify the per-file report
query_start = time.perf_counter()
self._db.query(f"IMPORT DATABASE '{import_dir}'")
payload = json.loads(self._db.query(f"IMPORT DATABASE '{import_dir}'"))
query_time = time.perf_counter() - query_start
_verify_import_counts(payload, expected_counts)

# Clean up import temp files
shutil.rmtree(import_dir)
Expand Down
57 changes: 57 additions & 0 deletions tests/test_issundb_import.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
"""Unit tests for the IssunDB import-report verification.

The check guards against the silent misclassification found in 0.1.0a15,
where an edge file with legacy endpoint keys imported as nodes and the
benchmark ran against an edgeless graph.
"""

import pytest

from graphbench.engines.issundb_engine import _verify_import_counts

EXPECTED = {
"Person": ("nodes", 3),
"FOLLOWS": ("relationships", 2),
}


def payload(records):
return {
"columns": ["target", "kind", "count"],
"records": [{"values": list(v)} for v in records],
}


def test_matching_report_passes():
_verify_import_counts(
payload([("Person", "nodes", 3), ("FOLLOWS", "relationships", 2)]),
EXPECTED,
)


def test_misclassified_edge_file_raises():
with pytest.raises(RuntimeError, match="FOLLOWS"):
_verify_import_counts(
payload([("Person", "nodes", 3), ("FOLLOWS", "nodes", 2)]),
EXPECTED,
)


def test_short_count_raises():
with pytest.raises(RuntimeError, match="FOLLOWS"):
_verify_import_counts(
payload([("Person", "nodes", 3), ("FOLLOWS", "relationships", 0)]),
EXPECTED,
)


def test_missing_target_raises():
with pytest.raises(RuntimeError, match="FOLLOWS"):
_verify_import_counts(payload([("Person", "nodes", 3)]), EXPECTED)


def test_legacy_result_shape_is_skipped():
_verify_import_counts(
{"columns": ["imported"], "records": [{"values": [True]}]},
EXPECTED,
)
Loading
Loading