143 lines
4.8 KiB
Python
143 lines
4.8 KiB
Python
from math import log
|
|
|
|
from ..reader import read_binary_uint64
|
|
from ..writer import write_binary_int64
|
|
from .base import Column
|
|
from .intcolumn import UInt8Column, UInt16Column, UInt32Column, UInt64Column
|
|
|
|
|
|
def create_low_cardinality_column(spec, column_by_spec_getter, column_options):
|
|
inner = spec[15:-1]
|
|
nested = column_by_spec_getter(inner)
|
|
return LowCardinalityColumn(nested, **column_options)
|
|
|
|
|
|
class LowCardinalityColumn(Column):
|
|
"""
|
|
Stores column as index (unique elements) and keys.
|
|
Good for de-duplication of large values with low cardinality.
|
|
"""
|
|
int_types = {
|
|
0: UInt8Column,
|
|
1: UInt16Column,
|
|
2: UInt32Column,
|
|
3: UInt64Column
|
|
}
|
|
|
|
# Need to read additional keys.
|
|
# Additional keys are stored before indexes as value N and N keys
|
|
# after them.
|
|
has_additional_keys_bit = 1 << 9
|
|
# Need to update dictionary.
|
|
# It means that previous granule has different dictionary.
|
|
need_update_dictionary = 1 << 10
|
|
|
|
serialization_type = has_additional_keys_bit | need_update_dictionary
|
|
|
|
def __init__(self, nested_column, **kwargs):
|
|
self.init_kwargs = kwargs
|
|
self.nested_column = nested_column
|
|
super(LowCardinalityColumn, self).__init__(**kwargs)
|
|
|
|
def read_state_prefix(self, buf):
|
|
super(LowCardinalityColumn, self).read_state_prefix(buf)
|
|
|
|
read_binary_uint64(buf)
|
|
|
|
def write_state_prefix(self, buf):
|
|
super(LowCardinalityColumn, self).write_state_prefix(buf)
|
|
|
|
# KeysSerializationVersion. See ClickHouse docs.
|
|
write_binary_int64(1, buf)
|
|
|
|
def _write_data(self, items, buf):
|
|
index, keys = [], []
|
|
key_by_index_element = {}
|
|
nested_is_nullable = False
|
|
|
|
if self.nested_column.nullable:
|
|
# First element represents NULL if column is nullable.
|
|
index.append(self.nested_column.null_value)
|
|
# Prevent null map writing. Reset nested column nullable flag.
|
|
self.nested_column.nullable = False
|
|
nested_is_nullable = True
|
|
|
|
for x in items:
|
|
if x is None:
|
|
# Zero element for null.
|
|
keys.append(0)
|
|
|
|
else:
|
|
key = key_by_index_element.get(x)
|
|
# Get key from index or add it to index.
|
|
if key is None:
|
|
key = len(key_by_index_element)
|
|
key_by_index_element[x] = key
|
|
index.append(x)
|
|
|
|
keys.append(key + 1)
|
|
else:
|
|
for x in items:
|
|
key = key_by_index_element.get(x)
|
|
|
|
# Get key from index or add it to index.
|
|
if key is None:
|
|
key = len(key_by_index_element)
|
|
key_by_index_element[x] = len(key_by_index_element)
|
|
index.append(x)
|
|
|
|
keys.append(key)
|
|
|
|
# Do not write anything for empty column.
|
|
# May happen while writing empty arrays.
|
|
if not len(index):
|
|
return
|
|
|
|
int_type = int(log(len(index), 2) / 8)
|
|
int_column = self.int_types[int_type](**self.init_kwargs)
|
|
|
|
serialization_type = self.serialization_type | int_type
|
|
|
|
write_binary_int64(serialization_type, buf)
|
|
write_binary_int64(len(index), buf)
|
|
|
|
if nested_is_nullable:
|
|
# Given we reset nested column nullable flag above,
|
|
# we need to write null map manually. If to invoke
|
|
# write_data method, it will cause an exception,
|
|
# because `prepare_data` may not be able to handle
|
|
# null value correctly.
|
|
self.nested_column.write_items(
|
|
[self.nested_column.null_value], buf)
|
|
# Remove null map from index, because it is already written.
|
|
index_to_write = index[1:]
|
|
self.nested_column.write_data(index_to_write, buf)
|
|
else:
|
|
self.nested_column.write_data(index, buf)
|
|
write_binary_int64(len(items), buf)
|
|
int_column.write_items(keys, buf)
|
|
|
|
def _read_data(self, n_items, buf, nulls_map=None):
|
|
if not n_items:
|
|
return tuple()
|
|
|
|
serialization_type = read_binary_uint64(buf)
|
|
|
|
# Lowest byte contains info about key type.
|
|
key_type = serialization_type & 0xf
|
|
keys_column = self.int_types[key_type](**self.init_kwargs)
|
|
|
|
nullable = self.nested_column.nullable
|
|
# Prevent null map reading. Reset nested column nullable flag.
|
|
self.nested_column.nullable = False
|
|
|
|
index_size = read_binary_uint64(buf)
|
|
index = self.nested_column.read_data(index_size, buf)
|
|
if nullable:
|
|
index = (None, ) + index[1:]
|
|
|
|
read_binary_uint64(buf) # number of keys
|
|
keys = keys_column.read_data(n_items, buf)
|
|
|
|
return tuple(index[x] for x in keys)
|