diff --git a/.gitignore b/.gitignore index 4b742a1..0f26cb4 100644 --- a/.gitignore +++ b/.gitignore @@ -232,3 +232,4 @@ __marimo__/ /wandb/ /mlruns/ /logs/ +/docs/ diff --git a/src/amlgraphx/__init__.py b/src/amlgraphx/__init__.py index 049b01a..afc4dda 100644 --- a/src/amlgraphx/__init__.py +++ b/src/amlgraphx/__init__.py @@ -1,2 +1,22 @@ +"""AMLGraphX public package exports.""" + +from .graphs import ( + AccountGraph, + TransactionGraph, + build_account_graph, + build_transaction_graph, +) + + def hello() -> str: + """Return the package greeting kept for backwards compatibility.""" return "Hello from amlgraphx!" + + +__all__ = [ + "AccountGraph", + "TransactionGraph", + "build_account_graph", + "build_transaction_graph", + "hello", +] diff --git a/src/amlgraphx/datasets/base.py b/src/amlgraphx/datasets/base.py index 9bf573f..d9efde7 100644 --- a/src/amlgraphx/datasets/base.py +++ b/src/amlgraphx/datasets/base.py @@ -135,6 +135,7 @@ def clean_lazy_frame( _SOURCE_ALIASES = ( "source", "sender", + "sender account", "from", "from account", "source id", @@ -146,8 +147,12 @@ def clean_lazy_frame( _TARGET_ALIASES = ( "target", "receiver", + "receiver account", "to", "to account", + "account.1", + "account 1", + "account duplicated 0", "target id", "dst", "namedest", diff --git a/src/amlgraphx/datasets/download.py b/src/amlgraphx/datasets/download.py index 2de22ac..cd8d664 100644 --- a/src/amlgraphx/datasets/download.py +++ b/src/amlgraphx/datasets/download.py @@ -51,7 +51,12 @@ def find_tabular_file(root: Path, preferred_terms: Sequence[str] = ()) -> Path: candidates = sorted( path for path in root.rglob("*") - if path.is_file() and path.suffix.lower() in {".csv", ".parquet"} + if ( + path.is_file() + and path.suffix.lower() in {".csv", ".parquet"} + and "__MACOSX" not in path.parts + and not path.name.startswith("._") + ) ) for term in preferred_terms: matches = [path for path in candidates if term.lower() in path.name.lower()] diff --git a/src/amlgraphx/datasets/paysim.py b/src/amlgraphx/datasets/paysim.py index 2b6cef5..60e7f0d 100644 --- a/src/amlgraphx/datasets/paysim.py +++ b/src/amlgraphx/datasets/paysim.py @@ -1,5 +1,6 @@ """PaySim dataset adapter.""" +from datetime import UTC, datetime from pathlib import Path import polars as pl @@ -86,12 +87,23 @@ def transaction_path(self) -> Path: return find_tabular_file(self.download(), ("log", "paysim")) def transactions(self) -> pl.LazyFrame: - """Return lazily scanned and cleaned PaySim transactions.""" - return clean_lazy_frame( + """Return lazily scanned PaySim transactions with a logical timestamp. + + PaySim's ``step`` is a simulated hour. The added ``timestamp`` anchors + that hour sequence at the Unix epoch solely to support temporal graph + construction; the original ``step`` column remains unchanged. + """ + frame = clean_lazy_frame( pl.scan_csv(self.transaction_path()), source_column="nameOrig", target_column="nameDest", ) + return frame.with_columns( + ( + pl.lit(datetime(1970, 1, 1, tzinfo=UTC)) + + pl.duration(hours=pl.col("step").cast(pl.Int64)) + ).alias("timestamp") + ) def _dataset_root(self) -> Path: if self.local_dir is not None: diff --git a/src/amlgraphx/graphs.py b/src/amlgraphx/graphs.py new file mode 100644 index 0000000..c7d88d0 --- /dev/null +++ b/src/amlgraphx/graphs.py @@ -0,0 +1,574 @@ +"""Graph views built from AMLGraphX transaction tables.""" + +from bisect import bisect_right +from collections import Counter +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta + +import polars as pl + +type TransactionTable = pl.DataFrame | pl.LazyFrame + +_SOURCE_ALIASES = ( + "source", + "sender", + "sender account", + "from", + "from account", + "account", + "source id", + "src", + "nameorig", + "origin account", + "origin", +) +_TARGET_ALIASES = ( + "target", + "receiver", + "receiver account", + "to", + "to account", + "account.1", + "account 1", + "account duplicated 0", + "target id", + "dst", + "namedest", + "destination account", + "destination", +) +_TIMESTAMP_ALIASES = ("timestamp", "datetime", "time", "date") +_TRANSACTION_ID_ALIASES = ( + "transaction_id", + "transaction id", + "tx id", + "id", +) +_ACCOUNT_ID_ALIASES = ( + "node_id", + "account number", + "account", + "account id", + "account_id", +) + +_ROW_INDEX = "__amlgraphx_row_index" +_TIMESTAMP_NS = "__amlgraphx_timestamp_ns" + + +@dataclass(frozen=True, slots=True) +class AccountGraph: + """Represent accounts as nodes and transactions as directed edges. + + Args: + nodes: Account table containing at least ``node_id``. + edges: Transaction table containing ``source`` and ``target``. + """ + + nodes: pl.DataFrame + edges: pl.DataFrame + + @property + def num_nodes(self) -> int: + """Return the number of account nodes.""" + return self.nodes.height + + @property + def num_edges(self) -> int: + """Return the number of transaction edges.""" + return self.edges.height + + @classmethod + def from_transactions( + cls, + transactions: TransactionTable, + *, + account_metadata: TransactionTable | None = None, + source_column: str | None = None, + target_column: str | None = None, + transaction_id_column: str | None = None, + account_id_column: str | None = None, + ) -> "AccountGraph": + """Build an account graph from a transaction table. + + Args: + transactions: Transaction rows as a Polars DataFrame or LazyFrame. + account_metadata: Optional account table to join onto ``nodes``. + source_column: Optional sender column override. + target_column: Optional receiver column override. + transaction_id_column: Optional transaction ID column override. + account_id_column: Optional account metadata ID column override. + + Returns: + An account graph. Repeated transactions remain separate edges. + + Raises: + ValueError: If source or target columns are missing. + TypeError: If the input is not a Polars table. + """ + frame, source, target = _prepare_transactions( + transactions, + source_column=source_column, + target_column=target_column, + transaction_id_column=transaction_id_column, + ) + + edge_columns = [ + column + for column in frame.columns + if column != _ROW_INDEX + ] + edges = frame.select(edge_columns) + nodes = pl.concat( + [ + frame.select(pl.col(source).alias("node_id")), + frame.select(pl.col(target).alias("node_id")), + ] + ).unique(subset=["node_id"], maintain_order=True).sort("node_id") + + if account_metadata is not None: + nodes = _join_account_metadata( + nodes, + account_metadata, + account_id_column=account_id_column, + ) + + return cls(nodes=nodes, edges=edges) + + +@dataclass(frozen=True, slots=True) +class TransactionGraph: + """Represent transactions as nodes linked by temporal money flow. + + An edge means that a later transaction starts from the earlier + transaction's receiver within the configured time window. It represents a + feasible temporal continuation, not proof of exact money tracing. + + Args: + nodes: Transaction table with canonical graph columns. + edges: Temporal succession edges between transaction IDs. + """ + + nodes: pl.DataFrame + edges: pl.DataFrame + + @property + def num_nodes(self) -> int: + """Return the number of transaction nodes.""" + return self.nodes.height + + @property + def num_edges(self) -> int: + """Return the number of temporal succession edges.""" + return self.edges.height + + @classmethod + def from_transactions( + cls, + transactions: TransactionTable, + *, + delta: timedelta, + source_column: str | None = None, + target_column: str | None = None, + timestamp_column: str | None = None, + transaction_id_column: str | None = None, + ) -> "TransactionGraph": + """Build a temporal transaction graph. + + Args: + transactions: Transaction rows as a Polars DataFrame or LazyFrame. + delta: Inclusive maximum time between two linked transactions. + source_column: Optional sender column override. + target_column: Optional receiver column override. + timestamp_column: Optional timestamp column override. + transaction_id_column: Optional transaction ID column override. + + Returns: + A graph whose nodes preserve the input transaction attributes and + whose edges contain transaction IDs, the via account, and the time + difference. + + Raises: + ValueError: If required columns are missing, timestamps are not + parseable, or ``delta`` is negative. + TypeError: If the input is not a Polars table. + """ + if not isinstance(delta, timedelta): + raise TypeError("delta must be a datetime.timedelta") + if delta < timedelta(0): + raise ValueError("delta must be non-negative") + + frame, source, target = _prepare_transactions( + transactions, + source_column=source_column, + target_column=target_column, + timestamp_column=timestamp_column, + transaction_id_column=transaction_id_column, + parse_timestamp=True, + ) + if "timestamp" not in frame.columns: + raise ValueError( + "Missing required timestamp column; expected one of: " + + ", ".join(_TIMESTAMP_ALIASES) + ) + if frame.height and frame["timestamp"].null_count() == frame.height: + raise ValueError( + "Timestamp column exists but contains no parseable timestamps" + ) + frame = frame.filter(pl.col("timestamp").is_not_null()) + + node_columns = [ + column + for column in frame.columns + if column not in {_ROW_INDEX, _TIMESTAMP_NS} + ] + nodes = frame.select(node_columns) + ordered = frame.with_columns( + _timestamp_nanoseconds(frame).alias(_TIMESTAMP_NS) + ).sort([_TIMESTAMP_NS, _ROW_INDEX]) + rows = list(ordered.iter_rows(named=True)) + delta_ns = _timedelta_to_nanoseconds(delta) + + outgoing: dict[str, list[tuple[int, int]]] = {} + outgoing_times: dict[str, list[int]] = {} + for position, row in enumerate(rows): + account = row[source] + timestamp = row[_TIMESTAMP_NS] + outgoing.setdefault(account, []).append((timestamp, position)) + outgoing_times.setdefault(account, []).append(timestamp) + + edge_records: list[dict[str, object]] = [] + for row in rows: + candidates = outgoing.get(row[target], []) + candidate_times = outgoing_times.get(row[target], []) + timestamp = row[_TIMESTAMP_NS] + start = bisect_right(candidate_times, timestamp) + end = bisect_right( + candidate_times, + timestamp + delta_ns, + ) + for _, successor_position in candidates[start:end]: + successor = rows[successor_position] + edge_records.append( + { + "source_transaction_id": row["transaction_id"], + "target_transaction_id": successor["transaction_id"], + "via_account": row[target], + "time_delta_ns": successor[_TIMESTAMP_NS] - timestamp, + } + ) + + edges = _transaction_edge_frame(edge_records) + return cls(nodes=nodes, edges=edges) + + +def build_account_graph( + transactions: TransactionTable, + *, + account_metadata: TransactionTable | None = None, + source_column: str | None = None, + target_column: str | None = None, + transaction_id_column: str | None = None, + account_id_column: str | None = None, +) -> AccountGraph: + """Build an account graph from transaction rows. + + Args: + transactions: Transaction rows as a Polars DataFrame or LazyFrame. + account_metadata: Optional account metadata table. + source_column: Optional sender column override. + target_column: Optional receiver column override. + transaction_id_column: Optional transaction ID column override. + account_id_column: Optional account metadata ID column override. + + Returns: + An account graph with one directed edge per transaction. + """ + return AccountGraph.from_transactions( + transactions, + account_metadata=account_metadata, + source_column=source_column, + target_column=target_column, + transaction_id_column=transaction_id_column, + account_id_column=account_id_column, + ) + + +def build_transaction_graph( + transactions: TransactionTable, + *, + delta: timedelta, + source_column: str | None = None, + target_column: str | None = None, + timestamp_column: str | None = None, + transaction_id_column: str | None = None, +) -> TransactionGraph: + """Build a temporal transaction graph from transaction rows. + + Args: + transactions: Transaction rows as a Polars DataFrame or LazyFrame. + delta: Inclusive maximum time between linked transactions. + source_column: Optional sender column override. + target_column: Optional receiver column override. + timestamp_column: Optional timestamp column override. + transaction_id_column: Optional transaction ID column override. + + Returns: + A transaction graph with all valid temporal succession edges. + """ + return TransactionGraph.from_transactions( + transactions, + delta=delta, + source_column=source_column, + target_column=target_column, + timestamp_column=timestamp_column, + transaction_id_column=transaction_id_column, + ) + + +def _prepare_transactions( + transactions: TransactionTable, + *, + source_column: str | None, + target_column: str | None, + timestamp_column: str | None = None, + transaction_id_column: str | None = None, + parse_timestamp: bool = False, +) -> tuple[pl.DataFrame, str, str]: + frame = _collect_frame(transactions) + source = _resolve_required_column( + frame.columns, + source_column, + _SOURCE_ALIASES, + "source account", + ) + target = _resolve_required_column( + frame.columns, + target_column, + _TARGET_ALIASES, + "target account", + ) + if _ROW_INDEX in frame.columns: + raise ValueError(f"Input column {_ROW_INDEX!r} is reserved") + + frame = frame.with_row_index(_ROW_INDEX) + id_column = _resolve_optional_column( + frame.columns, + transaction_id_column, + _TRANSACTION_ID_ALIASES, + ) + frame = frame.with_columns( + _make_transaction_ids(frame, id_column).alias("transaction_id") + ) + + frame = frame.with_columns( + pl.col(source) + .cast(pl.String) + .str.strip_chars() + .alias("source"), + pl.col(target) + .cast(pl.String) + .str.strip_chars() + .alias("target"), + ).filter( + pl.col("source").is_not_null() + & (pl.col("source") != "") + & pl.col("target").is_not_null() + & (pl.col("target") != "") + ) + + if parse_timestamp: + timestamp = _resolve_required_column( + frame.columns, + timestamp_column, + _TIMESTAMP_ALIASES, + "timestamp", + ) + frame = frame.with_columns( + _timestamp_expression(frame, timestamp).alias("timestamp") + ) + + return frame, "source", "target" + + +def _collect_frame(transactions: TransactionTable) -> pl.DataFrame: + if isinstance(transactions, pl.LazyFrame): + return transactions.collect() + if isinstance(transactions, pl.DataFrame): + return transactions.clone() + raise TypeError("transactions must be a polars.DataFrame or polars.LazyFrame") + + +def _make_transaction_ids( + frame: pl.DataFrame, + id_column: str | None, +) -> pl.Series: + if id_column is None: + return pl.Series( + "transaction_id", + [f"tx_{row_index}" for row_index in frame[_ROW_INDEX].to_list()], + dtype=pl.String, + ) + + values = ( + frame.get_column(id_column) + .cast(pl.String) + .str.strip_chars() + .to_list() + ) + counts = Counter(value for value in values if value not in (None, "")) + used = { + value for value, count in counts.items() if count == 1 and value is not None + } + transaction_ids: list[str] = [] + for row_index, value in zip(frame[_ROW_INDEX].to_list(), values, strict=True): + if value not in (None, "") and counts[value] == 1: + transaction_ids.append(value) + continue + + candidate = f"tx_{row_index}" + suffix = 1 + while candidate in used: + candidate = f"tx_{row_index}_{suffix}" + suffix += 1 + used.add(candidate) + transaction_ids.append(candidate) + + return pl.Series("transaction_id", transaction_ids, dtype=pl.String) + + +def _join_account_metadata( + nodes: pl.DataFrame, + account_metadata: TransactionTable, + *, + account_id_column: str | None, +) -> pl.DataFrame: + metadata = _collect_frame(account_metadata) + id_column = _resolve_required_column( + metadata.columns, + account_id_column, + _ACCOUNT_ID_ALIASES, + "account metadata ID", + ) + metadata = ( + metadata.with_columns( + pl.col(id_column).cast(pl.String).str.strip_chars().alias("node_id") + ) + .unique(subset=["node_id"], maintain_order=True) + ) + if id_column != "node_id": + metadata = metadata.drop(id_column) + return nodes.join(metadata, on="node_id", how="left") + + +def _resolve_required_column( + columns: Sequence[str], + requested: str | None, + aliases: Sequence[str], + logical_name: str, +) -> str: + column = _resolve_optional_column(columns, requested, aliases) + if column is None: + expected = ", ".join(aliases) + raise ValueError( + f"Missing required {logical_name} column; expected one of: {expected}" + ) + return column + + +def _resolve_optional_column( + columns: Sequence[str], + requested: str | None, + aliases: Sequence[str], +) -> str | None: + normalized = {_normalize_column(column): column for column in columns} + if requested is not None: + if requested in columns: + return requested + return normalized.get(_normalize_column(requested)) + for alias in aliases: + column = normalized.get(_normalize_column(alias)) + if column is not None: + return column + return None + + +def _normalize_column(column: str) -> str: + return " ".join(column.lower().replace("_", " ").replace(".", " ").split()) + + +def _timestamp_expression(frame: pl.DataFrame, column: str) -> pl.Expr: + dtype = frame.schema[column] + expression = pl.col(column) + if dtype == pl.Datetime: + return expression + if dtype == pl.Date: + return expression.cast(pl.Datetime) + if dtype == pl.String: + parsed_timestamp = _parse_datetime_strings(expression) + date_column = _resolve_optional_column( + frame.columns, + None, + ("date",), + ) + if date_column is not None and date_column != column: + date_only = pl.col(date_column).cast(pl.String).str.replace( + r"[T ].*$", "" + ) + return pl.coalesce( + parsed_timestamp, + _parse_datetime_strings( + pl.concat_str( + [date_only, pl.col(column)], + separator=" ", + ) + ), + ) + return parsed_timestamp + return expression.cast(pl.Datetime, strict=False) + + +def _timestamp_nanoseconds(frame: pl.DataFrame) -> pl.Expr: + time_unit = frame.schema["timestamp"].time_unit + scale = {"ms": 1_000_000, "us": 1_000, "ns": 1}[time_unit] + return pl.col("timestamp").cast(pl.Int64) * scale + + +def _timedelta_to_nanoseconds(value: timedelta) -> int: + return ( + value.days * 86_400_000_000_000_000 + + value.seconds * 1_000_000_000 + + value.microseconds * 1_000 + ) + + +def _parse_datetime_strings(expression: pl.Expr) -> pl.Expr: + normalized = expression.str.replace("T", " ") + return pl.coalesce( + normalized.str.to_datetime( + format="%Y-%m-%d %H:%M:%S%.f", + strict=False, + ), + normalized.str.to_datetime( + format="%Y-%m-%d %H:%M:%S", + strict=False, + ), + normalized.str.to_datetime(format="%Y-%m-%d %H:%M", strict=False), + normalized.str.to_datetime(format="%Y-%m-%d", strict=False), + ) + + +def _transaction_edge_frame(records: list[dict[str, object]]) -> pl.DataFrame: + if records: + return pl.DataFrame(records).with_columns( + pl.duration(nanoseconds=pl.col("time_delta_ns")).alias("time_delta") + ).drop("time_delta_ns") + return pl.DataFrame( + { + "source_transaction_id": pl.Series([], dtype=pl.String), + "target_transaction_id": pl.Series([], dtype=pl.String), + "via_account": pl.Series([], dtype=pl.String), + "time_delta": pl.Series([], dtype=pl.Duration("ns")), + } + ) diff --git a/tests/test_dataset_loading.py b/tests/test_dataset_loading.py index 0cebffe..21dcac8 100644 --- a/tests/test_dataset_loading.py +++ b/tests/test_dataset_loading.py @@ -1,5 +1,6 @@ """Tests for cache, extraction, cleaning, and the unified loader.""" +from datetime import UTC, datetime, timedelta from pathlib import Path from zipfile import ZipFile @@ -15,6 +16,7 @@ extract_zip, load_dataset, ) +from amlgraphx.graphs import build_account_graph, build_transaction_graph def test_extract_zip_rejects_invalid_archive(tmp_path: Path) -> None: @@ -84,3 +86,34 @@ def fake_download(**_: object) -> str: paysim_transactions.collect_schema().names() ) assert isinstance(samld.transactions(), pl.LazyFrame) + + +def test_paysim_transactions_support_both_graph_views( + monkeypatch: MonkeyPatch, tmp_path: Path +) -> None: + """PaySim skips macOS resource files and exposes a logical timestamp.""" + archive = tmp_path / "paysim.zip" + with ZipFile(archive, "w") as zip_file: + zip_file.writestr( + "paysim.csv", + "step,type,amount,nameOrig,nameDest,isFraud\n" + "1,PAYMENT,1,A,B,0\n" + "2,PAYMENT,2,B,C,1\n", + ) + zip_file.writestr("__MACOSX/._paysim.csv", "not,a,csv\n") + + monkeypatch.setattr( + "amlgraphx.datasets.download.hf_hub_download", + lambda **_: str(archive), + ) + transactions = PaySim(cache_dir=tmp_path / "cache").transactions() + + assert transactions.collect()["timestamp"].to_list() == [ + datetime(1970, 1, 1, 1, tzinfo=UTC), + datetime(1970, 1, 1, 2, tzinfo=UTC), + ] + assert build_account_graph(transactions).num_edges == 2 + assert build_transaction_graph( + transactions, + delta=timedelta(hours=1), + ).num_edges == 1 diff --git a/tests/test_graphs.py b/tests/test_graphs.py new file mode 100644 index 0000000..2957926 --- /dev/null +++ b/tests/test_graphs.py @@ -0,0 +1,312 @@ +"""Tests for AMLGraphX account and transaction graph views.""" + +from datetime import UTC, datetime, timedelta + +import polars as pl +import pytest + +from amlgraphx.datasets import clean_lazy_frame +from amlgraphx.graphs import ( + AccountGraph, + TransactionGraph, + build_account_graph, + build_transaction_graph, +) + + +def _transaction_frame( + rows: list[tuple[str, str, str, str, float, int]], +) -> pl.DataFrame: + return pl.DataFrame( + rows, + schema=[ + "transaction_id", + "source", + "target", + "timestamp", + "amount", + "label", + ], + orient="row", + ).with_columns(pl.col("timestamp").str.to_datetime()) + + +def test_account_graph_preserves_nodes_edges_and_attributes() -> None: + """Account graphs preserve accounts, repeated edges, and attributes.""" + frame = _transaction_frame( + [ + ("t1", "A", "B", "2025-01-01 09:00", 10.0, 0), + ("t2", "A", "B", "2025-01-01 09:05", 20.0, 1), + ("t3", "B", "C", "2025-01-01 09:10", 30.0, 0), + ] + ) + + graph = build_account_graph(frame) + + assert isinstance(graph, AccountGraph) + assert graph.num_nodes == 3 + assert graph.num_edges == 3 + assert graph.nodes["node_id"].to_list() == ["A", "B", "C"] + assert graph.edges.select("transaction_id").to_series().to_list() == [ + "t1", + "t2", + "t3", + ] + assert graph.edges["amount"].to_list() == [10.0, 20.0, 30.0] + assert graph.edges["label"].to_list() == [0, 1, 0] + + +def test_account_graph_joins_account_metadata() -> None: + """Account metadata is preserved on matching account nodes.""" + transactions = pl.DataFrame({"from": ["A"], "to": ["B"]}) + accounts = pl.DataFrame( + { + "Account Number": ["A", "B"], + "Bank ID": [1, 2], + "Entity ID": ["e1", "e2"], + } + ) + + graph = build_account_graph(transactions, account_metadata=accounts) + + assert graph.nodes["Bank ID"].to_list() == [1, 2] + assert graph.nodes["Entity ID"].to_list() == ["e1", "e2"] + + +def test_account_graph_strips_account_metadata_ids() -> None: + """Whitespace around metadata account IDs does not prevent a join.""" + transactions = pl.DataFrame({"from": ["A"], "to": ["B"]}) + accounts = pl.DataFrame({"Account Number": [" A ", " B "], "Bank ID": [1, 2]}) + + graph = build_account_graph(transactions, account_metadata=accounts) + + assert graph.nodes["Bank ID"].to_list() == [1, 2] + + +def test_account_graph_accepts_lazy_frames() -> None: + """Account graph construction accepts a lazy Polars frame.""" + frame = pl.LazyFrame({"Account": ["A"], "Account.1": ["B"]}) + + graph = AccountGraph.from_transactions(frame) + + assert graph.edges.select(["source", "target"]).to_dicts() == [ + {"source": "A", "target": "B"} + ] + + +def test_account_graph_preserves_string_timestamps() -> None: + """Account graph edges retain timestamp strings without parsing them.""" + frame = pl.DataFrame( + { + "source": ["A"], + "target": ["B"], + "timestamp": ["2022/09/01 00:20"], + } + ) + + graph = build_account_graph(frame) + + assert graph.edges["timestamp"].to_list() == ["2022/09/01 00:20"] + + +def test_account_graph_requires_endpoints() -> None: + """Missing endpoint columns produce a clear validation error.""" + with pytest.raises(ValueError, match="source account"): + build_account_graph(pl.DataFrame({"amount": [1.0]})) + + +def test_graphs_generate_deterministic_transaction_ids() -> None: + """Missing transaction IDs are generated from stable input row order.""" + frame = pl.DataFrame( + { + "source": ["A", "B"], + "target": ["B", "C"], + "timestamp": [ + datetime(2025, 1, 1, 9, tzinfo=UTC), + datetime(2025, 1, 1, 10, tzinfo=UTC), + ], + } + ) + + first = build_account_graph(frame) + second = build_account_graph(frame) + + assert first.edges["transaction_id"].to_list() == ["tx_0", "tx_1"] + assert first.edges["transaction_id"].to_list() == second.edges[ + "transaction_id" + ].to_list() + + +def test_transaction_graph_repairs_null_and_duplicate_ids() -> None: + """Invalid source IDs become deterministic unique transaction IDs.""" + frame = _transaction_frame( + [ + ("duplicate", "A", "B", "2025-01-01 09:00", 1.0, 0), + ("duplicate", "B", "C", "2025-01-01 09:10", 2.0, 0), + (None, "C", "D", "2025-01-01 09:20", 3.0, 0), + ] + ) + + graph = build_transaction_graph(frame, delta=timedelta(hours=1)) + node_ids = graph.nodes["transaction_id"].to_list() + + assert node_ids == ["tx_0", "tx_1", "tx_2"] + assert len(node_ids) == len(set(node_ids)) + assert all(value is not None for value in node_ids) + assert graph.edges.select( + ["source_transaction_id", "target_transaction_id"] + ).to_dicts() == [ + {"source_transaction_id": "tx_0", "target_transaction_id": "tx_1"}, + {"source_transaction_id": "tx_1", "target_transaction_id": "tx_2"}, + ] + + +def test_transaction_graph_combines_parsed_date_with_time() -> None: + """A parsed datetime date column combines correctly with a time string.""" + frame = pl.LazyFrame( + { + "source": ["A", "B"], + "target": ["B", "C"], + "Time": ["09:00:00", "09:30:00"], + "Date": ["2025-01-01", "2025-01-01"], + } + ) + cleaned = clean_lazy_frame( + frame, + source_column="source", + target_column="target", + timestamp_columns=("Date",), + ) + + graph = build_transaction_graph(cleaned, delta=timedelta(hours=1)) + + assert graph.num_edges == 1 + assert graph.edges["time_delta"].to_list() == [timedelta(minutes=30)] + + +def test_transaction_graph_keeps_full_timestamp_when_date_also_exists() -> None: + """Full timestamps are not combined again with a separate date column.""" + frame = pl.LazyFrame( + { + "source": ["A", "B"], + "target": ["B", "C"], + "Time": ["2025-01-01 09:00:00", "2025-01-01 09:30:00"], + "Date": ["2025-01-01", "2025-01-01"], + } + ) + + graph = build_transaction_graph(frame, delta=timedelta(hours=1)) + + assert graph.num_edges == 1 + assert graph.edges["time_delta"].to_list() == [timedelta(minutes=30)] + + +def test_transaction_graph_accepts_date_only_strings() -> None: + """Date-only timestamp columns retain their midnight timestamp.""" + graph = build_transaction_graph( + pl.DataFrame( + { + "source": ["A"], + "target": ["B"], + "Date": ["2025-01-01"], + } + ), + delta=timedelta(hours=1), + ) + + assert graph.num_nodes == 1 + + +def test_transaction_graph_preserves_nanosecond_ordering() -> None: + """Sub-microsecond timestamps remain ordered and measurable.""" + frame = pl.DataFrame( + { + "source": ["A", "B"], + "target": ["B", "C"], + "timestamp": [0, 500], + }, + schema={"source": pl.String, "target": pl.String, "timestamp": pl.Datetime("ns")}, + ) + + graph = build_transaction_graph(frame, delta=timedelta(microseconds=1)) + + assert graph.num_edges == 1 + assert graph.edges["time_delta"].cast(pl.Int64).to_list() == [500] + + +def test_transaction_graph_rejects_exceeded_delta() -> None: + """Successors beyond the inclusive time window are excluded.""" + frame = _transaction_frame( + [ + ("t1", "A", "B", "2025-01-01 09:00", 1.0, 0), + ("t2", "B", "C", "2025-01-01 10:01", 2.0, 0), + ] + ) + + graph = build_transaction_graph(frame, delta=timedelta(hours=1)) + + assert graph.num_edges == 0 + + +def test_transaction_graph_creates_directional_temporal_edges() -> None: + """Transaction edges follow account direction and the inclusive window.""" + frame = _transaction_frame( + [ + ("t1", "A", "B", "2025-01-01 09:00", 10.0, 0), + ("t2", "B", "C", "2025-01-01 09:20", 20.0, 0), + ("t3", "B", "D", "2025-01-01 10:00", 30.0, 1), + ("t4", "X", "B", "2025-01-01 09:30", 40.0, 0), + ] + ) + + graph = build_transaction_graph(frame, delta=timedelta(hours=1)) + + assert isinstance(graph, TransactionGraph) + assert graph.num_nodes == 4 + assert graph.edges.select( + ["source_transaction_id", "target_transaction_id"] + ).to_dicts() == [ + {"source_transaction_id": "t1", "target_transaction_id": "t2"}, + {"source_transaction_id": "t1", "target_transaction_id": "t3"}, + {"source_transaction_id": "t4", "target_transaction_id": "t3"}, + ] + assert graph.edges["via_account"].to_list() == ["B", "B", "B"] + assert graph.edges["time_delta"].to_list() == [ + timedelta(minutes=20), + timedelta(hours=1), + timedelta(minutes=30), + ] + + +def test_transaction_graph_rejects_same_or_backward_time() -> None: + """Only strictly later transactions may receive temporal edges.""" + frame = _transaction_frame( + [ + ("t1", "A", "B", "2025-01-01 10:00", 1.0, 0), + ("t2", "B", "C", "2025-01-01 10:00", 2.0, 0), + ("t3", "B", "D", "2025-01-01 09:00", 3.0, 0), + ] + ) + + graph = build_transaction_graph(frame, delta=timedelta(hours=2)) + + assert graph.num_edges == 0 + + +def test_transaction_graph_requires_timestamp_and_delta() -> None: + """Required timestamps and non-negative timedelta values are validated.""" + frame = pl.DataFrame({"source": ["A"], "target": ["B"]}) + + with pytest.raises(ValueError, match="timestamp"): + build_transaction_graph(frame, delta=timedelta(hours=1)) + with pytest.raises(ValueError, match="non-negative"): + build_transaction_graph( + pl.DataFrame( + { + "source": ["A"], + "target": ["B"], + "timestamp": [datetime(2025, 1, 1, tzinfo=UTC)], + } + ), + delta=timedelta(seconds=-1), + )