from __future__ import annotations
def test_union_accepts_dataframe(spark) -> None:
a = spark.createDataFrame([(1, "x")], ["id", "label"])
b = spark.createDataFrame([(2, "y")], ["id", "label"])
out = a.union(b)
rows = out.collect()
assert len(rows) == 2
assert rows[0]["id"] == 1 and rows[0]["label"] == "x"
assert rows[1]["id"] == 2 and rows[1]["label"] == "y"
def test_union_all_accepts_dataframe(spark) -> None:
a = spark.createDataFrame([(1,)], ["v"])
b = spark.createDataFrame([(2,)], ["v"])
out = a.unionAll(b)
rows = out.collect()
assert len(rows) == 2
assert rows[0]["v"] == 1 and rows[1]["v"] == 2