Select one or more partitions
| 3018 | |
| 3019 | |
| 3020 | class Partitions(Expr): |
| 3021 | """Select one or more partitions""" |
| 3022 | |
| 3023 | _parameters = ["frame", "partitions"] |
| 3024 | |
| 3025 | @functools.cached_property |
| 3026 | def _meta(self): |
| 3027 | return self.frame._meta |
| 3028 | |
| 3029 | def _divisions(self): |
| 3030 | divisions = [] |
| 3031 | for part in self.partitions: |
| 3032 | divisions.append(self.frame.divisions[part]) |
| 3033 | divisions.append(self.frame.divisions[part + 1]) |
| 3034 | return tuple(divisions) |
| 3035 | |
| 3036 | def _task(self, name: Key, index: int) -> Task: |
| 3037 | return Alias(name, (self.frame._name, self.partitions[index])) # type: ignore[return-value] |
| 3038 | |
| 3039 | def _simplify_down(self): |
| 3040 | from dask.dataframe.dask_expr import SetIndexBlockwise |
| 3041 | from dask.dataframe.tseries.resample import ResampleAggregation |
| 3042 | |
| 3043 | if isinstance(self.frame, Blockwise) and not isinstance( |
| 3044 | self.frame, (BlockwiseIO, Fused, SetIndexBlockwise, ResampleAggregation) |
| 3045 | ): |
| 3046 | operands = [ |
| 3047 | ( |
| 3048 | Partitions(op, self.partitions) |
| 3049 | if (isinstance(op, Expr) and not self.frame._broadcast_dep(op)) |
| 3050 | else op |
| 3051 | ) |
| 3052 | for op in self.frame.operands |
| 3053 | ] |
| 3054 | return type(self.frame)(*operands) |
| 3055 | elif isinstance(self.frame, PartitionsFiltered): |
| 3056 | if self.frame._partitions: |
| 3057 | partitions = [self.frame._partitions[p] for p in self.partitions] |
| 3058 | else: |
| 3059 | partitions = self.partitions |
| 3060 | # We assume that expressions defining a special "_partitions" |
| 3061 | # parameter can internally capture the same logic as `Partitions` |
| 3062 | return self.frame.substitute_parameters({"_partitions": partitions}) |
| 3063 | |
| 3064 | def _node_label_args(self): |
| 3065 | return [self.frame, self.partitions] |
| 3066 | |
| 3067 | |
| 3068 | class PartitionsFiltered(Expr): |