quasardb 3.14.2.dev4__cp313-cp313-win_amd64.whl

quasardb/extensions/writer.py

@@ -0,0 +1,193 @@
+import copy
+import quasardb
+import numpy as np
+import numpy.ma as ma
+
+__all__ = []
+
+def _ensure_ctype(self, idx, ctype):
+    assert 'table' in self._legacy_state
+    infos = self._legacy_state['table'].list_columns()
+    cinfo = infos[idx]
+
+    ctype_data = copy.copy(ctype)
+    ctype_column = copy.copy(cinfo.type)
+
+    if ctype_data == quasardb.ColumnType.Symbol:
+        ctype_data = quasardb.ColumnType.String
+
+    if ctype_column == quasardb.ColumnType.Symbol:
+        ctype_column = quasardb.ColumnType.String
+
+    if not ctype_data == ctype_column:
+        raise quasardb.IncompatibleTypeError()
+
+
+def _legacy_next_row(self, table):
+    if 'pending' not in self._legacy_state:
+        self._legacy_state['pending'] = []
+
+    if 'table' not in self._legacy_state:
+        self._legacy_state['table'] = table
+
+    self._legacy_state['pending'].append({'by_index': {}})
+
+    # Return reference to the row inside the buffer
+    return self._legacy_state['pending'][-1]
+
+
+def _legacy_current_row(self):
+    return self._legacy_state['pending'][-1]
+
+
+def _legacy_start_row(self, table, x):
+    row = _legacy_next_row(self, table)
+    assert '$timestamp' not in row
+    row['$timestamp'] = x
+
+
+def _legacy_set_double(self, idx, x):
+    _ensure_ctype(self, idx, quasardb.ColumnType.Double)
+    assert isinstance(x, float)
+    assert idx not in _legacy_current_row(self)['by_index']
+    _legacy_current_row(self)['by_index'][idx] = x
+
+
+def _legacy_set_int64(self, idx, x):
+    _ensure_ctype(self, idx, quasardb.ColumnType.Int64)
+    assert isinstance(x, int)
+    assert idx not in _legacy_current_row(self)['by_index']
+    _legacy_current_row(self)['by_index'][idx] = x
+
+
+def _legacy_set_timestamp(self, idx, x):
+    _ensure_ctype(self, idx, quasardb.ColumnType.Timestamp)
+    assert idx not in _legacy_current_row(self)['by_index']
+    _legacy_current_row(self)['by_index'][idx] = x
+
+
+def _legacy_set_string(self, idx, x):
+    _ensure_ctype(self, idx, quasardb.ColumnType.String)
+    assert isinstance(x, str)
+    assert idx not in _legacy_current_row(self)['by_index']
+
+    _legacy_current_row(self)['by_index'][idx] = x
+
+
+def _legacy_set_blob(self, idx, x):
+    _ensure_ctype(self, idx, quasardb.ColumnType.Blob)
+    assert isinstance(x, bytes)
+    assert idx not in _legacy_current_row(self)['by_index']
+
+    _legacy_current_row(self)['by_index'][idx] = x
+
+
+def _legacy_push(self):
+    if 'pending' not in self._legacy_state:
+        # Extremely likely default case, no "old" rows
+        return
+
+    assert 'table' in self._legacy_state
+    table = self._legacy_state['table']
+
+    # Some useful constants
+    dtype_by_ctype = {quasardb.ColumnType.Double: np.dtype('float64'),
+                      quasardb.ColumnType.Int64: np.dtype('int64'),
+                      quasardb.ColumnType.Timestamp: np.dtype('datetime64[ns]'),
+                      quasardb.ColumnType.String: np.dtype('unicode'),
+                      quasardb.ColumnType.Symbol: np.dtype('unicode'),
+                      quasardb.ColumnType.Blob: np.dtype('bytes')
+                    }
+
+    ctype_by_idx = {}
+    cinfos = table.list_columns()
+    for i in range(len(cinfos)):
+        ctype_by_idx[i] = cinfos[i].type
+
+    all_idx = set(ctype_by_idx.keys())
+
+    # Prepare data structure
+    pivoted = {'$timestamp': [],
+               'by_index': {}}
+    for i in all_idx:
+        pivoted['by_index'][i] = []
+
+    # Do the actual pivot
+    for row in self._legacy_state['pending']:
+        assert '$timestamp' in row
+        assert 'by_index' in row
+
+        pivoted['$timestamp'].append(row['$timestamp'])
+
+        for idx in pivoted['by_index'].keys():
+            val = row['by_index'].get(idx, None)
+            pivoted['by_index'][idx].append(val)
+
+    # Validation / verification, not strictly necessary. Effectively
+    # ensures that we have the exact same amount of values for every
+    # column
+    for xs in pivoted['by_index'].values():
+        assert len(xs) == len(pivoted['$timestamp'])
+
+    column_data = []
+
+    for idx,xs in pivoted['by_index'].items():
+        ctype = ctype_by_idx[idx]
+        dtype = dtype_by_ctype[ctype]
+
+        # None-mask works, because everything inside the list are just regular ojbects
+
+        mask = [x is None for x in xs]
+
+        xs_ = []
+        if all(mask):
+            xs_ = ma.masked_all(len(xs),
+                                dtype=dtype)
+        else:
+            xs_ = ma.masked_array(data=np.array(xs, dtype), mask=mask)
+
+        assert len(xs_) == len(pivoted['$timestamp'])
+
+        column_data.append(xs_)
+
+
+    push_data = quasardb.WriterData()
+    index = np.array(pivoted['$timestamp'], np.dtype('datetime64[ns]'))
+
+    push_data.append(table, index, column_data)
+
+
+    self._legacy_state = {}
+    return push_data
+
+
+def _wrap_fn(old_fn, replace_fn):
+
+    def wrapped(self, *args, **kwargs):
+        data = replace_fn(self)
+        if data:
+            return old_fn(self, data, *args, **kwargs)
+        else:
+            return old_fn(self, *args, **kwargs)
+
+    return wrapped
+
+
+def extend_writer(x):
+    """
+    Extends the writer with the "old", batch inserter API. This is purely
+    a backwards compatibility layer, and we want to avoid having to maintain that
+    in C++ with few benefits.
+    """
+
+    x.start_row      = _legacy_start_row
+    x.set_double     = _legacy_set_double
+    x.set_int64      = _legacy_set_int64
+    x.set_string     = _legacy_set_string
+    x.set_blob       = _legacy_set_blob
+    x.set_timestamp  = _legacy_set_timestamp
+
+    x.push           = _wrap_fn(x.push, _legacy_push)
+    x.push_fast      = _wrap_fn(x.push_fast, _legacy_push)
+    x.push_async     = _wrap_fn(x.push_async, _legacy_push)
+    x.push_truncate  = _wrap_fn(x.push_truncate, _legacy_push)

quasardb/firehose.py

@@ -0,0 +1,101 @@
+import time
+import quasardb
+import logging
+import numpy as np
+
+FIREHOSE_TABLE = "$qdb.firehose"
+POLL_INTERVAL = 0.1
+
+logger = logging.getLogger('quasardb.firehose')
+
+
+def _init():
+    """
+    Initialize our internal state.
+    """
+    return {'last': None,
+            'seen': set()}
+
+
+def _get_transactions_since(conn, table_name, last):
+    """
+    Retrieve all transactions since a certain timestamp. `last` is expected to be a dict
+    firehose row with at least a $timestamp attached.
+    """
+    if last is None:
+        q = "SELECT $timestamp, transaction_id, begin, end FROM \"{}\" WHERE table = '{}' ORDER BY $timestamp".format(
+            FIREHOSE_TABLE, table_name)
+    else:
+        q = "SELECT $timestamp, transaction_id, begin, end FROM \"{}\" IN RANGE ({}, +1y) WHERE table = '{}' ORDER BY $timestamp".format(
+            FIREHOSE_TABLE, last['$timestamp'], table_name)
+
+    return conn.query(q)
+
+
+def _get_transaction_data(conn, table_name, begin, end):
+    """
+    Gets all data from a certain table.
+    """
+    q = "SELECT * FROM \"{}\" IN RANGE ({}, {}) ".format(
+        table_name, begin, end)
+    return conn.query(q)
+
+
+def _get_next(conn, table_name, state):
+
+    # Our flow to retrieve new data is as follows:
+    # 1. Based on the state's last processed transaction, retrieve all transactions
+    #    that are logged into the firehose since then.
+    # 2. For each of the transactions, verify we haven't seen it before
+    # 3. For each of the transactions, pull in all data
+    # 4. Concatenate all this data (in order of quasardb transaction)
+
+    txs = _get_transactions_since(conn, table_name, state['last'])
+
+    xs = list()
+    for tx in txs:
+        txid = tx['transaction_id']
+
+        if state['last'] is not None and tx['$timestamp'] > state['last']['$timestamp']:
+            # At this point we are guaranteed that the transaction we encounter is
+            # 'new', will not conflict with any other transaction ids. It is thus
+            # safe to reset the txid set.
+            state['seen'] = set()
+
+        if txid not in state['seen']:
+            xs = xs + _get_transaction_data(conn,
+                                            table_name,
+                                            tx['begin'],
+                                            # The firehose logs transaction `end` span as
+                                            # end inclusive, while our bulk reader and/or query
+                                            # language are end exclusive.
+                                            tx['end'] + np.timedelta64(1, 'ns'))
+
+            # Because it is possible that multiple firehose changes are stored with the
+            # exact same $timestamp, we also keep track of the actually seen
+            # transaction ids.
+            state['seen'].add(txid)
+
+        state['last'] = tx
+
+    return (state, xs)
+
+
+def subscribe(conn, table_name):
+    state = _init()
+
+    while True:
+        # Note how this is effectively a never-ending fold loop
+        # that transforms state into a new state. This state effectively
+        # functions as a checkpoint.
+        #
+        # At a later point, we could choose to provide the user
+        # direct access to this 'state' object, so that they can
+        # implement e.g. mechanisms to replay from a certain checkpoint.
+        (state, xs) = _get_next(conn, table_name, state)
+
+        for x in xs:
+            yield x
+
+        # Our poll interval
+        time.sleep(POLL_INTERVAL)