[BEAM-9547] Roll forward #12858 (#12920)

* [BEAM-9547] Raise NotImplementedError and WontImplementError throughout DeferredDataframe (#12858)

* Draft of NotImplementedErrors for DataFrame

* Default to jira BEAM-9547

* pivot is NotImplementedError

* remove mistaken change to frames

* Don't use __ror__ to convert df to PCollection, it's already a DataFrame operator

Author: TheNeuralBit

sdks/python/apache_beam/dataframe/frame_base.py

@@ -251,6 +251,13 @@ def wrapper(self, *args, **kwargs):
   return wrapper
 
 
+def not_implemented_method(op, jira='BEAM-9547'):
+  def wrapper(self, *args, **kwargs):
+    raise NotImplementedError("'%s' is not yet supported (%s)" % (op, jira))
+
+  return wrapper
+
+
 def copy_and_mutate(func):
   def wrapper(self, *args, **kwargs):
     copy = self.copy()

sdks/python/apache_beam/dataframe/frames.py

@@ -70,6 +70,8 @@ def fillna(self, value, method):
             preserves_partition_by=partitionings.Singleton(),
             requires_partition_by=partitionings.Nothing()))
 
+  reindex = frame_base.not_implemented_method('reindex')
+
   to_numpy = to_string = frame_base.wont_implement_method('non-deferred value')
 
   transform = frame_base._elementwise_method(
@@ -117,6 +119,11 @@ def aggregate(self, func, axis=0, *args, **kwargs):
 
   head = tail = frame_base.wont_implement_method('order-sensitive')
 
+  memory_usage = frame_base.wont_implement_method('non-deferred value')
+
+  # In Series __contains__ checks the index
+  __contains__ = frame_base.wont_implement_method('non-deferred value')
+
   @frame_base.args_to_kwargs(pd.Series)
   @frame_base.populate_defaults(pd.Series)
   def nlargest(self, keep, **kwargs):
@@ -226,29 +233,6 @@ def str(self):
     return _DeferredStringMethods(expr)
 
 
-for base in ['add',
-             'sub',
-             'mul',
-             'div',
-             'truediv',
-             'floordiv',
-             'mod',
-             'pow',
-             'and',
-             'or']:
-  for p in ['%s', 'r%s', '__%s__', '__r%s__']:
-    # TODO: non-trivial level?
-    name = p % base
-    setattr(
-        DeferredSeries,
-        name,
-        frame_base._elementwise_method(name, restrictions={'level': None}))
-  setattr(
-      DeferredSeries,
-      '__i%s__' % base,
-      frame_base._elementwise_method('__i%s__' % base, inplace=True))
-for name in ['__lt__', '__le__', '__gt__', '__ge__', '__eq__', '__ne__']:
-  setattr(DeferredSeries, name, frame_base._elementwise_method(name))
 for name in ['apply', 'map', 'transform']:
   setattr(DeferredSeries, name, frame_base._elementwise_method(name))
 
@@ -259,6 +243,10 @@ class DeferredDataFrame(frame_base.DeferredFrame):
   def T(self):
     return self.transpose()
 
+  @property
+  def columns(self):
+    return self._expr.proxy().columns
+
   def groupby(self, by):
     # TODO: what happens to the existing index?
     # We set the columns to index as we have a notion of being partitioned by
@@ -280,13 +268,24 @@ def __getattr__(self, name):
 
   def __getitem__(self, key):
     # TODO: Replicate pd.DataFrame.__getitem__ logic
+    if isinstance(key, frame_base.DeferredBase):
+      # Fail early if key is a DeferredBase as it interacts surprisingly with
+      # key in self._expr.proxy().columns
+      raise NotImplementedError(
+          "Indexing with a deferred frame is not yet supported. Consider "
+          "using df.loc[...]")
+
     if (isinstance(key, list) and
         all(key_column in self._expr.proxy().columns
             for key_column in key)) or key in self._expr.proxy().columns:
       return self._elementwise(lambda df: df[key], 'get_column')
     else:
       raise NotImplementedError(key)
 
+  def __contains__(self, key):
+    # Checks if proxy has the given column
+    return self._expr.proxy().__contains__(key)
+
   def __setitem__(self, key, value):
     if isinstance(key, str):
       # yapf: disable
@@ -314,13 +313,37 @@ def set_index(self, keys, **kwargs):
           requires_partition_by=partitionings.Nothing(),
           preserves_partition_by=partitionings.Nothing()))
 
-  def at(self, *args, **kwargs):
-    raise NotImplementedError()
+  at = frame_base.not_implemented_method('at')
 
   @property
   def loc(self):
     return _DeferredLoc(self)
 
+  _get_index = _set_index = frame_base.not_implemented_method('index')
+  index = property(_get_index, _set_index)
+
+  @property
+  def axes(self):
+    return (self.index, self.columns)
+
+  apply = frame_base.not_implemented_method('apply')
+  explode = frame_base.not_implemented_method('explode')
+  isin = frame_base.not_implemented_method('isin')
+  assign = frame_base.not_implemented_method('assign')
+  append = frame_base.not_implemented_method('append')
+  combine = frame_base.not_implemented_method('combine')
+  combine_first = frame_base.not_implemented_method('combine_first')
+  cov = frame_base.not_implemented_method('cov')
+  corr = frame_base.not_implemented_method('corr')
+  count = frame_base.not_implemented_method('count')
+  dot = frame_base.not_implemented_method('dot')
+  drop = frame_base.not_implemented_method('drop')
+  eval = frame_base.not_implemented_method('eval')
+  reindex = frame_base.not_implemented_method('reindex')
+  melt = frame_base.not_implemented_method('melt')
+  pivot = frame_base.not_implemented_method('pivot')
+  pivot_table = frame_base.not_implemented_method('pivot_table')
+
   def aggregate(self, func, axis=0, *args, **kwargs):
     if axis is None:
       # Aggregate across all elements by first aggregating across columns,
@@ -383,6 +406,7 @@ def aggregate(self, func, axis=0, *args, **kwargs):
   applymap = frame_base._elementwise_method('applymap')
 
   memory_usage = frame_base.wont_implement_method('non-deferred value')
+  info = frame_base.wont_implement_method('non-deferred value')
 
   all = frame_base._agg_method('all')
   any = frame_base._agg_method('any')
@@ -398,6 +422,8 @@ def aggregate(self, func, axis=0, *args, **kwargs):
 
   def mode(self, axis=0, *args, **kwargs):
     if axis == 1 or axis == 'columns':
+      # Number of columns is max(number mode values for each row), so we can't
+      # determine how many there will be before looking at the data.
       raise frame_base.WontImplementError('non-deferred column values')
     return frame_base.DeferredFrame.wrap(
         expressions.ComputedExpression(
@@ -766,8 +792,7 @@ def sort_values(self, axis, **kwargs):
   transform = frame_base._elementwise_method(
       'transform', restrictions={'axis': 0})
 
-  def transpose(self, *args, **kwargs):
-    raise frame_base.WontImplementError('non-deferred column values')
+  transpose = frame_base.wont_implement_method('non-deferred column values')
 
   def unstack(self, *args, **kwargs):
     if self._expr.proxy().index.nlevels == 1:
@@ -799,7 +824,10 @@ def unstack(self, *args, **kwargs):
 class DeferredGroupBy(frame_base.DeferredFrame):
   def agg(self, fn):
     if not callable(fn):
-      raise NotImplementedError(fn)
+      # TODO: Add support for strings in (UN)LIFTABLE_AGGREGATIONS. Test by
+      # running doctests for pandas.core.groupby.generic
+      raise NotImplementedError('GroupBy.agg currently only supports callable '
+                                'arguments')
     return DeferredDataFrame(
         expressions.ComputedExpression(
             'agg',
@@ -963,3 +991,37 @@ class _DeferredStringMethods(frame_base.DeferredBase):
   setattr(_DeferredStringMethods,
           method,
           frame_base._elementwise_method(method))
+
+for base in ['add',
+             'sub',
+             'mul',
+             'div',
+             'truediv',
+             'floordiv',
+             'mod',
+             'pow',
+             'and',
+             'or']:
+  for p in ['%s', 'r%s', '__%s__', '__r%s__']:
+    # TODO: non-trivial level?
+    name = p % base
+    setattr(
+        DeferredSeries,
+        name,
+        frame_base._elementwise_method(name, restrictions={'level': None}))
+    setattr(
+        DeferredDataFrame,
+        name,
+        frame_base._elementwise_method(name, restrictions={'level': None}))
+  setattr(
+      DeferredSeries,
+      '__i%s__' % base,
+      frame_base._elementwise_method('__i%s__' % base, inplace=True))
+  setattr(
+      DeferredDataFrame,
+      '__i%s__' % base,
+      frame_base._elementwise_method('__i%s__' % base, inplace=True))
+
+for name in ['__lt__', '__le__', '__gt__', '__ge__', '__eq__', '__ne__']:
+  setattr(DeferredSeries, name, frame_base._elementwise_method(name))
+  setattr(DeferredDataFrame, name, frame_base._elementwise_method(name))

sdks/python/apache_beam/dataframe/io.py

@@ -40,8 +40,14 @@ def read_csv(path, *args, **kwargs):
   return _ReadFromPandas(pd.read_csv, path, args, kwargs, incremental=True)
 
 
+def _as_pc(df):
+  from apache_beam.dataframe import convert  # avoid circular import
+  # TODO(roberwb): Amortize the computation for multiple writes?
+  return convert.to_pcollection(df, yield_elements='pandas')
+
+
 def to_csv(df, path, *args, **kwargs):
-  return df | _WriteToPandas(
+  return _as_pc(df) | _WriteToPandas(
       pd.DataFrame.to_csv, path, args, kwargs, incremental=True, binary=False)
 
 
@@ -68,7 +74,7 @@ def to_json(df, path, orient=None, *args, **kwargs):
     else:
       raise frame_base.WontImplementError('not dataframes or series')
   kwargs['orient'] = orient
-  return df | _WriteToPandas(
+  return _as_pc(df) | _WriteToPandas(
       pd.DataFrame.to_json,
       path,
       args,
@@ -87,7 +93,7 @@ def read_html(path, *args, **kwargs):
 
 
 def to_html(df, path, *args, **kwargs):
-  return df | _WriteToPandas(
+  return _as_pc(df) | _WriteToPandas(
       pd.DataFrame.to_html,
       path,
       args,
@@ -109,7 +115,7 @@ def _binary_writer(format):
       lambda df,
       path,
       *args,
-      **kwargs: df | _WriteToPandas(func, path, args, kwargs))
+      **kwargs: _as_pc(df) | _WriteToPandas(func, path, args, kwargs))
 
 
 for format in ('excel', 'feather', 'parquet', 'stata'):
@@ -214,13 +220,6 @@ def __init__(
     self.incremental = incremental
     self.binary = binary
 
-  def __ror__(self, other, label=None):
-    if isinstance(other, frame_base.DeferredBase):
-      from apache_beam.dataframe import convert  # avoid circular import
-      # TODO(roberwb): Amortize the computation for multiple writes?
-      other = convert.to_pcollection(other, yield_elements='pandas')
-    return super(_WriteToPandas, self).__ror__(other, label)
-
   def expand(self, pcoll):
     dir, name = io.filesystems.FileSystems.split(self.path)
     return pcoll | fileio.WriteToFiles(