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:
            # re-raise exceptions to the client, returning
            # user-friendly Flight-specific errors when relevant
            try:
                f.result()
            except flight.FlightError as e:
                # NOTE: this strips the original message of everything
                # besides the original error message raised by the server
                msg = e.args[0].partition(" Detail:")[0]
                raise type(e)(msg, e.extra_info) from None

    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

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)

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"]