Skip to content

Commit 31a7abc

Browse files
[CHORE] connect: add tests for df.take() method (Eventual-Inc#3385)
1 parent 5b36bd0 commit 31a7abc

File tree

1 file changed

+19
-0
lines changed

1 file changed

+19
-0
lines changed

tests/connect/test_take.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
from __future__ import annotations
2+
3+
4+
def test_take(spark_session):
5+
# Create DataFrame with 10 rows
6+
df = spark_session.range(10)
7+
8+
# Take first 5 rows and collect
9+
result = df.take(5)
10+
11+
# Verify the expected values
12+
expected = df.limit(5).collect()
13+
14+
assert result == expected
15+
16+
# Test take with more rows than exist
17+
result_large = df.take(20)
18+
expected_large = df.collect()
19+
assert result_large == expected_large # Should return all existing rows

0 commit comments

Comments
 (0)