micromegas 0.2.3__py3-none-any.whl → 0.3.1__py3-none-any.whl

micromegas/flightsql/arrow_ipc_reader.py

@@ -0,0 +1,99 @@
+import pyarrow
+
+# based on https://github.com/apache/arrow-rs/blob/main/arrow-ipc/src/reader.rs
+
+class ArrayReader:
+    def __init__(self, schema, nodes, buffers):
+        self.schema = schema
+        self.nodes = nodes
+        self.current_node = 0
+        self.buffers = buffers
+        self.current_buffer = 0
+
+    def next_node(self):
+        assert self.current_node < len(self.nodes)
+        node = self.nodes[self.current_node]
+        self.current_node += 1
+        return node
+
+    def next_buffer(self):
+        assert self.current_buffer < len(self.buffers)
+        buffer = self.buffers[self.current_buffer]
+        self.current_buffer += 1
+        return buffer
+
+
+def create_primitive_array(node, data_type, null_buffer, data_buffer):
+    return pyarrow.NumericArray.from_buffers(
+        data_type, node.Length(), [null_buffer, data_buffer], node.NullCount()
+    )
+
+
+def create_string_array(node, data_type, null_buffer, offset_buffer, data_buffer):
+    return pyarrow.NumericArray.from_buffers(
+        data_type,
+        node.Length(),
+        [null_buffer, offset_buffer, data_buffer],
+        node.NullCount(),
+    )
+
+
+def read_column(reader, arrow_field):
+    if arrow_field.type in [
+        pyarrow.string(),
+        pyarrow.binary(),
+        pyarrow.large_binary(),
+        pyarrow.large_string(),
+    ]:
+        return create_string_array(
+            reader.next_node(),
+            arrow_field.type,
+            reader.next_buffer(),
+            reader.next_buffer(),
+            reader.next_buffer(),
+        )
+    elif pyarrow.types.is_primitive(arrow_field.type):
+        return create_primitive_array(
+            reader.next_node(),
+            arrow_field.type,
+            reader.next_buffer(),
+            reader.next_buffer(),
+        )
+    elif pyarrow.types.is_list(arrow_field.type):
+        list_node = reader.next_node()
+        list_buffers = [reader.next_buffer(), reader.next_buffer()]
+        values = read_column(reader, arrow_field.type.value_field)
+        return pyarrow.ListArray.from_buffers(
+            arrow_field.type,
+            list_node.Length(),
+            list_buffers,
+            list_node.NullCount(),
+            0,
+            [values],
+        )
+    elif pyarrow.types.is_struct(arrow_field.type):
+        struct_node = reader.next_node()
+        null_buffer = reader.next_buffer()
+        children = []
+        for child_field in arrow_field.type.fields:
+            child_column = read_column(reader, child_field)
+            children.append(child_column)
+        return pyarrow.StructArray.from_buffers(
+            arrow_field.type,
+            struct_node.Length(),
+            [null_buffer],
+            struct_node.NullCount(),
+            0,
+            children,
+        )
+    else:
+        raise RuntimeError("unsupported arrow field type {}".format(arrow_field.type))
+
+
+def read_record_batch(arrow_schema, nodes, buffers):
+    reader = ArrayReader(arrow_schema, nodes, buffers)
+    columns = []
+    for arrow_field in arrow_schema:
+        column = read_column(reader, arrow_field)
+        columns.append(column)
+    return pyarrow.RecordBatch.from_arrays(columns, schema=arrow_schema)

micromegas/flightsql/client.py

@@ -0,0 +1,308 @@
+import pyarrow
+import grpc
+from google.protobuf import any_pb2
+from . import Flight_pb2_grpc
+from . import FlightSql_pb2_grpc
+from . import FlightSql_pb2
+from . import Flight_pb2
+from . import arrow_flatbuffers
+from . import arrow_ipc_reader
+from . import time
+
+
+def fb_time_unit_to_string(fb_time_unit):
+    time_unit_enum = arrow_flatbuffers.TimeUnit
+    if fb_time_unit == time_unit_enum.SECOND:
+        return "s"
+    if fb_time_unit == time_unit_enum.MILLISECOND:
+        return "ms"
+    if fb_time_unit == time_unit_enum.MICROSECOND:
+        return "us"
+    if fb_time_unit == time_unit_enum.NANOSECOND:
+        return "ns"
+    raise RuntimeError("unsupported time unit {}".format(fb_time_unit))
+
+
+def fb_field_type_to_arrow(fb_field):
+    fb_type = fb_field.TypeType()
+    fb_type_enum = arrow_flatbuffers.Type
+    assert fb_type != fb_type_enum.NONE
+    if fb_type == fb_type_enum.Null:
+        return pyarrow.null()
+    elif fb_type == fb_type_enum.Int:
+        type_int = arrow_flatbuffers.Int()
+        field_type_table = fb_field.Type()
+        type_int.Init(field_type_table.Bytes, field_type_table.Pos)
+        if type_int.IsSigned():
+            if type_int.BitWidth() == 8:
+                return pyarrow.int8()
+            elif type_int.BitWidth() == 16:
+                return pyarrow.int16()
+            elif type_int.BitWidth() == 32:
+                return pyarrow.int32()
+            elif type_int.BitWidth() == 64:
+                return pyarrow.int64()
+            else:
+                raise RuntimeError(
+                    "unsupported int size {}".format(type_int.BitWidth())
+                )
+        else:
+            if type_int.BitWidth() == 8:
+                return pyarrow.uint8()
+            elif type_int.BitWidth() == 16:
+                return pyarrow.uint16()
+            elif type_int.BitWidth() == 32:
+                return pyarrow.uint32()
+            elif type_int.BitWidth() == 64:
+                return pyarrow.uint64()
+            else:
+                raise RuntimeError(
+                    "unsupported uint size {}".format(type_int.BitWidth())
+                )
+    elif fb_type == fb_type_enum.FloatingPoint:
+        return pyarrow.float64()
+    elif fb_type == fb_type_enum.Binary:
+        return pyarrow.binary()
+    elif fb_type == fb_type_enum.Utf8:
+        return pyarrow.utf8()
+    elif fb_type == fb_type_enum.Bool:
+        return pyarrow.bool()
+    elif fb_type == fb_type_enum.Timestamp:
+        ts_type = arrow_flatbuffers.Timestamp()
+        field_type_table = fb_field.Type()
+        ts_type.Init(field_type_table.Bytes, field_type_table.Pos)
+        return pyarrow.timestamp(
+            fb_time_unit_to_string(ts_type.Unit()), ts_type.Timezone()
+        )
+    elif fb_type == fb_type_enum.List:
+        assert 1 == fb_field.ChildrenLength()
+        child_field = fb_field_to_arrow(fb_field.Children(0))
+        return pyarrow.list_(child_field)
+    elif fb_type == fb_type_enum.Struct_:
+        struct_fields = []
+        for child_index in range(fb_field.ChildrenLength()):
+            child = fb_field_to_arrow(fb_field.Children(child_index))
+            struct_fields.append(child)
+        return pyarrow.struct(struct_fields)
+    raise RuntimeError("unknown flatbuffer type {}".format(fb_type))
+
+
+def fb_field_to_arrow(fb_field):
+    arrow_type = fb_field_type_to_arrow(fb_field)
+    return pyarrow.field(fb_field.Name(), arrow_type)
+
+
+def make_query_flight_descriptor(sql):
+    command_query = FlightSql_pb2.CommandStatementQuery(query=sql)
+    any_cmd = any_pb2.Any()
+    any_cmd.Pack(command_query)
+    desc = Flight_pb2.FlightDescriptor()
+    desc.type = Flight_pb2.FlightDescriptor.DescriptorType.CMD
+    desc.cmd = any_cmd.SerializeToString()
+    return desc
+
+
+def read_schema_from_flight_data(flight_data):
+    msg = arrow_flatbuffers.Message.GetRootAs(flight_data.data_header, 0)
+    assert msg.Version() == arrow_flatbuffers.MetadataVersion.V5
+    header = msg.Header()
+    assert msg.HeaderType() == arrow_flatbuffers.MessageHeader.Schema
+
+    schema = arrow_flatbuffers.Schema()
+    schema.Init(header.Bytes, header.Pos)
+    nb_fields = schema.FieldsLength()
+    arrow_fields = []
+    for x in range(nb_fields):
+        field = schema.Fields(x)
+        arrow_f = fb_field_to_arrow(field)
+        arrow_fields.append(arrow_f)
+    arrow_schema = pyarrow.schema(arrow_fields)
+    return arrow_schema
+
+
+def read_record_batch_from_flight_data(arrow_schema, flight_data):
+    msg = arrow_flatbuffers.Message.GetRootAs(flight_data.data_header, 0)
+    assert msg.HeaderType() == arrow_flatbuffers.MessageHeader.RecordBatch
+    header = msg.Header()
+    fb_record_batch = arrow_flatbuffers.RecordBatch()
+    fb_record_batch.Init(header.Bytes, header.Pos)
+    nodes = []
+    for node_index in range(fb_record_batch.NodesLength()):
+        node = fb_record_batch.Nodes(node_index)
+        nodes.append(node)
+
+    buffers = []
+    for buffer_index in range(fb_record_batch.BuffersLength()):
+        buffer = fb_record_batch.Buffers(buffer_index)
+        buffers.append(buffer)
+
+    body = pyarrow.py_buffer(flight_data.data_body)
+    arrow_buffers = []
+    for b in buffers:
+        s = body.slice(b.Offset(), b.Length())
+        arrow_buffers.append(s)
+    rb = arrow_ipc_reader.read_record_batch(arrow_schema, nodes, arrow_buffers)
+    return rb
+
+
+def channel_creds_from_token(token):
+    call_credentials = grpc.access_token_call_credentials(token)
+    channel_cred = grpc.composite_channel_credentials(
+        grpc.ssl_channel_credentials(), call_credentials
+    )
+    return channel_cred
+
+
+class FlightSQLAuthMetadataPlugin(grpc.AuthMetadataPlugin):
+    def __init__(self, headers):
+        # we transform the keys into lowercase to avoid illegal grpc metadata (like 'Authorization', for example)
+        self.__headers = [(k.lower(), v) for (k, v) in headers.items()]
+
+    def __call__(self, context, callback):
+        callback(self.__headers, None)
+
+
+def channel_creds_from_headers(headers):
+    auth_plugin = FlightSQLAuthMetadataPlugin(headers)
+    call_credentials = grpc.metadata_call_credentials(auth_plugin)
+    channel_cred = grpc.composite_channel_credentials(
+        grpc.ssl_channel_credentials(), call_credentials
+    )
+    return channel_cred
+
+
+class FlightSQLClient:
+    def __init__(self, host_port, channel_creds):
+        self.__host_port = host_port
+        self.__channel_creds = channel_creds
+
+    def make_channel(self):
+        if self.__channel_creds is None:
+            return grpc.insecure_channel(self.__host_port)
+        else:
+            return grpc.secure_channel(self.__host_port, self.__channel_creds)
+
+    def query(self, sql, begin=None, end=None):
+        metadata = []
+        if begin is not None:
+            metadata.append(("query_range_begin", time.format_datetime(begin)))
+        if end is not None:
+            metadata.append(("query_range_end", time.format_datetime(end)))
+
+        channel = self.make_channel()
+        stub = Flight_pb2_grpc.FlightServiceStub(channel)
+        desc = make_query_flight_descriptor(sql)
+        info = stub.GetFlightInfo(desc)
+        grpc_rdv = stub.DoGet(info.endpoint[0].ticket, metadata=metadata)
+        flight_data_list = list(grpc_rdv)
+        if len(flight_data_list) < 1:
+            raise RuntimeError("too few flightdata messages {}", len(flight_data_list))
+        schema_message = flight_data_list[0]
+        data_messages = flight_data_list[1:]
+        schema = read_schema_from_flight_data(schema_message)
+        record_batches = []
+        for msg in data_messages:
+            record_batches.append(read_record_batch_from_flight_data(schema, msg))
+        table = pyarrow.Table.from_batches(record_batches, schema)
+        return table.to_pandas()
+
+    def query_stream(self, sql, begin=None, end=None):
+        metadata = []
+        if begin is not None:
+            metadata.append(("query_range_begin", time.format_datetime(begin)))
+        if end is not None:
+            metadata.append(("query_range_end", time.format_datetime(end)))
+
+        channel = self.make_channel()
+        stub = Flight_pb2_grpc.FlightServiceStub(channel)
+        desc = make_query_flight_descriptor(sql)
+        info = stub.GetFlightInfo(desc)
+        grpc_rdv = stub.DoGet(info.endpoint[0].ticket, metadata=metadata)
+        schema_message = grpc_rdv.next()
+        schema = read_schema_from_flight_data(schema_message)
+        for msg in grpc_rdv:
+            yield read_record_batch_from_flight_data(schema, msg)
+
+    def retire_partitions(self, view_set_name, view_instance_id, begin, end):
+        sql = """
+          SELECT time, msg
+          FROM retire_partitions('{view_set_name}', '{view_instance_id}', '{begin}', '{end}')
+        """.format(
+            view_set_name=view_set_name,
+            view_instance_id=view_instance_id,
+            begin=begin.isoformat(),
+            end=end.isoformat(),
+        )
+        for rb in self.query_stream(sql):
+            for index, row in rb.to_pandas().iterrows():
+                print(row["time"], row["msg"])
+
+    def materialize_partitions(
+        self, view_set_name, view_instance_id, begin, end, partition_delta_seconds
+    ):
+        sql = """
+          SELECT time, msg
+          FROM materialize_partitions('{view_set_name}', '{view_instance_id}', '{begin}', '{end}', {partition_delta_seconds})
+        """.format(
+            view_set_name=view_set_name,
+            view_instance_id=view_instance_id,
+            begin=begin.isoformat(),
+            end=end.isoformat(),
+            partition_delta_seconds=partition_delta_seconds,
+        )
+        for rb in self.query_stream(sql):
+            for index, row in rb.to_pandas().iterrows():
+                print(row["time"], row["msg"])
+
+    def find_process(self, process_id):
+        sql = """
+            SELECT *
+            FROM processes
+            WHERE process_id='{process_id}';
+            """.format(
+            process_id=process_id
+        )
+        return self.query(sql)
+
+    def query_streams(self, begin, end, limit, process_id=None, tag_filter=None):
+        conditions = []
+        if process_id is not None:
+            conditions.append("process_id='{process_id}'".format(process_id=process_id))
+        if tag_filter is not None:
+            conditions.append(
+                "(array_position(tags, '{tag}') is not NULL)".format(tag=tag_filter)
+            )
+        where = ""
+        if len(conditions) > 0:
+            where = "WHERE " + " AND ".join(conditions)
+        sql = """
+            SELECT *
+            FROM streams
+            {where}
+            LIMIT {limit};
+            """.format(
+            where=where, limit=limit
+        )
+        return self.query(sql, begin, end)
+
+    def query_blocks(self, begin, end, limit, stream_id):
+        sql = """
+            SELECT *
+            FROM blocks
+            WHERE stream_id='{stream_id}'
+            LIMIT {limit};
+            """.format(
+                limit=limit,stream_id=stream_id
+        )
+        return self.query(sql, begin, end)
+
+    def query_spans(self, begin, end, limit, stream_id):
+        sql = """
+            SELECT *
+            FROM view_instance('thread_spans', '{stream_id}')
+            LIMIT {limit};
+            """.format(
+                limit=limit,stream_id=stream_id
+        )
+        return self.query(sql, begin, end)
+    

micromegas/flightsql/time.py

@@ -0,0 +1,17 @@
+import datetime
+import pandas
+
+def format_datetime(value):
+    nonetype = type(None)
+    value_type = type(value)
+    if value_type == datetime.datetime:
+        if value.tzinfo is None:
+            raise RuntimeError("datetime needs a valid time zone")
+        return value.isoformat()
+    elif value_type == pandas.Timestamp:
+        return value.isoformat()
+    elif value_type == str:
+        return format_datetime(datetime.datetime.fromisoformat(value))
+    elif value_type == type(None):
+        return None
+    raise RuntimeError("value of unknown type in format_datetime")

{micromegas-0.2.3.dist-info → micromegas-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: micromegas
-Version: 0.2.3
+Version: 0.3.1
 Summary: Python analytics client for https://github.com/madesroches/micromegas/
 Author: Marc-Antoine Desroches
 Author-email: madesroches@gmail.com
@@ -11,9 +11,11 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: cbor2 (>=5.6.3,<6.0.0)
 Requires-Dist: crc (>=7.0.0,<8.0.0)
+Requires-Dist: flatbuffers (>=24.12.23,<25.0.0)
+Requires-Dist: grpcio (>=1.69.0,<2.0.0)
 Requires-Dist: pandas (>=2.2.2,<3.0.0)
-Requires-Dist: protobuf (>=5.27.1,<6.0.0)
-Requires-Dist: pyarrow (>=16.0.0,<17.0.0)
+Requires-Dist: protobuf (>=5.29.1,<6.0.0)
+Requires-Dist: pyarrow (>=18.1.0,<19.0.0)
 Requires-Dist: requests (>=2.31.0,<3.0.0)
 Requires-Dist: tabulate (>=0.9.0,<0.10.0)
 Requires-Dist: tqdm (>=4.66.5,<5.0.0)
@@ -25,20 +27,22 @@ Python analytics client for https://github.com/madesroches/micromegas/
 
 ## Example usage
 
-Query the 2 most recent log entries from the analytics service
+Query the 2 most recent log entries from the flightsql service
 
 ```python
 import datetime
 import pandas as pd
 import micromegas
+import grpc
 
-BASE_URL = "http://localhost:8082/"
-client = micromegas.client.Client(BASE_URL)
+host_port = "localhost:50051"
+channel_cred = grpc.local_channel_credentials()
+client = micromegas.flightsql.client.FlightSQLClient(host_port, channel_cred)
 sql = """
 SELECT time, process_id, level, target, msg
 FROM log_entries
 WHERE level <= 4
-AND exe LIKE '%analytics%'
+AND exe LIKE '%flight%'
 ORDER BY time DESC
 LIMIT 2
 """

{micromegas-0.2.3.dist-info → micromegas-0.3.1.dist-info}/RECORD

@@ -1,5 +1,14 @@
-micromegas/__init__.py,sha256=E_j3LFxMk9rSMJunwDCi_90NsRHm1fKwjj_6KGMYCjQ,246
+micromegas/__init__.py,sha256=0O7EfbXxIVlykEWkGYGc9y2EhrPcVoi_k3VILKTZKd8,337
 micromegas/client.py,sha256=7d0aah179a5cfjmBRxcaq7jTPvmGg625D_ixzdvtgCw,3670
+micromegas/flightsql/FlightSql_pb2.py,sha256=s3CVAQ4W77ALQ9QYFlPZnPqVsRHTD5Ma0fzabQv_8Mc,28207
+micromegas/flightsql/FlightSql_pb2_grpc.py,sha256=dCQT-tKxsz3dQgD4Cr0-8F4H0zCwNkXLHzlrbWfMiBA,890
+micromegas/flightsql/Flight_pb2.py,sha256=nPe1utv8n59IdNsiv5sJYY10-K04pDVOhKDgV5t-h1A,12513
+micromegas/flightsql/Flight_pb2_grpc.py,sha256=mj9Nlla3wpf5xFAGhMZpLDX8vLBClWfZ4w03AbBvrgs,23483
+micromegas/flightsql/__init__.py,sha256=SRq2X59uKG-iuwFnSyT7wVfiTMSSryAAEWnQVjGWOM8,249
+micromegas/flightsql/arrow_flatbuffers.py,sha256=egpmS59sNFwWmtG2wMeE92MMIyGyZbsm9o24mUnY1MQ,100142
+micromegas/flightsql/arrow_ipc_reader.py,sha256=3rxyEgqo5100e0TT9ZKZxNe7lX6Lk0mS6yRoiIJtH6Q,3163
+micromegas/flightsql/client.py,sha256=IFSUlf67m_2Pk3aleiL34ZeU7iptnA_-9WZ9yqOxQ9M,11328
+micromegas/flightsql/time.py,sha256=EH3SUEpFvY0lNMj9mOcvfUJuSgrQ3YX4aJnwteK2qhk,582
 micromegas/perfetto.py,sha256=yuIe5iKvca61aWMBQNziSGM-DHcOEsiobtKx2SsNQ3E,7829
 micromegas/request.py,sha256=NV0urom5P3_P2q94gX51hxW_Fnrp_DDRorsP3mUb5NM,941
 micromegas/thirdparty/perfetto/protos/perfetto/common/android_energy_consumer_descriptor_pb2.py,sha256=l8QNXqnB-mJIkuFr2s1YoLQXHm3G-ZcOGp_OW_hQ0TE,1887
@@ -208,6 +217,6 @@ micromegas/thirdparty/perfetto/protos/perfetto/trace/translation/translation_tab
 micromegas/thirdparty/perfetto/protos/perfetto/trace/trigger_pb2.py,sha256=We7Yi8o3cEcrSNxY1zLUUO6tEWnD36C2f3O_s8_qv0I,1435
 micromegas/thirdparty/perfetto/protos/perfetto/trace/ui_state_pb2.py,sha256=Af-SXwhroNhRXMrtw6e2eU1liCImMRxSdmkt_AuSHf8,1752
 micromegas/time.py,sha256=eD9fWF2UHxaf-92yd1X2SEgUcpKypqPsvjBosLdpnQA,1026
-micromegas-0.2.3.dist-info/METADATA,sha256=aHUqasX8WUEKzqH0O9R0DJmPenBvdMM_SB_hYurAsFE,30467
-micromegas-0.2.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-micromegas-0.2.3.dist-info/RECORD,,
+micromegas-0.3.1.dist-info/METADATA,sha256=O_YfkAcg7sEMv3FeclRmF1t91QQja7C0zz52Q0ilNAs,30640
+micromegas-0.3.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+micromegas-0.3.1.dist-info/RECORD,,