diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py index 0098064451..8fc165d340 100644 --- a/narwhals/_arrow/expr.py +++ b/narwhals/_arrow/expr.py @@ -161,7 +161,7 @@ def func(df: ArrowDataFrame) -> Sequence[ArrowSeries]: # TODO(marco): is there a way to do this efficiently without # doing 2 sorts? Here we're sorting the dataframe and then # again calling `sort_indices`. `ArrowSeries.scatter` would also sort. - sorting_indices = pc.sort_indices(df[token].native) # type: ignore[call-overload] + sorting_indices = pc.sort_indices(df.get_column(token).native) # type: ignore[call-overload] return [s._with_native(s.native.take(sorting_indices)) for s in result] else: @@ -184,7 +184,7 @@ def func(df: ArrowDataFrame) -> Sequence[ArrowSeries]: right_on=partition_by, suffix="_right", ) - return [tmp[alias] for alias in aliases] + return [tmp.get_column(alias) for alias in aliases] return self.__class__( func, diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py index aa9931c901..efeb851e95 100644 --- a/narwhals/_arrow/series.py +++ b/narwhals/_arrow/series.py @@ -854,7 +854,7 @@ def mode(self: Self) -> ArrowSeries: return counts.filter( plx.col(col_token) == plx.col(col_token).max().broadcast(kind=ExprKind.AGGREGATION) - )[self.name] + ).get_column(self.name) def is_finite(self: Self) -> Self: return self._with_native(pc.is_finite(self.native)) diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py index 4d4493e2bb..752b047f33 100644 --- a/narwhals/_pandas_like/dataframe.py +++ b/narwhals/_pandas_like/dataframe.py @@ -953,7 +953,7 @@ def to_numpy(self: Self, dtype: Any = None, *, copy: bool | None = None) -> _2DA arr: Any = np.hstack( [ - self[col].to_numpy(copy=copy, dtype=None)[:, None] + self.get_column(col).to_numpy(copy=copy, dtype=None)[:, None] for col in self.columns ] ) diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py index a812a38653..ae6197a5cd 100644 --- a/narwhals/_pandas_like/expr.py +++ b/narwhals/_pandas_like/expr.py @@ -214,7 +214,7 @@ def func(df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]: *order_by, descending=False, nulls_last=False ) results = self(df.drop([token], strict=True)) - sorting_indices = df[token] + sorting_indices = df.get_column(token) for s in results: s._scatter_in_place(sorting_indices, s) return results @@ -257,14 +257,14 @@ def func(df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]: columns = list(set(partition_by).union(output_names).union(order_by)) token = generate_temporary_column_name(8, columns) df = ( - df[columns] + df.simple_select(*columns) .with_row_index(token) .sort(*order_by, descending=reverse, nulls_last=reverse) ) - sorting_indices = df[token] + sorting_indices = df.get_column(token) elif reverse: columns = list(set(partition_by).union(output_names)) - df = df[columns][::-1] + df = df.simple_select(*columns)[::-1] grouped = df._native_frame.groupby(partition_by) if function_name.startswith("rolling"): rolling = grouped[list(output_names)].rolling(**pandas_kwargs) @@ -287,7 +287,7 @@ def func(df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]: result_frame = df._with_native(res_native).rename( dict(zip(output_names, aliases)) ) - results = [result_frame[name] for name in aliases] + results = [result_frame.get_column(name) for name in aliases] if order_by: for s in results: s._scatter_in_place(sorting_indices, s)