ENH: store the 'codes' for a labelarray in the narrowest int type possible

This commit is contained in:
Joe Jevnik
2017-02-02 20:58:36 -05:00
parent d9321f4795
commit d378c9fca3
4 changed files with 236 additions and 80 deletions
+37 -1
View File
@@ -1,8 +1,10 @@
from itertools import product
from operator import eq, ne
import numpy as np
import warnings
import numpy as np
from toolz import take
from zipline.lib.labelarray import LabelArray
from zipline.testing import check_arrays, parameter_space, ZiplineTestCase
from zipline.utils.compat import unicode
@@ -337,3 +339,37 @@ class LabelArrayTestCase(ZiplineTestCase):
# Write the whole array.
arr[:] = orig_arr
check_arrays(arr, orig_arr)
def test_narrow_code_storage(self):
def create_categories(width, plus_one):
length = int(width / 8) + plus_one
return [
''.join(cs)
for cs in take(
2 ** width + plus_one,
product([chr(c) for c in range(256)], repeat=length),
)
]
# uint8
categories = create_categories(8, plus_one=False)
arr = LabelArray([], missing_value='', categories=categories)
self.assertEqual(arr.itemsize, 1)
# just over uint8
categories = create_categories(8, plus_one=True)
arr = LabelArray([], missing_value='', categories=categories)
self.assertEqual(arr.itemsize, 2)
# fits in uint16
categories = create_categories(16, plus_one=False)
arr = LabelArray([], missing_value='', categories=categories)
self.assertEqual(arr.itemsize, 2)
# just over uint16
categories = create_categories(16, plus_one=True)
arr = LabelArray([], missing_value='', categories=categories)
self.assertEqual(arr.itemsize, 4)
# NOTE: we could do this for 32 and 64; however, no one has enough RAM
# or time for that.