feat: improving array casting #1865
Conversation
| self: Self, inner: DType | type[DType], width: int | None = None | ||
| self: Self, | ||
| inner: DType | type[DType], | ||
| shape: int | tuple[int, ...] | None = None, |
There was a problem hiding this comment.
This is a breaking change. However by checking how narwhals is used on github, I could not find any instances of Array
| @@ -220,7 +224,7 @@ def broadcast_and_extract_dataframe_comparand( | |||
|
|
|||
| if isinstance(other, ArrowSeries): | |||
| len_other = len(other) | |||
| if len_other == 1: | |||
| if len_other == 1 and length != 1: | |||
There was a problem hiding this comment.
Otherwise for list/array types we end up getting the first element
| @@ -111,10 +111,13 @@ def native_to_narwhals_dtype(duckdb_dtype: str, version: Version) -> DType: | |||
| ) | |||
| if match_ := re.match(r"(.*)\[\]$", duckdb_dtype): | |||
| return dtypes.List(native_to_narwhals_dtype(match_.group(1), version)) | |||
| if match_ := re.match(r"(\w+)\[(\d+)\]", duckdb_dtype): | |||
| if match_ := re.match(r"(\w+)((?:\[\d+\])+)", duckdb_dtype): | |||
There was a problem hiding this comment.
Array type in duckdb can have multiple dimensions. The resulting type is: INNER[d1][d2][...]
With this new regex we can parse multiple instances of the dimensions
| duckdb_shape_fmt = "".join(f"[{item}]" for item in dtype.shape) # type: ignore[union-attr] | ||
| while isinstance(dtype.inner, dtypes.Array): # type: ignore[union-attr] | ||
| dtype = dtype.inner # type: ignore[union-attr] | ||
| inner = narwhals_to_native_dtype(dtype.inner, version) # type: ignore[union-attr] | ||
| return f"{inner}{duckdb_shape_fmt}" |
There was a problem hiding this comment.
First creates the shape [d1][d2]... then find the inner type recursively (first being non array)
| @@ -160,7 +160,7 @@ def broadcast_and_extract_dataframe_comparand(index: Any, other: Any) -> Any: | |||
| if isinstance(other, PandasLikeSeries): | |||
| len_other = other.len() | |||
|
|
|||
| if len_other == 1: | |||
| if len_other == 1 and len(index) != 1: | |||
There was a problem hiding this comment.
Similarly as for pyarrow, otherwise for list/array types we end up getting the first element
| ) | ||
| return NotImplementedError(msg) | ||
| if isinstance_or_issubclass(dtype, dtypes.Struct): | ||
| if isinstance_or_issubclass(dtype, (dtypes.Struct, dtypes.Array, dtypes.List)): |
There was a problem hiding this comment.
This diff is quite nice 😉
| def __repr__(self) -> str: | ||
| # Get leaf type | ||
| dtype = self.inner | ||
| while isinstance(dtype, Array): | ||
| dtype = dtype.inner | ||
|
|
||
| def __repr__(self: Self) -> str: | ||
| class_name = self.__class__.__name__ | ||
| return f"{class_name}({self.inner!r}, {self.width})" | ||
| return f"{class_name}({dtype!r}, shape={self.shape})" |
There was a problem hiding this comment.
This diff is that's causing marimo ci to fail - we could provide a fix there 🙈
What type of PR is this? (check all applicable)
Checklist
If you have comments or can explain your changes, please do so below