Improve st.experimental_memo exception on Snowpark/PySpark integration #5791
Conversation
329ad8c to
9d07678
Compare
There was a problem hiding this comment.
Looks great, but needs a test. (Looks like we didn't already have a test for the snowflake.snowpark.dataframe.DataFrame case, but we should have).
It might be useful to pull this set of type names into a tuple at the top of the file, e.g.:
# We show a special "UnevaluatedDataFrame" warning for cached funcs
# that attempt to return one of these unserializable types:
UNEVALUATED_DATAFRAME_TYPES = (
"snowflake.snowpark.table.Table", etc
)So that we can just iterate them here, and also in the test.
There was a problem hiding this comment.
I know this isn't part of the PR, but can we move this class declaration outside the function - to the top of the file, or wherever?
There was a problem hiding this comment.
Sure 👍 It's moved now into cache_errors.py file 👍
9d07678 to
e261a1a
Compare
|
@tconkling thank your for your review 🥇 This PR is now updated, can you please re-review it? 🙇 |
e261a1a to
18f5a19
Compare
tconkling
left a comment
There was a problem hiding this comment.
Great, thanks for the tests!
Before merging, can we get some clarity on which exception types should actually be caught to handle the UnevaluatedDataFrameError?
| class UnevaluatedDataFrameError(StreamlitAPIException): | ||
| pass | ||
|
|
||
| except (TypeError, RuntimeError, CacheError): |
There was a problem hiding this comment.
This seems overly broad - is it actually possible for any of these errors to be raised from an uncollected dataframe? (My assumption is that it's really just a CacheError that could be raised from a pickling failure?)
There was a problem hiding this comment.
The reason I've introduced 3 types of exceptions here:
- Snowpark triggers TypeError
- Spark triggers RunTime error in the Testing environment, and CacheError when run on normal build
Is that ok?
There was a problem hiding this comment.
I'd love to narrow this down - what are the operations that cause the TypeError and RuntimeError errors?
in get_or_create_cached_value, we’re catching exceptions thrown by cache.write_result(value_key, return_value, messages). The st.experimental_memo implementation of cache.write_result uses the pickle utility to get the bytes to write to the cache, and pickling errors get re-raised as CacheError:
try:
pickled_entry = pickle.dumps(multi_cache_results)
except pickle.PicklingError as exc:
raise CacheError(f"Failed to pickle {key}") from excIf we get TypeError and RuntimeError thrown from write_result, then it seems like it must be happening elsewhere in the function?
* develop: Singleton: `ttl` and `max_entries` support (#5821) Bump express from 4.17.1 to 4.18.2 in /frontend (#5823) Bump qs from 6.5.2 to 6.5.3 in /frontend (#5822) Bump decode-uri-component from 0.2.0 to 0.2.2 in /frontend (#5802) Improve st.experimental_memo exception on Snowpark/PySpark integration (#5791)
📚 Context
Please describe the project or issue background here
This PR shows custom
st.experimental_memoexception forsnowflake.snowpark.table.Tableandpyspark.sql.dataframe.DataFrame, additionally tosnowflake.snowpark.dataframe.DataFrame.What kind of change does this PR introduce?
🧠 Description of Changes
Add bullet points summarizing your changes here
Revised:
Insert screenshot of your updated UI/code here
Current:
Insert screenshot of existing UI/code here
🧪 Testing Done
🌐 References
Does this depend on other work, documents, or tickets?
Contribution License Agreement
By submitting this pull request you agree that all contributions to this project are made under the Apache 2.0 license.