This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 7a28e10b0b7 [FLINK-40198][python] Add Kafka source and sink APIs to
DataFrame API
7a28e10b0b7 is described below
commit 7a28e10b0b71d2efe41b9a407c837e3cfffd0b70
Author: niliushall <[email protected]>
AuthorDate: Thu Oct 8 16:21:18 2026 +0800
[FLINK-40198][python] Add Kafka source and sink APIs to DataFrame API
This closes #29439.
---
.../docs/reference/pyflink.dataframe/io.rst | 55 ++
flink-python/pyflink/dataframe/__init__.py | 9 +-
flink-python/pyflink/dataframe/dataframe.py | 115 ++++
flink-python/pyflink/dataframe/io.py | 534 +++++++++++++++-
.../pyflink/dataframe/tests/test_kafka_io.py | 677 +++++++++++++++++++++
5 files changed, 1387 insertions(+), 3 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/io.rst
b/flink-python/docs/reference/pyflink.dataframe/io.rst
index eaa059d4e3c..c1dd0172b20 100644
--- a/flink-python/docs/reference/pyflink.dataframe/io.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/io.rst
@@ -35,6 +35,7 @@ Readers
read_catalog_table
read_generic
read_json
+ read_kafka
read_parquet
Writers
@@ -48,6 +49,7 @@ Writers
DataFrame.write_catalog_table
DataFrame.write_generic
DataFrame.write_json
+ DataFrame.write_kafka
DataFrame.write_parquet
Writing to multiple sinks
@@ -80,6 +82,59 @@ The sinks share job backpressure and failover; this does not
guarantee cross-sin
Executing the statement set clears its staged writes. Client actions such as
``collect()``
and ``to_pandas()`` execute separately and are not added to the statement set.
+Kafka sources and sinks
+-----------------------
+
+``read_kafka`` and ``write_kafka`` use Flink's Kafka SQL connector. Make the
Kafka
+connector and the selected format JARs available to Flink before using them.
Readers
+require a schema and either ``topic`` (a name or list of names) or
``topic_pattern``.
+Use ``group_id`` to identify the consumer group. The default startup mode is
+``"group-offsets"``; ``startup_mode="earliest-offset"`` starts at the earliest
+available offset instead. Sources are unbounded by default; set
``bounded_mode``
+to stop reading at a selected offset or timestamp. Topic lists and
``key_fields``
+are converted to semicolon-separated connector options. Specific offsets can be
+supplied as ``{0: 42, 1: 300}`` or as the connector string
+``"partition:0,offset:42;partition:1,offset:300"`` for both startup and
bounded modes.
+Set ``topic_partition_discovery_interval=None`` to disable partition discovery;
+this explicitly sets ``scan.topic-partition-discovery.interval`` to ``"0"``.
+
+For example, read JSON records from one topic and write them to another::
+
+ import pyflink.dataframe as pf
+
+ events = pf.read_kafka(
+ "localhost:9092",
+ schema={"id": pf.DataType.int64(), "event": pf.DataType.string()},
+ topic="input-events",
+ group_id="dataframe-events",
+ startup_mode="earliest-offset",
+ )
+ events.write_kafka("localhost:9092", topic="output-events")
+
+Both methods default to JSON values. ``format`` and ``value_format`` select the
+value format; ``value_format`` takes precedence when both are provided. Pass
+format-specific settings through ``format_options`` or
``value_format_options``;
+use ``key_format`` and ``key_fields`` to configure keys. Additional Kafka
client
+settings go in ``properties``. Writers default to ``"at-least-once"`` delivery;
+``delivery_guarantee`` also supports ``"none"`` and ``"exactly-once"``.
+Exactly-once delivery requires checkpointing and ``transactional_id_prefix``
+(or ``sink.transactional-id-prefix`` in ``options``). Use a unique prefix for
each
+concurrently running Kafka sink. Use ``partitioner`` to select ``"default"``,
+``"fixed"``, ``"round-robin"``, or a custom ``FlinkKafkaPartitioner`` class
name;
+``"round-robin"`` requires records without keys. Set ``parallelism`` to
+configure the Kafka writer's parallelism.
+
+Both readers and writers accept additional raw connector settings through
+``options``, for example ``options={"scan.parallelism": "3"}`` for a reader or
+``options={"sink.parallelism": "3"}`` for a writer. Unlike ``properties``,
which
+adds the ``properties.`` prefix to Kafka client settings, ``options`` keys are
+passed through unchanged. The ``connector`` key is reserved. Conflicts with
options
+generated by other arguments, including default values, raise ``ValueError``;
+use the dedicated parameter to change such a setting.
+
+As with the other writers, pass ``statement_set=`` to stage a Kafka write
alongside
+other sinks instead of submitting it immediately.
+
Filesystem sources and sinks
----------------------------
diff --git a/flink-python/pyflink/dataframe/__init__.py
b/flink-python/pyflink/dataframe/__init__.py
index 0e67c1a9e54..e463bb19d0c 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -65,7 +65,13 @@ from pyflink.dataframe.context import (
from pyflink.dataframe.dataframe import DataFrame, GroupedDataFrame, col, lit
from pyflink.dataframe.dataframe_config import config
from pyflink.dataframe.datatype import DataType
-from pyflink.dataframe.io import read_catalog_table, read_generic, read_json,
read_parquet
+from pyflink.dataframe.io import (
+ read_catalog_table,
+ read_generic,
+ read_json,
+ read_kafka,
+ read_parquet,
+)
from pyflink.dataframe.iteration import CloseableIterator
from pyflink.dataframe.sql import sql
from pyflink.dataframe.udf import udf
@@ -89,6 +95,7 @@ __all__ = [
"read_catalog_table",
"read_generic",
"read_json",
+ "read_kafka",
"read_parquet",
"create_catalog",
"get_catalog",
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index 3404a4d1861..66152fb4ce9 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -2703,6 +2703,121 @@ class DataFrame:
statement_set=statement_set,
)
+ @PublicEvolving()
+ def write_kafka(
+ self,
+ bootstrap_servers: str,
+ *,
+ topic: str,
+ format: str = "json",
+ format_options: Optional[Dict[str, str]] = None,
+ key_format: Optional[str] = None,
+ key_format_options: Optional[Dict[str, str]] = None,
+ key_fields: Optional[List[str]] = None,
+ value_format: Optional[str] = None,
+ value_format_options: Optional[Dict[str, str]] = None,
+ value_fields_include: Literal["ALL", "EXCEPT_KEY"] = "ALL",
+ delivery_guarantee: Literal[
+ "none", "at-least-once", "exactly-once"
+ ] = "at-least-once",
+ transactional_id_prefix: Optional[str] = None,
+ partitioner: Optional[str] = None,
+ properties: Optional[Dict[str, str]] = None,
+ parallelism: Optional[int] = None,
+ options: Optional[Dict[str, str]] = None,
+ statement_set: Optional[StatementSet] = None,
+ ) -> None:
+ """
+ Write this DataFrame to Kafka using Flink's Kafka SQL connector.
+
+ The Kafka connector must be available to Flink. By default, the write
is
+ submitted immediately and waits for completion for local or MiniCluster
+ execution. Passing ``statement_set`` stages the write without
executing it.
+ Sink columns are derived from this DataFrame's schema.
+
+ ``format`` and ``value_format`` are aliases; ``value_format`` takes
+ precedence when both are set. ``format_options`` and
``value_format_options`` are both
+ normalized to ``value.<format>.<key>``; conflicting values are
rejected.
+
+ :param bootstrap_servers: Comma-separated Kafka bootstrap server
+ addresses.
+ :param topic: Kafka topic to write to.
+ :param format: Value format, for example ``"json"``, ``"csv"``,
+ ``"avro"``.
+ :param format_options: Value format options with string values, with
+ or without the ``<format>.`` prefix.
+ :param key_format: Key format, for example ``"json"`` or ``"avro"``.
+ :param key_format_options: Key format options with string values, with
+ or without the ``key.<key_format>.`` prefix.
+ :param key_fields: List of column names that make up the key.
+ :param value_format: Alternative to ``format`` for specifying the
+ value format.
+ :param value_format_options: Alternative to ``format_options`` for
+ specifying value format options.
+ :param value_fields_include: Whether value fields include key fields:
+ ``"ALL"`` or ``"EXCEPT_KEY"``.
+ :param delivery_guarantee: Delivery guarantee: ``"none"``,
+ ``"at-least-once"``, or ``"exactly-once"``. Exactly-once requires
+ checkpointing and a transaction id prefix.
+ :param transactional_id_prefix: Prefix for Kafka transaction
identifiers.
+ Required for ``"exactly-once"``, unless set via ``options``. Use a
unique
+ prefix for each concurrently running Kafka sink.
+ :param partitioner: Kafka sink partitioner, for example ``"default"``,
``"fixed"``,
+ ``"round-robin"``, or the fully qualified name of a custom
+ ``FlinkKafkaPartitioner`` subclass. ``None`` uses the connector
default.
+ :param properties: Additional Kafka producer properties with string
+ keys and values. Keys may be supplied with or without the
+ ``properties.`` prefix. ``properties.bootstrap.servers`` must use
+ the dedicated argument.
+ :param parallelism: Sink parallelism. The connector default is the
+ upstream parallelism.
+ :param options: Additional raw Kafka connector options with string
keys and values,
+ excluding ``connector``. Keys are passed through unchanged.
Conflicts with
+ options generated by other arguments (including defaults) raise
``ValueError``.
+ :param statement_set: Optional :class:`~pyflink.table.StatementSet` to
stage the write in.
+ It must use the same TableEnvironment as this DataFrame. Call its
``execute()``
+ method to submit all staged writes together.
+
+ :raises TypeError: If an argument has an invalid type.
+ :raises ValueError: If required arguments are missing, mutually
+ exclusive arguments conflict, or an option is invalid.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> events = pf.from_records(
+ ... [(1, "login")], schema=["user_id", "event"])
+ >>> events.write_kafka(
+ ... "localhost:9092",
+ ... topic="output-events",
+ ... format="json",
+ ... delivery_guarantee="at-least-once",
+ ... )
+
+ .. versionadded:: 2.4.0
+ """
+ from pyflink.dataframe.io import _build_kafka_sink_options
+
+ kafka_options = _build_kafka_sink_options(
+ bootstrap_servers,
+ topic=topic,
+ value_format=value_format if value_format is not None else format,
+ format=format,
+ format_options=format_options,
+ value_format_options=value_format_options,
+ key_format=key_format,
+ key_format_options=key_format_options,
+ key_fields=key_fields,
+ value_fields_include=value_fields_include,
+ delivery_guarantee=delivery_guarantee,
+ transactional_id_prefix=transactional_id_prefix,
+ partitioner=partitioner,
+ properties=properties,
+ parallelism=parallelism,
+ extra_options=options,
+ )
+ self._write("kafka", kafka_options, statement_set=statement_set)
+
@PublicEvolving()
def write_generic(
self,
diff --git a/flink-python/pyflink/dataframe/io.py
b/flink-python/pyflink/dataframe/io.py
index 14b2ba206a7..0c46c8eb3a5 100644
--- a/flink-python/pyflink/dataframe/io.py
+++ b/flink-python/pyflink/dataframe/io.py
@@ -16,7 +16,7 @@
# limitations under the License.
################################################################################
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Literal, Optional, Tuple, Union
from pyflink.dataframe.catalog import _validate_name
from pyflink.dataframe.context import get_or_create_table_environment
@@ -27,7 +27,7 @@ from pyflink.java_gateway import get_gateway
from pyflink.table import Schema, Table, TableDescriptor
from pyflink.util.api_stability_decorators import PublicEvolving
-__all__ = ["read_catalog_table", "read_generic", "read_json", "read_parquet"]
+__all__ = ["read_catalog_table", "read_generic", "read_json", "read_kafka",
"read_parquet"]
def _build_filesystem_options(
@@ -276,6 +276,427 @@ def read_json(
)
+def _normalize_kafka_properties(
+ properties: Optional[Dict[str, str]], *, reserved: tuple
+) -> Optional[Dict[str, str]]:
+ if properties is None:
+ return None
+ _validate_options(properties)
+ normalized = {
+ key if key.startswith("properties.") else f"properties.{key}": value
+ for key, value in properties.items()
+ }
+ for key in reserved:
+ if key in normalized:
+ raise ValueError(
+ f"{key!r} must be configured through its dedicated DataFrame "
+ f"argument, not through properties")
+ return normalized
+
+
+def _merge_kafka_format_options(
+ options: Dict[str, str],
+ value_format: str,
+ *,
+ format_options: Optional[Dict[str, str]],
+ value_format_options: Optional[Dict[str, str]],
+) -> None:
+ normalized_format_options = _normalize_kafka_format_options(
+ format_options, value_format=value_format)
+ normalized_value_format_options = _normalize_kafka_format_options(
+ value_format_options, value_format=value_format)
+ _merge_options(options, normalized_format_options)
+ _merge_options(options, normalized_value_format_options)
+
+
+def _merge_kafka_key_format_options(
+ options: Dict[str, str],
+ key_format: str,
+ key_format_options: Optional[Dict[str, str]],
+) -> None:
+ if key_format_options is None:
+ return
+ _validate_options(key_format_options)
+ normalized = {
+ key if key.startswith("key.") else f"key.{key_format}.{key}": value
+ for key, value in key_format_options.items()
+ }
+ _merge_options(options, normalized)
+
+
+def _normalize_kafka_format_options(
+ options: Optional[Dict[str, str]],
+ *,
+ value_format: str,
+) -> Dict[str, str]:
+ if options is None:
+ return {}
+ _validate_options(options)
+ value_prefix = f"value.{value_format}."
+ normalized: Dict[str, str] = {}
+ for key, value in options.items():
+ if key.startswith(value_prefix):
+ normalized[key] = value
+ elif key.startswith(f"{value_format}."):
+ normalized[value_prefix + key[len(f"{value_format}."):]] = value
+ else:
+ normalized[value_prefix + key] = value
+ return normalized
+
+
+def _build_kafka_options(
+ bootstrap_servers: str,
+ *,
+ topic: Optional[Union[str, List[str]]],
+ topic_pattern: Optional[str],
+ group_id: Optional[str],
+ value_format: str,
+ format_options: Optional[Dict[str, str]],
+ value_format_options: Optional[Dict[str, str]],
+ key_format: Optional[str],
+ key_format_options: Optional[Dict[str, str]],
+ key_fields: Optional[List[str]],
+ key_fields_prefix: Optional[str],
+ value_fields_include: str,
+ startup_mode: str,
+ startup_specific_offsets: Optional[Union[str, Dict[int, int]]],
+ startup_timestamp_millis: Optional[int],
+ topic_partition_discovery_interval: Optional[str],
+ bounded_mode: str,
+ bounded_timestamp_millis: Optional[int],
+ bounded_specific_offsets: Optional[Union[str, Dict[int, int]]],
+ properties: Optional[Dict[str, str]],
+ extra_options: Optional[Dict[str, str]] = None,
+) -> Dict[str, str]:
+ if not isinstance(bootstrap_servers, str):
+ raise TypeError("bootstrap_servers must be a string")
+ if not bootstrap_servers:
+ raise ValueError("bootstrap_servers must not be empty")
+
+ if topic is None and topic_pattern is None:
+ raise ValueError("either 'topic' or 'topic_pattern' must be specified")
+ if topic is not None and topic_pattern is not None:
+ raise ValueError("'topic' and 'topic_pattern' are mutually exclusive")
+ if topic is not None:
+ if isinstance(topic, str):
+ resolved_topic = topic
+ elif isinstance(topic, list):
+ if not topic:
+ raise ValueError("topic must not be an empty list")
+ if not all(isinstance(t, str) for t in topic):
+ raise TypeError("topic list elements must be strings")
+ resolved_topic = ";".join(topic)
+ else:
+ raise TypeError("topic must be a string or a list of strings")
+ else:
+ resolved_topic = None
+ if topic_pattern is not None and not isinstance(topic_pattern, str):
+ raise TypeError("topic_pattern must be a string")
+
+ if group_id is not None and not isinstance(group_id, str):
+ raise TypeError("group_id must be a string")
+
+ if not isinstance(value_format, str):
+ raise TypeError("format must be a string")
+ if not value_format:
+ raise ValueError("format must not be empty")
+
+ _validate_literal(value_fields_include, ("ALL", "EXCEPT_KEY"),
"value_fields_include")
+ _validate_literal(startup_mode, _STARTUP_MODES, "startup_mode")
+ _validate_literal(bounded_mode, _BOUNDED_MODES, "bounded_mode")
+
+ if startup_mode == "specific-offsets":
+ startup_specific_offsets = _convert_specific_offsets(
+ startup_specific_offsets, "startup_specific_offsets")
+ if startup_specific_offsets is None:
+ raise ValueError(
+ "startup_specific_offsets is required when "
+ "startup_mode='specific-offsets'")
+ else:
+ startup_specific_offsets = None
+ if startup_mode == "timestamp" and startup_timestamp_millis is None:
+ raise ValueError(
+ "startup_timestamp_millis is required when
startup_mode='timestamp'")
+ if (
+ startup_timestamp_millis is not None
+ and (
+ isinstance(startup_timestamp_millis, bool)
+ or not isinstance(startup_timestamp_millis, int)
+ )
+ ):
+ raise TypeError("startup_timestamp_millis must be an int")
+
+ if bounded_mode == "specific-offsets":
+ bounded_specific_offsets = _convert_specific_offsets(
+ bounded_specific_offsets, "bounded_specific_offsets")
+ if bounded_specific_offsets is None:
+ raise ValueError(
+ "bounded_specific_offsets is required when "
+ "bounded_mode='specific-offsets'")
+ else:
+ bounded_specific_offsets = None
+ if bounded_mode == "timestamp" and bounded_timestamp_millis is None:
+ raise ValueError(
+ "bounded_timestamp_millis is required when
bounded_mode='timestamp'")
+ if (
+ bounded_timestamp_millis is not None
+ and (
+ isinstance(bounded_timestamp_millis, bool)
+ or not isinstance(bounded_timestamp_millis, int)
+ )
+ ):
+ raise TypeError("bounded_timestamp_millis must be an int")
+
+ if (
+ topic_partition_discovery_interval is not None
+ and not isinstance(topic_partition_discovery_interval, str)
+ ):
+ raise TypeError("topic_partition_discovery_interval must be a string")
+
+ if key_format is not None and not isinstance(key_format, str):
+ raise TypeError("key_format must be a string")
+ if key_fields is not None:
+ if not isinstance(key_fields, list):
+ raise TypeError("key_fields must be a list of strings")
+ if not key_fields:
+ raise ValueError("key_fields must not be empty")
+ if not all(isinstance(f, str) for f in key_fields):
+ raise TypeError("key_fields elements must be strings")
+ if key_fields_prefix is not None and not isinstance(key_fields_prefix,
str):
+ raise TypeError("key_fields_prefix must be a string")
+
+ normalized_properties = _normalize_kafka_properties(
+ properties, reserved=("properties.bootstrap.servers",
"properties.group.id"))
+
+ options: Dict[str, str] = {
+ "properties.bootstrap.servers": bootstrap_servers,
+ "value.format": value_format,
+ "value.fields-include": value_fields_include,
+ }
+ if group_id is not None:
+ options["properties.group.id"] = group_id
+ if resolved_topic is not None:
+ options["topic"] = resolved_topic
+ if topic_pattern is not None:
+ options["topic-pattern"] = topic_pattern
+ options["scan.topic-partition-discovery.interval"] = (
+ topic_partition_discovery_interval
+ if topic_partition_discovery_interval is not None else "0")
+ options["scan.startup.mode"] = startup_mode
+ if startup_specific_offsets is not None:
+ options["scan.startup.specific-offsets"] = startup_specific_offsets
+ if startup_timestamp_millis is not None:
+ options["scan.startup.timestamp-millis"] =
str(startup_timestamp_millis)
+ if bounded_mode != "unbounded":
+ options["scan.bounded.mode"] = bounded_mode
+ if bounded_timestamp_millis is not None:
+ options["scan.bounded.timestamp-millis"] =
str(bounded_timestamp_millis)
+ if bounded_specific_offsets is not None:
+ options["scan.bounded.specific-offsets"] = bounded_specific_offsets
+
+ if key_format is not None:
+ options["key.format"] = key_format
+ _merge_kafka_key_format_options(options, key_format, key_format_options)
+ if key_fields is not None:
+ options["key.fields"] = ";".join(key_fields)
+ if key_fields_prefix is not None:
+ options["key.fields-prefix"] = key_fields_prefix
+
+ _merge_kafka_format_options(
+ options,
+ value_format,
+ format_options=format_options,
+ value_format_options=value_format_options,
+ )
+
+ if normalized_properties is not None:
+ _merge_options(options, normalized_properties)
+ if extra_options is not None:
+ _merge_options(options, extra_options)
+
+ return options
+
+
+_STARTUP_MODES = (
+ "earliest-offset", "latest-offset", "group-offsets",
+ "timestamp", "specific-offsets")
+_BOUNDED_MODES = (
+ "unbounded", "group-offsets", "latest-offset",
+ "timestamp", "specific-offsets")
+_DELIVERY_GUARANTEES = ("none", "at-least-once", "exactly-once")
+
+
+def _validate_literal(value: str, choices: tuple, name: str) -> None:
+ if not isinstance(value, str):
+ raise TypeError(f"{name} must be a string")
+ if value not in choices:
+ raise ValueError(
+ f"{name} must be one of {choices}, got {value!r}")
+
+
+def _convert_specific_offsets(
+ value: Optional[Union[str, Dict[int, int]]], name: str
+) -> Optional[str]:
+ if value is None:
+ return None
+ if isinstance(value, str):
+ return value
+ if isinstance(value, dict):
+ parts = []
+ for partition, offset in value.items():
+ if not isinstance(partition, int) or isinstance(partition, bool):
+ raise TypeError(f"{name} keys must be ints (partition
numbers)")
+ if not isinstance(offset, int) or isinstance(offset, bool):
+ raise TypeError(f"{name} values must be ints (offsets)")
+ parts.append(f"partition:{partition},offset:{offset}")
+ return ";".join(parts)
+ raise TypeError(f"{name} must be a str or a Dict[int, int]")
+
+
+@PublicEvolving()
+def read_kafka(
+ bootstrap_servers: str,
+ *,
+ schema: Dict[str, DataType],
+ topic: Optional[Union[str, List[str]]] = None,
+ topic_pattern: Optional[str] = None,
+ group_id: Optional[str] = None,
+ format: str = "json",
+ format_options: Optional[Dict[str, str]] = None,
+ key_format: Optional[str] = None,
+ key_format_options: Optional[Dict[str, str]] = None,
+ key_fields: Optional[List[str]] = None,
+ key_fields_prefix: Optional[str] = None,
+ value_format: Optional[str] = None,
+ value_format_options: Optional[Dict[str, str]] = None,
+ value_fields_include: Literal["ALL", "EXCEPT_KEY"] = "ALL",
+ startup_mode: Literal[
+ "earliest-offset", "latest-offset", "group-offsets",
+ "timestamp", "specific-offsets"
+ ] = "group-offsets",
+ startup_specific_offsets: Optional[Union[str, Dict[int, int]]] = None,
+ startup_timestamp_millis: Optional[int] = None,
+ topic_partition_discovery_interval: Optional[str] = "5 min",
+ bounded_mode: Literal[
+ "unbounded", "group-offsets", "latest-offset",
+ "timestamp", "specific-offsets"
+ ] = "unbounded",
+ bounded_timestamp_millis: Optional[int] = None,
+ bounded_specific_offsets: Optional[Union[str, Dict[int, int]]] = None,
+ properties: Optional[Dict[str, str]] = None,
+ options: Optional[Dict[str, str]] = None,
+ computed_columns: Optional[Dict[str, str]] = None,
+ watermark: Optional[Tuple[str, str]] = None,
+) -> DataFrame:
+ """
+ Read data from Kafka using Flink's Kafka SQL connector.
+
+ The Kafka connector must be available to Flink. Exactly one of ``topic`` or
+ ``topic_pattern`` must be specified. Physical columns in ``schema`` are
followed by
+ computed columns in dictionary insertion order. A watermark can reference
a physical
+ or computed timestamp column.
+
+ ``format`` and ``value_format`` are aliases; ``value_format`` takes
precedence when
+ both are set.
+ ``format_options`` and ``value_format_options`` are both normalized to
+ ``value.<format>.<key>``; conflicting values are rejected.
+
+ :param bootstrap_servers: Comma-separated Kafka bootstrap server addresses.
+ :param schema: Non-empty mapping of physical column names to DataFrame
data types.
+ :param topic: Kafka topic name or list of names. Mutually exclusive with
+ ``topic_pattern``.
+ :param topic_pattern: Regular expression matching Kafka topic names.
Mutually
+ exclusive with ``topic``.
+ :param group_id: Kafka consumer group id. Maps to ``properties.group.id``.
+ :param format: Value format, for example ``"json"``, ``"csv"``, ``"avro"``.
+ :param format_options: Value format options with string values, with or
without
+ the ``<format>.`` prefix.
+ :param key_format: Key format, for example ``"json"`` or ``"avro"``.
+ :param key_format_options: Key format options with string values, with or
without
+ the ``key.<key_format>.`` prefix.
+ :param key_fields: List of column names that make up the key.
+ :param key_fields_prefix: Prefix for key fields to avoid name clashes with
value
+ fields.
+ :param value_format: Alternative to ``format`` for specifying the value
format.
+ :param value_format_options: Alternative to ``format_options`` for
specifying
+ value format options.
+ :param value_fields_include: Whether value fields include key fields:
+ ``"ALL"`` or ``"EXCEPT_KEY"``.
+ :param startup_mode: Startup mode: ``"earliest-offset"``,
``"latest-offset"``,
+ ``"group-offsets"``, ``"timestamp"``, or ``"specific-offsets"``.
+ :param startup_specific_offsets: Specific offsets for ``startup_mode`` of
+ ``"specific-offsets"``, as a string
+ (``"partition:0,offset:42;partition:1,offset:300"``) or a dict
+ (``{partition: offset}``).
+ :param startup_timestamp_millis: Startup timestamp in milliseconds for
+ ``startup_mode`` of ``"timestamp"``.
+ :param topic_partition_discovery_interval: Interval for topic partition
+ discovery, for example ``"5 min"``. Set to ``None`` to disable by
setting
+ ``scan.topic-partition-discovery.interval`` to ``"0"``.
+ :param bounded_mode: Bounded mode: ``"unbounded"``, ``"group-offsets"``,
+ ``"latest-offset"``, ``"timestamp"``, or ``"specific-offsets"``.
+ :param bounded_timestamp_millis: Bounded timestamp in milliseconds for
+ ``bounded_mode`` of ``"timestamp"``.
+ :param bounded_specific_offsets: Specific offsets for ``bounded_mode`` of
+ ``"specific-offsets"``, using the same string or dict syntax as
+ ``startup_specific_offsets``.
+ :param properties: Additional Kafka consumer properties with string keys
and
+ values. Keys may be supplied with or without the ``properties.``
prefix.
+ ``properties.bootstrap.servers`` and ``properties.group.id`` must use
+ their dedicated arguments.
+ :param options: Additional raw Kafka connector options with string keys
and values,
+ excluding ``connector``. Keys are passed through unchanged. Conflicts
with
+ options generated by other arguments (including defaults) raise
``ValueError``.
+ :param computed_columns: Optional SQL expressions keyed by computed column
name.
+ :param watermark: Optional ``(column, expression)`` watermark declaration.
+ :return: A DataFrame backed by the Kafka source.
+ :raises TypeError: If an argument has an invalid type.
+ :raises ValueError: If required arguments are missing, mutually exclusive
+ arguments conflict, or an option is invalid.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> events = pf.read_kafka(
+ ... "localhost:9092",
+ ... schema={"user_id": pf.DataType.int64(), "event":
pf.DataType.string()},
+ ... topic="events",
+ ... format="json",
+ ... startup_mode="earliest-offset",
+ ... )
+
+ .. versionadded:: 2.4.0
+ """
+ effective_value_format = value_format or format
+ kafka_options = _build_kafka_options(
+ bootstrap_servers,
+ topic=topic,
+ topic_pattern=topic_pattern,
+ group_id=group_id,
+ value_format=effective_value_format,
+ format_options=format_options,
+ value_format_options=value_format_options,
+ key_format=key_format,
+ key_format_options=key_format_options,
+ key_fields=key_fields,
+ key_fields_prefix=key_fields_prefix,
+ value_fields_include=value_fields_include,
+ startup_mode=startup_mode,
+ startup_specific_offsets=startup_specific_offsets,
+ startup_timestamp_millis=startup_timestamp_millis,
+ topic_partition_discovery_interval=topic_partition_discovery_interval,
+ bounded_mode=bounded_mode,
+ bounded_timestamp_millis=bounded_timestamp_millis,
+ bounded_specific_offsets=bounded_specific_offsets,
+ properties=properties,
+ extra_options=options,
+ )
+ return _read(
+ "kafka", schema=schema, options=kafka_options,
+ computed_columns=computed_columns, watermark=watermark,
+ )
+
+
def _validate_connector(connector: str) -> None:
if not isinstance(connector, str):
raise TypeError("connector must be a string")
@@ -608,3 +1029,112 @@ def read_catalog_table(
if validated_columns or validated_watermark is not None:
table = _extend_catalog_table(table, validated_columns,
validated_watermark)
return DataFrame(table)
+
+
+def _build_kafka_sink_options(
+ bootstrap_servers: str,
+ *,
+ topic: str,
+ value_format: str,
+ format: str,
+ format_options: Optional[Dict[str, str]],
+ value_format_options: Optional[Dict[str, str]],
+ key_format: Optional[str],
+ key_format_options: Optional[Dict[str, str]],
+ key_fields: Optional[List[str]],
+ value_fields_include: str,
+ delivery_guarantee: str,
+ properties: Optional[Dict[str, str]],
+ parallelism: Optional[int],
+ transactional_id_prefix: Optional[str] = None,
+ partitioner: Optional[str] = None,
+ extra_options: Optional[Dict[str, str]] = None,
+) -> Dict[str, str]:
+ if not isinstance(bootstrap_servers, str):
+ raise TypeError("bootstrap_servers must be a string")
+ if not bootstrap_servers:
+ raise ValueError("bootstrap_servers must not be empty")
+
+ if not isinstance(topic, str):
+ raise TypeError("topic must be a string")
+ if not topic:
+ raise ValueError("topic must not be empty")
+
+ if not isinstance(value_format, str):
+ raise TypeError("format must be a string")
+ if not value_format:
+ raise ValueError("format must not be empty")
+
+ _validate_literal(value_fields_include, ("ALL", "EXCEPT_KEY"),
"value_fields_include")
+ _validate_literal(delivery_guarantee, _DELIVERY_GUARANTEES,
"delivery_guarantee")
+
+ if key_format is not None and not isinstance(key_format, str):
+ raise TypeError("key_format must be a string")
+ if key_fields is not None:
+ if not isinstance(key_fields, list):
+ raise TypeError("key_fields must be a list of strings")
+ if not key_fields:
+ raise ValueError("key_fields must not be empty")
+ if not all(isinstance(f, str) for f in key_fields):
+ raise TypeError("key_fields elements must be strings")
+ if key_format_options is not None:
+ _validate_options(key_format_options)
+
+ if parallelism is not None:
+ if isinstance(parallelism, bool) or not isinstance(parallelism, int):
+ raise TypeError("parallelism must be an int or None")
+
+ if transactional_id_prefix is not None:
+ if not isinstance(transactional_id_prefix, str):
+ raise TypeError("transactional_id_prefix must be a string")
+ if not transactional_id_prefix:
+ raise ValueError("transactional_id_prefix must not be empty")
+ if partitioner is not None:
+ if not isinstance(partitioner, str):
+ raise TypeError("partitioner must be a string")
+ if not partitioner:
+ raise ValueError("partitioner must not be empty")
+
+ normalized_properties = _normalize_kafka_properties(
+ properties, reserved=("properties.bootstrap.servers",))
+
+ options: Dict[str, str] = {
+ "properties.bootstrap.servers": bootstrap_servers,
+ "topic": topic,
+ "value.format": value_format,
+ "value.fields-include": value_fields_include,
+ "sink.delivery-guarantee": delivery_guarantee,
+ }
+ if transactional_id_prefix is not None:
+ options["sink.transactional-id-prefix"] = transactional_id_prefix
+ if partitioner is not None:
+ options["sink.partitioner"] = partitioner
+ if key_format is not None:
+ options["key.format"] = key_format
+ _merge_kafka_key_format_options(options, key_format, key_format_options)
+ if key_fields is not None:
+ options["key.fields"] = ";".join(key_fields)
+
+ _merge_kafka_format_options(
+ options,
+ value_format,
+ format_options=format_options,
+ value_format_options=value_format_options,
+ )
+
+ if parallelism is not None:
+ options["sink.parallelism"] = str(parallelism)
+
+ if normalized_properties is not None:
+ _merge_options(options, normalized_properties)
+ if extra_options is not None:
+ _merge_options(options, extra_options)
+
+ if (
+ options["sink.delivery-guarantee"] == "exactly-once"
+ and not options.get("sink.transactional-id-prefix")
+ ):
+ raise ValueError(
+ "sink.transactional-id-prefix is required when
delivery_guarantee='exactly-once'")
+
+ return options
diff --git a/flink-python/pyflink/dataframe/tests/test_kafka_io.py
b/flink-python/pyflink/dataframe/tests/test_kafka_io.py
new file mode 100644
index 00000000000..70f2fd85246
--- /dev/null
+++ b/flink-python/pyflink/dataframe/tests/test_kafka_io.py
@@ -0,0 +1,677 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+import inspect
+import unittest
+from unittest.mock import MagicMock, patch
+
+import pyflink.dataframe as pf
+from pyflink.dataframe import DataType
+from pyflink.dataframe.io import (
+ _build_kafka_options,
+ _build_kafka_sink_options,
+ _convert_specific_offsets,
+)
+from pyflink.testing.test_case_utils import PyFlinkDataFrameUTTestCase
+
+
+class KafkaOptionTests(unittest.TestCase):
+ def test_reader_schema_is_first_keyword_only_parameter(self):
+ parameters = inspect.signature(pf.read_kafka).parameters
+ self.assertEqual(list(parameters)[:2], ["bootstrap_servers", "schema"])
+ self.assertEqual(parameters["schema"].kind,
inspect.Parameter.KEYWORD_ONLY)
+ self.assertIs(parameters["schema"].default, inspect.Parameter.empty)
+
+ def test_writer_uses_parallelism_parameter(self):
+ parameters = inspect.signature(pf.DataFrame.write_kafka).parameters
+ self.assertIn("parallelism", parameters)
+ self.assertNotIn("sink_parallelism", parameters)
+ self.assertEqual(parameters["parallelism"].kind,
inspect.Parameter.KEYWORD_ONLY)
+ self.assertIsNone(parameters["parallelism"].default)
+
+ def _source_options(self, **overrides):
+ arguments = {
+ "bootstrap_servers": "localhost:9092", "topic": "events",
"topic_pattern": None,
+ "group_id": None, "value_format": "json", "format_options": None,
+ "value_format_options": None, "key_format": None,
"key_format_options": None,
+ "key_fields": None, "key_fields_prefix": None,
"value_fields_include": "ALL",
+ "startup_mode": "group-offsets", "startup_specific_offsets": None,
+ "startup_timestamp_millis": None,
"topic_partition_discovery_interval": "5 min",
+ "bounded_mode": "unbounded", "bounded_timestamp_millis": None,
+ "bounded_specific_offsets": None, "properties": None,
+ }
+ arguments.update(overrides)
+ return _build_kafka_options(**arguments)
+
+ def _sink_options(self, **overrides):
+ arguments = {
+ "bootstrap_servers": "localhost:9092", "topic": "events",
"value_format": "json",
+ "format": "json", "format_options": None, "value_format_options":
None,
+ "key_format": None, "key_format_options": None, "key_fields": None,
+ "value_fields_include": "ALL", "delivery_guarantee":
"at-least-once",
+ "properties": None, "parallelism": None,
+ }
+ arguments.update(overrides)
+ return _build_kafka_sink_options(**arguments)
+
+ def test_source_discovery_interval(self):
+ for interval, expected in ((None, "0"), ("0", "0"), ("5 min", "5
min")):
+ with self.subTest(interval=interval):
+ options =
self._source_options(topic_partition_discovery_interval=interval)
+
self.assertEqual(options["scan.topic-partition-discovery.interval"], expected)
+
+ def test_bounded_specific_offsets(self):
+ options = self._source_options(
+ bounded_mode="specific-offsets", bounded_specific_offsets={0: 10,
1: 20})
+ self.assertEqual(options["scan.bounded.specific-offsets"],
+ "partition:0,offset:10;partition:1,offset:20")
+
+ def test_sink_transactional_id_prefix_and_partitioner(self):
+ options = self._sink_options(
+ delivery_guarantee="exactly-once",
transactional_id_prefix="dataframe-events",
+ partitioner="fixed")
+ self.assertEqual(options["sink.transactional-id-prefix"],
"dataframe-events")
+ self.assertEqual(options["sink.partitioner"], "fixed")
+ defaults = self._sink_options()
+ self.assertNotIn("sink.transactional-id-prefix", defaults)
+ self.assertNotIn("sink.partitioner", defaults)
+ custom =
self._sink_options(partitioner="org.example.CustomPartitioner")
+ self.assertEqual(custom["sink.partitioner"],
"org.example.CustomPartitioner")
+
+ def test_sink_rejects_missing_or_invalid_transaction_options(self):
+ for overrides, error_type, message in (
+ ({"delivery_guarantee": "exactly-once"}, ValueError,
"transactional-id-prefix"),
+ ({"transactional_id_prefix": 1}, TypeError,
"transactional_id_prefix"),
+ ({"transactional_id_prefix": ""}, ValueError,
"transactional_id_prefix"),
+ ({"partitioner": 1}, TypeError, "partitioner"),
+ ({"partitioner": ""}, ValueError, "partitioner"),
+ ({"delivery_guarantee": "exactly-once",
+ "extra_options": {"sink.transactional-id-prefix": ""}},
ValueError,
+ "transactional-id-prefix"),
+ ):
+ with self.subTest(overrides=overrides):
+ with self.assertRaisesRegex(error_type, message):
+ self._sink_options(**overrides)
+
+ def test_raw_connector_options(self):
+ source = self._source_options(
+ extra_options={"scan.parallelism": "3", "properties.client.id":
"reader"})
+ self.assertEqual(source["scan.parallelism"], "3")
+ self.assertEqual(source["properties.client.id"], "reader")
+ sink = self._sink_options(delivery_guarantee="exactly-once",
extra_options={
+ "sink.transactional-id-prefix": "writer", "sink.partitioner":
"round-robin"})
+ self.assertEqual(sink["sink.transactional-id-prefix"], "writer")
+ self.assertEqual(sink["sink.partitioner"], "round-robin")
+
+ def test_raw_options_reject_invalid_or_conflicting_values(self):
+ for build in (self._source_options, self._sink_options):
+ raw_options = {"topic": "events", "properties.client.id": "reader"}
+ original_options = raw_options.copy()
+ build(extra_options=raw_options)
+ self.assertEqual(raw_options, original_options)
+ for extra, error_type in (
+ ([], TypeError), ({"scan.parallelism": 3}, TypeError), ({1:
"value"}, TypeError),
+ ({"": "value"}, ValueError), ({"connector": "filesystem"},
ValueError),
+ ({"topic": "other"}, ValueError), ({"value.format": "csv"},
ValueError),
+ ({"properties.bootstrap.servers": "other:9092"}, ValueError),
+ ):
+ with self.subTest(build=build.__name__, extra=extra):
+ with self.assertRaises(error_type):
+ build(extra_options=extra)
+ self.assertEqual(build(extra_options={"topic":
"events"})["topic"], "events")
+ with self.assertRaisesRegex(ValueError, "properties.client.id"):
+ build(properties={"client.id": "one"},
+ extra_options={"properties.client.id": "two"})
+ with self.assertRaisesRegex(ValueError, "sink.partitioner"):
+ self._sink_options(partitioner="fixed",
extra_options={"sink.partitioner": "default"})
+ with self.assertRaisesRegex(ValueError,
"sink.transactional-id-prefix"):
+ self._sink_options(transactional_id_prefix="one",
+ extra_options={"sink.transactional-id-prefix":
"two"})
+ with self.assertRaisesRegex(ValueError, "scan.startup.mode"):
+ self._source_options(extra_options={"scan.startup.mode":
"latest-offset"})
+
+ def test_builds_source_options(self):
+ options = _build_kafka_options(
+ "localhost:9092",
+ topic="events",
+ topic_pattern=None,
+ group_id="group",
+ value_format="json",
+ format_options={"ignore-parse-errors": "true"},
+ value_format_options=None,
+ key_format="json",
+ key_format_options={"ignore-parse-errors": "false"},
+ key_fields=["id", "name"],
+ key_fields_prefix="k_",
+ value_fields_include="EXCEPT_KEY",
+ startup_mode="specific-offsets",
+ startup_specific_offsets={0: 10, 1: 20},
+ startup_timestamp_millis=None,
+ topic_partition_discovery_interval="1 min",
+ bounded_mode="timestamp",
+ bounded_timestamp_millis=123456,
+ bounded_specific_offsets=None,
+ properties={"enable.idempotence": "true"},
+ )
+
+ self.assertEqual(options, {
+ "properties.bootstrap.servers": "localhost:9092",
+ "properties.group.id": "group",
+ "topic": "events",
+ "value.format": "json",
+ "value.json.ignore-parse-errors": "true",
+ "key.format": "json",
+ "key.json.ignore-parse-errors": "false",
+ "key.fields": "id;name",
+ "key.fields-prefix": "k_",
+ "value.fields-include": "EXCEPT_KEY",
+ "scan.startup.mode": "specific-offsets",
+ "scan.startup.specific-offsets":
"partition:0,offset:10;partition:1,offset:20",
+ "scan.topic-partition-discovery.interval": "1 min",
+ "scan.bounded.mode": "timestamp",
+ "scan.bounded.timestamp-millis": "123456",
+ "properties.enable.idempotence": "true",
+ })
+
+ def test_source_topic_list_is_semicolon_joined(self):
+ options = _build_kafka_options(
+ "localhost:9092",
+ topic=["events-1", "events-2"],
+ topic_pattern=None,
+ group_id=None,
+ value_format="json",
+ format_options=None,
+ value_format_options=None,
+ key_format=None,
+ key_format_options=None,
+ key_fields=None,
+ key_fields_prefix=None,
+ value_fields_include="ALL",
+ startup_mode="group-offsets",
+ startup_specific_offsets=None,
+ startup_timestamp_millis=None,
+ topic_partition_discovery_interval=None,
+ bounded_mode="unbounded",
+ bounded_timestamp_millis=None,
+ bounded_specific_offsets=None,
+ properties=None,
+ )
+
+ self.assertEqual(options["topic"], "events-1;events-2")
+ self.assertNotIn("topic-pattern", options)
+
+ def test_source_topic_pattern(self):
+ options = _build_kafka_options(
+ "localhost:9092",
+ topic=None,
+ topic_pattern="events-.*",
+ group_id=None,
+ value_format="json",
+ format_options=None,
+ value_format_options=None,
+ key_format=None,
+ key_format_options=None,
+ key_fields=None,
+ key_fields_prefix=None,
+ value_fields_include="ALL",
+ startup_mode="group-offsets",
+ startup_specific_offsets=None,
+ startup_timestamp_millis=None,
+ topic_partition_discovery_interval=None,
+ bounded_mode="unbounded",
+ bounded_timestamp_millis=None,
+ bounded_specific_offsets=None,
+ properties=None,
+ )
+
+ self.assertEqual(options["topic-pattern"], "events-.*")
+ self.assertNotIn("topic", options)
+
+ def test_source_rejects_conflicting_format_option_aliases(self):
+ with self.assertRaisesRegex(ValueError, "conflicting values"):
+ _build_kafka_options(
+ "localhost:9092",
+ topic="events",
+ topic_pattern=None,
+ group_id=None,
+ value_format="json",
+ format_options={"ignore-parse-errors": "true"},
+ value_format_options={"ignore-parse-errors": "false"},
+ key_format=None,
+ key_format_options=None,
+ key_fields=None,
+ key_fields_prefix=None,
+ value_fields_include="ALL",
+ startup_mode="group-offsets",
+ startup_specific_offsets=None,
+ startup_timestamp_millis=None,
+ topic_partition_discovery_interval=None,
+ bounded_mode="unbounded",
+ bounded_timestamp_millis=None,
+ bounded_specific_offsets=None,
+ properties=None,
+ )
+
+ def test_source_rejects_invalid_arguments(self):
+ base_arguments = {
+ "topic": "events",
+ "topic_pattern": None,
+ "group_id": None,
+ "value_format": "json",
+ "format_options": None,
+ "value_format_options": None,
+ "key_format": None,
+ "key_format_options": None,
+ "key_fields": None,
+ "key_fields_prefix": None,
+ "value_fields_include": "ALL",
+ "startup_mode": "group-offsets",
+ "startup_specific_offsets": None,
+ "startup_timestamp_millis": None,
+ "topic_partition_discovery_interval": None,
+ "bounded_mode": "unbounded",
+ "bounded_timestamp_millis": None,
+ "bounded_specific_offsets": None,
+ "properties": None,
+ }
+ cases = [
+ ({"bootstrap_servers": None}, TypeError, "bootstrap_servers must
be a string"),
+ ({"bootstrap_servers": ""}, ValueError, "bootstrap_servers must
not be empty"),
+ ({"topic": None}, ValueError, "either 'topic' or 'topic_pattern'"),
+ ({"topic": "a", "topic_pattern": "b"}, ValueError, "mutually
exclusive"),
+ ({"topic": []}, ValueError, "topic must not be an empty list"),
+ ({"topic": [1]}, TypeError, "topic list elements must be strings"),
+ ({"topic": 1}, TypeError, "topic must be a string"),
+ ({"topic": None, "topic_pattern": 1}, TypeError, "topic_pattern
must be a string"),
+ ({"group_id": 1}, TypeError, "group_id must be a string"),
+ ({"value_fields_include": "BAD"}, ValueError,
"value_fields_include"),
+ ({"startup_mode": "bad"}, ValueError, "startup_mode"),
+ ({"startup_mode": "specific-offsets"}, ValueError,
"startup_specific_offsets"),
+ ({"startup_mode": "timestamp"}, ValueError,
"startup_timestamp_millis"),
+ ({"startup_mode": "timestamp", "startup_timestamp_millis": True},
TypeError,
+ "startup_timestamp_millis must be an int"),
+ ({"bounded_mode": "bad"}, ValueError, "bounded_mode"),
+ ({"bounded_mode": "specific-offsets"}, ValueError,
"bounded_specific_offsets"),
+ ({"bounded_mode": "timestamp"}, ValueError,
"bounded_timestamp_millis"),
+ ({"bounded_mode": "timestamp", "bounded_timestamp_millis": True},
TypeError,
+ "bounded_timestamp_millis must be an int"),
+ ({"topic_partition_discovery_interval": 1}, TypeError,
+ "topic_partition_discovery_interval must be a string"),
+ ({"key_format": 1}, TypeError, "key_format must be a string"),
+ ({"key_fields": "id"}, TypeError, "key_fields must be a list"),
+ ({"key_fields": []}, ValueError, "key_fields must not be empty"),
+ ({"key_fields": [1]}, TypeError, "key_fields elements must be
strings"),
+ ({"key_fields_prefix": 1}, TypeError, "key_fields_prefix must be a
string"),
+ ({"properties": {"properties.bootstrap.servers": "other:9092"}},
ValueError,
+ "properties.bootstrap.servers"),
+ ({"properties": {"properties.group.id": "other"}}, ValueError,
+ "properties.group.id"),
+ ]
+
+ for overrides, error_type, message in cases:
+ with self.subTest(overrides=overrides):
+ arguments = {"bootstrap_servers": "localhost:9092"}
+ arguments.update(base_arguments)
+ arguments.update(overrides)
+ with self.assertRaisesRegex(error_type, message):
+ _build_kafka_options(**arguments)
+
+ def test_source_accepts_explicit_value_format_options(self):
+ arguments = {
+ "bootstrap_servers": "localhost:9092",
+ "topic": "events",
+ "topic_pattern": None,
+ "group_id": None,
+ "value_format": "json",
+ "format_options": {"ignore-parse-errors": "true"},
+ "value_format_options": {"ignore-parse-errors": "false"},
+ "key_format": None,
+ "key_format_options": None,
+ "key_fields": None,
+ "key_fields_prefix": None,
+ "value_fields_include": "ALL",
+ "startup_mode": "group-offsets",
+ "startup_specific_offsets": None,
+ "startup_timestamp_millis": None,
+ "topic_partition_discovery_interval": None,
+ "bounded_mode": "unbounded",
+ "bounded_timestamp_millis": None,
+ "bounded_specific_offsets": None,
+ "properties": None,
+ }
+ arguments["format_options"] = {"json.ignore-parse-errors": "true"}
+ arguments["value_format_options"] = {"json.ignore-parse-errors":
"true"}
+ options = _build_kafka_options(**arguments)
+
+ self.assertEqual(options["value.json.ignore-parse-errors"], "true")
+
+ def test_specific_offsets_are_converted(self):
+ self.assertEqual(_convert_specific_offsets({1: 2}, "offsets"),
"partition:1,offset:2")
+ self.assertEqual(_convert_specific_offsets("partition:0,offset:1",
"offsets"),
+ "partition:0,offset:1")
+ self.assertEqual(
+ _convert_specific_offsets({0: 10, 1: 20}, "offsets"),
+ "partition:0,offset:10;partition:1,offset:20")
+ self.assertIsNone(_convert_specific_offsets(None, "offsets"))
+ with self.assertRaisesRegex(TypeError, "offsets keys"):
+ _convert_specific_offsets({"1": 2}, "offsets")
+ with self.assertRaisesRegex(TypeError, "offsets values"):
+ _convert_specific_offsets({1: "2"}, "offsets")
+ with self.assertRaisesRegex(TypeError, "offsets must be"):
+ _convert_specific_offsets([], "offsets")
+
+ def test_source_prefixes_kafka_properties(self):
+ options = _build_kafka_options(
+ "localhost:9092",
+ topic="events",
+ topic_pattern=None,
+ group_id=None,
+ value_format="json",
+ format_options=None,
+ value_format_options=None,
+ key_format=None,
+ key_format_options=None,
+ key_fields=None,
+ key_fields_prefix=None,
+ value_fields_include="ALL",
+ startup_mode="group-offsets",
+ startup_specific_offsets=None,
+ startup_timestamp_millis=None,
+ topic_partition_discovery_interval=None,
+ bounded_mode="unbounded",
+ bounded_timestamp_millis=None,
+ bounded_specific_offsets=None,
+ properties={"security.protocol": "SASL_SSL"},
+ )
+
+ self.assertEqual(options["properties.security.protocol"], "SASL_SSL")
+
+ def test_builds_sink_options(self):
+ options = _build_kafka_sink_options(
+ "localhost:9092",
+ topic="events",
+ value_format="json",
+ format="json",
+ format_options={"ignore-parse-errors": "true"},
+ value_format_options=None,
+ key_format="json",
+ key_format_options={"ignore-parse-errors": "false"},
+ key_fields=["id", "name"],
+ value_fields_include="EXCEPT_KEY",
+ delivery_guarantee="exactly-once",
+ transactional_id_prefix="events-writer",
+ properties={"transaction.timeout.ms": "60000"},
+ parallelism=3,
+ )
+
+ self.assertEqual(options, {
+ "properties.bootstrap.servers": "localhost:9092",
+ "topic": "events",
+ "value.format": "json",
+ "value.json.ignore-parse-errors": "true",
+ "key.format": "json",
+ "key.json.ignore-parse-errors": "false",
+ "key.fields": "id;name",
+ "value.fields-include": "EXCEPT_KEY",
+ "sink.delivery-guarantee": "exactly-once",
+ "sink.transactional-id-prefix": "events-writer",
+ "properties.transaction.timeout.ms": "60000",
+ "sink.parallelism": "3",
+ })
+
+ def test_sink_value_format_takes_precedence(self):
+ options = _build_kafka_sink_options(
+ "localhost:9092",
+ topic="events",
+ value_format="csv",
+ format=None,
+ format_options=None,
+ value_format_options=None,
+ key_format=None,
+ key_format_options=None,
+ key_fields=None,
+ value_fields_include="ALL",
+ delivery_guarantee="at-least-once",
+ properties=None,
+ parallelism=None,
+ )
+
+ self.assertEqual(options["value.format"], "csv")
+
+ def test_sink_value_format_overrides_format(self):
+ options = _build_kafka_sink_options(
+ "localhost:9092",
+ topic="events",
+ value_format="csv",
+ format="json",
+ format_options=None,
+ value_format_options=None,
+ key_format=None,
+ key_format_options=None,
+ key_fields=None,
+ value_fields_include="ALL",
+ delivery_guarantee="at-least-once",
+ properties=None,
+ parallelism=None,
+ )
+
+ self.assertEqual(options["value.format"], "csv")
+
+ def test_sink_rejects_invalid_arguments(self):
+ base_arguments = {
+ "value_format": "json",
+ "format": "json",
+ "format_options": None,
+ "value_format_options": None,
+ "key_format": None,
+ "key_format_options": None,
+ "key_fields": None,
+ "value_fields_include": "ALL",
+ "delivery_guarantee": "at-least-once",
+ "properties": None,
+ "parallelism": None,
+ }
+ cases = [
+ ({"bootstrap_servers": None}, TypeError, "bootstrap_servers must
be a string"),
+ ({"topic": None}, TypeError, "topic must be a string"),
+ ({"topic": ""}, ValueError, "topic must not be empty"),
+ ({"value_fields_include": "BAD"}, ValueError,
"value_fields_include"),
+ ({"delivery_guarantee": "bad"}, ValueError, "delivery_guarantee"),
+ ({"key_format": 1}, TypeError, "key_format must be a string"),
+ ({"key_fields": "id"}, TypeError, "key_fields must be a list"),
+ ({"key_fields": []}, ValueError, "key_fields must not be empty"),
+ ({"key_fields": [1]}, TypeError, "key_fields elements must be
strings"),
+ ({"parallelism": "3"}, TypeError, "parallelism must be an int"),
+ ({"parallelism": True}, TypeError, "parallelism must be an int"),
+ ({"properties": {"properties.bootstrap.servers": "other:9092"}},
ValueError,
+ "properties.bootstrap.servers"),
+
+ ]
+
+ for overrides, error_type, message in cases:
+ with self.subTest(overrides=overrides):
+ arguments = {
+ "bootstrap_servers": "localhost:9092",
+ "topic": "events",
+ }
+ arguments.update(base_arguments)
+ arguments.update(overrides)
+ with self.assertRaisesRegex(error_type, message):
+ _build_kafka_sink_options(**arguments)
+
+
+class KafkaDescriptorTests(PyFlinkDataFrameUTTestCase):
+ _SCHEMA = {"id": DataType.int64(), "name": DataType.string()}
+
+ def test_read_kafka_builds_source_descriptor(self):
+ with patch.object(
+ self.t_env,
+ "from_descriptor",
+ wraps=self.t_env.from_descriptor,
+ ) as from_descriptor:
+ dataframe = pf.read_kafka(
+ "localhost:9092",
+ topic="events",
+ schema={
+ "id": DataType.int64(),
+ "ts_millis": DataType.int64(),
+ },
+ format="json",
+ startup_mode="earliest-offset",
+ computed_columns={
+ "event_time": "TO_TIMESTAMP_LTZ(ts_millis, 3)"
+ },
+ watermark=(
+ "event_time",
+ "event_time - INTERVAL '5' SECOND",
+ ),
+ )
+
+ descriptor = from_descriptor.call_args.args[0]
+ self.assertEqual(descriptor.get_options().get("connector"), "kafka")
+ self.assertEqual(descriptor.get_options().get("topic"), "events")
+ self.assertEqual(descriptor.get_options().get("value.format"), "json")
+ self.assertEqual(
+ descriptor.get_options().get("scan.startup.mode"),
"earliest-offset")
+ self.assert_dataframe_schema(dataframe, ["id", "ts_millis",
"event_time"])
+ watermark_specs =
dataframe._table.get_resolved_schema().get_watermark_specs()
+ self.assertEqual(len(watermark_specs), 1)
+ self.assertEqual(watermark_specs[0].get_rowtime_attribute(),
"event_time")
+
+ def test_read_kafka_uses_value_format(self):
+ with patch.object(
+ self.t_env,
+ "from_descriptor",
+ wraps=self.t_env.from_descriptor,
+ ) as from_descriptor:
+ pf.read_kafka(
+ "localhost:9092",
+ topic="events",
+ schema=self._SCHEMA,
+ format="json",
+ value_format="csv",
+ )
+
+ options = from_descriptor.call_args.args[0].get_options()
+ self.assertEqual(options.get("value.format"), "csv")
+
+ def test_read_kafka_forwards_raw_options(self):
+ with patch.object(
+ self.t_env, "from_descriptor", wraps=self.t_env.from_descriptor,
+ ) as from_descriptor:
+ pf.read_kafka(
+ "localhost:9092", schema=self._SCHEMA, topic=["events-1",
"events-2"],
+ key_format="json", key_fields=["id", "name"],
+ startup_mode="specific-offsets", startup_specific_offsets={0:
10, 1: 20},
+ topic_partition_discovery_interval=None,
options={"scan.parallelism": "3"},
+ )
+
+ options = from_descriptor.call_args.args[0].get_options()
+ self.assertEqual(options.get("topic"), "events-1;events-2")
+ self.assertEqual(options.get("key.fields"), "id;name")
+ self.assertEqual(options.get("scan.startup.specific-offsets"),
+ "partition:0,offset:10;partition:1,offset:20")
+
self.assertEqual(options.get("scan.topic-partition-discovery.interval"), "0")
+ self.assertEqual(options.get("scan.parallelism"), "3")
+
+ def test_write_kafka_forwards_transaction_and_raw_options(self):
+ dataframe = pf.from_records([(1, "a")], schema=["id", "name"])
+ for use_raw_options in (False, True):
+ with self.subTest(use_raw_options=use_raw_options):
+ raw_options = {"sink.parallelism": "3"}
+ transaction_options = {
+ "sink.transactional-id-prefix": "events-writer",
"sink.partitioner": "fixed",
+ }
+ if use_raw_options:
+ raw_options.update(transaction_options)
+ with patch.object(dataframe._table, "execute_insert") as
execute_insert:
+ dataframe.write_kafka(
+ "localhost:9092", topic="events",
delivery_guarantee="exactly-once",
+ transactional_id_prefix=None if use_raw_options else
"events-writer",
+ partitioner=None if use_raw_options else "fixed",
options=raw_options,
+ key_format="json", key_fields=["id", "name"],
+ )
+
+ options = execute_insert.call_args.args[0].get_options()
+ self.assertEqual(options.get("sink.transactional-id-prefix"),
"events-writer")
+ self.assertEqual(options.get("sink.partitioner"), "fixed")
+ self.assertEqual(options.get("sink.parallelism"), "3")
+ self.assertEqual(options.get("key.fields"), "id;name")
+
+ def test_write_kafka_stages_options_without_execution(self):
+ dataframe = pf.from_records([(1, "a")], schema=["id", "name"])
+ statement_set = pf.create_statement_set()
+ with patch.object(statement_set, "add_insert") as add_insert, \
+ patch.object(statement_set, "execute") as execute, \
+ patch.object(dataframe._table, "execute_insert") as
execute_insert:
+ dataframe.write_kafka(
+ "localhost:9092", topic="events",
delivery_guarantee="exactly-once",
+ transactional_id_prefix="events-writer", partitioner="fixed",
+ options={"sink.parallelism": "3"}, statement_set=statement_set,
+ )
+
+ add_insert.assert_called_once()
+ self.assertIs(add_insert.call_args.args[1], dataframe._table)
+ options = add_insert.call_args.args[0].get_options()
+ self.assertEqual(options.get("sink.transactional-id-prefix"),
"events-writer")
+ self.assertEqual(options.get("sink.partitioner"), "fixed")
+ self.assertEqual(options.get("sink.parallelism"), "3")
+ execute.assert_not_called()
+ execute_insert.assert_not_called()
+
+ def test_write_kafka_builds_sink_descriptor(self):
+ dataframe = pf.from_records([(1, "a")], schema=["id", "name"])
+
+ with patch.object(dataframe._table, "execute_insert") as
execute_insert:
+ execute_insert.return_value = MagicMock()
+ result = dataframe.write_kafka(
+ "localhost:9092",
+ topic="events",
+ format="json",
+ delivery_guarantee="at-least-once",
+ parallelism=3,
+ )
+
+ descriptor = execute_insert.call_args.args[0]
+ self.assertIsNone(result)
+ self.assertIsNone(descriptor.get_schema())
+ self.assertEqual(descriptor.get_options().get("connector"), "kafka")
+ self.assertEqual(descriptor.get_options().get("topic"), "events")
+ self.assertEqual(descriptor.get_options().get("value.format"), "json")
+ self.assertEqual(
+ descriptor.get_options().get("sink.delivery-guarantee"),
+ "at-least-once",
+ )
+ self.assertEqual(descriptor.get_options().get("sink.parallelism"), "3")
+ self.assertEqual(execute_insert.call_args.kwargs, {"overwrite": False})
+
+ def test_write_kafka_uses_value_format(self):
+ dataframe = pf.from_records([(1, "a")], schema=["id", "name"])
+ with patch.object(dataframe._table, "execute_insert") as
execute_insert:
+ dataframe.write_kafka(
+ "localhost:9092",
+ topic="events",
+ value_format="csv",
+ )
+
+ options = execute_insert.call_args.args[0].get_options()
+ self.assertEqual(options.get("value.format"), "csv")
+
+
+if __name__ == "__main__":
+ unittest.main()
