567 lines
16 KiB
Python
567 lines
16 KiB
Python
#
|
|
# vector.py
|
|
#
|
|
# This source file is part of the FoundationDB open source project
|
|
#
|
|
# Copyright 2013-2018 Apple Inc. and the FoundationDB project authors
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
"""FoundationDB Vector Layer.
|
|
|
|
Provides the Vector() class for storing and manipulating arrays
|
|
in FoundationDB.
|
|
"""
|
|
|
|
import fdb
|
|
import fdb.tuple
|
|
|
|
import threading
|
|
|
|
fdb.api_version(22)
|
|
|
|
|
|
###################################
|
|
# This defines a Subspace of keys #
|
|
###################################
|
|
|
|
|
|
class Subspace (object):
|
|
def __init__(self, prefixTuple, rawPrefix=""):
|
|
self.rawPrefix = rawPrefix + fdb.tuple.pack(prefixTuple)
|
|
|
|
def __getitem__(self, name):
|
|
return Subspace((name,), self.rawPrefix)
|
|
|
|
def key(self):
|
|
return self.rawPrefix
|
|
|
|
def pack(self, tuple):
|
|
return self.rawPrefix + fdb.tuple.pack(tuple)
|
|
|
|
def unpack(self, key):
|
|
assert key.startswith(self.rawPrefix)
|
|
return fdb.tuple.unpack(key[len(self.rawPrefix):])
|
|
|
|
def range(self, tuple=()):
|
|
p = fdb.tuple.range(tuple)
|
|
return slice(self.rawPrefix + p.start, self.rawPrefix + p.stop)
|
|
|
|
|
|
########################
|
|
# _ImplicitTransaction #
|
|
########################
|
|
|
|
# A local class which is used to allow vector operations to be performed without
|
|
# explicitly passing a transaction. It is created by vector.use_transaction
|
|
# and is used as follows:
|
|
#
|
|
# with vector.use_transaction(tr):
|
|
# vector[0] = 1
|
|
# vector.push(1)
|
|
# ...
|
|
|
|
|
|
class _ImplicitTransaction:
|
|
def __init__(self, vector, tr):
|
|
self.vector = vector
|
|
self.tr = tr
|
|
self.initialValue = self.vector.local.tr
|
|
|
|
def __enter__(self):
|
|
if self.initialValue is not None and self.vector.local.tr != self.tr:
|
|
raise Exception('use_transaction cannot be nested')
|
|
|
|
self.vector.local.tr = self.tr
|
|
|
|
def __exit__(self, type, value, traceback):
|
|
self.vector.local.tr = self.initialValue
|
|
|
|
|
|
##########
|
|
# Vector #
|
|
##########
|
|
|
|
# Vector stores each of its values using its index as the key.
|
|
# The size of a vector is equal to the index of its last key + 1.
|
|
##
|
|
# For indexes smaller than the vector's size that have no associated key
|
|
# in the database, the value will be the specified defaultValue.
|
|
##
|
|
# If the last value in the vector has the default value, its key will
|
|
# always be set so that size can be determined.
|
|
##
|
|
# By creating Vector with a Subspace, all kv pairs modified by the
|
|
# layer will have keys that start within that Subspace.
|
|
|
|
|
|
class Vector:
|
|
"""Represents a potentially sparse array in FoundationDB."""
|
|
|
|
# Public functions
|
|
|
|
def __init__(self, subspace, defaultValue=''):
|
|
self.subspace = subspace
|
|
self.defaultValue = defaultValue
|
|
self.local = threading.local()
|
|
self.local.tr = None
|
|
|
|
def use_transaction(self, tr):
|
|
"""
|
|
Get an object that can be used in a with statement to perform operations
|
|
on this vector without supplying a transaction as an argument to each operation.
|
|
|
|
For example:
|
|
|
|
with vector.use_transaction(tr):
|
|
vector[0] = 1
|
|
vector.push(1)
|
|
...
|
|
"""
|
|
return _ImplicitTransaction(self, tr)
|
|
|
|
def size(self, tr=None):
|
|
"""Get the number of items in the Vector. This number includes the sparsely represented items."""
|
|
return self._size(self._to_transaction(tr))
|
|
|
|
def push(self, val, tr=None):
|
|
"""Push a single item onto the end of the Vector."""
|
|
self._push(val, self._to_transaction(tr))
|
|
|
|
def back(self, tr=None):
|
|
"""Get the value of the last item in the Vector."""
|
|
return self._back(self._to_transaction(tr))
|
|
|
|
def front(self, tr=None):
|
|
"""Get the value of the first item in the Vector."""
|
|
return self._get(0, self._to_transaction(tr))
|
|
|
|
def pop(self, tr=None):
|
|
"""Get and pops the last item off the Vector."""
|
|
return self._pop(self._to_transaction(tr))
|
|
|
|
def swap(self, i1, i2, tr=None):
|
|
"""Swap the items at positions i1 and i2."""
|
|
self._swap(i1, i2, self._to_transaction(tr))
|
|
|
|
def get(self, index, tr=None):
|
|
"""Get the item at the specified index."""
|
|
return self._get(index, self._to_transaction(tr))
|
|
|
|
def get_range(self, startIndex=None, endIndex=None, step=None, tr=None):
|
|
"""Get a range of items in the Vector, returned as a generator."""
|
|
return self._get_range(startIndex, endIndex, step, self._to_transaction(tr))
|
|
|
|
def set(self, index, val, tr=None):
|
|
"""Set the value at a particular index in the Vector."""
|
|
self._set(index, val, self._to_transaction(tr))
|
|
|
|
def empty(self, tr=None):
|
|
"""Test whether the Vector is empty."""
|
|
return self._size(self._to_transaction(tr)) == 0
|
|
|
|
def resize(self, length, tr=None):
|
|
"""Grow or shrink the size of the Vector."""
|
|
self._resize(length, self._to_transaction(tr))
|
|
|
|
def clear(self, tr=None):
|
|
"""Remove all items from the Vector."""
|
|
self._clear(self._to_transaction(tr))
|
|
|
|
# Vector supports array notation when combined with use_transaction
|
|
def __setitem__(self, index, val):
|
|
"""Set the item at the specified index. Can only be used with use_transaction."""
|
|
self.set(index, val)
|
|
|
|
def __getitem__(self, index):
|
|
"""
|
|
Get the item(s) at the specified index or range. Ranges are returned as a generator. Can only
|
|
be used with use_transaction()
|
|
"""
|
|
if isinstance(index, slice):
|
|
return self.get_range(index.start, index.stop, index.step)
|
|
else:
|
|
return self.get(index)
|
|
|
|
# Private functions
|
|
|
|
@fdb.transactional
|
|
def _push(self, val, tr):
|
|
tr[self._key_at(self.size())] = fdb.tuple.pack((val,))
|
|
|
|
@fdb.transactional
|
|
def _back(self, tr):
|
|
keyRange = self.subspace.range()
|
|
last = tr.get_range(keyRange.start, keyRange.stop, 1, True)
|
|
for k, v in last:
|
|
return fdb.tuple.unpack(v)[0]
|
|
return None
|
|
|
|
@fdb.transactional
|
|
def _pop(self, tr):
|
|
keyRange = self.subspace.range()
|
|
|
|
# Read the last two entries so we can check if the second to last item
|
|
# is being represented sparsely. If so, we will be required to set it
|
|
# to the default value
|
|
lastTwo = list(tr.get_range(keyRange.start, keyRange.stop, 2, True))
|
|
indices = [self.subspace.unpack(kv.key)[0] for kv in lastTwo]
|
|
|
|
# Vector was empty
|
|
if len(lastTwo) == 0:
|
|
return None
|
|
|
|
# Vector has size one:
|
|
elif indices[0] == 0:
|
|
pass
|
|
|
|
# Second to last item is being represented sparsely
|
|
elif len(lastTwo) == 1 or indices[0] > indices[1] + 1:
|
|
tr[self._key_at(indices[0] - 1)] = fdb.tuple.pack((self.defaultValue,))
|
|
|
|
del tr[lastTwo[0].key]
|
|
return fdb.tuple.unpack(lastTwo[0].value)[0]
|
|
|
|
@fdb.transactional
|
|
def _swap(self, i1, i2, tr):
|
|
k1 = self._key_at(i1)
|
|
k2 = self._key_at(i2)
|
|
|
|
currentSize = self._size(tr)
|
|
v1 = tr[k1]
|
|
v2 = tr[k2]
|
|
|
|
if i1 > currentSize or i2 > currentSize or i1 < 0 or i2 < 0:
|
|
raise IndexError('vector.swap: indices (%d, %d) out of range' % (i1, i2))
|
|
|
|
if v2.present():
|
|
tr[k1] = v2
|
|
elif v1.present() and i1 < currentSize - 1:
|
|
del tr[k1]
|
|
|
|
if v1.present():
|
|
tr[k2] = v1
|
|
elif v2.present() and i2 < currentSize - 1:
|
|
del tr[k2]
|
|
|
|
@fdb.transactional
|
|
def _get(self, index, tr):
|
|
if index < 0:
|
|
raise IndexError('vector.get: index \'%d\' out of range' % index)
|
|
|
|
start = self._key_at(index)
|
|
end = self.subspace.range().stop
|
|
|
|
output = tr.get_range(start, end, 1)
|
|
for k, v in output:
|
|
# The requested index had an associated key
|
|
if(start == k):
|
|
return fdb.tuple.unpack(v)[0]
|
|
|
|
# The requested index is sparsely represented
|
|
return self.defaultValue
|
|
|
|
# We requested a value past the end of the vector
|
|
raise IndexError('vector.get: index \'%d\' out of range' % index)
|
|
|
|
def _get_range(self, startIndex, endIndex, step, tr):
|
|
size = self._size(tr)
|
|
if startIndex is not None and startIndex < 0:
|
|
startIndex = max(0, size + startIndex)
|
|
if endIndex is not None and endIndex < 0:
|
|
endIndex = max(0, size + endIndex)
|
|
|
|
if step is None:
|
|
if startIndex is None or endIndex is None or startIndex <= endIndex:
|
|
step = 1
|
|
else:
|
|
step = -1
|
|
elif step == 0:
|
|
raise ValueError('vector.get_range: step cannot be zero')
|
|
|
|
if startIndex is None:
|
|
if step > 0:
|
|
start = self.subspace.range().start
|
|
else:
|
|
end = self.subspace.range().stop
|
|
else:
|
|
if step > 0:
|
|
start = self._key_at(startIndex)
|
|
else:
|
|
end = self._key_at(startIndex + 1)
|
|
|
|
if endIndex is None:
|
|
if step > 0:
|
|
end = self.subspace.range().stop
|
|
else:
|
|
start = self.subspace.range().start
|
|
|
|
else:
|
|
if step > 0:
|
|
end = self._key_at(endIndex)
|
|
else:
|
|
start = self._key_at(endIndex + 1)
|
|
|
|
result = tr.get_range(start, end, 0, step < 0)
|
|
|
|
currentIndex = startIndex
|
|
if currentIndex is None:
|
|
if step > 0:
|
|
currentIndex = 0
|
|
else:
|
|
currentIndex = size - 1
|
|
elif currentIndex >= size:
|
|
currentIndex = size - 1
|
|
|
|
for k, v in result:
|
|
keyIndex = self.subspace.unpack(k)[0]
|
|
while (step > 0 and currentIndex < keyIndex) or (step < 0 and currentIndex > keyIndex):
|
|
currentIndex = currentIndex + step
|
|
yield self.defaultValue
|
|
|
|
if currentIndex == keyIndex:
|
|
currentIndex = currentIndex + step
|
|
yield fdb.tuple.unpack(v)[0]
|
|
|
|
@fdb.transactional
|
|
def _size(self, tr):
|
|
keyRange = self.subspace.range()
|
|
lastKey = tr.get_key(fdb.KeySelector.last_less_or_equal(keyRange.stop))
|
|
if lastKey < keyRange.start:
|
|
return 0
|
|
|
|
return self.subspace.unpack(lastKey)[0] + 1
|
|
|
|
@fdb.transactional
|
|
def _set(self, index, val, tr):
|
|
tr[self._key_at(index)] = fdb.tuple.pack((val,))
|
|
|
|
@fdb.transactional
|
|
def _resize(self, length, tr):
|
|
currentSize = self.size()
|
|
if(length == currentSize):
|
|
return
|
|
if(length < currentSize):
|
|
self._shrink(tr, length, currentSize)
|
|
else:
|
|
self._expand(tr, length, currentSize)
|
|
|
|
@fdb.transactional
|
|
def _shrink(self, tr, length, currentSize):
|
|
tr.clear_range(self._key_at(length), self.subspace.range().stop)
|
|
|
|
# Check if the new end of the vector was being sparsely represented
|
|
if self._size(tr) < length:
|
|
tr[self._key_at(length - 1)] = fdb.tuple.pack((self.defaultValue,))
|
|
|
|
@fdb.transactional
|
|
def _expand(self, tr, length, currentSize):
|
|
tr[self._key_at(length - 1)] = fdb.tuple.pack((self.defaultValue,))
|
|
|
|
@fdb.transactional
|
|
def _clear(self, tr):
|
|
del tr[self.subspace.range()]
|
|
|
|
def _to_transaction(self, tr):
|
|
if tr is None:
|
|
if self.local.tr is None:
|
|
raise Exception('No transaction specified and use_transaction has not been called')
|
|
else:
|
|
return self.local.tr
|
|
else:
|
|
return tr
|
|
|
|
def _key_at(self, index):
|
|
return self.subspace.pack((index,))
|
|
|
|
|
|
##################
|
|
# internal tests #
|
|
##################
|
|
|
|
|
|
# caution: modifies the database!
|
|
@fdb.transactional
|
|
def vector_test(tr):
|
|
vector = Vector(Subspace(('test_vector',)), 0)
|
|
|
|
with vector.use_transaction(tr):
|
|
print 'Clearing any previous values in the vector'
|
|
vector.clear()
|
|
|
|
print '\nMODIFIERS'
|
|
|
|
# Set + Push
|
|
vector[0] = 1
|
|
vector[1] = 2
|
|
vector.push(3)
|
|
_print_vector(vector, tr)
|
|
|
|
# Swap
|
|
vector.swap(0, 2)
|
|
_print_vector(vector, tr)
|
|
|
|
# Pop
|
|
print 'Popped:', vector.pop()
|
|
_print_vector(vector, tr)
|
|
|
|
# Clear
|
|
vector.clear()
|
|
|
|
print 'Pop empty:', vector.pop()
|
|
_print_vector(vector, tr)
|
|
|
|
vector.push('Foo')
|
|
print 'Pop size 1:', vector.pop()
|
|
_print_vector(vector, tr)
|
|
|
|
print '\nCAPACITY OPERATIONS'
|
|
|
|
# Capacity
|
|
print 'Size:', vector.size()
|
|
print 'Empty:', vector.empty()
|
|
|
|
print 'Resizing to length 5'
|
|
vector.resize(5)
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print 'Setting values'
|
|
vector[0] = 'The'
|
|
vector[1] = 'Quick'
|
|
vector[2] = 'Brown'
|
|
vector[3] = 'Fox'
|
|
vector[4] = 'Jumps'
|
|
vector[5] = 'Over'
|
|
_print_vector(vector, tr)
|
|
|
|
print '\nFRONT'
|
|
print vector.front()
|
|
|
|
print '\nBACK'
|
|
print vector.back()
|
|
|
|
print '\nELEMENT ACCESS'
|
|
print 'Index 0:', vector[0]
|
|
print 'Index 5:', vector[5]
|
|
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print '\nRESIZING'
|
|
print 'Resizing to 3'
|
|
vector.resize(3)
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print 'Resizing to 3 again'
|
|
vector.resize(3)
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print 'Resizing to 6'
|
|
vector.resize(6)
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print '\nSPARSE TESTS'
|
|
print 'Popping sparse vector'
|
|
vector.pop()
|
|
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print 'Resizing to 4'
|
|
vector.resize(4)
|
|
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print 'Adding "word" to index 10, resize to 25'
|
|
vector[10] = 'word'
|
|
vector.resize(25)
|
|
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print 'Popping sparse vector'
|
|
vector.pop()
|
|
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print 'Swapping with sparse element'
|
|
vector.swap(10, 15)
|
|
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
print 'Swapping sparse elements'
|
|
vector.swap(12, 13)
|
|
|
|
_print_vector(vector, tr)
|
|
print 'Size:', vector.size()
|
|
|
|
|
|
##############################
|
|
# Vector sample usage #
|
|
##############################
|
|
|
|
|
|
import sys
|
|
|
|
|
|
# caution: modifies the database!
|
|
@fdb.transactional
|
|
def vector_example(tr):
|
|
vector = Vector(Subspace(('my_vector',)), 0)
|
|
|
|
with vector.use_transaction(tr):
|
|
vector.clear()
|
|
|
|
for i in range(0, 5):
|
|
vector.push(i)
|
|
|
|
_print_vector(vector, tr)
|
|
|
|
print 'Pop last item: %d' % vector.pop()
|
|
_print_vector(vector, tr)
|
|
|
|
vector[1] = 10
|
|
vector.resize(11)
|
|
_print_vector(vector, tr)
|
|
|
|
vector.swap(1, 10)
|
|
_print_vector(vector, tr)
|
|
|
|
|
|
def _print_vector(vector, tr):
|
|
first = True
|
|
with vector.use_transaction(tr):
|
|
for v in vector:
|
|
if not first:
|
|
sys.stdout.write(',')
|
|
|
|
first = False
|
|
sys.stdout.write(repr(v))
|
|
|
|
print
|
|
|
|
|
|
# caution: modifies the database!
|
|
if __name__ == '__main__':
|
|
db = fdb.open()
|
|
vector_example(db)
|
|
# vector_test(db)
|