Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@
- `dt.daysinmonth`
- `sort_values`
- `to_datetime`
- `drop`
- Reuse row count from the relaxed query compiler in `get_axis_len`.

#### Bug Fixes
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12988,6 +12988,32 @@ def drop(
columns: Optional[Sequence[Hashable]] = None,
level: Optional[Level] = None,
errors: Literal["raise", "ignore"] = "raise",
) -> "SnowflakeQueryCompiler":
"""
Wrapper around _drop_internal to be supported in faster pandas.
"""
relaxed_query_compiler = None
if self._relaxed_query_compiler is not None and index is None:
relaxed_query_compiler = self._relaxed_query_compiler._drop_internal(
index=index,
columns=columns,
level=level,
errors=errors,
)
qc = self._drop_internal(
index=index,
columns=columns,
level=level,
errors=errors,
)
return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)

def _drop_internal(
self,
index: Optional[Sequence[Hashable]] = None,
columns: Optional[Sequence[Hashable]] = None,
level: Optional[Level] = None,
errors: Literal["raise", "ignore"] = "raise",
) -> "SnowflakeQueryCompiler":
"""
Drop specified rows or columns.
Expand Down
29 changes: 29 additions & 0 deletions tests/integ/modin/test_faster_pandas.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,35 @@ def test_read_filter_join_flag_disabled(session):
assert_frame_equal(snow_result, native_result)


@sql_count_checker(query_count=3)
def test_drop(session):
# create tables
table_name = Utils.random_name_for_temp_object(TempObjectType.TABLE)
session.create_dataframe(
native_pd.DataFrame([[2, True], [1, False], [3, False]], columns=["A", "B"])
).write.save_as_table(table_name, table_type="temp")

# create snow dataframes
df = pd.read_snowflake(table_name)
snow_result = df.drop(columns=["B"])

# verify that the input dataframe has a populated relaxed query compiler
assert df._query_compiler._relaxed_query_compiler is not None
assert df._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
# verify that the output dataframe also has a populated relaxed query compiler
assert snow_result._query_compiler._relaxed_query_compiler is not None
assert (
snow_result._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
)

# create pandas dataframes
native_df = df.to_pandas()
native_result = native_df.drop(columns=["B"])

# compare results
assert_frame_equal(snow_result, native_result)


@pytest.mark.parametrize("func", ["isna", "isnull", "notna", "notnull"])
@sql_count_checker(query_count=3)
def test_isna_notna(session, func):
Expand Down
Loading