[SPARK-59072] Close remaining pyspark v4.2.0 public API gaps in the Rust Spark Connect client drop-in - #81
Open
HyukjinKwon wants to merge 51 commits into
Open
Conversation
…udf/cume_dist/broadcast/column Expose the higher-order functions (transform, filter, exists, forall, aggregate, reduce, zip_with, transform_keys, transform_values, map_filter, map_zip_with) and cume_dist/broadcast/call_function/call_udf/column, mirroring pyspark.sql.connect.functions exactly. - Replace the fragile fixed-variable-name HOF bindings with faithful primitives (pyfunc_named_lambda_variable, pyfunc_lambda_function, pyfunc_invoke_function) and copy the official _create_lambda / _invoke_higher_order_function into functions.py (fresh unique lambda var names via a monotonic counter, arity 1-3 validation, return-Column validation) so nested lambdas no longer collide. - Fix CallFunctionWrapper (core) to carry its argument expressions; fix core call_function/call_udf to pass arguments (they previously dropped them). - Tests: hof_golden + functions_golden (511 cases) stay green; new structural test asserts call_function/call_udf carry args; offline Python tests cover every new function plus arity/return/type validation paths. Co-authored-by: Isaac <no-reply@databricks.com>
…/truncateTable/analyzeTable/getCreateTableString/getTableProperties/listPartitions/listViews Add the 10 catalog DDL/metadata operations that pyspark v4.2.0 exposes, mirroring pyspark.sql.connect.catalog exactly (the proto Catalog oneof variants already exist in catalog.proto): - core catalog.rs: create_database (with properties map), drop_database (cascade), drop_table (purge), drop_view, truncate_table, analyze_table (no_scan), get_create_table_string, get_table_properties (key/value rows), list_partitions, list_views (current-database fallback when a pattern is given without a db). - pyspark-rs: expose all 10 with the official camelCase signatures/defaults; getTableProperties returns a dict, listPartitions/listViews return DataFrames (matching the other list_* catalog methods). - Test: e2e catalog_ddl_v420_surface exercises every method against a live server (create db + partitioned table -> list/analyze/props/create-string/truncate/drop -> drop view -> drop db) and cleans up. Verified green against Spark 4.2.0. Co-authored-by: Isaac <no-reply@databricks.com>
…m, RuntimeConf.getAll, DataFrameStatFunctions.sampleBy
Mirror the official pyspark v4.2.0 implementations:
- DataFrame.repartitionById(numPartitions, partitionIdCol): fix core repartition_by_id
which IGNORED the column (it just called repartition); now wraps the column in a new
DirectShufflePartitionID expression (added to core, proto already had it) and
repartitions by it.
- DataFrame.zipWithIndex(indexColName="index"): select(col("*"),
distributed_sequence_id().alias(indexColName)) as in the reference.
- Column.transform(f): returns f(self).
- RuntimeConf.getAll: expose the property (dict) over core get_all.
- DataFrameStatFunctions.sampleBy: expose over core stat.sample_by.
Tests: e2e dataframe_extra_surface now exercises repartitionById + zipWithIndex against
a live server; offline tests cover Column.transform and method presence. Verified all
five end-to-end against Spark 4.2.0.
Co-authored-by: Isaac <no-reply@databricks.com>
…/DataStreamReader.changes, DataStreamReader.xml/name
Reader CDC + XML + a native (Rust) StreamingQueryManager listener bus:
- DataFrameReader.changes / DataStreamReader.changes: new core LogicalPlan::RelationChanges
(proto field already existed) -> spark.read.changes(table) / readStream.changes(table).
- DataStreamReader.xml(path, **options): format("xml") + load, mirroring the reference.
- DataStreamReader.name(source_name): expose the core source-name setter (with the
reference's ASCII validation).
- StreamingQueryManager listener bus is now implemented natively in the Rust core
(StreamingQueryListener trait + StreamingQueryListenerEvent + a background dispatch
thread) so Rust clients get the feature too; add_listener/remove_listener/close live
in the core. spark.streams returns the native manager; the Python side is a thin
adapter (PyListenerAdapter -> _dispatch_listener_event) that builds the typed event
and invokes the listener callback. The old server-side payload registration is kept as
register_python_listener/unregister_python_listener.
Tests: new e2e streaming_native_listener_bus drives the Rust listener trait against a
live server (asserts events dispatched, then remove + close); RelationChanges has a
deterministic plan-serialization test; DataFrameReader.changes verified reaching the
server (UNSUPPORTED_FEATURE.CHANGE_DATA_CAPTURE on a plain table proves the wire form);
offline tests cover method presence + the manager listener API. All e2e + offline green.
Co-authored-by: Isaac <no-reply@databricks.com>
…mpleString/needConversion/fromInternal/toInternal/fromDDL + StructType/StructField specifics)
Expose the DataType object-model on every type class, mirroring pyspark.sql.types:
- All DataType classes: json, jsonValue, typeName, needConversion, fromInternal/
toInternal (identity, the base contract), fromDDL (client-side DDL parse).
- StructType: typeName ("struct"), fieldNames, toDDL, treeString, toNullable, fromJson.
- StructField: simpleString, json, jsonValue, needConversion, fromJson,
getCollationMetadata; typeName raises (as in pyspark: "use typeName on its type").
- Core gains to_ddl/tree_string/to_nullable and from_json_str / StructField::from_json_str;
the Python fromJson/fromDDL route through the core (no serde in the extension).
Verified byte-for-byte against official pyspark for the context-free methods (e.g.
StructField.json == '{"metadata":{},"name":"a","nullable":true,"type":"integer"}'); our
toDDL/treeString/fromDDL additionally work without a live session (the classic path
requires one). Tests: offline datatype/structfield/structtype object-model tests +
core to_ddl/tree_string/to_nullable/from_json_str unit tests.
Co-authored-by: Isaac <no-reply@databricks.com>
…r.channelBuilder - functions.udtf / arrow_udtf: re-export the UDTF decorators (udtf.py gains arrow_udtf, which builds a UserDefinedTableFunction with SQL_ARROW_TABLE_UDF=301); udtf already existed and is now surfaced on pyspark.sql.functions. - functions.arrow_udf: new Arrow scalar UDF decorator (udf.py), the Arrow analogue of pandas_udf (SQL_SCALAR_ARROW_UDF=250 / SQL_SCALAR_ARROW_ITER_UDF=251), mirroring pyspark.sql.pandas.functions.arrow_udf. - SparkSession.Builder.channelBuilder: instead of raising, accept a Spark Connect ChannelBuilder and reconstruct an sc:// URL from its host/port (+ connection params) for the native Rust transport; a builder without a host raises a clear error. Tests: offline tests build udtf/arrow_udtf (assert evalType) + arrow_udf, and cover channelBuilder (endpoint reconstruction + error path). channelBuilder verified end-to-end producing a working session against the live server. Co-authored-by: Isaac <no-reply@databricks.com>
…Row.count/index, avro/protobuf/partitioning submodules - registerJavaFunction/registerJavaUDAF: core SparkSession.register_java_function builds a CommonInlineUserDefinedFunction carrying a JavaUDF and sends the RegisterFunction command (mirrors client.register_java); spark.udf now returns the real pyspark.sql.udf.UDFRegistration (bound to the session) instead of a minimal inline stub, so register + the two Java methods are all available. Verified the JavaUDF RegisterFunction reaches the server (class-load error on a bogus class proves the wire). - Row.count(value) / Row.index(value[, start[, stop]]): the inherited tuple methods. - pyspark.sql.avro.functions / pyspark.sql.protobuf.functions / pyspark.sql.functions. partitioning are now importable (re-export the flat implementations), matching the official import paths. Tests: offline coverage for Row.count/index, the submodule import paths, and the UDFRegistration Java methods; registerJavaFunction verified against the live server. Co-authored-by: Isaac <no-reply@databricks.com>
….random, StructField.fromDDL, StorageLevel/StreamingQueryListener re-exports
Close the smaller parity gaps surfaced by a follow-up audit:
- UserDefinedFunction.asNondeterministic(): returns a non-deterministic copy; the
deterministic flag is now threaded through pyfunc_make_udf to the UDF proto (was
hardcoded true).
- UserDefinedTableFunction.asDeterministic(): deterministic copy.
- functions.random: alias for rand (matches pyspark).
- StructField.fromDDL("a int" / "a: int" / "a array<int>"): single-field DDL parse.
- Re-exports: `from pyspark import StorageLevel` and
`from pyspark.sql.streaming import StreamingQueryListener` now resolve.
Tests: offline coverage for all of the above.
Co-authored-by: Isaac <no-reply@databricks.com>
…cType/IntegralType/FractionalType/DatetimeType/AnyTimeType/AnsiIntervalType/SpatialType) Give the type classes the reference MRO via real PyO3 inheritance so isinstance(dt, DataType) / isinstance(dt, NumericType) / issubclass(...) work and the base-class names import from pyspark.sql.types: - PyDataType is now the subclassable base (holds `inner`); the abstract intermediate bases (AtomicType/NumericType/IntegralType/FractionalType/DatetimeType/AnyTimeType/ AnsiIntervalType/SpatialType) are empty PyO3 classes in the extends chain. - Every concrete type extends the correct base and builds its initializer chain via an `init_chain!` macro, e.g. IntegerType -> IntegralType -> NumericType -> AtomicType -> DataType (matches reference __mro__ exactly). StructType.toNullable/fromJson construct through the chain (Py::new). Verified: reference MRO for all atomic + parameterized types, isinstance/issubclass up the whole chain, object model (simpleString/json/typeName), pickling, StructType methods, and flat createDataFrame all still work against the live server. Offline hierarchy tests added; core type tests + full offline suite green. Co-authored-by: Isaac <no-reply@databricks.com>
…CatalogMetadata/TablePartition) — Rust-backed + typed catalog methods Implement the catalog metadata result classes natively in the Rust core (structs in catalog.rs) with typed parsing methods (list_tables_typed/get_table_typed/ list_databases_typed/get_database_typed/list_functions_typed/get_function_typed/ list_columns_typed/list_catalogs_typed/list_partitions_typed/list_views_typed), and expose them as PyO3 classes. spark.catalog.listTables/listDatabases/getTable/... now return the typed List[Table]/List[Database]/... objects, matching reference pyspark (previously returned DataFrames). CatalogColumn is exposed as `Column` via pyspark.sql.catalog to avoid colliding with the expression Column. The core DataFrame-returning methods are retained (used elsewhere). Tests: native e2e catalog_typed_results_surface parses every result type against a live server; offline test asserts the classes are Rust-backed + importable. Verified live end-to-end (listCatalogs/Databases/Tables/Functions/Columns + get* + isinstance). Co-authored-by: Isaac <no-reply@databricks.com>
…y_to_data_type ordering for mutated StructType
- GeometryType(srid)/GeographyType(srid): expose the spatial types (core already had
DataType::Geometry/Geography + proto), extending SpatialType with the correct MRO
(GeometryType -> SpatialType -> AtomicType -> DataType). srid getter + simpleString
("geometry(4326)") match reference.
- Fix: py_to_data_type now matches the concrete subclasses (which read their live fields)
BEFORE the shared PyDataType base, so a StructType mutated via add() reflects the new
field in json()/jsonValue() (the base carries only the construction-time snapshot).
Tests: offline spatial-type MRO/isinstance/srid + StructType.add()->json regression.
Co-authored-by: Isaac <no-reply@databricks.com>
…em = impl Into<Column>>
Sweep every Rust core API that took Vec<Column> to the idiomatic
`fn f<C: Into<Column>>(x: impl IntoIterator<Item = C>)`, so callers can pass arrays
(`select([col("a"), col("b")])`), string column names (`select(["a", "b"])` via a new
`From<&str>/String/&String for Column`), iterators, or `vec![...]` (still compiles):
DataFrame select/group_by/rollup/cube/unpivot, Column.isin, the functions.* variadic
builders (array/concat/coalesce/create_map/struct/hash/stack/named_struct/elt/
grouping_id/map_concat/xxhash64/reflect/try_reflect/java_method/arrays_zip),
call_function/call_udf/variant_delete, tvf.json_tuple/stack, writer.partition_by,
table_arg.partition_by/order_by, plan::project, wasm call.
Empty literals need an element type (`Vec::<Column>::new()`), applied in the golden
tests. Updated all docs (README + docs/*.md) and examples/ to the array form. New test
covers arrays/&str/vec parity; golden + plan suites stay green.
Co-authored-by: Isaac <no-reply@databricks.com>
Add pyspark.util.PythonEvalType as a Rust-backed PyO3 class carrying every eval-type constant (SQL_BATCHED_UDF=100 ... SQL_SCALAR_ARROW_UDF=250 ... SQL_ARROW_TABLE_UDF=301, incl. the transform-with-state variants), mirroring pyspark.util.PythonEvalType. Co-authored-by: Isaac <no-reply@databricks.com>
…e SQL classes now 0 gaps vs v4.2.0 Add the two remaining StructField collation methods: getCollationsMap(metadata) parses the __COLLATIONS metadata key into a path->collation-name map (empty when absent), and schemaCollationValue(dt) returns "provider.collation" (UTF8_BINARY/UTF8_LCASE -> spark, else icu). With these, the introspection audit shows 0 method gaps across all core SQL classes (DataFrame/Column/SparkSession/GroupedData/Reader/Writer/Catalog/Window/ Observation/streaming/types) and the full functions module vs pyspark v4.2.0. Co-authored-by: Isaac <no-reply@databricks.com>
Expose the core TableArg as a PyO3 class and make DataFrame.asTable() return it (no longer a misnamed alias): partitionBy/orderBy/withSinglePartition chain and build the TABLE_ARG SubqueryExpression, mirroring pyspark.sql.connect.table_arg.TableArg. Core TableArg gains Clone for the chaining wrappers; to_column_list is now pub(crate). Co-authored-by: Isaac <no-reply@databricks.com>
…n/OrderingColumn/SelectedColumn/AnalyzeResult + SkipRestOfInputTableException) — Rust-backed Add the polymorphic-UDTF analyze data holders as Rust-backed PyO3 classes (frozen where the reference dataclass is frozen; AnalyzeResult is mutable), plus the SkipRestOfInputTableException, re-exported from pyspark.sql.udtf. Mirrors pyspark.sql.udtf. Co-authored-by: Isaac <no-reply@databricks.com>
…ked) + vendored variant codec
Add the value classes pyspark.sql.types.VariantVal / Geometry / Geography as Rust-backed
PyO3 holders (value/metadata or wkb/srid + accessors, fromWKB, equality). VariantVal's
toJson/toPython/parseJson delegate to the variant binary codec, vendored byte-identical
from upstream pyspark.sql.variant_utils (a pure client-side algorithm, identical on the
worker). Verified JSON round-trip ('{"a":1,"b":[2,3]}' -> toPython {'a':1,'b':[2,3]}).
Co-authored-by: Isaac <no-reply@databricks.com>
Formatting-only pass over the SPARK-59072 changes. Co-authored-by: Isaac <no-reply@databricks.com>
zhengruifeng
approved these changes
Aug 28, 2026
Reimplement pyspark.storagelevel.StorageLevel as a Rust-backed PyO3 class (fields + repr/str/eq + all preset class attributes NONE/DISK_ONLY[_2/_3]/MEMORY_ONLY[_2]/ MEMORY_AND_DISK[_2]/OFF_HEAP/MEMORY_AND_DISK_DESER); the Python module just re-exports it. Co-authored-by: Isaac <no-reply@databricks.com>
- Add pyspark.sql.types.UserDefinedType as a faithful Python subclass of the Rust-backed DataType (its jsonValue cloudpickles the concrete class and fromJson re-imports it — inherently a Python construct, like the pickling serializers). Wire py_to_data_type to lower a UDT instance into the core Udt (jsonValue + sqlType), and _parse_datatype_json_value to reconstruct it from the "udt" JSON form. - Add a core-DataType -> concrete-Python-type materializer (data_type_to_py), the inverse of py_to_data_type, so schema introspection yields the proper classes. - DataFrame.schema is now a property returning a real StructType (was a bare DataType method), so df.schema.fields / df.schema["c"] / df.schema.fieldNames() work. - StructType gains .fields / .names / __getitem__ (index, name, slice) / __iter__ / __len__, and .add() is chainable. StructField gains .dataType / .metadata getters, a faithful __repr__, and __eq__. - PyDataType gets a __new__ so pure-Python subclasses (UDTs) are instantiable. Co-authored-by: Isaac <no-reply@databricks.com>
The vendored Spark Connect .proto set already covers 100% of the v4.2.0 wire surface (it is the v4.2.0 protobuf plus the fork's documented local additions — Zip, nanosecond timestamp literals/types, SCD2 pipeline history — none of which v4.2.0 removed). Only the version markers lagged: bump PROTO_VERSION.txt and SPARK_SHA.txt to v4.2.0 (tag 32f7299), the user-agent SPARK_VERSION to "4.2.0", and the architecture doc's protocol note. Co-authored-by: Isaac <no-reply@databricks.com>
… logger Adds the Spark Declarative Pipelines client surface and the two official pure-Python pipeline parity tests, which now pass unmodified against this client: - pyspark.pipelines: api (@table/@materialized_view/@temporary_view/@append_flow/ @create_streaming_table/@create_auto_cdc_flow/@create_sink decorators), flow, output, graph_element_registry, source_code_location, type_error_utils, logging_utils — vendored from v4.2.0. The only edit is api.py's two lazy `pyspark.sql.connect.functions.builtin` imports (expr/col), repointed at our Rust-backed pyspark.sql.functions (no separate Connect functions module here). - pyspark.pipelines.tests: the in-memory LocalGraphElementRegistry + the official test_decorators (3) and test_graph_element_registry (10) — all 13 green. To get there, two real parity gaps are closed by vendoring the pure-Python upstream modules verbatim (they were stubs / missing): - pyspark.errors: faithful base.py (getCondition/getMessageParameters/getSqlState/ message-template formatting), utils.ErrorClassesReader, error_classes + error-conditions.json, exceptions/tblib — replacing the hand-written stub that lacked getCondition and error-class message rendering. - pyspark.logger: PySparkLogger structured logging (required by errors.base). Also adds a thin pyspark.sql.connect.functions[.builtin] compat shim re-exporting pyspark.sql.functions, so upstream connect-facing tests import unmodified. The Connect-server pipeline glue (create_dataflow_graph/start_run + the SparkConnectGraphElementRegistry) and its server-backed tests need a Rust PipelineCommand execution path and are the next step. Co-authored-by: Isaac <no-reply@databricks.com>
…AGENTS.md playbook
- scripts/check_vendored_upstream.sh now guards individual vendored files (not just
whole dirs): the byte-identical files under pyspark/{logger,errors,pipelines} are
diffed against the v4.2.0 tag in CI. Split into VENDORED_DIRS (whole trees, e.g.
cloudpickle) and VENDORED_FILES (files whose sibling files are fork-adapted — e.g.
pipelines/api.py's import repointing, our Rust-backed errors/exceptions/__init__.py,
upstream's py4j/grpc errors/exceptions/{captured,connect}.py which we omit). Verified
all 21 files byte-identical against a live apache/spark v4.2.0 clone.
- AGENTS.md: append the hard-won lessons from the v4.2.0 parity push so the 4.3.0
upgrade doesn't rediscover them — module/package-level audit (the class diff misses
whole modules), faithful errors+logger vendoring + import-closure tracing, connect.*
compat shims for running official tests directly, DataType<->Python materialization /
df.schema property / StructType introspection / UDT, the pipelines/SDP surface + the
Rust PipelineCommand seam the server path needs, and proto-superset verification.
Co-authored-by: Isaac <no-reply@databricks.com>
…mmand seam)
Implements the server-backed SDP path so a dataflow graph can be created, populated,
and run against a Spark Connect server through the Rust transport — verified end-to-end
against a live Spark 4.2.0 server (a materialized-view dry-run completes; a streaming
table fed by a batch query raises AnalysisException carrying
INVALID_FLOW_QUERY_TYPE.BATCH_RELATION_FOR_STREAMING_TABLE during event iteration, exactly
as the official pipelines/tests/test_spark_connect.py expects).
- crates/spark-connect/src/pipelines.rs: build + execute PipelineCommand.{CreateDataflowGraph
(returns graph id), DefineOutput (table/mv/view/sink details), DefineFlow (relation +
auto-CDC), DefineSqlGraphElements, StartRun}. StartRun returns a PipelineRunStream that
opens its gRPC stream lazily on first pull, so a fail-fast dry-run validation error
surfaces during iteration (not from start_run) — matching the reference generator.
- crates/pyspark-rs/src/pipelines.rs: module-level _pyspark.pipeline_* functions +
PipelineRunStream iterator (yields (message, timestamp_micros); maps server errors to the
pyspark exception type). Module functions rather than SparkSession methods since
multiple-pymethods is unavailable.
- python/pyspark/pipelines/spark_connect_{pipeline,graph_element_registry}.py: vendored then
rewired to call the Rust functions instead of assembling pb2.Command + the Python gRPC
client (df._plan.plan / col.to_plan / client.execute_command).
- python/pyspark/errors/exceptions/connect.py: re-export the canonical error classes under
the connect import path (the Rust transport already raises them), so upstream connect-facing
code/tests import unmodified.
Remaining for running test_spark_connect.py verbatim: ReusedConnectTestCase harness compat
(SparkConf, SparkSession.client, builder.config(conf=)).
Co-authored-by: Isaac <no-reply@databricks.com>
Replace the hand-written minimal pyspark/sql/__init__.py (which exported only a
handful of names, causing repeated "cannot import name ..." gaps) with the official
v4.2.0 __init__ verbatim, backed by thin submodules that re-export the Rust-backed
classes under their official import paths — so imports match upstream for a drop-in:
- new sql/{session,column,group,readwriter,merge}.py (re-export the Rust classes),
sql/context.py (legacy SQLContext/HiveContext + UDFRegistration/UDTFRegistration),
sql/utils.py (is_remote / is_timestamp_ntz_preferred), sql/internal.py (InternalFunction).
- dataframe.py also re-exports DataFrameNaFunctions/DataFrameStatFunctions.
- types.py: re-export Row; add module-level cast (upstream leaks typing.cast from this
module), _drop_metadata, and _create_row (builds our Rust Row from name/value pairs).
- vendor pyspark.sql.pandas pure helpers (utils/types/typehints/conversion/map_ops/__init__)
verbatim; connect-coupled group_ops/functions are thin re-exports of Rust
PandasCogroupedOps / pandas_udf.
pyspark.sql and its submodules import cleanly; 82 offline+pipeline tests stay green.
Note: pyspark.pandas still does not import (the tracked subset references data_type_ops
which was never included; full upstream additionally needs the canonical Row to be a tuple
subclass with mutable __fields__ — PyO3 cannot subclass tuple, so it requires adopting the
upstream Python Row and returning it from collect()/head()/take()). Tracked separately.
Co-authored-by: Isaac <no-reply@databricks.com>
…CI (fmt/docs)
- Ergonomics sweep (matching the earlier Vec<Column> -> impl IntoIterator work, which
had only covered Column lists): DataFrameWriter/V2 cluster_by/partition_by/bucket_by/
sort_by and DataFrame join_using/hint now take `impl IntoIterator<Item = impl Into<String>>`,
so `.partition_by(["year", "month"])` (array of &str) works, not just Vec<String>.
Existing Vec<String> callers still compile (Vec<String>: IntoIterator<Item=String>).
Empty-list call sites annotated as Vec::<String>::new().
- docs/data-sources.md: use the ergonomic `.partition_by(["year", "month"])` form.
- rustfmt: format the new pipelines.rs (+ incidental) so the blocking Lint gate passes.
Fixes the PR CI: Lint (rustfmt) was failing on pipelines.rs; the Rust CI / Coverage
`repartition_by_id` compile error was stale (that test file no longer exists; the current
e2e_coverage.rs already passes col("id")). All tests compile; 205 lib tests pass; all 51
doc snippets compile.
Co-authored-by: Isaac <no-reply@databricks.com>
…efault) - Replace the diverged 48-file pyspark.pandas subset (which referenced data_type_ops that was never included, so it never imported) with the full v4.2.0 upstream pandas-on-Spark package (67 files incl. data_type_ops/indexes/missing/plot/spark/ typedef/usage_logging). `import pyspark.pandas` now succeeds. Two pure upstream helpers it needs at import were added verbatim to pyspark/sql/utils.py (get_lit_sql_str, pyspark_column_op). - Real parity bug surfaced by pandas runtime and fixed: RuntimeConf.get(key, default) now returns the default for an unknown key (server GetWithDefault op) instead of raising SQL_CONF_NOT_FOUND. get(key) without a default still raises, mirroring RuntimeConfig.get. Remaining for pandas-on-Spark runtime (pandas parity tests): createDataFrame(pandas_df) Arrow ingestion (our createDataFrame currently accepts only Python lists), and likely further boundaries — tracked as follow-on work; import + config are done. Co-authored-by: Isaac <no-reply@databricks.com>
…ol, df.columns property Real parity gaps surfaced by exercising pyspark.pandas against a live server, each fixed mirroring pyspark semantics: - SparkSession.createDataFrame now accepts a pandas DataFrame (converts rows with NaN/NaT -> None via astype(object).where(notna, None), taking column names from the frame), in addition to lists/tuples/Rows. (A future optimization is Arrow-batch ingestion into the LocalRelation instead of a Python row round-trip.) - DataFrame._col(name): the private column accessor pandas-on-Spark's scol_for uses. - DataFrame.columns is now a property (was a method) — pandas-on-Spark does `x not in sdf.columns`; mirrors DataFrame.columns in pyspark. (Same property-vs-method class of bug as df.schema.) pyspark.pandas now imports and its DataFrame construction from pandas runs; further pandas-on-Spark runtime relies on connect internals (spark_column_equals / pyspark.sql.connect.column, InternalFrame column-equality) that are the next increments. Co-authored-by: Isaac <no-reply@databricks.com>
…+ AGENTS.md gotchas - .github/workflows/pandas-parity.yml: a separate "Pandas Connect Parity (Spark 4.2.0)" workflow that runs the official pandas-on-Spark connect tests through our Rust transport on their own (they take 30-60 min), with timeout-minutes: 90 and continue-on-error: true (advisory) until pandas-on-Spark runtime parity is complete — so it never blocks the main CI gate. Modeled on the existing connect-parity job. - python/pyspark/sql/connect/column.py: re-export the Rust-backed Column under the connect import path, so pandas-on-Spark's spark_column_equals (and connect-facing tests) get a working `isinstance(x, ConnectColumn)`. (It compares repr(); a faithful expression-string Column.__repr__ is the next pandas-runtime step — the core has no expression pretty-printer yet, so today repr is generic.) - AGENTS.md section 17: the continued v4.2.0 gotchas (property-vs-method bugs, RuntimeConf.get default, createDataFrame(pandas)+Arrow note, full pyspark.pandas vendoring, the pandas-runtime connect-internal chain, the Row-tuple blocker, drop-in module layout, Vec<String> sweep). Co-authored-by: Isaac <no-reply@databricks.com>
…rg repartition_by_id
The PR was 9 commits behind master, which added tests (e2e_ops_breadth.rs et al.). CI
tests the PR merged with master, so master's `df.repartition_by_id(2)` (old 1-arg call)
broke against this PR's corrected 2-arg `repartition_by_id(num_partitions, partition_id_col)`
(mirroring pyspark's repartitionById). Merged master and updated the test to
`repartition_by_id(2, col("id"))` and `hint("broadcast", Vec::<String>::new())` (the latter
for the Vec<String>->IntoIterator sweep). dataframe.rs/AGENTS.md/pyproject.toml auto-merged;
all tests compile, fmt clean, 51 doc snippets compile, 82 offline+pipeline tests pass.
Co-authored-by: Isaac <no-reply@databricks.com>
The merge-reconciliation edit left `df.hint("broadcast", Vec::<String>::new()).collect()
.unwrap();` on one over-length line; rustfmt (blocking Lint gate) wants it wrapped.
Co-authored-by: Isaac <no-reply@databricks.com>
… guidelines Consolidate the appended session-specific gotchas (former sections 10-17) into six general principle sections (10-15): parity is package/module/attribute-level not just class-method names; vendor pure-Python client modules faithfully and trace import closures; mirror pyspark method semantics (documented properties as getters, default/error behavior) and apply signature sweeps consistently; DataType picklability/materialization; version-tag vendoring + proto-superset verification; and the standing architectural constraints (Row tuple limitation, pandas-on-Spark runtime as a multi-step effort, the pipelines seam). Removes one-off file paths, line numbers, and "fixed this session" narrative so the guidance stays durable. Co-authored-by: Isaac <no-reply@databricks.com>
…arkSession.client, testing.utils Groundwork so official ReusedConnectTestCase-based suites (pipelines/pandas/...) can run against this drop-in: - SparkSession.Builder.config now accepts conf= (a SparkConf; applies its getAll() pairs), matching pyspark's config(conf=...). - SparkSession.client returns a minimal SparkConnectClientStub exposing the members test/util code touches (_server_session_id, _cleanup_ml_cache); the real transport is Rust. - Vendor pyspark/testing/utils.py from v4.2.0 (it was MISSING — connectutils/mlutils/pandasutils imported it and silently fell back to object stubs, causing "duplicate base class" errors). Its imports are all satisfied by our drop-in. Provides PySparkBaseTestCase / PySparkErrorTestUtils / should_test_connect / assertDataFrameEqual etc. Full pyspark.testing (connectutils' connect-internal deps) follows. Co-authored-by: Isaac <no-reply@databricks.com>
Record the guiding rule: every module in the pyspark-client connect_packages list (python/packaging/client/setup.py at the tracked version tag) must be either vendored verbatim (pure-Python) or exposed with equivalent Rust-backed logic under the same import path — no third option. Calls out the non-obvious public APIs (pyspark.testing, pandas, sql.pandas, sql.plot, logger, errors, ml, pipelines) and the one exception (connect.* gRPC internals, replaced by the Rust transport, with thin re-export shims where imported). Co-authored-by: Isaac <no-reply@databricks.com>
…-equivalent) v4.2.0 makes pyspark.sql.functions a package (builtin.py + partitioning.py). Our functions is a Rust-backed module that already registers a synthetic `partitioning` submodule; add a `builtin` synthetic submodule the same way so `import pyspark.sql.functions.builtin` and `from pyspark.sql.functions.builtin import col` work like reference pyspark, completing the connect_packages "pyspark.sql.functions" entry. Co-authored-by: Isaac <no-reply@databricks.com>
…ng/plot/worker/testing
Reconcile the bulk of the v4.2.0 pyspark-client connect_packages list — every entry must be
vendored verbatim (pure-Python) or Rust-equivalent. Vendored from v4.2.0 and made importable
against this drop-in (classic SparkContext/RDD-only deps guarded so modules still import,
matching what the connect-only wheel can do):
- pyspark.ml.{linalg,param,torch,deepspeed} (+ ml/util.py, ml/common.py, ml/dl_util.py helpers)
- pyspark.mllib (+ linalg fully functional, stat guarded)
- pyspark.streaming (DStream; classic deps guarded)
- pyspark.sql.plot (df.plot backend; + require_minimum_plotly_version/linspace helpers)
- pyspark.sql.worker (+ top-level accumulators.py, profiler.py, worker_util.py, memory_profiler_ext.py,
sql/conversion.py, sql/datasource_internal.py, sql/profiler.py; worker submodules import-blocked
only on classic local_connect_and_auth, as in a connect-only client)
- pyspark.testing: the full public API — vendored sqlutils/unittestutils/find_spark_home, working
assertDataFrameEqual/assertSchemaEqual + ReusedConnectTestCase/should_test_connect
- thin pyspark.sql.connect.* shims (dataframe/session/column/functions/plan/proto) re-exporting the
Rust-backed classes so official connect-facing code and tests import unmodified
Supporting fixes:
- PyDataFrame is now #[pyclass(subclass)] so Python can subclass the connect DataFrame
(testing.connectutils' MockDF), mirroring pyspark.
- util.py gains VersionUtils / _print_missing_jar (used by ml/streaming).
Remaining connect_packages: the deeper connect internals (sql.connect.{proto,client,avro,protobuf,
resource,shell,streaming[.worker]}) and sql.streaming.proto (generated protobuf) — follow-ups.
Co-authored-by: Isaac <no-reply@databricks.com>
…iles in drift check
- Thin shim packages for the rest of the connect_packages connect internals:
sql.connect.{avro,protobuf,resource,streaming(+worker),client,shell} re-export the
Rust-backed equivalents; sql.streaming.proto is a present placeholder (real protobuf
bindings are a follow-up). With the earlier dataframe/session/column/functions/plan/proto
shims, the connect import surface official code/tests use is now covered.
- scripts/check_vendored_upstream.sh: guard the byte-identical vendored files against the
v4.2.0 tag — 11 whole-tree dirs (pandas/*, mllib/linalg, sql/{worker,plot}, testing/tests)
and 82 individual files (logger/errors/pipelines/ml/mllib/pandas/sql.pandas/streaming/
testing/top-level). Adapted files (import-guarded / helper-augmented) are excluded.
Co-authored-by: Isaac <no-reply@databricks.com>
…show real pass/fail The advisory steps were vacuously green: raw pytest could not import rust_transport_plugin (scripts/ was not on PYTHONPATH) and continue-on-error hid the setup failure. Fix to mirror run_official_tests.py: PYTHONPATH=scripts:<spark-source>/python, working-directory at the Spark source python/ with a relative test path, RUST_PYSPARK_SO set, and remove the step-level continue-on-error so a broken setup or a real test failure turns the job red. These are separate, non-required workflows, so red is honest signal rather than a merge blocker. Co-authored-by: Isaac <no-reply@databricks.com>
…InPandas] - Vendor the arbitrary-stateful-processing Python API from v4.2.0 into pyspark/sql/streaming/: stateful_processor (StatefulProcessor, StatefulProcessorHandle, ValueState/ListState/MapState, TimerValues, ExpiredTimerInfo), stateful_processor_util (TransformWithStateInPandasUdfUtils), the value/list/map state clients + stateful_processor_api_client, state.py, and the driver worker (server-only imports guarded). These are cloudpickled and executed by the server's Python worker (like DataSource/UDFs), so they must be the upstream Python classes. - Rust-back the client-side method: GroupedData.transformWithState / transformWithStateInPandas now take a StatefulProcessor (+ outputStructType/outputMode/timeMode/initialState/ eventTimeColumnName), wrap it via TransformWithStateInPandasUdfUtils into the transform-with- state UDF, and build the GroupMap proto with transform_with_state_info in the Rust core (which already supported it). Verified end-to-end constructing the plan against a live server. Co-authored-by: Isaac <no-reply@databricks.com>
…rop BLOCKING prefix The full pipeline suite ran 40 passed / 27 skipped: the skipped ones are test_cli/test_init_cli, gated on have_yaml (PyYAML is a pyspark-client install dep we were not installing). Add pyyaml to the pandas + pipelines workflow deps so those run instead of skip. Also drop the "BLOCKING -"/ "ADVISORY"/"FULL SUITE" labels from the pipelines workflow step names/comments — every step is blocking now. Co-authored-by: Isaac <no-reply@databricks.com>
…n (not Rust) Record the principle behind UDF/UDT/DataSource/StatefulProcessor: they are cloudpickled and run on the server Python worker, so the classes must be the upstream Python ones (a Rust class breaks the cloudpickle round-trip and never runs on the client); only the client-side proto-building METHOD is Rust-backed. Proto-building -> Rust; cloudpickled server code -> Python. Co-authored-by: Isaac <no-reply@databricks.com>
VariantVal.toJson/toPython/parseJson (client-side variant decoding) delegate to the vendored, byte-identical upstream variant_utils.py, so variant client parity is already met. Guard it against drift like the other verbatim-vendored files. (A Rust-native variant codec would be a spark-connect crate enhancement for Rust API users, not a pyspark-client parity item.) Co-authored-by: Isaac <no-reply@databricks.com>
Mirror pyspark createDataFrame input forms that were missing: a list of dicts (field names = sorted keys, per _infer_schema) and a list of scalars (single-field rows, e.g. createDataFrame([1,2,3], IntegerType())), in addition to list/tuple/Row/pandas. Co-authored-by: Isaac <no-reply@databricks.com>
…ay/map/struct/Row) Rigorous fix after finding that createDataFrame with nested types was broken. Now handles array/map/struct (and nested Rows) end-to-end, verified round-tripping against a live server. - schema_to_arrow_fields now uses a new recursive datatype_to_arrow (Spark DataType -> Arrow DataType incl. List/Map/Struct), instead of a scalar-only match. - build_arrow_array is now schema-driven: a recursive values_to_arrow builds the column against its declared DataType, recursing into ListArray/MapArray/StructArray (offsets + child arrays + null buffers). The old value-driven scalar builder is kept only as a no-declared-type fallback. - coerce_value is now recursive into array/struct/map, so a nested Integer field built from a Python int (which decodes as Long) is coerced to match the Arrow schema (previously only top-level columns were coerced -> "Type mismatch in row data" for nested fields). - py_to_value: a pyspark Row value now becomes Value::Struct (nested Rows convert), and the binary check is bytes/bytearray-specific — `extract::<Vec<u8>>()` also matched a Python list of small ints like [1,2,3], mis-tagging arrays as binary. Co-authored-by: Isaac <no-reply@databricks.com>
… raw dict
Previously a VARIANT column materialized as a plain {value, metadata} dict because the
arrow->row conversion turned the physical struct<value binary, metadata binary> into a
generic struct. Now the conversion recognizes a variant column (its inner `metadata` arrow
field carries metadata {"variant": "true"}) and produces a new core Value::Variant carrying
the raw bytes, which pyspark-rs materializes as a pyspark.sql.types.VariantVal — so
row.var.toJson()/toPython() work, matching pyspark. Decoding stays lazy in Python
(variant_utils); only the VariantVal wrapping is done in Rust.
Co-authored-by: Isaac <no-reply@databricks.com>
…das.usage_logging import) pyspark.pandas.usage_logging imports pyspark.instrumentation_utils (a pure-Python top-level module, not itself in connect_packages but a transitive dep). Vendor it verbatim from v4.2.0 so usage_logging imports; guard it in the drift check. All 45 connect_packages now import. Co-authored-by: Isaac <no-reply@databricks.com>
…ma inference Closes the three remaining type gaps found by a rigorous round-trip audit: - Decimal: py_to_value now checks decimal.Decimal BEFORE int/float (Decimal defines __float__/__int__, so it was mis-extracted as Double), and values_to_arrow builds the Decimal128 array with the DECLARED precision/scale (was hardcoded to 38, mismatching a DecimalType(10,2) Arrow field). createDataFrame(Decimal, DecimalType(p,s)) now round-trips. - NullType: arrow_value_at handles an all-null NullArray -> Value::Null (collect no longer errors "Unsupported Arrow type Null"). - Nested schema inference (schema=None): value_to_datatype now infers Array/Struct/Map/Variant recursively, so createDataFrame([Row(id=1, info=Row(a=1, b=[1,2]))]) infers the nested schema. Verified: decimal/null/nested-Row-inferred/array/map/struct/date/timestamp all round-trip. Co-authored-by: Isaac <no-reply@databricks.com>
…npacks a list; _to_pandas
pandas-on-Spark's spark_column_equals compares repr(col) strings, so the generic
"Column()" repr broke column identity throughout the pandas layer (e.g. ps.range()
returned an empty (n, 0) frame). Fixes, all in Rust (no Python monkey-patches):
- Expression::render() mirrors pyspark.sql.connect.expressions.*.__repr__:
ColumnReference -> name, Literal -> value (NULL for null), UnresolvedFunction ->
infix "(a op b)" for binary operators / "(NOT x)" / "name(args)", Alias -> "x AS y",
Cast -> "CAST(x AS t)", SortOrder, CaseWhen, ExtractValue, UpdateFields, star.
Column.__repr__/__str__ now return Column<'<expr>'>.
- to_column_list unpacks a single list/tuple arg, so df.select(["a","b"]) works like
df.select("a","b") (matches pyspark) -- fixes every *cols method at once.
- DataFrame._to_pandas(**kwargs): the Arrow-based conversion pandas-on-Spark calls
(internal.py: sdf._to_pandas(pandasStructHandlingMode="row")); delegates to toPandas.
- connect.functions.builtin._invoke_function_over_columns now faithfully mirrors
upstream: builds an UnresolvedFunction by name (server resolves the internal
functions distributed_sequence_id / pandas_* etc.) instead of a wrong client-side
row_number-window reimplementation.
Reverts the earlier agent's vendored-file edit (pandas/internal.py) and the
sql/dataframe.py select monkey-patch in favor of these Rust fixes.
Co-authored-by: Isaac <no-reply@databricks.com>
…unpack gotchas General guidelines so future agents don't reintroduce the Column.__repr__ constant, the client-side distributed_sequence_id reimplementation, or per-method list unpacking, and never work around Rust gaps by editing vendored files or monkey-patching. Co-authored-by: Isaac <no-reply@databricks.com>
…n unpacks a list Fixes the two real pandas-on-Spark parity failures (get_dummies on datetime, Series.isin with a datetime), plus adds scipy to the pandas parity env for the corr tests. - to_column (the lit() path) only handled bool/int/float/str/None. A collected timestamp comes back as a pandas Timestamp (a datetime subclass whose __int__ returns nanoseconds), so the early int extract silently turned it into a huge integer literal that never matched the TIMESTAMP column -> get_dummies produced all-zeros. Now a shared scalar_literal() maps datetime -> Timestamp(micros), date -> Date(days), Decimal -> Decimal(value,precision,scale), bytes/bytearray -> Binary, checked BEFORE int/float, mirroring session::py_to_value so lit() and createDataFrame() encode identically. - Column.isin now unpacks a single list/tuple/set argument (col.isin([a,b]) == col.isin(a,b)), matching pyspark; pandas-on-Spark's Series.isin passes a single list of lit() columns. - pandas-parity workflow installs scipy (Apache's pandas test env has it; corr method= 'spearman'/'kendall' and corrwith need it). Co-authored-by: Isaac <no-reply@databricks.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
Closes the remaining public pyspark v4.2.0 API gaps in the Rust Spark Connect client
drop-in, found by a rigorous introspection audit of
pyspark.sql.connect.*vs thedrop-in. Everything is mirrored from the official implementation and backed by the Rust
core (with a thin PyO3/Python surface), and every addition has tests.
Highlights (per-commit):
transform/filter/exists/forall/aggregate/reduce/zip_with/transform_keys/transform_values/map_filter/map_zip_with) via afaithful
_create_lambda/_invoke_higher_order_functionport, pluscume_dist/broadcast/call_function/call_udf/column/random;udtf/arrow_udtf/arrow_udfdecorators.Fixed
CallFunction/call_function/call_udfdropping their arguments.truncateTable/analyzeTable/getCreateTableString/getTableProperties/listPartitions/listViews)
and the result classes (Table/Database/Function/Column/CatalogMetadata/TablePartition),
so
listTables()etc. return typed objects like reference pyspark.needConversion/fromInternal/toInternal/fromDDL; StructType toDDL/treeString/toNullable/
fromJson/fieldNames; StructField collation methods), and the real class hierarchy
(DataType/AtomicType/NumericType/IntegralType/FractionalType/DatetimeType/AnyTimeType/
AnsiIntervalType/SpatialType) with the reference MRO so
isinstance(dt, NumericType)works.New
GeometryType/GeographyType, and theVariantVal/Geometry/Geographyvalue holders.and a native Rust
StreamingQueryManagerlistener bus (trait + background dispatch) soaddListener/removeListener/close work (Python is a thin adapter).
UserDefinedTableFunction.asDeterministic.
the UDTF analyze classes (AnalyzeResult/AnalyzeArgument/PartitioningColumn/OrderingColumn/
SelectedColumn + SkipRestOfInputTableException);
SparkSession.Builder.channelBuildernowreconstructs an sc:// URL from the builder instead of raising; re-export paths
(
pyspark.StorageLevel,pyspark.sql.streaming.StreamingQueryListener,pyspark.sql.{avro,protobuf}.functions,pyspark.sql.functions.partitioning).Vec<Column>API now takesimpl IntoIterator<Item = impl Into<Column>>(arrays,vec!, iterators, and&strcolumnnames), with
From<&str>/String for Column. Docs/README/examples updated to the array form.After this PR the introspection audit shows 0 method gaps across all core SQL classes
(DataFrame/Column/SparkSession/GroupedData/Reader/Writer/Catalog/Window/Observation/streaming/
types) and 0 gaps in the functions module vs pyspark v4.2.0.
Why are the changes needed?
The drop-in under-exposed the (already near-complete) Rust core and diverged on several
signatures; some newer 4.2.0 types/classes were missing entirely. This brings the client to
method-for-method parity with pyspark v4.2.0's public Connect surface.
Does this PR introduce any user-facing change?
Yes — many previously-missing public APIs are now available; behavior mirrors pyspark v4.2.0.
The only intentional non-additions are internal transport plumbing (
Column.to_plan,SparkSession.client) that don't apply to a native Rust transport.How was this patch tested?
Rust golden/plan/unit tests + server-gated e2e tests (run against a live Spark 4.2.0 Connect
server) + offline Python drop-in tests, all green;
cargo fmt/no-stub audit pass. Typeobject-model output was verified byte-for-byte against official pyspark.
Known follow-ups (not in this PR)
UserDefinedTypecustom-UDT protocol (Python-subclassing extensibility; coreDataType::Udtalready exists).
StatefulProcessor+DataFrame.transformWithState/transformWithStateInPandas(a large newstreaming stateful engine).
VariantVal.toJson/toPython/parseJsoncurrently use the vendored, byte-identical upstream
variant_utils).This pull request and its description were written by Isaac.