[5/n subset refactor] [serdes] Enable serializing mappings with non-scalar keys #18057

clairelin135 · 2023-11-16T01:22:03Z

This PR enables serializing mappings with non-scalar keys, intended to be used for serializing mappings keyed by asset key. Previously this was not possible because seven.json.dumps can only serialize dictionaries keyed by a primitive or a string. Without this change, we would need a custom FieldSerializer when serializing any dict keyed by a non-scalar.

clairelin135 · 2023-11-16T01:22:20Z

Current dependencies on/for this PR:

master
- PR [1/n subset refactor] Split TimeWindowPartitionsSubset #17684
  - PR [2/n subset refactor] Make TimeWindowPartitionsDefinition serializable #17660
    - PR [3/n subset refactor] Add whitelist_for_serdes to TimeWindowPartitionsSubset #17702
      - PR [4/n subset refactor] Add whitelist_for_serdes to DefaultPartitionsSubset #17703
        
        PR [5/n subset refactor] [serdes] Enable serializing mappings with non-scalar keys #18057 👈
        
        PR [6/n subset refactor] Serialize AssetGraphSubset and AssetBackfillData with whitelist_for_serdes #17844
        
        PR [7/n subset refactor] Use new asset backfill data serialization format #17929

This stack of pull requests is managed by Graphite.

clairelin135 · 2023-11-16T19:23:22Z

python_modules/dagster/dagster/_core/definitions/asset_key.py

@@ -0,0 +1,189 @@
+import re


directly copies over AssetKey code from the other file to avoid circular imports

OwenKephart · 2023-11-16T19:42:56Z

python_modules/dagster/dagster/_serdes/serdes.py

+                    for key, value in dict_val.items()
+                }
+            }
+
        return {
            key: _pack_value(value, whitelist_map, f"{descent_path}.{key}")


apologies for snooping, but the title of this PR was very exciting as this has been a huge pain in the past -- but what happens if _pack_value() is called on both the key and the value in all cases (so we don't need to special-case this asset key thing)?

_pack_value() is a no-op on ints, bools, floats, strs, etc., and so this shouldn't have any impact on existing dictionaries (I think), but transparently allows fancier objects to get serde'd properly

when you call _pack_value on an asset key, it serializes it as a dictionary, which cannot exist as the key of a dictionary

alangenfeld

I can help you come up with some benchmarks to verify, but suspect doing these type checks against all the keys all the time is too costly.

Maybe we can cook up a special type (AssetMap[TVal] or something) that implements a dict like interface and then special case serdes against that?

alangenfeld · 2023-11-16T21:47:06Z

python_modules/dagster/dagster/_serdes/serdes.py

+
+        dict_val = cast(dict, val)
+
+        if dict_val and all(type(key) is AssetKey for key in dict_val.keys()):


suspect this is a potentially rough perf regression (related comment on 675)

sryza · 2023-11-20T22:12:55Z

the idea is that the custom type implements Mapping[AssetKey, TVal] so the only places that would have to know about the new type would be the decl of the serializable NamedTuples, any callsites would just treat it as the interface Mapping[AssetKey, Tval]

Not the end of the world, but I think this would make our codebase less grokkable. A lot of our core entities are these serializable NamedTuples, and someone ramping up and looking at them needs to wonder what the difference is between and AssetMap and a Mapping[AssetKey, ...].

alangenfeld · 2023-11-20T22:22:29Z

I think you could limit the only reference to creating the instance in the body of __new__ (with a comment). The typehints on NamedTuple and __new__ could stay Mapping

alangenfeld · 2023-11-20T22:26:13Z

python_modules/dagster/dagster/_serdes/serdes.py

+        from dagster._core.events import AssetKey
+
+        return {
+            AssetKey.from_user_string(key): _unpack_value(value, whitelist_map, context)


to_user_string / from_user_string doesn't correctly roundtrip multipart asset keys that contain / in them

>>> from dagster import AssetKey >>> a = AssetKey(['this/sucks', 'a////bit']) >>> a AssetKey(['this/sucks', 'a////bit']) >>> a.to_user_string() 'this/sucks/a////bit' >>> b = AssetKey.from_user_string(a.to_user_string()) >>> b AssetKey(['this', 'sucks', 'a', '', '', '', 'bit']) >>> a == b False

ahh this is a good catch. Initially I picked to_user_string because I thought it would likely be more performant than using JSON, but probably a good idea to instead call asset_key.to_string()

github-actions · 2023-11-20T23:32:38Z

Deploy preview for dagit-core-storybook ready!

✅ Preview
https://dagit-core-storybook-3c1y26zsp-elementl.vercel.app
https://11-15-enable-serializing-dicts-keyed-by-asset-key.core-storybook.dagster-docs.io

Built with commit 793517c.
This pull request is being automatically deployed with vercel-action

github-actions · 2023-11-20T23:33:03Z

Deploy preview for dagster-docs ready!

Preview available at https://dagster-docs-3u5jr7akg-elementl.vercel.app
https://11-15-enable-serializing-dicts-keyed-by-asset-key.dagster.dagster-docs.io

Direct link to changed pages:

clairelin135 · 2023-11-20T23:37:37Z

This PR has been updated to contain a custom AssetKeyMap object. Since it subclasses Mapping, we indeed can just call this object within __new__ so we don't need to change the field types to use AssetKeyMap. Interested in thoughts so far.

I'm still looking, but haven't found an easy way to get the annotated key type of a mapping.

sryza · 2023-11-20T23:53:28Z

This PR has been updated to contain a custom AssetKeyMap object. Since it subclasses Mapping, we indeed can just call this object within new so we don't need to change the field types to use AssetKeyMap. Interested in thoughts so far.

I'd ideally like to end up in a world where we don't have __new__, using dataclasses.

clairelin135 · 2023-11-21T00:00:23Z

I'd ideally like to end up in a world where we don't have new, using dataclasses.

I think this is ok, if we switched to using dataclasses we could always implement the __post_init to enforce that the mapping is an AssetKeyMap

sryza · 2023-11-21T00:09:27Z

I'm still looking, but haven't found an easy way to get the annotated key type of a mapping.

We do this when inferring Dagster types, right?

sryza · 2023-11-21T00:21:00Z

From a quick benchmark, performing the type check takes about 5% of the time that it takes to serialize the dict.

import datetime
import json
import random
import string
from typing import Mapping

from dagster import AssetKey


def random_string() -> str:
    return "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(10))


def make_asset_key_dict() -> Mapping[AssetKey, str]:
    return {AssetKey(random_string()): random_string() for _ in range(10000)}


asset_key_dicts = [make_asset_key_dict() for _ in range(10)]

start = datetime.datetime.now()
for d in asset_key_dicts:
    all(type(k) == AssetKey for k in d.keys())

end = datetime.datetime.now()

print(end - start)


start = datetime.datetime.now()
for d in asset_key_dicts:
    json.dumps({str(k): v for k, v in d.items()})

end = datetime.datetime.now()
print(end - start)

Output:

0:00:00.006121
0:00:00.105415

alangenfeld

From a quick benchmark, performing the type check takes about 5% of the time that it takes to serialize the dict.

So the performance concern with the type check is in the context of serializing all the other not AssetKey-keyed dict s flowing through and the extra cost on the check to see if they should get the special non scalar key handling.

With that focus, making the target not a dict via this custom type AssetKeyMap scheme works well as it incurs ~0 cost to the regular dict s.

If we are not feeling this marker type approach, we could:

Get cute / clever with the code to minimize the impact to the scalar key dict s . Checking only one key is one flavor of this that risks some weird errors if someone tries to flow a heterogenous keyed structure through.
Use the type hint information - it should be possible, the get_origin / get_args APIs for poking at generics are a little goofy but we've made it work else where. Still would want benchmark and tweak how / when these checks fire to minimize their cost.

alangenfeld · 2023-11-21T16:02:21Z

python_modules/dagster/dagster/_serdes/serdes.py

@@ -681,10 +708,18 @@ def _pack_value(
            _pack_value(item, whitelist_map, f"{descent_path}[{idx}]")
            for idx, item in enumerate(cast(list, val))
        ]
+    if tval is AssetKeyMap:


nit: you can double check in the benchmark but moving this if check down since the condition is more rare may be measurably beneficial

alangenfeld · 2023-11-21T16:05:09Z

python_modules/dagster/dagster/_serdes/serdes.py

+_V = TypeVar("_V")
+
+
+class AssetKeyMap(Mapping["AssetKey", _V]):


its not much of a lift to instead of just supporting AssetKey keys, support any serdes-able keys.

I was thinking Instead of the to_string-ed key -> value approach currently used at [1], we could pack them in the style of the items API and do something like {'__map_items__': [[_pack(k), _pack(v)] for k, v in map.items()]}

this idea applies even if we change the detection scheme

alangenfeld · 2023-11-21T16:05:22Z

python_modules/dagster/dagster/_serdes/serdes.py

+    if tval is AssetKeyMap:
+        return {
+            "__asset_key_map__": {
+                key.to_string(): _pack_value(value, whitelist_map, f"{descent_path}.{key}")


sryza · 2023-11-21T16:39:52Z

So the performance concern with the type check is in the context of serializing all the other not AssetKey-keyed dict s flowing through and the extra cost on the check to see if they should get the special non scalar key handling.

I'd expect it to be much faster in those cases, because we'd only check the first key and notice that it's not an AssetKey.

sryza · 2023-11-21T16:40:59Z

I'm still concerned about the AssetKeyMap approach, because it adds to the lore you need to know about the codebase, to do basic data modeling.

alangenfeld · 2023-11-21T17:02:43Z

I'd expect it to be much faster in those cases

Right agree the keys check failing should be faster than the keys check passing, but even then this code path is hot enough that adding that key check is impactful.

From this internal discussion around benchmarking

https://dagsterlabs.slack.com/archives/C03CA4TVCAW/p1700520085656059?thread_ts=1700517125.390159&cid=C03CA4TVCAW

seeing the serialization times increase by about 1.6x

Requested changes addressed -- enabled serializing mappings with generic non-scalar keys

alangenfeld · 2023-11-21T21:03:43Z

python_modules/dagster/dagster/_serdes/serdes.py

+    """Wrapper class for non-scalar key mappings, used to performantly type check when serializing
+    without having to access types of specific keys.


without having to access types of specific keys.

nit: maybe something like

"without impacting the performance of serializing the more common scalar key dicts. May be replaceable with a different clever scheme"

alangenfeld · 2023-11-21T21:06:40Z

python_modules/dagster/dagster/_serdes/serdes.py

@@ -686,6 +711,16 @@ def _pack_value(
            key: _pack_value(value, whitelist_map, f"{descent_path}.{key}")
            for key, value in cast(dict, val).items()
        }
+    if tval is SerializableNonScalarKeyMapping:
+        return {
+            "__non_scalar_key_mapping_items__": [


nit: maybe something a little more terse like __mapping_items__

clairelin135 mentioned this pull request Nov 16, 2023

[7/n subset refactor] Use new asset backfill data serialization format #17929

Merged

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from dbfc898 to a5942f3 Compare November 16, 2023 18:43

clairelin135 changed the base branch from claire/default-subset-serialization to master November 16, 2023 18:43

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch 3 times, most recently from b5792a1 to 59f6dbf Compare November 16, 2023 18:57

clairelin135 commented Nov 16, 2023

View reviewed changes

clairelin135 changed the base branch from master to claire/default-subset-serialization November 16, 2023 19:25

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from 59f6dbf to deb962a Compare November 16, 2023 19:25

clairelin135 changed the title ~~[4.5/n subset refactor] Enable serializing dicts keyed by asset key~~ [5/n subset refactor] Enable serializing dicts keyed by asset key Nov 16, 2023

OwenKephart reviewed Nov 16, 2023

View reviewed changes

clairelin135 changed the title ~~[5/n subset refactor] Enable serializing dicts keyed by asset key~~ [5/n subset refactor] [serdes] Enable serializing dicts keyed by asset key Nov 16, 2023

clairelin135 marked this pull request as ready for review November 16, 2023 20:38

clairelin135 requested review from sryza and alangenfeld November 16, 2023 20:38

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from deb962a to 8be2bba Compare November 16, 2023 20:48

clairelin135 force-pushed the claire/default-subset-serialization branch from 78a91dd to e6c5ca9 Compare November 16, 2023 20:49

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from 8be2bba to 3e699f0 Compare November 16, 2023 20:49

alangenfeld previously requested changes Nov 16, 2023

View reviewed changes

clairelin135 force-pushed the claire/default-subset-serialization branch from e6c5ca9 to 66d7ba9 Compare November 16, 2023 22:36

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from 3e699f0 to 31698f7 Compare November 16, 2023 22:36

clairelin135 force-pushed the claire/default-subset-serialization branch from 66d7ba9 to 3195b7e Compare November 16, 2023 23:08

Base automatically changed from claire/default-subset-serialization to master November 16, 2023 23:37

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from 31698f7 to 39bcd25 Compare November 17, 2023 00:11

alangenfeld reviewed Nov 20, 2023

View reviewed changes

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from 39bcd25 to 769d52c Compare November 20, 2023 23:29

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from 769d52c to 656cd32 Compare November 20, 2023 23:38

clairelin135 requested a review from alangenfeld November 20, 2023 23:40

alangenfeld reviewed Nov 21, 2023

View reviewed changes

clairelin135 added 2 commits November 21, 2023 10:47

enable serializing dicts keyed by asset key

fc3f2de

amend to use custom AssetKeyMap type

0bf1331

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from 656cd32 to 793517c Compare November 21, 2023 19:41

generalize for non-scalar values

bd4854f

clairelin135 force-pushed the 11-15-enable_serializing_dicts_keyed_by_asset_key branch from 793517c to bd4854f Compare November 21, 2023 19:44

clairelin135 changed the title ~~[5/n subset refactor] [serdes] Enable serializing dicts keyed by asset key~~ [5/n subset refactor] [serdes] Enable serializing mappings with non-scalar keys Nov 21, 2023

clairelin135 requested a review from alangenfeld November 21, 2023 20:58

alangenfeld approved these changes Nov 21, 2023

View reviewed changes

pr feedback

c52e54b

clairelin135 merged commit dde3a5e into master Nov 21, 2023

clairelin135 deleted the 11-15-enable_serializing_dicts_keyed_by_asset_key branch November 21, 2023 22:26

alangenfeld mentioned this pull request Feb 5, 2024

[pyright] 1.1.339 -> 1.1.349 #19482

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[5/n subset refactor] [serdes] Enable serializing mappings with non-scalar keys #18057

[5/n subset refactor] [serdes] Enable serializing mappings with non-scalar keys #18057

clairelin135 commented Nov 16, 2023 •

edited

Loading

clairelin135 commented Nov 16, 2023 •

edited

Loading

clairelin135 Nov 16, 2023

OwenKephart Nov 16, 2023

clairelin135 Nov 16, 2023

alangenfeld left a comment

alangenfeld Nov 16, 2023

sryza commented Nov 20, 2023

alangenfeld commented Nov 20, 2023

alangenfeld Nov 20, 2023 •

edited

Loading

clairelin135 Nov 20, 2023 •

edited

Loading

github-actions bot commented Nov 20, 2023 •

edited

Loading

github-actions bot commented Nov 20, 2023 •

edited

Loading

clairelin135 commented Nov 20, 2023

sryza commented Nov 20, 2023

clairelin135 commented Nov 21, 2023

sryza commented Nov 21, 2023

sryza commented Nov 21, 2023

alangenfeld left a comment

alangenfeld Nov 21, 2023

alangenfeld Nov 21, 2023 •

edited

Loading

alangenfeld Nov 21, 2023

sryza commented Nov 21, 2023

sryza commented Nov 21, 2023

alangenfeld commented Nov 21, 2023

alangenfeld Nov 21, 2023

alangenfeld Nov 21, 2023


		dict_val = cast(dict, val)

		if dict_val and all(type(key) is AssetKey for key in dict_val.keys()):

		_V = TypeVar("_V")


		class AssetKeyMap(Mapping["AssetKey", _V]):

		"""Wrapper class for non-scalar key mappings, used to performantly type check when serializing
		without having to access types of specific keys.

[5/n subset refactor] [serdes] Enable serializing mappings with non-scalar keys #18057

[5/n subset refactor] [serdes] Enable serializing mappings with non-scalar keys #18057

Conversation

clairelin135 commented Nov 16, 2023 • edited Loading

clairelin135 commented Nov 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alangenfeld left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sryza commented Nov 20, 2023

alangenfeld commented Nov 20, 2023

alangenfeld Nov 20, 2023 • edited Loading

Choose a reason for hiding this comment

clairelin135 Nov 20, 2023 • edited Loading

Choose a reason for hiding this comment

github-actions bot commented Nov 20, 2023 • edited Loading

github-actions bot commented Nov 20, 2023 • edited Loading

clairelin135 commented Nov 20, 2023

sryza commented Nov 20, 2023

clairelin135 commented Nov 21, 2023

sryza commented Nov 21, 2023

sryza commented Nov 21, 2023

alangenfeld left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alangenfeld Nov 21, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sryza commented Nov 21, 2023

sryza commented Nov 21, 2023

alangenfeld commented Nov 21, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

clairelin135 commented Nov 16, 2023 •

edited

Loading

clairelin135 commented Nov 16, 2023 •

edited

Loading

alangenfeld Nov 20, 2023 •

edited

Loading

clairelin135 Nov 20, 2023 •

edited

Loading

github-actions bot commented Nov 20, 2023 •

edited

Loading

github-actions bot commented Nov 20, 2023 •

edited

Loading

alangenfeld Nov 21, 2023 •

edited

Loading