Class MetadataFlightServer

python/pyarrow/tests/test_flight.py:172–210

A Flight server that numbers incoming/outgoing data.

class MetadataFlightServer(FlightServerBase):
    """A Flight server that numbers incoming/outgoing data."""

    def __init__(self, options=None, **kwargs):
        super().__init__(**kwargs)
        self.options = options

    def do_get(self, context, ticket):
        data = [
            pa.array([-10, -5, 0, 5, 10])
        ]
        table = pa.Table.from_arrays(data, names=['a'])
        return flight.GeneratorStream(
            table.schema,
            self.number_batches(table),
            options=self.options)

    def do_put(self, context, descriptor, reader, writer):
        counter = 0
        expected_data = [-10, -5, 0, 5, 10]
        assert reader.stats.num_messages == 1
        for batch, buf in reader:
            assert batch.equals(pa.RecordBatch.from_arrays(
                [pa.array([expected_data[counter]])],
                ['a']
            ))
            assert buf is not None
            client_counter, = struct.unpack('<i', buf.to_pybytes())
            assert counter == client_counter
            writer.write(struct.pack('<i', counter))
            counter += 1
        assert reader.stats.num_messages == 6
        assert reader.stats.num_record_batches == 5

    @staticmethod
    def number_batches(table):
        for idx, batch in enumerate(table.to_batches()):
            buf = struct.pack('<i', idx)
            yield batch, buf


class EchoFlightServer(FlightServerBase):