Select one or more partitions
| 2999 | |
| 3000 | |
| 3001 | class Partitions(Expr): |
| 3002 | """Select one or more partitions""" |
| 3003 | |
| 3004 | _parameters = ["frame", "partitions"] |
| 3005 | |
| 3006 | @functools.cached_property |
| 3007 | def _meta(self): |
| 3008 | return self.frame._meta |
| 3009 | |
| 3010 | def _divisions(self): |
| 3011 | divisions = [] |
| 3012 | for part in self.partitions: |
| 3013 | divisions.append(self.frame.divisions[part]) |
| 3014 | divisions.append(self.frame.divisions[part + 1]) |
| 3015 | return tuple(divisions) |
| 3016 | |
| 3017 | def _task(self, name: Key, index: int) -> Task: |
| 3018 | return Alias(name, (self.frame._name, self.partitions[index])) # type: ignore[return-value] |
| 3019 | |
| 3020 | def _simplify_down(self): |
| 3021 | from dask.dataframe.dask_expr import SetIndexBlockwise |
| 3022 | from dask.dataframe.tseries.resample import ResampleAggregation |
| 3023 | |
| 3024 | if isinstance(self.frame, Blockwise) and not isinstance( |
| 3025 | self.frame, (BlockwiseIO, Fused, SetIndexBlockwise, ResampleAggregation) |
| 3026 | ): |
| 3027 | operands = [ |
| 3028 | ( |
| 3029 | Partitions(op, self.partitions) |
| 3030 | if (isinstance(op, Expr) and not self.frame._broadcast_dep(op)) |
| 3031 | else op |
| 3032 | ) |
| 3033 | for op in self.frame.operands |
| 3034 | ] |
| 3035 | return type(self.frame)(*operands) |
| 3036 | elif isinstance(self.frame, PartitionsFiltered): |
| 3037 | if self.frame._partitions: |
| 3038 | partitions = [self.frame._partitions[p] for p in self.partitions] |
| 3039 | else: |
| 3040 | partitions = self.partitions |
| 3041 | # We assume that expressions defining a special "_partitions" |
| 3042 | # parameter can internally capture the same logic as `Partitions` |
| 3043 | return self.frame.substitute_parameters({"_partitions": partitions}) |
| 3044 | |
| 3045 | def _node_label_args(self): |
| 3046 | return [self.frame, self.partitions] |
| 3047 | |
| 3048 | |
| 3049 | class PartitionsFiltered(Expr): |