From 447edf2902569708e01a7b956b732fbf2bc29907 Mon Sep 17 00:00:00 2001 From: beetle0915 <120192315+beetle0915@users.noreply.github.com> Date: Tue, 15 Sep 2026 23:39:17 +0800 Subject: [PATCH] [FLINK-40435][python] Add explode support to DataFrame API Generated-by: OpenAI Codex CLI 0.154.0-alpha.6.2 --- .../reference/pyflink.dataframe/dataframe.rst | 1 + flink-python/pyflink/dataframe/dataframe.py | 99 +++++++ .../pyflink/dataframe/tests/test_dataframe.py | 265 +++++++++++++++++- 3 files changed, 364 insertions(+), 1 deletion(-) diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst index e78bbaf0d93c5..cd6562d32f56c 100644 --- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst +++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst @@ -60,6 +60,7 @@ Transformations DataFrame.rename DataFrame.filter DataFrame.where + DataFrame.explode DataFrame.drop_duplicates DataFrame.distinct DataFrame.unique diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 2e1b8fd8208e9..fca652796d1d8 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -46,6 +46,7 @@ lit as table_lit, ) from pyflink.table.table import Table +from pyflink.table.types import ArrayType, MapType, MultisetType, RowType from pyflink.util.api_stability_decorators import PublicEvolving __all__ = ["DataFrame", "GroupedDataFrame", "col", "lit"] @@ -622,6 +623,104 @@ def top_n( distinct = drop_duplicates unique = drop_duplicates + @PublicEvolving() + def explode( + self, + column: Union[str, Expression], + output_column: Optional[Union[str, List[str]]] = None, + ignore_empty_and_null: bool = False, + ) -> "DataFrame": + """ + Expand an ARRAY, MAP, or MULTISET into rows, preserving duplicate occurrences. + + A referenced input column is removed. Other input columns are retained, followed by + the expanded fields. For a computed collection expression, all input columns are retained. + MAP values yield key and value fields; ROW elements yield one field per ROW field. + Empty and null collections produce a row with null output fields unless + ``ignore_empty_and_null`` is true. + + :param column: Collection column name or row-wise expression to expand. + :param output_column: Output name or list of names. Required for multiple fields; + a single field defaults to the selected column name. Names must be unique and must + not conflict with retained input columns. + :param ignore_empty_and_null: Whether to drop rows with empty or null collections. + :return: A new DataFrame with the expanded rows. + :raises TypeError: If an argument has an unsupported type or the input is not a collection. + :raises ValueError: If the expression is not row-wise, selects multiple columns, + or output names are invalid. + + Example:: + + >>> import pyflink.dataframe as pf + >>> df = pf.from_dict({"id": [1, 2], "tags": [["a", "b"], []]}) + >>> result = df.explode("tags") + >>> result = df.explode("tags", "tag", ignore_empty_and_null=True) + + .. versionadded:: 2.4.0 + """ + if not isinstance(column, (str, Expression)): + raise TypeError("column must be a column name or expression") + if not isinstance(ignore_empty_and_null, bool): + raise TypeError("ignore_empty_and_null must be a boolean") + if output_column is not None and not isinstance(output_column, (str, list)): + raise TypeError("output_column must be a string or list of strings") + + expression = table_col(column) if isinstance(column, str) else column + selected = self._table.select(expression) + schema = selected.get_resolved_schema() + if len(schema.get_column_names()) != 1: + raise ValueError("column must select a single column") + projection = selected._j_table.getQueryOperation() + # Aggregates insert an intermediate operation whose field indexes refer to its result. + if not projection.getChildren().get(0).equals(self._table._j_table.getQueryOperation()): + raise ValueError("column must be a row-wise expression, not an aggregation") + data_type = schema.get_column_data_types()[0] + if isinstance(data_type, MapType): + field_count = 2 + elif isinstance(data_type, (ArrayType, MultisetType)): + element_type = data_type.element_type + field_count = len(element_type.fields) if isinstance(element_type, RowType) else 1 + else: + raise TypeError("column must have an ARRAY, MAP, or MULTISET type") + + if output_column is None: + if field_count != 1: + raise ValueError("output_column is required for multiple output fields") + output_names = [schema.get_column_names()[0]] + else: + output_names = [output_column] if isinstance(output_column, str) else output_column + if not all(isinstance(name, str) for name in output_names): + raise TypeError("output_column must contain only strings") + if len(output_names) != field_count: + raise ValueError("output_column must contain %d name(s)" % field_count) + if any(not name for name in output_names) or len(set(output_names)) != field_count: + raise ValueError("output_column names must be non-empty and unique") + + table = self._table + columns = list(table.get_resolved_schema().get_column_names()) + resolved = projection.getProjectList().get(0) + # Resolve the input field by index so an alias does not hide the column to remove. + if resolved.getClass().getSimpleName() == "FieldReferenceExpression": + collection_name = columns.pop(resolved.getFieldIndex()) + else: + taken = set(columns) | set(output_names) + collection_name = _unique_name("__pf_explode", taken) + table = table.add_columns(expression.alias(collection_name)) + if set(output_names).intersection(columns): + raise ValueError("output_column names conflict with retained input columns") + + projections = ["src." + _quote_identifier(name) for name in columns] + projections.extend("expanded." + _quote_identifier(name) for name in output_names) + query = "SELECT %s FROM %s AS src %s UNNEST(src.%s) AS expanded(%s)%s" % ( + ", ".join(projections), + _quote_identifier(str(table)), + "CROSS JOIN" if ignore_empty_and_null else "LEFT JOIN", + _quote_identifier(collection_name), + ", ".join(_quote_identifier(name) for name in output_names), + "" if ignore_empty_and_null else " ON TRUE", + ) + return DataFrame(table._t_env.sql_query(query)) + # ======================== Filtering & Ordering ======================== @PublicEvolving() diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py b/flink-python/pyflink/dataframe/tests/test_dataframe.py index 736fdbc17d763..9e272205132e5 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -23,13 +23,14 @@ import pandas as pd import pyarrow as pa import unittest +from collections import Counter from datetime import date, datetime, time, timedelta, timezone from py4j.protocol import Py4JJavaError from typing import NamedTuple from unittest.mock import Mock, patch import pyflink.dataframe as pf -from pyflink.common import Row +from pyflink.common import Row, RowKind from pyflink.dataframe.dataframe import ( _resolve_window_time_column, _to_interval_expression, @@ -38,6 +39,7 @@ from pyflink.table import ( DataTypes as TableDataTypes, EnvironmentSettings, + Table, TableEnvironment, TableSchema, ) @@ -332,6 +334,101 @@ def test_set_operations_preserve_streaming_restrictions(self): getattr(left, method)(right) +class DataFrameExplodeTests(PyFlinkDataFrameUTTestCase): + def setUp(self): + super().setUp() + self.df = pf.from_table(self.t_env.sql_query( + "SELECT 1 AS id, ARRAY[1, 2] AS items, 'A' AS label")) + + def test_explode_is_lazy_and_preserves_input(self): + table = self.df.to_table() + with patch.object(Table, "execute") as execute: + result = self.df.explode("items") + execute.assert_not_called() + self.assertIsNot(result, self.df) + self.assertIs(result.to_table()._t_env, self.t_env) + self.assertIs(self.df.to_table(), table) + self.assert_dataframe_schema(self.df, ["id", "items", "label"]) + self.assert_dataframe_schema(result, ["id", "label", "items"], [ + TableDataTypes.INT().not_null(), TableDataTypes.CHAR(1).not_null(), + TableDataTypes.INT(), + ]) + + def test_explode_supports_column_expressions_and_aliases(self): + for column, output, expected in [ + ("items", None, "items"), + (pf.col("items"), None, "items"), + (self.df["items"], "item", "item"), + (pf.col("items").alias("renamed"), None, "renamed"), + ("items", ["item"], "item"), + ]: + with self.subTest(column=str(column), output=output): + self.assert_dataframe_schema( + self.df.explode(column, output), ["id", "label", expected]) + + def test_explode_supports_computed_collection_expressions(self): + from pyflink.table.expressions import array + + result = self.df.explode(array(pf.col("id"), pf.lit(2)), "value") + self.assert_dataframe_schema(result, ["id", "items", "label", "value"]) + + def test_explode_resolves_collection_output_types(self): + for sql, names, types in [ + ("SELECT MAP['a', 1] AS items", ["key", "value"], + [TableDataTypes.CHAR(1), TableDataTypes.INT()]), + ("SELECT ARRAY[ROW(1, 'a')] AS items", ["number", "text"], + [TableDataTypes.INT(), TableDataTypes.CHAR(1)]), + ("SELECT ARRAY[ROW(1)] AS items", None, [TableDataTypes.INT()]), + ("SELECT MULTISET[1, 1, 2] AS items", None, [TableDataTypes.INT()]), + ]: + with self.subTest(sql=sql): + df = pf.from_table(self.t_env.sql_query(sql)) + self.assert_dataframe_schema( + df.explode("items", names), names or ["items"], types) + + def test_explode_rejects_invalid_arguments(self): + for column in [None, 1, ["items"]]: + with self.subTest(column=column): + with self.assertRaisesRegex(TypeError, "column"): + self.df.explode(column) + for flag in [None, 1, "true"]: + with self.subTest(flag=flag): + with self.assertRaisesRegex(TypeError, "ignore_empty_and_null"): + self.df.explode("items", ignore_empty_and_null=flag) + for output in [1, ("item",), [1]]: + with self.subTest(output=output): + with self.assertRaisesRegex(TypeError, "output_column"): + self.df.explode("items", output) + + def test_explode_rejects_missing_and_non_collection_columns(self): + for column in ["missing", pf.col("missing")]: + with self.subTest(column=str(column)): + with self.assertRaisesRegex(Py4JJavaError, "missing"): + self.df.explode(column) + with self.assertRaisesRegex(TypeError, "ARRAY, MAP, or MULTISET"): + self.df.explode("id") + with self.assertRaisesRegex(ValueError, "single column"): + self.df.explode(pf.col("*")) + + def test_explode_rejects_aggregate_expressions(self): + df = pf.from_table(self.t_env.sql_query("SELECT ARRAY[9] AS items, 1 AS id")) + with self.assertRaisesRegex(ValueError, "row-wise"): + df.explode(pf.col("id").collect, "value") + + def test_explode_rejects_invalid_output_names(self): + for output in [[], ["a", "b"], "", [""]]: + with self.subTest(output=output): + with self.assertRaisesRegex(ValueError, "output_column"): + self.df.explode("items", output) + with self.assertRaisesRegex(ValueError, "conflict"): + self.df.explode("items", "id") + df = pf.from_table(self.t_env.sql_query("SELECT MAP['a', 1] AS items")) + for output in [None, "item", ["item"], ["item", "item"]]: + with self.subTest(output=output): + with self.assertRaisesRegex(ValueError, "output_column"): + df.explode("items", output) + + class DataFrameSortingTests(PyFlinkDataFrameUTTestCase): def setUp(self): super().setUp() @@ -2543,6 +2640,172 @@ def test_union_all_retains_duplicates(self): ) +class DataFrameExplodeITTests(PyFlinkITTestCase): + def setUp(self): + self.t_env = TableEnvironment.create(EnvironmentSettings.in_batch_mode()) + + def test_explode_arrays_preserves_duplicates_and_null_elements(self): + df = pf.from_table(self.t_env.from_elements( + [(1, [2, 2, None]), (2, []), (3, None)], + TableDataTypes.ROW([ + TableDataTypes.FIELD("id", TableDataTypes.INT()), + TableDataTypes.FIELD("items", TableDataTypes.ARRAY(TableDataTypes.INT())), + ]))) + for ignore in [False, True]: + with self.subTest(ignore=ignore): + expected = [Row(1, 2), Row(1, 2), Row(1, None)] + if not ignore: + expected += [Row(2, None), Row(3, None)] + self.assertCountEqual( + df.explode("items", ignore_empty_and_null=ignore).collect(), expected) + + def test_explode_maps(self): + df = pf.from_table(self.t_env.from_elements( + [(1, {"a": 2, "b": 3}), (2, {}), (3, None)], + TableDataTypes.ROW([ + TableDataTypes.FIELD("id", TableDataTypes.INT()), + TableDataTypes.FIELD("items", TableDataTypes.MAP( + TableDataTypes.STRING(), TableDataTypes.INT())), + ]))) + for ignore in [False, True]: + with self.subTest(ignore=ignore): + expected = [Row(1, "a", 2), Row(1, "b", 3)] + if not ignore: + expected += [Row(2, None, None), Row(3, None, None)] + self.assertCountEqual( + df.explode("items", ["key", "value"], ignore).collect(), expected) + nullable_values = pf.from_table(self.t_env.sql_query( + "SELECT MAP['a', CAST(NULL AS INT)] AS items")) + self.assertEqual( + nullable_values.explode("items", ["key", "value"], True).collect(), + [Row("a", None)]) + + def test_explode_multisets(self): + df = pf.from_table(self.t_env.sql_query( + "SELECT id, COLLECT(item) AS items FROM " + "(VALUES (1, 2), (1, 2), (2, CAST(NULL AS INT))) AS T(id, item) GROUP BY id " + "UNION ALL SELECT 3, CAST(NULL AS INT MULTISET)")) + for ignore in [False, True]: + with self.subTest(ignore=ignore): + expected = [Row(1, 2), Row(1, 2)] + if not ignore: + expected += [Row(2, None), Row(3, None)] + self.assertCountEqual( + df.explode("items", ignore_empty_and_null=ignore).collect(), expected) + + def test_explode_row_elements(self): + array_table = self.t_env.from_elements( + [(1, [Row(2, "a"), Row(2, "a")]), (2, []), (3, None)], + TableDataTypes.ROW([ + TableDataTypes.FIELD("id", TableDataTypes.INT()), + TableDataTypes.FIELD("items", TableDataTypes.ARRAY(TableDataTypes.ROW([ + TableDataTypes.FIELD("n", TableDataTypes.INT()), + TableDataTypes.FIELD("s", TableDataTypes.STRING()), + ]))), + ])) + multiset_table = self.t_env.sql_query( + "SELECT id, COLLECT(item) AS items FROM " + "(VALUES (1, ROW(2, 'a')), (1, ROW(2, 'a')), " + "(2, CAST(NULL AS ROW))) AS T(id, item) GROUP BY id " + "UNION ALL SELECT 3, CAST(NULL AS ROW MULTISET)") + for table in [array_table, multiset_table]: + df = pf.from_table(table) + for ignore in [False, True]: + with self.subTest(schema=str(table.get_resolved_schema()), ignore=ignore): + expected = [Row(1, 2, "a"), Row(1, 2, "a")] + if not ignore: + expected += [Row(2, None, None), Row(3, None, None)] + self.assertCountEqual( + df.explode("items", ["number", "text"], ignore).collect(), expected) + + def test_explode_computed_expressions(self): + from pyflink.table.expressions import array + + df = pf.from_table(self.t_env.sql_query("SELECT 1 AS __pf_explode")) + result = df.explode(array(pf.col("__pf_explode"), pf.lit(2)).alias("pair"), "value") + self.assertEqual(result.columns, ["__pf_explode", "value"]) + self.assertCountEqual(result.collect(), [Row(1, 1), Row(1, 2)]) + + def test_explode_after_projection_and_with_aliased_column(self): + df = pf.from_table(self.t_env.sql_query( + "SELECT 1 AS id, ARRAY[2, 3] AS items")) + projected = df.select(pf.col("items").alias("values"), "id") + result = projected.explode(pf.col("values").alias("value")) + self.assertEqual(result.columns, ["id", "value"]) + self.assertCountEqual(result.collect(), [Row(1, 2), Row(1, 3)]) + + def test_explode_quotes_identifiers(self): + df = pf.from_table(self.t_env.sql_query( + "SELECT 1 AS `__pf_explode`, 2 AS `select`, ARRAY[3, 4] AS `a``b`")) + result = df.explode(pf.col("a`b"), "value` name") + self.assertEqual(result.columns, ["__pf_explode", "select", "value` name"]) + self.assertCountEqual(result.collect(), [Row(1, 2, 3), Row(1, 2, 4)]) + + def test_explode_collection_only_and_empty_inputs(self): + df = pf.from_table(self.t_env.sql_query("SELECT ARRAY[ROW(1), ROW(2)] AS items")) + self.assertCountEqual(df.explode("items").collect(), [Row(1), Row(2)]) + for ignore in [False, True]: + with self.subTest(ignore=ignore): + self.assertEqual( + df.filter(pf.lit(False)).explode( + "items", ignore_empty_and_null=ignore).collect(), []) + + def test_explode_composes_with_filter_and_aggregation(self): + df = pf.from_table(self.t_env.sql_query( + "SELECT * FROM (VALUES (1, ARRAY[2, 2, 3]), (2, ARRAY[3])) AS T(id, items)")) + result = (df.explode("items", "value") + .filter(pf.col("value") > 2) + .group_by("value") + .agg(total=pf.col("id").count)) + self.assertEqual(result.collect(), [Row(3, 2)]) + + +class DataFrameExplodeStreamITTests(PyFlinkStreamDataFrameTestCase): + def test_explode_arrays_in_streaming_mode(self): + df = pf.from_table(self.t_env.from_elements( + [(1, [2, 2]), (2, []), (3, None)], + TableDataTypes.ROW([ + TableDataTypes.FIELD("id", TableDataTypes.INT()), + TableDataTypes.FIELD("items", TableDataTypes.ARRAY(TableDataTypes.INT())), + ]))) + for ignore in [False, True]: + with self.subTest(ignore=ignore): + expected = [Row(1, 2), Row(1, 2)] + if not ignore: + expected += [Row(2, None), Row(3, None)] + self.assertCountEqual( + df.explode("items", ignore_empty_and_null=ignore).collect(), expected) + + def test_explode_maps_and_row_elements_in_streaming_mode(self): + for sql, outputs, expected in [ + ("SELECT MAP['a', 1, 'b', 2] AS items", ["key", "value"], + [Row("a", 1), Row("b", 2)]), + ("SELECT ARRAY[ROW(1, 'a'), ROW(2, 'b')] AS items", ["n", "s"], + [Row(1, "a"), Row(2, "b")]), + ]: + for ignore in [False, True]: + with self.subTest(sql=sql, ignore=ignore): + df = pf.from_table(self.t_env.sql_query(sql)) + self.assertCountEqual(df.explode("items", outputs, ignore).collect(), expected) + + def test_explode_multiset_changelog_preserves_multiplicities(self): + df = pf.from_table(self.t_env.sql_query( + "SELECT id, COLLECT(item) AS items FROM " + "(VALUES (1, 2), (1, 2), (2, CAST(NULL AS INT))) AS T(id, item) GROUP BY id")) + for ignore in [False, True]: + with self.subTest(ignore=ignore): + counts = Counter() + for row in df.explode("items", ignore_empty_and_null=ignore).collect(): + if row.get_row_kind() in (RowKind.INSERT, RowKind.UPDATE_AFTER): + counts[tuple(row)] += 1 + else: + counts[tuple(row)] -= 1 + expected = {(1, 2): 2} + if not ignore: + expected[(2, None)] = 1 + self.assertEqual({row: count for row, count in counts.items() if count}, expected) + + class DataFrameWindowITTests(PyFlinkStreamDataFrameTestCase): @classmethod def setUpClass(cls):