tkati-node-dedup 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- tkati_node_dedup-0.3.0/PKG-INFO +106 -0
- tkati_node_dedup-0.3.0/README.md +94 -0
- tkati_node_dedup-0.3.0/pyproject.toml +34 -0
- tkati_node_dedup-0.3.0/setup.cfg +4 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup/__init__.py +0 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup/__main__.py +4 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup/main.py +109 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup/py.typed +0 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup/settings.py +23 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup/store.py +224 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup.egg-info/PKG-INFO +106 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup.egg-info/SOURCES.txt +16 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup.egg-info/dependency_links.txt +1 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup.egg-info/entry_points.txt +2 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup.egg-info/requires.txt +5 -0
- tkati_node_dedup-0.3.0/src/tkati_node_dedup.egg-info/top_level.txt +1 -0
- tkati_node_dedup-0.3.0/tests/test_node_dedup.py +228 -0
- tkati_node_dedup-0.3.0/tests/test_store.py +114 -0
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tkati-node-dedup
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: Kafka-to-Kafka streaming node that deduplicates events by a configurable field within a rolling processing-time window
|
|
5
|
+
Requires-Python: >=3.13
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
Requires-Dist: tkati-core==0.3.0
|
|
8
|
+
Requires-Dist: loguru>=0.7.0
|
|
9
|
+
Requires-Dist: pydantic-settings>=2.11.0
|
|
10
|
+
Requires-Dist: pyarrow>=21.0.0
|
|
11
|
+
Requires-Dist: rocksdict>=0.3.0
|
|
12
|
+
|
|
13
|
+
# tkati-node-dedup — streaming deduplication node
|
|
14
|
+
|
|
15
|
+
Reads batches from a Kafka input topic, drops events that are duplicates of an
|
|
16
|
+
event seen on the same `field` within a rolling processing-time window, and
|
|
17
|
+
writes the deduplicated batch to a configurable output. Duplicate state is
|
|
18
|
+
tracked in an embedded, on-disk RocksDB store local to this process — no
|
|
19
|
+
external dedup service is required.
|
|
20
|
+
|
|
21
|
+
## Configuration
|
|
22
|
+
|
|
23
|
+
Settings are loaded from a TOML file. Set the `SETTINGS_FILE` environment
|
|
24
|
+
variable to point to it (defaults to `settings.toml`).
|
|
25
|
+
|
|
26
|
+
```toml
|
|
27
|
+
[input]
|
|
28
|
+
type = "kafka"
|
|
29
|
+
|
|
30
|
+
[input.connection]
|
|
31
|
+
broker = "redpanda:29092"
|
|
32
|
+
|
|
33
|
+
[input.topic]
|
|
34
|
+
name = "raw_event"
|
|
35
|
+
|
|
36
|
+
[input.topic.schema]
|
|
37
|
+
uid = "string"
|
|
38
|
+
time = "timestamp[ms]"
|
|
39
|
+
# … other columns
|
|
40
|
+
|
|
41
|
+
[input.consumer]
|
|
42
|
+
group_id = "node-dedup-group"
|
|
43
|
+
batch_size = 1000
|
|
44
|
+
batch_timeout_sec = 10
|
|
45
|
+
auto_offset_reset = "latest"
|
|
46
|
+
|
|
47
|
+
[output]
|
|
48
|
+
type = "kafka"
|
|
49
|
+
|
|
50
|
+
[output.connection]
|
|
51
|
+
broker = "redpanda:29092"
|
|
52
|
+
|
|
53
|
+
[output.topic]
|
|
54
|
+
name = "raw_event_deduped"
|
|
55
|
+
|
|
56
|
+
[dedup]
|
|
57
|
+
field = "uid" # column in input.topic.schema to dedup by
|
|
58
|
+
window_hours = 3 # rolling dedup window
|
|
59
|
+
bucket_hours = 1 # on-disk bucket granularity (effective window is
|
|
60
|
+
# window_hours .. window_hours + bucket_hours)
|
|
61
|
+
store_dir = "/var/lib/tkati-node-dedup/store"
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
Output and DLQ follow the same `OutputSettings` shape as `tkati-node-el`
|
|
65
|
+
(`"kafka"` or `"clickhouse"`) — see that package's README for the full
|
|
66
|
+
connection/table config shape.
|
|
67
|
+
|
|
68
|
+
## Delivery & dedup guarantees
|
|
69
|
+
|
|
70
|
+
**Delivery: at-least-once.** Offsets are committed only after (1) the
|
|
71
|
+
filtered batch is produced and confirmed delivered (`produce_arrow` followed
|
|
72
|
+
by a blocking `flush`), and (2) the surviving keys are durably written to the
|
|
73
|
+
current RocksDB bucket. If the process crashes between steps, the same input
|
|
74
|
+
batch is re-read at restart; because the keys from a completed produce are
|
|
75
|
+
already marked seen, re-processing that batch is a no-op (or reproduces only
|
|
76
|
+
the genuinely-new subset) rather than losing data.
|
|
77
|
+
|
|
78
|
+
**On any internal dedup-store failure — a bucket won't open, a lookup errors,
|
|
79
|
+
a disk I/O error — the node treats the event as NOT a duplicate and forwards
|
|
80
|
+
it.** This node will occasionally forward a duplicate it should have caught,
|
|
81
|
+
but will never silently drop a real event because of dedup-store trouble.
|
|
82
|
+
|
|
83
|
+
**The window is approximate, not exact.** Because state is bucketed in
|
|
84
|
+
`bucket_hours` increments (default 1h) rather than a true sliding window, the
|
|
85
|
+
effective dedup window is between `window_hours` and
|
|
86
|
+
`window_hours + bucket_hours`. Stale buckets are deleted from disk
|
|
87
|
+
automatically once they fall outside the window — checked once per iteration,
|
|
88
|
+
so state never grows unbounded.
|
|
89
|
+
|
|
90
|
+
**Bucketing is by processing time**, not any timestamp field in the event
|
|
91
|
+
payload — an event's bucket is when this node handles it, not when it
|
|
92
|
+
happened upstream.
|
|
93
|
+
|
|
94
|
+
## IMPORTANT: dedup state is local to this process
|
|
95
|
+
|
|
96
|
+
The RocksDB store lives on local disk at `store_dir` and is **not shared**
|
|
97
|
+
between instances. Running multiple concurrent instances of this node against
|
|
98
|
+
the same input topic (e.g. multiple consumers in the same consumer group, or
|
|
99
|
+
multiple replicas) will **not** dedup correctly across instances unless the
|
|
100
|
+
input is partitioned such that all events sharing a `field` value are always
|
|
101
|
+
routed to the *same* instance (e.g. Kafka partitioning keyed on `field`, one
|
|
102
|
+
node instance per partition or partition subset it exclusively owns). Running
|
|
103
|
+
this node with more parallelism than that will let duplicates leak through
|
|
104
|
+
across instance boundaries. This is a direct consequence of choosing an
|
|
105
|
+
embedded local-file store instead of a shared external one — evaluate this
|
|
106
|
+
tradeoff before scaling this node horizontally.
|
|
@@ -0,0 +1,94 @@
|
|
|
1
|
+
# tkati-node-dedup — streaming deduplication node
|
|
2
|
+
|
|
3
|
+
Reads batches from a Kafka input topic, drops events that are duplicates of an
|
|
4
|
+
event seen on the same `field` within a rolling processing-time window, and
|
|
5
|
+
writes the deduplicated batch to a configurable output. Duplicate state is
|
|
6
|
+
tracked in an embedded, on-disk RocksDB store local to this process — no
|
|
7
|
+
external dedup service is required.
|
|
8
|
+
|
|
9
|
+
## Configuration
|
|
10
|
+
|
|
11
|
+
Settings are loaded from a TOML file. Set the `SETTINGS_FILE` environment
|
|
12
|
+
variable to point to it (defaults to `settings.toml`).
|
|
13
|
+
|
|
14
|
+
```toml
|
|
15
|
+
[input]
|
|
16
|
+
type = "kafka"
|
|
17
|
+
|
|
18
|
+
[input.connection]
|
|
19
|
+
broker = "redpanda:29092"
|
|
20
|
+
|
|
21
|
+
[input.topic]
|
|
22
|
+
name = "raw_event"
|
|
23
|
+
|
|
24
|
+
[input.topic.schema]
|
|
25
|
+
uid = "string"
|
|
26
|
+
time = "timestamp[ms]"
|
|
27
|
+
# … other columns
|
|
28
|
+
|
|
29
|
+
[input.consumer]
|
|
30
|
+
group_id = "node-dedup-group"
|
|
31
|
+
batch_size = 1000
|
|
32
|
+
batch_timeout_sec = 10
|
|
33
|
+
auto_offset_reset = "latest"
|
|
34
|
+
|
|
35
|
+
[output]
|
|
36
|
+
type = "kafka"
|
|
37
|
+
|
|
38
|
+
[output.connection]
|
|
39
|
+
broker = "redpanda:29092"
|
|
40
|
+
|
|
41
|
+
[output.topic]
|
|
42
|
+
name = "raw_event_deduped"
|
|
43
|
+
|
|
44
|
+
[dedup]
|
|
45
|
+
field = "uid" # column in input.topic.schema to dedup by
|
|
46
|
+
window_hours = 3 # rolling dedup window
|
|
47
|
+
bucket_hours = 1 # on-disk bucket granularity (effective window is
|
|
48
|
+
# window_hours .. window_hours + bucket_hours)
|
|
49
|
+
store_dir = "/var/lib/tkati-node-dedup/store"
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
Output and DLQ follow the same `OutputSettings` shape as `tkati-node-el`
|
|
53
|
+
(`"kafka"` or `"clickhouse"`) — see that package's README for the full
|
|
54
|
+
connection/table config shape.
|
|
55
|
+
|
|
56
|
+
## Delivery & dedup guarantees
|
|
57
|
+
|
|
58
|
+
**Delivery: at-least-once.** Offsets are committed only after (1) the
|
|
59
|
+
filtered batch is produced and confirmed delivered (`produce_arrow` followed
|
|
60
|
+
by a blocking `flush`), and (2) the surviving keys are durably written to the
|
|
61
|
+
current RocksDB bucket. If the process crashes between steps, the same input
|
|
62
|
+
batch is re-read at restart; because the keys from a completed produce are
|
|
63
|
+
already marked seen, re-processing that batch is a no-op (or reproduces only
|
|
64
|
+
the genuinely-new subset) rather than losing data.
|
|
65
|
+
|
|
66
|
+
**On any internal dedup-store failure — a bucket won't open, a lookup errors,
|
|
67
|
+
a disk I/O error — the node treats the event as NOT a duplicate and forwards
|
|
68
|
+
it.** This node will occasionally forward a duplicate it should have caught,
|
|
69
|
+
but will never silently drop a real event because of dedup-store trouble.
|
|
70
|
+
|
|
71
|
+
**The window is approximate, not exact.** Because state is bucketed in
|
|
72
|
+
`bucket_hours` increments (default 1h) rather than a true sliding window, the
|
|
73
|
+
effective dedup window is between `window_hours` and
|
|
74
|
+
`window_hours + bucket_hours`. Stale buckets are deleted from disk
|
|
75
|
+
automatically once they fall outside the window — checked once per iteration,
|
|
76
|
+
so state never grows unbounded.
|
|
77
|
+
|
|
78
|
+
**Bucketing is by processing time**, not any timestamp field in the event
|
|
79
|
+
payload — an event's bucket is when this node handles it, not when it
|
|
80
|
+
happened upstream.
|
|
81
|
+
|
|
82
|
+
## IMPORTANT: dedup state is local to this process
|
|
83
|
+
|
|
84
|
+
The RocksDB store lives on local disk at `store_dir` and is **not shared**
|
|
85
|
+
between instances. Running multiple concurrent instances of this node against
|
|
86
|
+
the same input topic (e.g. multiple consumers in the same consumer group, or
|
|
87
|
+
multiple replicas) will **not** dedup correctly across instances unless the
|
|
88
|
+
input is partitioned such that all events sharing a `field` value are always
|
|
89
|
+
routed to the *same* instance (e.g. Kafka partitioning keyed on `field`, one
|
|
90
|
+
node instance per partition or partition subset it exclusively owns). Running
|
|
91
|
+
this node with more parallelism than that will let duplicates leak through
|
|
92
|
+
across instance boundaries. This is a direct consequence of choosing an
|
|
93
|
+
embedded local-file store instead of a shared external one — evaluate this
|
|
94
|
+
tradeoff before scaling this node horizontally.
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "tkati-node-dedup"
|
|
3
|
+
version = "0.3.0"
|
|
4
|
+
description = "Kafka-to-Kafka streaming node that deduplicates events by a configurable field within a rolling processing-time window"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.13"
|
|
7
|
+
dependencies = [
|
|
8
|
+
"tkati-core==0.3.0",
|
|
9
|
+
"loguru>=0.7.0",
|
|
10
|
+
"pydantic-settings>=2.11.0",
|
|
11
|
+
"pyarrow>=21.0.0",
|
|
12
|
+
"rocksdict>=0.3.0",
|
|
13
|
+
]
|
|
14
|
+
|
|
15
|
+
[project.scripts]
|
|
16
|
+
tkati-node-dedup = "tkati_node_dedup.main:main"
|
|
17
|
+
|
|
18
|
+
[dependency-groups]
|
|
19
|
+
dev = ["pytest>=9.0.1", "confluent-kafka>=2.11.0"]
|
|
20
|
+
|
|
21
|
+
[build-system]
|
|
22
|
+
requires = ["setuptools", "wheel"]
|
|
23
|
+
build-backend = "setuptools.build_meta"
|
|
24
|
+
|
|
25
|
+
[tool.setuptools.packages.find]
|
|
26
|
+
where = ["src"]
|
|
27
|
+
include = ["tkati_node_dedup*"]
|
|
28
|
+
|
|
29
|
+
[tool.uv]
|
|
30
|
+
package = true
|
|
31
|
+
|
|
32
|
+
[tool.uv-workspace-codegen]
|
|
33
|
+
generate = true
|
|
34
|
+
template_type = ["test", "publish"]
|
|
File without changes
|
|
@@ -0,0 +1,109 @@
|
|
|
1
|
+
import pyarrow as pa
|
|
2
|
+
from loguru import logger
|
|
3
|
+
from tkati_core import Consumer, Producer, build_consumer, build_producer
|
|
4
|
+
|
|
5
|
+
from tkati_node_dedup.settings import AppSettings
|
|
6
|
+
from tkati_node_dedup.store import BucketedDedupStore
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def _dedupe_batch(
|
|
10
|
+
batch: pa.Table, field: str, store: BucketedDedupStore
|
|
11
|
+
) -> tuple[pa.Table, list[bytes]]:
|
|
12
|
+
"""Filter out rows whose dedup key was already seen (in-batch or in the store).
|
|
13
|
+
|
|
14
|
+
Returns (filtered_batch, keys_to_mark_seen). Null values in `field` always
|
|
15
|
+
pass through and are never added to the store — we can't dedup on nothing.
|
|
16
|
+
|
|
17
|
+
Encoding and the store lookup are both batched (one pass over the column,
|
|
18
|
+
one RocksDB round trip per open bucket) rather than done per row.
|
|
19
|
+
"""
|
|
20
|
+
keys = store.encode_keys(batch.column(field))
|
|
21
|
+
keep_mask, new_keys = store.filter_duplicates(keys)
|
|
22
|
+
filtered = batch.filter(keep_mask)
|
|
23
|
+
return filtered, new_keys
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def run_one_iteration(
|
|
27
|
+
consumer: Consumer,
|
|
28
|
+
producer: Producer,
|
|
29
|
+
store: BucketedDedupStore,
|
|
30
|
+
settings: AppSettings,
|
|
31
|
+
) -> None:
|
|
32
|
+
# Runs first, every iteration (even if no batch arrives), and can never
|
|
33
|
+
# raise. Buckets must be fresh *before* the dedupe check below runs —
|
|
34
|
+
# doing this only after commit would leave a just-expired bucket open and
|
|
35
|
+
# checked against for one extra iteration, and an idle node (no messages,
|
|
36
|
+
# read_arrow returns None below) would never clean up at all.
|
|
37
|
+
try:
|
|
38
|
+
store.cleanup_expired()
|
|
39
|
+
except Exception:
|
|
40
|
+
logger.exception("dedup store cleanup failed; will retry next iteration")
|
|
41
|
+
|
|
42
|
+
batch = consumer.read_arrow(
|
|
43
|
+
num_messages=settings.input.consumer.batch_size,
|
|
44
|
+
timeout=settings.input.consumer.batch_timeout_sec,
|
|
45
|
+
)
|
|
46
|
+
if batch is None:
|
|
47
|
+
return
|
|
48
|
+
|
|
49
|
+
field = settings.dedup.field
|
|
50
|
+
if field not in batch.column_names:
|
|
51
|
+
logger.warning(
|
|
52
|
+
f"Dedup field '{field}' missing from batch schema; passing batch through unfiltered"
|
|
53
|
+
)
|
|
54
|
+
filtered, new_keys = batch, []
|
|
55
|
+
else:
|
|
56
|
+
filtered, new_keys = _dedupe_batch(batch, field, store)
|
|
57
|
+
|
|
58
|
+
dropped = len(batch) - len(filtered)
|
|
59
|
+
|
|
60
|
+
if len(filtered) > 0:
|
|
61
|
+
producer.produce_arrow(filtered)
|
|
62
|
+
# Block until actually delivered before marking anything "seen" or
|
|
63
|
+
# committing. Required here even though tkati-node-el's loop skips it:
|
|
64
|
+
# KafkaProducer.produce_arrow() only enqueues (non-blocking), and
|
|
65
|
+
# marking a key seen before it's durably delivered would risk losing
|
|
66
|
+
# the event permanently on a crash. ClickhouseProducer.flush() is a
|
|
67
|
+
# no-op since its inserts are already synchronous.
|
|
68
|
+
producer.flush()
|
|
69
|
+
|
|
70
|
+
# Only after a confirmed-successful produce: mark these keys seen.
|
|
71
|
+
store.add_many(new_keys)
|
|
72
|
+
|
|
73
|
+
# Only after mark-seen: commit. If we crash before this line, the batch is
|
|
74
|
+
# re-read at restart; those keys are already in the store, so re-processing
|
|
75
|
+
# it drops what was already produced — a harmless duplicate at worst, never
|
|
76
|
+
# a lost event.
|
|
77
|
+
consumer.commit()
|
|
78
|
+
|
|
79
|
+
logger.info(
|
|
80
|
+
f"Batch of {len(batch)} rows: produced {len(filtered)}, "
|
|
81
|
+
f"deduped {dropped} ({len(new_keys)} newly marked seen)"
|
|
82
|
+
)
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def main() -> None:
|
|
86
|
+
settings = AppSettings()
|
|
87
|
+
|
|
88
|
+
consumer = build_consumer(settings.input)
|
|
89
|
+
|
|
90
|
+
dlq_producer: Producer | None = None
|
|
91
|
+
if settings.dlq is not None:
|
|
92
|
+
dlq_producer = build_producer(settings.dlq)
|
|
93
|
+
|
|
94
|
+
producer = build_producer(settings.output, dlq_producer=dlq_producer)
|
|
95
|
+
|
|
96
|
+
store = BucketedDedupStore(
|
|
97
|
+
root_dir=settings.dedup.store_dir,
|
|
98
|
+
window_hours=settings.dedup.window_hours,
|
|
99
|
+
bucket_hours=settings.dedup.bucket_hours,
|
|
100
|
+
)
|
|
101
|
+
|
|
102
|
+
try:
|
|
103
|
+
while True:
|
|
104
|
+
run_one_iteration(consumer, producer, store, settings)
|
|
105
|
+
finally:
|
|
106
|
+
consumer.close()
|
|
107
|
+
if dlq_producer is not None:
|
|
108
|
+
dlq_producer.close()
|
|
109
|
+
store.close()
|
|
File without changes
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
from pydantic import BaseModel, field_validator
|
|
2
|
+
from tkati_core.settings import InputSettings, OutputSettings, TomlBaseSettings
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class DedupSettings(BaseModel):
|
|
6
|
+
field: str
|
|
7
|
+
window_hours: int = 3
|
|
8
|
+
bucket_hours: int = 1
|
|
9
|
+
store_dir: str = "./dedup_store"
|
|
10
|
+
|
|
11
|
+
@field_validator("window_hours", "bucket_hours")
|
|
12
|
+
@classmethod
|
|
13
|
+
def _positive(cls, v: int) -> int:
|
|
14
|
+
if v <= 0:
|
|
15
|
+
raise ValueError("must be a positive number of hours")
|
|
16
|
+
return v
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class AppSettings(TomlBaseSettings):
|
|
20
|
+
input: InputSettings
|
|
21
|
+
output: OutputSettings
|
|
22
|
+
dlq: OutputSettings | None = None
|
|
23
|
+
dedup: DedupSettings
|
|
@@ -0,0 +1,224 @@
|
|
|
1
|
+
"""Bucketed, embedded RocksDB store for windowed key deduplication.
|
|
2
|
+
|
|
3
|
+
One RocksDB database directory per wall-clock-aligned hour bucket. A key is
|
|
4
|
+
considered "seen" if it exists in any bucket currently inside the window.
|
|
5
|
+
Existence-only store: values are always empty bytes, only key presence matters.
|
|
6
|
+
|
|
7
|
+
Lookups and writes are batched against RocksDB (one call per open bucket for
|
|
8
|
+
reads, one WriteBatch for writes) rather than one call per key — at realistic
|
|
9
|
+
batch sizes and millions of runs a day, per-row FFI calls into RocksDB are the
|
|
10
|
+
dominant cost, and rocksdict supports genuine batching for both directions.
|
|
11
|
+
|
|
12
|
+
Failure policy throughout: any per-bucket open/lookup/write failure is caught,
|
|
13
|
+
logged, and treated as "not seen" / skipped — never raised. This favors
|
|
14
|
+
forwarding a possible duplicate over silently dropping a real event.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
import math
|
|
18
|
+
import shutil
|
|
19
|
+
import time
|
|
20
|
+
from collections.abc import Iterable
|
|
21
|
+
from pathlib import Path
|
|
22
|
+
from typing import cast
|
|
23
|
+
|
|
24
|
+
import pyarrow as pa
|
|
25
|
+
import pyarrow.compute as pc
|
|
26
|
+
from loguru import logger
|
|
27
|
+
from rocksdict import Options, Rdict, WriteBatch
|
|
28
|
+
|
|
29
|
+
_BUCKET_PREFIX = "bucket-"
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _now() -> float:
|
|
33
|
+
"""Indirection over time.time() so tests can freeze this store's clock
|
|
34
|
+
without patching the global time module (which would also freeze
|
|
35
|
+
unrelated code, e.g. the Kafka consumer's poll-timeout bookkeeping)."""
|
|
36
|
+
return time.time()
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def _bucket_index(ts: float, bucket_seconds: int) -> int:
|
|
40
|
+
return int(ts // bucket_seconds)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class BucketedDedupStore:
|
|
44
|
+
def __init__(self, root_dir: str, window_hours: int, bucket_hours: int = 1) -> None:
|
|
45
|
+
self.root = Path(root_dir)
|
|
46
|
+
self.root.mkdir(parents=True, exist_ok=True)
|
|
47
|
+
self.bucket_seconds = bucket_hours * 3600
|
|
48
|
+
self.num_buckets = math.ceil(window_hours / bucket_hours)
|
|
49
|
+
self._opts = Options(raw_mode=True)
|
|
50
|
+
self._dbs: dict[int, Rdict] = {}
|
|
51
|
+
self._recover_existing_buckets()
|
|
52
|
+
|
|
53
|
+
def _bucket_path(self, bucket: int) -> Path:
|
|
54
|
+
return self.root / f"{_BUCKET_PREFIX}{bucket:012d}"
|
|
55
|
+
|
|
56
|
+
def _min_live_bucket(self) -> int:
|
|
57
|
+
current = _bucket_index(_now(), self.bucket_seconds)
|
|
58
|
+
return current - self.num_buckets + 1
|
|
59
|
+
|
|
60
|
+
def _recover_existing_buckets(self) -> None:
|
|
61
|
+
"""Reopen on-disk buckets still inside the window; destroy stale ones.
|
|
62
|
+
|
|
63
|
+
Makes dedup state survive a node restart.
|
|
64
|
+
"""
|
|
65
|
+
min_live = self._min_live_bucket()
|
|
66
|
+
for entry in sorted(self.root.glob(f"{_BUCKET_PREFIX}*")):
|
|
67
|
+
if not entry.is_dir():
|
|
68
|
+
continue
|
|
69
|
+
try:
|
|
70
|
+
bucket = int(entry.name.removeprefix(_BUCKET_PREFIX))
|
|
71
|
+
except ValueError:
|
|
72
|
+
logger.warning(f"Ignoring unrecognized entry in dedup store dir: {entry}")
|
|
73
|
+
continue
|
|
74
|
+
if bucket < min_live:
|
|
75
|
+
logger.info(f"Startup: removing stale dedup bucket {bucket} ({entry})")
|
|
76
|
+
self._destroy_path(entry)
|
|
77
|
+
continue
|
|
78
|
+
try:
|
|
79
|
+
self._dbs[bucket] = Rdict(str(entry), options=self._opts)
|
|
80
|
+
logger.info(f"Startup: reopened dedup bucket {bucket} from {entry}")
|
|
81
|
+
except Exception:
|
|
82
|
+
logger.exception(
|
|
83
|
+
f"Failed to reopen dedup bucket {bucket} at {entry}; "
|
|
84
|
+
"starting empty for this bucket (may pass through some duplicates)"
|
|
85
|
+
)
|
|
86
|
+
|
|
87
|
+
def _ensure_current_open(self) -> Rdict | None:
|
|
88
|
+
bucket = _bucket_index(_now(), self.bucket_seconds)
|
|
89
|
+
if bucket in self._dbs:
|
|
90
|
+
return self._dbs[bucket]
|
|
91
|
+
try:
|
|
92
|
+
db = Rdict(str(self._bucket_path(bucket)), options=self._opts)
|
|
93
|
+
except Exception:
|
|
94
|
+
logger.exception(f"Failed to open current dedup bucket {bucket}")
|
|
95
|
+
return None
|
|
96
|
+
self._dbs[bucket] = db
|
|
97
|
+
return db
|
|
98
|
+
|
|
99
|
+
def encode_keys(self, values: pa.Array | pa.ChunkedArray) -> list[bytes | None]:
|
|
100
|
+
"""Vectorized byte-key encoding for a column of scalar dedup-field values.
|
|
101
|
+
|
|
102
|
+
Cast to string via pyarrow compute (fast, vectorized — no per-row
|
|
103
|
+
Python type dispatch), then UTF-8 encode each non-null value. Nulls
|
|
104
|
+
stay None: they're never queried or stored, always passed through.
|
|
105
|
+
|
|
106
|
+
Float or timestamp dedup fields are discouraged: their string
|
|
107
|
+
representation isn't guaranteed stable across producers.
|
|
108
|
+
"""
|
|
109
|
+
strings = pc.cast(values, pa.string())
|
|
110
|
+
return [None if s is None else s.encode("utf-8") for s in strings.to_pylist()]
|
|
111
|
+
|
|
112
|
+
def filter_duplicates(
|
|
113
|
+
self, keys: list[bytes | None]
|
|
114
|
+
) -> tuple[pa.BooleanArray, list[bytes]]:
|
|
115
|
+
"""
|
|
116
|
+
Given per-row encoded keys (None = no key, always kept), returns
|
|
117
|
+
(keep_mask, keys_to_mark_seen). keep_mask[i] corresponds to keys[i]
|
|
118
|
+
and is directly usable with pyarrow.Table.filter().
|
|
119
|
+
|
|
120
|
+
In-batch duplicates (two rows with the same key, neither yet in the
|
|
121
|
+
store) are resolved locally; the remaining unique candidates are
|
|
122
|
+
checked against the store in a single batched round trip per open
|
|
123
|
+
bucket, not one lookup per key.
|
|
124
|
+
"""
|
|
125
|
+
keep_mask = [False] * len(keys)
|
|
126
|
+
seen_in_batch: set[bytes] = set()
|
|
127
|
+
to_check: list[bytes] = []
|
|
128
|
+
to_check_idx: list[int] = []
|
|
129
|
+
|
|
130
|
+
for i, key in enumerate(keys):
|
|
131
|
+
if key is None:
|
|
132
|
+
keep_mask[i] = True
|
|
133
|
+
continue
|
|
134
|
+
if key in seen_in_batch:
|
|
135
|
+
continue
|
|
136
|
+
seen_in_batch.add(key)
|
|
137
|
+
to_check.append(key)
|
|
138
|
+
to_check_idx.append(i)
|
|
139
|
+
|
|
140
|
+
already_seen = self._batch_contains(to_check)
|
|
141
|
+
|
|
142
|
+
new_keys: list[bytes] = []
|
|
143
|
+
for idx, key, seen in zip(to_check_idx, to_check, already_seen, strict=True):
|
|
144
|
+
if not seen:
|
|
145
|
+
keep_mask[idx] = True
|
|
146
|
+
new_keys.append(key)
|
|
147
|
+
|
|
148
|
+
return pa.array(keep_mask, type=pa.bool_()), new_keys
|
|
149
|
+
|
|
150
|
+
def _batch_contains(self, keys: list[bytes]) -> list[bool]:
|
|
151
|
+
if not keys:
|
|
152
|
+
return []
|
|
153
|
+
found = [False] * len(keys)
|
|
154
|
+
for bucket, db in list(self._dbs.items()):
|
|
155
|
+
try:
|
|
156
|
+
# rocksdict's stub declares `List[...]` invariantly, so
|
|
157
|
+
# list[bytes] isn't accepted as-is, and doesn't distinguish
|
|
158
|
+
# the list-in/list-out overload from the scalar one for the
|
|
159
|
+
# return type either — get() genuinely returns a list here
|
|
160
|
+
# since `keys` is a list.
|
|
161
|
+
keys_arg = cast("list[str | int | float | bytes]", keys)
|
|
162
|
+
results = cast("list[bytes | None]", db.get(keys_arg))
|
|
163
|
+
except Exception:
|
|
164
|
+
logger.exception(f"Batch lookup failed against dedup bucket {bucket}; skipping it")
|
|
165
|
+
continue
|
|
166
|
+
for i, r in enumerate(results):
|
|
167
|
+
if r is not None:
|
|
168
|
+
found[i] = True
|
|
169
|
+
return found
|
|
170
|
+
|
|
171
|
+
def add_many(self, keys: Iterable[bytes]) -> None:
|
|
172
|
+
keys = list(keys)
|
|
173
|
+
if not keys:
|
|
174
|
+
return
|
|
175
|
+
db = self._ensure_current_open()
|
|
176
|
+
if db is None:
|
|
177
|
+
return
|
|
178
|
+
try:
|
|
179
|
+
wb = WriteBatch(raw_mode=True)
|
|
180
|
+
for key in keys:
|
|
181
|
+
wb.put(key, b"")
|
|
182
|
+
db.write(wb)
|
|
183
|
+
except Exception:
|
|
184
|
+
logger.exception("Failed to batch-write keys into current dedup bucket")
|
|
185
|
+
|
|
186
|
+
def contains(self, key: bytes) -> bool:
|
|
187
|
+
return self._batch_contains([key])[0]
|
|
188
|
+
|
|
189
|
+
def add(self, key: bytes) -> None:
|
|
190
|
+
self.add_many([key])
|
|
191
|
+
|
|
192
|
+
def _destroy_path(self, path: Path) -> None:
|
|
193
|
+
try:
|
|
194
|
+
Rdict.destroy(str(path), self._opts)
|
|
195
|
+
except Exception:
|
|
196
|
+
logger.exception(f"Rdict.destroy failed for {path}; falling back to rmtree")
|
|
197
|
+
shutil.rmtree(path, ignore_errors=True)
|
|
198
|
+
|
|
199
|
+
def cleanup_expired(self) -> None:
|
|
200
|
+
"""Close+delete on-disk buckets that fell out of the window.
|
|
201
|
+
|
|
202
|
+
Called at the start of every iteration (before the dedupe check runs
|
|
203
|
+
against possibly-stale buckets), and cheap to call every time: only
|
|
204
|
+
does real I/O once per hour rollover, since _min_live_bucket() only
|
|
205
|
+
changes then.
|
|
206
|
+
"""
|
|
207
|
+
min_live = self._min_live_bucket()
|
|
208
|
+
for bucket in list(self._dbs.keys()):
|
|
209
|
+
if bucket < min_live:
|
|
210
|
+
db = self._dbs.pop(bucket)
|
|
211
|
+
try:
|
|
212
|
+
db.close()
|
|
213
|
+
except Exception:
|
|
214
|
+
logger.exception(f"Error closing expired dedup bucket {bucket}")
|
|
215
|
+
self._destroy_path(self._bucket_path(bucket))
|
|
216
|
+
logger.info(f"Expired dedup bucket {bucket} removed")
|
|
217
|
+
|
|
218
|
+
def close(self) -> None:
|
|
219
|
+
for bucket, db in self._dbs.items():
|
|
220
|
+
try:
|
|
221
|
+
db.close()
|
|
222
|
+
except Exception:
|
|
223
|
+
logger.exception(f"Error closing dedup bucket {bucket} on shutdown")
|
|
224
|
+
self._dbs.clear()
|
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tkati-node-dedup
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: Kafka-to-Kafka streaming node that deduplicates events by a configurable field within a rolling processing-time window
|
|
5
|
+
Requires-Python: >=3.13
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
Requires-Dist: tkati-core==0.3.0
|
|
8
|
+
Requires-Dist: loguru>=0.7.0
|
|
9
|
+
Requires-Dist: pydantic-settings>=2.11.0
|
|
10
|
+
Requires-Dist: pyarrow>=21.0.0
|
|
11
|
+
Requires-Dist: rocksdict>=0.3.0
|
|
12
|
+
|
|
13
|
+
# tkati-node-dedup — streaming deduplication node
|
|
14
|
+
|
|
15
|
+
Reads batches from a Kafka input topic, drops events that are duplicates of an
|
|
16
|
+
event seen on the same `field` within a rolling processing-time window, and
|
|
17
|
+
writes the deduplicated batch to a configurable output. Duplicate state is
|
|
18
|
+
tracked in an embedded, on-disk RocksDB store local to this process — no
|
|
19
|
+
external dedup service is required.
|
|
20
|
+
|
|
21
|
+
## Configuration
|
|
22
|
+
|
|
23
|
+
Settings are loaded from a TOML file. Set the `SETTINGS_FILE` environment
|
|
24
|
+
variable to point to it (defaults to `settings.toml`).
|
|
25
|
+
|
|
26
|
+
```toml
|
|
27
|
+
[input]
|
|
28
|
+
type = "kafka"
|
|
29
|
+
|
|
30
|
+
[input.connection]
|
|
31
|
+
broker = "redpanda:29092"
|
|
32
|
+
|
|
33
|
+
[input.topic]
|
|
34
|
+
name = "raw_event"
|
|
35
|
+
|
|
36
|
+
[input.topic.schema]
|
|
37
|
+
uid = "string"
|
|
38
|
+
time = "timestamp[ms]"
|
|
39
|
+
# … other columns
|
|
40
|
+
|
|
41
|
+
[input.consumer]
|
|
42
|
+
group_id = "node-dedup-group"
|
|
43
|
+
batch_size = 1000
|
|
44
|
+
batch_timeout_sec = 10
|
|
45
|
+
auto_offset_reset = "latest"
|
|
46
|
+
|
|
47
|
+
[output]
|
|
48
|
+
type = "kafka"
|
|
49
|
+
|
|
50
|
+
[output.connection]
|
|
51
|
+
broker = "redpanda:29092"
|
|
52
|
+
|
|
53
|
+
[output.topic]
|
|
54
|
+
name = "raw_event_deduped"
|
|
55
|
+
|
|
56
|
+
[dedup]
|
|
57
|
+
field = "uid" # column in input.topic.schema to dedup by
|
|
58
|
+
window_hours = 3 # rolling dedup window
|
|
59
|
+
bucket_hours = 1 # on-disk bucket granularity (effective window is
|
|
60
|
+
# window_hours .. window_hours + bucket_hours)
|
|
61
|
+
store_dir = "/var/lib/tkati-node-dedup/store"
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
Output and DLQ follow the same `OutputSettings` shape as `tkati-node-el`
|
|
65
|
+
(`"kafka"` or `"clickhouse"`) — see that package's README for the full
|
|
66
|
+
connection/table config shape.
|
|
67
|
+
|
|
68
|
+
## Delivery & dedup guarantees
|
|
69
|
+
|
|
70
|
+
**Delivery: at-least-once.** Offsets are committed only after (1) the
|
|
71
|
+
filtered batch is produced and confirmed delivered (`produce_arrow` followed
|
|
72
|
+
by a blocking `flush`), and (2) the surviving keys are durably written to the
|
|
73
|
+
current RocksDB bucket. If the process crashes between steps, the same input
|
|
74
|
+
batch is re-read at restart; because the keys from a completed produce are
|
|
75
|
+
already marked seen, re-processing that batch is a no-op (or reproduces only
|
|
76
|
+
the genuinely-new subset) rather than losing data.
|
|
77
|
+
|
|
78
|
+
**On any internal dedup-store failure — a bucket won't open, a lookup errors,
|
|
79
|
+
a disk I/O error — the node treats the event as NOT a duplicate and forwards
|
|
80
|
+
it.** This node will occasionally forward a duplicate it should have caught,
|
|
81
|
+
but will never silently drop a real event because of dedup-store trouble.
|
|
82
|
+
|
|
83
|
+
**The window is approximate, not exact.** Because state is bucketed in
|
|
84
|
+
`bucket_hours` increments (default 1h) rather than a true sliding window, the
|
|
85
|
+
effective dedup window is between `window_hours` and
|
|
86
|
+
`window_hours + bucket_hours`. Stale buckets are deleted from disk
|
|
87
|
+
automatically once they fall outside the window — checked once per iteration,
|
|
88
|
+
so state never grows unbounded.
|
|
89
|
+
|
|
90
|
+
**Bucketing is by processing time**, not any timestamp field in the event
|
|
91
|
+
payload — an event's bucket is when this node handles it, not when it
|
|
92
|
+
happened upstream.
|
|
93
|
+
|
|
94
|
+
## IMPORTANT: dedup state is local to this process
|
|
95
|
+
|
|
96
|
+
The RocksDB store lives on local disk at `store_dir` and is **not shared**
|
|
97
|
+
between instances. Running multiple concurrent instances of this node against
|
|
98
|
+
the same input topic (e.g. multiple consumers in the same consumer group, or
|
|
99
|
+
multiple replicas) will **not** dedup correctly across instances unless the
|
|
100
|
+
input is partitioned such that all events sharing a `field` value are always
|
|
101
|
+
routed to the *same* instance (e.g. Kafka partitioning keyed on `field`, one
|
|
102
|
+
node instance per partition or partition subset it exclusively owns). Running
|
|
103
|
+
this node with more parallelism than that will let duplicates leak through
|
|
104
|
+
across instance boundaries. This is a direct consequence of choosing an
|
|
105
|
+
embedded local-file store instead of a shared external one — evaluate this
|
|
106
|
+
tradeoff before scaling this node horizontally.
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
src/tkati_node_dedup/__init__.py
|
|
4
|
+
src/tkati_node_dedup/__main__.py
|
|
5
|
+
src/tkati_node_dedup/main.py
|
|
6
|
+
src/tkati_node_dedup/py.typed
|
|
7
|
+
src/tkati_node_dedup/settings.py
|
|
8
|
+
src/tkati_node_dedup/store.py
|
|
9
|
+
src/tkati_node_dedup.egg-info/PKG-INFO
|
|
10
|
+
src/tkati_node_dedup.egg-info/SOURCES.txt
|
|
11
|
+
src/tkati_node_dedup.egg-info/dependency_links.txt
|
|
12
|
+
src/tkati_node_dedup.egg-info/entry_points.txt
|
|
13
|
+
src/tkati_node_dedup.egg-info/requires.txt
|
|
14
|
+
src/tkati_node_dedup.egg-info/top_level.txt
|
|
15
|
+
tests/test_node_dedup.py
|
|
16
|
+
tests/test_store.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
tkati_node_dedup
|
|
@@ -0,0 +1,228 @@
|
|
|
1
|
+
import time
|
|
2
|
+
from unittest.mock import MagicMock
|
|
3
|
+
|
|
4
|
+
import orjson
|
|
5
|
+
import pyarrow as pa
|
|
6
|
+
import pytest
|
|
7
|
+
from confluent_kafka import Consumer as RawConsumer
|
|
8
|
+
from confluent_kafka import Producer as RawProducer
|
|
9
|
+
from tkati_core.kafka.consumer import KafkaConsumer
|
|
10
|
+
from tkati_core.kafka.producer import KafkaProducer
|
|
11
|
+
from tkati_core.kafka.settings import KafkaOutputSettings
|
|
12
|
+
from tkati_node_dedup.main import run_one_iteration
|
|
13
|
+
from tkati_node_dedup.settings import AppSettings
|
|
14
|
+
from tkati_node_dedup.store import BucketedDedupStore
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def _make_consumer(test_settings: AppSettings) -> KafkaConsumer:
|
|
18
|
+
assert test_settings.input.type == "kafka"
|
|
19
|
+
return KafkaConsumer(
|
|
20
|
+
kafka_config={
|
|
21
|
+
"bootstrap.servers": test_settings.input.connection.broker,
|
|
22
|
+
"group.id": test_settings.input.consumer.group_id,
|
|
23
|
+
"auto.offset.reset": "earliest",
|
|
24
|
+
"enable.auto.commit": False,
|
|
25
|
+
},
|
|
26
|
+
topic_name=test_settings.input.topic.name,
|
|
27
|
+
input_schema=test_settings.input.topic.schema,
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def _make_producer(test_settings: AppSettings) -> KafkaProducer:
|
|
32
|
+
assert isinstance(test_settings.output, KafkaOutputSettings)
|
|
33
|
+
return KafkaProducer.from_output_settings(test_settings.output)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def _make_store(test_settings: AppSettings) -> BucketedDedupStore:
|
|
37
|
+
return BucketedDedupStore(
|
|
38
|
+
root_dir=test_settings.dedup.store_dir,
|
|
39
|
+
window_hours=test_settings.dedup.window_hours,
|
|
40
|
+
bucket_hours=test_settings.dedup.bucket_hours,
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _drain_output(test_settings: AppSettings, expected: int, timeout: float = 10.0) -> list[dict]:
|
|
45
|
+
assert isinstance(test_settings.output, KafkaOutputSettings)
|
|
46
|
+
consumer = RawConsumer(
|
|
47
|
+
{
|
|
48
|
+
"bootstrap.servers": test_settings.output.connection.broker,
|
|
49
|
+
"group.id": f"verify-{test_settings.output.topic.name}",
|
|
50
|
+
"auto.offset.reset": "earliest",
|
|
51
|
+
"enable.auto.commit": False,
|
|
52
|
+
}
|
|
53
|
+
)
|
|
54
|
+
consumer.subscribe([test_settings.output.topic.name])
|
|
55
|
+
rows: list[dict] = []
|
|
56
|
+
deadline = time.time() + timeout
|
|
57
|
+
try:
|
|
58
|
+
while len(rows) < expected and time.time() < deadline:
|
|
59
|
+
msg = consumer.poll(1.0)
|
|
60
|
+
if msg is None or msg.error():
|
|
61
|
+
continue
|
|
62
|
+
value = msg.value()
|
|
63
|
+
assert value is not None
|
|
64
|
+
rows.append(orjson.loads(value))
|
|
65
|
+
finally:
|
|
66
|
+
consumer.close()
|
|
67
|
+
return rows
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def _event(uid: str | None, val: int) -> dict:
|
|
71
|
+
return {"uid": uid, "time": int(time.time() * 1000), "val": val}
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def test_basic_in_batch_dedup(
|
|
75
|
+
kafka_producer: RawProducer, test_settings: AppSettings
|
|
76
|
+
) -> None:
|
|
77
|
+
"""Two messages with the same uid produced before one poll: only one survives."""
|
|
78
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event("dup-1", 1)))
|
|
79
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event("dup-1", 2)))
|
|
80
|
+
kafka_producer.flush()
|
|
81
|
+
|
|
82
|
+
consumer = _make_consumer(test_settings)
|
|
83
|
+
producer = _make_producer(test_settings)
|
|
84
|
+
store = _make_store(test_settings)
|
|
85
|
+
try:
|
|
86
|
+
run_one_iteration(consumer, producer, store, test_settings)
|
|
87
|
+
finally:
|
|
88
|
+
consumer.close()
|
|
89
|
+
store.close()
|
|
90
|
+
|
|
91
|
+
rows = _drain_output(test_settings, expected=1)
|
|
92
|
+
assert len(rows) == 1
|
|
93
|
+
assert rows[0]["uid"] == "dup-1"
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def test_cross_batch_dedup(kafka_producer: RawProducer, test_settings: AppSettings) -> None:
|
|
97
|
+
"""Same uid produced across two separate iterations: only the first survives."""
|
|
98
|
+
consumer = _make_consumer(test_settings)
|
|
99
|
+
producer = _make_producer(test_settings)
|
|
100
|
+
store = _make_store(test_settings)
|
|
101
|
+
try:
|
|
102
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event("dup-2", 1)))
|
|
103
|
+
kafka_producer.flush()
|
|
104
|
+
run_one_iteration(consumer, producer, store, test_settings)
|
|
105
|
+
|
|
106
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event("dup-2", 2)))
|
|
107
|
+
kafka_producer.flush()
|
|
108
|
+
run_one_iteration(consumer, producer, store, test_settings)
|
|
109
|
+
finally:
|
|
110
|
+
consumer.close()
|
|
111
|
+
store.close()
|
|
112
|
+
|
|
113
|
+
rows = _drain_output(test_settings, expected=1)
|
|
114
|
+
assert len(rows) == 1
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def test_bucket_rollover_lets_key_through_again(
|
|
118
|
+
kafka_producer: RawProducer, test_settings: AppSettings, monkeypatch
|
|
119
|
+
) -> None:
|
|
120
|
+
test_settings.dedup.window_hours = 1
|
|
121
|
+
test_settings.dedup.bucket_hours = 1
|
|
122
|
+
|
|
123
|
+
now = [1_000_000.0]
|
|
124
|
+
monkeypatch.setattr("tkati_node_dedup.store._now", lambda: now[0])
|
|
125
|
+
|
|
126
|
+
consumer = _make_consumer(test_settings)
|
|
127
|
+
producer = _make_producer(test_settings)
|
|
128
|
+
store = _make_store(test_settings)
|
|
129
|
+
try:
|
|
130
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event("dup-3", 1)))
|
|
131
|
+
kafka_producer.flush()
|
|
132
|
+
run_one_iteration(consumer, producer, store, test_settings)
|
|
133
|
+
|
|
134
|
+
# Advance well past window_hours + bucket_hours so the bucket ages out.
|
|
135
|
+
now[0] += 5 * 3600
|
|
136
|
+
|
|
137
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event("dup-3", 2)))
|
|
138
|
+
kafka_producer.flush()
|
|
139
|
+
run_one_iteration(consumer, producer, store, test_settings)
|
|
140
|
+
finally:
|
|
141
|
+
consumer.close()
|
|
142
|
+
store.close()
|
|
143
|
+
|
|
144
|
+
rows = _drain_output(test_settings, expected=2)
|
|
145
|
+
assert len(rows) == 2
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
def test_null_dedup_field_passes_through(
|
|
149
|
+
kafka_producer: RawProducer, test_settings: AppSettings
|
|
150
|
+
) -> None:
|
|
151
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event(None, 1)))
|
|
152
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event(None, 2)))
|
|
153
|
+
kafka_producer.flush()
|
|
154
|
+
|
|
155
|
+
consumer = _make_consumer(test_settings)
|
|
156
|
+
producer = _make_producer(test_settings)
|
|
157
|
+
store = _make_store(test_settings)
|
|
158
|
+
try:
|
|
159
|
+
run_one_iteration(consumer, producer, store, test_settings)
|
|
160
|
+
finally:
|
|
161
|
+
consumer.close()
|
|
162
|
+
store.close()
|
|
163
|
+
|
|
164
|
+
rows = _drain_output(test_settings, expected=2)
|
|
165
|
+
assert len(rows) == 2
|
|
166
|
+
|
|
167
|
+
|
|
168
|
+
def test_missing_dedup_field_in_schema(
|
|
169
|
+
kafka_producer: RawProducer, test_settings: AppSettings, caplog
|
|
170
|
+
) -> None:
|
|
171
|
+
test_settings.dedup.field = "does_not_exist"
|
|
172
|
+
|
|
173
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event("dup-4", 1)))
|
|
174
|
+
kafka_producer.produce(test_settings.input.topic.name, value=orjson.dumps(_event("dup-4", 2)))
|
|
175
|
+
kafka_producer.flush()
|
|
176
|
+
|
|
177
|
+
consumer = _make_consumer(test_settings)
|
|
178
|
+
producer = _make_producer(test_settings)
|
|
179
|
+
store = _make_store(test_settings)
|
|
180
|
+
try:
|
|
181
|
+
run_one_iteration(consumer, producer, store, test_settings)
|
|
182
|
+
finally:
|
|
183
|
+
consumer.close()
|
|
184
|
+
store.close()
|
|
185
|
+
|
|
186
|
+
# Both rows pass through unfiltered — there's no column to dedup by.
|
|
187
|
+
rows = _drain_output(test_settings, expected=2)
|
|
188
|
+
assert len(rows) == 2
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
def test_crash_before_flush_does_not_mark_seen_or_commit(tmp_path) -> None:
|
|
192
|
+
"""If produce/flush fails, the key must not be marked seen and the offset
|
|
193
|
+
must not be committed — re-processing the same message afterward must not
|
|
194
|
+
treat it as a duplicate."""
|
|
195
|
+
batch = pa.table({"uid": ["crash-uid"], "val": [1]})
|
|
196
|
+
|
|
197
|
+
consumer = MagicMock()
|
|
198
|
+
consumer.read_arrow.return_value = batch
|
|
199
|
+
|
|
200
|
+
producer = MagicMock()
|
|
201
|
+
producer.produce_arrow = MagicMock()
|
|
202
|
+
producer.flush = MagicMock(side_effect=RuntimeError("boom"))
|
|
203
|
+
|
|
204
|
+
store = BucketedDedupStore(str(tmp_path), window_hours=3, bucket_hours=1)
|
|
205
|
+
settings = MagicMock()
|
|
206
|
+
settings.input.consumer.batch_size = 100
|
|
207
|
+
settings.input.consumer.batch_timeout_sec = 5
|
|
208
|
+
settings.dedup.field = "uid"
|
|
209
|
+
|
|
210
|
+
with pytest.raises(RuntimeError, match="boom"):
|
|
211
|
+
run_one_iteration(consumer, producer, store, settings)
|
|
212
|
+
|
|
213
|
+
consumer.commit.assert_not_called()
|
|
214
|
+
assert store.contains(b"crash-uid") is False
|
|
215
|
+
|
|
216
|
+
# Simulate a restart: same batch re-read, this time produce succeeds.
|
|
217
|
+
producer2 = MagicMock()
|
|
218
|
+
producer2.produce_arrow = MagicMock()
|
|
219
|
+
producer2.flush = MagicMock()
|
|
220
|
+
|
|
221
|
+
run_one_iteration(consumer, producer2, store, settings)
|
|
222
|
+
|
|
223
|
+
produced_table = producer2.produce_arrow.call_args[0][0]
|
|
224
|
+
assert len(produced_table) == 1 # not dropped as a duplicate
|
|
225
|
+
consumer.commit.assert_called_once()
|
|
226
|
+
assert store.contains(b"crash-uid") is True
|
|
227
|
+
|
|
228
|
+
store.close()
|
|
@@ -0,0 +1,114 @@
|
|
|
1
|
+
from typing import cast
|
|
2
|
+
|
|
3
|
+
import pyarrow as pa
|
|
4
|
+
from rocksdict import Rdict
|
|
5
|
+
from tkati_node_dedup.store import BucketedDedupStore
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def _key(value: str) -> bytes:
|
|
9
|
+
return value.encode("utf-8")
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def test_add_and_contains(tmp_path) -> None:
|
|
13
|
+
store = BucketedDedupStore(str(tmp_path), window_hours=3, bucket_hours=1)
|
|
14
|
+
key = _key("abc123")
|
|
15
|
+
assert store.contains(key) is False
|
|
16
|
+
store.add(key)
|
|
17
|
+
assert store.contains(key) is True
|
|
18
|
+
assert store.contains(_key("other")) is False
|
|
19
|
+
store.close()
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def test_bucket_rollover_expiry(tmp_path, monkeypatch) -> None:
|
|
23
|
+
now = [1_000_000.0]
|
|
24
|
+
monkeypatch.setattr("tkati_node_dedup.store._now", lambda: now[0])
|
|
25
|
+
|
|
26
|
+
store = BucketedDedupStore(str(tmp_path), window_hours=2, bucket_hours=1)
|
|
27
|
+
key = _key("abc123")
|
|
28
|
+
store.add(key)
|
|
29
|
+
assert store.contains(key) is True
|
|
30
|
+
|
|
31
|
+
# Advance past window_hours + bucket_hours so the bucket fully ages out.
|
|
32
|
+
now[0] += 4 * 3600
|
|
33
|
+
store.cleanup_expired()
|
|
34
|
+
|
|
35
|
+
assert store.contains(key) is False
|
|
36
|
+
remaining = list(tmp_path.glob("bucket-*"))
|
|
37
|
+
assert remaining == []
|
|
38
|
+
store.close()
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def test_restart_resumes_existing_buckets(tmp_path) -> None:
|
|
42
|
+
store = BucketedDedupStore(str(tmp_path), window_hours=3, bucket_hours=1)
|
|
43
|
+
key = _key("abc123")
|
|
44
|
+
store.add(key)
|
|
45
|
+
store.close()
|
|
46
|
+
|
|
47
|
+
store2 = BucketedDedupStore(str(tmp_path), window_hours=3, bucket_hours=1)
|
|
48
|
+
assert store2.contains(key) is True
|
|
49
|
+
store2.close()
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def test_restart_discards_stale_on_disk_buckets(tmp_path, monkeypatch) -> None:
|
|
53
|
+
now = [1_000_000.0]
|
|
54
|
+
monkeypatch.setattr("tkati_node_dedup.store._now", lambda: now[0])
|
|
55
|
+
|
|
56
|
+
store = BucketedDedupStore(str(tmp_path), window_hours=1, bucket_hours=1)
|
|
57
|
+
key = _key("abc123")
|
|
58
|
+
store.add(key)
|
|
59
|
+
store.close()
|
|
60
|
+
|
|
61
|
+
now[0] += 10 * 3600 # far outside the window
|
|
62
|
+
|
|
63
|
+
store2 = BucketedDedupStore(str(tmp_path), window_hours=1, bucket_hours=1)
|
|
64
|
+
assert store2.contains(key) is False
|
|
65
|
+
assert list(tmp_path.glob("bucket-*")) == []
|
|
66
|
+
store2.close()
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def test_lookup_failure_is_treated_as_not_seen(tmp_path) -> None:
|
|
70
|
+
store = BucketedDedupStore(str(tmp_path), window_hours=1, bucket_hours=1)
|
|
71
|
+
key = _key("abc123")
|
|
72
|
+
|
|
73
|
+
class BrokenDict:
|
|
74
|
+
def get(self, _keys: list[bytes]) -> list[bytes | None]:
|
|
75
|
+
raise RuntimeError("boom")
|
|
76
|
+
|
|
77
|
+
def close(self) -> None:
|
|
78
|
+
pass
|
|
79
|
+
|
|
80
|
+
store._dbs[0] = cast(Rdict, BrokenDict())
|
|
81
|
+
|
|
82
|
+
assert store.contains(key) is False # must not raise
|
|
83
|
+
store.close()
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def test_encode_keys_vectorized_cast_and_nulls(tmp_path) -> None:
|
|
87
|
+
store = BucketedDedupStore(str(tmp_path), window_hours=3, bucket_hours=1)
|
|
88
|
+
|
|
89
|
+
strings = store.encode_keys(pa.array(["abc", None, "def"], type=pa.string()))
|
|
90
|
+
assert strings == [b"abc", None, b"def"]
|
|
91
|
+
|
|
92
|
+
ints = store.encode_keys(pa.array([1, None, 42], type=pa.int64()))
|
|
93
|
+
assert ints == [b"1", None, b"42"]
|
|
94
|
+
|
|
95
|
+
store.close()
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def test_filter_duplicates_in_batch_and_store_duplicates(tmp_path) -> None:
|
|
99
|
+
store = BucketedDedupStore(str(tmp_path), window_hours=3, bucket_hours=1)
|
|
100
|
+
store.add(_key("already-seen"))
|
|
101
|
+
|
|
102
|
+
keys = [
|
|
103
|
+
_key("already-seen"), # duplicate of a key already in the store
|
|
104
|
+
_key("fresh"), # new, kept
|
|
105
|
+
_key("fresh"), # duplicate of the row above, within this same batch
|
|
106
|
+
None, # no key, always kept, never stored
|
|
107
|
+
]
|
|
108
|
+
|
|
109
|
+
keep_mask, new_keys = store.filter_duplicates(keys)
|
|
110
|
+
|
|
111
|
+
assert keep_mask.to_pylist() == [False, True, False, True]
|
|
112
|
+
assert new_keys == [_key("fresh")]
|
|
113
|
+
|
|
114
|
+
store.close()
|