dianfu commented on code in PR #28797: URL: https://github.com/apache/flink/pull/28797#discussion_r3636645007
########## flink-python/pyflink/dataframe/tests/test_dataframe.py: ########## @@ -0,0 +1,400 @@ +################################################################################ +# 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 unittest +from typing import NamedTuple + +import pyflink.dataframe as pf +from py4j.protocol import Py4JJavaError +from pyflink.common import Row +from pyflink.table import DataTypes as TableDataTypes +from pyflink.table.expression import Expression +from pyflink.testing.test_case_utils import PyFlinkStreamDataFrameTestCase + +_PREEXISTING_TABLE_ENVIRONMENT = object() + + +class _Point(NamedTuple): + x: int + y: int + + +class _OtherPoint(NamedTuple): + x: int + z: int + + +class _CloseableIterator: + def __init__(self, values=None, error=None): + self._values = iter(values or []) + self._error = error + self.closed = False + + def __iter__(self): + return self + + def __next__(self): + if self._error is not None: + raise self._error + return next(self._values) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.closed = True + + +class _TableResult: + def __init__(self, iterator): + self._iterator = iterator + + def collect(self): + return self._iterator + + +class _Table: + def __init__(self, iterator): + self._iterator = iterator + + def execute(self): + return _TableResult(self._iterator) + + +class DataFrameCollectTests(unittest.TestCase): + def test_collect_returns_all_rows_and_closes_iterator(self): + iterator = _CloseableIterator([Row(1, "Alice")]) + dataframe = pf.DataFrame(_Table(iterator)) + + self.assertEqual(dataframe.collect(), [Row(1, "Alice")]) + self.assertTrue(iterator.closed) + + def test_collect_closes_iterator_when_iteration_fails(self): + iterator = _CloseableIterator(error=RuntimeError("iteration failed")) + dataframe = pf.DataFrame(_Table(iterator)) + + with self.assertRaisesRegex(RuntimeError, "iteration failed"): + dataframe.collect() + + self.assertTrue(iterator.closed) + + +class DataFrameValidationTests(unittest.TestCase): + def setUp(self): + self.dataframe = pf.DataFrame(_Table(_CloseableIterator())) + + def test_filter_rejects_non_expression(self): + with self.assertRaisesRegex(TypeError, "predicate must be an Expression"): + self.dataframe.filter(True) + + def test_filter_requires_a_condition(self): + with self.assertRaisesRegex(ValueError, "requires at least one predicate"): + self.dataframe.filter() + + def test_filter_rejects_callable_returning_non_expression(self): + with self.assertRaisesRegex( + TypeError, "callable predicates must return an Expression" + ): + self.dataframe.filter(lambda df: True) + + def test_filter_rejects_expression_class(self): + with self.assertRaisesRegex(TypeError, "predicate must be an Expression"): + self.dataframe.filter(Expression) + + def test_with_column_rejects_non_expression(self): + with self.assertRaisesRegex(TypeError, "expr must be an Expression"): + self.dataframe.with_column("answer", 42) + + def test_with_column_rejects_callable_returning_non_expression(self): + with self.assertRaisesRegex(TypeError, "expr must be an Expression"): + self.dataframe.with_column("answer", lambda df: 42) + + def test_with_column_rejects_expression_class(self): + with self.assertRaisesRegex(TypeError, "expr must be an Expression"): + self.dataframe.with_column("answer", Expression) + + def test_with_column_rejects_non_string_name(self): + with self.assertRaisesRegex(TypeError, "name must be a string"): + self.dataframe.with_column(42, object()) + + def test_select_rejects_non_string_column(self): + with self.assertRaisesRegex(TypeError, "columns must be strings"): + self.dataframe.select("id", 42) + + def test_select_rejects_non_expression_projection(self): + with self.assertRaisesRegex(TypeError, "projections must be expressions"): + self.dataframe.select(answer=42) + + def test_getitem_rejects_unsupported_key(self): + with self.assertRaisesRegex(TypeError, "key must be a string, list"): + self.dataframe[42] + + def test_lit_rejects_non_dataframe_data_type(self): + with self.assertRaisesRegex( + TypeError, "data_type must be a pyflink.dataframe.DataType" + ): + pf.lit(1, object()) + + +class DataFrameEndToEndTests(PyFlinkStreamDataFrameTestCase): + @classmethod + def setUpClass(cls): + pf.set_table_environment(_PREEXISTING_TABLE_ENVIRONMENT) + super(DataFrameEndToEndTests, cls).setUpClass() + + @classmethod + def tearDownClass(cls): + try: + super(DataFrameEndToEndTests, cls).tearDownClass() + if pf.get_table_environment() is not _PREEXISTING_TABLE_ENVIRONMENT: + raise AssertionError("the previous table environment was not restored") + finally: + pf.set_table_environment(None) + + def test_uses_current_test_environment(self): + self.assertIs(pf.get_table_environment(), self.t_env) + + def test_complete_minimal_job(self): + df = pf.from_records( + [ + (1, "Alice", 30), + (2, "Bob", 17), + ], + schema=["id", "name", "age"], + ) + + result = ( + df.filter(pf.col("age") >= 18) + .with_column("age_next_year", pf.col("age") + 1) + .select("id", "name", "age_next_year") + ) + + self.assertEqual(result.collect(), [Row(1, "Alice", 31)]) + + def test_from_dict_respects_schema_order_and_subset(self): Review Comment: It has introduced about 30 it cases. Please check if we could merge them. Besides, it's also worth to check if it's possible to convert part of them into unit tests. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
