flight

Arrow Flight utilities.

ChainedFlightReader

ChainedFlightReader(endpoints, initial_url, metadata=None)

Bases: StreamReader

Sequential reader for time-ordered endpoint slices.

Groups endpoints by time slice (based on the start time in their tickets), orders slices chronologically, and reads each slice to completion before starting the next. Concurrent endpoints within the same slice are handled by a MultiFlightReader.

Source code in arrakis/flight.py

def __init__(
    self,
    endpoints: list[flight.FlightEndpoint],
    initial_url: str | flight.FlightClient,
    metadata: dict[str, Channel] | None = None,
):
    self.initial_url = initial_url
    self.metadata = metadata

    # group endpoints by start time
    slices: dict[int | None, list[flight.FlightEndpoint]] = {}
    for endpoint in endpoints:
        start = _endpoint_start_time(endpoint)
        slices.setdefault(start, []).append(endpoint)

    # order slices chronologically (None sorts first as live/unbounded)
    self._slices = sorted(
        slices.items(),
        key=lambda item: item[0] if item[0] is not None else float("inf"),
    )

    # single slice: delegate directly to one MultiFlightReader
    # and avoid the chaining/busy-wait overhead entirely
    if len(self._slices) == 1:
        _, endpoints = self._slices[0]
        self._delegate = self._make_reader(
            endpoints, self.initial_url, self.metadata
        )
    else:
        self._delegate = None

    self._current_reader: MultiFlightReader | None = None
    self._slice_index = 0

MultiFlightReader

MultiFlightReader(endpoints, initial_url, metadata=None)

Bases: StreamReader

Multi-threaded Arrow Flight endpoint stream iterator context manager

Given a list of endpoints, connect to all of them in parallel and stream data from them all interleaved.

Source code in arrakis/flight.py

def __init__(
    self,
    endpoints: list[flight.FlightEndpoint],
    initial_url: str | flight.FlightClient,
    metadata: dict[str, Channel] | None = None,
):
    """initialize with list of endpoints and an reusable flight client"""
    self.endpoints = endpoints
    self.initial_url = initial_url
    self.metadata = metadata

    self.queue: queue.SimpleQueue = queue.SimpleQueue()
    # FIXME: this quit event could be replaced with the
    # Queue.shutdown method after python 3.13
    self.quit_event: threading.Event = threading.Event()
    self.done_events = [threading.Event() for _ in self.endpoints]
    self.executor = concurrent.futures.ThreadPoolExecutor(
        max_workers=len(self.endpoints),
    )
    self.futures: list[concurrent.futures.Future] = []

close

close()

close all streams

Source code in arrakis/flight.py

def close(self):
    """close all streams"""
    self.quit_event.set()
    if self.futures is not None:
        for f in self.futures:
            self._reraise(f)

    self.executor.shutdown(cancel_futures=True)
    self.futures = []

done

done()

returns True when all threads are done and the queue is empty

Source code in arrakis/flight.py

def done(self) -> bool:
    """returns True when all threads are done and the queue is empty"""
    return all(e.is_set() for e in self.done_events) and self.queue.empty()

read

read(*, convert_blocks=False)

read all available data from all streams

Yields:

Type	Description
tuple of (stream_id, data)

Source code in arrakis/flight.py

def read(
    self,
    *,
    convert_blocks: bool = False,
) -> Generator[
    tuple[str, Any],
    None,
    None,
]:
    """read all available data from all streams

    Yields
    ------
    tuple of (stream_id, data)

    """
    if convert_blocks:
        assert self.metadata, (
            "Must be initialized with channel metadata to convert to blocks."
        )
    while True:
        try:
            stream_id, batch = self.queue.get(block=False)
            if convert_blocks:
                assert self.metadata  # needed for type checking consistency
                data = SeriesBlock.from_column_batch(
                    batch,
                    self.metadata,
                )
            else:
                data = batch
            yield stream_id, data
        except queue.Empty:
            break
    self._raise_failed_endpoints()

RequestValidator

RequestValidator()

A validator for JSON-encoded requests.

Source code in arrakis/flight.py

def __init__(self) -> None:
    self._validators: dict[RequestType, jsonschema.Draft7Validator] = {}

    # load generic descriptor schema
    schema = arrakis_schema.load_schema("descriptor.json")
    self._generic_validator = jsonschema.Draft7Validator(schema)

validate

validate(payload)

Validate a JSON-encoded request.

Parameters:

Name	Type	Description	Default
payload	Request	A dictionary with a 'request' and an 'args' key encoding the given Flight request.	required

Raises:

Type	Description
ValidationError	If the request does not match the expected schema.

Source code in arrakis/flight.py

def validate(self, payload: Request) -> None:
    """Validate a JSON-encoded request.

    Parameters
    ----------
    payload : Request
        A dictionary with a 'request' and an 'args' key encoding
        the given Flight request.

    Raises
    ------
    ValidationError
        If the request does not match the expected schema.

    """
    self._generic_validator.validate(payload)
    request = RequestType[payload["request"]]

    # load schema on demand
    if request not in self._validators:
        schema = arrakis_schema.load_schema(f"{request.name.lower()}.json")
        self._validators[request] = jsonschema.Draft7Validator(schema)

    self._validators[request].validate(payload)

connect

connect(location, **kwargs)

Connect to a Flight server with arrakis default channel options.

A drop-in wrapper for :func:pyarrow.flight.connect that applies KEEPALIVE_OPTIONS unless the caller passes its own generic_options.

Source code in arrakis/flight.py

def connect(location, **kwargs) -> flight.FlightClient:
    """Connect to a Flight server with arrakis default channel options.

    A drop-in wrapper for :func:`pyarrow.flight.connect` that applies
    ``KEEPALIVE_OPTIONS`` unless the caller passes its own
    ``generic_options``.
    """
    kwargs.setdefault("generic_options", list(KEEPALIVE_OPTIONS))
    return flight.connect(location, **kwargs)

create_command

create_command(request_type, *, validator, **kwargs)

Create a Flight command containing a JSON-encoded request.

Parameters:

Name	Type	Description	Default
request_type	RequestType	The type of request.	required
validator	RequestValidator	A validator to validate that the command matches the expected schema.	required
**kwargs	dict	Extra arguments corresponding to the specific request.	{}

Returns:

Type	Description
bytes	The JSON-encoded request.

Raises:

Type	Description
ValidationError	If the request does not match the expected schema.

Source code in arrakis/flight.py

def create_command(
    request_type: RequestType, *, validator: RequestValidator, **kwargs
) -> bytes:
    """Create a Flight command containing a JSON-encoded request.

    Parameters
    ----------
    request_type : RequestType
        The type of request.
    validator : RequestValidator
        A validator to validate that the command matches the expected schema.
    **kwargs : dict, optional
        Extra arguments corresponding to the specific request.

    Returns
    -------
    bytes
        The JSON-encoded request.

    Raises
    ------
    ValidationError
        If the request does not match the expected schema.

    """
    cmd: Request = {
        "request": request_type.name,
        "args": kwargs,
    }
    validator.validate(cmd)
    return json.dumps(cmd).encode("utf-8")

create_descriptor

create_descriptor(request_type, *, validator, **kwargs)

Create a Flight descriptor given a request.

Parameters:

Name	Type	Description	Default
request_type	RequestType	The type of request.	required
validator	RequestValidator	A validator to validate that the command matches the expected schema.	required
**kwargs	dict	Extra arguments corresponding to the specific request.	{}

Returns:

Type	Description
FlightDescriptor	A Flight Descriptor containing the request.

Raises:

Type	Description
ValidationError	If the request does not match the expected schema.

Source code in arrakis/flight.py

def create_descriptor(
    request_type: RequestType, *, validator: RequestValidator, **kwargs
) -> flight.FlightDescriptor:
    """Create a Flight descriptor given a request.

    Parameters
    ----------
    request_type : RequestType
        The type of request.
    validator : RequestValidator
        A validator to validate that the command matches the expected schema.
    **kwargs : dict, optional
        Extra arguments corresponding to the specific request.

    Returns
    -------
    flight.FlightDescriptor
        A Flight Descriptor containing the request.

    Raises
    ------
    ValidationError
        If the request does not match the expected schema.

    """
    cmd = create_command(request_type, validator=validator, **kwargs)
    return flight.FlightDescriptor.for_command(cmd)

parse_command

parse_command(cmd, *, validator)

Parse a Flight command into a request.

Parameters:

Name	Type	Description	Default
cmd	bytes	The JSON-encoded request.	required
validator	RequestValidator	A validator to validate that the command matches the expected schema.	required

Returns:

Name	Type	Description
request_type	RequestType	The type of request.
kwargs	dict	Arguments corresponding to the specific request.

Raises:

Type	Description
JSONDecodeError	If the command does not decode to valid JSON.
ValidationError	If the request does not match the expected schema.

Source code in arrakis/flight.py

def parse_command(
    cmd: bytes, *, validator: RequestValidator
) -> tuple[RequestType, dict]:
    """Parse a Flight command into a request.

    Parameters
    ----------
    cmd : bytes
        The JSON-encoded request.
    validator : RequestValidator
        A validator to validate that the command matches the expected schema.

    Returns
    -------
    request_type : RequestType
        The type of request.
    kwargs : dict
        Arguments corresponding to the specific request.

    Raises
    ------
    JSONDecodeError
        If the command does not decode to valid JSON.
    ValidationError
        If the request does not match the expected schema.

    """
    try:
        parsed = json.loads(cmd.decode("utf-8"))
    except json.JSONDecodeError as e:
        msg = "Command does not decode to valid JSON"
        raise json.JSONDecodeError(msg, e.doc, e.pos) from e
    else:
        validator.validate(parsed)
        return RequestType[parsed["request"]], parsed["args"]