def resize_to_new_table(self, new_size_per_shard, new_dir_name):
self.increment_rpc_count_by(n=2)
assert not FileUtil.does_dir_exist(dir_name=new_dir_name) or FileUtil.is_dir_empty(dir_name=new_dir_name)
new_sptable_storage = ShardedProtoTableStorage(size_per_shard=new_size_per_shard)
new_sptable_storage.initialize_from_dir(dir_name=new_dir_name)
for shard in range(self.get_num_shards()):
related_proto_file = self._shard_to_file(shard=shard)
proto_table = ProtoTableStorage()
proto_table.initialize_from_file(file_name=related_proto_file)
new_sptable_storage.write(data=proto_table.read_all())
self.increment_rpc_count_by(n=proto_table.get_rpc_call_count_and_reset())
return new_sptable_storage
def write(self, data, params=None):
if not params:
params = {}
if 'overwrite' not in params:
params['overwrite'] = True
assert isinstance(data, dict)
exising_shard_to_data_map = defaultdict(dict)
new_data = {}
for key, val in data.items():
if key in self._index_map.index_map:
exising_shard_to_data_map[self._index_map.index_map[key]][key] = val
else:
new_data[key] = val
try:
for shard, existing_data in exising_shard_to_data_map.items():
related_proto_file = self._shard_to_file(shard=shard)
proto_table = ProtoTableStorage()
proto_table.initialize_from_file(file_name=related_proto_file)
proto_table.write(data=existing_data, params=params)
self.increment_rpc_count_by(n=proto_table.get_rpc_call_count_and_reset())
if new_data:
all_new_keys = list(new_data.keys())
latest_shard = self.get_latest_shard()
latest_proto_file = self._shard_to_file(shard=latest_shard)
proto_table = ProtoTableStorage()
proto_table.initialize_from_file(file_name=latest_proto_file)
latest_proto_table_size = proto_table.get_num_entries()
proto_table.write(
data={key: new_data[key] for key in all_new_keys[:self._size_per_shard - latest_proto_table_size]},
params=params
)
self.increment_rpc_count_by(n=proto_table.get_rpc_call_count_and_reset())
for key in all_new_keys[:self._size_per_shard - latest_proto_table_size]:
self._index_map.index_map[key] = latest_shard
if len(all_new_keys) > self._size_per_shard - latest_proto_table_size:
remaining_new_keys = all_new_keys[self._size_per_shard - latest_proto_table_size:]
start_index = 0
while start_index < len(remaining_new_keys):
data_for_new_shard = remaining_new_keys[start_index:start_index + self._size_per_shard]
latest_shard += 1
self._index_map.cur_shard += 1
proto_file = self._shard_to_file(shard=latest_shard)
self._logger.info("Write to new file with name [" + proto_file + '] and shard [' +
str(latest_shard) + '].')
proto_table = ProtoTableStorage()
proto_table.initialize_from_file(file_name=proto_file)
proto_table.write(
data={key: new_data[key] for key in data_for_new_shard},
params=params
)
self.increment_rpc_count_by(n=proto_table.get_rpc_call_count_and_reset())
for key in data_for_new_shard:
self._index_map.index_map[key] = latest_shard
start_index += self._size_per_shard
self._logger.info("Writing the index map to [" + self._index_map_file + '].')
self.increment_rpc_count_by(n=1)
FileUtil.write_proto_to_file(
proto=self._index_map,
file_name=self._index_map_file
)
except Exception as err:
self._SYS_LOGGER.error("Write to dir [" + self.get_dir_name() + "] got exception: " + str(err) + '.')
self._logger.error("Write to dir [" + self.get_dir_name() + "] got exception: " + str(err) + '.')
raise StorageWriteException("Write to dir [" + self.get_dir_name() + "] got exception: " + str(err) + '.')