feat: add 'index', 'pad', 'nearest' interpolate methods (#162)

* feat: add 'index', 'pad', 'nearest' interpolate methods

Author: TrevorBergeron

bigframes/core/block_transforms.py

@@ -22,6 +22,7 @@
 import bigframes.core.blocks as blocks
 import bigframes.core.ordering as ordering
 import bigframes.core.window_spec as windows
+import bigframes.dtypes as dtypes
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
 
@@ -106,67 +107,59 @@ def indicate_duplicates(
 
 
 def interpolate(block: blocks.Block, method: str = "linear") -> blocks.Block:
- if method != "linear":
+ supported_methods = [
+ "linear",
+ "values",
+ "index",
+ "nearest",
+ "zero",
+ "slinear",
+ ]
+ if method not in supported_methods:
  raise NotImplementedError(
- f"Only 'linear' interpolate method supported. {constants.FEEDBACK_LINK}"
+ f"Method {method} not supported, following interpolate methods supported: {', '.join(supported_methods)}. {constants.FEEDBACK_LINK}"
  )
- backwards_window = windows.WindowSpec(following=0)
- forwards_window = windows.WindowSpec(preceding=0)
-
  output_column_ids = []
 
  original_columns = block.value_columns
  original_labels = block.column_labels
- block, offsets = block.promote_offsets()
+
+ if method == "linear": # Assumes evenly spaced, ignore index
+ block, xvalues = block.promote_offsets()
+ else:
+ index_columns = block.index_columns
+ if len(index_columns) != 1:
+ raise ValueError("only method 'linear' supports multi-index")
+ xvalues = block.index_columns[0]
+ if block.index_dtypes[0] not in dtypes.NUMERIC_BIGFRAMES_TYPES:
+ raise ValueError("Can only interpolate on numeric index.")
+
  for column in original_columns:
  # null in same places column is null
  should_interpolate = block._column_type(column) in [
  pd.Float64Dtype(),
  pd.Int64Dtype(),
  ]
  if should_interpolate:
- block, notnull = block.apply_unary_op(column, ops.notnull_op)
- block, masked_offsets = block.apply_binary_op(
- offsets, notnull, ops.partial_arg3(ops.where_op, None)
- )
-
- block, previous_value = block.apply_window_op(
- column, agg_ops.LastNonNullOp(), backwards_window
- )
- block, next_value = block.apply_window_op(
- column, agg_ops.FirstNonNullOp(), forwards_window
- )
- block, previous_value_offset = block.apply_window_op(
- masked_offsets,
- agg_ops.LastNonNullOp(),
- backwards_window,
- skip_reproject_unsafe=True,
- )
- block, next_value_offset = block.apply_window_op(
- masked_offsets,
- agg_ops.FirstNonNullOp(),
- forwards_window,
- skip_reproject_unsafe=True,
- )
-
- block, prediction_id = _interpolate(
+ interpolate_method_map = {
+ "linear": "linear",
+ "values": "linear",
+ "index": "linear",
+ "slinear": "linear",
+ "zero": "ffill",
+ "nearest": "nearest",
+ }
+ extrapolating_methods = ["linear", "values", "index"]
+ interpolate_method = interpolate_method_map[method]
+ do_extrapolate = method in extrapolating_methods
+ block, interpolated = _interpolate_column(
  block,
- previous_value_offset,
- previous_value,
- next_value_offset,
- next_value,
- offsets,
+ column,
+ xvalues,
+ interpolate_method=interpolate_method,
+ do_extrapolate=do_extrapolate,
  )
-
- block, interpolated_column = block.apply_binary_op(
- column, prediction_id, ops.fillna_op
- )
- # Pandas performs ffill-like behavior to extrapolate forwards
- block, interpolated_and_ffilled = block.apply_binary_op(
- interpolated_column, previous_value, ops.fillna_op
- )
-
- output_column_ids.append(interpolated_and_ffilled)
+ output_column_ids.append(interpolated)
  else:
  output_column_ids.append(column)
 
@@ -175,7 +168,80 @@ def interpolate(block: blocks.Block, method: str = "linear") -> blocks.Block:
  return block.with_column_labels(original_labels)
 
 
-def _interpolate(
+def _interpolate_column(
+ block: blocks.Block,
+ column: str,
+ x_values: str,
+ interpolate_method: str,
+ do_extrapolate: bool = True,
+) -> typing.Tuple[blocks.Block, str]:
+ if interpolate_method not in ["linear", "nearest", "ffill"]:
+ raise ValueError("interpolate method not supported")
+ window_ordering = (ordering.OrderingColumnReference(x_values),)
+ backwards_window = windows.WindowSpec(following=0, ordering=window_ordering)
+ forwards_window = windows.WindowSpec(preceding=0, ordering=window_ordering)
+
+ # Note, this method may
+ block, notnull = block.apply_unary_op(column, ops.notnull_op)
+ block, masked_offsets = block.apply_binary_op(
+ x_values, notnull, ops.partial_arg3(ops.where_op, None)
+ )
+
+ block, previous_value = block.apply_window_op(
+ column, agg_ops.LastNonNullOp(), backwards_window
+ )
+ block, next_value = block.apply_window_op(
+ column, agg_ops.FirstNonNullOp(), forwards_window
+ )
+ block, previous_value_offset = block.apply_window_op(
+ masked_offsets,
+ agg_ops.LastNonNullOp(),
+ backwards_window,
+ skip_reproject_unsafe=True,
+ )
+ block, next_value_offset = block.apply_window_op(
+ masked_offsets,
+ agg_ops.FirstNonNullOp(),
+ forwards_window,
+ skip_reproject_unsafe=True,
+ )
+
+ if interpolate_method == "linear":
+ block, prediction_id = _interpolate_points_linear(
+ block,
+ previous_value_offset,
+ previous_value,
+ next_value_offset,
+ next_value,
+ x_values,
+ )
+ elif interpolate_method == "nearest":
+ block, prediction_id = _interpolate_points_nearest(
+ block,
+ previous_value_offset,
+ previous_value,
+ next_value_offset,
+ next_value,
+ x_values,
+ )
+ else: # interpolate_method == 'ffill':
+ block, prediction_id = _interpolate_points_ffill(
+ block,
+ previous_value_offset,
+ previous_value,
+ next_value_offset,
+ next_value,
+ x_values,
+ )
+ if do_extrapolate:
+ block, prediction_id = block.apply_binary_op(
+ prediction_id, previous_value, ops.fillna_op
+ )
+
+ return block.apply_binary_op(column, prediction_id, ops.fillna_op)
+
+
+def _interpolate_points_linear(
  block: blocks.Block,
  x0_id: str,
  y0_id: str,
@@ -196,6 +262,53 @@ def _interpolate(
  return block, prediction_id
 
 
+def _interpolate_points_nearest(
+ block: blocks.Block,
+ x0_id: str,
+ y0_id: str,
+ x1_id: str,
+ y1_id: str,
+ xpredict_id: str,
+) -> typing.Tuple[blocks.Block, str]:
+ """Interpolate by taking the y value of the nearest x value"""
+ block, left_diff = block.apply_binary_op(xpredict_id, x0_id, ops.sub_op)
+ block, right_diff = block.apply_binary_op(x1_id, xpredict_id, ops.sub_op)
+ # If diffs equal, choose left
+ block, choose_left = block.apply_binary_op(left_diff, right_diff, ops.le_op)
+ block, choose_left = block.apply_unary_op(
+ choose_left, ops.partial_right(ops.fillna_op, False)
+ )
+
+ block, nearest = block.apply_ternary_op(y0_id, choose_left, y1_id, ops.where_op)
+
+ block, y0_exists = block.apply_unary_op(y0_id, ops.notnull_op)
+ block, y1_exists = block.apply_unary_op(y1_id, ops.notnull_op)
+ block, is_interpolation = block.apply_binary_op(y0_exists, y1_exists, ops.and_op)
+
+ block, prediction_id = block.apply_binary_op(
+ nearest, is_interpolation, ops.partial_arg3(ops.where_op, None)
+ )
+
+ return block, prediction_id
+
+
+def _interpolate_points_ffill(
+ block: blocks.Block,
+ x0_id: str,
+ y0_id: str,
+ x1_id: str,
+ y1_id: str,
+ xpredict_id: str,
+) -> typing.Tuple[blocks.Block, str]:
+ """Interpolates by using the preceding values"""
+ # check for existance of y1, otherwise we are extrapolating instead of interpolating
+ block, y1_exists = block.apply_unary_op(y1_id, ops.notnull_op)
+ block, prediction_id = block.apply_binary_op(
+ y0_id, y1_exists, ops.partial_arg3(ops.where_op, None)
+ )
+ return block, prediction_id
+
+
 def drop_duplicates(
  block: blocks.Block, columns: typing.Sequence[str], keep: str = "first"
 ) -> blocks.Block:

bigframes/dataframe.py

@@ -1450,6 +1450,8 @@ def reindex_like(self, other: DataFrame, *, validate: typing.Optional[bool] = No
  return self.reindex(index=other.index, columns=other.columns, validate=validate)
 
  def interpolate(self, method: str = "linear") -> DataFrame:
+ if method == "pad":
+ return self.ffill()
  result = block_ops.interpolate(self._block, method)
  return DataFrame(result)
 

bigframes/series.py

@@ -475,6 +475,8 @@ def replace(
  return Series(block.select_column(result_col))
 
  def interpolate(self, method: str = "linear") -> Series:
+ if method == "pad":
+ return self.ffill()
  result = block_ops.interpolate(self._block, method)
  return Series(result)
 

tests/system/small/test_series.py

@@ -273,21 +273,26 @@ def test_series_replace_list_scalar(scalars_dfs):
 
 
 @pytest.mark.parametrize(
- ("values",),
+ ("method",),
  (
- ([None, 1, 2, None, None, 16, None],),
- ([None, None, 3.6, None],),
- ([403.2, None, 352.1, None, None, 111.9],),
+ ("linear",),
+ ("values",),
+ ("slinear",),
+ ("nearest",),
+ ("zero",),
+ ("pad",),
  ),
 )
-def test_series_interpolate(values):
- pd_series = pd.Series(values)
+def test_series_interpolate(method):
+ values = [None, 1, 2, None, None, 16, None]
+ index = [-3.2, 11.4, 3.56, 4, 4.32, 5.55, 76.8]
+ pd_series = pd.Series(values, index)
  bf_series = series.Series(pd_series)
 
  # Pandas can only interpolate on "float64" columns
  # https://github.com/pandas-dev/pandas/issues/40252
- pd_result = pd_series.astype("float64").interpolate()
- bf_result = bf_series.interpolate().to_pandas()
+ pd_result = pd_series.astype("float64").interpolate(method=method)
+ bf_result = bf_series.interpolate(method=method).to_pandas()
 
  # pd uses non-null types, while bf uses nullable types
  pd.testing.assert_series_equal(

third_party/bigframes_vendored/pandas/core/frame.py

@@ -2872,17 +2872,6 @@ def interpolate(self, method: str = "linear"):
  """
  Fill NaN values using an interpolation method.
 
- Args:
- method (str, default 'linear'):
- Interpolation technique to use. Only 'linear' supported.
- 'linear': Ignore the index and treat the values as equally spaced.
- This is the only method supported on MultiIndexes.
-
- Returns:
- DataFrame:
- Returns the same object type as the caller, interpolated at
- some or all ``NaN`` values
-
  **Examples:**
 
  >>> import bigframes.pandas as bpd
@@ -2891,17 +2880,41 @@ def interpolate(self, method: str = "linear"):
  >>> df = bpd.DataFrame({
  ... 'A': [1, 2, 3, None, None, 6],
  ... 'B': [None, 6, None, 2, None, 3],
- ... })
+ ... }, index=[0, 0.1, 0.3, 0.7, 0.9, 1.0])
  >>> df.interpolate()
- A B
- 0 1.0 <NA>
- 1 2.0 6.0
- 2 3.0 4.0
- 3 4.0 2.0
- 4 5.0 2.5
- 5 6.0 3.0
+  A B
+ 0.0 1.0 <NA>
+ 0.1 2.0 6.0
+ 0.3 3.0 4.0
+ 0.7 4.0 2.0
+ 0.9 5.0 2.5
+ 1.0 6.0 3.0
  <BLANKLINE>
  [6 rows x 2 columns]
+ >>> df.interpolate(method="values")
+ A B
+ 0.0 1.0 <NA>
+ 0.1 2.0 6.0
+ 0.3 3.0 4.666667
+ 0.7 4.714286 2.0
+ 0.9 5.571429 2.666667
+ 1.0 6.0 3.0
+ <BLANKLINE>
+ [6 rows x 2 columns]
+
+ Args:
+ method (str, default 'linear'):
+ Interpolation technique to use. Only 'linear' supported.
+ 'linear': Ignore the index and treat the values as equally spaced.
+ This is the only method supported on MultiIndexes.
+ 'index', 'values': use the actual numerical values of the index.
+ 'pad': Fill in NaNs using existing values.
+ 'nearest', 'zero', 'slinear': Emulates `scipy.interpolate.interp1d`
+
+ Returns:
+ DataFrame:
+ Returns the same object type as the caller, interpolated at
+ some or all ``NaN`` values
  """
  raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)