From 7becfc68638d24c9ce5789a18ceeac908fe496e4 Mon Sep 17 00:00:00 2001 From: Andrew Gazelka Date: Tue, 19 Nov 2024 22:56:42 -0800 Subject: [PATCH] [FEAT] connect: `df.get_attr` --- tests/connect/test_get_attr.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 tests/connect/test_get_attr.py diff --git a/tests/connect/test_get_attr.py b/tests/connect/test_get_attr.py new file mode 100644 index 0000000000..ea17639708 --- /dev/null +++ b/tests/connect/test_get_attr.py @@ -0,0 +1,14 @@ +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_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