Skip to content

Commit

Permalink
[TEST] connect: df.dtypes
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewgazelka committed Dec 4, 2024
1 parent 62ae066 commit f1f2bfc
Showing 1 changed file with 27 additions and 0 deletions.
27 changes: 27 additions & 0 deletions tests/connect/test_dtypes.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
from __future__ import annotations

from pyspark.sql.functions import col


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

# Add a column that will have repeated values for grouping
df = df.withColumn("group", col("id") % 3)

# Check dtypes of the DataFrame
expected_dtypes = [
("id", "bigint"),
("group", "bigint")
]

# Get actual dtypes
actual_dtypes = df.dtypes

# Verify the dtypes match expected
assert actual_dtypes == expected_dtypes

# Also check individual column types
assert df.schema["id"].dataType.simpleString() == "bigint"
assert df.schema["group"].dataType.simpleString() == "bigint"

0 comments on commit f1f2bfc

Please sign in to comment.