feat: Kafka as a first-class batch source (JSON → Delta Lake via Unity Catalog)

[malon64]

Context

Add Kafka as a first-class batch source for floe run, without changing Floe's batch architecture. Each run processes a bounded snapshot of Kafka messages (up to max_messages), writes them to Delta Lake, and saves per-partition offsets so the next run continues from where it left off.

Event-triggered scheduling is handled outside Floe by a Dagster sensor that checks consumer group lag and yields a RunRequest only when lag is detected. This keeps the Floe binary simple and single-purpose.

Message format for Phase 1: JSON only (Avro deferred — apache-avro is already in deps).

---

Architecture: Materialization Strategy

Kafka messages are consumed and written to a temp NDJSON file inside resolve_local_inputs. The rest of the pipeline (schema validation, type casting, rejected-row routing, Delta write) is inherited unchanged. To the engine, "one Kafka batch = one NDJSON file."

• source.path carries the broker+topic URL: kafka://broker1:9092,broker2:9092/topic.name
• Kafka-specific config lives in a new KafkaSourceOptions struct (not SourceOptions — those are CSV/JSON-specific and validated with validate_known_keys)
• Offset tracking via new IncrementalMode::Offset backed by a KafkaOffsetState JSON file, analogous to how IncrementalMode::File tracks file URIs

---

YAML Config Example

version: "1"

catalogs:
  definitions:
    - name: unity_prod
      type: unity
      host: https://my-workspace.azuredatabricks.net
      catalog: my_catalog
      schema: raw_events
      token: ${DATABRICKS_TOKEN}
      create_schema_if_missing: true

entities:
  - name: kafka_orders
    incremental_mode: offset          # new variant; tracks per-partition Kafka offsets
    state:
      path: .floe/state/kafka_orders/state.json

    source:
      format: kafka
      path: kafka://broker1:9092,broker2:9092/orders
      options:
        group_id: floe-orders-consumer
        security_protocol: SASL_SSL
        sasl_mechanism: SCRAM-SHA-256
        sasl_username: ${KAFKA_USERNAME}
        sasl_password: ${KAFKA_PASSWORD}
        max_messages: 50000           # hard cap per run
        timeout_ms: 30000             # consumer poll timeout
        max_wait_ms: 5000             # stop early if no new messages
        message_format: json          # json only in Phase 1
        include_metadata: true        # inject _kafka_offset, _kafka_partition, etc.

    sink:
      write_mode: append
      accepted:
        format: delta
        path: s3://my-bucket/delta/orders
        storage: my_s3_storage
        delta:
          catalog: unity_prod
          table: orders

    policy:
      severity: reject

    schema:
      columns:
        - name: order_id
          type: string
          nullable: false
        - name: amount
          type: float64
        - name: _kafka_offset
          type: int64
        - name: _kafka_partition
          type: int32
        - name: _kafka_timestamp
          type: int64
        - name: _kafka_topic
          type: string

---

Scheduling: Dagster Sensor

A small Python sensor triggers floe run when Kafka lag is detected:

from dagster import sensor, RunRequest, SkipReason
from confluent_kafka.admin import AdminClient
from confluent_kafka import Consumer, TopicPartition

@sensor(job=floe_kafka_job, minimum_interval_seconds=30)
def kafka_lag_sensor(context):
    admin = AdminClient({"bootstrap.servers": "broker1:9092"})
    consumer = Consumer({"bootstrap.servers": "broker1:9092", "group.id": "floe-orders-consumer"})
    topic = "orders"

    metadata = admin.list_topics(topic)
    partitions = list(metadata.topics[topic].partitions.keys())
    committed = consumer.committed([TopicPartition(topic, p) for p in partitions])
    watermarks = {p: consumer.get_watermark_offsets(TopicPartition(topic, p)) for p in partitions}
    consumer.close()

    lag = sum(
        max(0, watermarks[c.partition][1] - (c.offset if c.offset >= 0 else 0))
        for c in committed
    )
    if lag > 0:
        yield RunRequest(run_key=str(context.cursor or 0), run_config={})
    else:
        yield SkipReason(f"No lag on {topic}")

@op
def run_floe(context):
    subprocess.run(["floe", "run", "--config", "config.yml", "--entity", "kafka_orders"], check=True)

This sensor file lives in examples/dagster/kafka_sensor.py (not compiled into the Floe binary).

---

Files to Create

File	Purpose
crates/floe-core/src/io/read/kafka.rs	KafkaInputAdapter implementation
examples/dagster/kafka_sensor.py	Dagster sensor + op that triggers floe run on lag

Files to Modify

File	Change
crates/floe-core/Cargo.toml	Add rdkafka = { version = "0.37", features = ["cmake-build","ssl","sasl"], optional = true } + kafka feature
crates/floe-core/src/config/types.rs	Add KafkaSourceOptions struct; add kafka: Option<KafkaSourceOptions> to SourceConfig; add IncrementalMode::Offset
crates/floe-core/src/config/parse.rs	Add parse_kafka_source_options(); patch parse_source() for format="kafka"; add "offset" to parse_incremental_mode()
crates/floe-core/src/config/validate.rs	Validate kafka:// path prefix; constrain incremental_mode to offset or none for kafka
crates/floe-core/src/state/mod.rs	Add KafkaOffsetState, load_kafka_offset_state(), save_kafka_offset_state()
crates/floe-core/src/io/read/mod.rs	Add #[cfg(feature = "kafka")] pub mod kafka;
crates/floe-core/src/io/read/json.rs	Expose read_ndjson_columns as pub(crate) so kafka adapter can delegate
crates/floe-core/src/io/format.rs	Add "kafka" arm in input_adapter() (feature-gated)
crates/floe-core/src/io/storage/core/extensions.rs	Add "kafka" => Ok(vec![]) to suffixes_for_format()
crates/floe-cli/Cargo.toml	Enable kafka feature for floe-core

---

Key Implementation Details

KafkaInputAdapter (crates/floe-core/src/io/read/kafka.rs)

struct KafkaInputAdapter;
static KAFKA_INPUT_ADAPTER: KafkaInputAdapter = KafkaInputAdapter;
pub(crate) fn kafka_input_adapter() -> &'static dyn InputAdapter { &KAFKA_INPUT_ADAPTER }

resolve_local_inputs() — all Kafka I/O lives here:

1. Parse source.path: kafka://brokers/topic
2. Load KafkaOffsetState from disk (per-partition committed offsets)
3. Build rdkafka::consumer::BaseConsumer (synchronous — not StreamConsumer) from KafkaSourceOptions
4. Assign per-partition start offsets via TopicPartitionList
5. Poll loop: collect up to max_messages, stop early if max_wait_ms expires with no new messages
6. If include_metadata: true, inject _kafka_offset, _kafka_partition, _kafka_timestamp, _kafka_topic into each JSON object
7. Write all messages as NDJSON to a temp file (std::env::temp_dir()/floe-kafka/{entity}-{uuid}.ndjson)
8. Commit partition → next_offset to KafkaOffsetState (save to disk)
9. Return ResolvedLocalInputs { files: vec![temp_path], mode: LocalInputMode::File }

read_inputs() — use LazyJsonLineReader::new(temp_path).finish()?.collect()? to build the DataFrame. Uniqueness checks are automatically skipped when no unique columns are declared in the schema.

Critical notes:

• Use BaseConsumer (sync) — resolve_local_inputs is called from sync Rust
• If polling returns zero messages, return files: vec![] — the engine handles this as a no-op
• Polars with_streaming(true) is not applicable: every adapter collects to DataFrame immediately; this is the most memory-efficient approach within the current design

KafkaOffsetState (crates/floe-core/src/state/mod.rs)

pub struct KafkaOffsetState {
    pub schema: String,              // "floe.state.kafka-offset.v1"
    pub entity: String,
    pub updated_at: Option<String>,
    pub offsets: BTreeMap<String, i64>,  // "0" -> next_offset_to_consume
}

Load/save via atomic JSON write (same pattern as EntityState in the same file).

KafkaSourceOptions (crates/floe-core/src/config/types.rs)

pub struct KafkaSourceOptions {
    pub group_id: Option<String>,
    pub security_protocol: Option<String>,  // PLAINTEXT | SSL | SASL_SSL | SASL_PLAINTEXT
    pub sasl_mechanism: Option<String>,     // PLAIN | SCRAM-SHA-256 | SCRAM-SHA-512
    pub sasl_username: Option<String>,
    pub sasl_password: Option<String>,
    pub ssl_ca_location: Option<String>,
    pub max_messages: Option<u64>,          // default: 10_000
    pub timeout_ms: Option<u64>,            // default: 30_000
    pub max_wait_ms: Option<u64>,           // default: 5_000
    pub message_format: Option<String>,     // "json" only in Phase 1
    pub include_metadata: Option<bool>,     // default: false
}

In parse_source(), when format == "kafka", parse source.options YAML map into KafkaSourceOptions (not SourceOptions), set options: None, set kafka: Some(...).

---

Implementation Order

1. crates/floe-core/Cargo.toml — add rdkafka optional dep + kafka feature
2. config/types.rs — add KafkaSourceOptions, extend SourceConfig, add IncrementalMode::Offset
3. state/mod.rs — add KafkaOffsetState + load/save functions
4. config/parse.rs — add parse_kafka_source_options(), patch parse_source(), add "offset" mode
5. config/validate.rs — kafka path prefix guard + incremental mode constraint
6. io/storage/core/extensions.rs — add "kafka" => Ok(vec![]) to suffixes_for_format
7. io/read/json.rs — make NDJSON column reader pub(crate)
8. io/read/kafka.rs — implement KafkaInputAdapter
9. io/read/mod.rs — add #[cfg(feature = "kafka")] pub mod kafka;
10. io/format.rs — add "kafka" match arm (feature-gated)
11. examples/dagster/kafka_sensor.py — Dagster sensor + op example

---

Verification

• cargo check --features kafka — no compile errors
• cargo test --features kafka — all existing tests pass
• Unit test: parse a format: kafka YAML fixture, assert source.format == "kafka" and source.kafka.is_some()
• Unit test: KafkaOffsetState JSON round-trip
• Integration test: pre-built NDJSON temp file → kafka_input_adapter().read_inputs() → assert correct DataFrame
• End-to-end (local): Redpanda (Docker) + test topic → floe run → assert Delta table written

---

Risks & Notes

• Offset commit timing: Offsets are committed after polling but before the sink write (at-most-once). At-least-once can be added later by committing only after a confirmed sink write.
• rdkafka cmake-build feature vendors librdkafka — longer compile but no system library dependency. Correct for Docker/Kubernetes.
• Empty batch: Zero messages → files: vec![] → engine no-op (verify against MAX_RESOLVED_INPUTS guard in run/mod.rs).
• Avro support deferred — apache-avro is already in deps.