From bc60a5fb43e04dd4d22c1d42bf60813f5d31ee9b Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 22:20:33 -0700 Subject: [PATCH 1/3] [SPARK-60681][CONNECT] Add DataFrame inspection completions --- crates/pyspark-rs/src/dataframe.rs | 15 ++ python/tests/test_dataframe_completion.py | 192 ++++++++++++++++++++++ 2 files changed, 207 insertions(+) create mode 100644 python/tests/test_dataframe_completion.py diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 7101aadf..33a8c774 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1129,6 +1129,21 @@ impl PyDataFrame { Ok(PyColumn::new(spark_connect::functions::col(name))) } + /// Class attributes and all column names, matching Python Connect. + fn __dir__(&self, py: Python<'_>) -> PyResult> { + // Inspect the class rather than the instance to avoid calling __dir__ again. + let mut attrs = py.get_type::().dir()?.extract::>()?; + attrs.extend(py.detach(|| self.dataframe.columns()).to_pyerr()?); + attrs.sort(); + attrs.dedup(); + Ok(attrs) + } + + /// Bracket completion keeps column order, duplicates and non-identifiers. + fn _ipython_key_completions_(&self, py: Python<'_>) -> PyResult> { + py.detach(|| self.dataframe.columns()).to_pyerr() + } + /// Aggregate over the whole DataFrame (shorthand for `groupBy().agg(...)`). /// /// Mirrors `pyspark.sql.DataFrame.agg(*exprs)`. diff --git a/python/tests/test_dataframe_completion.py b/python/tests/test_dataframe_completion.py new file mode 100644 index 00000000..06c3d8e3 --- /dev/null +++ b/python/tests/test_dataframe_completion.py @@ -0,0 +1,192 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DataFrame completion through the real client transport, without a Spark server. + +The fixture serves only AnalyzePlan. Tiny wire encoders keep the tests independent +of the reference PySpark client and its generated protobuf modules. +""" +from concurrent.futures import ThreadPoolExecutor +import multiprocessing + +import pytest + +grpc = pytest.importorskip("grpc") + +from pyspark.sql import Column, SparkSession + + +def _varint(value): + result = bytearray() + while value > 127: + result.append((value & 127) | 128) + value >>= 7 + result.append(value) + return bytes(result) + + +def _message(field, value): + if isinstance(value, str): + value = value.encode() + return _varint(field * 8 + 2) + _varint(len(value)) + value + + +def _schema(names=("id",)): + # DataType.struct=21, Struct.fields=1, DataType.long=7. + return _message(21, b"".join(_message(1, _field(name, _message(7, b""))) + for name in names)) + + +def _field(name, data_type): + return (_message(1, name) + _message(2, data_type) + b"\x18\x01" + + _message(4, '{"note":"original"}')) + + +def _serve(connection): + # A separate process also exercises columns/dtypes, whose first access may + # hold the client's GIL while waiting for the RPC. + class Service: + calls = 0 + data_type = _schema() + failure = None + + def analyze(self, request, context): + self.calls += 1 + if self.failure == "rpc": + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "schema unavailable") + if self.failure == "no-result": + return b"" + if self.failure == "missing-schema": + return _message(2, b"") + if self.failure == "invalid-type": + return _message(2, _message(1, b"")) + # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1. + return _message(2, _message(1, self.data_type)) + + service = Service() + with ThreadPoolExecutor(max_workers=4) as executor: + server = grpc.server(executor) + handler = grpc.unary_unary_rpc_method_handler(service.analyze) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", {"AnalyzePlan": handler}),)) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + connection.send(port) + try: + while True: + command = connection.recv() + if command[0] == "stop": + break + if command[0] == "get": + connection.send(getattr(service, command[1])) + elif command[0] == "set": + setattr(service, command[1], command[2]) + connection.send(None) + finally: + server.stop(0).wait() + + +class _RemoteService: + def __init__(self, connection): + object.__setattr__(self, "connection", connection) + + def __getattr__(self, name): + self.connection.send(("get", name)) + assert self.connection.poll(10), "Mock server did not respond" + return self.connection.recv() + + def __setattr__(self, name, value): + self.connection.send(("set", name, value)) + assert self.connection.poll(10), "Mock server did not respond" + self.connection.recv() + + +@pytest.fixture +def analyzed(): + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_serve, args=(child,)) + process.start() + try: + assert parent.poll(20), "Mock server did not start" + port = parent.recv() + spark = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() + yield spark, _RemoteService(parent) + finally: + if process.is_alive(): + parent.send(("stop",)) + process.join(10) + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + child.close() + + +@pytest.mark.parametrize("names", [ + ["z", "a"], + ["with space", "a.b", "with`tick", "", "123", "class", "雪", "_private"], + ["id", "id", "select", "columns", "__dir__", "_ipython_key_completions_"], + [], +]) +def test_completion_matches_python_connect(analyzed, names): + spark, service = analyzed + service.data_type = _schema(names) + df = spark.range(10) + assert df._ipython_key_completions_() == names + assert service.calls == 1 + expected = sorted(set(dir(type(df))) | set(names)) + assert df.__dir__() == expected + assert dir(df) == expected + assert service.calls == 3 # No schema caching in this independent change. + assert callable(df.select) + assert callable(df.__dir__) + assert callable(df._ipython_key_completions_) + assert df.columns == names + for name in names: + assert isinstance(df[name], Column) + assert df.schema.fieldNames() == names + + +def test_completion_lists_are_independent(analyzed): + spark, service = analyzed + df = spark.range(10) + df._ipython_key_completions_().append("extra") + df.__dir__().clear() + assert df._ipython_key_completions_() == ["id"] + assert "select" in dir(df) + service.data_type = _schema(["changed"]) + assert df._ipython_key_completions_() == ["changed"] + + +@pytest.mark.parametrize("method", ["__dir__", "_ipython_key_completions_"]) +@pytest.mark.parametrize("failure", ["rpc", "no-result", "missing-schema", "invalid-type"]) +def test_completion_propagates_schema_errors_and_can_retry(analyzed, method, failure): + spark, service = analyzed + df = spark.range(10) + service.failure = failure + message = { + "rpc": "schema unavailable", + "no-result": "no schema in response", + "missing-schema": "Schema is missing", + "invalid-type": "DataType kind not set", + }[failure] + for count in (1, 2): + with pytest.raises(Exception, match=message): + getattr(df, method)() + assert service.calls == count + service.failure = None + assert "id" in getattr(df, method)() + assert service.calls == 3 From 4019957a9124c80e78bb3210a1fd1c48305cdec4 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 00:38:07 -0700 Subject: [PATCH 2/3] [SPARK-60681][CONNECT] Decouple completion tests from schema caching --- python/tests/test_dataframe_completion.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/tests/test_dataframe_completion.py b/python/tests/test_dataframe_completion.py index 06c3d8e3..db6f003a 100644 --- a/python/tests/test_dataframe_completion.py +++ b/python/tests/test_dataframe_completion.py @@ -148,9 +148,11 @@ def test_completion_matches_python_connect(analyzed, names): assert df._ipython_key_completions_() == names assert service.calls == 1 expected = sorted(set(dir(type(df))) | set(names)) - assert df.__dir__() == expected - assert dir(df) == expected - assert service.calls == 3 # No schema caching in this independent change. + # Each first access must analyze its own fresh plan, with or without caching. + assert spark.range(10).__dir__() == expected + assert service.calls == 2 + assert dir(spark.range(10)) == expected + assert service.calls == 3 assert callable(df.select) assert callable(df.__dir__) assert callable(df._ipython_key_completions_) @@ -168,7 +170,7 @@ def test_completion_lists_are_independent(analyzed): assert df._ipython_key_completions_() == ["id"] assert "select" in dir(df) service.data_type = _schema(["changed"]) - assert df._ipython_key_completions_() == ["changed"] + assert spark.range(10)._ipython_key_completions_() == ["changed"] @pytest.mark.parametrize("method", ["__dir__", "_ipython_key_completions_"]) From 61d3ea4886939950a62550e23b7e377dda5bf6b5 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 00:42:17 -0700 Subject: [PATCH 3/3] [SPARK-60681][CONNECT] Run completion tests in the offline CI suite --- python/tests/test_dataframe_completion.py | 194 ---------------------- python/tests/test_dropin_offline.py | 177 ++++++++++++++++++++ 2 files changed, 177 insertions(+), 194 deletions(-) delete mode 100644 python/tests/test_dataframe_completion.py diff --git a/python/tests/test_dataframe_completion.py b/python/tests/test_dataframe_completion.py deleted file mode 100644 index db6f003a..00000000 --- a/python/tests/test_dataframe_completion.py +++ /dev/null @@ -1,194 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""DataFrame completion through the real client transport, without a Spark server. - -The fixture serves only AnalyzePlan. Tiny wire encoders keep the tests independent -of the reference PySpark client and its generated protobuf modules. -""" -from concurrent.futures import ThreadPoolExecutor -import multiprocessing - -import pytest - -grpc = pytest.importorskip("grpc") - -from pyspark.sql import Column, SparkSession - - -def _varint(value): - result = bytearray() - while value > 127: - result.append((value & 127) | 128) - value >>= 7 - result.append(value) - return bytes(result) - - -def _message(field, value): - if isinstance(value, str): - value = value.encode() - return _varint(field * 8 + 2) + _varint(len(value)) + value - - -def _schema(names=("id",)): - # DataType.struct=21, Struct.fields=1, DataType.long=7. - return _message(21, b"".join(_message(1, _field(name, _message(7, b""))) - for name in names)) - - -def _field(name, data_type): - return (_message(1, name) + _message(2, data_type) + b"\x18\x01" - + _message(4, '{"note":"original"}')) - - -def _serve(connection): - # A separate process also exercises columns/dtypes, whose first access may - # hold the client's GIL while waiting for the RPC. - class Service: - calls = 0 - data_type = _schema() - failure = None - - def analyze(self, request, context): - self.calls += 1 - if self.failure == "rpc": - context.abort(grpc.StatusCode.INVALID_ARGUMENT, "schema unavailable") - if self.failure == "no-result": - return b"" - if self.failure == "missing-schema": - return _message(2, b"") - if self.failure == "invalid-type": - return _message(2, _message(1, b"")) - # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1. - return _message(2, _message(1, self.data_type)) - - service = Service() - with ThreadPoolExecutor(max_workers=4) as executor: - server = grpc.server(executor) - handler = grpc.unary_unary_rpc_method_handler(service.analyze) - server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( - "spark.connect.SparkConnectService", {"AnalyzePlan": handler}),)) - port = server.add_insecure_port("127.0.0.1:0") - server.start() - connection.send(port) - try: - while True: - command = connection.recv() - if command[0] == "stop": - break - if command[0] == "get": - connection.send(getattr(service, command[1])) - elif command[0] == "set": - setattr(service, command[1], command[2]) - connection.send(None) - finally: - server.stop(0).wait() - - -class _RemoteService: - def __init__(self, connection): - object.__setattr__(self, "connection", connection) - - def __getattr__(self, name): - self.connection.send(("get", name)) - assert self.connection.poll(10), "Mock server did not respond" - return self.connection.recv() - - def __setattr__(self, name, value): - self.connection.send(("set", name, value)) - assert self.connection.poll(10), "Mock server did not respond" - self.connection.recv() - - -@pytest.fixture -def analyzed(): - context = multiprocessing.get_context("spawn") - parent, child = context.Pipe() - process = context.Process(target=_serve, args=(child,)) - process.start() - try: - assert parent.poll(20), "Mock server did not start" - port = parent.recv() - spark = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() - yield spark, _RemoteService(parent) - finally: - if process.is_alive(): - parent.send(("stop",)) - process.join(10) - if process.is_alive(): - process.terminate() - process.join(10) - parent.close() - child.close() - - -@pytest.mark.parametrize("names", [ - ["z", "a"], - ["with space", "a.b", "with`tick", "", "123", "class", "雪", "_private"], - ["id", "id", "select", "columns", "__dir__", "_ipython_key_completions_"], - [], -]) -def test_completion_matches_python_connect(analyzed, names): - spark, service = analyzed - service.data_type = _schema(names) - df = spark.range(10) - assert df._ipython_key_completions_() == names - assert service.calls == 1 - expected = sorted(set(dir(type(df))) | set(names)) - # Each first access must analyze its own fresh plan, with or without caching. - assert spark.range(10).__dir__() == expected - assert service.calls == 2 - assert dir(spark.range(10)) == expected - assert service.calls == 3 - assert callable(df.select) - assert callable(df.__dir__) - assert callable(df._ipython_key_completions_) - assert df.columns == names - for name in names: - assert isinstance(df[name], Column) - assert df.schema.fieldNames() == names - - -def test_completion_lists_are_independent(analyzed): - spark, service = analyzed - df = spark.range(10) - df._ipython_key_completions_().append("extra") - df.__dir__().clear() - assert df._ipython_key_completions_() == ["id"] - assert "select" in dir(df) - service.data_type = _schema(["changed"]) - assert spark.range(10)._ipython_key_completions_() == ["changed"] - - -@pytest.mark.parametrize("method", ["__dir__", "_ipython_key_completions_"]) -@pytest.mark.parametrize("failure", ["rpc", "no-result", "missing-schema", "invalid-type"]) -def test_completion_propagates_schema_errors_and_can_retry(analyzed, method, failure): - spark, service = analyzed - df = spark.range(10) - service.failure = failure - message = { - "rpc": "schema unavailable", - "no-result": "no schema in response", - "missing-schema": "Schema is missing", - "invalid-type": "DataType kind not set", - }[failure] - for count in (1, 2): - with pytest.raises(Exception, match=message): - getattr(df, method)() - assert service.calls == count - service.failure = None - assert "id" in getattr(df, method)() - assert service.calls == 3 diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 67ccaa4d..7311de74 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -2755,3 +2755,180 @@ def test_create_dataframe_width_mismatch_is_a_normal_exception(): else: raise AssertionError("expected a width validation error") """], check=True) + + +# ------------------------------------------------------- DataFrame completion + +def _completion_varint(value): + result = bytearray() + while value > 127: + result.append((value & 127) | 128) + value >>= 7 + result.append(value) + return bytes(result) + + +def _completion_message(field, value): + if isinstance(value, str): + value = value.encode() + return _completion_varint(field * 8 + 2) + _completion_varint(len(value)) + value + + +def _completion_schema(names=("id",)): + # DataType.struct=21, Struct.fields=1, DataType.long=7. + return _completion_message(21, b"".join(_completion_message(1, _completion_field(name, _completion_message(7, b""))) + for name in names)) + + +def _completion_field(name, data_type): + return (_completion_message(1, name) + _completion_message(2, data_type) + b"\x18\x01" + + _completion_message(4, '{"note":"original"}')) + + +def _completion_serve(connection): + from concurrent.futures import ThreadPoolExecutor + import grpc + + # A separate process also exercises columns/dtypes, whose first access may + # hold the client's GIL while waiting for the RPC. + class Service: + calls = 0 + data_type = _completion_schema() + failure = None + + def analyze(self, request, context): + self.calls += 1 + if self.failure == "rpc": + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "schema unavailable") + if self.failure == "no-result": + return b"" + if self.failure == "missing-schema": + return _completion_message(2, b"") + if self.failure == "invalid-type": + return _completion_message(2, _completion_message(1, b"")) + # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1. + return _completion_message(2, _completion_message(1, self.data_type)) + + service = Service() + with ThreadPoolExecutor(max_workers=4) as executor: + server = grpc.server(executor) + handler = grpc.unary_unary_rpc_method_handler(service.analyze) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", {"AnalyzePlan": handler}),)) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + connection.send(port) + try: + while True: + command = connection.recv() + if command[0] == "stop": + break + if command[0] == "get": + connection.send(getattr(service, command[1])) + elif command[0] == "set": + setattr(service, command[1], command[2]) + connection.send(None) + finally: + server.stop(0).wait() + + +class _CompletionRemoteService: + def __init__(self, connection): + object.__setattr__(self, "connection", connection) + + def __getattr__(self, name): + self.connection.send(("get", name)) + assert self.connection.poll(10), "Mock server did not respond" + return self.connection.recv() + + def __setattr__(self, name, value): + self.connection.send(("set", name, value)) + assert self.connection.poll(10), "Mock server did not respond" + self.connection.recv() + + +@pytest.fixture +def completion_analyzed(): + import multiprocessing + from pyspark.sql import SparkSession + + pytest.importorskip("grpc") + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_completion_serve, args=(child,)) + process.start() + try: + assert parent.poll(20), "Mock server did not start" + port = parent.recv() + spark = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() + yield spark, _CompletionRemoteService(parent) + finally: + if process.is_alive(): + parent.send(("stop",)) + process.join(10) + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + child.close() + + +@pytest.mark.parametrize("names", [ + ["z", "a"], + ["with space", "a.b", "with`tick", "", "123", "class", "雪", "_private"], + ["id", "id", "select", "columns", "__dir__", "_ipython_key_completions_"], + [], +]) +def test_completion_matches_python_connect(completion_analyzed, names): + from pyspark.sql import Column + + spark, service = completion_analyzed + service.data_type = _completion_schema(names) + df = spark.range(10) + assert df._ipython_key_completions_() == names + assert service.calls == 1 + expected = sorted(set(dir(type(df))) | set(names)) + # Each first access must analyze its own fresh plan, with or without caching. + assert spark.range(10).__dir__() == expected + assert service.calls == 2 + assert dir(spark.range(10)) == expected + assert service.calls == 3 + assert callable(df.select) + assert callable(df.__dir__) + assert callable(df._ipython_key_completions_) + assert df.columns == names + for name in names: + assert isinstance(df[name], Column) + assert df.schema.fieldNames() == names + + +def test_completion_lists_are_independent(completion_analyzed): + spark, service = completion_analyzed + df = spark.range(10) + df._ipython_key_completions_().append("extra") + df.__dir__().clear() + assert df._ipython_key_completions_() == ["id"] + assert "select" in dir(df) + service.data_type = _completion_schema(["changed"]) + assert spark.range(10)._ipython_key_completions_() == ["changed"] + + +@pytest.mark.parametrize("method", ["__dir__", "_ipython_key_completions_"]) +@pytest.mark.parametrize("failure", ["rpc", "no-result", "missing-schema", "invalid-type"]) +def test_completion_propagates_schema_errors_and_can_retry(completion_analyzed, method, failure): + spark, service = completion_analyzed + df = spark.range(10) + service.failure = failure + message = { + "rpc": "schema unavailable", + "no-result": "no schema in response", + "missing-schema": "Schema is missing", + "invalid-type": "DataType kind not set", + }[failure] + for count in (1, 2): + with pytest.raises(Exception, match=message): + getattr(df, method)() + assert service.calls == count + service.failure = None + assert "id" in getattr(df, method)() + assert service.calls == 3