TheNeuralBit commented on a change in pull request #11766: URL: https://github.com/apache/beam/pull/11766#discussion_r435426414
########## File path: sdks/python/apache_beam/dataframe/partitionings.py ########## @@ -0,0 +1,136 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import absolute_import + +from typing import Any +from typing import Iterable +from typing import TypeVar + +import pandas as pd + +Frame = TypeVar('Frame', bound=pd.core.generic.NDFrame) + + +class Partitioning(object): + """A class representing a (consistent) partitioning of dataframe objects. + """ + def is_subpartitioning_of(self, other): + # type: (Partitioning) -> bool + + """Returns whether self is a sub-partition of other. + + Specifically, returns whether something partitioned by self is necissarily + also partitioned by other. + """ + raise NotImplementedError + + def partition_fn(self, df): + # type: (Frame) -> Iterable[Tuple[Any, Frame]] + + """A callable that actually performs the partitioning of a Frame df. + + This will be invoked via a FlatMap in conjunction with a GroupKey to + achieve the desired partitioning. + """ + raise NotImplementedError + + +class Index(Partitioning): + """A partitioning by index (either fully or partially). + + If the set of "levels" of the index to consider is not specified, the entire + index is used. + + These form a partial order, given by + + Nothing() < Index([i]) < Index([i, j]) < ... < Index() < Singleton() + + The ordering is implemented via the is_subpartitioning_of method, where the + examples on the right are subpartitionings of the examples on the left above. + """ + + _INDEX_PARTITIONS = 10 + + def __init__(self, levels=None): + self._levels = levels + + def __eq__(self, other): + return type(self) == type(other) and self._levels == other._levels + + def __hash__(self): + if self._levels: + return hash(tuple(sorted(self._levels))) + else: + return hash(type(self)) + + def is_subpartitioning_of(self, other): + if isinstance(other, Nothing): + return True + elif isinstance(other, Index): + if self._levels is None: + return True + elif other._levels is None: + return False + else: + return all(level in other._levels for level in self._levels) + else: + return False + + def partition_fn(self, df): + if self._levels is None: + levels = list(range(df.index.nlevels)) + else: + levels = self._levels + hashes = sum( + pd.util.hash_array(df.index.get_level_values(level)) + for level in levels) + for key in range(self._INDEX_PARTITIONS): + yield key, df[hashes % self._INDEX_PARTITIONS == key] + + +class Singleton(Partitioning): + """A partitioning co-locating all data to a singleton partition. + """ + def __eq__(self, other): + return type(self) == type(other) + + def __hash__(self): + return hash(type(self)) + + def is_subpartitioning_of(self, other): + return True + + def partition_fn(self, df): + yield None, df + + +class Nothing(Partitioning): + """A partitioning imposing no constraints on the actual partitioning. + """ + def __eq__(self, other): + return type(self) == type(other) + + def __hash__(self): + return hash(type(self)) + + def is_subpartitioning_of(self, other): + return not other + + def __bool__(self): + return False + + __nonzero__ = __bool__ Review comment: I think that making Nothing falsy and relying on that in logic elsewhere harms readability. What do you think about dropping this and just explicitly checking for Nothing when needed? ########## File path: sdks/python/apache_beam/dataframe/transforms.py ########## @@ -138,36 +140,35 @@ def evaluate(partition, stage=self.stage): class Stage(object): """Used to build up a set of operations that can be fused together. """ - def __init__(self, inputs, is_grouping): + def __init__(self, inputs, partitioning): self.inputs = set(inputs) - self.is_grouping = is_grouping or len(self.inputs) > 1 + if len(self.inputs) > 1 and not partitioning: + # We have to shuffle to co-locate, might as well partition. + self.partitioning = partitionings.Index() + else: + self.partitioning = partitioning self.ops = [] self.outputs = set() # First define some helper functions. - def output_is_partitioned_by_index(expr, stage): - if expr in stage.inputs: - return stage.is_grouping - elif expr.preserves_partition_by_index(): - if expr.requires_partition_by_index(): + def output_is_partitioned_by(expr, stage, partitioning): + if not partitioning: + return True + elif stage.partitioning == partitionings.Singleton(): + # Within a stage, the singleton partitioning is trivially preserved. + return True + elif expr in stage.inputs: + return stage.partitioning.is_subpartitioning_of(partitioning) + elif expr.preserves_partition_by().is_subpartitioning_of(partitioning): + if expr.requires_partition_by().is_subpartitioning_of(partitioning): Review comment: Had trouble justifying this logic to myself, ended up writing a little proof which I'm a bit embarrassed to share with a Math PhD: ``` output partitioning of expr = min(expr.preserves, input partitioning) input partitioning >= expr.requires thus if expr.requires >= required output AND expr.preserves >= required output then output partitioning of expr >= required output Otherwise we need to go up the tree of inputs to figure out their partitionings ``` This may be the least concise way to express this so I don't know if it's worth putting in a comment verbatim, but something to that effect would be helpful (assuming I've got it right) ########## File path: sdks/python/apache_beam/dataframe/partitionings.py ########## @@ -0,0 +1,136 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import absolute_import + +from typing import Any +from typing import Iterable +from typing import TypeVar + +import pandas as pd + +Frame = TypeVar('Frame', bound=pd.core.generic.NDFrame) + + +class Partitioning(object): + """A class representing a (consistent) partitioning of dataframe objects. + """ + def is_subpartitioning_of(self, other): + # type: (Partitioning) -> bool + + """Returns whether self is a sub-partition of other. + + Specifically, returns whether something partitioned by self is necissarily + also partitioned by other. + """ + raise NotImplementedError + + def partition_fn(self, df): + # type: (Frame) -> Iterable[Tuple[Any, Frame]] + + """A callable that actually performs the partitioning of a Frame df. + + This will be invoked via a FlatMap in conjunction with a GroupKey to + achieve the desired partitioning. + """ + raise NotImplementedError + + +class Index(Partitioning): + """A partitioning by index (either fully or partially). + + If the set of "levels" of the index to consider is not specified, the entire + index is used. + + These form a partial order, given by + + Nothing() < Index([i]) < Index([i, j]) < ... < Index() < Singleton() + + The ordering is implemented via the is_subpartitioning_of method, where the + examples on the right are subpartitionings of the examples on the left above. + """ + + _INDEX_PARTITIONS = 10 + + def __init__(self, levels=None): + self._levels = levels + + def __eq__(self, other): + return type(self) == type(other) and self._levels == other._levels + + def __hash__(self): + if self._levels: + return hash(tuple(sorted(self._levels))) + else: + return hash(type(self)) + + def is_subpartitioning_of(self, other): + if isinstance(other, Nothing): + return True + elif isinstance(other, Index): + if self._levels is None: + return True + elif other._levels is None: + return False + else: + return all(level in other._levels for level in self._levels) + else: + return False + + def partition_fn(self, df): + if self._levels is None: + levels = list(range(df.index.nlevels)) + else: + levels = self._levels + hashes = sum( + pd.util.hash_array(df.index.get_level_values(level)) + for level in levels) + for key in range(self._INDEX_PARTITIONS): + yield key, df[hashes % self._INDEX_PARTITIONS == key] + + +class Singleton(Partitioning): + """A partitioning co-locating all data to a singleton partition. Review comment: Isn't `Singleton` completely partitioning the data into one element per partition? This description doesn't seem consistent to me, maybe I'm misunderstanding ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
