From caa99c54071e9e84043d58b46a693aad87a029ba Mon Sep 17 00:00:00 2001 From: Nikhil Benesch Date: Fri, 21 Aug 2026 09:56:20 -0400 Subject: [PATCH 1/2] Limit transparent vector encoding to write upserts. Co-authored-by: Cursor --- src/turbopuffer/_utils/_transform.py | 49 +++++++++++++++++++++------- tests/custom/test_vectors.py | 31 ++++++++++-------- 2 files changed, 56 insertions(+), 24 deletions(-) diff --git a/src/turbopuffer/_utils/_transform.py b/src/turbopuffer/_utils/_transform.py index b6459b50..806534fc 100644 --- a/src/turbopuffer/_utils/_transform.py +++ b/src/turbopuffer/_utils/_transform.py @@ -32,6 +32,7 @@ strip_annotated_type, ) from ..lib.vector import b64encode_vector +from ..types.namespace_write_params import NamespaceWriteParams _T = TypeVar("_T") @@ -115,7 +116,7 @@ def _vector_keys_from_write_body(data: object) -> frozenset[str]: # turbopuffer: Simple transform without expensive type introspection. -def _turbopuffer_transform(obj: object, *, vector_keys: frozenset[str]) -> object: +def _turbopuffer_transform(obj: object) -> object: if obj is None or isinstance(obj, (int, float, bool, str)): return obj if isinstance(obj, dict): @@ -124,22 +125,42 @@ def _turbopuffer_transform(obj: object, *, vector_keys: frozenset[str]) -> objec # Strip Omit and NotGiven values if isinstance(v, (Omit, NotGiven)): continue - if k in vector_keys: - result[k] = _encode_vector(v) - else: - result[k] = _turbopuffer_transform(v, vector_keys=vector_keys) + result[k] = _turbopuffer_transform(v) return result if isinstance(obj, list): - return [_turbopuffer_transform(i, vector_keys=vector_keys) for i in cast(List[object], obj)] + return [_turbopuffer_transform(i) for i in cast(List[object], obj)] if isinstance(obj, tuple): - return tuple(_turbopuffer_transform(i, vector_keys=vector_keys) for i in cast(tuple[object, ...], obj)) + return tuple(_turbopuffer_transform(i) for i in cast(tuple[object, ...], obj)) return obj +def _encode_write_vectors(data: object) -> object: + """Base64-encode vector attrs on write bodies (top-level document fields only).""" + if not isinstance(data, dict): + return data + body = cast(dict[str, object], data) + vector_keys = _vector_keys_from_write_body(body) + result = dict(body) + + def encode_attrs(attrs: dict[str, object]) -> dict[str, object]: + return {k: _encode_vector(v) if k in vector_keys else v for k, v in attrs.items()} + + rows = result.get("upsert_rows") + if is_iterable(rows) and not isinstance(rows, (str, bytes, dict)): + result["upsert_rows"] = [ + encode_attrs(cast(dict[str, object], row)) if isinstance(row, dict) else row + for row in cast(Any, rows) + ] + columns = result.get("upsert_columns") + if isinstance(columns, dict): + result["upsert_columns"] = encode_attrs(cast(dict[str, object], columns)) + return result + + # Wrapper over _transform_recursive providing fake types def transform( data: _T, - expected_type: object, # noqa: ARG001 - kept for API compatibility + expected_type: object, ) -> _T: """Transform dictionaries based off of type information from the given type, for example: @@ -157,7 +178,10 @@ class Params(TypedDict, total=False): It should be noted that the transformations that this function does are not represented in the type system. """ # turbopuffer: Use simple vector encoding instead of generic type-based transform. - return cast(_T, _turbopuffer_transform(data, vector_keys=_vector_keys_from_write_body(data))) + data = cast(_T, _turbopuffer_transform(data)) + if expected_type is NamespaceWriteParams: + return cast(_T, _encode_write_vectors(data)) + return data @lru_cache(maxsize=8096) @@ -358,7 +382,7 @@ async def async_maybe_transform( async def async_transform( data: _T, - expected_type: object, # noqa: ARG001 - kept for API compatibility + expected_type: object, ) -> _T: """Transform dictionaries based off of type information from the given type, for example: @@ -376,7 +400,10 @@ class Params(TypedDict, total=False): It should be noted that the transformations that this function does are not represented in the type system. """ # turbopuffer: Use simple vector encoding instead of generic type-based transform. - return cast(_T, _turbopuffer_transform(data, vector_keys=_vector_keys_from_write_body(data))) + data = cast(_T, _turbopuffer_transform(data)) + if expected_type is NamespaceWriteParams: + return cast(_T, _encode_write_vectors(data)) + return data async def _async_transform_recursive( diff --git a/tests/custom/test_vectors.py b/tests/custom/test_vectors.py index 7fe5d443..d1ebafd8 100644 --- a/tests/custom/test_vectors.py +++ b/tests/custom/test_vectors.py @@ -10,10 +10,9 @@ from turbopuffer.types import ( Row, Vector, - RowParam, - ColumnsParam, QueryBilling, VectorEncoding, + NamespaceWriteParams, NamespaceQueryResponse, namespace_query_params, ) @@ -566,11 +565,13 @@ def test_transparent_vector_encoding(): # Due to the nature of the hack, there's a high risk of a future refactoring # to the Stainless specification breaking the fast path. - transformed = transform({"id": 1, "vector": [0.1, 0.2, 0.3]}, RowParam) - assert transformed == {"id": 1, "vector": "zczMPc3MTD6amZk+"} + transformed = transform({"upsert_rows": [{"id": 1, "vector": [0.1, 0.2, 0.3]}]}, NamespaceWriteParams) + assert transformed == {"upsert_rows": [{"id": 1, "vector": "zczMPc3MTD6amZk+"}]} - transformed = transform({"id": [1, 2], "vector": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}, ColumnsParam) - assert transformed == {"id": [1, 2], "vector": ["zczMPc3MTD6amZk+", "zczMPgAAAD+amRk/"]} + transformed = transform( + {"upsert_columns": {"id": [1, 2], "vector": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}}, NamespaceWriteParams + ) + assert transformed == {"upsert_columns": {"id": [1, 2], "vector": ["zczMPc3MTD6amZk+", "zczMPgAAAD+amRk/"]}} # Named vector columns are encoded when schema sets ann; ann: False does not encode. body = { @@ -589,7 +590,7 @@ def test_transparent_vector_encoding(): "other_embedding": {"type": "[3]f32", "ann": False}, }, } - transformed = transform(body, object) + transformed = transform(body, NamespaceWriteParams) assert transformed == { "upsert_rows": [ { @@ -612,7 +613,7 @@ def test_transparent_vector_encoding(): "upsert_rows": [{"id": 1, "title_embedding": [0.1, 0.2, 0.3]}], "schema": {"title_embedding": {"type": "[3]f32", "ann": dict[str, object]()}}, } - transformed = transform(body_empty_ann, object) + transformed = transform(body_empty_ann, NamespaceWriteParams) assert transformed == { "upsert_rows": [{"id": 1, "title_embedding": "zczMPc3MTD6amZk+"}], "schema": {"title_embedding": {"type": "[3]f32", "ann": {}}}, @@ -625,11 +626,15 @@ async def test_transparent_vector_encoding_async(): # Due to the nature of the hack, there's a high risk of a future refactoring # to the Stainless specification breaking the fast path. - transformed = await async_transform({"id": 1, "vector": [0.1, 0.2, 0.3]}, RowParam) - assert transformed == {"id": 1, "vector": "zczMPc3MTD6amZk+"} + transformed = await async_transform( + {"upsert_rows": [{"id": 1, "vector": [0.1, 0.2, 0.3]}]}, NamespaceWriteParams + ) + assert transformed == {"upsert_rows": [{"id": 1, "vector": "zczMPc3MTD6amZk+"}]} - transformed = await async_transform({"id": [1, 2], "vector": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}, ColumnsParam) - assert transformed == {"id": [1, 2], "vector": ["zczMPc3MTD6amZk+", "zczMPgAAAD+amRk/"]} + transformed = await async_transform( + {"upsert_columns": {"id": [1, 2], "vector": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}}, NamespaceWriteParams + ) + assert transformed == {"upsert_columns": {"id": [1, 2], "vector": ["zczMPc3MTD6amZk+", "zczMPgAAAD+amRk/"]}} body = { "upsert_rows": [ @@ -647,7 +652,7 @@ async def test_transparent_vector_encoding_async(): "other_embedding": {"type": "[3]f32", "ann": False}, }, } - transformed = await async_transform(body, object) + transformed = await async_transform(body, NamespaceWriteParams) assert transformed == { "upsert_rows": [ { From 01d90759468c1dd2be7f9c32a2f67015b2460be9 Mon Sep 17 00:00:00 2001 From: Nikhil Benesch Date: Fri, 21 Aug 2026 09:58:04 -0400 Subject: [PATCH 2/2] Fix ruff formatting. Co-authored-by: Cursor --- src/turbopuffer/_utils/_transform.py | 3 +-- tests/custom/test_vectors.py | 4 +--- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/src/turbopuffer/_utils/_transform.py b/src/turbopuffer/_utils/_transform.py index 806534fc..7444189b 100644 --- a/src/turbopuffer/_utils/_transform.py +++ b/src/turbopuffer/_utils/_transform.py @@ -148,8 +148,7 @@ def encode_attrs(attrs: dict[str, object]) -> dict[str, object]: rows = result.get("upsert_rows") if is_iterable(rows) and not isinstance(rows, (str, bytes, dict)): result["upsert_rows"] = [ - encode_attrs(cast(dict[str, object], row)) if isinstance(row, dict) else row - for row in cast(Any, rows) + encode_attrs(cast(dict[str, object], row)) if isinstance(row, dict) else row for row in cast(Any, rows) ] columns = result.get("upsert_columns") if isinstance(columns, dict): diff --git a/tests/custom/test_vectors.py b/tests/custom/test_vectors.py index d1ebafd8..5e3ed80c 100644 --- a/tests/custom/test_vectors.py +++ b/tests/custom/test_vectors.py @@ -626,9 +626,7 @@ async def test_transparent_vector_encoding_async(): # Due to the nature of the hack, there's a high risk of a future refactoring # to the Stainless specification breaking the fast path. - transformed = await async_transform( - {"upsert_rows": [{"id": 1, "vector": [0.1, 0.2, 0.3]}]}, NamespaceWriteParams - ) + transformed = await async_transform({"upsert_rows": [{"id": 1, "vector": [0.1, 0.2, 0.3]}]}, NamespaceWriteParams) assert transformed == {"upsert_rows": [{"id": 1, "vector": "zczMPc3MTD6amZk+"}]} transformed = await async_transform(