tkati-node-dedup 0.3.0__py3-none-any.whl

tkati_node_dedup/__main__.py

@@ -0,0 +1,4 @@
+from .main import main
+
+if __name__ == "__main__":
+    main()

tkati_node_dedup/main.py

@@ -0,0 +1,109 @@
+import pyarrow as pa
+from loguru import logger
+from tkati_core import Consumer, Producer, build_consumer, build_producer
+
+from tkati_node_dedup.settings import AppSettings
+from tkati_node_dedup.store import BucketedDedupStore
+
+
+def _dedupe_batch(
+    batch: pa.Table, field: str, store: BucketedDedupStore
+) -> tuple[pa.Table, list[bytes]]:
+    """Filter out rows whose dedup key was already seen (in-batch or in the store).
+
+    Returns (filtered_batch, keys_to_mark_seen). Null values in `field` always
+    pass through and are never added to the store — we can't dedup on nothing.
+
+    Encoding and the store lookup are both batched (one pass over the column,
+    one RocksDB round trip per open bucket) rather than done per row.
+    """
+    keys = store.encode_keys(batch.column(field))
+    keep_mask, new_keys = store.filter_duplicates(keys)
+    filtered = batch.filter(keep_mask)
+    return filtered, new_keys
+
+
+def run_one_iteration(
+    consumer: Consumer,
+    producer: Producer,
+    store: BucketedDedupStore,
+    settings: AppSettings,
+) -> None:
+    # Runs first, every iteration (even if no batch arrives), and can never
+    # raise. Buckets must be fresh *before* the dedupe check below runs —
+    # doing this only after commit would leave a just-expired bucket open and
+    # checked against for one extra iteration, and an idle node (no messages,
+    # read_arrow returns None below) would never clean up at all.
+    try:
+        store.cleanup_expired()
+    except Exception:
+        logger.exception("dedup store cleanup failed; will retry next iteration")
+
+    batch = consumer.read_arrow(
+        num_messages=settings.input.consumer.batch_size,
+        timeout=settings.input.consumer.batch_timeout_sec,
+    )
+    if batch is None:
+        return
+
+    field = settings.dedup.field
+    if field not in batch.column_names:
+        logger.warning(
+            f"Dedup field '{field}' missing from batch schema; passing batch through unfiltered"
+        )
+        filtered, new_keys = batch, []
+    else:
+        filtered, new_keys = _dedupe_batch(batch, field, store)
+
+    dropped = len(batch) - len(filtered)
+
+    if len(filtered) > 0:
+        producer.produce_arrow(filtered)
+        # Block until actually delivered before marking anything "seen" or
+        # committing. Required here even though tkati-node-el's loop skips it:
+        # KafkaProducer.produce_arrow() only enqueues (non-blocking), and
+        # marking a key seen before it's durably delivered would risk losing
+        # the event permanently on a crash. ClickhouseProducer.flush() is a
+        # no-op since its inserts are already synchronous.
+        producer.flush()
+
+    # Only after a confirmed-successful produce: mark these keys seen.
+    store.add_many(new_keys)
+
+    # Only after mark-seen: commit. If we crash before this line, the batch is
+    # re-read at restart; those keys are already in the store, so re-processing
+    # it drops what was already produced — a harmless duplicate at worst, never
+    # a lost event.
+    consumer.commit()
+
+    logger.info(
+        f"Batch of {len(batch)} rows: produced {len(filtered)}, "
+        f"deduped {dropped} ({len(new_keys)} newly marked seen)"
+    )
+
+
+def main() -> None:
+    settings = AppSettings()
+
+    consumer = build_consumer(settings.input)
+
+    dlq_producer: Producer | None = None
+    if settings.dlq is not None:
+        dlq_producer = build_producer(settings.dlq)
+
+    producer = build_producer(settings.output, dlq_producer=dlq_producer)
+
+    store = BucketedDedupStore(
+        root_dir=settings.dedup.store_dir,
+        window_hours=settings.dedup.window_hours,
+        bucket_hours=settings.dedup.bucket_hours,
+    )
+
+    try:
+        while True:
+            run_one_iteration(consumer, producer, store, settings)
+    finally:
+        consumer.close()
+        if dlq_producer is not None:
+            dlq_producer.close()
+        store.close()

tkati_node_dedup/settings.py

@@ -0,0 +1,23 @@
+from pydantic import BaseModel, field_validator
+from tkati_core.settings import InputSettings, OutputSettings, TomlBaseSettings
+
+
+class DedupSettings(BaseModel):
+    field: str
+    window_hours: int = 3
+    bucket_hours: int = 1
+    store_dir: str = "./dedup_store"
+
+    @field_validator("window_hours", "bucket_hours")
+    @classmethod
+    def _positive(cls, v: int) -> int:
+        if v <= 0:
+            raise ValueError("must be a positive number of hours")
+        return v
+
+
+class AppSettings(TomlBaseSettings):
+    input: InputSettings
+    output: OutputSettings
+    dlq: OutputSettings | None = None
+    dedup: DedupSettings

tkati_node_dedup/store.py

@@ -0,0 +1,224 @@
+"""Bucketed, embedded RocksDB store for windowed key deduplication.
+
+One RocksDB database directory per wall-clock-aligned hour bucket. A key is
+considered "seen" if it exists in any bucket currently inside the window.
+Existence-only store: values are always empty bytes, only key presence matters.
+
+Lookups and writes are batched against RocksDB (one call per open bucket for
+reads, one WriteBatch for writes) rather than one call per key — at realistic
+batch sizes and millions of runs a day, per-row FFI calls into RocksDB are the
+dominant cost, and rocksdict supports genuine batching for both directions.
+
+Failure policy throughout: any per-bucket open/lookup/write failure is caught,
+logged, and treated as "not seen" / skipped — never raised. This favors
+forwarding a possible duplicate over silently dropping a real event.
+"""
+
+import math
+import shutil
+import time
+from collections.abc import Iterable
+from pathlib import Path
+from typing import cast
+
+import pyarrow as pa
+import pyarrow.compute as pc
+from loguru import logger
+from rocksdict import Options, Rdict, WriteBatch
+
+_BUCKET_PREFIX = "bucket-"
+
+
+def _now() -> float:
+    """Indirection over time.time() so tests can freeze this store's clock
+    without patching the global time module (which would also freeze
+    unrelated code, e.g. the Kafka consumer's poll-timeout bookkeeping)."""
+    return time.time()
+
+
+def _bucket_index(ts: float, bucket_seconds: int) -> int:
+    return int(ts // bucket_seconds)
+
+
+class BucketedDedupStore:
+    def __init__(self, root_dir: str, window_hours: int, bucket_hours: int = 1) -> None:
+        self.root = Path(root_dir)
+        self.root.mkdir(parents=True, exist_ok=True)
+        self.bucket_seconds = bucket_hours * 3600
+        self.num_buckets = math.ceil(window_hours / bucket_hours)
+        self._opts = Options(raw_mode=True)
+        self._dbs: dict[int, Rdict] = {}
+        self._recover_existing_buckets()
+
+    def _bucket_path(self, bucket: int) -> Path:
+        return self.root / f"{_BUCKET_PREFIX}{bucket:012d}"
+
+    def _min_live_bucket(self) -> int:
+        current = _bucket_index(_now(), self.bucket_seconds)
+        return current - self.num_buckets + 1
+
+    def _recover_existing_buckets(self) -> None:
+        """Reopen on-disk buckets still inside the window; destroy stale ones.
+
+        Makes dedup state survive a node restart.
+        """
+        min_live = self._min_live_bucket()
+        for entry in sorted(self.root.glob(f"{_BUCKET_PREFIX}*")):
+            if not entry.is_dir():
+                continue
+            try:
+                bucket = int(entry.name.removeprefix(_BUCKET_PREFIX))
+            except ValueError:
+                logger.warning(f"Ignoring unrecognized entry in dedup store dir: {entry}")
+                continue
+            if bucket < min_live:
+                logger.info(f"Startup: removing stale dedup bucket {bucket} ({entry})")
+                self._destroy_path(entry)
+                continue
+            try:
+                self._dbs[bucket] = Rdict(str(entry), options=self._opts)
+                logger.info(f"Startup: reopened dedup bucket {bucket} from {entry}")
+            except Exception:
+                logger.exception(
+                    f"Failed to reopen dedup bucket {bucket} at {entry}; "
+                    "starting empty for this bucket (may pass through some duplicates)"
+                )
+
+    def _ensure_current_open(self) -> Rdict | None:
+        bucket = _bucket_index(_now(), self.bucket_seconds)
+        if bucket in self._dbs:
+            return self._dbs[bucket]
+        try:
+            db = Rdict(str(self._bucket_path(bucket)), options=self._opts)
+        except Exception:
+            logger.exception(f"Failed to open current dedup bucket {bucket}")
+            return None
+        self._dbs[bucket] = db
+        return db
+
+    def encode_keys(self, values: pa.Array | pa.ChunkedArray) -> list[bytes | None]:
+        """Vectorized byte-key encoding for a column of scalar dedup-field values.
+
+        Cast to string via pyarrow compute (fast, vectorized — no per-row
+        Python type dispatch), then UTF-8 encode each non-null value. Nulls
+        stay None: they're never queried or stored, always passed through.
+
+        Float or timestamp dedup fields are discouraged: their string
+        representation isn't guaranteed stable across producers.
+        """
+        strings = pc.cast(values, pa.string())
+        return [None if s is None else s.encode("utf-8") for s in strings.to_pylist()]
+
+    def filter_duplicates(
+        self, keys: list[bytes | None]
+    ) -> tuple[pa.BooleanArray, list[bytes]]:
+        """
+        Given per-row encoded keys (None = no key, always kept), returns
+        (keep_mask, keys_to_mark_seen). keep_mask[i] corresponds to keys[i]
+        and is directly usable with pyarrow.Table.filter().
+
+        In-batch duplicates (two rows with the same key, neither yet in the
+        store) are resolved locally; the remaining unique candidates are
+        checked against the store in a single batched round trip per open
+        bucket, not one lookup per key.
+        """
+        keep_mask = [False] * len(keys)
+        seen_in_batch: set[bytes] = set()
+        to_check: list[bytes] = []
+        to_check_idx: list[int] = []
+
+        for i, key in enumerate(keys):
+            if key is None:
+                keep_mask[i] = True
+                continue
+            if key in seen_in_batch:
+                continue
+            seen_in_batch.add(key)
+            to_check.append(key)
+            to_check_idx.append(i)
+
+        already_seen = self._batch_contains(to_check)
+
+        new_keys: list[bytes] = []
+        for idx, key, seen in zip(to_check_idx, to_check, already_seen, strict=True):
+            if not seen:
+                keep_mask[idx] = True
+                new_keys.append(key)
+
+        return pa.array(keep_mask, type=pa.bool_()), new_keys
+
+    def _batch_contains(self, keys: list[bytes]) -> list[bool]:
+        if not keys:
+            return []
+        found = [False] * len(keys)
+        for bucket, db in list(self._dbs.items()):
+            try:
+                # rocksdict's stub declares `List[...]` invariantly, so
+                # list[bytes] isn't accepted as-is, and doesn't distinguish
+                # the list-in/list-out overload from the scalar one for the
+                # return type either — get() genuinely returns a list here
+                # since `keys` is a list.
+                keys_arg = cast("list[str | int | float | bytes]", keys)
+                results = cast("list[bytes | None]", db.get(keys_arg))
+            except Exception:
+                logger.exception(f"Batch lookup failed against dedup bucket {bucket}; skipping it")
+                continue
+            for i, r in enumerate(results):
+                if r is not None:
+                    found[i] = True
+        return found
+
+    def add_many(self, keys: Iterable[bytes]) -> None:
+        keys = list(keys)
+        if not keys:
+            return
+        db = self._ensure_current_open()
+        if db is None:
+            return
+        try:
+            wb = WriteBatch(raw_mode=True)
+            for key in keys:
+                wb.put(key, b"")
+            db.write(wb)
+        except Exception:
+            logger.exception("Failed to batch-write keys into current dedup bucket")
+
+    def contains(self, key: bytes) -> bool:
+        return self._batch_contains([key])[0]
+
+    def add(self, key: bytes) -> None:
+        self.add_many([key])
+
+    def _destroy_path(self, path: Path) -> None:
+        try:
+            Rdict.destroy(str(path), self._opts)
+        except Exception:
+            logger.exception(f"Rdict.destroy failed for {path}; falling back to rmtree")
+            shutil.rmtree(path, ignore_errors=True)
+
+    def cleanup_expired(self) -> None:
+        """Close+delete on-disk buckets that fell out of the window.
+
+        Called at the start of every iteration (before the dedupe check runs
+        against possibly-stale buckets), and cheap to call every time: only
+        does real I/O once per hour rollover, since _min_live_bucket() only
+        changes then.
+        """
+        min_live = self._min_live_bucket()
+        for bucket in list(self._dbs.keys()):
+            if bucket < min_live:
+                db = self._dbs.pop(bucket)
+                try:
+                    db.close()
+                except Exception:
+                    logger.exception(f"Error closing expired dedup bucket {bucket}")
+                self._destroy_path(self._bucket_path(bucket))
+                logger.info(f"Expired dedup bucket {bucket} removed")
+
+    def close(self) -> None:
+        for bucket, db in self._dbs.items():
+            try:
+                db.close()
+            except Exception:
+                logger.exception(f"Error closing dedup bucket {bucket} on shutdown")
+        self._dbs.clear()

tkati_node_dedup-0.3.0.dist-info/METADATA

@@ -0,0 +1,106 @@
+Metadata-Version: 2.4
+Name: tkati-node-dedup
+Version: 0.3.0
+Summary: Kafka-to-Kafka streaming node that deduplicates events by a configurable field within a rolling processing-time window
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+Requires-Dist: tkati-core==0.3.0
+Requires-Dist: loguru>=0.7.0
+Requires-Dist: pydantic-settings>=2.11.0
+Requires-Dist: pyarrow>=21.0.0
+Requires-Dist: rocksdict>=0.3.0
+
+# tkati-node-dedup — streaming deduplication node
+
+Reads batches from a Kafka input topic, drops events that are duplicates of an
+event seen on the same `field` within a rolling processing-time window, and
+writes the deduplicated batch to a configurable output. Duplicate state is
+tracked in an embedded, on-disk RocksDB store local to this process — no
+external dedup service is required.
+
+## Configuration
+
+Settings are loaded from a TOML file. Set the `SETTINGS_FILE` environment
+variable to point to it (defaults to `settings.toml`).
+
+```toml
+[input]
+type = "kafka"
+
+[input.connection]
+broker = "redpanda:29092"
+
+[input.topic]
+name = "raw_event"
+
+[input.topic.schema]
+uid  = "string"
+time = "timestamp[ms]"
+# … other columns
+
+[input.consumer]
+group_id          = "node-dedup-group"
+batch_size        = 1000
+batch_timeout_sec = 10
+auto_offset_reset = "latest"
+
+[output]
+type = "kafka"
+
+[output.connection]
+broker = "redpanda:29092"
+
+[output.topic]
+name = "raw_event_deduped"
+
+[dedup]
+field        = "uid"     # column in input.topic.schema to dedup by
+window_hours = 3         # rolling dedup window
+bucket_hours = 1         # on-disk bucket granularity (effective window is
+                          # window_hours .. window_hours + bucket_hours)
+store_dir    = "/var/lib/tkati-node-dedup/store"
+```
+
+Output and DLQ follow the same `OutputSettings` shape as `tkati-node-el`
+(`"kafka"` or `"clickhouse"`) — see that package's README for the full
+connection/table config shape.
+
+## Delivery & dedup guarantees
+
+**Delivery: at-least-once.** Offsets are committed only after (1) the
+filtered batch is produced and confirmed delivered (`produce_arrow` followed
+by a blocking `flush`), and (2) the surviving keys are durably written to the
+current RocksDB bucket. If the process crashes between steps, the same input
+batch is re-read at restart; because the keys from a completed produce are
+already marked seen, re-processing that batch is a no-op (or reproduces only
+the genuinely-new subset) rather than losing data.
+
+**On any internal dedup-store failure — a bucket won't open, a lookup errors,
+a disk I/O error — the node treats the event as NOT a duplicate and forwards
+it.** This node will occasionally forward a duplicate it should have caught,
+but will never silently drop a real event because of dedup-store trouble.
+
+**The window is approximate, not exact.** Because state is bucketed in
+`bucket_hours` increments (default 1h) rather than a true sliding window, the
+effective dedup window is between `window_hours` and
+`window_hours + bucket_hours`. Stale buckets are deleted from disk
+automatically once they fall outside the window — checked once per iteration,
+so state never grows unbounded.
+
+**Bucketing is by processing time**, not any timestamp field in the event
+payload — an event's bucket is when this node handles it, not when it
+happened upstream.
+
+## IMPORTANT: dedup state is local to this process
+
+The RocksDB store lives on local disk at `store_dir` and is **not shared**
+between instances. Running multiple concurrent instances of this node against
+the same input topic (e.g. multiple consumers in the same consumer group, or
+multiple replicas) will **not** dedup correctly across instances unless the
+input is partitioned such that all events sharing a `field` value are always
+routed to the *same* instance (e.g. Kafka partitioning keyed on `field`, one
+node instance per partition or partition subset it exclusively owns). Running
+this node with more parallelism than that will let duplicates leak through
+across instance boundaries. This is a direct consequence of choosing an
+embedded local-file store instead of a shared external one — evaluate this
+tradeoff before scaling this node horizontally.

tkati_node_dedup-0.3.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+tkati_node_dedup/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tkati_node_dedup/__main__.py,sha256=Vdhw8YA1K3wPMlbJQYL5WqvRzAKVeZ16mZQFO9VRmCo,62
+tkati_node_dedup/main.py,sha256=86DHLsGrhyYhWOp8CkqiCBYCqDrSFSISGh7YxgloyOQ,3914
+tkati_node_dedup/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tkati_node_dedup/settings.py,sha256=FyIUGM_pTeyAgXnOw33yhfbhbOYqQRrEfm06mi7rfqc,633
+tkati_node_dedup/store.py,sha256=e5eD792Auk3hyB1KiduUcCft2uxTmOY0gyXi1JTzcK4,8904
+tkati_node_dedup-0.3.0.dist-info/METADATA,sha256=S1echBtd7htveJBcrall8z5l6LpJIY5UdFpgUQwpgRQ,4167
+tkati_node_dedup-0.3.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+tkati_node_dedup-0.3.0.dist-info/entry_points.txt,sha256=ivsESwayqwUo1rkYBFKUZ_FJD4lMj1-WFbd5S9ZpLto,64
+tkati_node_dedup-0.3.0.dist-info/top_level.txt,sha256=dpxwzDwQ4E175J1C2ajLZZ0ya2SSPs6QT5YJNXi-b2I,17
+tkati_node_dedup-0.3.0.dist-info/RECORD,,

tkati_node_dedup-0.3.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

tkati_node_dedup-0.3.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+tkati-node-dedup = tkati_node_dedup.main:main

tkati_node_dedup-0.3.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+tkati_node_dedup