Add experimental native pyspark.sql.DataFrame tracking

[w-martin]

Adds experimental tracking of pyspark.sql.DataFrame as a native type, alongside pandas and polars. Today Spark is only reachable through the SQL-connector path (spark.sql(sql).toPandas(), which ends in a pandas frame); this covers the case where the frame stays a Spark frame all the way through.

Why this one is different from a pandas-mirror backend

Spark isn't a pandas-API clone, so almost nothing could be reused by name:

pandas/polars	PySpark
assign(col=…)	withColumn(name, expr) / withColumns({…})
rename(columns={…})	withColumnRenamed(old, new) — two positional args, not a mapping
drop(columns=[…])	drop(*cols) — bare varargs
pl.col("x")	F.col("x") / functions.col("x") / sf.col("x")
read_csv(usecols=[…])	no read-time column subset at all

What's tracked

• extract_spark_col_name — the F.col("x") analogue of extract_pl_col_name, wired into the shared collect_pl_col_names so every existing validator picks it up for free. The bare from pyspark.sql.functions import col form already worked via the existing name-only branch.
• Structural ops — withColumn/withColumns, withColumnRenamed/withColumnsRenamed, varargs select (including .alias() output names), varargs drop, toDF, union/unionAll/unionByName.
• Row-passthrough — where, limit, offset, orderBy, distinct, dropDuplicates, sortWithinPartitions, repartition, coalesce, cache/persist/unpersist, checkpoint, alias, hint. Each checked against its real signature for whether it actually returns a column-identical DataFrame.
• Read/schema inference — from an explicit schema (DDL string, StructType, or list of names; inline or held in a variable, which is the idiomatic spelling), or from a .select(...) chained onto an undeclared read. A bare spark.read.csv(path) is genuinely unknowable and reports untracked-dataframe, reusing the existing warning and coverage accounting rather than inventing a new code.
• spark.sql(...) kept native — columns from the SELECT list via the existing register_sql_dataframe.
• Annotated[pyspark.sql.DataFrame, Schema] — already worked (the recognizer keys on the DataFrame attribute name, not the module); now covered by tests so it stays working.

The judgment calls worth a second look

join — left untracked, deliberately. Its output depends on the join type (left_semi/left_anti return only the left side's columns), on whether on is a name list (which collapses the join keys into one column) or a condition (which keeps both), and on how duplicate names across the two sides resolve. That is squarely the runtime-dependent bucket. This is more conservative than pandas' merge, which is tracked as a union — I think correctly so, but it's the call I'd most want reviewed.

union/unionAll — tracked as the LEFT schema, not a union. PySpark documents these as resolving columns by position, "following the standard behavior in SQL", so the result carries the left frame's names whatever the right frame calls its columns. Spark rejects a mismatched column count outright, so for any program that runs at all, left's schema is the result's. This is more precise than pandas' concat handling (sort+dedup of both), not a weaker version of it.

unionByName — left schema, except with allowMissingColumns=True. That flag is the one case where the result really is the union of both sides, and it's tracked only when the right side is itself tracked; otherwise the binding is dropped rather than guessed at.

Unresolvable ops drop the binding rather than keeping it. df = df.withColumn(runtime_name, …) can't be resolved, and leaving the previous schema in place would then report the newly added column as unknown — a false positive on the most common Spark shape there is. untrack_targets exists for exactly this.

Not done, and why

• selectExpr(...) — output names come from SQL expression strings; parsing them wasn't worth the false-positive risk.
• UDFs (pandas_udf, .rdd.map, a lambda in withColumn) — untracked by design, per the project's stated philosophy. withColumn still adds the column when its name is a literal; only the contents are opaque.
• Chained receivers — df.filter(…).limit(5) doesn't track, because the row-passthrough branch requires a plain Name receiver. This is a pre-existing limitation shared with pandas (df.query(…).head() behaves the same), not new here; fixing it means recursive chain resolution across the whole linter, which felt out of scope.
• spark.read.text(path) — schema is documented as value plus any partition columns, so neither "unknown" nor ["value"] is accurate. Left out rather than half-modelled.
• No per-variable backend. The linter tracks a column set, not which library produced it, so Spark's head/tail/first (which return Row/list[Row]) still fall under the shared ROW_PASSTHROUGH_METHODS that pandas needs. Noted in the code and the method matrix.

Coexistence with the existing SQL path

The two never see the same AST node: the SQL-connector path matches on the outer .toPandas() call, this one only when the .sql(...) call is itself the assigned value. Native .sql() is additionally gated on the receiver being a recognized SparkSession (from a getOrCreate() chain or a SparkSession-annotated parameter) — without that gate, DuckDB's duckdb.sql(q) returns a relation, and registering one would make rel.fetchall() read as an unknown column.

Evidence it's genuinely untouched: I ran the checker from this branch and from main over all 101 tracked .py files (examples, tests, src) and diffed the full output. Byte-identical. One regression surfaced that way during development — session.createDataFrame(rows, [...]) in examples/sql_connectors/databricks/ made seed_df.write read as an unknown column — and is fixed by SPARK_RESERVED_ATTRIBUTES.

One shared-path change is intentional: drop() now reads varargs, which also fixes polars' drop("a", "b") (previously only the first column was dropped). pandas' deprecated positional-axis drop("a", 1) still resolves exactly as before, and there's a test pinning that.

Verification status — please read

No Spark was executed. This machine has no JVM (/usr/bin/java is the macOS stub: "Unable to locate a Java Runtime"), so nothing here is execution-verified against a running Spark session.

• Source-verified — I installed pyspark==4.2.0 and read the actual signatures and docstrings in pyspark/sql/dataframe.py, readwriter.py, session.py and types.py. The load-bearing claims all come from there, not from memory: union resolving by position and unionByName by name; drop and withColumnRenamed being documented no-ops when the column is absent; csv/json taking schema as the second positional argument and load as the third; head/tail returning Row/list[Row]; text's value-plus-partitions schema.
• Execution-verified — the checker's behaviour is, thoroughly: 41 new Rust tests, plus two hand-written Spark files run through the built binary covering every branch, plus the corpus diff above.

So: what the checker does with Spark source is well tested. Whether Spark at runtime agrees is docs/source-verified only. The withColumnRenamed-is-a-no-op claim drives an unknown-column error, so it's the one most worth a sanity check by someone with a JVM.

Checks

uv run inv all passes end to end — ruff format/check, ty, bandit, complexipy, cargo fmt --check, cargo clippy -D warnings, 301 Python tests at the 100% branch-coverage gate, 235 Rust tests, licensecheck. No lint-ignore rules or bandit skips were added; the example uses from pyspark.sql import functions rather than the conventional as F purely because the aliased form trips ruff's N812 and I wasn't going to add an ignore for it (the checker accepts F, sf, functions, and bare col identically — noted in the example's docstring).

https://claude.ai/code/session_013nRj6HvRbwkzAoj3QPbiSj