refactor: port DaftContext to rust side #3767

universalmind303 · 2025-02-04T20:54:26Z

ports the DaftContext to the rust side so we can more easily reuse it for connect & catalog work.

notes for reviewers:

after removing daft-local-execution from daft-connect, it made cargo check fail, so some of the changes to unrelated Cargo.toml's was to properly feature flag them. We missed a couple that were causing things to break.
@rchowell I requested review from you specifically to take a look at context.rs. This contains the "context"/"session", and the global state that was previously only available in python. I added the catalog to this state as well.

codspeed-hq · 2025-02-04T22:49:47Z

CodSpeed Performance Report

Merging #3767 will improve performances by 54.97%

_{Comparing universalmind303:rs-context (144ad60) with main (86adc44)}

Summary

⚡ 1 improvements
✅ 26 untouched benchmarks

Benchmarks breakdown

	Benchmark	`BASE`	`HEAD`	Change
⚡	`test_show[100 Small Files]`	24.2 ms	15.6 ms	+54.97%

universalmind303 · 2025-02-05T20:07:43Z

src/daft-context/src/lib.rs

+static DAFT_CONTEXT: OnceCell<DaftContext> = OnceCell::new();
+
+#[cfg(feature = "python")]
+pub fn get_context() -> DaftContext {


Currently we only have a singleton context as this is how it was done in py, but we may want to think about of ways that a user could initialize multiple contexts/sessions.

The global singleton context is pretty useful, I wonder if there are other ways we can produce the required functionality? In my current mental model, all user interactions with the system go through a session/connection whether that is implicit or explicit. Whereas the Context I think of a global program state for the process itself. I'm curious what functionality you had in mind where multiple contexts would be necessary which cannot be achieved by pulling that necessary state into the session?

I know that you have said that currently "context" and "session" are synonymous, but perhaps this is the reason to no longer make them so?

context – process state

session – connection state

We are pushing up against the same concepts and patterns described in SQL standard Part 1 where our context represents state for the SQL-environment (singleton), whereas our session should represent state for the SQL-connection. This is not meant to be definitive, but hopefully this helps us tease out the role of context vs the role of session.

hmm yeah I was actually thinking a bit about this yesterday as well. I think that seems like a reasonable distinction to make.

Since we're going to make a hard distinction between the context and the session, I'll update the PR remove the catalog from the context.

codecov · 2025-02-06T17:03:15Z

Codecov Report

Attention: Patch coverage is 84.18972% with 80 lines in your changes missing coverage. Please review.

Project coverage is 77.82%. Comparing base (62d6b97) to head (144ad60).
Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
src/daft-context/src/lib.rs	80.00%	40 Missing ⚠️
src/daft-context/src/python.rs	76.36%	26 Missing ⚠️
src/daft-py-runners/src/lib.rs	90.98%	11 Missing ⚠️
daft/context.py	91.66%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3767      +/-   ##
==========================================
+ Coverage   75.88%   77.82%   +1.93%     
==========================================
  Files         733      737       +4     
  Lines       96091    93361    -2730     
==========================================
- Hits        72919    72656     -263     
+ Misses      23172    20705    -2467

Files with missing lines	Coverage Δ
src/daft-connect/src/execute.rs	`89.33% <100.00%> (+3.42%)`	⬆️
src/daft-connect/src/lib.rs	`94.44% <ø> (ø)`
src/daft-connect/src/session.rs	`100.00% <100.00%> (ø)`
src/daft-connect/src/spark_analyzer.rs	`73.97% <100.00%> (+1.33%)`	⬆️
src/lib.rs	`96.15% <100.00%> (+0.04%)`	⬆️
daft/context.py	`83.33% <91.66%> (-4.33%)`	⬇️
src/daft-py-runners/src/lib.rs	`90.98% <90.98%> (ø)`
src/daft-context/src/python.rs	`76.36% <76.36%> (ø)`
src/daft-context/src/lib.rs	`80.00% <80.00%> (ø)`

... and 47 files with indirect coverage changes

rchowell · 2025-02-06T16:28:07Z

daft/context.py

+    def __init__(self, ctx: PyDaftContext | None = None):
+        if ctx is not None:
+            self._ctx = ctx
+        else:
+            self._ctx = PyDaftContext()


Is there a reason to stop using __new__ with the lock? I see how the other methods are protected by the mutex on the Rust side, but I'm not sure I follow how the singleton is thread safe.

we don't need the __new__ because having multiple python instances doesn't matter anymore. they are all backed by the single rust based context now.

Understood, my question is about how the Rust singleton's construction is protected from multiple instantiations.

Oh got it. It's protected by making the constructor and the fields private so there is no way to create a DaftContext outside of using get_context.

Ok, and get_context is protected by the OnceCell. So now it's possible to have multiple python context objects, but they are always backed by the same Rust context. I think that's fine? The original issue was in regards to having multiple runners, so this is probably ok.

rchowell · 2025-02-06T16:48:58Z

src/daft-context/src/lib.rs

+static DAFT_CONTEXT: OnceCell<DaftContext> = OnceCell::new();
+
+#[cfg(feature = "python")]
+pub fn get_context() -> DaftContext {


The global singleton context is pretty useful, I wonder if there are other ways we can produce the required functionality? In my current mental model, all user interactions with the system go through a session/connection whether that is implicit or explicit. Whereas the Context I think of a global program state for the process itself. I'm curious what functionality you had in mind where multiple contexts would be necessary which cannot be achieved by pulling that necessary state into the session?

I know that you have said that currently "context" and "session" are synonymous, but perhaps this is the reason to no longer make them so?

context – process state

session – connection state

We are pushing up against the same concepts and patterns described in SQL standard Part 1 where our context represents state for the SQL-environment (singleton), whereas our session should represent state for the SQL-connection. This is not meant to be definitive, but hopefully this helps us tease out the role of context vs the role of session.

universalmind303 · 2025-02-06T18:24:06Z

I'm curious what functionality you had in mind where multiple contexts would be necessary which cannot be achieved by pulling that necessary state into the session?

Ideally i think the config objects should be pulled into the session though I'm not sure yet what makes the most sense for the `runner'.

I'm totally open to adapting this to our needs for session as things evolve. This PR is solely to get a 1-1 mapping with our existing python DaftContext that can be used within connect and potentially elsewhere.

rchowell

This PR is solely to get a 1-1 mapping with our existing python DaftContext

Understood, and this is a huge help for when I lower catalog things into Rust. For this switchover, I believe the proof is in the tests, no concerns beyond the thread safe singleton which you've addressed.

universalmind303 · 2025-02-06T19:01:03Z

cc @raunakab, just a heads up, this PR will cause conflicts with your PR #3756

raunakab · 2025-02-06T19:05:25Z

@universalmind303 Looks like we're moving a lot of global state over to the Rust side of things? If so, I don't mind having this merge into main and then merging main into my changes!

I believe this should be a simple update. I will just need to see how to you've proposed global state in Rust and follow that pattern for daft broadcasting.

raunakab · 2025-02-06T19:12:28Z

src/daft-context/src/lib.rs

@universalmind303 What's the purpose of decorating functions with #[cfg(not(feature = "python"))] and having their bodies remain as unimplemented!()? I see this pattern has been used for all methods / functions defined in this file.

pretty much all of the context depends on python as the runners are implemented in python, so this is just needed to not break things when running with cargo test|check|clippy --no-default-features.

honestly, i'd like to just remove the python feature flag all together, but it makes rust testing less portable as it needs to link to python.

@rchowell

ports the `DaftContext` to the rust side so we can more easily reuse it for connect & catalog work. notes for reviewers: - after removing `daft-local-execution` from daft-connect, it made `cargo check` fail, so some of the changes to unrelated `Cargo.toml`'s was to properly feature flag them. We missed a couple that were causing things to break. - @rchowell I requested review from you specifically to take a look at [context.rs](https://github.com/Eventual-Inc/Daft/pull/3767/files#diff-5b536482a8303505c0c91a561a632218591bf8f02fff21f9f8a3535a3f0ff8a5). This contains the "context"/"session", and the global state that was previously only available in python. I added the `catalog` to this state as well.

universalmind303 added 7 commits February 4, 2025 11:56

wip

ecf2c85

wip

b21698e

wip

99999f8

wip

ca95b31

wip

e865e4d

wip

f96ce9e

port daftcontext to rust

a583c03

github-actions bot added the refactor label Feb 4, 2025

universalmind303 added 3 commits February 4, 2025 15:00

code cleanup

f65b197

code cleanup

dde9d59

idk

5afa072

universalmind303 added 7 commits February 5, 2025 09:14

try to fix linking error

013f0e1

try to fix linking error

1f73fca

try to fix linking error

3751b7f

try to fix linking error

60ed15d

fix more py feature flags

cf8b0ed

fix more py feature flags

01df38b

fix runner setting

5df866e

universalmind303 requested review from jaychia, samster25 and rchowell February 5, 2025 16:48

universalmind303 added 4 commits February 5, 2025 10:53

revert test

49b3e8e

fix runner issue

c269de5

code cleanup

89d79e2

unused deps

b74ee62

universalmind303 commented Feb 5, 2025

View reviewed changes

universalmind303 added 2 commits February 5, 2025 16:50

fix tests

3160f97

fix tests

cf7a951

rchowell reviewed Feb 6, 2025

View reviewed changes

universalmind303 added 2 commits February 6, 2025 12:04

remove catalog from context

cccd4b0

make state more private

144ad60

rchowell approved these changes Feb 6, 2025

View reviewed changes

universalmind303 merged commit d40aef8 into Eventual-Inc:main Feb 6, 2025
41 checks passed

raunakab reviewed Feb 6, 2025

View reviewed changes

raunakab mentioned this pull request Feb 6, 2025

feat: Broadcast plan over port 3238 #3756

Closed

rchowell mentioned this pull request Feb 7, 2025

feat(catalog): Defines a session for connection state #3782

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

refactor: port DaftContext to rust side #3767

refactor: port DaftContext to rust side #3767

universalmind303 commented Feb 4, 2025 •

edited

Loading

codspeed-hq bot commented Feb 4, 2025 •

edited

Loading

universalmind303 Feb 5, 2025

rchowell Feb 6, 2025

universalmind303 Feb 6, 2025

universalmind303 Feb 6, 2025

codecov bot commented Feb 6, 2025 •

edited

Loading

rchowell Feb 6, 2025

universalmind303 Feb 6, 2025

rchowell Feb 6, 2025

universalmind303 Feb 6, 2025 •

edited

Loading

rchowell Feb 6, 2025

rchowell Feb 6, 2025

universalmind303 commented Feb 6, 2025 •

edited

Loading

rchowell left a comment

universalmind303 commented Feb 6, 2025

raunakab commented Feb 6, 2025

raunakab Feb 6, 2025

universalmind303 Feb 6, 2025

universalmind303 Feb 6, 2025

refactor: port DaftContext to rust side #3767

refactor: port DaftContext to rust side #3767

Conversation

universalmind303 commented Feb 4, 2025 • edited Loading

codspeed-hq bot commented Feb 4, 2025 • edited Loading

Merging #3767 will improve performances by 54.97%

Summary

Benchmarks breakdown

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Feb 6, 2025 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

universalmind303 Feb 6, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

universalmind303 commented Feb 6, 2025 • edited Loading

rchowell left a comment

Choose a reason for hiding this comment

universalmind303 commented Feb 6, 2025

raunakab commented Feb 6, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

universalmind303 commented Feb 4, 2025 •

edited

Loading

codspeed-hq bot commented Feb 4, 2025 •

edited

Loading

codecov bot commented Feb 6, 2025 •

edited

Loading

universalmind303 Feb 6, 2025 •

edited

Loading

universalmind303 commented Feb 6, 2025 •

edited

Loading