| 22 | |
| 23 | |
| 24 | class Rechunk(ArrayExpr): |
| 25 | _parameters = [ |
| 26 | "array", |
| 27 | "_chunks", |
| 28 | "threshold", |
| 29 | "block_size_limit", |
| 30 | "balance", |
| 31 | "method", |
| 32 | ] |
| 33 | |
| 34 | _defaults = { |
| 35 | "_chunks": "auto", |
| 36 | "threshold": None, |
| 37 | "block_size_limit": None, |
| 38 | "balance": None, |
| 39 | "method": None, |
| 40 | } |
| 41 | |
| 42 | @property |
| 43 | def _meta(self): |
| 44 | return self.array._meta |
| 45 | |
| 46 | @property |
| 47 | def _name(self): |
| 48 | return "rechunk-merge-" + tokenize(*self.operands) |
| 49 | |
| 50 | @cached_property |
| 51 | def chunks(self): |
| 52 | x = self.array |
| 53 | chunks = self.operand("_chunks") |
| 54 | |
| 55 | # don't rechunk if array is empty |
| 56 | if x.ndim > 0 and all(s == 0 for s in x.shape): |
| 57 | return x.chunks |
| 58 | |
| 59 | if isinstance(chunks, dict): |
| 60 | chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()} |
| 61 | for i in range(x.ndim): |
| 62 | if i not in chunks: |
| 63 | chunks[i] = x.chunks[i] |
| 64 | elif chunks[i] is None: |
| 65 | chunks[i] = x.chunks[i] |
| 66 | if isinstance(chunks, (tuple, list)): |
| 67 | chunks = tuple( |
| 68 | lc if lc is not None else rc for lc, rc in zip(chunks, x.chunks) |
| 69 | ) |
| 70 | chunks = normalize_chunks( |
| 71 | chunks, |
| 72 | x.shape, |
| 73 | limit=self.block_size_limit, |
| 74 | dtype=x.dtype, |
| 75 | previous_chunks=x.chunks, |
| 76 | ) |
| 77 | |
| 78 | if not len(chunks) == x.ndim: |
| 79 | raise ValueError("Provided chunks are not consistent with shape") |
| 80 | |
| 81 | if self.balance: |
no outgoing calls
no test coverage detected
searching dependent graphs…