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 b5508f59e02 [FLINK-40429][python] Add StatementSet support for
DataFrame writes
b5508f59e02 is described below
commit b5508f59e02d22c28722e35f3292ea5154d23348
Author: beetle0915 <[email protected]>
AuthorDate: Tue Oct 6 17:28:22 2026 +0800
[FLINK-40429][python] Add StatementSet support for DataFrame writes
Expose create_statement_set and stage DataFrame writes through the existing
Table API StatementSet when requested. Preserve eager writes by default and
cover multi-sink execution in batch and streaming modes.
This closes #29400.
Generated-by: OpenAI Codex (GPT-6)
---
.../reference/pyflink.dataframe/environment.rst | 1 +
.../docs/reference/pyflink.dataframe/io.rst | 30 ++++
flink-python/pyflink/dataframe/__init__.py | 2 +
flink-python/pyflink/dataframe/context.py | 28 +++
flink-python/pyflink/dataframe/dataframe.py | 104 ++++++++---
.../pyflink/dataframe/tests/test_statement_set.py | 198 +++++++++++++++++++++
6 files changed, 342 insertions(+), 21 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/environment.rst
b/flink-python/docs/reference/pyflink.dataframe/environment.rst
index e0e4d08f925..241de31022e 100644
--- a/flink-python/docs/reference/pyflink.dataframe/environment.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/environment.rst
@@ -30,3 +30,4 @@ Functions for setting, inspecting, and creating the
environment used by DataFram
set_table_environment
get_table_environment
get_or_create_table_environment
+ create_statement_set
diff --git a/flink-python/docs/reference/pyflink.dataframe/io.rst
b/flink-python/docs/reference/pyflink.dataframe/io.rst
index a92bd5a13fd..eaa059d4e3c 100644
--- a/flink-python/docs/reference/pyflink.dataframe/io.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/io.rst
@@ -50,6 +50,36 @@ Writers
DataFrame.write_json
DataFrame.write_parquet
+Writing to multiple sinks
+-------------------------
+
+Writers submit their jobs immediately by default. To submit several writes as
one job,
+create a statement set and pass it to each writer with ``statement_set=``.
These calls
+only stage the writes; ``execute()`` submits them together and returns a
+:class:`~pyflink.table.TableResult`. Use ``wait()`` on that result to await
completion::
+
+ import pyflink.dataframe as pf
+
+ pf.config.set("execution.runtime-mode", "batch")
+ events = pf.from_records([(1, "login")], schema=["id", "event"])
+ statement_set = pf.create_statement_set()
+ events.write_json("file:///tmp/events-json", statement_set=statement_set)
+ events.write_generic(
+ "filesystem",
+ options={"path": "file:///tmp/events-csv", "format": "csv"},
+ statement_set=statement_set,
+ )
+ result = statement_set.execute()
+ result.wait()
+
+``statement_set.explain()`` inspects the combined plan without executing or
clearing the staged
+writes.
+The statement set and all its DataFrames must use the same TableEnvironment.
Flink can
+reuse shared sources and sub-plans across the sinks, subject to optimizer
configuration.
+The sinks share job backpressure and failover; this does not guarantee
cross-sink transactions.
+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.
+
Filesystem sources and sinks
----------------------------
diff --git a/flink-python/pyflink/dataframe/__init__.py
b/flink-python/pyflink/dataframe/__init__.py
index 2ac311e66b4..0a30d752d19 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -57,6 +57,7 @@ from pyflink.dataframe.convert import (
range,
)
from pyflink.dataframe.context import (
+ create_statement_set,
get_or_create_table_environment,
get_table_environment,
set_table_environment,
@@ -100,4 +101,5 @@ __all__ = [
"set_table_environment",
"get_table_environment",
"get_or_create_table_environment",
+ "create_statement_set",
]
diff --git a/flink-python/pyflink/dataframe/context.py
b/flink-python/pyflink/dataframe/context.py
index 3abdcdaf403..4219643b9a7 100644
--- a/flink-python/pyflink/dataframe/context.py
+++ b/flink-python/pyflink/dataframe/context.py
@@ -21,6 +21,7 @@ from typing import Optional
from pyflink.common import Configuration
from pyflink.java_gateway import get_gateway
from pyflink.table import EnvironmentSettings, StreamTableEnvironment,
TableEnvironment
+from pyflink.table.statement_set import StatementSet
from pyflink.util.api_stability_decorators import PublicEvolving
from pyflink.util.java_utils import get_j_env_configuration
@@ -28,11 +29,38 @@ __all__ = [
"set_table_environment",
"get_table_environment",
"get_or_create_table_environment",
+ "create_statement_set",
]
_global_table_environment: Optional[TableEnvironment] = None
+@PublicEvolving()
+def create_statement_set() -> StatementSet:
+ """
+ Create a statement set for submitting multiple DataFrame writes as one job.
+
+ Pass the set as ``statement_set`` to a DataFrame writer to stage the write
without
+ executing it. Call :meth:`~pyflink.table.StatementSet.execute` to submit
the staged
+ writes, or :meth:`~pyflink.table.StatementSet.explain` to inspect their
execution plan.
+ The set and its DataFrames must belong to the same TableEnvironment.
+
+ :return: A statement set using the configured or newly created DataFrame
environment.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> events = pf.from_records([(1, "login")], schema=["id", "event"])
+ >>> statement_set = pf.create_statement_set()
+ >>> events.write_json("file:///tmp/events-json",
statement_set=statement_set)
+ >>> events.write_generic("blackhole", options={},
statement_set=statement_set)
+ >>> result = statement_set.execute()
+
+ .. versionadded:: 2.4.0
+ """
+ return get_or_create_table_environment().create_statement_set()
+
+
@PublicEvolving()
def set_table_environment(t_env: Optional[TableEnvironment]) -> None:
"""
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index a2694e562a0..7a9679166c8 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -52,6 +52,7 @@ from pyflink.table.expressions import (
lit as table_lit,
)
from pyflink.table.table import Table
+from pyflink.table.statement_set import StatementSet
from pyflink.table.types import ArrayType, MapType, MultisetType, RowType
from pyflink.table.table_descriptor import TableDescriptor
from pyflink.util.api_stability_decorators import PublicEvolving
@@ -1823,12 +1824,14 @@ class DataFrame:
partition_commit_policy_kind: Optional[str] = None,
connector_options: Optional[Dict[str, str]] = None,
format_options: Optional[Dict[str, str]] = None,
+ statement_set: Optional[StatementSet] = None,
) -> None:
"""
Write Parquet files using Flink's filesystem connector.
- The filesystem connector and Parquet format must be available to
Flink. The write
- is submitted immediately and waits for completion for local or
MiniCluster execution.
+ The filesystem connector and Parquet format 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. Writes default
to append in both
batch and streaming execution. Explicit overwrite requires batch
execution and is rejected
by the filesystem connector in streaming execution. Partitioned
overwrite replaces only
@@ -1888,10 +1891,14 @@ class DataFrame:
leave dictionary values unchanged; conflicting explicit values are
rejected.
For INT64 timestamp encoding, use ``{"write.int64.timestamp":
"true",
"timestamp.time.unit": "micros"}``. The default encoding is INT96.
+ :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 the path is empty, the write mode is
unsupported, the partition
specification is empty or contains empty or duplicate names, or
options conflict
- or contain reserved, empty or duplicate keys.
+ or contain reserved, empty or duplicate keys, or the statement set
belongs to a
+ different TableEnvironment.
Example::
@@ -1932,7 +1939,13 @@ class DataFrame:
extra_connector_options=connector_options,
extra_format_options=format_options,
)
- self._write("filesystem", options, mode=mode,
partition_by=partition_by)
+ self._write(
+ "filesystem",
+ options,
+ mode=mode,
+ partition_by=partition_by,
+ statement_set=statement_set,
+ )
@PublicEvolving()
def write_json(
@@ -1957,12 +1970,14 @@ class DataFrame:
partition_commit_policy_kind: Optional[str] = None,
connector_options: Optional[Dict[str, str]] = None,
format_options: Optional[Dict[str, str]] = None,
+ statement_set: Optional[StatementSet] = None,
) -> None:
"""
Write newline-delimited JSON files using Flink's filesystem connector.
- The filesystem connector and JSON format must be available to Flink.
The write
- is submitted immediately and waits for completion for local or
MiniCluster execution.
+ The filesystem connector and JSON format 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. Writes default
to append in both
batch and streaming execution. Explicit overwrite requires batch
execution and is rejected
by the filesystem connector in streaming execution. Partitioned
overwrite replaces only
@@ -2025,10 +2040,14 @@ class DataFrame:
the ``json.`` prefix, for example ``{"timestamp-format.standard":
"ISO-8601"}``.
``None`` parameters leave dictionary values unchanged; conflicting
explicit values
are rejected.
+ :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 the path is empty, the write mode is
unsupported, the partition
specification is empty or contains empty or duplicate names, or
options conflict
- or contain reserved, empty or duplicate keys.
+ or contain reserved, empty or duplicate keys, or the statement set
belongs to a
+ different TableEnvironment.
Example::
@@ -2071,22 +2090,39 @@ class DataFrame:
extra_connector_options=connector_options,
extra_format_options=format_options,
)
- self._write("filesystem", options, mode=mode,
partition_by=partition_by)
+ self._write(
+ "filesystem",
+ options,
+ mode=mode,
+ partition_by=partition_by,
+ statement_set=statement_set,
+ )
@PublicEvolving()
- def write_generic(self, connector: str, *, options: Dict[str, str]) ->
None:
+ def write_generic(
+ self,
+ connector: str,
+ *,
+ options: Dict[str, str],
+ statement_set: Optional[StatementSet] = None,
+ ) -> None:
"""
Write this DataFrame using a connector and its raw Table connector
options.
- The connector must be available through Flink's factory discovery
mechanism. The write is
- submitted immediately and waits for completion when using local or
MiniCluster execution.
+ The connector must be available through Flink's factory discovery
mechanism. 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 output schema.
:param connector: Factory identifier used as the ``connector`` Table
option.
:param options: Connector options, excluding the reserved
``connector`` option.
+ :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 the connector or an option key is empty, if
``options`` contains
- the reserved ``connector`` key, or if Flink rejects the connector
or its options.
+ the reserved ``connector`` key, if Flink rejects the connector or
its options, or
+ if the statement set belongs to a different TableEnvironment.
Example::
@@ -2102,7 +2138,7 @@ class DataFrame:
.. versionadded:: 2.4.0
"""
- self._write(connector, options)
+ self._write(connector, options, statement_set=statement_set)
def _write(
self,
@@ -2111,6 +2147,7 @@ class DataFrame:
*,
mode: str = "append",
partition_by: Optional[Union[str, List[str]]] = None,
+ statement_set: Optional[StatementSet] = None,
) -> None:
from pyflink.dataframe.errors import _raise_as_value_error
from pyflink.dataframe.io import _build_generic_descriptor
@@ -2121,27 +2158,38 @@ class DataFrame:
raise ValueError("mode must be 'append' or 'overwrite'")
descriptor = _build_generic_descriptor(connector, options,
partition_by=partition_by)
try:
- self._execute_insert(descriptor, mode == "overwrite")
+ self._execute_insert(descriptor, mode == "overwrite",
statement_set=statement_set)
except Exception as error:
_raise_as_value_error(error)
@PublicEvolving()
- def write_catalog_table(self, path: str, *, overwrite: bool = False) ->
None:
+ def write_catalog_table(
+ self,
+ path: str,
+ *,
+ overwrite: bool = False,
+ statement_set: Optional[StatementSet] = None,
+ ) -> None:
"""
Write this DataFrame to a table registered in a catalog.
``path`` is ``table_name``, ``db_name.table_name``, or
``catalog_name.db_name.table_name``.
Missing parts are resolved against the current catalog and database,
see
:func:`~pyflink.dataframe.use_catalog` and
:func:`~pyflink.dataframe.use_database`. The
- write runs right away. On a local or MiniCluster setup the call blocks
until the write is
- done.
+ write runs right away by default. On a local or MiniCluster setup the
call blocks until
+ the write is done. Passing ``statement_set`` stages the write without
executing it.
:param path: Path of the catalog table.
:param overwrite: Whether existing data should be replaced, like
``INSERT OVERWRITE``.
Not every connector supports overwriting.
- :raises TypeError: If ``path`` is not a string or ``overwrite`` is not
a bool.
+ :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 ``path`` is not a string, ``overwrite`` is not a
bool, or
+ ``statement_set`` is neither a StatementSet nor ``None``.
:raises ValueError: If ``path`` is empty, malformed, or does not name
a table, or if the
- DataFrame's columns do not match the table.
+ DataFrame's columns do not match the table, or the statement set
belongs to a
+ different TableEnvironment.
Example::
@@ -2160,13 +2208,27 @@ class DataFrame:
if not isinstance(overwrite, bool):
raise TypeError("overwrite must be a bool")
try:
- self._execute_insert(path, overwrite)
+ self._execute_insert(path, overwrite, statement_set=statement_set)
except Exception as error:
_raise_as_value_error(error)
def _execute_insert(
- self, target: Union[str, TableDescriptor], overwrite: bool = False
+ self,
+ target: Union[str, TableDescriptor],
+ overwrite: bool = False,
+ *,
+ statement_set: Optional[StatementSet] = None,
) -> None:
+ if statement_set is not None:
+ if not isinstance(statement_set, StatementSet):
+ raise TypeError("statement_set must be a StatementSet or None")
+ if self._table._t_env._j_tenv != statement_set._t_env._j_tenv:
+ raise ValueError(
+ "DataFrame and statement_set must belong to the same
TableEnvironment"
+ )
+ statement_set.add_insert(target, self._table, overwrite=overwrite)
+ return
+
result = self._table.execute_insert(target, overwrite=overwrite)
execution_target = self._table._t_env.get_config().get(
"execution.target", None
diff --git a/flink-python/pyflink/dataframe/tests/test_statement_set.py
b/flink-python/pyflink/dataframe/tests/test_statement_set.py
new file mode 100644
index 00000000000..a589aafa6fd
--- /dev/null
+++ b/flink-python/pyflink/dataframe/tests/test_statement_set.py
@@ -0,0 +1,198 @@
+################################################################################
+# 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 json
+import os
+import unittest
+from unittest.mock import patch
+
+import pyflink.dataframe as pf
+from pyflink.datastream import RuntimeExecutionMode
+from pyflink.table import (
+ EnvironmentSettings,
+ StatementSet,
+ StreamTableEnvironment,
+ TableEnvironment,
+)
+from pyflink.testing.test_case_utils import PyFlinkDataFrameUTTestCase
+
+
+class StatementSetTests(PyFlinkDataFrameUTTestCase):
+ def test_factory_uses_dataframe_environment(self):
+ statement_set = pf.create_statement_set()
+
+ self.assertIsInstance(statement_set, StatementSet)
+ self.assertIs(statement_set._t_env, self.t_env)
+ self.assertIsNot(statement_set, pf.create_statement_set())
+ self.assertIn("create_statement_set", pf.__all__)
+
+ def test_factory_creates_default_environment(self):
+ pf.set_table_environment(None)
+
+ statement_set = pf.create_statement_set()
+
+ self.assertIs(statement_set._t_env, pf.get_table_environment())
+ self.assertIsNotNone(pf.get_table_environment())
+
+ def _writers(self, dataframe):
+ return [
+ (dataframe.write_generic, ("blackhole",), {"options": {}}),
+ (dataframe.write_json, (os.path.join(self.tempdir, "json"),), {}),
+ (dataframe.write_parquet, (os.path.join(self.tempdir,
"parquet"),), {}),
+ (dataframe.write_catalog_table, ("sink_table",), {}),
+ ]
+
+ def test_writers_stage_without_execution(self):
+ dataframe = pf.from_records([(1, "a")], schema=["id", "name"])
+ for writer, arguments, options in self._writers(dataframe):
+ with self.subTest(writer=writer.__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:
+ result = writer(*arguments, statement_set=statement_set,
**options)
+
+ self.assertIsNone(result)
+ add_insert.assert_called_once()
+ self.assertIs(add_insert.call_args.args[1], dataframe._table)
+ self.assertEqual(add_insert.call_args.kwargs, {"overwrite":
False})
+ if writer.__name__ == "write_catalog_table":
+ self.assertEqual(add_insert.call_args.args[0],
"sink_table")
+ else:
+ descriptor = add_insert.call_args.args[0]
+ connector = "blackhole" if writer.__name__ ==
"write_generic" else "filesystem"
+ self.assertEqual(descriptor.get_options()["connector"],
connector)
+ execute.assert_not_called()
+ execute_insert.assert_not_called()
+
+ def test_staged_writes_preserve_overwrite_and_partition_options(self):
+ dataframe = pf.from_records([(1, "a")], schema=["id", "name"])
+ for writer in (dataframe.write_json, dataframe.write_parquet):
+ with self.subTest(writer=writer.__name__):
+ statement_set = pf.create_statement_set()
+ with patch.object(statement_set, "add_insert") as add_insert:
+ writer("/output", mode="overwrite", partition_by="name",
+ statement_set=statement_set)
+
+ descriptor = add_insert.call_args.args[0]
+ self.assertEqual(list(descriptor.get_partition_keys()),
["name"])
+ self.assertEqual(add_insert.call_args.kwargs, {"overwrite":
True})
+
+ statement_set = pf.create_statement_set()
+ with patch.object(statement_set, "add_insert") as add_insert:
+ dataframe.write_catalog_table("sink_table", overwrite=True,
+ statement_set=statement_set)
+ self.assertEqual(add_insert.call_args.kwargs, {"overwrite": True})
+
+ def test_writers_reject_invalid_statement_set(self):
+ dataframe = pf.from_records([(1, "a")], schema=["id", "name"])
+ for writer, arguments, options in self._writers(dataframe):
+ with self.subTest(writer=writer.__name__):
+ with patch.object(dataframe._table, "execute_insert") as
execute_insert:
+ with self.assertRaisesRegex(TypeError, "statement_set must
be a StatementSet"):
+ writer(*arguments, statement_set=object(), **options)
+ execute_insert.assert_not_called()
+
+ def test_writers_reject_different_environments(self):
+ dataframe = pf.from_records([(1, "a")], schema=["id", "name"])
+ other =
TableEnvironment.create(EnvironmentSettings.in_streaming_mode())
+ statement_set = other.create_statement_set()
+ for writer, arguments, options in self._writers(dataframe):
+ with self.subTest(writer=writer.__name__):
+ with patch.object(statement_set, "add_insert") as add_insert:
+ with self.assertRaisesRegex(ValueError, "same
TableEnvironment"):
+ writer(*arguments, statement_set=statement_set,
**options)
+ add_insert.assert_not_called()
+
+ def test_explain_keeps_staged_writes_without_execution(self):
+ dataframe = pf.from_records([(1,)], schema=["id"])
+ for name in ("sink_one", "sink_two"):
+ self.t_env.execute_sql(
+ f"CREATE TABLE {name} (id BIGINT) WITH ('connector' =
'blackhole')"
+ )
+ statement_set = pf.create_statement_set()
+ with patch.object(dataframe._table, "execute_insert") as
execute_insert, \
+ patch.object(statement_set, "execute") as execute:
+ dataframe.write_catalog_table("sink_one",
statement_set=statement_set)
+ dataframe.write_catalog_table("sink_two",
statement_set=statement_set)
+ for _ in range(2):
+ plan = statement_set.explain()
+ self.assertIn("sink_one", plan)
+ self.assertIn("sink_two", plan)
+
+ execute_insert.assert_not_called()
+ execute.assert_not_called()
+
+ def test_staged_catalog_errors_use_existing_translation(self):
+ dataframe = pf.from_records([(1,)], schema=["id"])
+ statement_set = pf.create_statement_set()
+ with self.assertRaises(ValueError):
+ dataframe.write_catalog_table("missing_sink",
statement_set=statement_set)
+
+
+class StatementSetITTests(PyFlinkDataFrameUTTestCase):
+ def test_multi_sink_batch_write(self):
+ self._check_multi_sink_write(RuntimeExecutionMode.BATCH)
+
+ def test_multi_sink_streaming_write(self):
+ self._check_multi_sink_write(RuntimeExecutionMode.STREAMING)
+
+ def _check_multi_sink_write(self, mode):
+ self.env.set_runtime_mode(mode)
+ self.env.set_parallelism(1)
+ settings = (EnvironmentSettings.in_batch_mode() if mode ==
RuntimeExecutionMode.BATCH
+ else EnvironmentSettings.in_streaming_mode())
+ self.t_env = StreamTableEnvironment.create(self.env,
environment_settings=settings)
+ pf.set_table_environment(self.t_env)
+ json_path = os.path.join(self.tempdir, "json-" + mode.name)
+ csv_path = os.path.join(self.tempdir, "csv-" + mode.name)
+ self.t_env.execute_sql(
+ "CREATE TABLE csv_sink (id BIGINT, name STRING) WITH ("
+ "'connector' = 'filesystem', "
+ f"'path' = '{csv_path}', 'format' = 'csv')"
+ )
+ dataframe = pf.from_records([(1, "a"), (2, "b"), (3, "c")],
schema=["id", "name"])
+ dataframe = dataframe.filter(pf.col("id") > 1)
+ statement_set = pf.create_statement_set()
+ dataframe.write_json(json_path, statement_set=statement_set)
+ dataframe.write_catalog_table("csv_sink", statement_set=statement_set)
+ self.assertFalse(os.path.exists(json_path))
+ self.assertFalse(os.path.exists(csv_path))
+
+ result = statement_set.execute()
+ self.assertIsNotNone(result.get_job_client())
+ result.wait()
+
+ json_rows = [json.loads(line) for line in self._read_lines(json_path)]
+ self.assertEqual(sorted((row["id"], row["name"]) for row in json_rows),
+ [(2, "b"), (3, "c")])
+ self.assertEqual(sorted(self._read_lines(csv_path)), ["2,b", "3,c"])
+
+ @staticmethod
+ def _read_lines(path):
+ lines = []
+ for directory, _, files in os.walk(path):
+ for name in files:
+ if not name.startswith((".", "_")):
+ with open(os.path.join(directory, name), encoding="utf-8")
as output:
+ lines.extend(line.rstrip("\n") for line in output)
+ return lines
+
+
+if __name__ == "__main__":
+ unittest.main()
