228 lines
6.2 KiB
Python
228 lines
6.2 KiB
Python
from .columns.util import get_inner_spec, get_inner_columns_with_types
|
|
from .reader import read_varint, read_binary_uint8, read_binary_int32
|
|
from .varint import write_varint
|
|
from .writer import write_binary_uint8, write_binary_int32
|
|
|
|
|
|
class BlockInfo(object):
|
|
is_overflows = False
|
|
bucket_num = -1
|
|
|
|
def write(self, buf):
|
|
# Set of pairs (`FIELD_NUM`, value) in binary form. Then 0.
|
|
write_varint(1, buf)
|
|
write_binary_uint8(self.is_overflows, buf)
|
|
|
|
write_varint(2, buf)
|
|
write_binary_int32(self.bucket_num, buf)
|
|
|
|
write_varint(0, buf)
|
|
|
|
def read(self, buf):
|
|
while True:
|
|
field_num = read_varint(buf)
|
|
if not field_num:
|
|
break
|
|
|
|
if field_num == 1:
|
|
self.is_overflows = bool(read_binary_uint8(buf))
|
|
|
|
elif field_num == 2:
|
|
self.bucket_num = read_binary_int32(buf)
|
|
|
|
|
|
class BaseBlock(object):
|
|
def __init__(self, columns_with_types=None, data=None,
|
|
info=None, types_check=False):
|
|
self.columns_with_types = columns_with_types or []
|
|
self.types_check = types_check
|
|
self.info = info or BlockInfo()
|
|
self.data = self.normalize(data or [])
|
|
|
|
super(BaseBlock, self).__init__()
|
|
|
|
def normalize(self, data):
|
|
return data
|
|
|
|
@property
|
|
def num_columns(self):
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def num_rows(self):
|
|
raise NotImplementedError
|
|
|
|
def get_columns(self):
|
|
raise NotImplementedError
|
|
|
|
def get_rows(self):
|
|
raise NotImplementedError
|
|
|
|
def get_column_by_index(self, index):
|
|
raise NotImplementedError
|
|
|
|
def transposed(self):
|
|
return list(zip(*self.data))
|
|
|
|
|
|
class ColumnOrientedBlock(BaseBlock):
|
|
def normalize(self, data):
|
|
if not data:
|
|
return []
|
|
|
|
self._check_number_of_columns(data)
|
|
self._check_all_columns_equal_length(data)
|
|
return data
|
|
|
|
@property
|
|
def num_columns(self):
|
|
return len(self.data)
|
|
|
|
@property
|
|
def num_rows(self):
|
|
return len(self.data[0]) if self.num_columns else 0
|
|
|
|
def get_columns(self):
|
|
return self.data
|
|
|
|
def get_rows(self):
|
|
return self.transposed()
|
|
|
|
def get_column_by_index(self, index):
|
|
return self.data[index]
|
|
|
|
def _check_number_of_columns(self, data):
|
|
expected_row_len = len(self.columns_with_types)
|
|
|
|
got = len(data)
|
|
if expected_row_len != got:
|
|
msg = 'Expected {} columns, got {}'.format(expected_row_len, got)
|
|
raise ValueError(msg)
|
|
|
|
def _check_all_columns_equal_length(self, data):
|
|
expected = len(data[0])
|
|
|
|
for column in data:
|
|
got = len(column)
|
|
if got != expected:
|
|
msg = 'Expected {} rows, got {}'.format(expected, got)
|
|
raise ValueError(msg)
|
|
|
|
|
|
class RowOrientedBlock(BaseBlock):
|
|
dict_row_types = (dict, )
|
|
tuple_row_types = (list, tuple)
|
|
supported_row_types = dict_row_types + tuple_row_types
|
|
|
|
def normalize(self, data):
|
|
if not data:
|
|
return []
|
|
|
|
# Guessing about whole data format by first row.
|
|
first_row = data[0]
|
|
|
|
if self.types_check:
|
|
self._check_row_type(first_row)
|
|
|
|
if isinstance(first_row, dict):
|
|
self._mutate_dicts_to_rows(data)
|
|
else:
|
|
self._check_rows(data)
|
|
|
|
return data
|
|
|
|
@property
|
|
def num_columns(self):
|
|
if self.columns_with_types is not None:
|
|
return len(self.columns_with_types)
|
|
|
|
return len(self.data[0]) if self.num_rows else 0
|
|
|
|
@property
|
|
def num_rows(self):
|
|
return len(self.data)
|
|
|
|
def get_columns(self):
|
|
return self.transposed()
|
|
|
|
def get_rows(self):
|
|
return self.data
|
|
|
|
def get_column_by_index(self, index):
|
|
return [row[index] for row in self.data]
|
|
|
|
def _mutate_dicts_to_rows(self, data):
|
|
check_row_type = False
|
|
if self.types_check:
|
|
check_row_type = self._check_dict_row_type
|
|
|
|
return self._pure_mutate_dicts_to_rows(
|
|
data,
|
|
self.columns_with_types,
|
|
check_row_type,
|
|
)
|
|
|
|
def _pure_mutate_dicts_to_rows(
|
|
self,
|
|
data,
|
|
columns_with_types,
|
|
check_row_type,
|
|
):
|
|
columns_with_cwt = []
|
|
for name, type_ in columns_with_types:
|
|
cwt = None
|
|
if type_.startswith('Nested'):
|
|
inner_spec = get_inner_spec('Nested', type_)
|
|
cwt = get_inner_columns_with_types(inner_spec)
|
|
columns_with_cwt.append((name, cwt))
|
|
|
|
for i, row in enumerate(data):
|
|
if check_row_type:
|
|
check_row_type(row)
|
|
|
|
new_data = []
|
|
for name, cwt in columns_with_cwt:
|
|
if cwt is None:
|
|
new_data.append(row[name])
|
|
else:
|
|
new_data.append(self._pure_mutate_dicts_to_rows(
|
|
row[name], cwt, check_row_type
|
|
))
|
|
data[i] = new_data
|
|
# return for recursion
|
|
return data
|
|
|
|
def _check_rows(self, data):
|
|
expected_row_len = len(self.columns_with_types)
|
|
|
|
got = len(data[0])
|
|
if expected_row_len != got:
|
|
msg = 'Expected {} columns, got {}'.format(expected_row_len, got)
|
|
raise ValueError(msg)
|
|
|
|
if self.types_check:
|
|
check_row_type = self._check_tuple_row_type
|
|
for row in data:
|
|
check_row_type(row)
|
|
|
|
def _check_row_type(self, row):
|
|
if not isinstance(row, self.supported_row_types):
|
|
raise TypeError(
|
|
'Unsupported row type: {}. dict, list or tuple is expected.'
|
|
.format(type(row))
|
|
)
|
|
|
|
def _check_tuple_row_type(self, row):
|
|
if not isinstance(row, self.tuple_row_types):
|
|
raise TypeError(
|
|
'Unsupported row type: {}. list or tuple is expected.'
|
|
.format(type(row))
|
|
)
|
|
|
|
def _check_dict_row_type(self, row):
|
|
if not isinstance(row, self.dict_row_types):
|
|
raise TypeError(
|
|
'Unsupported row type: {}. dict is expected.'
|
|
.format(type(row))
|
|
)
|