robertwb commented on a change in pull request #11766:
URL: https://github.com/apache/beam/pull/11766#discussion_r435632933
##########
File path: sdks/python/apache_beam/dataframe/expressions.py
##########
@@ -85,16 +87,10 @@ def evaluate_at(self, session): # type: (Session) -> T
"""Returns the result of self with the bindings given in session."""
raise NotImplementedError(type(self))
- def requires_partition_by_index(self): # type: () -> bool
- """Whether this expression requires its argument(s) to be partitioned
- by index."""
- # TODO: It might be necessary to support partitioning by part of the index,
- # for some args, which would require returning more than a boolean here.
+ def requires_partition_by(self): # type: () -> Partitioning
raise NotImplementedError(type(self))
- def preserves_partition_by_index(self): # type: () -> bool
- """Whether the result of this expression will be partitioned by index
- whenever all of its inputs are partitioned by index."""
+ def preserves_partition_by(self): # type: () -> Partitioning
Review comment:
Docstring comments added.
##########
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:
Reworded.
##########
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:
Done.
##########
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:
Yeah, fleshing this out with more comments.
----------------------------------------------------------------
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]
