This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new ed2a28099e15 [MINOR][SQL][PYTHON][TEST] Add tests for dependent
columns in withColumns
ed2a28099e15 is described below
commit ed2a28099e15db7e913922a761e7fe8c947f5978
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 26 11:26:17 2026 +0800
[MINOR][SQL][PYTHON][TEST] Add tests for dependent columns in withColumns
### What changes were proposed in this pull request?
This PR adds direct regression coverage for dependent columns in
`DataFrame.withColumns`.
The JVM test exercises the internal `withColumns(Seq[String], Seq[Column])`
overload. The PySpark test covers a multi-level dependency graph in both
classic and Connect parity suites, plus a negative forward-reference case where
the later column is otherwise valid.
### Why are the changes needed?
Dependent columns currently work through lateral column alias resolution,
but `withColumns` had no direct test guarding this behavior or its ordering
constraint.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- `./build/sbt` targeted `DataFrameSuite` test
- Classic PySpark targeted
`DataFrameTests.test_with_columns_with_dependencies`
- Spark Connect targeted
`DataFrameParityTests.test_with_columns_with_dependencies`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58280 from zhengruifeng/sql-py-withcolumns-dependent-tests-dev-1.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/sql/tests/test_dataframe.py | 35 ++++++++++++++++++++++
.../org/apache/spark/sql/DataFrameSuite.scala | 14 +++++++--
2 files changed, 47 insertions(+), 2 deletions(-)
diff --git a/python/pyspark/sql/tests/test_dataframe.py
b/python/pyspark/sql/tests/test_dataframe.py
index d7fe514c6052..70c52404fc4f 100644
--- a/python/pyspark/sql/tests/test_dataframe.py
+++ b/python/pyspark/sql/tests/test_dataframe.py
@@ -30,6 +30,7 @@ from pyspark.errors import (
IllegalArgumentException,
PySparkTypeError,
PySparkValueError,
+ QueryContextType,
)
from pyspark.sql import DataFrame, Row, functions
from pyspark.sql.functions import (
@@ -425,6 +426,40 @@ class DataFrameTestsMixin:
self.assertRaises(TypeError, self.df.withColumns, ["key"])
self.assertRaises(Exception, self.df.withColumns)
+ def test_with_columns_with_dependencies(self):
+ df = self.spark.range(3).withColumns(
+ {
+ "a": col("id") + 1,
+ "b": col("a") + 1,
+ "c": col("a") + col("b"),
+ "d": col("a") + col("b") + col("c"),
+ }
+ )
+
+ assertDataFrameEqual(
+ df,
+ [
+ Row(id=0, a=1, b=2, c=3, d=6),
+ Row(id=1, a=2, b=3, c=5, d=10),
+ Row(id=2, a=3, b=4, c=7, d=14),
+ ],
+ )
+
+ with self.assertRaises(AnalysisException) as pe:
+ self.spark.range(1).withColumns(
+ {
+ "a": col("b") - 1,
+ "b": col("id") + 1,
+ }
+ ).collect()
+ self.check_error(
+ exception=pe.exception,
+ errorClass="UNRESOLVED_COLUMN.WITH_SUGGESTION",
+ messageParameters={"objectName": "`b`", "proposal": "`id`"},
+ query_context_type=QueryContextType.DataFrame,
+ fragment="col",
+ )
+
def test_generic_hints(self):
df1 = self.spark.range(10e10).toDF("id")
df2 = self.spark.range(10e10).toDF("id")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index 577ae025f1df..423dcc83dbed 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -636,6 +636,18 @@ class DataFrameSuite extends SharedSparkSession
assert(df.schema.map(_.name) === Seq("key", "value", "newCol1", "newCol2"))
}
+ test("withColumns: internal method with dependent columns") {
+ val df = testData.toDF().withColumns(
+ Seq("newCol1", "newCol2"),
+ Seq(col("key") + 1, col("newCol1") + 1))
+ checkAnswer(
+ df,
+ testData.collect().map { case Row(key: Int, value: String) =>
+ Row(key, value, key + 1, key + 2)
+ }.toSeq)
+ assert(df.schema.map(_.name) === Seq("key", "value", "newCol1", "newCol2"))
+ }
+
test("withColumns: internal method") {
val df = testData.toDF().withColumns(Seq("newCol1", "newCol2"),
Seq(col("key") + 1, col("key") + 2))
@@ -2822,5 +2834,3 @@ case class OutputListAwareConstraintsTestPlan(
override def newInstance(): LogicalPlan = copy(outputList =
outputList.map(_.newInstance()))
}
-
-
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
