Skip to content

Commit

Permalink
Feat/add collect schema to interchange dfs (#1646)
Browse files Browse the repository at this point in the history
  • Loading branch information
Dhanunjaya-Elluri authored Dec 23, 2024
1 parent 45f48f4 commit ee6ac29
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 0 deletions.
8 changes: 8 additions & 0 deletions narwhals/_duckdb/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,3 +151,11 @@ def _change_version(self: Self, version: Version) -> Self:

def _from_native_frame(self: Self, df: Any) -> Self:
return self.__class__(df, version=self._version)

def collect_schema(self) -> dict[str, DType]:
return {
column_name: native_to_narwhals_dtype(str(duckdb_dtype), self._version)
for column_name, duckdb_dtype in zip(
self._native_frame.columns, self._native_frame.types
)
}
6 changes: 6 additions & 0 deletions narwhals/_ibis/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -129,3 +129,9 @@ def _change_version(self: Self, version: Version) -> Self:

def _from_native_frame(self: Self, df: Any) -> Self:
return self.__class__(df, version=self._version)

def collect_schema(self) -> dict[str, DType]:
return {
column_name: native_to_narwhals_dtype(ibis_dtype, self._version)
for column_name, ibis_dtype in self._native_frame.schema().items()
}
2 changes: 2 additions & 0 deletions tests/frame/interchange_schema_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,7 @@ def test_interchange_schema_ibis(
assert result == expected
assert df["a"].dtype == nw.Int64
assert df.columns == list(expected.keys())
assert df.collect_schema() == expected


def test_interchange_schema_duckdb() -> None:
Expand Down Expand Up @@ -221,6 +222,7 @@ def test_interchange_schema_duckdb() -> None:
assert result == expected
assert df["a"].dtype == nw.Int64
assert df.columns == list(expected.keys())
assert df.collect_schema() == expected


def test_invalid() -> None:
Expand Down

0 comments on commit ee6ac29

Please sign in to comment.