Skip to content

Commit

Permalink
[CHORE] connect: add tests for df.take() method (#3385)
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewgazelka authored Nov 23, 2024
1 parent 5b36bd0 commit 31a7abc
Showing 1 changed file with 19 additions and 0 deletions.
19 changes: 19 additions & 0 deletions tests/connect/test_take.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 31a7abc

Please sign in to comment.