make sure all records have the same fields.

Browse Source
This commit is contained in:
Friedrich Lindenberg 2018-04-04 10:43:19 +02:00
parent 9d99ed3e29
commit 66fd7ce173
1 changed files with 12 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
dataset/chunked.py
View File

@ -10,18 +10,25 @@ class ChunkedInsert(object):
def __init__(self, table, chunksize=1000): def __init__(self, table, chunksize=1000):
self.queue = [] self.queue = []
self.fields = set()
self.table = table self.table = table
self.chunksize = chunksize self.chunksize = chunksize
def flush(self):
for item in self.queue:
for field in self.fields:
item[field] = item.get(field)
self.table.insert_many(self.queue)
self.queue = []
def insert(self, item): def insert(self, item):
self.fields.update(item.keys())
self.queue.append(item) self.queue.append(item)
if len(self.queue) > self.chunksize: if len(self.queue) >= self.chunksize:
self.table.insert_many(self.queue) self.flush()
self.queue = []
def __enter__(self): def __enter__(self):
return self return self
def __exit__(self, exc_type, exc_val, exc_tb): def __exit__(self, exc_type, exc_val, exc_tb):
self.table.insert_many(self.queue) self.flush()
self.queue = []