-
Notifications
You must be signed in to change notification settings - Fork 175
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
14228e3
commit 9c90503
Showing
2 changed files
with
37 additions
and
13 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
from __future__ import annotations | ||
|
||
# import time | ||
import pytest | ||
from pyspark.sql import SparkSession | ||
|
||
|
||
@pytest.fixture | ||
def spark_session(): | ||
"""Fixture to create and clean up a Spark session.""" | ||
from daft.daft import connect_start | ||
|
||
# Start Daft Connect server | ||
server = connect_start("sc://localhost:50051") | ||
|
||
# Initialize Spark Connect session | ||
session = SparkSession.builder.appName("DaftConfigTest").remote("sc://localhost:50051").getOrCreate() | ||
|
||
yield session | ||
|
||
# Cleanup | ||
server.shutdown() | ||
session.stop() | ||
# time.sleep(2) # Allow time for session cleanup | ||
|
||
|
||
def test_range_first(spark_session): | ||
spark_range = spark_session.range(10) | ||
first_row = spark_range.first() | ||
assert first_row["id"] == 0, "First row should have id=0" | ||
|
||
|
||
def test_range_limit(spark_session): | ||
spark_range = spark_session.range(10) | ||
limited_df = spark_range.limit(5).toPandas() | ||
assert len(limited_df) == 5, "Limited DataFrame should have 5 rows" | ||
assert list(limited_df["id"]) == list(range(5)), "Limited DataFrame should contain values 0-4" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters