SNOW-2432713: Ensure that dummy_row_pos_mode passed to OrderedDataFrame.join() matches SnowflakeQueryCompiler._dummy_row_pos_mode (#3901)

Author: sfc-gh-helmeleegy

src/snowflake/snowpark/modin/plugin/_internal/align_utils.py

@@ -13,7 +13,10 @@
 
 
 def align_axis_0_left(
-    frame: InternalFrame, other_frame: InternalFrame, join: str
+    frame: InternalFrame,
+    other_frame: InternalFrame,
+    join: str,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Gets the left align results.
@@ -27,7 +30,9 @@ def align_axis_0_left(
         New InternalFrame representing aligned left frame.
     """
     if join == "right":
-        left_result, left_column_mapper = align_on_index(other_frame, frame, how="left")
+        left_result, left_column_mapper = align_on_index(
+            other_frame, frame, dummy_row_pos_mode, how="left"
+        )
         left_frame_data_ids = left_column_mapper.map_right_quoted_identifiers(
             frame.data_column_snowflake_quoted_identifiers
         )
@@ -36,7 +41,9 @@ def align_axis_0_left(
             left_frame_data_ids + left_index_ids
         )
     else:
-        left_result, left_column_mapper = align_on_index(frame, other_frame, how=join)
+        left_result, left_column_mapper = align_on_index(
+            frame, other_frame, dummy_row_pos_mode, how=join
+        )
         left_frame_data_ids = left_column_mapper.map_left_quoted_identifiers(
             frame.data_column_snowflake_quoted_identifiers
         )
@@ -58,7 +65,10 @@ def align_axis_0_left(
 
 
 def align_axis_0_right(
-    frame: InternalFrame, other_frame: InternalFrame, join: str
+    frame: InternalFrame,
+    other_frame: InternalFrame,
+    join: str,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Gets the right align results.
@@ -72,7 +82,9 @@ def align_axis_0_right(
         New InternalFrame representing aligned right frame.
     """
     if join == "left":
-        right_result, right_column_mapper = align_on_index(frame, other_frame, how=join)
+        right_result, right_column_mapper = align_on_index(
+            frame, other_frame, dummy_row_pos_mode, how=join
+        )
         right_frame_data_ids = right_column_mapper.map_right_quoted_identifiers(
             other_frame.data_column_snowflake_quoted_identifiers
         )
@@ -82,7 +94,7 @@ def align_axis_0_right(
         )
     elif join == "right":
         right_result, right_column_mapper = align_on_index(
-            other_frame, frame, how="left"
+            other_frame, frame, dummy_row_pos_mode, how="left"
         )
         right_frame_data_ids = right_column_mapper.map_left_quoted_identifiers(
             other_frame.data_column_snowflake_quoted_identifiers
@@ -92,7 +104,9 @@ def align_axis_0_right(
             right_frame_data_ids + right_index_ids
         )
     else:
-        right_result, right_column_mapper = align_on_index(other_frame, frame, how=join)
+        right_result, right_column_mapper = align_on_index(
+            other_frame, frame, dummy_row_pos_mode, how=join
+        )
         right_frame_data_ids = right_column_mapper.map_left_quoted_identifiers(
             other_frame.data_column_snowflake_quoted_identifiers
         )

src/snowflake/snowpark/modin/plugin/_internal/cumulative_utils.py

@@ -86,6 +86,7 @@ def get_groupby_cumagg_frame_axis0(
     groupby_kwargs: dict[str, Any],
     cumagg_func: Callable,
     cumagg_func_name: str,
+    dummy_row_pos_mode: bool,
     ascending: bool = True,
 ) -> InternalFrame:
     """
@@ -124,7 +125,7 @@ def get_groupby_cumagg_frame_axis0(
         )
 
     qc, by_list = resample_and_extract_groupby_column_pandas_labels(
-        query_compiler, by, level
+        query_compiler, by, level, dummy_row_pos_mode
     )
 
     if numeric_only:

src/snowflake/snowpark/modin/plugin/_internal/cut_utils.py

@@ -151,8 +151,8 @@ def compute_bin_indices(
     values_frame: InternalFrame,
     cuts_frame: InternalFrame,
     n_cuts: int,
+    dummy_row_pos_mode: bool,
     right: bool = True,
-    dummy_row_pos_mode: bool = False,
 ) -> InternalFrame:
     """
     Given a frame of cuts, i.e. borders of bins (strictly increasing) compute for the data in values_frame the index of the bin they fall into.
@@ -190,6 +190,7 @@ def compute_bin_indices(
         values_frame,
         cuts_frame,
         how="asof",
+        dummy_row_pos_mode=dummy_row_pos_mode,
         left_match_col=values_frame.data_column_snowflake_quoted_identifiers[0],
         right_match_col=cuts_frame.data_column_snowflake_quoted_identifiers[0],
         match_comparator=MatchComparator.LESS_THAN_OR_EQUAL_TO

src/snowflake/snowpark/modin/plugin/_internal/frame.py

@@ -1356,7 +1356,7 @@ def select_active_columns(self) -> "InternalFrame":
     def strip_duplicates(
         self: "InternalFrame",
         quoted_identifiers: list[str],
-        dummy_row_pos_mode: bool = False,
+        dummy_row_pos_mode: bool,
     ) -> "InternalFrame":
         """
         When assigning frames via index operations for duplicates only the last entry is used, as entries are repeatedly overwritten.
@@ -1403,10 +1403,10 @@ def strip_duplicates(
 
         joined_ordered_dataframe = frame.ordered_dataframe.join(
             right=relevant_last_value_row_positions,
+            dummy_row_pos_mode=dummy_row_pos_mode,
             left_on_cols=[frame.row_position_snowflake_quoted_identifier],
             right_on_cols=[relevant_last_value_row_positions_quoted_identifier],
             how="inner",
-            dummy_row_pos_mode=dummy_row_pos_mode,
         )
 
         # Because we reuse row position to select the relevant columns, we need to

src/snowflake/snowpark/modin/plugin/_internal/get_dummies_utils.py

@@ -228,7 +228,7 @@ def get_dummies_helper(
     prefixes: list[Hashable],
     prefix_sep: str,
     dtype: Any,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Helper function for get dummies to perform encoding on given columns
@@ -348,6 +348,7 @@ def get_dummies_helper(
             left_on=result_internal_frame.index_column_snowflake_quoted_identifiers,
             right_on=pivoted_internal_frame.index_column_snowflake_quoted_identifiers,
             how="inner",
+            dummy_row_pos_mode=dummy_row_pos_mode,
         ).result_frame
 
     # optimization: keep the original row position column as the result ordered frame

src/snowflake/snowpark/modin/plugin/_internal/groupby_utils.py

@@ -471,6 +471,7 @@ def resample_and_extract_groupby_column_pandas_labels(
     query_compiler: "snowflake_query_compiler.SnowflakeQueryCompiler",
     by: Any,
     level: Optional[IndexLabel],
+    dummy_row_pos_mode: bool,
     *,
     skip_resample: bool = False,
 ) -> tuple[
@@ -614,6 +615,7 @@ def find_resample_columns(
                 right=expected_resample_bins_frame,
                 # Perform an outer join to preserve additional index columns.
                 how="outer",
+                dummy_row_pos_mode=dummy_row_pos_mode,
                 # identifier might get mangled by binning operation; look it up again
                 left_on=binned_frame.get_snowflake_quoted_identifiers_group_by_pandas_labels(
                     [original_label]
@@ -658,6 +660,7 @@ def get_frame_with_groupby_columns_as_index(
     by: Any,
     level: Optional[Union[Hashable, int]],
     dropna: bool,
+    dummy_row_pos_mode: bool,
 ) -> Optional["snowflake_query_compiler.SnowflakeQueryCompiler"]:
     """
     Returns a new dataframe with the following properties:
@@ -753,7 +756,7 @@ def get_frame_with_groupby_columns_as_index(
     )
 
     query_compiler, by_list = resample_and_extract_groupby_column_pandas_labels(
-        query_compiler, by, level
+        query_compiler, by, level, dummy_row_pos_mode
     )
 
     if by_list is None:
@@ -918,6 +921,7 @@ def fill_missing_groupby_resample_bins_for_frame(
     by_list: list,
     orig_datetime_index_col_label: str,
     datetime_index_col_identifier: str,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Returns a new InternalFrame created using 2 rules.
@@ -1028,6 +1032,7 @@ def fill_missing_groupby_resample_bins_for_frame(
         frame,
         multi_expected_resample_bins_snowpark_frame,
         how="right",
+        dummy_row_pos_mode=dummy_row_pos_mode,
         left_on=frame.index_column_snowflake_quoted_identifiers,
         right_on=multi_expected_resample_bins_snowpark_frame.index_column_snowflake_quoted_identifiers,
         sort=False,

src/snowflake/snowpark/modin/plugin/_internal/indexing_utils.py

@@ -219,7 +219,7 @@ def validate_out_of_bound(key_max: Any, key_min: Any, axis_len: int) -> None:
 def get_frame_by_row_pos_frame(
     internal_frame: InternalFrame,
     key: InternalFrame,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Select rows from this internal_frame by row positions in the key frame
@@ -269,7 +269,7 @@ def get_frame_by_row_pos_frame(
 def _get_frame_by_row_pos_boolean_frame(
     internal_frame: InternalFrame,
     key: InternalFrame,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Select rows using the boolean frame positional key. The two frames will be inner joined on their row position column
@@ -320,7 +320,7 @@ def _get_frame_by_row_pos_boolean_frame(
 def _get_frame_by_row_pos_int_frame(
     internal_frame: InternalFrame,
     key: InternalFrame,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Select rows using the int frame positional key. The two frames will be inner joined on the internal_frame's row
@@ -374,7 +374,7 @@ def _get_frame_by_row_pos_int_frame(
 def _get_adjusted_key_frame_by_row_pos_int_frame(
     internal_frame: InternalFrame,
     key: InternalFrame,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Return the key frame with any negative row positions adjusted by the internal frame.  For example, if the original
@@ -1100,7 +1100,7 @@ def get_index_frame_by_row_label_slice(
 def get_frame_by_row_label(
     internal_frame: InternalFrame,
     key: Union[InternalFrame, slice, tuple],
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Select rows by labels in the key.
@@ -1472,7 +1472,7 @@ def generate_bound_column(
 def _get_frame_by_row_label_boolean_frame(
     internal_frame: InternalFrame,
     key: InternalFrame,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Select rows with boolean frame key. Here, if the frame and key's index are aligned, then the join is on their row
@@ -1490,6 +1490,7 @@ def _get_frame_by_row_label_boolean_frame(
     joined_frame, result_column_mapper = align_on_index(
         internal_frame,
         key,
+        dummy_row_pos_mode,
         "coalesce",
         dummy_row_pos_mode,
     )
@@ -1525,7 +1526,7 @@ def _get_frame_by_row_label_boolean_frame(
 def _get_frame_by_row_label_non_boolean_frame(
     internal_frame: InternalFrame,
     key: InternalFrame,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Select rows where its index is equal to the index in the key value.
@@ -1644,7 +1645,7 @@ def _set_2d_labels_helper_for_frame_item(
     matching_item_rows_by_label: bool,
     col_info: LocSetColInfo,
     index_is_bool_indexer: bool,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     This set 2d label helper method handles df[index, columns] = item where index is a non-boolean indexer and item is a
@@ -1679,12 +1680,14 @@ def _set_2d_labels_helper_for_frame_item(
         result_frame = align_on_index(
             internal_frame,
             index,
+            dummy_row_pos_mode,
             "coalesce",
         ).result_frame
 
         return align_on_index(
             result_frame,
             item,
+            dummy_row_pos_mode,
             "coalesce",
         ).result_frame
 
@@ -1694,6 +1697,7 @@ def _set_2d_labels_helper_for_frame_item(
             return align_on_index(
                 internal_frame,
                 item,
+                dummy_row_pos_mode,
                 "coalesce",
             ).result_frame
         else:
@@ -1852,7 +1856,7 @@ def _set_2d_labels_helper_for_non_frame_item(
     internal_frame: InternalFrame,
     index: Union[slice, Scalar, InternalFrame],
     index_is_bool_indexer: bool,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     The helper method for the case where item is not an internal frame
@@ -1894,6 +1898,7 @@ def _set_2d_labels_helper_for_non_frame_item(
         return align_on_index(
             internal_frame,
             index,
+            dummy_row_pos_mode,
             "coalesce",
         ).result_frame
     elif isinstance(index, InternalFrame):
@@ -1924,7 +1929,7 @@ def _set_2d_labels_helper_for_single_column_wise_item(
     item_data_column_pandas_labels: list[Hashable],
     index_is_bool_indexer: bool,
     enforce_match_item_by_row_labels: bool,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     # If it's a single column with an item list, then we set the item values column-wise, for example,
@@ -2140,7 +2145,7 @@ def set_frame_2d_labels(
     index_is_bool_indexer: bool,
     deduplicate_columns: bool,
     frame_is_df_and_item_is_series: bool,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Helper function to handle the general loc set functionality. The general idea here is to join the key from ``index``
@@ -2355,6 +2360,7 @@ def set_frame_2d_labels(
                     item_data_column_pandas_labels,
                     index_is_bool_indexer,
                     enforce_match_item_by_row_labels,
+                    dummy_row_pos_mode,
                 )
                 # we convert bool indexer to non-bool one above so set it to False now
                 index_is_bool_indexer = False
@@ -2608,7 +2614,7 @@ def set_frame_2d_positional(
     set_as_coords: bool,
     item: Union[InternalFrame, Scalar],
     is_item_series: bool,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Helper function to handle the general (worst case) 2-join case where index (aka row_key) and item are both frames.
@@ -2837,7 +2843,7 @@ def set_frame_2d_positional(
 def get_kv_frame_from_index_and_item_frames(
     index: InternalFrame,
     item: InternalFrame,
-    dummy_row_pos_mode: bool = False,
+    dummy_row_pos_mode: bool,
 ) -> InternalFrame:
     """
     Return the key-value frame from the key (index) and item (values) frames by aligning on the row positions.