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_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