Skip to content

Commit

Permalink
[CHORE] connect test: df.get_attr (#3349)
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewgazelka authored Nov 26, 2024
1 parent b99038e commit b201c61
Showing 1 changed file with 15 additions and 0 deletions.
15 changes: 15 additions & 0 deletions tests/connect/test_get_attr.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from __future__ import annotations


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

# Get column using df[...]
# df.get_attr("id") is equivalent to df["id"]
df_col = df["id"]

# Check that column values match expected range
values = df.select(df_col).collect() # Changed to select column first
assert len(values) == 10
assert [row[0] for row in values] == list(range(10)) # Need to extract values from Row objects

0 comments on commit b201c61

Please sign in to comment.