You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by vi...@apache.org on 2026/08/30 03:11:46 UTC
(spark-connect-rust) branch master updated: [SPARK-59117][PYTHON] Set __module__ on the PyO3 classes to the reference pyspark path
This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark-connect-rust.git
The following commit(s) were added to refs/heads/master by this push:
new 18c7abd [SPARK-59117][PYTHON] Set __module__ on the PyO3 classes to the reference pyspark path
18c7abd is described below
commit 18c7abd8d69397d27b08322dec73d6dbc06c7254
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Sat Aug 29 20:11:37 2026 -0700
[SPARK-59117][PYTHON] Set __module__ on the PyO3 classes to the reference pyspark path
### What changes were proposed in this pull request?
Add an explicit `module = "..."` to every `#[pyclass]` in `crates/pyspark-rs/src` (99
of them, including the `abstract_type!` macro that generates the 8 intermediate
DataType bases), and set the module on `create_exception!` for
`SkipRestOfInputTableException`.
PyO3 defaults a `#[pyclass]`'s `__module__` to `builtins` when no `module` argument is
given, and none of ours set it.
The class-to-module mapping is derived from where our own drop-in skin
(`python/pyspark/**`) already re-exports each class, cross-checked against Apache Spark
v4.2.0 (commit `32f72996011` — the same commit pinned in
`crates/spark-connect-proto/proto/PROTO_SHA.txt`).
Two mapping notes:
- `DataFrame` and `Column` map to `pyspark.sql.connect.*`, **not** `pyspark.sql.*`.
Upstream's `pyspark/sql/dataframe.py:78` and `pyspark/sql/column.py:41` hold the
*abstract* parent; the concrete class a Connect session hands the user is
`pyspark/sql/connect/dataframe.py:112` / `.../connect/column.py:109`, which is what
`type(df).__module__` reports on a real Connect session.
- `RustRpcError` deliberately keeps `_pyspark`: it is an internal transport error with
no upstream counterpart, so asserting a pyspark path for it would be false.
The hand-written `__reduce__` impls in `types.rs` are **kept**. They carry the type's
parameters (`DecimalType(12,3)`, `ArrayType(...)`), which a by-reference pickle of the
bare class would not restore.
### Why are the changes needed?
`__module__ == "builtins"` is a behavioral divergence from the reference client, not
just a cosmetic one:
1. **The class itself was unpicklable.** pickle serializes a class *by reference* —
writing `__module__` + `__qualname__` and re-importing on load — so the lookup
landed in `builtins` and failed:
```
>>> import pickle
>>> from pyspark.sql.types import IntegerType
>>> pickle.dumps(IntegerType)
PicklingError: Can't pickle <class 'builtins.IntegerType'>:
attribute lookup IntegerType on builtins failed
```
Real pyspark pickles these fine. Reproduced for `DataFrame`, `Column`, and every
`DataType` class.
2. **Error messages, reprs, and `__module__`-based dispatch showed `builtins.DataFrame`**
instead of the pyspark path — user-visible, and it undercuts the drop-in premise.
Instance pickling already worked, but only because each DataType hand-writes a
`__reduce__` routing through `pyspark.sql.types._parse_datatype_json_string`. The cost
of the wrong `__module__` was therefore already being paid, worked around per class
rather than fixed at the root.
This is the same "declared but never applied" gap as SPARK-59032 (parsed gRPC keepalive
never applied to the channel) and SPARK-59037 (128 MiB max message size declared but
unused): the intent existed, but the wiring that makes it take effect was missing.
### Does this PR introduce _any_ user-facing change?
Yes, and it is the point of the change: the drop-in's classes now report the same
`__module__` as the reference client (`pyspark.sql.types.IntegerType` rather than
`builtins.IntegerType`), so pickling a class works and error messages/reprs show the
pyspark path. No API surface is added or removed.
### How was this patch tested?
Four new offline tests in `python/tests/test_dropin_offline.py`:
- the module paths for `DataFrame` / `Column` / `SparkSession` / `DataType` /
`IntegerType` / `StructType` / `NumericType` (macro-generated base) / `Row`,
- that the DataType classes pickle by reference and round-trip to the *identical*
class object,
- and, as a guard against a future `#[pyclass]` landing without `module`, that no class
exported by `pyspark._pyspark` reports `builtins`.
Verified manually against a locally built extension: all **108** exported classes are
covered, none left on `builtins`. Regression-checked that the existing `__reduce__`
instance path is unaffected — `DecimalType(12,3)` → `decimal(12,3)`, nested
`StructType`, the UDF cloudpickle payload, `isinstance`, and the MRO all behave as
before.
`cargo test` passes (93 core tests + all suites, 0 failures). `cargo clippy -p
pyspark-rs` reports **389 warnings both before and after** this change — i.e.
unchanged from master, nothing introduced (all pre-existing `Default::default()` /
large-`Err` lints). `cargo fmt --check` is clean.
Not run in my local environment (no `pytest`/`pandas` on the system Python): the full
`test_dropin_offline.py` suite and the parity gates — those need CI. I verified the new
tests' assertions by executing their logic directly against the built extension.
Closes #89 from viirya/pyclass-module-attr.
Authored-by: Liang-Chi Hsieh <vi...@gmail.com>
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
---
crates/pyspark-rs/src/catalog.rs | 14 ++++----
crates/pyspark-rs/src/column.rs | 2 +-
crates/pyspark-rs/src/conf.rs | 2 +-
crates/pyspark-rs/src/dataframe.rs | 18 +++++------
crates/pyspark-rs/src/datasource.rs | 2 +-
crates/pyspark-rs/src/eval_type.rs | 2 +-
crates/pyspark-rs/src/group.rs | 4 +--
crates/pyspark-rs/src/ml.rs | 31 +++++++++++-------
crates/pyspark-rs/src/observation.rs | 2 +-
crates/pyspark-rs/src/pipelines.rs | 5 ++-
crates/pyspark-rs/src/profiler.rs | 2 +-
crates/pyspark-rs/src/readwriter.rs | 2 +-
crates/pyspark-rs/src/resource.rs | 8 ++---
crates/pyspark-rs/src/row.rs | 2 +-
crates/pyspark-rs/src/session.rs | 10 ++++--
crates/pyspark-rs/src/stat.rs | 2 +-
crates/pyspark-rs/src/storagelevel.rs | 7 +++-
crates/pyspark-rs/src/streaming.rs | 28 +++++++++++-----
crates/pyspark-rs/src/tablearg.rs | 2 +-
crates/pyspark-rs/src/transport.rs | 4 +--
crates/pyspark-rs/src/tvf.rs | 2 +-
crates/pyspark-rs/src/types.rs | 60 ++++++++++++++++++-----------------
crates/pyspark-rs/src/udtf_analyze.rs | 19 +++++++----
crates/pyspark-rs/src/values.rs | 6 ++--
crates/pyspark-rs/src/window.rs | 6 ++--
python/tests/test_dropin_offline.py | 51 +++++++++++++++++++++++++++++
26 files changed, 193 insertions(+), 100 deletions(-)
diff --git a/crates/pyspark-rs/src/catalog.rs b/crates/pyspark-rs/src/catalog.rs
index ea2f675..e4736f2 100644
--- a/crates/pyspark-rs/src/catalog.rs
+++ b/crates/pyspark-rs/src/catalog.rs
@@ -9,7 +9,7 @@ use crate::dataframe::PyDataFrame;
use crate::errors::ResultExt;
/// `pyspark.sql.catalog.CatalogMetadata`.
-#[pyclass(name = "CatalogMetadata")]
+#[pyclass(name = "CatalogMetadata", module = "pyspark.sql.catalog")]
pub struct PyCatalogMetadata {
inner: CatalogMetadata,
}
@@ -32,7 +32,7 @@ impl PyCatalogMetadata {
}
/// `pyspark.sql.catalog.Database`.
-#[pyclass(name = "Database")]
+#[pyclass(name = "Database", module = "pyspark.sql.catalog")]
pub struct PyDatabase {
inner: Database,
}
@@ -60,7 +60,7 @@ impl PyDatabase {
}
/// `pyspark.sql.catalog.Table`.
-#[pyclass(name = "Table")]
+#[pyclass(name = "Table", module = "pyspark.sql.catalog")]
pub struct PyTable {
inner: Table,
}
@@ -108,7 +108,7 @@ impl PyTable {
/// `pyspark.sql.catalog.Column` (exposed as `CatalogColumn` in the extension to avoid
/// colliding with the expression `Column`; re-exported as `Column` by pyspark.sql.catalog).
-#[pyclass(name = "CatalogColumn")]
+#[pyclass(name = "CatalogColumn", module = "pyspark.sql.catalog")]
pub struct PyCatalogColumn {
inner: Column,
}
@@ -151,7 +151,7 @@ impl PyCatalogColumn {
}
/// `pyspark.sql.catalog.Function`.
-#[pyclass(name = "Function")]
+#[pyclass(name = "Function", module = "pyspark.sql.catalog")]
pub struct PyFunction {
inner: Function,
}
@@ -190,7 +190,7 @@ impl PyFunction {
}
/// `pyspark.sql.catalog.TablePartition`.
-#[pyclass(name = "TablePartition")]
+#[pyclass(name = "TablePartition", module = "pyspark.sql.catalog")]
pub struct PyTablePartition {
inner: TablePartition,
}
@@ -206,7 +206,7 @@ impl PyTablePartition {
}
/// Python wrapper for Spark Catalog.
-#[pyclass(name = "Catalog")]
+#[pyclass(name = "Catalog", module = "pyspark.sql.catalog")]
pub struct PyCatalog {
pub(crate) catalog: Catalog,
}
diff --git a/crates/pyspark-rs/src/column.rs b/crates/pyspark-rs/src/column.rs
index 75b8643..bbd1860 100644
--- a/crates/pyspark-rs/src/column.rs
+++ b/crates/pyspark-rs/src/column.rs
@@ -30,7 +30,7 @@ fn raise_pyspark(
}
/// Python wrapper for a Spark Column.
-#[pyclass(name = "Column")]
+#[pyclass(name = "Column", module = "pyspark.sql.connect.column")]
pub struct PyColumn {
pub(crate) column: Column,
}
diff --git a/crates/pyspark-rs/src/conf.rs b/crates/pyspark-rs/src/conf.rs
index cf67532..999959b 100644
--- a/crates/pyspark-rs/src/conf.rs
+++ b/crates/pyspark-rs/src/conf.rs
@@ -6,7 +6,7 @@ use spark_connect::conf::RuntimeConf;
use crate::errors::ResultExt;
/// Python wrapper for the session runtime configuration.
-#[pyclass(name = "RuntimeConf")]
+#[pyclass(name = "RuntimeConf", module = "pyspark.sql.connect.conf")]
pub struct PyRuntimeConf {
conf: RuntimeConf,
}
diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs
index 00be354..b06f4d3 100644
--- a/crates/pyspark-rs/src/dataframe.rs
+++ b/crates/pyspark-rs/src/dataframe.rs
@@ -44,7 +44,7 @@ fn flatten_str_cols(cols: Vec<Bound<'_, PyAny>>) -> PyResult<Vec<String>> {
}
/// Python wrapper for a Spark DataFrame.
-#[pyclass(name = "DataFrame", subclass)]
+#[pyclass(name = "DataFrame", module = "pyspark.sql.connect.dataframe", subclass)]
pub struct PyDataFrame {
pub(crate) dataframe: DataFrame,
}
@@ -1462,7 +1462,7 @@ impl PyDataFrame {
}
/// Python wrapper for `DataFrameNaFunctions` (`df.na`).
-#[pyclass(name = "DataFrameNaFunctions")]
+#[pyclass(name = "DataFrameNaFunctions", module = "pyspark.sql.dataframe")]
pub struct PyDataFrameNaFunctions {
inner: spark_connect::group::NaFunctions,
df: spark_connect::dataframe::DataFrame,
@@ -1582,7 +1582,7 @@ fn build_replacement_pairs(
}
/// Python wrapper for `MergeIntoWriter` (`df.mergeInto(...)`).
-#[pyclass(name = "MergeIntoWriter")]
+#[pyclass(name = "MergeIntoWriter", module = "pyspark.sql.merge")]
pub struct PyMergeIntoWriter {
inner: Option<spark_connect::MergeIntoWriter>,
}
@@ -1647,7 +1647,7 @@ fn assignments_to_map(
}
/// `whenMatched(...)` clause builder.
-#[pyclass(name = "WhenMatched")]
+#[pyclass(name = "WhenMatched", module = "pyspark.sql.merge")]
pub struct PyWhenMatched {
inner: Option<spark_connect::merge::WhenMatched>,
}
@@ -1691,7 +1691,7 @@ impl PyWhenMatched {
}
/// `whenNotMatched(...)` clause builder.
-#[pyclass(name = "WhenNotMatched")]
+#[pyclass(name = "WhenNotMatched", module = "pyspark.sql.merge")]
pub struct PyWhenNotMatched {
inner: Option<spark_connect::merge::WhenNotMatched>,
}
@@ -1725,7 +1725,7 @@ impl PyWhenNotMatched {
}
/// `whenNotMatchedBySource(...)` clause builder.
-#[pyclass(name = "WhenNotMatchedBySource")]
+#[pyclass(name = "WhenNotMatchedBySource", module = "pyspark.sql.merge")]
pub struct PyWhenNotMatchedBySource {
inner: Option<spark_connect::merge::WhenNotMatchedBySource>,
}
@@ -1769,7 +1769,7 @@ impl PyWhenNotMatchedBySource {
}
/// Python wrapper for the V1 `DataFrameWriter` (`df.write`).
-#[pyclass(name = "DataFrameWriter")]
+#[pyclass(name = "DataFrameWriter", module = "pyspark.sql.readwriter")]
pub struct PyDataFrameWriter {
inner: Option<spark_connect::readwriter::DataFrameWriter>,
}
@@ -2079,7 +2079,7 @@ impl PyDataFrameWriter {
}
/// Python wrapper for the V2 `DataFrameWriter` (`df.writeTo`).
-#[pyclass(name = "DataFrameWriterV2")]
+#[pyclass(name = "DataFrameWriterV2", module = "pyspark.sql.readwriter")]
pub struct PyDataFrameWriterV2 {
inner: Option<spark_connect::readwriter::DataFrameWriterV2>,
}
@@ -2270,7 +2270,7 @@ fn build_side_effect_udf(
/// Python wrapper for LocalRowIterator yielding Row objects lazily.
///
/// Implements `__iter__` and `__next__` to make it a proper Python iterator.
-#[pyclass(name = "LocalRowIterator")]
+#[pyclass(name = "LocalRowIterator", module = "pyspark.sql.dataframe")]
pub struct PyLocalRowIterator {
iterator: LocalRowIterator,
}
diff --git a/crates/pyspark-rs/src/datasource.rs b/crates/pyspark-rs/src/datasource.rs
index db02870..c6abc2a 100644
--- a/crates/pyspark-rs/src/datasource.rs
+++ b/crates/pyspark-rs/src/datasource.rs
@@ -12,7 +12,7 @@ use crate::session::PySparkSession;
/// Python accessor for data source registration.
///
/// Exposed as `SparkSession.dataSource` to register custom data sources.
-#[pyclass(name = "DataSourceRegistration")]
+#[pyclass(name = "DataSourceRegistration", module = "pyspark.sql.datasource")]
pub struct PyDataSourceRegistration {
pub(crate) session: PySparkSession,
}
diff --git a/crates/pyspark-rs/src/eval_type.rs b/crates/pyspark-rs/src/eval_type.rs
index 2257817..dd36068 100644
--- a/crates/pyspark-rs/src/eval_type.rs
+++ b/crates/pyspark-rs/src/eval_type.rs
@@ -5,7 +5,7 @@
use pyo3::prelude::*;
/// Evaluation-type constants for Python UDFs/UDTFs. Mirrors `pyspark.util.PythonEvalType`.
-#[pyclass(name = "PythonEvalType")]
+#[pyclass(name = "PythonEvalType", module = "pyspark.util")]
pub struct PyPythonEvalType;
#[allow(non_upper_case_globals)]
diff --git a/crates/pyspark-rs/src/group.rs b/crates/pyspark-rs/src/group.rs
index 4f129a0..8621926 100644
--- a/crates/pyspark-rs/src/group.rs
+++ b/crates/pyspark-rs/src/group.rs
@@ -9,7 +9,7 @@ use crate::errors::ResultExt;
use crate::functions::to_column;
/// Python wrapper for Spark GroupedData.
-#[pyclass(name = "GroupedData")]
+#[pyclass(name = "GroupedData", module = "pyspark.sql.group")]
pub struct PyGroupedData {
pub(crate) grouped_data: RustGroupedData,
}
@@ -363,7 +363,7 @@ fn twus_func<'py>(
}
/// Python wrapper for cogrouped pandas/arrow ops (`pyspark.sql.group.PandasCogroupedOps`).
-#[pyclass(name = "PandasCogroupedOps")]
+#[pyclass(name = "PandasCogroupedOps", module = "pyspark.sql.pandas.group_ops")]
pub struct PyCoGroupedData {
pub(crate) inner: RustCoGroupedData,
}
diff --git a/crates/pyspark-rs/src/ml.rs b/crates/pyspark-rs/src/ml.rs
index c8af505..252b38c 100644
--- a/crates/pyspark-rs/src/ml.rs
+++ b/crates/pyspark-rs/src/ml.rs
@@ -16,7 +16,7 @@ use crate::dataframe::PyDataFrame;
use crate::errors::ResultExt;
/// A fitted ML model (result of `Estimator.fit`), wrapping a boxed core `Model`.
-#[pyclass(name = "MLModel")]
+#[pyclass(name = "MLModel", module = "pyspark.ml.connect")]
pub struct PyMLModel {
model: Box<dyn Model>,
}
@@ -32,7 +32,7 @@ impl PyMLModel {
}
/// StandardScaler estimator.
-#[pyclass(name = "StandardScaler")]
+#[pyclass(name = "StandardScaler", module = "pyspark.ml.connect.feature")]
pub struct PyStandardScaler {
inner: StandardScaler,
}
@@ -72,7 +72,7 @@ impl PyStandardScaler {
}
/// MaxAbsScaler estimator.
-#[pyclass(name = "MaxAbsScaler")]
+#[pyclass(name = "MaxAbsScaler", module = "pyspark.ml.connect.feature")]
pub struct PyMaxAbsScaler {
inner: MaxAbsScaler,
}
@@ -112,7 +112,7 @@ impl PyMaxAbsScaler {
}
/// StringIndexer estimator.
-#[pyclass(name = "StringIndexer")]
+#[pyclass(name = "StringIndexer", module = "pyspark.ml.connect.feature")]
pub struct PyStringIndexer {
inner: StringIndexer,
}
@@ -152,7 +152,7 @@ impl PyStringIndexer {
}
/// VectorAssembler transformer.
-#[pyclass(name = "VectorAssembler")]
+#[pyclass(name = "VectorAssembler", module = "pyspark.ml.connect.feature")]
pub struct PyVectorAssembler {
inner: VectorAssembler,
}
@@ -194,7 +194,10 @@ impl PyVectorAssembler {
}
/// LogisticRegression estimator.
-#[pyclass(name = "LogisticRegression")]
+#[pyclass(
+ name = "LogisticRegression",
+ module = "pyspark.ml.connect.classification"
+)]
pub struct PyLogisticRegression {
inner: LogisticRegression,
}
@@ -257,7 +260,7 @@ impl PyLogisticRegression {
}
/// RegressionEvaluator.
-#[pyclass(name = "RegressionEvaluator")]
+#[pyclass(name = "RegressionEvaluator", module = "pyspark.ml.connect.evaluation")]
pub struct PyRegressionEvaluator {
inner: RegressionEvaluator,
}
@@ -308,7 +311,10 @@ impl PyRegressionEvaluator {
}
/// BinaryClassificationEvaluator.
-#[pyclass(name = "BinaryClassificationEvaluator")]
+#[pyclass(
+ name = "BinaryClassificationEvaluator",
+ module = "pyspark.ml.connect.evaluation"
+)]
pub struct PyBinaryClassificationEvaluator {
inner: BinaryClassificationEvaluator,
}
@@ -355,7 +361,7 @@ impl PyBinaryClassificationEvaluator {
}
/// Pipeline estimator (chains stages by name).
-#[pyclass(name = "Pipeline")]
+#[pyclass(name = "Pipeline", module = "pyspark.ml.connect.pipeline")]
pub struct PyPipeline {
inner: Pipeline,
}
@@ -387,7 +393,10 @@ impl PyPipeline {
}
/// MulticlassClassificationEvaluator.
-#[pyclass(name = "MulticlassClassificationEvaluator")]
+#[pyclass(
+ name = "MulticlassClassificationEvaluator",
+ module = "pyspark.ml.connect.evaluation"
+)]
pub struct PyMulticlassClassificationEvaluator {
inner: MulticlassClassificationEvaluator,
}
@@ -438,7 +447,7 @@ impl PyMulticlassClassificationEvaluator {
}
/// CrossValidator (k-fold tuning).
-#[pyclass(name = "CrossValidator")]
+#[pyclass(name = "CrossValidator", module = "pyspark.ml.connect.tuning")]
pub struct PyCrossValidator {
inner: CrossValidator,
}
diff --git a/crates/pyspark-rs/src/observation.rs b/crates/pyspark-rs/src/observation.rs
index 0e1f170..9ec3538 100644
--- a/crates/pyspark-rs/src/observation.rs
+++ b/crates/pyspark-rs/src/observation.rs
@@ -8,7 +8,7 @@ use spark_connect::observation::Observation;
static OBS_COUNTER: AtomicU64 = AtomicU64::new(0);
/// Python wrapper for a named metric Observation.
-#[pyclass(name = "Observation")]
+#[pyclass(name = "Observation", module = "pyspark.sql.observation")]
pub struct PyObservation {
pub(crate) inner: Observation,
}
diff --git a/crates/pyspark-rs/src/pipelines.rs b/crates/pyspark-rs/src/pipelines.rs
index a892a21..40c41b8 100644
--- a/crates/pyspark-rs/src/pipelines.rs
+++ b/crates/pyspark-rs/src/pipelines.rs
@@ -245,7 +245,10 @@ pub fn pipeline_define_sql_graph_elements(
/// `(message: str, timestamp_micros: int)` per pipeline event; command-result responses
/// are skipped; a server-side failure is raised (mapped to the pyspark exception type,
/// e.g. `AnalysisException`) during iteration.
-#[pyclass(name = "PipelineRunStream")]
+#[pyclass(
+ name = "PipelineRunStream",
+ module = "pyspark.pipelines.spark_connect_pipeline"
+)]
pub struct PyPipelineRunStream {
inner: PipelineRunStream,
}
diff --git a/crates/pyspark-rs/src/profiler.rs b/crates/pyspark-rs/src/profiler.rs
index febd51f..652fde8 100644
--- a/crates/pyspark-rs/src/profiler.rs
+++ b/crates/pyspark-rs/src/profiler.rs
@@ -9,7 +9,7 @@ use crate::errors::ResultExt;
/// Python wrapper for the ProfilerCollector.
///
/// Exposed as `SparkSession.profile` to access profiler results.
-#[pyclass(name = "ProfilerCollector")]
+#[pyclass(name = "ProfilerCollector", module = "pyspark.sql.profiler")]
pub struct PyProfilerCollector {
pub(crate) inner: Arc<ProfilerCollector>,
}
diff --git a/crates/pyspark-rs/src/readwriter.rs b/crates/pyspark-rs/src/readwriter.rs
index 39d3270..75177fd 100644
--- a/crates/pyspark-rs/src/readwriter.rs
+++ b/crates/pyspark-rs/src/readwriter.rs
@@ -11,7 +11,7 @@ use crate::dataframe::PyDataFrame;
/// Python wrapper for the batch DataFrameReader. The core reader is a consuming
/// builder, so each step takes the inner value and returns a fresh wrapper (mirrors
/// the existing PyDataStreamReader).
-#[pyclass(name = "DataFrameReader")]
+#[pyclass(name = "DataFrameReader", module = "pyspark.sql.readwriter")]
pub struct PyDataFrameReader {
inner: Option<DataFrameReader>,
}
diff --git a/crates/pyspark-rs/src/resource.rs b/crates/pyspark-rs/src/resource.rs
index bbca466..425e9f0 100644
--- a/crates/pyspark-rs/src/resource.rs
+++ b/crates/pyspark-rs/src/resource.rs
@@ -6,7 +6,7 @@ use spark_connect::resource::{
};
/// Python wrapper for ExecutorResourceRequests.
-#[pyclass(name = "ExecutorResourceRequests")]
+#[pyclass(name = "ExecutorResourceRequests", module = "pyspark.resource.profile")]
pub struct PyExecutorResourceRequests {
pub(crate) inner: ExecutorResourceRequests,
}
@@ -61,7 +61,7 @@ impl PyExecutorResourceRequests {
}
/// Python wrapper for TaskResourceRequests.
-#[pyclass(name = "TaskResourceRequests")]
+#[pyclass(name = "TaskResourceRequests", module = "pyspark.resource.profile")]
pub struct PyTaskResourceRequests {
pub(crate) inner: TaskResourceRequests,
}
@@ -92,7 +92,7 @@ impl PyTaskResourceRequests {
}
/// Python wrapper for ResourceProfileBuilder.
-#[pyclass(name = "ResourceProfileBuilder")]
+#[pyclass(name = "ResourceProfileBuilder", module = "pyspark.resource.profile")]
pub struct PyResourceProfileBuilder {
pub(crate) inner: ResourceProfileBuilder,
}
@@ -133,7 +133,7 @@ impl PyResourceProfileBuilder {
}
/// Python wrapper for ResourceProfile.
-#[pyclass(name = "ResourceProfile")]
+#[pyclass(name = "ResourceProfile", module = "pyspark.resource.profile")]
pub struct PyResourceProfile {
pub(crate) inner: ResourceProfile,
}
diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs
index 4f4b4da..524bb73 100644
--- a/crates/pyspark-rs/src/row.rs
+++ b/crates/pyspark-rs/src/row.rs
@@ -6,7 +6,7 @@ use pyo3::types::{PyBytes, PyDict, PyList};
use spark_connect::row::{Row, Value};
/// Python wrapper for a Spark Row.
-#[pyclass(name = "Row")]
+#[pyclass(name = "Row", module = "pyspark.sql.types")]
pub struct PyRow {
pub(crate) row: Row,
}
diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs
index e319015..03175f7 100644
--- a/crates/pyspark-rs/src/session.rs
+++ b/crates/pyspark-rs/src/session.rs
@@ -19,7 +19,11 @@ use crate::streaming::{PyDataStreamReader, PyStreamingQueryManager};
///
/// Mirrors `pyspark.sql.SparkSession.Builder`: chainable and reached via the
/// `SparkSession.builder` class attribute.
-#[pyclass(name = "SparkSessionBuilder", from_py_object)]
+#[pyclass(
+ name = "SparkSessionBuilder",
+ module = "pyspark.sql.session",
+ from_py_object
+)]
#[derive(Clone)]
pub struct PySparkSessionBuilder {
remote_url: Option<String>,
@@ -200,7 +204,7 @@ fn active_slot() -> &'static std::sync::Mutex<Option<SparkSession>> {
}
/// Python wrapper for a Spark session.
-#[pyclass(name = "SparkSession")]
+#[pyclass(name = "SparkSession", module = "pyspark.sql.session")]
pub struct PySparkSession {
pub(crate) session: SparkSession,
}
@@ -779,7 +783,7 @@ impl PySparkSession {
/// Minimal stand-in for the reference `SparkConnectClient`, returned by
/// `SparkSession.client`. The real client is the Rust transport; this only surfaces the
/// members that test/utility code (e.g. `ReusedConnectTestCase`) touches.
-#[pyclass(name = "SparkConnectClientStub")]
+#[pyclass(name = "SparkConnectClientStub", module = "pyspark.sql.connect.client")]
pub struct PyConnectClientStub {
session_id: String,
}
diff --git a/crates/pyspark-rs/src/stat.rs b/crates/pyspark-rs/src/stat.rs
index 9556216..c545fbf 100644
--- a/crates/pyspark-rs/src/stat.rs
+++ b/crates/pyspark-rs/src/stat.rs
@@ -7,7 +7,7 @@ use crate::dataframe::PyDataFrame;
use crate::errors::ResultExt;
/// Python wrapper for DataFrame statistical functions.
-#[pyclass(name = "DataFrameStatFunctions")]
+#[pyclass(name = "DataFrameStatFunctions", module = "pyspark.sql.dataframe")]
pub struct PyStatFunctions {
stat: StatFunctions,
}
diff --git a/crates/pyspark-rs/src/storagelevel.rs b/crates/pyspark-rs/src/storagelevel.rs
index c677d76..3fb410a 100644
--- a/crates/pyspark-rs/src/storagelevel.rs
+++ b/crates/pyspark-rs/src/storagelevel.rs
@@ -4,7 +4,12 @@ use pyo3::prelude::*;
/// Flags controlling storage of a persisted DataFrame (disk/memory/off-heap,
/// (de)serialized, replication). Mirrors `pyspark.storagelevel.StorageLevel`.
-#[pyclass(name = "StorageLevel", get_all, from_py_object)]
+#[pyclass(
+ name = "StorageLevel",
+ module = "pyspark.storagelevel",
+ get_all,
+ from_py_object
+)]
#[derive(Clone)]
pub struct PyStorageLevel {
pub useDisk: bool,
diff --git a/crates/pyspark-rs/src/streaming.rs b/crates/pyspark-rs/src/streaming.rs
index 742e033..da40991 100644
--- a/crates/pyspark-rs/src/streaming.rs
+++ b/crates/pyspark-rs/src/streaming.rs
@@ -41,7 +41,7 @@ impl StreamingQueryListener for PyListenerAdapter {
}
/// Python wrapper for DataStreamReader.
-#[pyclass(name = "DataStreamReader")]
+#[pyclass(name = "DataStreamReader", module = "pyspark.sql.streaming.readwriter")]
pub struct PyDataStreamReader {
inner: Option<DataStreamReader>,
}
@@ -180,7 +180,11 @@ impl PyDataStreamReader {
}
/// Python wrapper for Trigger.
-#[pyclass(name = "Trigger", from_py_object)]
+#[pyclass(
+ name = "Trigger",
+ module = "pyspark.sql.streaming.readwriter",
+ from_py_object
+)]
#[derive(Clone)]
pub struct PyTrigger {
inner: Trigger,
@@ -220,7 +224,7 @@ impl PyTrigger {
}
/// Python wrapper for DataStreamWriter.
-#[pyclass(name = "DataStreamWriter")]
+#[pyclass(name = "DataStreamWriter", module = "pyspark.sql.streaming.readwriter")]
pub struct PyDataStreamWriter {
inner: Option<DataStreamWriter>,
}
@@ -402,7 +406,11 @@ impl PyDataStreamWriter {
}
/// Python wrapper for StreamingQueryStatus.
-#[pyclass(name = "StreamingQueryStatus", from_py_object)]
+#[pyclass(
+ name = "StreamingQueryStatus",
+ module = "pyspark.sql.streaming.query",
+ from_py_object
+)]
#[derive(Clone)]
pub struct PyStreamingQueryStatus {
inner: StreamingQueryStatus,
@@ -438,7 +446,11 @@ impl PyStreamingQueryStatus {
}
/// Python wrapper for StreamingQueryException.
-#[pyclass(name = "StreamingQueryException", from_py_object)]
+#[pyclass(
+ name = "StreamingQueryException",
+ module = "pyspark.sql.streaming.query",
+ from_py_object
+)]
#[derive(Clone)]
pub struct PyStreamingQueryException {
inner: StreamingQueryException,
@@ -464,7 +476,7 @@ impl PyStreamingQueryException {
}
/// Python wrapper for StreamingQuery.
-#[pyclass(name = "StreamingQuery")]
+#[pyclass(name = "StreamingQuery", module = "pyspark.sql.streaming.query")]
pub struct PyStreamingQuery {
inner: StreamingQuery,
}
@@ -538,7 +550,7 @@ impl PyStreamingQuery {
}
/// Python wrapper for StreamingQueryManager.
-#[pyclass(name = "StreamingQueryManager")]
+#[pyclass(name = "StreamingQueryManager", module = "pyspark.sql.streaming.query")]
pub struct PyStreamingQueryManager {
inner: StreamingQueryManager,
}
@@ -609,7 +621,7 @@ impl PyStreamingQueryManager {
/// Python wrapper for ListenerEventStream.
/// Implements `__iter__` and `__next__` to yield (event_type: i32, event_json: String) tuples.
-#[pyclass(name = "ListenerEventStream")]
+#[pyclass(name = "ListenerEventStream", module = "pyspark.sql.streaming.query")]
pub struct PyListenerEventStream {
inner: Option<ListenerEventStream>,
}
diff --git a/crates/pyspark-rs/src/tablearg.rs b/crates/pyspark-rs/src/tablearg.rs
index 024b028..0723d3f 100644
--- a/crates/pyspark-rs/src/tablearg.rs
+++ b/crates/pyspark-rs/src/tablearg.rs
@@ -9,7 +9,7 @@ use crate::errors::ResultExt;
/// A table argument for a table-valued function, mirroring
/// `pyspark.sql.connect.table_arg.TableArg`. Built via `df.asTable()`, then optionally
/// `partitionBy`/`orderBy`/`withSinglePartition` (each returns a new TableArg).
-#[pyclass(name = "TableArg")]
+#[pyclass(name = "TableArg", module = "pyspark.sql.table_arg")]
pub struct PyTableArg {
pub(crate) inner: TableArg,
}
diff --git a/crates/pyspark-rs/src/transport.rs b/crates/pyspark-rs/src/transport.rs
index 4210715..2e48d94 100644
--- a/crates/pyspark-rs/src/transport.rs
+++ b/crates/pyspark-rs/src/transport.rs
@@ -56,7 +56,7 @@ pyo3::create_exception!(_pyspark, RustRpcError, pyo3::exceptions::PyException);
/// so the reference client's reattachable iterator can drive it correctly - detecting an
/// early stream end (e.g. an interrupted operation) and reattaching, rather than us
/// eagerly collecting the whole stream and reporting partial results as success.
-#[pyclass]
+#[pyclass(module = "pyspark.sql.connect.client")]
pub struct ResponseStream {
stream: Option<tonic::Streaming<Vec<u8>>>,
// Cancellation shared with the owning stub: when the client is closed (mirroring
@@ -140,7 +140,7 @@ impl ResponseStream {
///
/// Each method mirrors one `SparkConnectServiceStub` RPC. Requests/responses cross
/// the boundary as protobuf bytes so the Python side can keep using its pb2 types.
-#[pyclass(name = "RustConnectStub")]
+#[pyclass(name = "RustConnectStub", module = "pyspark.sql.connect.client")]
pub struct RustConnectStub {
client: SparkConnectClient,
// Shared cancellation, cloned into every ResponseStream this stub hands out. close()
diff --git a/crates/pyspark-rs/src/tvf.rs b/crates/pyspark-rs/src/tvf.rs
index ffa0a58..22cf760 100644
--- a/crates/pyspark-rs/src/tvf.rs
+++ b/crates/pyspark-rs/src/tvf.rs
@@ -8,7 +8,7 @@ use crate::errors::ResultExt;
use crate::functions::to_column;
/// Python wrapper for the table-valued-function namespace (`spark.tvf`).
-#[pyclass(name = "TableValuedFunction")]
+#[pyclass(name = "TableValuedFunction", module = "pyspark.sql.tvf")]
pub struct PyTableValuedFunction {
inner: TableValuedFunction,
}
diff --git a/crates/pyspark-rs/src/types.rs b/crates/pyspark-rs/src/types.rs
index bc61069..ad89a8b 100644
--- a/crates/pyspark-rs/src/types.rs
+++ b/crates/pyspark-rs/src/types.rs
@@ -18,7 +18,7 @@ fn type_reduce(py: Python<'_>, json: &str) -> PyResult<(Py<PyAny>, (String,))> {
/// (mirroring pyspark's `DataType`); concrete type classes extend it (directly or via
/// the intermediate abstract bases below), so `isinstance(dt, DataType)` and the object
/// model (simpleString/typeName/json/...) are inherited from here via `self.inner`.
-#[pyclass(subclass, name = "DataType")]
+#[pyclass(subclass, name = "DataType", module = "pyspark.sql.types")]
pub struct PyDataType {
pub(crate) inner: DataType,
}
@@ -29,7 +29,9 @@ pub struct PyDataType {
/// existing so `isinstance(dt, NumericType)` etc. work with the reference MRO.
macro_rules! abstract_type {
($ty:ident, $name:literal, $parent:ty) => {
- #[pyclass(subclass, name = $name, extends = $parent)]
+ // `module` matches where the drop-in re-exports these (pyspark.sql.types), so
+ // `__module__` reports the reference path instead of PyO3's default `builtins`.
+ #[pyclass(subclass, name = $name, module = "pyspark.sql.types", extends = $parent)]
pub struct $ty;
};
}
@@ -201,7 +203,7 @@ impl PyDataType {
}
// Define concrete type classes
-#[pyclass(name = "NullType", extends = PyDataType)]
+#[pyclass(name = "NullType", module = "pyspark.sql.types", extends = PyDataType)]
pub struct PyNullType;
#[pymethods]
@@ -258,7 +260,7 @@ impl PyNullType {
}
}
-#[pyclass(name = "BooleanType", extends = PyAtomicType)]
+#[pyclass(name = "BooleanType", module = "pyspark.sql.types", extends = PyAtomicType)]
pub struct PyBooleanType;
#[pymethods]
@@ -315,7 +317,7 @@ impl PyBooleanType {
}
}
-#[pyclass(name = "ByteType", extends = PyIntegralType)]
+#[pyclass(name = "ByteType", module = "pyspark.sql.types", extends = PyIntegralType)]
pub struct PyByteType;
#[pymethods]
@@ -376,7 +378,7 @@ impl PyByteType {
}
}
-#[pyclass(name = "ShortType", extends = PyIntegralType)]
+#[pyclass(name = "ShortType", module = "pyspark.sql.types", extends = PyIntegralType)]
pub struct PyShortType;
#[pymethods]
@@ -437,7 +439,7 @@ impl PyShortType {
}
}
-#[pyclass(name = "IntegerType", extends = PyIntegralType)]
+#[pyclass(name = "IntegerType", module = "pyspark.sql.types", extends = PyIntegralType)]
pub struct PyIntegerType;
#[pymethods]
@@ -498,7 +500,7 @@ impl PyIntegerType {
}
}
-#[pyclass(name = "LongType", extends = PyIntegralType)]
+#[pyclass(name = "LongType", module = "pyspark.sql.types", extends = PyIntegralType)]
pub struct PyLongType;
#[pymethods]
@@ -559,7 +561,7 @@ impl PyLongType {
}
}
-#[pyclass(name = "FloatType", extends = PyFractionalType)]
+#[pyclass(name = "FloatType", module = "pyspark.sql.types", extends = PyFractionalType)]
pub struct PyFloatType;
#[pymethods]
@@ -620,7 +622,7 @@ impl PyFloatType {
}
}
-#[pyclass(name = "DoubleType", extends = PyFractionalType)]
+#[pyclass(name = "DoubleType", module = "pyspark.sql.types", extends = PyFractionalType)]
pub struct PyDoubleType;
#[pymethods]
@@ -681,7 +683,7 @@ impl PyDoubleType {
}
}
-#[pyclass(name = "DecimalType", extends = PyFractionalType)]
+#[pyclass(name = "DecimalType", module = "pyspark.sql.types", extends = PyFractionalType)]
pub struct PyDecimalType {
pub precision: i32,
pub scale: i32,
@@ -760,7 +762,7 @@ impl PyDecimalType {
}
}
-#[pyclass(name = "StringType", extends = PyAtomicType)]
+#[pyclass(name = "StringType", module = "pyspark.sql.types", extends = PyAtomicType)]
pub struct PyStringType {
pub collation: String,
}
@@ -868,7 +870,7 @@ impl PyStringType {
}
}
-#[pyclass(name = "BinaryType", extends = PyAtomicType)]
+#[pyclass(name = "BinaryType", module = "pyspark.sql.types", extends = PyAtomicType)]
pub struct PyBinaryType;
#[pymethods]
@@ -925,7 +927,7 @@ impl PyBinaryType {
}
}
-#[pyclass(name = "DateType", extends = PyDatetimeType)]
+#[pyclass(name = "DateType", module = "pyspark.sql.types", extends = PyDatetimeType)]
pub struct PyDateType;
#[pymethods]
@@ -982,7 +984,7 @@ impl PyDateType {
}
}
-#[pyclass(name = "TimestampType", extends = PyDatetimeType)]
+#[pyclass(name = "TimestampType", module = "pyspark.sql.types", extends = PyDatetimeType)]
pub struct PyTimestampType;
#[pymethods]
@@ -1043,7 +1045,7 @@ impl PyTimestampType {
}
}
-#[pyclass(name = "TimestampNTZType", extends = PyDatetimeType)]
+#[pyclass(name = "TimestampNTZType", module = "pyspark.sql.types", extends = PyDatetimeType)]
pub struct PyTimestampNTZType;
#[pymethods]
@@ -1104,7 +1106,7 @@ impl PyTimestampNTZType {
}
}
-#[pyclass(name = "ArrayType", extends = PyDataType)]
+#[pyclass(name = "ArrayType", module = "pyspark.sql.types", extends = PyDataType)]
pub struct PyArrayType {
pub element_type: DataType,
pub contains_null: bool,
@@ -1231,7 +1233,7 @@ impl PyArrayType {
}
}
-#[pyclass(name = "MapType", extends = PyDataType)]
+#[pyclass(name = "MapType", module = "pyspark.sql.types", extends = PyDataType)]
pub struct PyMapType {
pub key_type: DataType,
pub value_type: DataType,
@@ -1385,7 +1387,7 @@ impl PyMapType {
}
}
-#[pyclass(name = "StructField")]
+#[pyclass(name = "StructField", module = "pyspark.sql.types")]
pub struct PyStructField {
pub(crate) field: StructField,
}
@@ -1583,7 +1585,7 @@ impl PyStructField {
}
}
-#[pyclass(name = "StructType", extends = PyDataType)]
+#[pyclass(name = "StructType", module = "pyspark.sql.types", extends = PyDataType)]
pub struct PyStructType {
pub(crate) fields: Vec<StructField>,
}
@@ -1872,7 +1874,7 @@ impl PyStructType {
}
}
-#[pyclass(name = "CharType", extends = PyAtomicType)]
+#[pyclass(name = "CharType", module = "pyspark.sql.types", extends = PyAtomicType)]
pub struct PyCharType {
pub length: i32,
}
@@ -1939,7 +1941,7 @@ impl PyCharType {
}
}
-#[pyclass(name = "VarcharType", extends = PyAtomicType)]
+#[pyclass(name = "VarcharType", module = "pyspark.sql.types", extends = PyAtomicType)]
pub struct PyVarcharType {
pub length: i32,
}
@@ -2006,7 +2008,7 @@ impl PyVarcharType {
}
}
-#[pyclass(name = "TimeType", extends = PyAnyTimeType)]
+#[pyclass(name = "TimeType", module = "pyspark.sql.types", extends = PyAnyTimeType)]
pub struct PyTimeType {
pub precision: i32,
}
@@ -2075,7 +2077,7 @@ impl PyTimeType {
}
}
-#[pyclass(name = "CalendarIntervalType", extends = PyDataType)]
+#[pyclass(name = "CalendarIntervalType", module = "pyspark.sql.types", extends = PyDataType)]
pub struct PyCalendarIntervalType;
#[pymethods]
@@ -2132,7 +2134,7 @@ impl PyCalendarIntervalType {
}
}
-#[pyclass(name = "YearMonthIntervalType", extends = PyAnsiIntervalType)]
+#[pyclass(name = "YearMonthIntervalType", module = "pyspark.sql.types", extends = PyAnsiIntervalType)]
pub struct PyYearMonthIntervalType {
pub start_field: i32,
pub end_field: i32,
@@ -2211,7 +2213,7 @@ impl PyYearMonthIntervalType {
}
}
-#[pyclass(name = "DayTimeIntervalType", extends = PyAnsiIntervalType)]
+#[pyclass(name = "DayTimeIntervalType", module = "pyspark.sql.types", extends = PyAnsiIntervalType)]
pub struct PyDayTimeIntervalType {
pub start_field: i32,
pub end_field: i32,
@@ -2290,7 +2292,7 @@ impl PyDayTimeIntervalType {
}
}
-#[pyclass(name = "VariantType", extends = PyAtomicType)]
+#[pyclass(name = "VariantType", module = "pyspark.sql.types", extends = PyAtomicType)]
pub struct PyVariantType;
#[pymethods]
@@ -2347,7 +2349,7 @@ impl PyVariantType {
}
}
-#[pyclass(name = "GeometryType", extends = PySpatialType)]
+#[pyclass(name = "GeometryType", module = "pyspark.sql.types", extends = PySpatialType)]
pub struct PyGeometryType {
pub srid: i32,
}
@@ -2381,7 +2383,7 @@ impl PyGeometryType {
}
}
-#[pyclass(name = "GeographyType", extends = PySpatialType)]
+#[pyclass(name = "GeographyType", module = "pyspark.sql.types", extends = PySpatialType)]
pub struct PyGeographyType {
pub srid: i32,
}
diff --git a/crates/pyspark-rs/src/udtf_analyze.rs b/crates/pyspark-rs/src/udtf_analyze.rs
index b5937d7..ac0863d 100644
--- a/crates/pyspark-rs/src/udtf_analyze.rs
+++ b/crates/pyspark-rs/src/udtf_analyze.rs
@@ -6,14 +6,16 @@
use pyo3::prelude::*;
+// The module argument sets `__module__`; use the reference path (upstream defines this
+// in pyspark/sql/udtf.py) rather than the extension's own name.
pyo3::create_exception!(
- _pyspark,
+ pyspark.sql.udtf,
SkipRestOfInputTableException,
pyo3::exceptions::PyException
);
/// `pyspark.sql.udtf.AnalyzeArgument`.
-#[pyclass(name = "AnalyzeArgument", get_all)]
+#[pyclass(name = "AnalyzeArgument", module = "pyspark.sql.udtf", get_all)]
pub struct PyAnalyzeArgument {
pub dataType: Py<PyAny>,
pub value: Py<PyAny>,
@@ -40,7 +42,12 @@ impl PyAnalyzeArgument {
}
/// `pyspark.sql.udtf.PartitioningColumn` (frozen).
-#[pyclass(name = "PartitioningColumn", get_all, frozen)]
+#[pyclass(
+ name = "PartitioningColumn",
+ module = "pyspark.sql.udtf",
+ get_all,
+ frozen
+)]
pub struct PyPartitioningColumn {
pub name: String,
}
@@ -53,7 +60,7 @@ impl PyPartitioningColumn {
}
/// `pyspark.sql.udtf.OrderingColumn` (frozen).
-#[pyclass(name = "OrderingColumn", get_all, frozen)]
+#[pyclass(name = "OrderingColumn", module = "pyspark.sql.udtf", get_all, frozen)]
pub struct PyOrderingColumn {
pub name: String,
pub ascending: bool,
@@ -74,7 +81,7 @@ impl PyOrderingColumn {
}
/// `pyspark.sql.udtf.SelectedColumn` (frozen).
-#[pyclass(name = "SelectedColumn", get_all, frozen)]
+#[pyclass(name = "SelectedColumn", module = "pyspark.sql.udtf", get_all, frozen)]
pub struct PySelectedColumn {
pub name: String,
pub alias: String,
@@ -89,7 +96,7 @@ impl PySelectedColumn {
}
/// `pyspark.sql.udtf.AnalyzeResult` (mutable dataclass).
-#[pyclass(name = "AnalyzeResult", get_all, set_all)]
+#[pyclass(name = "AnalyzeResult", module = "pyspark.sql.udtf", get_all, set_all)]
pub struct PyAnalyzeResult {
pub schema: Py<PyAny>,
pub withSinglePartition: bool,
diff --git a/crates/pyspark-rs/src/values.rs b/crates/pyspark-rs/src/values.rs
index 6a3d4ee..0cea8f8 100644
--- a/crates/pyspark-rs/src/values.rs
+++ b/crates/pyspark-rs/src/values.rs
@@ -6,7 +6,7 @@ use pyo3::types::PyBytes;
/// A Variant value (the `value`/`metadata` binary components). Mirrors
/// `pyspark.sql.types.VariantVal`.
-#[pyclass(name = "VariantVal")]
+#[pyclass(name = "VariantVal", module = "pyspark.sql.types")]
pub struct PyVariantVal {
value: Vec<u8>,
metadata: Vec<u8>,
@@ -70,7 +70,7 @@ impl PyVariantVal {
}
/// A Geography value (WKB bytes + SRID). Mirrors `pyspark.sql.types.Geography`.
-#[pyclass(name = "Geography")]
+#[pyclass(name = "Geography", module = "pyspark.sql.types")]
pub struct PyGeography {
wkb: Vec<u8>,
srid: i32,
@@ -104,7 +104,7 @@ impl PyGeography {
}
/// A Geometry value (WKB bytes + SRID). Mirrors `pyspark.sql.types.Geometry`.
-#[pyclass(name = "Geometry")]
+#[pyclass(name = "Geometry", module = "pyspark.sql.types")]
pub struct PyGeometry {
wkb: Vec<u8>,
srid: i32,
diff --git a/crates/pyspark-rs/src/window.rs b/crates/pyspark-rs/src/window.rs
index c81aef3..8fd508d 100644
--- a/crates/pyspark-rs/src/window.rs
+++ b/crates/pyspark-rs/src/window.rs
@@ -34,7 +34,7 @@ fn to_sort_orders(cols: Vec<Bound<'_, PyAny>>) -> PyResult<Vec<SortOrder>> {
}
/// Python wrapper for a window frame bound.
-#[pyclass(name = "FrameBound", from_py_object)]
+#[pyclass(name = "FrameBound", module = "pyspark.sql.window", from_py_object)]
#[derive(Clone)]
pub struct PyFrameBound {
pub(crate) bound: FrameBound,
@@ -82,7 +82,7 @@ impl PyFrameBound {
}
/// Python wrapper for a WindowSpec.
-#[pyclass(name = "WindowSpec", from_py_object)]
+#[pyclass(name = "WindowSpec", module = "pyspark.sql.window", from_py_object)]
#[derive(Clone)]
pub struct PyWindowSpec {
pub(crate) spec: WindowSpec,
@@ -142,7 +142,7 @@ impl PyWindowSpec {
}
/// Python wrapper for the Window static builder.
-#[pyclass(name = "Window")]
+#[pyclass(name = "Window", module = "pyspark.sql.window")]
pub struct PyWindow;
#[pymethods]
diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py
index 14e236c..07e2e66 100644
--- a/python/tests/test_dropin_offline.py
+++ b/python/tests/test_dropin_offline.py
@@ -2572,3 +2572,54 @@ def test_pandas_udf_eval_type_inferred_from_type_hints():
def agg(s: pd.Series) -> float:
return float(s.mean())
assert agg.evalType == 202, agg.evalType
+
+
+# --------------------------------------------------------- __module__ / class pickling
+# PyO3 defaults a #[pyclass]'s __module__ to "builtins", which diverges from the
+# reference client (where each class reports the module that defines it) and makes the
+# class itself unpicklable: pickle serializes a class BY REFERENCE, writing
+# __module__ + __qualname__ and re-importing on load, so "builtins.IntegerType" fails
+# the lookup. Every #[pyclass] therefore carries an explicit `module = "..."`.
+# NB: DataFrame/Column report the *connect* module: upstream pyspark.sql.{dataframe,
+# column} hold the abstract parent, while the concrete class a Connect session hands
+# out is pyspark.sql.connect.*.
+
+
+@pytest.mark.parametrize(
+ "cls,expected",
+ [
+ (_DF, "pyspark.sql.connect.dataframe"),
+ (_Col, "pyspark.sql.connect.column"),
+ (_SS, "pyspark.sql.session"),
+ (T.DataType, "pyspark.sql.types"),
+ (T.IntegerType, "pyspark.sql.types"),
+ (T.StructType, "pyspark.sql.types"),
+ (T.NumericType, "pyspark.sql.types"), # macro-generated abstract base
+ (T.Row, "pyspark.sql.types"),
+ ],
+)
+def test_class_module_matches_reference_path(cls, expected):
+ assert cls.__module__ == expected, f"{cls.__name__}.__module__ = {cls.__module__}"
+
+
+@pytest.mark.parametrize(
+ "cls",
+ [T.DataType, T.IntegerType, T.StringType, T.ArrayType, T.MapType, T.StructType,
+ T.DecimalType, T.NumericType, T.AtomicType],
+)
+def test_datatype_class_itself_is_picklable(cls):
+ # Pickling the CLASS (not an instance) round-trips to the identical object. A
+ # __module__ of "builtins" raises PicklingError here.
+ assert pickle.loads(pickle.dumps(cls)) is cls
+
+
+def test_no_pyclass_reports_builtins_module():
+ # Guard against a new #[pyclass] landing without `module = "..."`.
+ import pyspark._pyspark as _rust
+
+ offenders = sorted(
+ name
+ for name, obj in vars(_rust).items()
+ if isinstance(obj, type) and obj.__module__ == "builtins"
+ )
+ assert offenders == [], f"missing module=: {offenders}"
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org