from collections.abc import Hashable
[docs]
class CategoricalColumnMap:
"""
Structure for assigning auto-incrementing indexes for unique values and creates a mapping of values to
indexes.
These values can be used to construct categorical columns to save memory.
Example::
>>> mapping = CategoricalColumnMap()
>>> mapping.get_index('a1')
0
>>> mapping.get_index('a2')
1
>>> mapping.get_index('a1')
0
>>> mapping.get_index('a2')
1
>>> mapping.get_index('a3')
2
>>> mapping.values
{'a1': 0, 'a2': 1, 'a3': 2}
"""
def __init__(self):
self._next_index = 0
self._values = {}
@property
def values(self) -> dict[Hashable, int]:
"""
:return: Dictionary with currently assigned values.
"""
return self._values
[docs]
def get_index(self, value: Hashable) -> int:
"""
Adds a new value to the mapping. If the value already exists, previous index is returned. If not, a new one
is assigned.
:param value: Value to add.
:return: Index assigned to the value.
"""
if value in self._values:
return self._values[value]
else:
index = self._next_index
self._next_index += 1
self._values[value] = index
return index