robin-sparkless 4.0.0

"""
Unit tests for Issue #336: WindowFunction comparison operators. Uses get_spark_imports from fixture only.
"""

import pytest

from tests.fixtures.spark_imports import get_spark_imports

_imports = get_spark_imports()
SparkSession = _imports.SparkSession
F = _imports.F
Window = _imports.Window
StructType = _imports.StructType
StructField = _imports.StructField
StringType = _imports.StringType
IntegerType = _imports.IntegerType
LongType = _imports.LongType
DoubleType = _imports.DoubleType
BooleanType = _imports.BooleanType


class TestIssue336WindowFunctionComparison:
    """Test WindowFunction comparison operators."""

    def test_window_function_gt_comparison(self):
        """Test WindowFunction > comparison."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "GT-Zero",
                F.when(F.row_number().over(w) > 0, F.lit(True)).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 2
            assert rows[0]["GT-Zero"] is True
            assert rows[1]["GT-Zero"] is True
        finally:
            spark.stop()

    def test_window_function_lt_comparison(self):
        """Test WindowFunction < comparison."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "LT-Five",
                F.when(F.row_number().over(w) < 5, F.lit(True)).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 2
            assert rows[0]["LT-Five"] is True
            assert rows[1]["LT-Five"] is True
        finally:
            spark.stop()

    def test_window_function_ge_comparison(self):
        """Test WindowFunction >= comparison."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "GE-One",
                F.when(F.row_number().over(w) >= 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 2
            assert rows[0]["GE-One"] is True
            assert rows[1]["GE-One"] is True
        finally:
            spark.stop()

    def test_window_function_le_comparison(self):
        """Test WindowFunction <= comparison."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "LE-One",
                F.when(F.row_number().over(w) <= 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 2
            assert rows[0]["LE-One"] is True
            assert rows[1]["LE-One"] is True
        finally:
            spark.stop()

    def test_window_function_eq_comparison(self):
        """Test WindowFunction == comparison."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "EQ-One",
                F.when(F.row_number().over(w) == 1, F.lit("First")).otherwise(
                    F.lit("Other")
                ),
            )
            rows = result.collect()

            assert len(rows) == 2
            assert rows[0]["EQ-One"] == "First"
            assert rows[1]["EQ-One"] == "First"
        finally:
            spark.stop()

    def test_window_function_ne_comparison(self):
        """Test WindowFunction != comparison."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "NE-Zero",
                F.when(F.row_number().over(w) != 0, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 2
            assert rows[0]["NE-Zero"] is True
            assert rows[1]["NE-Zero"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_filter(self):
        """Window functions in filter are not allowed (PySpark parity)."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            # In PySpark this raises AnalysisException: it is not allowed to use
            # window functions inside WHERE clause.
            with pytest.raises(Exception) as exc_info:
                (
                    df.filter(F.row_number().over(w) == 1)
                    .select("Name", "Type", "Score")
                    .collect()
                )

            msg = str(exc_info.value)
            assert "window functions inside WHERE clause" in msg
        finally:
            spark.stop()

    def test_window_function_comparison_with_multiple_conditions(self):
        """Test WindowFunction comparison with multiple when conditions."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                    {"Name": "Charlie", "Type": "C"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "Category",
                F.when(F.row_number().over(w) == 1, F.lit("First"))
                .when(F.row_number().over(w) == 2, F.lit("Second"))
                .otherwise(F.lit("Other")),
            )
            rows = result.collect()

            assert len(rows) == 3
            assert rows[0]["Category"] == "First"
            assert rows[1]["Category"] == "First"
            assert rows[2]["Category"] == "First"
        finally:
            spark.stop()

    def test_window_function_comparison_with_rank(self):
        """Test WindowFunction comparison with rank()."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 100},
                    {"Name": "Charlie", "Type": "A", "Score": 90},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "TopRank",
                F.when(F.rank().over(w) <= 2, F.lit(True)).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Both Alice and Bob should have rank 1, so both should be True
            assert rows[0]["TopRank"] is True
            assert rows[1]["TopRank"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_dense_rank(self):
        """Test WindowFunction comparison with dense_rank()."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 100},
                    {"Name": "Charlie", "Type": "A", "Score": 90},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "TopDenseRank",
                F.when(F.dense_rank().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # With Score desc ordering and correct dense_rank semantics:
            # Alice and Bob both have Score 100, so they share dense_rank 1.
            # Charlie has Score 90, so has dense_rank 2.
            true_count = sum(1 for row in rows if row["TopDenseRank"] is True)
            # Both Alice and Bob should have dense_rank == 1 => two True values.
            assert true_count == 2
        finally:
            spark.stop()

    def test_window_function_comparison_with_percent_rank(self):
        """Test WindowFunction comparison with percent_rank()."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "TopPercent",
                F.when(F.percent_rank().over(w) == 0.0, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Find Alice (highest score) - she should have percent_rank 0.0
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["TopPercent"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_lag(self):
        """Test WindowFunction comparison with lag()."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "HasPrevious",
                F.when(F.isnull(F.lag("Score", 1).over(w)), F.lit(False)).otherwise(
                    F.lit(True)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Find Alice (highest score, first in order) - should not have previous
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["HasPrevious"] is False
            # Other rows should have previous
            for row in rows:
                if row["Name"] != "Alice":
                    assert row["HasPrevious"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_lead(self):
        """Test WindowFunction comparison with lead()."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "HasNext",
                F.when(F.isnull(F.lead("Score", 1).over(w)), F.lit(False)).otherwise(
                    F.lit(True)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Find Charlie (lowest score, last in order) - should not have next
            charlie_row = next(row for row in rows if row["Name"] == "Charlie")
            assert charlie_row["HasNext"] is False
            # Other rows should have next
            other_rows = [row for row in rows if row["Name"] != "Charlie"]
            for row in other_rows:
                assert row["HasNext"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_sum(self):
        """Test WindowFunction comparison with sum()."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "HighRunningSum",
                F.when(F.sum("Score").over(w) > 150, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Find Alice (highest score, first) - has sum 100
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["HighRunningSum"] is False
            # Other rows should have running sum > 150
            other_rows = [row for row in rows if row["Name"] != "Alice"]
            for row in other_rows:
                assert row["HighRunningSum"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_direct_filter(self):
        """Window functions used directly in filter raise (PySpark parity)."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            with pytest.raises(Exception) as exc_info:
                (
                    df.filter(F.row_number().over(w) == 1)
                    .select("Name", "Type", "Score")
                    .collect()
                )

            msg = str(exc_info.value)
            assert "window functions inside WHERE clause" in msg
        finally:
            spark.stop()

    def test_window_function_comparison_with_eqNullSafe(self):
        """Test WindowFunction eqNullSafe comparison."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "EQ-One-NullSafe",
                F.when(F.row_number().over(w).eqNullSafe(1), F.lit("First")).otherwise(
                    F.lit("Other")
                ),
            )
            rows = result.collect()

            assert len(rows) == 2
            assert rows[0]["EQ-One-NullSafe"] == "First"
            assert rows[1]["EQ-One-NullSafe"] == "First"
        finally:
            spark.stop()

    def test_window_function_comparison_with_isnotnull(self):
        """Test WindowFunction isnotnull() semantics (via F.isnotnull)."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "HasNext",
                F.when(F.isnotnull(F.lead("Score", 1).over(w)), F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Find Charlie (lowest score, last) - should not have next
            charlie_row = next(row for row in rows if row["Name"] == "Charlie")
            assert charlie_row["HasNext"] is False
            # Other rows should have next
            for row in rows:
                if row["Name"] != "Charlie":
                    assert row["HasNext"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_null_values(self):
        """Test WindowFunction comparison when window function returns null."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": None},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "HasPrevious",
                F.when(F.isnull(F.lag("Score", 1).over(w)), F.lit("NoPrev")).otherwise(
                    F.lit("HasPrev")
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # First row should have no previous
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["HasPrevious"] == "NoPrev"
        finally:
            spark.stop()

    def test_window_function_comparison_with_empty_dataframe(self):
        """Test WindowFunction comparison with empty DataFrame."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame([], schema="Name string, Type string")

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "GT-Zero",
                F.when(F.row_number().over(w) > 0, F.lit(True)).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 0
        finally:
            spark.stop()

    def test_window_function_comparison_with_single_row(self):
        """Test WindowFunction comparison with single row."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame([{"Name": "Alice", "Type": "A"}])

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "EQ-One",
                F.when(F.row_number().over(w) == 1, F.lit("First")).otherwise(
                    F.lit("Other")
                ),
            )
            rows = result.collect()

            assert len(rows) == 1
            assert rows[0]["EQ-One"] == "First"
        finally:
            spark.stop()

    def test_window_function_comparison_with_large_dataset(self):
        """Test WindowFunction comparison with larger dataset."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            data = [
                {"Name": f"Person{i}", "Type": "A", "Score": 100 - i} for i in range(20)
            ]
            df = spark.createDataFrame(data)

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "TopTen",
                F.when(F.row_number().over(w) <= 10, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 20
            top_ten_count = sum(1 for row in rows if row["TopTen"] is True)
            assert top_ten_count == 10
        finally:
            spark.stop()

    def test_window_function_comparison_with_multiple_window_functions(self):
        """Test WindowFunction comparison with multiple different window functions."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            ).withColumn(
                "TopRank",
                F.when(F.rank().over(w) == 1, F.lit(True)).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Both columns should exist
            assert "IsFirst" in rows[0]
            assert "TopRank" in rows[0]
        finally:
            spark.stop()

    def test_window_function_comparison_with_arithmetic_operations(self):
        """Test WindowFunction comparison after arithmetic operations."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            # First compute the arithmetic operation as a column
            result = df.withColumn(
                "RankPlusOne", (F.row_number().over(w) + 1)
            ).withColumn(
                "GT-Two",
                F.when(F.col("RankPlusOne") > 2, F.lit(True)).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 3
            # First row: row_number = 1, so row_number + 1 = 2, which is not > 2
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["GT-Two"] is False
            # Other rows should have row_number + 1 > 2
            for row in rows:
                if row["Name"] != "Alice":
                    assert row["GT-Two"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_select(self):
        """Test WindowFunction comparison followed by select."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            ).select("Name", "Type", "IsFirst")
            rows = result.collect()

            assert len(rows) == 2
            assert "IsFirst" in rows[0]
        finally:
            spark.stop()

    def test_window_function_comparison_with_orderby(self):
        """Test WindowFunction comparison followed by orderBy."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            ).orderBy("Name")
            rows = result.collect()

            assert len(rows) == 3
            # Should be ordered by Name
            assert rows[0]["Name"] in ["Alice", "Bob", "Charlie"]
        finally:
            spark.stop()

    def test_window_function_comparison_with_groupby(self):
        """Test WindowFunction comparison followed by groupBy."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "B", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = (
                df.withColumn(
                    "IsFirst",
                    F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                        F.lit(False)
                    ),
                )
                .groupBy("Type")
                .agg(F.max("Score").alias("MaxScore"))
            )
            rows = result.collect()

            assert len(rows) == 2
            types = {row["Type"] for row in rows}
            assert types == {"A", "B"}
        finally:
            spark.stop()

    def test_window_function_comparison_with_join(self):
        """Test WindowFunction comparison in join operations."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df1 = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                ]
            )

            df2 = spark.createDataFrame(
                [
                    {"Type": "A", "Dept": "Engineering"},
                    {"Type": "B", "Dept": "Sales"},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df1.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            ).join(df2, on="Type", how="left")
            rows = result.collect()

            assert len(rows) == 2
            assert rows[0]["Dept"] == "Engineering"
        finally:
            spark.stop()

    def test_window_function_comparison_with_union(self):
        """Test WindowFunction comparison in union operations."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df1 = spark.createDataFrame([{"Name": "Alice", "Type": "A", "Score": 100}])

            df2 = spark.createDataFrame([{"Name": "Bob", "Type": "B", "Score": 90}])

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result1 = df1.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            result2 = df2.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            combined = result1.unionByName(result2, allowMissingColumns=True)
            rows = combined.collect()

            assert len(rows) == 2
            names = {row["Name"] for row in rows}
            assert names == {"Alice", "Bob"}
        finally:
            spark.stop()

    def test_window_function_comparison_with_distinct(self):
        """Test WindowFunction comparison followed by distinct."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Alice", "Type": "A", "Score": 100},  # Duplicate
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = (
                df.withColumn(
                    "IsFirst",
                    F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                        F.lit(False)
                    ),
                )
                .select("Name", "Type", "Score")
                .distinct()
            )
            rows = result.collect()

            assert len(rows) == 1
            assert rows[0]["Name"] == "Alice"
        finally:
            spark.stop()

    def test_window_function_comparison_with_limit(self):
        """Test WindowFunction comparison followed by limit."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            ).limit(2)
            rows = result.collect()

            assert len(rows) == 2
        finally:
            spark.stop()

    def test_window_function_comparison_chained_operations(self):
        """Test WindowFunction comparison in chained DataFrame operations."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "B", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = (
                df.withColumn(
                    "IsFirst",
                    F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                        F.lit(False)
                    ),
                )
                .filter(F.col("IsFirst"))
                .select("Name", "Type", "Score")
            )
            rows = result.collect()

            assert len(rows) == 2  # One per Type partition
            types = {row["Type"] for row in rows}
            assert types == {"A", "B"}
        finally:
            spark.stop()

    def test_window_function_comparison_with_nested_select(self):
        """Test WindowFunction comparison with nested select operations."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = (
                df.withColumn(
                    "IsFirst",
                    F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                        F.lit(False)
                    ),
                )
                .select("Name", "IsFirst")
                .select("Name")
            )
            rows = result.collect()

            assert len(rows) == 2
            assert "Name" in rows[0]
        finally:
            spark.stop()

    def test_window_function_comparison_with_case_when_chain(self):
        """Test WindowFunction comparison with complex case/when chains."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "Category",
                F.when(F.row_number().over(w) == 1, F.lit("Gold"))
                .when(F.row_number().over(w) == 2, F.lit("Silver"))
                .when(F.row_number().over(w) == 3, F.lit("Bronze"))
                .otherwise(F.lit("Other")),
            )
            rows = result.collect()

            assert len(rows) == 3
            # All should have valid categories
            for row in rows:
                assert row["Category"] in ["Gold", "Silver", "Bronze", "Other"]
        finally:
            spark.stop()

    def test_window_function_comparison_with_coalesce(self):
        """Test WindowFunction comparison with coalesce."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            # First compute the window function comparison as a column
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(1)).otherwise(F.lit(None)),
            )
            # Then apply coalesce on the computed column (not nested)
            result = result.withColumn(
                "RankOrOne", F.coalesce(F.col("IsFirst"), F.lit(0))
            )
            rows = result.collect()

            assert len(rows) == 2
            # First row should have 1, second should have 0
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["RankOrOne"] == 1
            bob_row = next(row for row in rows if row["Name"] == "Bob")
            assert bob_row["RankOrOne"] == 0
        finally:
            spark.stop()

    def test_window_function_comparison_with_cast(self):
        """Test WindowFunction comparison with cast operation."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "RankStr",
                F.when(F.row_number().over(w) == 1, F.lit("1")).otherwise(F.lit("0")),
            ).withColumn("RankInt", F.col("RankStr").cast("int"))
            rows = result.collect()

            assert len(rows) == 2
            assert "RankInt" in rows[0]
            assert isinstance(rows[0]["RankInt"], (int, type(None)))
        finally:
            spark.stop()

    def test_window_function_comparison_with_avg(self):
        """Test WindowFunction comparison with avg() window function."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "AboveAvg",
                F.when(F.avg("Score").over(w) > 85, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Average is 90, so all should be True (since we're checking > 85)
            for row in rows:
                assert row["AboveAvg"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_max(self):
        """Test WindowFunction comparison with max() window function."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsMax",
                F.when(F.max("Score").over(w) == F.col("Score"), F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Alice has max score
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["IsMax"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_min(self):
        """Test WindowFunction comparison with min() window function."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsMin",
                F.when(F.min("Score").over(w) == F.col("Score"), F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Charlie has min score
            charlie_row = next(row for row in rows if row["Name"] == "Charlie")
            assert charlie_row["IsMin"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_count(self):
        """Test WindowFunction comparison with count() window function (PySpark semantics)."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "HasMultiple",
                F.when(F.count("Score").over(w) > 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # In PySpark, the running count in this window is [1, 2, 3] for
            # [Alice, Bob, Charlie], so HasMultiple is [False, True, True].
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            other_rows = [row for row in rows if row["Name"] != "Alice"]
            assert alice_row["HasMultiple"] is False
            for row in other_rows:
                assert row["HasMultiple"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_ntile(self):
        """Test WindowFunction comparison with ntile() window function."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                    {"Name": "David", "Type": "A", "Score": 70},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "TopHalf",
                F.when(F.ntile(2).over(w) == 1, F.lit(True)).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 4
            # Top half should have ntile 1
            top_half_count = sum(1 for row in rows if row["TopHalf"] is True)
            assert top_half_count == 2
        finally:
            spark.stop()

    def test_window_function_comparison_with_cume_dist(self):
        """Test WindowFunction comparison with cume_dist() window function."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "HighCumeDist",
                F.when(F.cume_dist().over(w) > 0.5, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # All should have valid cume_dist values
            for row in rows:
                assert "HighCumeDist" in row
        finally:
            spark.stop()

    def test_window_function_comparison_with_first_value(self):
        """Test WindowFunction comparison with first_value() window function."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirstValue",
                F.when(
                    F.first_value("Score").over(w) == F.col("Score"), F.lit(True)
                ).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Alice has the first value (highest score)
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["IsFirstValue"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_last_value(self):
        """Test WindowFunction comparison with last_value() window function."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsLastValue",
                F.when(
                    F.last_value("Score").over(w) == F.col("Score"), F.lit(True)
                ).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 3
            # All rows should have last_value equal to their own score
            # (because default frame is UNBOUNDED PRECEDING AND CURRENT ROW)
            for row in rows:
                assert row["IsLastValue"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_countDistinct(self):
        """Distinct window aggregates are not supported in PySpark."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 100},
                    {"Name": "Charlie", "Type": "A", "Score": 90},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            with pytest.raises(Exception) as exc_info:
                df.withColumn(
                    "HasMultipleDistinct",
                    F.when(F.countDistinct("Score").over(w) > 1, F.lit(True)).otherwise(
                        F.lit(False)
                    ),
                ).collect()

            msg = str(exc_info.value)
            assert "Distinct window functions are not supported" in msg or (
                "DISTINCT_WINDOW_FUNCTION_UNSUPPORTED" in msg
            )
        finally:
            spark.stop()

    def test_window_function_comparison_with_string_values(self):
        """Test WindowFunction comparison with string column values."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Category": "High"},
                    {"Name": "Bob", "Type": "A", "Category": "Medium"},
                    {"Name": "Charlie", "Type": "A", "Category": "Low"},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Category").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # One should be True (first row after ordering)
            true_count = sum(1 for row in rows if row["IsFirst"] is True)
            assert true_count == 1
            # Verify all rows have the IsFirst column
            for row in rows:
                assert "IsFirst" in row
        finally:
            spark.stop()

    def test_window_function_comparison_with_float_values(self):
        """Test WindowFunction comparison with float column values."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100.5},
                    {"Name": "Bob", "Type": "A", "Score": 90.3},
                    {"Name": "Charlie", "Type": "A", "Score": 80.7},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Alice has highest score
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["IsFirst"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_schema_verification(self):
        """Test that WindowFunction comparison produces correct schema."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "GT-Zero",
                F.when(F.row_number().over(w) > 0, F.lit(True)).otherwise(F.lit(False)),
            )

            schema = result.schema
            field_names = [field.name for field in schema.fields]

            # Should contain original columns plus new window function comparison column
            assert "Name" in field_names
            assert "Type" in field_names
            assert "GT-Zero" in field_names
        finally:
            spark.stop()

    def test_window_function_comparison_with_complex_filter(self):
        """Test WindowFunction comparison in complex filter conditions."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                    {"Name": "David", "Type": "B", "Score": 95},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            # First apply window function, then filter
            result = (
                df.withColumn("Rank", F.row_number().over(w))
                .filter((F.col("Rank") == 1) & (F.col("Type") == "A"))
                .select("Name", "Type", "Score")
            )
            rows = result.collect()

            assert len(rows) == 1
            assert rows[0]["Name"] == "Alice"
        finally:
            spark.stop()

    def test_window_function_comparison_with_multiple_partitions(self):
        """Test WindowFunction comparison with multiple partitions."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "B", "Score": 95},
                    {"Name": "David", "Type": "B", "Score": 85},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 4
            # Should have one True per partition
            true_count = sum(1 for row in rows if row["IsFirst"] is True)
            assert true_count == 2  # One for Type A, one for Type B
        finally:
            spark.stop()

    def test_window_function_comparison_with_no_partition(self):
        """Test WindowFunction comparison without partitionBy."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Only one should be True (the first in the entire dataset)
            true_count = sum(1 for row in rows if row["IsFirst"] is True)
            assert true_count == 1
            # Find the row with IsFirst=True
            first_row = next(row for row in rows if row["IsFirst"] is True)
            # Verify it's one of the rows (actual ordering may vary)
            assert first_row["Name"] in ["Alice", "Bob", "Charlie"]
        finally:
            spark.stop()

    def test_window_function_comparison_with_rowsBetween(self):
        """Test WindowFunction comparison with rowsBetween boundaries."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = (
                Window()
                .partitionBy("Type")
                .orderBy(F.col("Score").desc())
                .rowsBetween(Window.unboundedPreceding, Window.currentRow)
            )
            result = df.withColumn(
                "HighRunningSum",
                F.when(F.sum("Score").over(w) > 150, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # First row has sum 100, others should have sum > 150
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["HighRunningSum"] is False
        finally:
            spark.stop()

    def test_window_function_comparison_with_rangeBetween(self):
        """Test WindowFunction comparison with rangeBetween boundaries."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = (
                Window()
                .partitionBy("Type")
                .orderBy(F.col("Score").desc())
                .rangeBetween(Window.unboundedPreceding, Window.currentRow)
            )
            result = df.withColumn(
                "HighRunningSum",
                F.when(F.sum("Score").over(w) > 150, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # All should have valid values
            for row in rows:
                assert "HighRunningSum" in row
        finally:
            spark.stop()

    def test_window_function_comparison_with_negative_values(self):
        """Test WindowFunction comparison with negative values."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": -10},
                    {"Name": "Bob", "Type": "A", "Score": -20},
                    {"Name": "Charlie", "Type": "A", "Score": -30},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Alice has highest (least negative) score
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["IsFirst"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_zero_values(self):
        """Test WindowFunction comparison with zero values."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 0},
                    {"Name": "Bob", "Type": "A", "Score": 0},
                    {"Name": "Charlie", "Type": "A", "Score": 0},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # One should be True (first row)
            true_count = sum(1 for row in rows if row["IsFirst"] is True)
            assert true_count == 1
        finally:
            spark.stop()

    def test_window_function_comparison_with_duplicate_scores(self):
        """Test WindowFunction comparison with duplicate scores."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 100},
                    {"Name": "Charlie", "Type": "A", "Score": 100},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Only one should be True (row_number assigns unique numbers)
            true_count = sum(1 for row in rows if row["IsFirst"] is True)
            assert true_count == 1
        finally:
            spark.stop()

    def test_window_function_comparison_with_all_null_partition(self):
        """Test WindowFunction comparison when partition has all null values."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            # Use explicit schema to avoid inference issues with all nulls
            schema = StructType(
                [
                    StructField("Name", StringType(), True),
                    StructField("Type", StringType(), True),
                    StructField("Score", IntegerType(), True),
                ]
            )
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": None},
                    {"Name": "Bob", "Type": "A", "Score": None},
                ],
                schema=schema,
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 2
            # One should be True
            true_count = sum(1 for row in rows if row["IsFirst"] is True)
            assert true_count == 1
        finally:
            spark.stop()

    def test_window_function_comparison_with_mixed_types(self):
        """Test WindowFunction comparison with mixed data types."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100, "Category": "High"},
                    {"Name": "Bob", "Type": "A", "Score": 90, "Category": "Medium"},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit("Yes")).otherwise(
                    F.lit("No")
                ),
            )
            rows = result.collect()

            assert len(rows) == 2
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["IsFirst"] == "Yes"
        finally:
            spark.stop()

    def test_window_function_comparison_with_desc_ordering(self):
        """Test WindowFunction comparison with descending order."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Alice should be first (highest score with desc)
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            assert alice_row["IsFirst"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_asc_ordering(self):
        """Test WindowFunction comparison with ascending order."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").asc())
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Charlie should be first (lowest score with asc)
            charlie_row = next(row for row in rows if row["Name"] == "Charlie")
            assert charlie_row["IsFirst"] is True
        finally:
            spark.stop()

    def test_window_function_comparison_with_multiple_orderby_columns(self):
        """Test WindowFunction comparison with multiple orderBy columns."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100, "Age": 25},
                    {"Name": "Bob", "Type": "A", "Score": 100, "Age": 30},
                    {"Name": "Charlie", "Type": "A", "Score": 90, "Age": 20},
                ]
            )

            w = (
                Window()
                .partitionBy("Type")
                .orderBy(F.col("Score").desc(), F.col("Age").asc())
            )
            result = df.withColumn(
                "IsFirst",
                F.when(F.row_number().over(w) == 1, F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # One should be True (first row after ordering by Score desc, then Age asc)
            true_count = sum(1 for row in rows if row["IsFirst"] is True)
            assert true_count == 1
            # Find the first row
            first_row = next(row for row in rows if row["IsFirst"] is True)
            # Should have Score 100 (highest) and be one of Alice or Bob
            assert first_row["Score"] == 100
            assert first_row["Name"] in ["Alice", "Bob"]
        finally:
            spark.stop()

    def test_window_function_comparison_with_chained_filters(self):
        """Window function in chained filters raises (PySpark parity)."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                    {"Name": "David", "Type": "B", "Score": 95},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            with pytest.raises(Exception) as exc_info:
                (
                    df.filter(F.row_number().over(w) == 1)
                    .filter(F.col("Type") == "A")
                    .select("Name", "Type", "Score")
                    .collect()
                )

            msg = str(exc_info.value)
            assert "window functions inside WHERE clause" in msg
        finally:
            spark.stop()

    def test_window_function_comparison_with_window_function_in_value(self):
        """Test WindowFunction comparison where value is also a window function."""
        spark = SparkSession.builder.appName("issue-336").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w1 = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            w2 = Window().partitionBy("Type").orderBy(F.col("Score").asc())
            # First compute both window functions as columns
            df_with_ranks = df.withColumn("Rank1", F.row_number().over(w1)).withColumn(
                "Rank2", F.row_number().over(w2)
            )
            # Then compare the columns (not window functions directly)
            result = df_with_ranks.withColumn(
                "RankMatch",
                F.when(F.col("Rank1") == F.col("Rank2"), F.lit(True)).otherwise(
                    F.lit(False)
                ),
            )
            rows = result.collect()

            assert len(rows) == 3
            # Middle row (Bob) should match (rank 2 in both)
            bob_row = next(row for row in rows if row["Name"] == "Bob")
            assert bob_row["RankMatch"] is True
        finally:
            spark.stop()