From 055a478aa2cedbba66270df50f37cca53a4b4ece Mon Sep 17 00:00:00 2001 From: beetle0915 <120192315+beetle0915@users.noreply.github.com> Date: Sat, 12 Sep 2026 15:26:04 +0800 Subject: [PATCH] [FLINK-40418][python] Support attribute-based column access on DataFrame Generated-by: OpenAI Codex (codex-cli 0.154.0-alpha.6.2) --- .../reference/pyflink.dataframe/dataframe.rst | 6 + flink-python/pyflink/dataframe/dataframe.py | 41 +++++++ .../pyflink/dataframe/tests/test_dataframe.py | 111 ++++++++++++++++++ 3 files changed, 158 insertions(+) diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst index c4f36fb8ae6e48..1f8591a237ee0d 100644 --- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst +++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst @@ -25,6 +25,10 @@ Transformation methods return new DataFrames and support fluent chaining. They b plans lazily without starting a Flink job; execution is triggered by an action such as ``DataFrame.collect`` or ``DataFrame.to_pandas``. +Columns can also be referenced as attributes, such as ``df.name``, when their names are valid +Python identifiers, are not keywords, and do not conflict with existing DataFrame attributes. +Use bracket access for other names, such as ``df["select"]`` or ``df["first name"]``. + Example:: >>> import pyflink.dataframe as pf @@ -32,6 +36,7 @@ Example:: >>> result = df.select("id", "name") \ ... .with_column("id_doubled", pf.col("id") * 2) \ ... .filter(pf.col("id") > 0) + >>> names = df.select(df.name) DataFrame --------- @@ -69,6 +74,7 @@ Transformations DataFrame.offset DataFrame.head DataFrame.__getitem__ + DataFrame.__getattr__ Aggregations ------------ diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 426e3a5640d68c..904ef1bbd2d010 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -17,6 +17,7 @@ ################################################################################ import datetime +import keyword from typing import ( TYPE_CHECKING, Any, @@ -1070,6 +1071,46 @@ def __getitem__( return self.filter(key) raise TypeError("key must be a string, list, tuple, or Expression") + @PublicEvolving() + def __getattr__(self, name: str) -> Expression: + """ + Return a column expression for an attribute name. + + The name must be a valid Python identifier, must not be a Python keyword, and must + identify an existing column. Existing DataFrame attributes take precedence over columns. + Use ``df["column name"]`` for names that cannot be accessed as attributes, or + ``df["select"]`` for columns that conflict with existing attributes. + + This method resolves the schema without executing a Flink job. + + :param name: Name of the referenced column. + :return: An expression referencing the column. + :raises AttributeError: If the name is invalid or the column does not exist. + + Example:: + + >>> import pyflink.dataframe as pf + >>> df = pf.from_records([{"id": 1, "name": "Alice"}]) + >>> selected = df.select(df.name) + >>> filtered = df.filter(df.id > 0) + + .. versionadded:: 2.4.0 + """ + if not name.isidentifier() or keyword.iskeyword(name): + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") + + # Avoid re-entering __getattr__ if the underlying table has not been initialized. + try: + table = object.__getattribute__(self, "_table") + except AttributeError: + raise AttributeError( + f"'{type(self).__name__}' object has no attribute '{name}'" + ) from None + + if name not in table.get_resolved_schema().get_column_names(): + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") + return table_col(name) + # ======================== Composition ======================== @PublicEvolving() diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py b/flink-python/pyflink/dataframe/tests/test_dataframe.py index f104db9242f437..e3fa52dc0696b9 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py @@ -1048,6 +1048,108 @@ def test_getitem_rejects_unsupported_key(self): self.dataframe[42] +class DataFrameGetAttrTests(PyFlinkDataFrameUTTestCase): + def setUp(self): + super().setUp() + self.dataframe = pf.from_records( + [(1, "Alice"), (2, "Bob")], schema=["id", "name"] + ) + + def test_getattr_returns_column_expression(self): + self.assertIsInstance(self.dataframe.name, Expression) + self.assertEqual(str(self.dataframe.name), str(self.dataframe["name"])) + self.assert_dataframe_schema( + self.dataframe.select(self.dataframe.name), ["name"] + ) + + def test_getattr_composes_with_filter_and_select(self): + result = self.dataframe.filter(lambda df: df.id > 1).select( + self.dataframe.name, (self.dataframe.id + 1).alias("next_id") + ) + self.assert_dataframe_schema(result, ["name", "next_id"]) + + def test_missing_column_raises_attribute_error(self): + with self.assertRaisesRegex(AttributeError, "missing"): + self.dataframe.missing + self.assertFalse(hasattr(self.dataframe, "missing")) + default = object() + self.assertIs(getattr(self.dataframe, "missing", default), default) + self.assertTrue(hasattr(self.dataframe, "name")) + + def test_existing_attributes_take_precedence_over_columns(self): + names = ["select", "filter", "columns", "schema", "_table", "__class__"] + dataframe = pf.from_records([tuple(range(len(names)))], schema=names) + self.assertIs(dataframe.select.__func__, pf.DataFrame.select) + self.assertIs(dataframe.filter.__func__, pf.DataFrame.filter) + self.assertEqual(dataframe.columns, names) + self.assertIsInstance(dataframe.schema, TableSchema) + self.assertIs(dataframe._table, dataframe.to_table()) + self.assertIs(dataframe.__class__, pf.DataFrame) + for name in names: + with self.subTest(name=name): + self.assert_dataframe_schema(dataframe.select(dataframe[name]), [name]) + + def test_instance_attributes_take_precedence_over_columns(self): + marker = object() + self.dataframe.name = marker + self.assertIs(self.dataframe.name, marker) + self.assertIsInstance(self.dataframe["name"], Expression) + + def test_invalid_identifiers_and_keywords_are_not_attributes(self): + for name in ("first name", "first-name", "1name", "class", "None"): + with self.subTest(name=name): + dataframe = pf.from_records([(1,)], schema=[name]) + with self.assertRaises(AttributeError): + getattr(dataframe, name) + self.assertIsInstance(dataframe[name], Expression) + + def test_valid_identifiers_include_unicode_and_underscores(self): + for name in ("_name", "name_2", "\u540d\u5b57", "match"): + with self.subTest(name=name): + dataframe = pf.from_records([(1,)], schema=[name]) + self.assert_dataframe_schema(dataframe.select(getattr(dataframe, name)), [name]) + + def test_getattr_uses_the_transformed_schema(self): + renamed = self.dataframe.rename_columns({"name": "label"}) + self.assertIsInstance(renamed.label, Expression) + self.assertFalse(hasattr(renamed, "name")) + projected = self.dataframe.select("id") + self.assertFalse(hasattr(projected, "name")) + self.assertIsInstance(self.dataframe.name, Expression) + + +class DataFrameGetAttrValidationTests(unittest.TestCase): + def test_invalid_names_do_not_resolve_schema(self): + table = Mock() + for name in ("", "first name", "class"): + with self.subTest(name=name): + with self.assertRaises(AttributeError): + getattr(pf.DataFrame(table), name) + table.get_resolved_schema.assert_not_called() + + def test_uninitialized_dataframe_does_not_recurse(self): + dataframe = object.__new__(pf.DataFrame) + for name in ("_table", "name", "__setstate__"): + with self.subTest(name=name): + with self.assertRaises(AttributeError): + getattr(dataframe, name) + + def test_column_access_does_not_execute_a_job(self): + table = Mock() + table.get_resolved_schema.return_value.get_column_names.return_value = ["name"] + expression = object() + with patch("pyflink.dataframe.dataframe.table_col", return_value=expression) as col: + self.assertIs(pf.DataFrame(table).name, expression) + col.assert_called_once_with("name") + table.execute.assert_not_called() + + def test_schema_errors_are_not_hidden(self): + table = Mock() + table.get_resolved_schema.side_effect = RuntimeError("schema unavailable") + with self.assertRaisesRegex(RuntimeError, "schema unavailable"): + pf.DataFrame(table).name + + class DataFrameLiteralTests(PyFlinkDataFrameUTTestCase): def setUp(self): super().setUp() @@ -2221,6 +2323,15 @@ def setUp(self): self.addCleanup(pf.set_table_environment, previous_environment) self.t_env = TableEnvironment.create(EnvironmentSettings.in_batch_mode()) + def test_attribute_column_access_executes(self): + dataframe = pf.from_table(self.t_env.sql_query( + "SELECT * FROM (VALUES (1, 'Alice'), (2, 'Bob')) AS T(id, name)" + )) + result = dataframe.filter(lambda df: df.id > 1).select( + dataframe.name, (dataframe.id + 10).alias("next_id") + ) + self.assertEqual(result.collect(), [Row("Bob", 12)]) + def _ordered_dataframe(self): table = self.t_env.sql_query( "SELECT * FROM (VALUES (3, 'C'), (1, 'A'), (4, 'D'), (2, 'B')) "