This is an automated email from the ASF dual-hosted git repository.
uros-b pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 4ecd8768731b [SPARK-55886][PYTHON][TEST][FOLLOWUP] Add tests for
DataFrame.col resolution through DataFrame.zip
4ecd8768731b is described below
commit 4ecd8768731b81d8fc626ec578ee95c59475b315
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Sat Jun 20 01:24:02 2026 +0200
[SPARK-55886][PYTHON][TEST][FOLLOWUP] Add tests for DataFrame.col
resolution through DataFrame.zip
### What changes were proposed in this pull request?
This PR adds PySpark tests for resolving a DataFrame-scoped column
reference (e.g. `df1.zip(df2).select(df1.some_col)`) through `DataFrame.zip`.
The tests live in the column test suites, next to the existing `test_resolve_*`
resolve-through-operator tests, named `test_resolve_after_zip*` following the
`after_*` convention used for binary combinators (union, intersect):
- `python/pyspark/sql/tests/test_column.py` (`ColumnTestsMixin`): 9 tests
asserting the Classic behavior - selecting each side by its originating
DataFrame, reordering, disambiguating duplicate column names, shared-producer
dedup, use in an expression, use in a filter, through chained projections, and
two base-side cases (`df.zip(right).select(df.a)`).
- `python/pyspark/sql/tests/connect/test_parity_column.py`
(`ColumnParityTests`): 8 overrides asserting the Spark Connect behavior.
On Classic, `df.col` resolves by attribute id, which `ResolveZip` preserves
in the merged `Project`, so the reference resolves. On Spark Connect, `df.col`
resolves by plan id; `ResolveZip` merges the two sides into a single plan and
the per-DataFrame plan-id tags do not survive, so the reference raises
`CANNOT_RESOLVE_DATAFRAME_COLUMN` in both strict and lenient resolution modes.
This mirrors the existing `test_resolve_after_union` divergence, where a tagged
reference also cannot be f [...]
Note on the lenient mode
(`spark.sql.analyzer.strictDataFrameColumnResolution=false`): one might expect
the name-based fallback to resolve the reference there, since the zip output
does contain a column with the requested name. It does not: when the tagged
plan id is not found anywhere in the subtree, `resolveDataFrameColumn` throws
`CANNOT_RESOLVE_DATAFRAME_COLUMN` before any name-based fallback is attempted
(same as the `test_resolve_after_union` case). The overrides therefore live
[...]
Note on base-side references (`df.zip(right).select(df.a)`): on Connect the
outcome depends on the base plan's root node, i.e. the node carrying the
DataFrame's plan-id tag. `createDataFrame` analyzes to a `Project` over a
`LocalRelation`; that tagged `Project` is dissolved into the merged chain by
`ResolveZip`, so the reference raises like the projected-side cases
(`test_resolve_after_zip_base_side`, overridden in the parity suite). A base
whose root is the relation itself (`range` a [...]
### Why are the changes needed?
`DataFrame.zip` is a newly added API (SPARK-55886), and its interaction
with DataFrame-scoped column references (`DataFrame.col`) was not covered by
tests. These tests document the supported Classic behavior and pin down the
current Spark Connect limitation, so any future change in resolution behavior
is caught.
### Does this PR introduce _any_ user-facing change?
No. This is a test-only change.
### How was this patch tested?
New unit tests, run locally:
- `pyspark.sql.tests.test_column` (Classic `ColumnTests`) - the 9 new tests
pass.
- `pyspark.sql.tests.connect.test_parity_column` (`ColumnParityTests` and
`ColumnParityTestsWithNonStrictDFColResolution`) - all 9 pass under both strict
and lenient DataFrame column resolution modes (18 invocations: 8 overrides plus
the inherited bare-base test).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 4.8
Closes #56398 from zhengruifeng/zip-df-col-tests-dev7.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Uros Bojanic <[email protected]>
(cherry picked from commit 6d75d7ae3a4b278891a9eab53ebd61cc384d7bd1)
Signed-off-by: Uros Bojanic <[email protected]>
---
.../sql/tests/connect/test_parity_column.py | 68 +++++++++++
python/pyspark/sql/tests/test_column.py | 124 +++++++++++++++++++++
2 files changed, 192 insertions(+)
diff --git a/python/pyspark/sql/tests/connect/test_parity_column.py
b/python/pyspark/sql/tests/connect/test_parity_column.py
index a2b00d7955ee..06eb26069578 100644
--- a/python/pyspark/sql/tests/connect/test_parity_column.py
+++ b/python/pyspark/sql/tests/connect/test_parity_column.py
@@ -50,6 +50,74 @@ class ColumnParityTests(ColumnTestsMixin,
ReusedConnectTestCase):
with self.assertRaisesRegex(AnalysisException,
"CANNOT_RESOLVE_DATAFRAME_COLUMN"):
df1.union(df2).select(df1.c).collect()
+ # zip merges the two column-projected sides into a single plan, so the
+ # per-DataFrame plan-id tags do not survive ResolveZip. A tagged left/right
+ # reference can no longer be found and raises in both strict and lenient
+ # modes - the throw precedes any name-based fallback - unlike Classic,
which
+ # resolves by the attribute id that ResolveZip preserves. A base-side
+ # reference (df.zip(right).select(df.a)) raises too when the base's tagged
+ # plan root is a Project that the rewrite dissolves (createDataFrame,
+ # overridden below); only a bare relation root like range survives as the
+ # merged base and resolves (test_resolve_after_zip_bare_base_side,
+ # inherited with no override).
+
+ def test_resolve_after_zip(self):
+ df = self.spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ["a",
"b"])
+ left = df.select((df.a + 1).alias("x"))
+ right = df.select((df.b * 2).alias("y"))
+ with self.assertRaisesRegex(AnalysisException,
"CANNOT_RESOLVE_DATAFRAME_COLUMN"):
+ left.zip(right).select(left.x).collect()
+
+ def test_resolve_after_zip_reordered(self):
+ df = self.spark.createDataFrame([(1, 10), (2, 20)], ["a", "b"])
+ left = df.select(df.a.alias("x"))
+ right = df.select(df.b.alias("y"))
+ with self.assertRaisesRegex(AnalysisException,
"CANNOT_RESOLVE_DATAFRAME_COLUMN"):
+ left.zip(right).select(right.y, left.x).collect()
+
+ def test_resolve_after_zip_duplicate_names(self):
+ df = self.spark.createDataFrame([(1, 2), (3, 4)], ["a", "b"])
+ left = df.select(df.a.alias("v"))
+ right = df.select(df.b.alias("v"))
+ with self.assertRaisesRegex(AnalysisException,
"CANNOT_RESOLVE_DATAFRAME_COLUMN"):
+ left.zip(right).select(left.v).collect()
+
+ def test_resolve_after_zip_shared_producer(self):
+ df = self.spark.createDataFrame([(1, 2), (3, 4)], ["a", "b"])
+ left = df.select(df.a.alias("v"))
+ right = df.select(df.a.alias("v"))
+ with self.assertRaisesRegex(AnalysisException,
"CANNOT_RESOLVE_DATAFRAME_COLUMN"):
+ left.zip(right).select(right.v).collect()
+
+ def test_resolve_after_zip_in_expression(self):
+ df = self.spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ["a",
"b"])
+ left = df.select(df.a.alias("x"))
+ right = df.select(df.b.alias("y"))
+ with self.assertRaisesRegex(AnalysisException,
"CANNOT_RESOLVE_DATAFRAME_COLUMN"):
+ left.zip(right).select((left.x + right.y).alias("s")).collect()
+
+ def test_resolve_after_zip_in_filter(self):
+ df = self.spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ["a",
"b"])
+ left = df.select(df.a.alias("x"))
+ right = df.select(df.b.alias("y"))
+ with self.assertRaisesRegex(AnalysisException,
"CANNOT_RESOLVE_DATAFRAME_COLUMN"):
+ left.zip(right).filter(left.x >= 2).collect()
+
+ def test_resolve_after_zip_chained(self):
+ df = self.spark.createDataFrame([(1, 2, 3), (4, 5, 6)], ["a", "b",
"c"])
+ left0 = df.select("a", "b")
+ left = left0.select((left0.a + 1).alias("x"))
+ right0 = df.select("b", "c")
+ right = right0.select((right0.c * 2).alias("y"))
+ with self.assertRaisesRegex(AnalysisException,
"CANNOT_RESOLVE_DATAFRAME_COLUMN"):
+ left.zip(right).select(left.x, right.y).collect()
+
+ def test_resolve_after_zip_base_side(self):
+ df = self.spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ["a",
"b"])
+ right = df.select((df.b * 2).alias("y"))
+ with self.assertRaisesRegex(AnalysisException,
"CANNOT_RESOLVE_DATAFRAME_COLUMN"):
+ df.zip(right).select(df.a).collect()
+
def test_df_col_resolution_mode(self):
self.assertEqual(
self.spark.conf.get("spark.sql.analyzer.strictDataFrameColumnResolution"),
diff --git a/python/pyspark/sql/tests/test_column.py
b/python/pyspark/sql/tests/test_column.py
index c53775250389..28333cb3f622 100644
--- a/python/pyspark/sql/tests/test_column.py
+++ b/python/pyspark/sql/tests/test_column.py
@@ -928,6 +928,130 @@ class ColumnTestsMixin:
rows = df1.intersect(df2).select(df1.c).collect()
self.assertEqual([r.c for r in rows], [2])
+ def test_resolve_after_zip(self):
+ # zip merges two column-projected DataFrames side by side. Classic
+ # resolves the tagged left/right reference by attribute id, which
+ # ResolveZip preserves in the merged Project, so it succeeds. Connect
+ # resolves by plan id, but ResolveZip collapses the two sides into one
+ # plan and drops the per-DataFrame plan-id tags, so the tagged
+ # reference is never found and it raises (overridden in the parity
suite).
+ df = self.spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ["a",
"b"])
+ left = df.select((df.a + 1).alias("x"))
+ right = df.select((df.b * 2).alias("y"))
+ zipped = left.zip(right)
+ self.assertEqual(zipped.columns, ["x", "y"])
+ self.assertEqual(sorted(r.x for r in zipped.select(left.x).collect()),
[2, 3, 4])
+ self.assertEqual(sorted(r.y for r in
zipped.select(right.y).collect()), [20, 40, 60])
+ self.assertEqual(
+ sorted((r.x, r.y) for r in zipped.select(left.x,
right.y).collect()),
+ [(2, 20), (3, 40), (4, 60)],
+ )
+
+ def test_resolve_after_zip_reordered(self):
+ # The originating DataFrame controls which side each column reads
from, in
+ # any order. Classic resolves by attribute id; Connect raises.
+ df = self.spark.createDataFrame([(1, 10), (2, 20)], ["a", "b"])
+ left = df.select(df.a.alias("x"))
+ right = df.select(df.b.alias("y"))
+ zipped = left.zip(right)
+ self.assertEqual(
+ sorted((r.y, r.x) for r in zipped.select(right.y,
left.x).collect()),
+ [(10, 1), (20, 2)],
+ )
+
+ def test_resolve_after_zip_duplicate_names(self):
+ # Both sides expose a column named "v" from different sources, so the
+ # merged schema has two "v"s. A bare "v" is ambiguous, but the tagged
+ # left/right reference disambiguates by attribute id on Classic.
+ # Connect raises (no plan-id node survives the merge).
+ df = self.spark.createDataFrame([(1, 2), (3, 4)], ["a", "b"])
+ left = df.select(df.a.alias("v"))
+ right = df.select(df.b.alias("v"))
+ zipped = left.zip(right)
+ self.assertEqual(zipped.columns, ["v", "v"])
+ self.assertEqual(sorted(r.v for r in zipped.select(left.v).collect()),
[1, 3])
+ self.assertEqual(sorted(r.v for r in
zipped.select(right.v).collect()), [2, 4])
+
+ def test_resolve_after_zip_shared_producer(self):
+ # Both sides project the same producer under the same name "v";
+ # ResolveZip dedups the producer but keeps one output column per side,
+ # each still selectable by its originating DataFrame on Classic.
+ df = self.spark.createDataFrame([(1, 2), (3, 4)], ["a", "b"])
+ left = df.select(df.a.alias("v"))
+ right = df.select(df.a.alias("v"))
+ zipped = left.zip(right)
+ self.assertEqual(zipped.columns, ["v", "v"])
+ self.assertEqual(sorted(r.v for r in zipped.select(left.v).collect()),
[1, 3])
+ self.assertEqual(sorted(r.v for r in
zipped.select(right.v).collect()), [1, 3])
+
+ def test_resolve_after_zip_in_expression(self):
+ # A tagged reference from each side can be combined in one expression.
+ df = self.spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ["a",
"b"])
+ left = df.select(df.a.alias("x"))
+ right = df.select(df.b.alias("y"))
+ zipped = left.zip(right)
+ self.assertEqual(
+ sorted(r.s for r in zipped.select((left.x +
right.y).alias("s")).collect()),
+ [11, 22, 33],
+ )
+
+ def test_resolve_after_zip_in_filter(self):
+ # A tagged reference resolves in a filter over the zip result too.
+ df = self.spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ["a",
"b"])
+ left = df.select(df.a.alias("x"))
+ right = df.select(df.b.alias("y"))
+ zipped = left.zip(right)
+ result = zipped.filter(left.x >= 2).select(left.x, right.y)
+ self.assertEqual(
+ sorted((r.x, r.y) for r in result.collect()),
+ [(2, 20), (3, 30)],
+ )
+
+ def test_resolve_after_zip_chained(self):
+ # Each side is its own chain of projections off the shared base; the
+ # tagged reference still resolves on Classic.
+ df = self.spark.createDataFrame([(1, 2, 3), (4, 5, 6)], ["a", "b",
"c"])
+ left0 = df.select("a", "b")
+ left = left0.select((left0.a + 1).alias("x"))
+ right0 = df.select("b", "c")
+ right = right0.select((right0.c * 2).alias("y"))
+ zipped = left.zip(right)
+ self.assertEqual(
+ sorted((r.x, r.y) for r in zipped.select(left.x,
right.y).collect()),
+ [(2, 6), (5, 12)],
+ )
+
+ def test_resolve_after_zip_base_side(self):
+ # A base-side reference: df.zip(right).select(df.a). Classic resolves
+ # by attribute id like the projected-side cases. On Connect the result
+ # depends on the base plan's root node: createDataFrame analyzes to a
+ # Project over a LocalRelation, and that Project - the node carrying
+ # the DataFrame's plan-id tag - is dissolved into the merged chain by
+ # ResolveZip, so the tagged reference raises (overridden in the parity
+ # suite). See test_resolve_after_zip_bare_base_side for a base whose
+ # tagged root survives the merge.
+ df = self.spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ["a",
"b"])
+ right = df.select((df.b * 2).alias("y"))
+ zipped = df.zip(right)
+ self.assertEqual(zipped.columns, ["a", "b", "y"])
+ self.assertEqual(sorted(r.a for r in zipped.select(df.a).collect()),
[1, 2, 3])
+ self.assertEqual(
+ sorted((r.b, r.a) for r in zipped.select(df.b, df.a).collect()),
+ [(10, 1), (20, 2), (30, 3)],
+ )
+
+ def test_resolve_after_zip_bare_base_side(self):
+ # When the base DataFrame's plan root is the base relation itself
+ # (range analyzes to a bare Range node), ResolveZip reuses that node
+ # unchanged as the merged plan's base, so its plan-id tag survives and
+ # the tagged reference resolves on Connect in all modes as well - the
+ # one zip shape with no Classic/Connect divergence for df.col.
+ df = self.spark.range(3)
+ right = df.select((df.id * 2).alias("y"))
+ zipped = df.zip(right)
+ self.assertEqual(zipped.columns, ["id", "y"])
+ self.assertEqual(sorted(r.id for r in zipped.select(df.id).collect()),
[0, 1, 2])
+
def test_resolve_self_join_alias(self):
# Both self-join sides originate from the same plan-id-tagged
# ancestor, yielding two equal-depth candidates with the same
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
