Skip to content

Commit

Permalink
[FEAT]: connect: df.where
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewgazelka committed Nov 21, 2024
1 parent ffc943b commit 5a5553d
Showing 1 changed file with 16 additions and 0 deletions.
16 changes: 16 additions & 0 deletions tests/connect/test_where.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
from __future__ import annotations

from pyspark.sql.functions import col


def test_where(spark_session):
# Create DataFrame from range(10)
df = spark_session.range(10)

# Filter the DataFrame where 'id' is greater than 5
df_filtered = df.where(col("id") > 5)

# Verify the filter was applied correctly by checking the expected data
df_filtered_pandas = df_filtered.toPandas()
expected_data = [6, 7, 8, 9]
assert df_filtered_pandas["id"].tolist() == expected_data, "Filtered data does not match expected data"

0 comments on commit 5a5553d

Please sign in to comment.