tkati-core 0.1.0__tar.gz

tkati_core-0.1.0/PKG-INFO

@@ -0,0 +1,118 @@
+Metadata-Version: 2.4
+Name: tkati-core
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+Requires-Dist: confluent-kafka>=2.11.0
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: orjson>=3.9.0
+Requires-Dist: pyarrow>=21.0.0
+Requires-Dist: pydantic-settings>=2.11.0
+
+# tkati-core
+
+For now we assume that each node gets one input and produces one output in a
+form of kafka stream.
+
+## Settings
+
+General form of settings is:
+
+```toml
+[input.topic]
+# definition of input stream:
+# - broker
+# - topic name
+# - message schema
+# - message format = "json" / "arrow-batch"
+
+[input.consumer]
+# parameters local to this consumer
+# - group_id
+# - batch_size
+# - batch_timeout_sec
+# - auto_offset_reset
+
+[output.topic]
+# definition of output stream
+# - broker
+# - topic name
+# - message schema
+# - message format = "json" / "arrow-batch"
+# - key_column (optional) = column to use as the Kafka message key
+
+[...]
+# settings specific to node function
+```
+
+## Usage
+
+### Constructing a consumer from settings
+
+Use `KafkaArrowConsumer.from_input_settings` to construct a consumer directly from
+`KafkaInputSettings` — no need to manually map fields to Confluent Kafka config keys.
+
+```python
+from tkati_core.settings import TomlBaseSettings, KafkaInputSettings
+from tkati_core.consumer import KafkaArrowConsumer
+
+class AppSettings(TomlBaseSettings):
+    input: KafkaInputSettings
+    # ...
+
+settings = AppSettings()
+consumer = KafkaArrowConsumer.from_input_settings(settings.input)
+
+# Read a batch
+table = consumer.read_to_pyarrow(
+    aggregation_interval_seconds=settings.input.consumer.batch_timeout_sec,
+    max_events_to_aggregate=settings.input.consumer.batch_size,
+)
+consumer.commit()
+```
+
+The factory method sets `enable.auto.commit=False` — offsets must be committed explicitly
+via `consumer.commit()`.
+
+### Constructing a producer from settings
+
+Use `KafkaArrowProducer.from_output_settings` to construct a producer directly from
+`KafkaOutputSettings`. It accepts PyArrow tables or record batches and handles
+serialization according to the topic's `format` setting.
+
+```python
+from tkati_core.settings import TomlBaseSettings, KafkaOutputSettings
+from tkati_core.producer import KafkaArrowProducer
+
+class AppSettings(TomlBaseSettings):
+    output: KafkaOutputSettings
+    # ...
+
+settings = AppSettings()
+producer = KafkaArrowProducer.from_output_settings(settings.output)
+
+# Produce a PyArrow table (one message per row for "json" format)
+producer.produce(table)
+producer.flush()
+producer.close()  # flushes and releases resources
+```
+
+**Formats** — controlled by `output.topic.format` in `settings.toml`:
+
+- `"json"` *(default)*: each row becomes a separate Kafka message serialized with orjson.
+- `"arrow-batch"`: the entire table is serialized as a single Arrow IPC stream message.
+
+**Message keys** — controlled by `output.topic.key_column` in `settings.toml`:
+
+```toml
+[output.topic]
+broker = "localhost:9092"
+name = "my-output-topic"
+key_column = "customer_id"   # column whose value becomes the Kafka message key
+```
+
+`key_column` is optional. When omitted (or `None`), messages are produced without a key.
+When set, the value of that column for each row is used as the Kafka message key
+(JSON format only — ignored for `"arrow-batch"`). This determines which Kafka partition
+each message is routed to.

tkati_core-0.1.0/README.md

@@ -0,0 +1,106 @@
+# tkati-core
+
+For now we assume that each node gets one input and produces one output in a
+form of kafka stream.
+
+## Settings
+
+General form of settings is:
+
+```toml
+[input.topic]
+# definition of input stream:
+# - broker
+# - topic name
+# - message schema
+# - message format = "json" / "arrow-batch"
+
+[input.consumer]
+# parameters local to this consumer
+# - group_id
+# - batch_size
+# - batch_timeout_sec
+# - auto_offset_reset
+
+[output.topic]
+# definition of output stream
+# - broker
+# - topic name
+# - message schema
+# - message format = "json" / "arrow-batch"
+# - key_column (optional) = column to use as the Kafka message key
+
+[...]
+# settings specific to node function
+```
+
+## Usage
+
+### Constructing a consumer from settings
+
+Use `KafkaArrowConsumer.from_input_settings` to construct a consumer directly from
+`KafkaInputSettings` — no need to manually map fields to Confluent Kafka config keys.
+
+```python
+from tkati_core.settings import TomlBaseSettings, KafkaInputSettings
+from tkati_core.consumer import KafkaArrowConsumer
+
+class AppSettings(TomlBaseSettings):
+    input: KafkaInputSettings
+    # ...
+
+settings = AppSettings()
+consumer = KafkaArrowConsumer.from_input_settings(settings.input)
+
+# Read a batch
+table = consumer.read_to_pyarrow(
+    aggregation_interval_seconds=settings.input.consumer.batch_timeout_sec,
+    max_events_to_aggregate=settings.input.consumer.batch_size,
+)
+consumer.commit()
+```
+
+The factory method sets `enable.auto.commit=False` — offsets must be committed explicitly
+via `consumer.commit()`.
+
+### Constructing a producer from settings
+
+Use `KafkaArrowProducer.from_output_settings` to construct a producer directly from
+`KafkaOutputSettings`. It accepts PyArrow tables or record batches and handles
+serialization according to the topic's `format` setting.
+
+```python
+from tkati_core.settings import TomlBaseSettings, KafkaOutputSettings
+from tkati_core.producer import KafkaArrowProducer
+
+class AppSettings(TomlBaseSettings):
+    output: KafkaOutputSettings
+    # ...
+
+settings = AppSettings()
+producer = KafkaArrowProducer.from_output_settings(settings.output)
+
+# Produce a PyArrow table (one message per row for "json" format)
+producer.produce(table)
+producer.flush()
+producer.close()  # flushes and releases resources
+```
+
+**Formats** — controlled by `output.topic.format` in `settings.toml`:
+
+- `"json"` *(default)*: each row becomes a separate Kafka message serialized with orjson.
+- `"arrow-batch"`: the entire table is serialized as a single Arrow IPC stream message.
+
+**Message keys** — controlled by `output.topic.key_column` in `settings.toml`:
+
+```toml
+[output.topic]
+broker = "localhost:9092"
+name = "my-output-topic"
+key_column = "customer_id"   # column whose value becomes the Kafka message key
+```
+
+`key_column` is optional. When omitted (or `None`), messages are produced without a key.
+When set, the value of that column for each row is used as the Kafka message key
+(JSON format only — ignored for `"arrow-batch"`). This determines which Kafka partition
+each message is routed to.

tkati_core-0.1.0/pyproject.toml

@@ -0,0 +1,28 @@
+[project]
+name = "tkati-core"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = [
+    "confluent-kafka>=2.11.0",
+    "loguru>=0.7.3",
+    "orjson>=3.9.0",
+    "pyarrow>=21.0.0",
+    "pydantic-settings>=2.11.0",
+]
+
+[dependency-groups]
+dev = ["pytest>=9.0.1"]
+
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["tkati_core"]
+
+[tool.uv-workspace-codegen]
+generate = true
+template_type = "package"

tkati_core-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tkati_core-0.1.0/tests/test_consumer.py

@@ -0,0 +1,172 @@
+import orjson
+import pyarrow as pa
+import pytest
+from confluent_kafka import Producer
+from tkati_core.consumer import KafkaConsumer
+from tkati_core.settings import KafkaInputSettings
+
+
+def test_from_input_settings_sets_attributes(input_settings: KafkaInputSettings):
+    consumer = KafkaConsumer.from_input_settings(input_settings)
+    try:
+        assert consumer.topic_name == input_settings.topic.name
+        assert consumer.input_schema == input_settings.topic.schema
+    finally:
+        consumer.close()
+
+
+def test_from_input_settings_builds_arrow_schemas(input_settings: KafkaInputSettings):
+    consumer = KafkaConsumer.from_input_settings(input_settings)
+    try:
+        assert consumer.parse_schema.field("id").type == pa.string()
+        assert consumer.cast_schema.field("value").type == pa.int64()
+    finally:
+        consumer.close()
+
+
+def test_read_arrow_returns_none_on_empty_topic(
+    input_settings: KafkaInputSettings, kafka_input_topic: str
+):
+    consumer = KafkaConsumer.from_input_settings(input_settings)
+    try:
+        result = consumer.read_arrow(
+            aggregation_interval_seconds=2,
+            max_events_to_aggregate=10,
+        )
+        assert result is None
+    finally:
+        consumer.close()
+
+
+def test_read_arrow_reads_json_messages(
+    input_settings: KafkaInputSettings,
+    kafka_input_topic: str,
+    raw_producer: Producer,
+):
+    events = [
+        {"id": "a", "value": 1},
+        {"id": "b", "value": 2},
+        {"id": "c", "value": 3},
+    ]
+    for event in events:
+        raw_producer.produce(kafka_input_topic, value=orjson.dumps(event))
+    raw_producer.flush()
+
+    consumer = KafkaConsumer.from_input_settings(input_settings)
+    try:
+        table = consumer.read_arrow(
+            aggregation_interval_seconds=5,
+            max_events_to_aggregate=10,
+        )
+    finally:
+        consumer.close()
+
+    assert table is not None
+    assert len(table) == 3
+    assert table.schema.field("id").type == pa.string()
+    assert table.schema.field("value").type == pa.int64()
+    assert sorted(table.column("id").to_pylist()) == ["a", "b", "c"]  # type: ignore
+    assert sorted(table.column("value").to_pylist()) == [1, 2, 3]  # type: ignore
+
+
+def test_read_arrow_respects_max_events(
+    input_settings: KafkaInputSettings,
+    kafka_input_topic: str,
+    raw_producer: Producer,
+):
+    for i in range(10):
+        raw_producer.produce(kafka_input_topic, value=orjson.dumps({"id": str(i), "value": i}))
+    raw_producer.flush()
+
+    consumer = KafkaConsumer.from_input_settings(input_settings)
+    try:
+        table = consumer.read_arrow(
+            aggregation_interval_seconds=5,
+            max_events_to_aggregate=3,
+        )
+    finally:
+        consumer.close()
+
+    assert table is not None
+    assert len(table) == 3
+
+
+def test_read_arrow_timestamp_casting(
+    kafka_input_topic: str,
+    raw_producer: Producer,
+    run_id: str,
+):
+    """Verify timestamp[ms] fields are parsed as int64 and cast to pa.timestamp('ms')."""
+    from tkati_core.settings import (
+        KafkaConsumerSettings,
+        KafkaInputSettings,
+        KafkaTopicSettings,
+    )
+
+    settings = KafkaInputSettings(
+        topic=KafkaTopicSettings(
+            broker="localhost:9092",
+            name=kafka_input_topic,
+            schema={"ts": "timestamp[ms]"},
+        ),
+        consumer=KafkaConsumerSettings(
+            group_id=f"test-ts-{run_id}",
+            auto_offset_reset="earliest",
+        ),
+    )
+
+    raw_producer.produce(kafka_input_topic, value=orjson.dumps({"ts": 1_700_000_000_000}))
+    raw_producer.flush()
+
+    consumer = KafkaConsumer.from_input_settings(settings)
+    try:
+        table = consumer.read_arrow(aggregation_interval_seconds=5, max_events_to_aggregate=1)
+    finally:
+        consumer.close()
+
+    assert table is not None
+    assert table.schema.field("ts").type == pa.timestamp("ms")
+    assert table.column("ts")[0].as_py().timestamp() * 1000 == pytest.approx(1_700_000_000_000)
+
+
+def test_read_pylist_returns_none_on_empty_topic(
+    input_settings: KafkaInputSettings, kafka_input_topic: str
+):
+    consumer = KafkaConsumer.from_input_settings(input_settings)
+    try:
+        result = consumer.read_pylist(
+            aggregation_interval_seconds=2,
+            max_events_to_aggregate=10,
+        )
+        assert result is None
+    finally:
+        consumer.close()
+
+
+def test_read_pylist_reads_json_messages(
+    input_settings: KafkaInputSettings,
+    kafka_input_topic: str,
+    raw_producer: Producer,
+):
+    events = [
+        {"id": "a", "value": 1},
+        {"id": "b", "value": 2},
+        {"id": "c", "value": 3},
+    ]
+    for event in events:
+        raw_producer.produce(kafka_input_topic, value=orjson.dumps(event))
+    raw_producer.flush()
+
+    consumer = KafkaConsumer.from_input_settings(input_settings)
+    try:
+        rows = consumer.read_pylist(
+            aggregation_interval_seconds=5,
+            max_events_to_aggregate=10,
+        )
+    finally:
+        consumer.close()
+
+    assert rows is not None
+    assert len(rows) == 3
+    assert sorted(r["id"] for r in rows) == ["a", "b", "c"]
+    assert sorted(r["value"] for r in rows) == [1, 2, 3]

tkati_core-0.1.0/tests/test_producer.py

@@ -0,0 +1,198 @@
+import time
+
+import orjson
+import pyarrow as pa
+from confluent_kafka import Consumer
+from tkati_core.producer import KafkaProducer
+from tkati_core.settings import KafkaOutputSettings, KafkaTopicSettings
+
+
+def _consume_all(consumer: Consumer, topic: str, count: int, timeout: float = 10.0) -> list:
+    """Consume exactly `count` messages from `topic`, returning raw confluent Message objects."""
+    consumer.subscribe([topic])
+    messages = []  # type: ignore
+    deadline = time.time() + timeout
+    while len(messages) < count and time.time() < deadline:
+        batch = consumer.consume(num_messages=count - len(messages), timeout=1.0)
+        messages.extend(m for m in batch if not m.error())
+    return messages
+
+
+def test_from_output_settings_sets_attributes(output_settings: KafkaOutputSettings):
+    producer = KafkaProducer.from_output_settings(output_settings)
+    try:
+        assert producer.topic_name == output_settings.topic.name
+        assert producer.format == output_settings.topic.format
+        assert producer.key_column == output_settings.topic.key_column
+    finally:
+        producer.close()
+
+
+def test_produce_json_format(
+    output_settings: KafkaOutputSettings,
+    kafka_output_topic: str,
+    raw_consumer: Consumer,
+):
+    table = pa.table({"name": ["alice", "bob"], "score": [10, 20]})
+
+    producer = KafkaProducer.from_output_settings(output_settings)
+    try:
+        producer.produce_arrow(table)
+        producer.flush()
+    finally:
+        producer.close()
+
+    messages = _consume_all(raw_consumer, kafka_output_topic, count=2)
+    assert len(messages) == 2
+    parsed = [orjson.loads(m.value()) for m in messages]
+    names = {r["name"] for r in parsed}
+    scores = {r["score"] for r in parsed}
+    assert names == {"alice", "bob"}
+    assert scores == {10, 20}
+
+
+def test_produce_json_format_no_message_key_by_default(
+    output_settings: KafkaOutputSettings,
+    kafka_output_topic: str,
+    raw_consumer: Consumer,
+):
+    table = pa.table({"x": [1]})
+
+    producer = KafkaProducer.from_output_settings(output_settings)
+    try:
+        producer.produce_arrow(table)
+        producer.flush()
+    finally:
+        producer.close()
+
+    messages = _consume_all(raw_consumer, kafka_output_topic, count=1)
+    assert messages[0].key() is None
+
+
+def test_produce_json_format_with_key_column(
+    kafka_output_topic: str,
+    raw_consumer: Consumer,
+    run_id: str,
+):
+    settings = KafkaOutputSettings(
+        topic=KafkaTopicSettings(
+            broker="localhost:9092",
+            name=kafka_output_topic,
+            key_column="user_id",
+        )
+    )
+    table = pa.table({"user_id": ["u1", "u2"], "value": [100, 200]})
+
+    producer = KafkaProducer.from_output_settings(settings)
+    try:
+        producer.produce_arrow(table)
+        producer.flush()
+    finally:
+        producer.close()
+
+    messages = _consume_all(raw_consumer, kafka_output_topic, count=2)
+    assert len(messages) == 2
+    keys = {m.key().decode() for m in messages}
+    assert keys == {"u1", "u2"}
+
+
+def test_produce_arrow_batch_format(
+    kafka_output_topic: str,
+    raw_consumer: Consumer,
+    run_id: str,
+):
+    settings = KafkaOutputSettings(
+        topic=KafkaTopicSettings(
+            broker="localhost:9092",
+            name=kafka_output_topic,
+            format="arrow-batch",
+        )
+    )
+    original = pa.table({"id": ["x", "y", "z"], "n": [1, 2, 3]})
+
+    producer = KafkaProducer.from_output_settings(settings)
+    try:
+        producer.produce_arrow(original)
+        producer.flush()
+    finally:
+        producer.close()
+
+    messages = _consume_all(raw_consumer, kafka_output_topic, count=1)
+    assert len(messages) == 1
+
+    reader = pa.ipc.open_stream(messages[0].value())
+    recovered = reader.read_all()
+
+    assert recovered.schema == original.schema
+    assert recovered.equals(original)
+
+
+def test_produce_record_batch(
+    output_settings: KafkaOutputSettings,
+    kafka_output_topic: str,
+    raw_consumer: Consumer,
+):
+    """produce() accepts pa.RecordBatch as well as pa.Table."""
+    batch = pa.record_batch({"x": [7, 8]})
+
+    producer = KafkaProducer.from_output_settings(output_settings)
+    try:
+        producer.produce_arrow(batch)
+        producer.flush()
+    finally:
+        producer.close()
+
+    messages = _consume_all(raw_consumer, kafka_output_topic, count=2)
+    assert len(messages) == 2
+    values = {orjson.loads(m.value())["x"] for m in messages}
+    assert values == {7, 8}
+
+
+def test_produce_pylist_json_format(
+    output_settings: KafkaOutputSettings,
+    kafka_output_topic: str,
+    raw_consumer: Consumer,
+):
+    rows = [{"name": "alice", "score": 10}, {"name": "bob", "score": 20}]
+
+    producer = KafkaProducer.from_output_settings(output_settings)
+    try:
+        producer.produce_pylist(rows)
+        producer.flush()
+    finally:
+        producer.close()
+
+    messages = _consume_all(raw_consumer, kafka_output_topic, count=2)
+    assert len(messages) == 2
+    parsed = [orjson.loads(m.value()) for m in messages]
+    names = {r["name"] for r in parsed}
+    scores = {r["score"] for r in parsed}
+    assert names == {"alice", "bob"}
+    assert scores == {10, 20}
+
+
+def test_produce_pylist_with_key_column(
+    kafka_output_topic: str,
+    raw_consumer: Consumer,
+    run_id: str,
+):
+    settings = KafkaOutputSettings(
+        topic=KafkaTopicSettings(
+            broker="localhost:9092",
+            name=kafka_output_topic,
+            key_column="user_id",
+        )
+    )
+    rows = [{"user_id": "u1", "value": 100}, {"user_id": "u2", "value": 200}]
+
+    producer = KafkaProducer.from_output_settings(settings)
+    try:
+        producer.produce_pylist(rows)
+        producer.flush()
+    finally:
+        producer.close()
+
+    messages = _consume_all(raw_consumer, kafka_output_topic, count=2)
+    assert len(messages) == 2
+    keys = {m.key().decode() for m in messages}
+    assert keys == {"u1", "u2"}

tkati_core-0.1.0/tkati_core/consumer.py

@@ -0,0 +1,268 @@
+"""Kafka consumer utilities for reading messages into PyArrow tables."""
+
+import time
+from io import BytesIO
+from typing import TYPE_CHECKING
+
+import orjson
+import pyarrow as pa
+from confluent_kafka import Consumer
+from loguru import logger
+from pyarrow import json as pa_json
+
+if TYPE_CHECKING:
+    from tkati_core.settings import KafkaInputSettings
+
+
+class KafkaConsumer:
+    """
+    A Kafka consumer wrapper that reads messages into PyArrow tables or Python lists.
+
+    This class manages the Kafka consumer lifecycle, topic subscription,
+    and provides a convenient interface for reading messages as PyArrow tables
+    or plain Python dicts.
+    """
+
+    @classmethod
+    def from_input_settings(cls, settings: "KafkaInputSettings") -> "KafkaConsumer":
+        """
+        Construct a KafkaConsumer from a KafkaInputSettings instance.
+
+        Sets enable.auto.commit=False — offsets must be committed explicitly via .commit().
+        """
+        kafka_config: dict[str, str | bool] = {
+            "bootstrap.servers": settings.topic.broker,
+            "group.id": settings.consumer.group_id,
+            "auto.offset.reset": settings.consumer.auto_offset_reset,
+            "enable.auto.commit": False,
+        }
+        return cls(
+            kafka_config=kafka_config,
+            topic_name=settings.topic.name,
+            input_schema=settings.topic.schema,
+        )
+
+    def __init__(
+        self,
+        kafka_config: dict[str, str | bool],
+        topic_name: str,
+        input_schema: dict[str, str],
+    ) -> None:
+        """
+        Initialize the Kafka consumer with the provided configuration.
+
+        Args:
+            kafka_config: Dictionary of Kafka consumer configuration parameters.
+                         Common keys include:
+                         - 'bootstrap.servers': Kafka broker addresses
+                         - 'group.id': Consumer group ID
+                         - 'auto.offset.reset': Offset reset behavior
+                         - 'enable.auto.commit': Whether to auto-commit offsets
+        """
+        self.consumer = Consumer(kafka_config)
+        self.topic_name = topic_name
+        self.input_schema = input_schema
+
+        self.consumer.subscribe([self.topic_name])
+
+        # Create PyArrow schema based on input_schema
+        # type -> (parse_type, cast_type)
+        type_mapping: dict[str, tuple[pa.DataType, pa.DataType]] = {
+            "string": (pa.string(), pa.string()),
+            "int32": (pa.int32(), pa.int32()),
+            "int64": (pa.int64(), pa.int64()),
+            "uint32": (pa.uint32(), pa.uint32()),
+            "uint64": (pa.uint64(), pa.uint64()),
+            "uint8": (pa.uint8(), pa.uint8()),
+            "int": (pa.int32(), pa.int32()),
+            "timestamp[ms]": (pa.int64(), pa.timestamp("ms")),
+        }
+
+        parse_schema_fields = []
+        cast_schema_fields = []
+
+        for field_name, field_type in input_schema.items():
+            types = type_mapping.get(field_type)
+
+            if types is not None:
+                parse_type, cast_type = types
+            else:
+                logger.warning(
+                    f"Unsupported field type '{field_type}' for field '{field_name}'. Defaulting to string."
+                )
+                parse_type, cast_type = (pa.string(), pa.string())
+
+            parse_schema_fields.append(pa.field(field_name, parse_type))
+            cast_schema_fields.append(pa.field(field_name, cast_type))
+
+        self.parse_schema = pa.schema(parse_schema_fields)
+        self.cast_schema = pa.schema(cast_schema_fields)
+
+        logger.info(
+            f"Initialized KafkaConsumer with config: {kafka_config} and topic: {topic_name}"
+        )
+
+    def _consume_batch(
+        self,
+        aggregation_interval_seconds: int,
+        max_events_to_aggregate: int,
+    ) -> tuple[list, int]:
+        """
+        Consume raw messages from Kafka within the given time and count limits.
+
+        Returns a tuple of (messages, events_read) where messages is a list of
+        confluent_kafka Message objects (without errors).
+        """
+        if self.topic_name:
+            logger.info(
+                f"Consuming events from topic(s): {self.topic_name} for up to {aggregation_interval_seconds}s or {max_events_to_aggregate} events"
+            )
+        else:
+            logger.info(
+                f"Consuming events for up to {aggregation_interval_seconds}s or {max_events_to_aggregate} events"
+            )
+
+        start_time = time.time()
+        events_read = 0
+        poll_timeout = 10
+        valid_messages = []
+
+        while events_read < max_events_to_aggregate:
+            elapsed = time.time() - start_time
+            remaining_time = aggregation_interval_seconds - elapsed
+
+            if remaining_time <= 0:
+                logger.info(f"Reached time limit of {aggregation_interval_seconds}s")
+                break
+
+            remaining_messages = max_events_to_aggregate - events_read
+            batch_timeout = min(poll_timeout, remaining_time)
+            messages = self.consumer.consume(
+                num_messages=min(remaining_messages, 1_000_000),
+                timeout=batch_timeout,
+            )
+
+            if not messages:
+                continue
+
+            for msg in messages:
+                if msg.error():
+                    logger.info(f"Consumer error: {msg.error()}")
+                    continue
+                valid_messages.append(msg)
+                events_read += 1
+
+        elapsed_total = time.time() - start_time
+        logger.info(f"Consumed {events_read} events in {elapsed_total:.2f}s")
+        return valid_messages, events_read
+
+    # WARNING: This function breaks if any single message is malformed JSON. We may
+    # want to enhance it to handle individual message errors more gracefully.
+    def read_arrow(
+        self,
+        aggregation_interval_seconds: int,
+        max_events_to_aggregate: int,
+    ) -> pa.Table | None:
+        """
+        Read messages from subscribed topics into a PyArrow table.
+
+        Args:
+            aggregation_interval_seconds: Maximum time in seconds to consume messages.
+            max_events_to_aggregate: Maximum number of events to consume.
+
+        Returns:
+            A PyArrow Table containing the parsed events, or None if no data was consumed.
+
+        Notes:
+            - Does NOT commit offsets. The caller is responsible for managing consumer lifecycle.
+            - Does NOT subscribe to topics. The consumer must be pre-subscribed.
+            - Raises exceptions on JSON parsing errors.
+            - Uses permissive parsing that ignores unexpected fields in JSON messages.
+        """
+        valid_messages, events_read = self._consume_batch(
+            aggregation_interval_seconds, max_events_to_aggregate
+        )
+
+        if events_read == 0:
+            logger.info("No data consumed from topic.")
+            return None
+
+        buffer = BytesIO()
+        for msg in valid_messages:
+            buffer.write(msg.value())
+            buffer.write(b"\n")
+        buffer.seek(0)
+
+        parse_options = pa_json.ParseOptions(
+            explicit_schema=self.parse_schema,
+            unexpected_field_behavior="ignore",
+        )
+
+        try:
+            table = pa_json.read_json(buffer, parse_options=parse_options)
+            table = table.cast(self.cast_schema)
+            actual_rows = len(table)
+
+            if actual_rows != events_read:
+                logger.warning(
+                    f"Row count mismatch: consumed {events_read} messages, but parsed {actual_rows} rows. {events_read - actual_rows} messages may have been skipped."
+                )
+            else:
+                logger.info(
+                    f"Successfully parsed {actual_rows} rows matching {events_read} consumed messages"
+                )
+
+        except Exception as e:
+            logger.error(f"Failed to parse JSON with PyArrow: {e}")
+            raise
+
+        return table
+
+    def read_pylist(
+        self,
+        aggregation_interval_seconds: int,
+        max_events_to_aggregate: int,
+    ) -> list[dict] | None:
+        """
+        Read messages from subscribed topics into a list of dicts.
+
+        Same batching semantics as read_arrow (time + count limits).
+        Messages that fail JSON parsing are skipped (logged as errors).
+
+        Returns:
+            A list of parsed event dicts, or None if no data was consumed.
+
+        Notes:
+            - Does NOT commit offsets. The caller is responsible for managing consumer lifecycle.
+        """
+        valid_messages, events_read = self._consume_batch(
+            aggregation_interval_seconds, max_events_to_aggregate
+        )
+
+        if events_read == 0:
+            logger.info("No data consumed from topic.")
+            return None
+
+        rows = []
+        for msg in valid_messages:
+            try:
+                rows.append(orjson.loads(msg.value()))
+            except Exception as e:
+                logger.error(f"Error parsing message from topic {msg.topic()}: {e}")
+
+        logger.info(f"Successfully parsed {len(rows)} rows")
+        return rows if rows else None
+
+    def commit(self) -> None:
+        """
+        Commit the current offsets for all subscribed topics.
+        """
+        self.consumer.commit()
+        logger.debug("Committed offsets")
+
+    def close(self) -> None:
+        """
+        Close the Kafka consumer and release resources.
+        """
+        self.consumer.close()
+        logger.info("Closed KafkaConsumer")

tkati_core-0.1.0/tkati_core/producer.py

@@ -0,0 +1,105 @@
+"""Kafka producer utilities for writing PyArrow tables as messages."""
+
+from typing import TYPE_CHECKING, Literal
+
+import orjson
+import pyarrow as pa
+from confluent_kafka import Producer
+from loguru import logger
+
+if TYPE_CHECKING:
+    from tkati_core.settings import KafkaOutputSettings
+
+
+class KafkaProducer:
+    """
+    A Kafka producer wrapper that writes data as messages.
+
+    Supports producing from PyArrow tables/batches or plain Python dicts.
+
+    For Arrow-based production, two serialization formats are controlled by the
+    topic's ``format`` setting:
+    - ``"json"``: produces one Kafka message per row, serialized with orjson.
+    - ``"arrow-batch"``: produces the entire table as a single Arrow IPC message.
+
+    The optional ``key_column`` setting (from ``KafkaTopicSettings``) names the
+    column whose value is used as the Kafka message key for each row (JSON format only).
+    """
+
+    def __init__(
+        self,
+        kafka_config: dict[str, str],
+        topic_name: str,
+        format: Literal["json", "arrow-batch"] = "json",
+        key_column: str | None = None,
+    ) -> None:
+        self.producer = Producer(kafka_config)
+        self.topic_name = topic_name
+        self.format = format
+        self.key_column = key_column
+        logger.info(
+            f"Initialized KafkaProducer with topic: {topic_name}, format: {format}"
+        )
+
+    @classmethod
+    def from_output_settings(cls, settings: "KafkaOutputSettings") -> "KafkaProducer":
+        """
+        Construct a KafkaProducer from a KafkaOutputSettings instance.
+        """
+        return cls(
+            kafka_config={"bootstrap.servers": settings.topic.broker},
+            topic_name=settings.topic.name,
+            format=settings.topic.format,
+            key_column=settings.topic.key_column,
+        )
+
+    def produce_arrow(self, data: pa.Table | pa.RecordBatch) -> None:
+        """
+        Produce data to the configured topic.
+
+        For ``"json"`` format each row becomes a separate Kafka message serialized
+        with orjson. If ``key_column`` is set, its value is used as the message key.
+
+        For ``"arrow-batch"`` format the entire table is serialized as a single
+        Arrow IPC stream message.
+        """
+        if self.format == "json":
+            self.produce_pylist(data.to_pylist())
+        elif self.format == "arrow-batch":
+            table = (
+                data if isinstance(data, pa.Table) else pa.Table.from_batches([data])
+            )
+            buf = pa.BufferOutputStream()
+            with pa.ipc.new_stream(buf, table.schema) as writer:
+                for batch in table.to_batches():
+                    writer.write_batch(batch)
+            self.producer.produce(self.topic_name, value=buf.getvalue().to_pybytes())
+
+    def produce_pylist(self, rows: list[dict]) -> None:
+        """
+        Produce a list of dicts to the configured topic as JSON messages.
+
+        Each dict becomes a separate Kafka message serialized with orjson.
+        If ``key_column`` is set, its value is used as the Kafka message key.
+        """
+        for row in rows:
+            key = (
+                str(row[self.key_column])
+                if self.key_column and self.key_column in row
+                else None
+            )
+            self.producer.produce(self.topic_name, value=orjson.dumps(row), key=key)
+
+    def flush(self) -> None:
+        """
+        Block until all queued messages have been delivered.
+        """
+        self.producer.flush()
+        logger.debug("Flushed KafkaProducer")
+
+    def close(self) -> None:
+        """
+        Flush pending messages and release resources.
+        """
+        self.producer.flush()
+        logger.info("Closed KafkaProducer")

tkati_core-0.1.0/tkati_core/settings.py

@@ -0,0 +1,66 @@
+import os
+from typing import Literal
+
+from loguru import logger
+from pydantic import BaseModel, Field
+from pydantic_settings import (
+    BaseSettings,
+    PydanticBaseSettingsSource,
+    SettingsConfigDict,
+    TomlConfigSettingsSource,
+)
+
+SETTINGS_FILE = os.getenv("SETTINGS_FILE", "settings.toml")
+
+
+class KafkaTopicSettings(BaseModel):
+    broker: str
+    name: str
+    schema: dict[str, str] = Field(default_factory=dict)  # type: ignore
+    format: Literal["json", "arrow-batch"] = "json"
+    key_column: str | None = None
+
+
+class KafkaConsumerSettings(BaseModel):
+    group_id: str
+    batch_size: int = 1000
+    batch_timeout_sec: int = 5
+    auto_offset_reset: str = "latest"
+
+
+class KafkaInputSettings(BaseModel):
+    topic: KafkaTopicSettings
+    consumer: KafkaConsumerSettings
+
+
+class KafkaOutputSettings(BaseModel):
+    topic: KafkaTopicSettings
+
+
+logger.info(f"Using settings file: {os.path.abspath(SETTINGS_FILE)}")
+
+
+class TomlBaseSettings(BaseSettings):
+    model_config = SettingsConfigDict(
+        toml_file=SETTINGS_FILE,
+        env_file=".env",
+        extra="ignore",
+        env_nested_delimiter="__",
+    )
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        return (
+            init_settings,
+            env_settings,
+            dotenv_settings,
+            file_secret_settings,
+            TomlConfigSettingsSource(settings_cls),
+        )

tkati_core-0.1.0/tkati_core/testing.py

@@ -0,0 +1,11 @@
+from collections.abc import Generator
+
+import pytest
+from confluent_kafka.admin import AdminClient
+
+
+@pytest.fixture(scope="function")
+def kafka_admin_client() -> Generator[AdminClient, None, None]:
+    """Provides a Kafka AdminClient for integration tests. Requires Redpanda on localhost:9092."""
+    admin = AdminClient({"bootstrap.servers": "localhost:9092"})
+    yield admin

tkati_core-0.1.0/tkati_core.egg-info/PKG-INFO

@@ -0,0 +1,118 @@
+Metadata-Version: 2.4
+Name: tkati-core
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+Requires-Dist: confluent-kafka>=2.11.0
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: orjson>=3.9.0
+Requires-Dist: pyarrow>=21.0.0
+Requires-Dist: pydantic-settings>=2.11.0
+
+# tkati-core
+
+For now we assume that each node gets one input and produces one output in a
+form of kafka stream.
+
+## Settings
+
+General form of settings is:
+
+```toml
+[input.topic]
+# definition of input stream:
+# - broker
+# - topic name
+# - message schema
+# - message format = "json" / "arrow-batch"
+
+[input.consumer]
+# parameters local to this consumer
+# - group_id
+# - batch_size
+# - batch_timeout_sec
+# - auto_offset_reset
+
+[output.topic]
+# definition of output stream
+# - broker
+# - topic name
+# - message schema
+# - message format = "json" / "arrow-batch"
+# - key_column (optional) = column to use as the Kafka message key
+
+[...]
+# settings specific to node function
+```
+
+## Usage
+
+### Constructing a consumer from settings
+
+Use `KafkaArrowConsumer.from_input_settings` to construct a consumer directly from
+`KafkaInputSettings` — no need to manually map fields to Confluent Kafka config keys.
+
+```python
+from tkati_core.settings import TomlBaseSettings, KafkaInputSettings
+from tkati_core.consumer import KafkaArrowConsumer
+
+class AppSettings(TomlBaseSettings):
+    input: KafkaInputSettings
+    # ...
+
+settings = AppSettings()
+consumer = KafkaArrowConsumer.from_input_settings(settings.input)
+
+# Read a batch
+table = consumer.read_to_pyarrow(
+    aggregation_interval_seconds=settings.input.consumer.batch_timeout_sec,
+    max_events_to_aggregate=settings.input.consumer.batch_size,
+)
+consumer.commit()
+```
+
+The factory method sets `enable.auto.commit=False` — offsets must be committed explicitly
+via `consumer.commit()`.
+
+### Constructing a producer from settings
+
+Use `KafkaArrowProducer.from_output_settings` to construct a producer directly from
+`KafkaOutputSettings`. It accepts PyArrow tables or record batches and handles
+serialization according to the topic's `format` setting.
+
+```python
+from tkati_core.settings import TomlBaseSettings, KafkaOutputSettings
+from tkati_core.producer import KafkaArrowProducer
+
+class AppSettings(TomlBaseSettings):
+    output: KafkaOutputSettings
+    # ...
+
+settings = AppSettings()
+producer = KafkaArrowProducer.from_output_settings(settings.output)
+
+# Produce a PyArrow table (one message per row for "json" format)
+producer.produce(table)
+producer.flush()
+producer.close()  # flushes and releases resources
+```
+
+**Formats** — controlled by `output.topic.format` in `settings.toml`:
+
+- `"json"` *(default)*: each row becomes a separate Kafka message serialized with orjson.
+- `"arrow-batch"`: the entire table is serialized as a single Arrow IPC stream message.
+
+**Message keys** — controlled by `output.topic.key_column` in `settings.toml`:
+
+```toml
+[output.topic]
+broker = "localhost:9092"
+name = "my-output-topic"
+key_column = "customer_id"   # column whose value becomes the Kafka message key
+```
+
+`key_column` is optional. When omitted (or `None`), messages are produced without a key.
+When set, the value of that column for each row is used as the Kafka message key
+(JSON format only — ignored for `"arrow-batch"`). This determines which Kafka partition
+each message is routed to.

tkati_core-0.1.0/tkati_core.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+README.md
+pyproject.toml
+tests/test_consumer.py
+tests/test_producer.py
+tkati_core/consumer.py
+tkati_core/producer.py
+tkati_core/py.typed
+tkati_core/settings.py
+tkati_core/testing.py
+tkati_core.egg-info/PKG-INFO
+tkati_core.egg-info/SOURCES.txt
+tkati_core.egg-info/dependency_links.txt
+tkati_core.egg-info/requires.txt
+tkati_core.egg-info/top_level.txt

tkati_core-0.1.0/tkati_core.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

tkati_core-0.1.0/tkati_core.egg-info/requires.txt

@@ -0,0 +1,5 @@
+confluent-kafka>=2.11.0
+loguru>=0.7.3
+orjson>=3.9.0
+pyarrow>=21.0.0
+pydantic-settings>=2.11.0

tkati_core-0.1.0/tkati_core.egg-info/top_level.txt

@@ -0,0 +1 @@
+tkati_core