fix key prefix

usc-isi-i2 · Apr 12, 2018 · 5d30a5d · 5d30a5d
1 parent 6a57c4b
commit 5d30a5d
Showing 1 changed file with 15 additions and 19 deletions.
diff --git a/rltk/io/adapter/hbase_adapter.py b/rltk/io/adapter/hbase_adapter.py
@@ -6,25 +6,24 @@
 
 
 class HBaseAdapter(KeyValueAdapter):
-    def __init__(self, host, table, serializer: Serializer=None, key_format='{record_id}', **kwargs):
+    def __init__(self, host, table, serializer: Serializer=None, key_prefix='', **kwargs):
         if not serializer:
             serializer = PickleSerializer()
         self._conn = happybase.Connection(host=host, **kwargs)
         self._serializer = serializer
-        self._key_format = key_format
-        self._family_name = 'rltk'
-        self._col_name = 'obj'
-        if bytes(table, 'utf-8') not in self._conn.tables():
+        self._key_prefix = key_prefix.encode('utf-8')
+        self._family_name = b'rltk'
+        self._col_name = b'obj'
+        self._fam_col_name = '{}:{}'.format(
+            self._family_name.decode('utf-8'), self._col_name.decode('utf-8')).encode('utf-8')
+
+        if table.encode('utf-8') not in self._conn.tables():
             self._create_table(table)
         self._table = self._conn.table(table)
 
-        try:
-            self._get_key('test_id')
-        except:
-            raise ValueError('Invalid key_format.')
-
     def _get_key(self, record_id):
-        return bytes(self._key_format.format(record_id=record_id), 'utf-8')
+        return '{prefix}_{record_id}'.format(
+            prefix=self._key_prefix.decode('utf-8'), record_id=record_id).encode('utf-8')
 
     def __del__(self):
         try:
@@ -33,20 +32,17 @@ def __del__(self):
             pass
 
     def _create_table(self, table_name):
-        self._conn.create_table(table_name, {bytes(self._family_name, 'utf-8'):dict()})
+        self._conn.create_table(table_name, {self._family_name:dict()})
 
     def get(self, record_id) -> Record:
-        col = bytes('{}:{}'.format(self._family_name, self._col_name), 'utf-8')
-        return self._serializer.loads(self._table.rows(self._get_key(record_id))[col])
+        return self._serializer.loads(self._table.rows(self._get_key(record_id))[self._fam_col_name])
 
     def set(self, record_id, record: Record):
-        col = bytes('{}:{}'.format(self._family_name, self._col_name), 'utf-8')
-        return self._table.put(self._get_key(record_id), {col: self._serializer.dumps(record)})
+        return self._table.put(self._get_key(record_id), {self._fam_col_name: self._serializer.dumps(record)})
 
     def __iter__(self):
         return self.__next__()
 
     def __next__(self):
-        col = bytes('{}:{}'.format(self._family_name, self._col_name), 'utf-8')
-        for key, data in self._table.scan(self._get_key('*')):
-            yield self._serializer.loads(data[col])
+        for key, data in self._table.scan(row_prefix=self._key_prefix):
+            yield self._serializer.loads(data[self._fam_col_name])