Optimize numeric array conversion and serialize directly into Python bytes
Some checks failed
Build / wheels (x86_64, ubuntu-latest) (push) Failing after 1s
Build / source (push) Failing after 1s
Build / lint (push) Failing after 1s
Build / wheels (aarch64, ubuntu-24.04-arm) (push) Has been cancelled
Build / wheels (arm64, macos-15) (push) Has been cancelled

This commit is contained in:
2026-09-21 22:48:21 -05:00
parent caee18aeef
commit 05cc505ef4
10 changed files with 328 additions and 13 deletions

View File

@@ -0,0 +1,61 @@
{
"source_commit": "caee18aeef0a9a54ad8f98292a325a1675a43bec",
"cases": [
{
"size": 0,
"text_size": 0,
"sha256": "5d1e945d489c17931e0bd82f6daba3f8657795d6abbe81bded5dd34e3c4bab6d",
"bytes": 256,
"segments": 1
},
{
"size": 1,
"text_size": 1,
"sha256": "c5f683aec8691d1049acf8913bfdea5d35a126ee4c2745814aebbeea070bde16",
"bytes": 264,
"segments": 1
},
{
"size": 7,
"text_size": 7,
"sha256": "18478024be5786d8d07e68355a5f81f66dbd1ee48d328e4e375d04ca2bda9682",
"bytes": 272,
"segments": 1
},
{
"size": 1024,
"text_size": 1024,
"sha256": "fc40c903202d5c6eee2ffa7bd3be3e651badbc8623dbebe3fbc7361010bed360",
"bytes": 3328,
"segments": 1
},
{
"size": 8192,
"text_size": 8192,
"sha256": "c4eefda6030b879873ff39dc6a0d17cdc8f059af8ba24869a8b8f3802928dd89",
"bytes": 24856,
"segments": 3
},
{
"size": 65536,
"text_size": 65536,
"sha256": "a846392964ce60df5e15ff8c50d400b9783ca2d8b07216f85e9dc24c0706f7fa",
"bytes": 196888,
"segments": 3
},
{
"size": 1048576,
"text_size": 1048576,
"sha256": "ef3d77a4e1f80f2cf839f73c04688410957d6948f284ae0fd46f29a679f973fe",
"bytes": 3146008,
"segments": 3
},
{
"size": 8192,
"text_size": 0,
"sha256": "bf2e44b6eaa518a054ab8a260fcb03e154b0e07866096fa7c222c018abb7baaf",
"bytes": 8464,
"segments": 2
}
]
}

127
test/test_bulk_numeric.py Normal file
View File

@@ -0,0 +1,127 @@
# Copyright (c) 2026 IQ.Lvbs. All rights reserved.
import gc
import capnp
import numpy as np
import pytest
@pytest.fixture(scope="module")
def numeric_schema(tmp_path_factory):
path = tmp_path_factory.mktemp("bulk-schema") / "numeric.capnp"
path.write_text("""# Copyright (c) 2026 IQ.Lvbs. All rights reserved.
@0xe7ab11f278d8b04a;
struct Numbers {
f32 @0 :List(Float32);
f64 @1 :List(Float64);
i64 @2 :List(Int64);
u64 @3 :List(UInt64);
text @4 :List(Text);
nested @5 :List(List(Float32));
}
""")
return capnp.load(str(path)).Numbers
@pytest.mark.parametrize(
("field", "values"),
[
("f32", []),
("f32", [0.0, -0.0, 1.25, -2.5, float("inf"), -float("inf"), float("nan")]),
("f64", [1e-300, 1e300, 1.23456789012345, -0.0]),
("i64", [-(2**63), -1, 0, 2**63 - 1]),
("u64", [0, 2**53 + 1, 2**64 - 1]),
],
)
def test_bulk_values_and_lifetime(numeric_schema, field, values):
builder = numeric_schema.new_message(**{field: values})
encoded = builder.to_bytes()
with numeric_schema.from_bytes(encoded) as reader:
data = getattr(reader, field)
expected = np.array(data, dtype=np.float64)
buffer = data.to_float64_bytes()
actual = np.frombuffer(buffer, dtype=np.float64)
np.testing.assert_array_equal(actual, expected)
np.testing.assert_array_equal(np.signbit(actual), np.signbit(expected))
assert not actual.flags.writeable
del reader, data, encoded, builder, buffer
gc.collect()
np.testing.assert_array_equal(actual, expected)
def test_bulk_rejects_text(numeric_schema):
builder = numeric_schema.new_message(text=["not numeric"])
with numeric_schema.from_bytes(builder.to_bytes()) as reader:
with pytest.raises(TypeError):
reader.text.to_float64_bytes()
def test_bulk_is_an_owned_snapshot(numeric_schema):
builder = numeric_schema.new_message(f32=[1.0, 2.0])
reader = builder.as_reader()
array = np.frombuffer(reader.f32.to_float64_bytes(), dtype=np.float64)
builder.f32[0] = 9.0
np.testing.assert_array_equal(array, [1.0, 2.0])
np.testing.assert_array_equal(np.frombuffer(reader.f32.to_float64_bytes(), dtype=np.float64), [9.0, 2.0])
@pytest.mark.parametrize(
("field", "values"),
[
("f32", []),
("f32", [0.0, -0.0, 1.25, float("inf"), float("nan")]),
("f64", [1e-300, 1e300, -0.0]),
("i64", [-(2**63), 0, 2**63 - 1]),
("u64", [0, 2**64 - 1]),
("text", ["a", "bc"]),
],
)
@pytest.mark.parametrize("dtype", [None, np.float64, np.complex128, "U32"])
@pytest.mark.parametrize("reader", [False, True])
def test_array_protocol(numeric_schema, field, values, dtype, reader):
builder = numeric_schema.new_message(**{field: values})
data = getattr(builder.as_reader() if reader else builder, field)
if field == "text" and dtype in (np.float64, np.complex128):
with pytest.raises(ValueError):
np.array(data, dtype=dtype)
return
expected = np.array(list(data), dtype=dtype)
actual = np.array(data, dtype=dtype)
np.testing.assert_array_equal(actual, expected)
assert actual.dtype == expected.dtype
assert actual.flags.writeable
assert actual.flags.owndata
with pytest.raises(ValueError):
np.asarray(data, copy=False)
def test_float32_array_ownership(numeric_schema):
builder = numeric_schema.new_message(f32=[1.25, -0.0, float("nan")])
actual = np.asarray(builder.as_reader().f32, dtype=np.float32)
expected = np.asarray(list(builder.f32), dtype=np.float32)
np.testing.assert_array_equal(actual, expected)
np.testing.assert_array_equal(np.signbit(actual), np.signbit(expected))
builder.f32[0] = 99.0
del builder
gc.collect()
actual[2] = 7.0
assert actual[0] == 1.25
assert actual[2] == 7.0
def test_array_cast_overflow(numeric_schema):
builder = numeric_schema.new_message(f64=[1e300])
with pytest.warns(RuntimeWarning, match="overflow"):
expected = np.array(list(builder.f64), dtype=np.float32)
with pytest.warns(RuntimeWarning, match="overflow"):
actual = np.array(builder.as_reader().f64, dtype=np.float32)
np.testing.assert_array_equal(actual, expected)
def test_nested_arrays(numeric_schema):
builder = numeric_schema.new_message(nested=[[1.0, 2.0], [3.0, 4.0]])
for message in (builder, builder.as_reader()):
actual = np.array(message.nested)
expected = np.array([list(row) for row in message.nested])
np.testing.assert_array_equal(actual, expected)
assert actual.dtype == expected.dtype

View File

@@ -0,0 +1,26 @@
# Copyright (c) 2026 IQ.Lvbs. All rights reserved.
import hashlib
import json
import struct
from pathlib import Path
import capnp
import pytest
ROOT = Path(__file__).parent
CASES = json.loads((ROOT / "serialization-reference.json").read_text())["cases"]
@pytest.mark.parametrize("case", CASES, ids=lambda case: str(case["size"]))
def test_serialization_identity(case):
schema = capnp.load(str(ROOT / "all_types.capnp")).TestAllTypes
size = case["size"]
text_size = case["text_size"]
message = schema.new_message(dataField=b"Z" * size, textField="iq" * text_size, float64List=[1.25, -0.0, 123.0])
data = message.to_bytes()
assert len(data) == case["bytes"]
assert struct.unpack_from("<I", data)[0] + 1 == case["segments"]
assert hashlib.sha256(data).hexdigest() == case["sha256"]
with schema.from_bytes(data) as reader:
assert reader.dataField == b"Z" * size
assert reader.textField == "iq" * text_size