perf: use cached property for schema/columns lazy frames (#2093)

Author: MarcoGorelli

narwhals/_dask/dataframe.py

@@ -47,6 +47,7 @@ def __init__(
         self._backend_version = backend_version
         self._implementation = Implementation.DASK
         self._version = version
+        self._cached_schema: dict[str, DType] | None = None
         validate_backend_version(self._implementation, self._backend_version)
 
     def __native_namespace__(self: Self) -> ModuleType:
@@ -132,7 +133,7 @@ def collect(
 
     @property
     def columns(self: Self) -> list[str]:
-        return self._native_frame.columns.tolist()
+        return list(self.schema)
 
     def filter(self: Self, predicate: DaskExpr) -> Self:
         # `[0]` is safe as the predicate's expression only returns a single column
@@ -182,13 +183,15 @@ def drop_nulls(self: Self, subset: list[str] | None) -> Self:
 
     @property
     def schema(self: Self) -> dict[str, DType]:
-        native_dtypes = self._native_frame.dtypes
-        return {
-            col: native_to_narwhals_dtype(
-                native_dtypes[col], self._version, self._implementation
-            )
-            for col in self._native_frame.columns
-        }
+        if self._cached_schema is None:
+            native_dtypes = self._native_frame.dtypes
+            self._cached_schema = {
+                col: native_to_narwhals_dtype(
+                    native_dtypes[col], self._version, self._implementation
+                )
+                for col in self._native_frame.columns
+            }
+        return self._cached_schema
 
     def collect_schema(self: Self) -> dict[str, DType]:
         return self.schema

narwhals/_duckdb/dataframe.py

@@ -55,6 +55,7 @@ def __init__(
         self._native_frame: duckdb.DuckDBPyRelation = df
         self._version = version
         self._backend_version = backend_version
+        self._cached_schema: dict[str, DType] | None = None
         validate_backend_version(self._implementation, self._backend_version)
 
     def __narwhals_dataframe__(self: Self) -> Self:  # pragma: no cover
@@ -189,16 +190,20 @@ def filter(self: Self, predicate: DuckDBExpr) -> Self:
 
     @property
     def schema(self: 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
-            )
-        }
+        if self._cached_schema is None:
+            # Note: prefer `self._cached_schema` over `functools.cached_property`
+            # due to Python3.13 failures.
+            self._cached_schema = {
+                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
+                )
+            }
+        return self._cached_schema
 
     @property
     def columns(self: Self) -> list[str]:
-        return self._native_frame.columns
+        return list(self.schema)
 
     def to_pandas(self: Self) -> pd.DataFrame:
         # only if version is v1, keep around for backcompat

narwhals/_spark_like/dataframe.py

@@ -50,6 +50,7 @@ def __init__(
         self._backend_version = backend_version
         self._implementation = implementation
         self._version = version
+        self._cached_schema: dict[str, DType] | None = None
         validate_backend_version(self._implementation, self._backend_version)
 
     @property
@@ -178,7 +179,7 @@ def _collect_to_arrow(self) -> pa.Table:
 
     @property
     def columns(self: Self) -> list[str]:
-        return self._native_frame.columns  # type: ignore[no-any-return]
+        return list(self.schema)
 
     def collect(
         self: Self,
@@ -266,14 +267,16 @@ def filter(self: Self, predicate: SparkLikeExpr) -> Self:
 
     @property
     def schema(self: Self) -> dict[str, DType]:
-        return {
-            field.name: native_to_narwhals_dtype(
-                dtype=field.dataType,
-                version=self._version,
-                spark_types=self._native_dtypes,
-            )
-            for field in self._native_frame.schema
-        }
+        if self._cached_schema is None:
+            self._cached_schema = {
+                field.name: native_to_narwhals_dtype(
+                    dtype=field.dataType,
+                    version=self._version,
+                    spark_types=self._native_dtypes,
+                )
+                for field in self._native_frame.schema
+            }
+        return self._cached_schema
 
     def collect_schema(self: Self) -> dict[str, DType]:
         return self.schema
@@ -356,11 +359,6 @@ def join(
         left_columns = self.columns
         right_columns = other.columns
 
-        if isinstance(left_on, str):
-            left_on = [left_on]
-        if isinstance(right_on, str):
-            right_on = [right_on]
-
         # create a mapping for columns on other
         # `right_on` columns will be renamed as `left_on`
         # the remaining columns will be either added the suffix or left unchanged.

narwhals/_spark_like/utils.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from functools import lru_cache
 from typing import TYPE_CHECKING
 from typing import Any
 
@@ -20,7 +19,6 @@
     from narwhals.utils import Version
 
 
-@lru_cache(maxsize=16)
 def native_to_narwhals_dtype(
     dtype: pyspark_types.DataType,
     version: Version,