From aa3f7305c1ec5c7b99beedc553cb414dff988f13 Mon Sep 17 00:00:00 2001 From: Shatrughna Upadhyay Date: Sat, 3 Oct 2026 22:02:34 -0700 Subject: [PATCH] GH-34485: [Python][FlightRPC] Add experimental chunked IPC exchange example --- python/examples/flight/chunked_ipc/README.md | 168 ++++++++ .../examples/flight/chunked_ipc/benchmark.py | 120 ++++++ .../flight/chunked_ipc/chunked_ipc.py | 347 ++++++++++++++++ .../flight/chunked_ipc/test_chunked_ipc.py | 384 ++++++++++++++++++ 4 files changed, 1019 insertions(+) create mode 100644 python/examples/flight/chunked_ipc/README.md create mode 100644 python/examples/flight/chunked_ipc/benchmark.py create mode 100644 python/examples/flight/chunked_ipc/chunked_ipc.py create mode 100644 python/examples/flight/chunked_ipc/test_chunked_ipc.py diff --git a/python/examples/flight/chunked_ipc/README.md b/python/examples/flight/chunked_ipc/README.md new file mode 100644 index 000000000000..0c6f9ad7fee2 --- /dev/null +++ b/python/examples/flight/chunked_ipc/README.md @@ -0,0 +1,168 @@ + + +# Experimental chunked IPC over Flight + +A single binary or string value can exceed a gRPC message limit. Splitting a +record batch into fewer rows cannot split that value. This example serializes +an Arrow IPC stream incrementally and carries its bytes in bounded, +metadata-only `DoExchange` messages. The server spools the upload to a temporary +file, then echoes the bytes; the client reconstructs and decodes the table. + +This is a private application protocol for exploring +[GH-34485](https://github.com/apache/arrow/issues/34485). It does not modify +`Flight.proto`, implement native fragmented `FlightData`, or make ordinary +Flight clients understand the fragments. Both endpoints must opt in. + +## Run + +Install a Flight-enabled PyArrow and pytest in your development environment. +From the Arrow repository root: + +```console +python -m pytest python/examples/flight/chunked_ipc/test_chunked_ipc.py -q +python python/examples/flight/chunked_ipc/benchmark.py --value-mib 16 --frame-kib 64 --iterations 3 +``` + +The tests and benchmark launch a real localhost Flight server on an ephemeral +port. The tests send 4 MiB values in one row while the client's gRPC send and +receive caps are 128 KiB. A control test shows an ordinary Flight write of that +row fails at the same cap. Binary, string, large binary, large string, +dictionary replacement, schema/field metadata and zero-batch streams are +covered. The benchmark accepts `--kind string` and emits JSON lines. + +From this directory, a minimal call looks like: + +```python +import pyarrow as pa +import pyarrow.flight as flight +from chunked_ipc import ChunkedEchoServer, Limits, exchange + +table = pa.table({"value": [b"x" * (4 * 1024 * 1024)]}) +limits = Limits(max_frame_bytes=16 * 1024, max_total_bytes=8 * 1024 * 1024) +with ChunkedEchoServer(limits=limits) as server: + with flight.FlightClient(("localhost", server.port)) as client: + result = exchange(client, table.schema, iter(table.to_batches()), + limits, flight.FlightCallOptions(timeout=30)) + assert result.table.equals(table) + print(result.sent, result.received) +``` + +The `batches` argument can be a generator. The sender asks for the next batch +after writing the previous batch, and closes the iterator after sending starts +if it exposes `close()`. A huge input value still has to exist in the source +batch; this is fragmentation of its serialized representation. + +## Example protocol + +1. `DoAction("arrow.example.chunked-ipc.v1.capabilities")` returns the protocol + name and server limits. The client takes the smaller local/server limits. +2. A command descriptor contains the protocol name, newline, and a JSON object + with the selected `max_frame_bytes` and `max_total_bytes`. The server checks + that both are within its policy before creating a temporary file. +3. Each `FlightData.app_metadata` contains a 21-byte header followed by payload. + No outer Flight schema or record batches are sent. +4. The client sends the IPC stream and an explicit `END`, then half-closes the + upload. Only after a clean incoming EOF does the server replay the stream. +5. The client validates response framing, requires a clean EOF, then decodes. + +| Header field | Encoding | Meaning | +| --- | --- | --- | +| Magic | 4 bytes, `CIP1` | Versioned example frame | +| Kind | 1 byte | `0` = DATA; `1` = END | +| Sequence | unsigned 64-bit, big endian | Starts at zero; increments per frame | +| Offset | unsigned 64-bit, big endian | Number of preceding IPC payload bytes | + +DATA payloads must be nonempty. END has no payload and its offset must equal +the received byte count. A frame's entire metadata field, including header, +must fit the selected frame cap. This is not the size of the enclosing protobuf +or gRPC message; transport configuration must leave overhead headroom. + +Out-of-order/duplicate frames, offset gaps/overlaps, unknown versions/kinds, +oversized frames, excessive total bytes, missing END, and frames after END are +rejected. A modern IPC end marker is also required. Decoding rejects trailing +bytes after the IPC stream. The echo server treats the IPC body as opaque; +semantic IPC validation occurs at the client. Framing has no checksum or +resume support and relies on the underlying reliable stream for delivery. + +## Memory, disk and copying + +The bounds are **per exchange**, not a process-wide admission policy: + +* The receiver's temporary file contains at most `max_total_bytes` encoded IPC + payload bytes. It does not allocate a Python buffer from an advertised total + length or retain a list of fragments. The default total cap is 256 MiB. +* Framing scratch space is proportional to `max_frame_bytes`: a sender copies + each payload slice into a header-plus-payload message, and a receiver borrows + an incoming metadata view while writing it to the file. Replay reads one + payload chunk at a time. Python file buffering and object overhead also exist. +* This accounting excludes the source batch, IPC/Python bridge buffers, incoming + Flight message allocations, queued native gRPC writes, OS file cache and the + complete decoded result. Returning a table requires memory proportional to + its contents. Native IPC/gRPC allocations are not bounded by these checks; + an oversized message reaches the transport before Python can reject it. +* The sender uses uncompressed IPC. The decoder is the ordinary IPC reader and + does not enforce a decompressed-allocation cap against hostile IPC input. + Use this example with trusted peers. Disk-full and decoding failures propagate + while context managers close temporary files. + +This trades zero-copy transfer for small messages and disk-backed staging. +Each direction copies frame payloads and writes/reads a temporary file. The +server waits for the complete upload before responding, so it demonstrates +bounded-message transport, not a low-latency pipelined service. Concurrent +exchanges can each consume their full cap. A deployment would need disk quotas, +admission control, authentication, deadlines and an agreed protocol. + +One IPC writer/reader spans the complete transfer. Schema metadata, dictionary +messages, dictionary replacements and record-batch boundaries remain inside +that stream; fragment boundaries can cut through any of them. A fragment is +not an independently readable Arrow record batch. + +## Negotiation, fallback and cancellation + +An unsupported capability action fails before consuming source batches. There +is no silent fallback: a caller may explicitly use normal Flight with smaller +batches where row slicing is sufficient, or report that a single value exceeds +the transport limit. A native Flight implementation would need a separately +agreed capability and wire format. Capability discovery here is application +convention, not authentication or a standardized Flight feature. + +On source or receive failure, the client cancels the reader and closes the +writer. Server and client reassembly files use context managers; cancellation, +deadline expiration, malformed input and normal completion release them. Tests +cover cancellation and timeout while the server is blocked reading a partial +upload. A process crash and arbitrary user generators blocking outside Flight +are outside this cleanup guarantee. + +## Interpreting the benchmark + +The reported IPC bytes and metadata frame counts are measured, but exclude +protobuf envelopes, descriptors, capability messages and transport framing. +The input table is created before timing/tracing. Timing covers capability +discovery, upload, server staging, download and decoding, with both peers in +one process. Equality validation runs after timing. + +`tracemalloc` reports Python allocations during that phase; it does not measure +native Arrow/gRPC allocations. Arrow's pool counters do not include Python +buffers that may back decoded arrays. RSS is the process **lifetime high-water +mark** and includes both peers, libraries, allocator retention and input. It is +not per-call retained memory or a gRPC-only peak. Run each configuration in a +fresh process for comparisons; later iterations share caches and high-water +marks. These local measurements establish feasibility, not production speed +or zero-copy equivalence. diff --git a/python/examples/flight/chunked_ipc/benchmark.py b/python/examples/flight/chunked_ipc/benchmark.py new file mode 100644 index 000000000000..b0a7d3354518 --- /dev/null +++ b/python/examples/flight/chunked_ipc/benchmark.py @@ -0,0 +1,120 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Reproducible localhost round-trip benchmark; prints one JSON row per run. + +RSS is a lifetime process high-water mark covering BOTH client and server. +tracemalloc measures Python allocations during the exchange, not native gRPC +or Arrow allocation. The input table is constructed before measurement. Use +a fresh process per configuration when comparing peak RSS. +""" + +import argparse +from dataclasses import asdict +import json +import platform +import sys +import time +import tracemalloc + +import pyarrow as pa +import pyarrow.flight as flight + +from chunked_ipc import ChunkedEchoServer, Limits, exchange + + +def peak_rss_bytes(): + if sys.platform not in ("darwin", "linux"): + return None + import resource + value = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return value if sys.platform == "darwin" else value * 1024 + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--value-mib", type=int, default=16) + parser.add_argument("--frame-kib", type=int, default=64) + parser.add_argument("--iterations", type=int, default=3) + parser.add_argument("--kind", choices=["binary", "string"], + default="binary") + args = parser.parse_args() + if min(args.value_mib, args.frame_kib, args.iterations) <= 0: + parser.error("sizes and iteration count must be positive") + value_bytes = args.value_mib * 1024 * 1024 + frame_bytes = args.frame_kib * 1024 + value = (b"x" if args.kind == "binary" else "x") * value_bytes + table = pa.table({"value": [value]}) + del value + limits = Limits(frame_bytes, value_bytes + 1024 * 1024) + # Leave explicit headroom for Flight/protobuf overhead above the metadata + # cap. A frame cap is not the size of the entire serialized gRPC message. + grpc_limit = frame_bytes + 64 * 1024 + grpc_options = [("grpc.max_send_message_length", grpc_limit), + ("grpc.max_receive_message_length", grpc_limit)] + with ChunkedEchoServer(limits=limits) as server: + with flight.FlightClient( + ("localhost", server.port), + generic_options=grpc_options) as client: + for iteration in range(args.iterations): + rss_before = peak_rss_bytes() + arrow_before = pa.total_allocated_bytes() + tracemalloc.start() + start = time.perf_counter() + result = exchange( + client, table.schema, iter(table.to_batches()), limits, + flight.FlightCallOptions(timeout=120)) + elapsed = time.perf_counter() - start + python_current, python_peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + rss_after = peak_rss_bytes() + arrow_after = pa.total_allocated_bytes() + assert result.table.equals(table) + print(json.dumps({ + "iteration": iteration, + "python": platform.python_version(), + "pyarrow": pa.__version__, + "platform": platform.platform(), + "kind": args.kind, + "rows": table.num_rows, + "input_arrow_bytes": table.nbytes, + "limits": asdict(limits), + "grpc_message_limit_bytes": grpc_limit, + "sent": asdict(result.sent), + "received": asdict(result.received), + "elapsed_seconds": elapsed, + "roundtrip_ipc_mib_per_second": ( + 2 * result.sent.ipc_bytes / 1024**2 / elapsed), + "python_tracemalloc_current_bytes": python_current, + "python_tracemalloc_peak_bytes": python_peak, + "process_peak_rss_before_bytes": rss_before, + "process_peak_rss_after_bytes": rss_after, + "arrow_pool_live_before_bytes": arrow_before, + "arrow_pool_live_after_bytes": arrow_after, + "arrow_pool_lifetime_peak_bytes": ( + pa.default_memory_pool().max_memory()), + "memory_scope": ( + "combined client/server process; " + "input precedes tracing; " + "RSS and Arrow peaks are lifetime high-water marks; " + "tracemalloc excludes native gRPC/Arrow allocations"), + }), flush=True) + del result + + +if __name__ == "__main__": + main() diff --git a/python/examples/flight/chunked_ipc/chunked_ipc.py b/python/examples/flight/chunked_ipc/chunked_ipc.py new file mode 100644 index 000000000000..1792f26155df --- /dev/null +++ b/python/examples/flight/chunked_ipc/chunked_ipc.py @@ -0,0 +1,347 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Experimental IPC fragmentation over metadata-only Flight DoExchange. + +This is an application protocol, not a change to Flight or its protobuf schema. +Both peers must opt in. Reassembly uses a temporary file; decoded Arrow data, +source batches, native IPC/gRPC allocations and OS file cache are not covered +by the bounded frame scratch space. See README.md for the protocol and limits. +""" + +import contextlib +from dataclasses import asdict, dataclass +import io +import json +import struct +import tempfile + +import pyarrow as pa +import pyarrow.flight as flight + + +PROTOCOL = "arrow.example.chunked-ipc.v1" +CAPABILITIES_ACTION = PROTOCOL + ".capabilities" +MAGIC = b"CIP1" +DATA = 0 +END = 1 +HEADER = struct.Struct("!4sBQQ") # magic, kind, sequence, byte offset +IPC_EOS = b"\xff\xff\xff\xff\x00\x00\x00\x00" + + +class ProtocolError(ValueError): + """An invalid or unsupported example protocol message.""" + + +@dataclass(frozen=True) +class Limits: + """Limits apply to metadata bytes and encoded IPC bytes, respectively.""" + + max_frame_bytes: int = 64 * 1024 + max_total_bytes: int = 256 * 1024 * 1024 + + def __post_init__(self): + for value in (self.max_frame_bytes, self.max_total_bytes): + if type(value) is not int or value <= 0 or value >= 2**63: + raise ProtocolError("limits must be positive 63-bit integers") + if self.max_frame_bytes <= HEADER.size: + raise ProtocolError("frame limit must leave room for a payload") + + +@dataclass +class Stats: + ipc_bytes: int = 0 + metadata_bytes: int = 0 + frames: int = 0 + largest_frame_bytes: int = 0 + + def add(self, frame_bytes, payload_bytes): + self.ipc_bytes += payload_bytes + self.metadata_bytes += frame_bytes + self.frames += 1 + self.largest_frame_bytes = max(self.largest_frame_bytes, frame_bytes) + + +class FrameSink(io.RawIOBase): + """Serialize one bounded metadata frame at a time, without a frame list. + + The callback must consume a frame synchronously. Flight may make its own + copies or queue writes; this object does not control transport allocation. + """ + + def __init__(self, send_metadata, limits): + super().__init__() + self.send_metadata = send_metadata + self.limits = limits + self.stats = Stats() + self.finished = False + + def writable(self): + return True + + def write(self, data): + if self.closed or self.finished: + raise ProtocolError("write after end of stream") + view = memoryview(data).cast("B") + if self.stats.ipc_bytes + len(view) > self.limits.max_total_bytes: + raise ProtocolError("encoded IPC stream exceeds total byte limit") + payload_limit = self.limits.max_frame_bytes - HEADER.size + for start in range(0, len(view), payload_limit): + part = view[start:start + payload_limit] + header = HEADER.pack(MAGIC, DATA, self.stats.frames, + self.stats.ipc_bytes) + frame = header + part.tobytes() + self.send_metadata(frame) + self.stats.add(len(frame), len(part)) + return len(view) + + def finish(self): + if self.closed or self.finished: + raise ProtocolError("duplicate end of stream") + frame = HEADER.pack(MAGIC, END, self.stats.frames, + self.stats.ipc_bytes) + self.send_metadata(frame) + self.stats.add(len(frame), 0) + self.finished = True + + +def send_batches(send_metadata, schema, batches, limits): + """Write a single uncompressed IPC stream from a lazy batch iterable. + + The iterable is owned for this call and closed on failure or completion + when it exposes close(). An input batch, including one huge row, is still + held in full by the caller/IPC writer while its bytes are fragmented. + """ + batches = iter(batches) + try: + with FrameSink(send_metadata, limits) as sink: + options = pa.ipc.IpcWriteOptions( + use_legacy_format=False, compression=None) + with pa.ipc.new_stream(sink, schema, options=options) as writer: + for batch in batches: + writer.write_batch(batch) + sink.finish() + return sink.stats + finally: + close = getattr(batches, "close", None) + if close is not None: + close() + + +class Reassembler: + """Validate framing and spool a capped encoded stream to a temporary file. + + No advertised size causes an allocation. accept() holds a borrowed frame + view while writing it to disk. IPC decoding can allocate memory + proportional to the complete output, separately in read_table(). + """ + + def __init__(self, limits, temp_dir=None): + self.limits = limits + self.stats = Stats() + self.ended = False + self.complete = False + self.file = tempfile.TemporaryFile(dir=temp_dir) + + def __enter__(self): + return self + + def __exit__(self, *args): + self.file.close() + + def accept(self, metadata): + if self.ended or self.complete: + raise ProtocolError("frame after END") + frame = memoryview(metadata).cast("B") + if len(frame) > self.limits.max_frame_bytes: + raise ProtocolError("frame exceeds metadata byte limit") + if len(frame) < HEADER.size: + raise ProtocolError("truncated frame header") + magic, kind, sequence, offset = HEADER.unpack_from(frame) + if magic != MAGIC or kind not in (DATA, END): + raise ProtocolError("unsupported frame magic or kind") + if sequence != self.stats.frames: + raise ProtocolError("out-of-order or duplicate frame") + if offset != self.stats.ipc_bytes: + raise ProtocolError("noncontiguous byte offset") + payload = frame[HEADER.size:] + if kind == END: + if payload: + raise ProtocolError("END must not contain a payload") + self.ended = True + else: + if not payload: + raise ProtocolError("DATA must contain a payload") + if offset + len(payload) > self.limits.max_total_bytes: + raise ProtocolError( + "encoded IPC stream exceeds total byte limit") + self.file.write(payload) + self.stats.add(len(frame), len(payload)) + + def finish(self): + """Call only after the Flight stream has reached a clean EOF.""" + if not self.ended: + raise ProtocolError("truncated stream: missing END") + if self.stats.ipc_bytes < len(IPC_EOS): + raise ProtocolError("truncated IPC stream") + self.file.seek(-len(IPC_EOS), io.SEEK_END) + if self.file.read() != IPC_EOS: + raise ProtocolError("missing IPC end-of-stream marker") + self.complete = True + + def read_table(self): + if not self.complete: + raise ProtocolError("stream has not reached a clean EOF") + self.file.seek(0) + with pa.ipc.open_stream(self.file) as reader: + # Compression is deliberately not part of this example protocol. + # A total encoded-byte cap is not a decoded-allocation guarantee. + table = reader.read_all() + if self.file.tell() != self.stats.ipc_bytes: + raise ProtocolError("trailing bytes after IPC stream") + return table + + def replay(self, send_metadata): + """Echo opaque IPC bytes. The client performs semantic IPC decoding.""" + if not self.complete: + raise ProtocolError("stream has not reached a clean EOF") + self.file.seek(0) + with FrameSink(send_metadata, self.limits) as sink: + while chunk := self.file.read( + self.limits.max_frame_bytes - HEADER.size): + sink.write(chunk) + sink.finish() + return sink.stats + + +def _parse_limits(value): + if not isinstance(value, dict) or set(value) != { + "max_frame_bytes", "max_total_bytes"}: + raise ProtocolError("invalid limits object") + return Limits(**value) + + +def exchange_descriptor(limits): + command = PROTOCOL.encode() + b"\n" + json.dumps(asdict(limits)).encode() + return flight.FlightDescriptor.for_command(command) + + +def negotiate(client, requested, options=None): + """Require opt-in before upload; never fall back automatically.""" + try: + results = iter(client.do_action( + flight.Action(CAPABILITIES_ACTION, b""), options=options)) + first = next(results, None) + if first is None or len(first.body) > 1024: + raise ProtocolError("invalid capability response") + document = json.loads(first.body.to_pybytes()) + if next(results, None) is not None: + raise ProtocolError("expected one capability response") + except pa.ArrowNotImplementedError as exc: + raise ProtocolError("server does not support " + PROTOCOL) from exc + if not isinstance(document, dict) or document.get("protocol") != PROTOCOL: + raise ProtocolError("server does not support " + PROTOCOL) + offered = _parse_limits(document.get("limits")) + return Limits(min(requested.max_frame_bytes, offered.max_frame_bytes), + min(requested.max_total_bytes, offered.max_total_bytes)) + + +def _receive(reader, receiver): + for chunk in reader: + if chunk.data is not None or chunk.app_metadata is None: + raise ProtocolError("expected metadata-only FlightData") + receiver.accept(chunk.app_metadata) + receiver.finish() + + +class ChunkedEchoServer(flight.FlightServerBase): + """Disk-backed echo endpoint for trusted localhost experiments.""" + + def __init__(self, location=None, limits=None, temp_dir=None): + self.limits = limits or Limits() + self.temp_dir = temp_dir + super().__init__(location) + + def list_actions(self, context): + return [(CAPABILITIES_ACTION, "Experimental chunked IPC limits")] + + def do_action(self, context, action): + if action.type != CAPABILITIES_ACTION: + raise pa.ArrowNotImplementedError("unsupported action") + document = {"protocol": PROTOCOL, "limits": asdict(self.limits)} + yield flight.Result(json.dumps(document).encode()) + + def do_exchange(self, context, descriptor, reader, writer): + prefix = PROTOCOL.encode() + b"\n" + if (descriptor.descriptor_type != flight.DescriptorType.CMD or + len(descriptor.command) > 512 or + not descriptor.command.startswith(prefix)): + raise pa.ArrowNotImplementedError("unsupported exchange command") + limits = _parse_limits(json.loads(descriptor.command[len(prefix):])) + if (limits.max_frame_bytes > self.limits.max_frame_bytes or + limits.max_total_bytes > self.limits.max_total_bytes): + raise ProtocolError("requested limits exceed server capabilities") + + def send(metadata): + if context.is_cancelled(): + raise flight.FlightCancelledError("exchange cancelled") + writer.write_metadata(metadata) + + with Reassembler(limits, self.temp_dir) as receiver: + for chunk in reader: + if context.is_cancelled(): + raise flight.FlightCancelledError("exchange cancelled") + if chunk.data is not None or chunk.app_metadata is None: + raise ProtocolError("expected metadata-only FlightData") + receiver.accept(chunk.app_metadata) + receiver.finish() + receiver.replay(send) + + +@dataclass +class ExchangeResult: + table: pa.Table + sent: Stats + received: Stats + limits: Limits + + +def exchange(client, schema, batches, limits=None, options=None, + temp_dir=None): + """Negotiate, send lazy batches, half-close, then receive the echoed table. + + The peer must drain the complete upload before replying, as this helper is + deliberately sequential. Both client and server temporary files close on + success, malformed input, source failure, timeout, and cancellation. + """ + limits = negotiate(client, limits or Limits(), options) + writer, reader = client.do_exchange(exchange_descriptor(limits), + options=options) + try: + sent = send_batches(writer.write_metadata, schema, batches, limits) + writer.done_writing() + with Reassembler(limits, temp_dir) as receiver: + _receive(reader, receiver) + table = receiver.read_table() + received = receiver.stats + writer.close() + return ExchangeResult(table, sent, received, limits) + except BaseException: + reader.cancel() + with contextlib.suppress(Exception): + writer.close() + raise diff --git a/python/examples/flight/chunked_ipc/test_chunked_ipc.py b/python/examples/flight/chunked_ipc/test_chunked_ipc.py new file mode 100644 index 000000000000..8ab2a02e5e38 --- /dev/null +++ b/python/examples/flight/chunked_ipc/test_chunked_ipc.py @@ -0,0 +1,384 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Pure framing tests and real localhost Flight RPC integration tests.""" + +import contextlib +import dataclasses +import threading + +import pyarrow as pa +import pyarrow.flight as flight +import pytest + +import chunked_ipc as example + + +def frame(kind=example.DATA, sequence=0, offset=0, payload=b"x", + magic=example.MAGIC): + return example.HEADER.pack(magic, kind, sequence, offset) + payload + + +def encoded(table, limits=None): + frames = [] + stats = example.send_batches( + frames.append, table.schema, table.to_batches(), + limits or example.Limits()) + return frames, stats + + +def decode(frames, limits): + with example.Reassembler(limits) as receiver: + for metadata in frames: + receiver.accept(metadata) + receiver.finish() + return receiver.read_table() + + +@pytest.mark.parametrize("kwargs", [ + {"max_frame_bytes": 0}, {"max_frame_bytes": example.HEADER.size}, + {"max_frame_bytes": True}, {"max_total_bytes": -1}, + {"max_total_bytes": 2**63}, {"max_total_bytes": 1.5}, +]) +def test_invalid_limits(kwargs): + with pytest.raises(example.ProtocolError): + example.Limits(**kwargs) + + +def test_exact_frame_and_total_bounds(): + table = pa.table({"value": [b"x" * 10000]}) + frames, initial = encoded(table, example.Limits(128, 20000)) + limits = example.Limits(128, initial.ipc_bytes) + frames, stats = encoded(table, limits) + assert max(map(len, frames)) == limits.max_frame_bytes + assert stats.metadata_bytes == sum(map(len, frames)) + assert stats.ipc_bytes == stats.metadata_bytes - len(frames) * 21 + assert stats.frames == len(frames) + assert decode(frames, limits).equals(table) + with pytest.raises(example.ProtocolError, match="total byte limit"): + encoded(table, example.Limits(128, initial.ipc_bytes - 1)) + with pytest.raises(example.ProtocolError, match="total byte limit"): + decode(frames, example.Limits(128, initial.ipc_bytes - 1)) + + +@pytest.mark.parametrize("frames, message", [ + ([b"short"], "truncated frame header"), + ([frame(magic=b"BAD!")], "magic or kind"), + ([frame(kind=2)], "magic or kind"), + ([frame(sequence=1)], "out-of-order"), + ([frame(offset=1)], "byte offset"), + ([frame(payload=b"")], "DATA must contain"), + ([frame(payload=b"x" * 200)], "frame exceeds"), + ([frame(), frame()], "out-of-order"), + ([frame(), frame(sequence=1, offset=0)], "byte offset"), + ([frame(), frame(sequence=2, offset=1)], "out-of-order"), + ([frame(kind=example.END)], "END must not contain"), + ([frame(kind=example.END, payload=b""), frame()], "frame after END"), +]) +def test_malformed_frames(frames, message): + with example.Reassembler(example.Limits(128, 10000)) as receiver: + with pytest.raises(example.ProtocolError, match=message): + for metadata in frames: + receiver.accept(metadata) + assert receiver.file.closed + + +def test_truncated_flight_and_ipc_streams(): + table = pa.table({"value": [1, 2]}) + frames, _ = encoded(table) + with pytest.raises(example.ProtocolError, match="missing END"): + decode(frames[:-1], example.Limits()) + with pytest.raises(example.ProtocolError, match="truncated IPC"): + decode([frame(kind=example.END, payload=b"")], example.Limits()) + raw = b"".join(f[example.HEADER.size:] for f in frames[:-1]) + # A syntactically framed END is insufficient if IPC itself was truncated. + for damaged in (raw[:-8], b"not an IPC stream" + example.IPC_EOS, + raw + example.IPC_EOS): + frames = [frame(payload=damaged), + frame(kind=example.END, sequence=1, + offset=len(damaged), payload=b"")] + with pytest.raises((example.ProtocolError, pa.ArrowInvalid, + pa.ArrowIOError)): + decode(frames, example.Limits()) + + +def test_no_decode_or_replay_before_eof(): + frames, _ = encoded(pa.table({"v": [1]})) + with example.Reassembler(example.Limits()) as receiver: + for metadata in frames: + receiver.accept(metadata) + with pytest.raises(example.ProtocolError, match="clean EOF"): + receiver.read_table() + with pytest.raises(example.ProtocolError, match="clean EOF"): + receiver.replay(lambda metadata: None) + + +def test_sink_rejects_writes_after_end(): + with example.FrameSink(lambda metadata: None, example.Limits()) as sink: + sink.finish() + with pytest.raises(example.ProtocolError): + sink.write(b"x") + with pytest.raises(example.ProtocolError): + sink.finish() + + +def test_lazy_batches_and_source_cleanup_on_send_failure(): + events = [] + schema = pa.schema([("v", pa.int64())]) + + def batches(): + try: + events.append("first") + yield pa.record_batch([[1]], schema=schema) + # The previous batch was sent before the next batch was requested. + assert "send" in events + events.append("second") + yield pa.record_batch([[2]], schema=schema) + finally: + events.append("closed") + + frames = [] + + def send(metadata): + events.append("send") + frames.append(metadata) + + example.send_batches(send, schema, batches(), example.Limits()) + assert events.count("closed") == 1 + assert decode(frames, example.Limits()).column(0).to_pylist() == [1, 2] + + events.clear() + + def fail(metadata): + raise RuntimeError("send failed") + + with pytest.raises(RuntimeError, match="send failed"): + example.send_batches(fail, schema, batches(), example.Limits()) + assert events[-1] == "closed" + + +def client_for(server): + return flight.FlightClient( + ("localhost", server.port), generic_options=[ + ("grpc.max_send_message_length", 128 * 1024), + ("grpc.max_receive_message_length", 128 * 1024)]) + + +@pytest.mark.parametrize("arrow_type, value", [ + (pa.binary(), b"x" * (4 * 1024 * 1024)), + (pa.string(), "y" * (4 * 1024 * 1024)), + (pa.large_binary(), b"z" * (4 * 1024 * 1024)), + (pa.large_string(), "w" * (4 * 1024 * 1024)), +]) +def test_large_single_row_over_real_rpc(arrow_type, value): + table = pa.table({"value": pa.array([value], type=arrow_type)}) + # Slicing cannot shrink the one value; fragment its encoded bytes instead. + assert table.slice(0, 1).nbytes > 128 * 1024 + limits = example.Limits(16 * 1024, 8 * 1024 * 1024) + with example.ChunkedEchoServer(limits=limits) as server: + with client_for(server) as client: + result = example.exchange( + client, table.schema, iter(table.to_batches()), limits, + flight.FlightCallOptions(timeout=20)) + assert result.table.equals(table, check_metadata=True) + assert result.sent.ipc_bytes == result.received.ipc_bytes + assert result.sent.frames > 256 + assert result.sent.largest_frame_bytes <= limits.max_frame_bytes + assert result.received.largest_frame_bytes <= limits.max_frame_bytes + + +class PlainEchoServer(flight.FlightServerBase): + def do_exchange(self, context, descriptor, reader, writer): + writer.begin(reader.schema) + for chunk in reader: + writer.write_batch(chunk.data) + + +def test_ordinary_single_row_write_hits_transport_limit(): + batch = pa.record_batch([pa.array([b"x" * (4 * 1024 * 1024)])], + names=["value"]) + with PlainEchoServer() as server, client_for(server) as client: + writer, reader = client.do_exchange( + flight.FlightDescriptor.for_command(b"echo"), + options=flight.FlightCallOptions(timeout=10)) + try: + with pytest.raises((pa.ArrowInvalid, flight.FlightError), + match="[Ll]arger than max|[Ee]xceeds|[Ss]ize"): + writer.begin(batch.schema) + writer.write_batch(batch) + writer.done_writing() + reader.read_all() + finally: + reader.cancel() + with contextlib.suppress(Exception): + writer.close() + + +def test_schema_dictionary_replacements_and_empty_table(): + dtype = pa.dictionary(pa.int8(), pa.string()) + schema = pa.schema([pa.field("value", dtype, metadata={b"field": b"yes"})], + metadata={b"schema": b"preserved"}) + batches = [pa.record_batch([pa.array(values, type=dtype)], schema=schema) + for values in (["a", "b", None], ["b", "c"])] + with example.ChunkedEchoServer() as server, client_for(server) as client: + result = example.exchange(client, schema, iter(batches), + example.Limits(1024, 100000)) + assert result.table.equals(pa.Table.from_batches(batches), + check_metadata=True) + empty = example.exchange(client, schema, iter(()), + example.Limits(1024, 100000)) + assert empty.table.num_rows == 0 + assert empty.table.schema.equals(schema, check_metadata=True) + + +def test_negotiation_uses_smaller_limits(): + server_limits = example.Limits(2048, 8192) + requested = example.Limits(1024, 16384) + table = pa.table({"value": [b"x" * 4000]}) + with example.ChunkedEchoServer(limits=server_limits) as server: + with client_for(server) as client: + result = example.exchange(client, table.schema, table.to_batches(), + requested) + assert result.limits == example.Limits(1024, 8192) + assert result.received.largest_frame_bytes <= 1024 + + +def test_unsupported_server_fails_before_source_consumption(): + consumed = [] + + def batches(): + consumed.append(True) + yield pa.record_batch([[1]], names=["v"]) + + with flight.FlightServerBase() as server, client_for(server) as client: + with pytest.raises(example.ProtocolError, match="does not support"): + example.exchange(client, pa.schema([("v", pa.int64())]), batches()) + assert not consumed + + +def tracked_tempfiles(monkeypatch): + files = [] + created = threading.Event() + original = example.tempfile.TemporaryFile + + def factory(*args, **kwargs): + file = original(*args, **kwargs) + files.append(file) + created.set() + return file + + monkeypatch.setattr(example.tempfile, "TemporaryFile", factory) + return files, created + + +def test_cancel_partial_rpc_closes_server_spool(monkeypatch): + files, created = tracked_tempfiles(monkeypatch) + with example.ChunkedEchoServer() as server, client_for(server) as client: + writer, reader = client.do_exchange( + example.exchange_descriptor(example.Limits()), + options=flight.FlightCallOptions(timeout=10)) + writer.write_metadata(frame(payload=b"partial IPC bytes")) + assert created.wait(5) + reader.cancel() + with contextlib.suppress(Exception): + writer.close() + assert files and all(file.closed for file in files) + + +def test_rpc_truncation_and_source_failure_close_spools(monkeypatch): + files, _ = tracked_tempfiles(monkeypatch) + table = pa.table({"value": [1, 2]}) + frames, _ = encoded(table) + with example.ChunkedEchoServer() as server, client_for(server) as client: + writer, reader = client.do_exchange( + example.exchange_descriptor(example.Limits()), + options=flight.FlightCallOptions(timeout=10)) + try: + for metadata in frames[:-1]: + writer.write_metadata(metadata) + writer.done_writing() + with pytest.raises(pa.ArrowInvalid, match="missing END"): + list(reader) + finally: + reader.cancel() + with contextlib.suppress(Exception): + writer.close() + + def broken_source(): + yield from table.to_batches() + raise RuntimeError("source failed") + + with pytest.raises(RuntimeError, match="source failed"): + example.exchange(client, table.schema, broken_source(), + options=flight.FlightCallOptions(timeout=10)) + assert files and all(file.closed for file in files) + + +@pytest.mark.parametrize("metadata, message", [ + (frame(sequence=1), "out-of-order"), + (frame(offset=100), "byte offset"), + (frame(payload=b"x" * 2048), "frame exceeds"), +]) +def test_rpc_rejects_bad_frames(monkeypatch, metadata, message): + files, _ = tracked_tempfiles(monkeypatch) + limits = example.Limits(1024, 4096) + with example.ChunkedEchoServer(limits=limits) as server: + with client_for(server) as client: + writer, reader = client.do_exchange( + example.exchange_descriptor(limits), + options=flight.FlightCallOptions(timeout=10)) + try: + writer.write_metadata(metadata) + writer.done_writing() + with pytest.raises(pa.ArrowInvalid, match=message): + list(reader) + finally: + reader.cancel() + with contextlib.suppress(Exception): + writer.close() + assert files and all(file.closed for file in files) + + +def test_rpc_timeout_closes_server_spool(monkeypatch): + files, created = tracked_tempfiles(monkeypatch) + with example.ChunkedEchoServer() as server, client_for(server) as client: + writer, reader = client.do_exchange( + example.exchange_descriptor(example.Limits()), + options=flight.FlightCallOptions(timeout=0.5)) + try: + writer.write_metadata(frame(payload=b"incomplete")) + assert created.wait(5) + # Leave the upload open. The server must release its temporary file + # when the client deadline cancels a blocked incoming read. + with pytest.raises(flight.FlightTimedOutError): + reader.read_chunk() + finally: + reader.cancel() + with contextlib.suppress(Exception): + writer.close() + assert files and all(file.closed for file in files) + + +@pytest.mark.parametrize("change", [ + {"max_frame_bytes": 1}, {"max_frame_bytes": True}, + {"max_total_bytes": 0}, {"unexpected": 1}, +]) +def test_capability_limits_are_validated(change): + limits = dataclasses.asdict(example.Limits()) + limits.update(change) + with pytest.raises(example.ProtocolError): + example._parse_limits(limits)