WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-31 12:19:58 +02:00
Code Issues Packages Projects Releases Wiki Activity

serialize/deserialize pandas with pickle fallback (#5057)

Browse Source
This commit is contained in:
Sydney Runkle
2025-06-12 15:14:00 -04:00
committed by GitHub
parent 116b5d1cac
commit 3ad061f0d7
9 changed files with 1055 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/docs/concepts/persistence.md
+17
View File
@@ -473,6 +473,23 @@ If the checkpointer is used with asynchronous graph execution (i.e. executing th
When checkpointers save the graph state, they need to serialize the channel values in the state. This is done using serializer objects.
`langgraph_checkpoint` defines [protocol][langgraph.checkpoint.serde.base.SerializerProtocol] for implementing serializers provides a default implementation ([JsonPlusSerializer][langgraph.checkpoint.serde.jsonplus.JsonPlusSerializer]) that handles a wide variety of types, including LangChain and LangGraph primitives, datetimes, enums and more.
#### Serialization with `pickle`
The default serializer, [`JsonPlusSerializer`][langgraph.checkpoint.serde.jsonplus.JsonPlusSerializer], uses ormsgpack and JSON under the hood, which is not suitable for all types of objects.
If you want to fallback to pickle for objects not currently supported by our msgpack encoder (such as Pandas dataframes),
you can use the `pickle_fallback` argument of the `JsonPlusSerializer`:
```python
from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
# ... Define the graph ...
graph.compile(
checkpointer=MemorySaver(serde=JsonPlusSerializer(pickle_fallback=True))
)
```
#### Encryption
Checkpointers can optionally encrypt all persisted state. To enable this, pass an instance of [`EncryptedSerializer`][langgraph.checkpoint.serde.encrypted.EncryptedSerializer] to the `serde` argument of any `BaseCheckpointSaver` implementation. The easiest way to create an encrypted serializer is via [`from_pycryptodome_aes`][langgraph.checkpoint.serde.encrypted.EncryptedSerializer.from_pycryptodome_aes], which reads the AES key from the `LANGGRAPH_AES_KEY` environment variable (or accepts a `key` argument):
libs/checkpoint-postgres/uv.lock
Generated
+1
View File
@@ -327,6 +327,7 @@ dev = [
{ name = "dataclasses-json" },
{ name = "mypy" },
{ name = "numpy" },
{ name = "pandas" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
{ name = "pytest-mock" },
libs/checkpoint-sqlite/uv.lock
Generated
+1
View File
@@ -339,6 +339,7 @@ dev = [
{ name = "dataclasses-json" },
{ name = "mypy" },
{ name = "numpy" },
{ name = "pandas" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
{ name = "pytest-mock" },
libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
+1
View File
@@ -475,6 +475,7 @@ def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
buf = obj.tobytes(order="A")
meta = (obj.dtype.str, obj.shape, order, buf)
return ormsgpack.Ext(EXT_NUMPY_ARRAY, _msgpack_enc(meta))
elif isinstance(obj, BaseException):
return repr(obj)
else:
libs/checkpoint/pyproject.toml
+1
View File
@@ -30,6 +30,7 @@ dev = [
"mypy",
"dataclasses-json",
"numpy",
"pandas",
]
[tool.hatch.build.targets.wheel]
libs/checkpoint/tests/test_jsonplus.py
+133 -12
View File
@@ -12,6 +12,7 @@ from zoneinfo import ZoneInfo
import dataclasses_json
import numpy as np
import pandas as pd
import pytest
from pydantic import BaseModel, SecretStr
from pydantic.v1 import BaseModel as BaseModelV1
@@ -332,19 +333,139 @@ def test_serde_jsonplus_numpy_array_json_hook(arr: np.ndarray) -> None:
assert result == arr.tolist()
def test_loads_cannot_find() -> None:
serde = JsonPlusSerializer()
@pytest.mark.parametrize(
"df",
[
pd.DataFrame(),
pd.DataFrame({"int_col": [1, 2, 3]}),
pd.DataFrame({"float_col": [1.1, 2.2, 3.3]}),
pd.DataFrame({"str_col": ["a", "b", "c"]}),
pd.DataFrame({"bool_col": [True, False, True]}),
pd.DataFrame(
{
"datetime_col": [
datetime(2024, 1, 1),
datetime(2024, 1, 2),
datetime(2024, 1, 3),
]
}
),
pd.DataFrame(
{
"int_col": [1, 2, 3],
"float_col": [1.1, 2.2, 3.3],
"str_col": ["a", "b", "c"],
}
),
pd.DataFrame(
{
"int_col": [1, 2, None],
"float_col": [1.1, None, 3.3],
"str_col": ["a", None, "c"],
}
),
pd.DataFrame({"cat_col": pd.Categorical(["a", "b", "a", "c"])}),
pd.DataFrame(
{
"int8": pd.array([1, 2, 3], dtype="int8"),
"int16": pd.array([10, 20, 30], dtype="int16"),
"int32": pd.array([100, 200, 300], dtype="int32"),
"int64": pd.array([1000, 2000, 3000], dtype="int64"),
"float32": pd.array([1.1, 2.2, 3.3], dtype="float32"),
"float64": pd.array([10.1, 20.2, 30.3], dtype="float64"),
}
),
pd.DataFrame({"value": [1, 2, 3]}, index=["x", "y", "z"]),
pd.DataFrame(
[[1, 2, 3, 4]],
columns=pd.MultiIndex.from_tuples(
[("A", "X"), ("A", "Y"), ("B", "X"), ("B", "Y")]
),
),
pd.DataFrame(
{"value": [1, 2, 3]}, index=pd.date_range("2024-01-01", periods=3, freq="D")
),
pd.DataFrame(
{
"col1": range(1000),
"col2": [f"str_{i}" for i in range(1000)],
"col3": np.random.rand(1000),
}
),
pd.DataFrame(
{"tz_datetime": pd.date_range("2024-01-01", periods=3, freq="D", tz="UTC")}
),
pd.DataFrame({"timedelta": pd.to_timedelta([1, 2, 3], unit="D")}),
pd.DataFrame({"period": pd.period_range("2024-01", periods=3, freq="M")}),
pd.DataFrame({"interval": pd.interval_range(start=0, end=3, periods=3)}),
pd.DataFrame({"unicode": ["Hello 🌍", "Python 🐍", "Data 📊"]}),
pd.DataFrame({"mixed": [1, "string", [1, 2, 3], {"key": "value"}]}),
pd.DataFrame({"a": [1], "b": ["test"], "c": [3.14]}),
pd.DataFrame({"single": [42]}),
pd.DataFrame(
{
"small": [sys.float_info.min, 0, sys.float_info.max],
"large_int": [-(2**63), 0, 2**63 - 1],
}
),
pd.DataFrame({"special_strings": ["", "null", "None", "NaN", "inf", "-inf"]}),
pd.DataFrame({"bytes_col": [b"hello", b"world", b"\x00\x01\x02"]}),
],
)
def test_serde_jsonplus_pandas_dataframe(df: pd.DataFrame) -> None:
serde = JsonPlusSerializer(pickle_fallback=True)
dumped = (
"json",
b'{"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydanticccc"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}',
)
dumped = serde.dumps_typed(df)
assert dumped[0] == "pickle"
result = serde.loads_typed(dumped)
assert result.equals(df)
assert serde.loads_typed(dumped) is None, "Should return None if cannot find class"
dumped = (
"json",
b'{"lc": 2, "type": "constructor", "id": ["tests", "test_jsonpluss", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}',
)
@pytest.mark.parametrize(
"series",
[
pd.Series([]),
pd.Series([1, 2, 3]),
pd.Series([1.1, 2.2, 3.3]),
pd.Series(["a", "b", "c"]),
pd.Series([True, False, True]),
pd.Series([datetime(2024, 1, 1), datetime(2024, 1, 2), datetime(2024, 1, 3)]),
pd.Series([1, 2, None]),
pd.Series([1.1, None, 3.3]),
pd.Series(["a", None, "c"]),
pd.Series(pd.Categorical(["a", "b", "a", "c"])),
pd.Series([1, 2, 3], dtype="int8"),
pd.Series([10, 20, 30], dtype="int16"),
pd.Series([100, 200, 300], dtype="int32"),
pd.Series([1000, 2000, 3000], dtype="int64"),
pd.Series([1.1, 2.2, 3.3], dtype="float32"),
pd.Series([10.1, 20.2, 30.3], dtype="float64"),
pd.Series([1, 2, 3], index=["x", "y", "z"]),
pd.Series([1, 2, 3], index=pd.date_range("2024-01-01", periods=3, freq="D")),
pd.Series(range(1000)),
pd.Series(pd.date_range("2024-01-01", periods=3, freq="D", tz="UTC")),
pd.Series(pd.to_timedelta([1, 2, 3], unit="D")),
pd.Series(pd.period_range("2024-01", periods=3, freq="M")),
pd.Series(pd.interval_range(start=0, end=3, periods=3)),
pd.Series(["Hello 🌍", "Python 🐍", "Data 📊"]),
pd.Series([1, "string", [1, 2, 3], {"key": "value"}]),
pd.Series([42], name="single"),
pd.Series([sys.float_info.min, 0, sys.float_info.max]),
pd.Series([-(2**63), 0, 2**63 - 1]),
pd.Series(["", "null", "None", "NaN", "inf", "-inf"]),
pd.Series([b"hello", b"world", b"\x00\x01\x02"]),
pd.Series([1, 2, 3], name="named_series"),
pd.Series(
[10, 20],
index=pd.MultiIndex.from_tuples([("a", 1), ("b", 2)], names=["x", "y"]),
),
],
)
def test_serde_jsonplus_pandas_series(series: pd.Series) -> None:
serde = JsonPlusSerializer(pickle_fallback=True)
dumped = serde.dumps_typed(series)
assert serde.loads_typed(dumped) is None, "Should return None if cannot find module"
assert dumped[0] == "pickle"
result = serde.loads_typed(dumped)
assert result.equals(series)
libs/checkpoint/uv.lock
Generated
+899 -799
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/uv.lock
Generated
+1
View File
@@ -1329,6 +1329,7 @@ dev = [
{ name = "dataclasses-json" },
{ name = "mypy" },
{ name = "numpy" },
{ name = "pandas" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
{ name = "pytest-mock" },
libs/prebuilt/uv.lock
Generated
+1
View File
@@ -390,6 +390,7 @@ dev = [
{ name = "dataclasses-json" },
{ name = "mypy" },
{ name = "numpy" },
{ name = "pandas" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
{ name = "pytest-mock" },