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