Support the label loop over Spark Connect (RelationPlugin + shared-session fix)#1350
Open
Chandanydv1234 wants to merge 4 commits into
Open
Conversation
The command plugin reused the CLI's init/execute/stop pattern, but over Spark Connect the SparkSession is owned by the long-running server and shared across requests. client.stop() -> session.stop() tore it down after the first phase, so later phases failed with "Cannot call methods on a stopped SparkContext". Do not stop the session in the plugin. Fixes zinggAI#1349
A CommandPlugin can only signal success/failure and cannot return rows, so the label/findAndLabel phases could not run over Spark Connect. Add ZinggRelationPlugin, which returns Labeller.getUnmarkedRecords() as a LogicalPlan so Spark streams the unmarked pairs back to the client, and Zingg.getUnmarkedPairs() to fetch them as a PyArrow table. Registered via spark.connect.extensions.relation.classes.
Completes the label loop: after fetching pairs with getUnmarkedPairs() and marking them on the client, the labels need to reach the server so train() can use them. Add Zingg.writeMarkedPairs() (backed by write_marked_pairs()), which reads the unmarked pairs server-side to keep their exact schema, joins the client's labels by cluster id, overwrites z_isMatch, and appends to the model's marked training data -- plain Spark Connect DataFrame ops, no extra server plugin. Paths are derived from Arguments, so nothing is hardcoded. examples/label_over_connect.py runs the whole loop (fetch -> mark -> write back) with all paths as arguments.
findAndLabel = findTrainingData + label. The relation plugin previously only read existing unmarked pairs. For findAndLabel it now runs the finder first (SparkFindAndLabeller.getFinder().execute()) to generate fresh pairs, then returns them via the labeller -- so it works over Connect in one call. Plain label is unchanged.
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
Enables the interactive
labelandfindAndLabelworkflows over Spark Connect, and fixes a bug that limited the server to one phase per start.Why
A
CommandPlugincan only report success/failure — it can't return rows to the client. But labelling needs the server to hand record pairs back to the user to mark, solabel/findAndLabelcouldn't run over Spark Connect at all.Changes
Server
ZinggRelationPlugin(new): returns the unmarked training pairs as aLogicalPlan, so Spark streams the rows back to the client — the two-way channel aCommandPluginlacks. Handles bothlabel(reads existing pairs) andfindAndLabel(runs the finder first, then returns fresh pairs).ZinggCommandPlugin: no longer callsclient.stop(), which was stopping the server's sharedSparkSessionafter the first phase — every later phase then failed with "Cannot call methods on a stopped SparkContext". (Fixes Spark Connect: server only handles one phase per start (ZinggCommandPlugin stops the shared session) #1349)Client
Zingg.getUnmarkedPairs(): fetches the pairs as a PyArrow table.Zingg.writeMarkedPairs(): writes the labels back to the model's marked training data (paths derived fromArguments), sotrainpicks them up.examples/label_over_connect.py: runs the fullfetch → mark → write-backloop.Registration
The server also needs
protobuf-javaon the classpath (the plugins use the unshaded protobuf types).Testing
Ran end to end over Spark Connect on Spark 3.5.1:
findTrainingData → label → train → matchfindAndLabel → train → matchIn both, the labels are written back,
trainpicks them up, andmatchproduces clustered output.