1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
|
import numpy as np
import pandas as pd
def column_chunks(columns, n):
for column in columns:
if not isinstance(column, (np.ndarray, pd.DatetimeIndex)):
raise TypeError(
'Unsupported column type: {}. '
'ndarray/DatetimeIndex is expected.'
.format(type(column))
)
# create chunk generator for every column
chunked = [
iter(np.array_split(c, len(c) // n) if len(c) > n else [c])
for c in columns
]
while True:
# get next chunk for every column
item = [next(column, []) for column in chunked]
if not any(len(x) for x in item):
break
yield item
|
