Skip to content

Commit

Permalink
[FEAT] connect: collect (#3326)
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewgazelka authored Nov 21, 2024
1 parent 5fee192 commit d9d916b
Showing 1 changed file with 14 additions and 0 deletions.
14 changes: 14 additions & 0 deletions tests/connect/test_collect.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
from __future__ import annotations


def test_range_collect(spark_session):
# Create a range using Spark
# For example, creating a range from 0 to 9
spark_range = spark_session.range(10) # Creates DataFrame with numbers 0 to 9

# Collect the data
collected_rows = spark_range.collect()

# Verify the collected data has expected values
assert len(collected_rows) == 10, "Should have 10 rows"
assert [row["id"] for row in collected_rows] == list(range(10)), "Should contain values 0-9"

0 comments on commit d9d916b

Please sign in to comment.