diff --git a/tests/connect/test_take.py b/tests/connect/test_take.py new file mode 100644 index 0000000000..2e7809f232 --- /dev/null +++ b/tests/connect/test_take.py @@ -0,0 +1,19 @@ +from __future__ import annotations + + +def test_take(spark_session): + # Create DataFrame with 10 rows + df = spark_session.range(10) + + # Take first 5 rows and collect + result = df.take(5) + + # Verify the expected values + expected = df.limit(5).collect() + + assert result == expected + + # Test take with more rows than exist + result_large = df.take(20) + expected_large = df.collect() + assert result_large == expected_large # Should return all existing rows