From 099354adfb73f4e89b194e69e7b27e2a3171a0fb Mon Sep 17 00:00:00 2001 From: Andrew Gazelka Date: Mon, 18 Nov 2024 13:49:32 -0800 Subject: [PATCH] feat: connect: collect --- tests/connect/test_collect.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 tests/connect/test_collect.py diff --git a/tests/connect/test_collect.py b/tests/connect/test_collect.py new file mode 100644 index 0000000000..0a9387dd0b --- /dev/null +++ b/tests/connect/test_collect.py @@ -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"