diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 7055b2a067..b2acea46c8 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -150,6 +150,8 @@ jobs: python-version: ["3.14t"] os: [ubuntu-latest] runs-on: ${{ matrix.os }} + # Backstop against deadlocks in multithreaded tests: fail instead of hanging. + timeout-minutes: 10 env: PYTHON_GIL: 0 UV_LOCKED: 0 @@ -164,5 +166,20 @@ jobs: enable-cache: "true" cache-suffix: python-314t-${{ matrix.python-version }} cache-dependency-glob: "pyproject.toml" - - name: Run pytest - run: make run-ci DEPS="--pre --group tests --group plugins --extra pandas --extra pyarrow" CMD="pytest tests --cov=src --cov=tests --cov-fail-under=50 --runslow --durations=30 --constructors=pandas,pandas[nullable],pandas[pyarrow],pyarrow" + # Run the suite again with every test executed by 4 threads at once + # with pytest-run-parallel to surface thread-safety issues. + # This catches global-state races across the whole suite; shared-object races + # are covered by the dedicated tests in tests/free_threading_test.py. + # Tests marked `thread_unsafe` run single-threaded. + - name: Run pytest (parallel threads) + env: + # NOTE: Clear PYTEST_ADDOPTS since xdist forks processes, + # and the point here is threads within one process + PYTEST_ADDOPTS: "" + # TODO(Unassigned): Add other backends (duckdb, polars, ...) as they rollout free-threaded support + # polars: https://github.com/pola-rs/polars/issues/27955 + run: | + make run-ci \ + DEPS="--pre --group tests --group plugins --extra pandas --extra pyarrow" \ + RUN_ONLY="--with pytest-run-parallel" \ + CMD="pytest tests --parallel-threads=4 --constructors=pandas,pandas[nullable],pandas[pyarrow],pyarrow" diff --git a/_typos.toml b/_typos.toml index 74e66bbe0a..c9f57e937a 100644 --- a/_typos.toml +++ b/_typos.toml @@ -5,6 +5,7 @@ ba = "ba" # Used as column name in docstring examples (way too much?) iy = "iy" # Used as column name (once in a test) pn = "pn" # Used in docs: pn = PandasLikeNamespace(...) TYP = "TYP" # Used in flake8 rule +tpe = "tpe" # Alias for ThreadPoolExecutor [files] extend-exclude = ["tests/data/*"] diff --git a/docs/concepts/thread_safety.md b/docs/concepts/thread_safety.md new file mode 100644 index 0000000000..64050c9db3 --- /dev/null +++ b/docs/concepts/thread_safety.md @@ -0,0 +1,122 @@ +# Thread safety + +## TL;DR + +- **Narwhals objects are safe to share between threads for reading**: methods return new + objects; nothing mutates the receiver, nor the native object passed to + [`from_native`](../api-reference/narwhals.md#narwhals.from_native). +- **Writing is the backend's business**: Narwhals has no in-place mutation API, but once you + reach the native object (`to_native`, `to_numpy`, ...) the backend's rules apply. +- **DuckDB needs one connection per thread**: the only backend with a rule that changes how + you write Narwhals code, see [DuckDB](#duckdb). + +!!! tip + All of this holds on both the GIL-enabled and the [free-threaded] build. A race is a bug + either way; free-threading just makes it likely enough to notice. + +## Global state inside Narwhals + +| State | Thread-safe? | Notes | +| ----- | ------------ | ----- | +| Function caches
(`backend_version`, dtype conversions, plugin discovery, ...) | Yes | [`lru_cache` is threadsafe][lru_cache] and the cached functions are pure,
so a concurrent cold cache costs at most a duplicated computation | +| Constant lookup tables
(dtype mappings, interval units, ...) | Yes | Read-only after import | +| Configuration / registries / mutable singletons | N/A | Narwhals has none | +| `narwhals.sql`'s DuckDB catalog | Yes, per thread | See [`narwhals.sql`](#narwhalssql) | + +### Backend detection + +Narwhals does not import your backend in order to recognise its objects: you must already +have imported it to have an object to pass, so it looks the module up in `sys.modules` +instead. It *does* import on your behalf when you name a backend +(`nw.from_dict(..., backend="pandas")`) or convert between them (`to_polars`, `to_arrow`, +...). Since a module is [visible in `sys.modules` before its body finishes executing][loading], +import your backends on the main thread rather than racing that first import from a pool. + +## Sharing Narwhals objects + +All of the following are safe to share between threads: + +| Object | Why | +| --- | --- | +| `DataFrame`, `LazyFrame`, `Series` | Every method returns a new object.
Lazily-cached metadata (schema, column names) is computed idempotently,
so a concurrent first access only repeats work | +| `Expr` | A chain of immutable nodes. `Expr.over` copies before it rewrites,
so reusing one expression in several queries at once cannot corrupt it | +| `GroupBy` (the result of `DataFrame.group_by(...)`) | Aggregation state is passed along the call, not stashed on the object | +| `Schema` | A `dict` subclass: safe to read concurrently, but do not mutate a shared instance | + +!!! warning + Sharing is safe because nobody writes: if a worker hands the native object to something + that mutates it in place, Narwhals' view goes stale too - `from_native` does not copy. + +## Backends + +| Backend | Reads | Writes to the native object | +| --- | --- | --- | +| Polars | Safe | Operations return new objects, though immutability is [not enforced][polars-immutable] | +| PyArrow | Safe | [Arrow data is immutable][arrow-immutable] | +| pandas (NumPy-backed) | Safe | **Not safe** through the NumPy buffer (`.values`, `.to_numpy()`):
[mutating a shared array races][numpy-thread-safety]. pandas' own writes go
through [copy-on-write], the default since pandas 3.0 | +| pandas (Arrow-backed) | Safe | [copy-on-write]: a write produces new buffers rather than mutating shared ones | +| DuckDB | One connection per thread, see below | N/A | +| PySpark, SQLFrame, Ibis, Dask | Delegated to the session / scheduler | Consult the backend's own docs | + +## DuckDB + +Every thread needs [its own `.cursor()`, a thread-local connection to the same +database][multiple threads]: + +```python +import duckdb +import narwhals as nw + + +def worker(con: duckdb.DuckDBPyConnection) -> None: + cursor = con.cursor() # one per thread + lf = nw.from_native(cursor.sql("select * from my_table")) + print(lf.filter(nw.col("a") > 1).collect()) +``` + +1. **A cursor does not inherit `LOCAL`-[scoped][duckdb-config] settings**: `TimeZone` is + one of them, so use `SET GLOBAL`, or set it on every cursor. It matters in Narwhals + because DuckDB keeps the time zone in the connection rather than the dtype: + `collect_schema()` reports a time-zone-aware column using *the cursor's* time zone. +2. **A relation belongs to the connection that created it**: combining relations from two + connections is [rejected outright][duckdb-join-relation], so build frames you intend to + `join` or `concat` in the thread that uses them. +3. **`collect_schema()` on frames from a shared connection is safe, executing on one is + not**: a time-zone-aware dtype is the one case where Narwhals queries your connection + *implicitly*, and it serializes that query. Explicit execution (`collect`, `to_arrow`, + ...) runs on the connection you gave it and follows DuckDB's rules. + +[`scan_csv`](../api-reference/narwhals.md#narwhals.scan_csv) and +[`scan_parquet`](../api-reference/narwhals.md#narwhals.scan_parquet) read through DuckDB's +process-global default connection unless you pass `connection=con.cursor()`. + +### `narwhals.sql` + +[`narwhals.sql.table`](../api-reference/sql.md) keeps a module-level DuckDB connection and +gives each thread its own cursor on it. Since a cursor is a connection to the same database, +the catalog is shared process-wide and `name` must be unique across threads; and by (2) +above, tables created in *different* threads cannot be joined. + +## Free-threaded Python + +Narwhals is pure Python with no compiled extensions, so importing it never re-enables the +GIL. Whether *your* stack supports the [free-threaded] build depends on the backend: check +that its wheels are built for `cp314t` (or whichever version you run). + +CI covers it in two ways: the whole suite runs again under [pytest-run-parallel] with +`--parallel-threads=4`, which surfaces races through global state, and +`tests/free_threading_test.py` stresses every guarantee on this page. If you find a race, +please [open an issue](https://github.com/narwhals-dev/narwhals/issues) with a script that +reproduces it, ideally on a free-threaded build or with `sys.setswitchinterval(1e-7)`. + +[arrow-immutable]: https://arrow.apache.org/docs/python/data.html +[copy-on-write]: https://pandas.pydata.org/docs/user_guide/copy_on_write.html +[duckdb-config]: https://duckdb.org/docs/stable/configuration/overview +[duckdb-join-relation]: https://github.com/duckdb/duckdb/blob/fabf1d60bb0565032ad7d48e64f689fdbf616719/src/main/relation/join_relation.cpp#L25-L26 +[free-threaded]: https://docs.python.org/3/howto/free-threading-python.html +[loading]: https://docs.python.org/3/reference/import.html#loading +[lru_cache]: https://docs.python.org/3/library/functools.html#functools.lru_cache +[multiple threads]: https://duckdb.org/docs/stable/guides/python/multiple_threads +[numpy-thread-safety]: https://numpy.org/doc/stable/reference/thread_safety.html +[polars-immutable]: https://github.com/pola-rs/polars/issues/17447 +[pytest-run-parallel]: https://github.com/Quansight-Labs/pytest-run-parallel diff --git a/pyproject.toml b/pyproject.toml index 4a6a9d6cc1..9cec7002fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -415,7 +415,12 @@ filterwarnings = [ "ignore:.*interchange protocol is deprecated.*:DeprecationWarning" ] xfail_strict = true -markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"] +markers = [ + "slow: marks tests as slow (deselect with '-m \"not slow\"')", + # Registered manually so the mark is valid even when pytest-run-parallel is not installed; + # the plugin (python-314t job) uses it to run marked tests single-threaded. + "thread_unsafe(reason): marks tests that must not run in multiple threads at once", +] env = [ "MODIN_ENGINE=python", "PYARROW_IGNORE_TIMEZONE=1", diff --git a/src/narwhals/_duckdb/dataframe.py b/src/narwhals/_duckdb/dataframe.py index 18d0fb750b..489eef0dc8 100644 --- a/src/narwhals/_duckdb/dataframe.py +++ b/src/narwhals/_duckdb/dataframe.py @@ -16,6 +16,7 @@ join_column_names, lit, native_to_narwhals_dtype, + temporary_view_name, window_expression, ) from narwhals._sql.dataframe import SQLLazyFrame @@ -388,13 +389,19 @@ def join_asof( else f'rhs."{name}"' for name in keep_cols ) - query = f""" + # `rhs` is resolved via replacement scan. + # `rel.query` creates a view, so this fails on read-only databases (#3567); + # unavoidable until the relational API gains ASOF join support. + view = temporary_view_name() + joined = lhs.query( + view, + f""" SELECT {",".join(rhs_select)} - FROM lhs + FROM {view} AS lhs ASOF LEFT JOIN rhs ON {condition} - """ # noqa: S608 - joined = duckdb.sql(query) + """, # noqa: S608 + ) select = [col(name) for name in lhs.columns] select.extend( @@ -457,25 +464,23 @@ def sort(self, *by: str, descending: bool | Sequence[bool], nulls_last: bool) -> return self._with_native(self.native.sort(*it)) def top_k(self, k: int, *, by: Iterable[str], reverse: bool | Sequence[bool]) -> Self: - _rel = self.native by = list(by) if isinstance(reverse, bool): descending = extend_bool(not reverse, len(by)) else: descending = tuple(not rev for rev in reverse) + tmp_name = generate_temporary_column_name(8, self.columns, prefix="row_number_") expr = window_expression( F("row_number"), order_by=by, descending=descending, nulls_last=extend_bool(True, len(by)), ) - condition = expr <= lit(k) - query = f""" - SELECT * - FROM _rel - QUALIFY {condition} - """ # noqa: S608 - return self._with_native(duckdb.sql(query)) + return self._with_native( + self.native.select(StarExpression(), expr.alias(tmp_name)).filter( + col(tmp_name) <= lit(k) + ) + ).drop([tmp_name], strict=False) def drop_nulls(self, subset: Sequence[str] | None) -> Self: subset_ = subset if subset is not None else self.columns @@ -542,19 +547,21 @@ def unpivot( raise NotImplementedError(msg) unpivot_on = join_column_names(*on_) - _rel = self.native + # `rel.query` creates a view, so this fails on read-only databases (#3567). # Replace with Python API once # https://github.com/duckdb/duckdb/discussions/16980 is addressed. - query = f""" - unpivot _rel + view = temporary_view_name() + unpivoted = self.native.query( + view, + f""" + unpivot {view} on {unpivot_on} into name {col(variable_name)} value {col(value_name)} - """ - return self._with_native( - duckdb.sql(query).select(*[*index_, variable_name, value_name]) + """, ) + return self._with_native(unpivoted.select(*index_, variable_name, value_name)) @requires.backend_version((1, 3)) def with_row_index(self, name: str, order_by: Sequence[str]) -> Self: diff --git a/src/narwhals/_duckdb/namespace.py b/src/narwhals/_duckdb/namespace.py index 7247864eef..50d9b5e44b 100644 --- a/src/narwhals/_duckdb/namespace.py +++ b/src/narwhals/_duckdb/namespace.py @@ -20,6 +20,7 @@ lit, narwhals_to_native_dtype, sql_expression, + temporary_view_name, when, window_expression, ) @@ -73,12 +74,18 @@ def scan_csv( self, source: NormalizedPath, *, separator: str = ",", **kwds: Any ) -> DuckDBLazyFrame: validate_separators(separator, ("delimiter", "delim", "sep"), kwds) - native = duckdb.read_csv(source, delimiter=separator, **kwds) - return self._lazyframe.from_native(native, context=self) + # Without an explicit `connection`, DuckDB reads through the + # process-global default connection. + reader = kwds.pop("connection", None) or duckdb + native_frame = reader.read_csv(source, delimiter=separator, **kwds) + return self._lazyframe.from_native(native_frame, context=self) def scan_parquet(self, source: NormalizedPath, **kwds: Any) -> DuckDBLazyFrame: - native = duckdb.read_parquet(source, **kwds) - return self._lazyframe.from_native(native, context=self) + # Without an explicit `connection`, DuckDB reads through the + # process-global default connection. + reader = kwds.pop("connection", None) or duckdb + native_frame = reader.read_parquet(source, **kwds) + return self._lazyframe.from_native(native_frame, context=self) def _function(self, name: str, *args: Expression) -> Expression: # type: ignore[override] return function(name, *args) @@ -113,9 +120,13 @@ def concat( res = first.native for _item in native_items[1:]: # TODO(unassigned): use relational API when available https://github.com/duckdb/duckdb/discussions/16996 - res = duckdb.sql(""" - from res select * union all by name from _item select * - """) + # `_item` is resolved via replacement scan. + # `rel.query` creates a view, so this fails on read-only databases (#3567). + view = temporary_view_name() + res = res.query( + view, + f"from {view} select * union all by name from _item select *", # noqa: S608 + ) return first._with_native(res) res = reduce(lambda x, y: x.union(y), native_items) return first._with_native(res) diff --git a/src/narwhals/_duckdb/utils.py b/src/narwhals/_duckdb/utils.py index ae84fe1d99..84914b6b65 100644 --- a/src/narwhals/_duckdb/utils.py +++ b/src/narwhals/_duckdb/utils.py @@ -1,13 +1,20 @@ from __future__ import annotations import operator +import threading from functools import lru_cache from typing import TYPE_CHECKING, Any import duckdb from duckdb import Expression -from narwhals._utils import Implementation, Version, extend_bool, isinstance_or_issubclass +from narwhals._utils import ( + Implementation, + Version, + extend_bool, + generate_temporary_column_name, + isinstance_or_issubclass, +) from narwhals.exceptions import ColumnNotFoundError if TYPE_CHECKING: @@ -199,10 +206,24 @@ def native_to_narwhals_dtype( return _non_nested_native_to_narwhals_dtype(duckdb_dtype_id, version) +_TIME_ZONE_LOCK = threading.Lock() +"""Serializes the only query Narwhals issues *implicitly* on a user's connection. + +A connection must not be used concurrently (see [multiple threads]), and a relation +gives us no way to reach its own connection to open a per-thread `.cursor()`. +This keeps `collect_schema()` safe on relations sharing a connection; explicit execution +(`collect`, ...) still follows DuckDB's rules, see [docs/concepts/thread_safety.md]. + +[multiple threads]: https://duckdb.org/docs/stable/guides/python/multiple_threads +""" + + def fetch_rel_time_zone(rel: duckdb.DuckDBPyRelation) -> str: - result = rel.query( - "duckdb_settings()", "select value from duckdb_settings() where name = 'TimeZone'" - ).fetchone() + with _TIME_ZONE_LOCK: + result = rel.query( + "duckdb_settings()", + "select value from duckdb_settings() where name = 'TimeZone'", + ).fetchone() assert result is not None # noqa: S101 return result[0] # type: ignore[no-any-return] @@ -337,6 +358,26 @@ def join_column_names(*names: str) -> str: return ", ".join(str(col(name)) for name in names) +def temporary_view_name() -> str: + """Unique name for `DuckDBPyRelation.query`'s `virtual_table_name`. + + `rel.query(view, sql)` registers `rel` as a view named `view` and runs `sql` on + `rel`'s own connection (see [relational api]). We prefer it to `duckdb.sql(statement)`, + which uses the global default connection and cannot see relations from other + connections (see [using connections in parallel pythonprograms]). + + The name must be unique per call: views persist and the lazy result re-binds by + name on execution, so a reused name shadows earlier views and corrupts their plans. + `rel.query` must also run in the frame whose locals the SQL references, as + [replacement scans] only see the caller's frame [3]. + + [relational api]: https://duckdb.org/docs/current/clients/python/relational_api + [using connections in parallel pythonprograms]: https://duckdb.org/docs/current/clients/python/overview#using-connections-in-parallel-python-programs + [replacement scans]: https://duckdb.org/docs/current/clients/c/replacement_scans + """ + return generate_temporary_column_name(8, [], prefix="_narwhals_") + + def generate_order_by_sql( *order_by: str | Expression, descending: Sequence[bool], nulls_last: Sequence[bool] ) -> str: diff --git a/src/narwhals/_pandas_like/group_by.py b/src/narwhals/_pandas_like/group_by.py index 8b2572cc75..7def377157 100644 --- a/src/narwhals/_pandas_like/group_by.py +++ b/src/narwhals/_pandas_like/group_by.py @@ -112,9 +112,10 @@ def with_expand_names(self, group_by: PandasLikeGroupBy, /) -> AggExpr: ) return self - def _getitem_aggs(self, group_by: PandasLikeGroupBy) -> pd.DataFrame | pd.Series[Any]: + def _getitem_aggs( + self, group_by: PandasLikeGroupBy, grouped: NativeGroupBy + ) -> pd.DataFrame | pd.Series[Any]: """Evaluate the wrapped expression as a group_by operation.""" - grouped = group_by._grouped result: pd.DataFrame | pd.Series[Any] names = self.output_names if self.is_len() and self.is_top_level_function(): @@ -296,13 +297,13 @@ def agg(self, *exprs: PandasLikeExpr) -> PandasLikeDataFrame: # noqa: PLR0912 ).groupby(self._keys.copy(), **self._group_by_kwargs) else: grouped = self._native.groupby(self._keys.copy(), **self._group_by_kwargs) - self._grouped = grouped - + # NOTE: `grouped` varies per `agg` call, so pass it along rather than stashing it + # on `self`: a shared `GroupBy` would let one call read another's grouping. if all_aggs_are_simple: result: pd.DataFrame if agg_exprs: ns = self.compliant.__narwhals_namespace__() - result = ns._concat_horizontal(self._getitem_aggs(agg_exprs)) + result = ns._concat_horizontal(self._getitem_aggs(agg_exprs, grouped)) else: result = self.compliant.__native_namespace__().DataFrame( list(grouped.groups), columns=self._keys @@ -337,9 +338,9 @@ def _select_results( ) def _getitem_aggs( - self, exprs: Iterable[AggExpr], / + self, exprs: Iterable[AggExpr], grouped: NativeGroupBy, / ) -> list[pd.DataFrame | pd.Series[Any]]: - return [e._getitem_aggs(self) for e in exprs] + return [e._getitem_aggs(self, grouped) for e in exprs] def _apply_aggs( self, grouped: NativeGroupBy, exprs: Iterable[PandasLikeExpr] diff --git a/src/narwhals/functions.py b/src/narwhals/functions.py index fb661554b3..f30b7a9713 100644 --- a/src/narwhals/functions.py +++ b/src/narwhals/functions.py @@ -692,6 +692,24 @@ def scan_csv( For the libraries that do not support lazy dataframes, the function reads a csv file eagerly and then converts the resulting dataframe to a lazyframe. + Note: + Spark-like backends require a `session` object to be passed in `kwargs`. + + For instance: + + ```py + import narwhals as nw + from sqlframe.duckdb import DuckDBSession + + nw.scan_csv(source, backend="sqlframe", session=DuckDBSession()) + ``` + + Note: + For the DuckDB backend, a `connection` object can be passed in `kwargs` + to read through a specific `DuckDBPyConnection` instead of the + process-global default connection. Reading through a per-thread connection + is what makes concurrent reads safe, see [thread safety](../concepts/thread_safety.md). + Arguments: source: Path to a file. backend: The eager backend for DataFrame creation. @@ -806,7 +824,7 @@ def scan_parquet( a parquet file eagerly and then converts the resulting dataframe to a lazyframe. Note: - Spark like backends require a session object to be passed in `kwargs`. + Spark like backends require a `session` object to be passed in `kwargs`. For instance: @@ -817,6 +835,12 @@ def scan_parquet( nw.scan_parquet(source, backend="sqlframe", session=DuckDBSession()) ``` + Note: + For the DuckDB backend, a `connection` object can be passed in `kwargs` + to read through a specific `DuckDBPyConnection` instead of the + process-global default connection. Reading through a per-thread connection + is what makes concurrent reads safe, see [thread safety](../concepts/thread_safety.md). + Arguments: source: Path to a file. backend: The eager backend for DataFrame creation. diff --git a/src/narwhals/sql.py b/src/narwhals/sql.py index bfa888c145..f701602a1e 100644 --- a/src/narwhals/sql.py +++ b/src/narwhals/sql.py @@ -1,5 +1,6 @@ from __future__ import annotations +import threading from typing import TYPE_CHECKING, Literal from narwhals._duckdb.utils import DeferredTimeZone, narwhals_to_native_dtype @@ -22,9 +23,22 @@ raise ModuleNotFoundError(msg) from _exc CONN = duckdb.connect() -TZ = DeferredTimeZone( - CONN.sql("select value from duckdb_settings() where name = 'TimeZone'") -) +_LOCAL = threading.local() + + +def _cursor() -> duckdb.DuckDBPyConnection: + """Return current thread's cursor on `CONN`, creating it on first use. + + DuckDB connections must not be used concurrently from multiple threads. + Citing from [Multiple Python Threads](https://duckdb.org/docs/stable/guides/python/multiple_threads): + + > Each thread must use the `.cursor()` method to create a thread-local + > connection to the same DuckDB file based on the original connection + """ + # duckdb 1.4 and older don't keep cursors alive + if (cursor := getattr(_LOCAL, "cursor", None)) is None: + cursor = _LOCAL.cursor = CONN.cursor() + return cursor class SQLTable(LazyFrame[duckdb.DuckDBPyRelation]): @@ -71,6 +85,13 @@ def table(name: str, schema: IntoSchema) -> SQLTable: Note that this requires DuckDB to be installed. + Note: + Tables are created in a module-level DuckDB catalog, shared by the whole process, + so `name` must be unique across threads. Each thread gets its own cursor on that + catalog, and DuckDB rejects combining relations from different connections, so + tables created in *different* threads cannot be joined or concatenated - see + [thread safety](../concepts/thread_safety.md). + Parameters: name: Table name. schema: Table schema. @@ -90,16 +111,20 @@ def table(name: str, schema: IntoSchema) -> SQLTable: | 0 rows | └────────────────────────────┘ """ + cursor = _cursor() + tz = DeferredTimeZone( + cursor.sql("select value from duckdb_settings() where name = 'TimeZone'") + ) column_mapping = { - col: narwhals_to_native_dtype(dtype, Version.MAIN, TZ) + col: narwhals_to_native_dtype(dtype, Version.MAIN, tz) for col, dtype in Schema(schema).items() } dtypes = ", ".join(f'"{col}" {dtype}' for col, dtype in column_mapping.items()) - CONN.sql(f""" + cursor.sql(f""" CREATE TABLE "{name}" ({dtypes}); """) - lf = from_native(CONN.table(name)) + lf = from_native(cursor.table(name)) return SQLTable(lf._compliant_frame, level=lf._level) diff --git a/src/narwhals/stable/v2/__init__.py b/src/narwhals/stable/v2/__init__.py index 3eea1b8b4e..fedf26350c 100644 --- a/src/narwhals/stable/v2/__init__.py +++ b/src/narwhals/stable/v2/__init__.py @@ -1079,6 +1079,24 @@ def scan_csv( For the libraries that do not support lazy dataframes, the function reads a csv file eagerly and then converts the resulting dataframe to a lazyframe. + Note: + Spark-like backends require a `session` object to be passed in `kwargs`. + + For instance: + + ```py + import narwhals as nw + from sqlframe.duckdb import DuckDBSession + + nw.scan_csv(source, backend="sqlframe", session=DuckDBSession()) + ``` + + Note: + For the DuckDB backend, a `connection` object can be passed in `kwargs` + to read through a specific `DuckDBPyConnection` instead of the + process-global default connection. Reading through a per-thread connection + is what makes concurrent reads safe, see [thread safety](../concepts/thread_safety.md). + Arguments: source: Path to a file. backend: The eager backend for DataFrame creation. @@ -1128,7 +1146,7 @@ def scan_parquet( a parquet file eagerly and then converts the resulting dataframe to a lazyframe. Note: - Spark like backends require a session object to be passed in `kwargs`. + Spark like backends require a `session` object to be passed in `kwargs`. For instance: @@ -1139,6 +1157,12 @@ def scan_parquet( nw.scan_parquet(source, backend="sqlframe", session=DuckDBSession()) ``` + Note: + For the DuckDB backend, a `connection` object can be passed in `kwargs` + to read through a specific `DuckDBPyConnection` instead of the + process-global default connection. Reading through a per-thread connection + is what makes concurrent reads safe, see [thread safety](../concepts/thread_safety.md). + Arguments: source: Path to a file. backend: The eager backend for DataFrame creation. diff --git a/tests/conftest.py b/tests/conftest.py index 0b7dec95bd..b208ead6c0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import threading import uuid from copy import deepcopy from functools import lru_cache @@ -176,16 +177,29 @@ def polars_lazy_constructor(obj: Data) -> pl.LazyFrame: return pl.LazyFrame(obj) +_duckdb_local = threading.local() + + def duckdb_lazy_constructor(obj: dict[str, Any]) -> NativeDuckDB: + """One DuckDB connection per thread. + + DuckDB connections are not safe for concurrent use, and pytest-run-parallel runs + the same test in several threads at once. Relations from different connections + cannot be combined, so frames built within one test (thread) must share a connection. + + See: https://duckdb.org/docs/lts/guides/python/multiple_threads + """ pytest.importorskip("duckdb") pytest.importorskip("pyarrow") import duckdb import pyarrow as pa - duckdb.sql("""set timezone = 'UTC'""") + if (conn := getattr(_duckdb_local, "conn", None)) is None: + conn = _duckdb_local.conn = duckdb.connect() + conn.sql("""set timezone = 'UTC'""") _df = pa.table(obj) - return duckdb.sql("select * from _df") + return conn.sql("select * from _df") def dask_lazy_p1_constructor(obj: Data) -> NativeDask: # pragma: no cover diff --git a/tests/dtypes/dtypes_test.py b/tests/dtypes/dtypes_test.py index ed675ceceb..74fbdc275b 100644 --- a/tests/dtypes/dtypes_test.py +++ b/tests/dtypes/dtypes_test.py @@ -299,7 +299,8 @@ def test_huge_int() -> None: else: # pragma: no cover pass - rel = duckdb.sql(""" + conn = duckdb.connect() + rel = conn.sql(""" select cast(a as int128) as a from df """) @@ -307,7 +308,7 @@ def test_huge_int() -> None: result = nw.from_native(rel).schema assert result["a"] == nw.Int128 - rel = duckdb.sql(""" + rel = conn.sql(""" select cast(a as uint128) as a from df """) @@ -380,7 +381,8 @@ def test_decimal() -> None: df = pl.DataFrame({"a": [1]}, schema={"a": pl.Decimal}) result = nw.from_native(df).schema assert result["a"] == nw.Decimal - rel = duckdb.sql(""" + con = duckdb.connect() + rel = con.sql(""" select * from df """) @@ -467,7 +469,8 @@ def test_huge_int_to_native() -> None: ) assert df_casted.schema["a_int"] == pl.Int128 - rel = duckdb.sql(""" + con = duckdb.connect() + rel = con.sql(""" select cast(a as int64) as a from df """) @@ -656,18 +659,19 @@ def test_datetime_w_tz_duckdb() -> None: pytest.importorskip("duckdb") import duckdb - duckdb.sql("""set timezone = 'Europe/Amsterdam'""") + conn = duckdb.connect() + conn.sql("""set timezone = 'Europe/Amsterdam'""") df = nw.from_native( - duckdb.sql("""select * from values (timestamptz '2020-01-01')df(a)""") + conn.sql("""select * from values (timestamptz '2020-01-01')df(a)""") ) result = df.collect_schema() assert result["a"] == nw.Datetime("us", "Europe/Amsterdam") - duckdb.sql("""set timezone = 'Asia/Kathmandu'""") + conn.sql("""set timezone = 'Asia/Kathmandu'""") result = df.collect_schema() assert result["a"] == nw.Datetime("us", "Asia/Kathmandu") df = nw.from_native( - duckdb.sql( + conn.sql( """select * from values (timestamptz '2020-01-01', [[timestamptz '2020-01-02']])df(a,b)""" ) ) diff --git a/tests/expr_and_series/dt/convert_time_zone_test.py b/tests/expr_and_series/dt/convert_time_zone_test.py index 65d1a6e3b6..7b7d3c72bb 100644 --- a/tests/expr_and_series/dt/convert_time_zone_test.py +++ b/tests/expr_and_series/dt/convert_time_zone_test.py @@ -141,8 +141,9 @@ def test_convert_time_zone_to_connection_tz_duckdb() -> None: import duckdb - duckdb.sql("set timezone = 'Asia/Kolkata'") - rel = duckdb.sql("""select * from values (timestamptz '2020-01-01') df(a)""") + conn = duckdb.connect() + conn.sql("set timezone = 'Asia/Kolkata'") + rel = conn.sql("""select * from values (timestamptz '2020-01-01') df(a)""") result = nw.from_native(rel).with_columns( nw.col("a").dt.convert_time_zone("Asia/Kolkata") ) diff --git a/tests/expr_and_series/dt/replace_time_zone_test.py b/tests/expr_and_series/dt/replace_time_zone_test.py index 1c9dff7d59..ec9701ee1e 100644 --- a/tests/expr_and_series/dt/replace_time_zone_test.py +++ b/tests/expr_and_series/dt/replace_time_zone_test.py @@ -129,8 +129,9 @@ def test_replace_time_zone_to_connection_tz_duckdb() -> None: import duckdb - duckdb.sql("set timezone = 'Asia/Kolkata'") - rel = duckdb.sql("""select * from values (timestamptz '2020-01-01') df(a)""") + conn = duckdb.connect() + conn.sql("set timezone = 'Asia/Kolkata'") + rel = conn.sql("""select * from values (timestamptz '2020-01-01') df(a)""") result = nw.from_native(rel).with_columns( nw.col("a").dt.replace_time_zone("Asia/Kolkata") ) diff --git a/tests/expr_and_series/dt/truncate_test.py b/tests/expr_and_series/dt/truncate_test.py index 40ce18d428..9540fbecd0 100644 --- a/tests/expr_and_series/dt/truncate_test.py +++ b/tests/expr_and_series/dt/truncate_test.py @@ -204,12 +204,12 @@ def test_truncate_tz_aware_duckdb() -> None: import duckdb - duckdb.sql("""set timezone = 'Europe/Amsterdam'""") - rel = duckdb.sql("""select * from values (timestamptz '2020-10-25') df(a)""") + conn = duckdb.connect() + conn.sql("""set timezone = 'Europe/Amsterdam'""") + rel = conn.sql("""select * from values (timestamptz '2020-10-25') df(a)""") result = nw.from_native(rel).with_columns(a_truncated=nw.col("a").dt.truncate("1mo")) expected = { "a": [datetime(2020, 10, 25, tzinfo=ZoneInfo("Europe/Amsterdam"))], "a_truncated": [datetime(2020, 10, 1, tzinfo=ZoneInfo("Europe/Amsterdam"))], } assert_equal_data(result, expected) - duckdb.sql("""set timezone = 'UTC'""") diff --git a/tests/expr_and_series/rolling_sum_test.py b/tests/expr_and_series/rolling_sum_test.py index df7b481826..30fb264929 100644 --- a/tests/expr_and_series/rolling_sum_test.py +++ b/tests/expr_and_series/rolling_sum_test.py @@ -1,6 +1,7 @@ from __future__ import annotations import random +from functools import partial from typing import Any import hypothesis.strategies as st @@ -174,21 +175,26 @@ def test_rolling_sum_series(constructor_eager: ConstructorEager) -> None: ( -1, None, - pytest.raises( - ValueError, match="window_size must be greater or equal than 1" + partial( + pytest.raises, + ValueError, + match="window_size must be greater or equal than 1", ), ), ( 2, -1, - pytest.raises( - ValueError, match="min_samples must be greater or equal than 1" + partial( + pytest.raises, + ValueError, + match="min_samples must be greater or equal than 1", ), ), ( 1, 2, - pytest.raises( + partial( + pytest.raises, InvalidOperationError, match="`min_samples` must be less or equal than `window_size`", ), @@ -196,12 +202,20 @@ def test_rolling_sum_series(constructor_eager: ConstructorEager) -> None: ( 4.2, None, - pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+window_size="), + partial( + pytest.raises, + TypeError, + match=r"Expected '.+?', got: '.+?'\s+window_size=", + ), ), ( 2, 4.2, - pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+min_samples="), + partial( + pytest.raises, + TypeError, + match=r"Expected '.+?', got: '.+?'\s+min_samples=", + ), ), ], ) @@ -213,7 +227,7 @@ def test_rolling_sum_expr_invalid_params( ) -> None: df = nw.from_native(constructor_eager(data)) - with context: + with context(): df.select( nw.col("a").rolling_sum(window_size=window_size, min_samples=min_samples) ) @@ -228,21 +242,26 @@ def test_rolling_sum_expr_invalid_params( ( -1, None, - pytest.raises( - ValueError, match="window_size must be greater or equal than 1" + partial( + pytest.raises, + ValueError, + match="window_size must be greater or equal than 1", ), ), ( 2, -1, - pytest.raises( - ValueError, match="min_samples must be greater or equal than 1" + partial( + pytest.raises, + ValueError, + match="min_samples must be greater or equal than 1", ), ), ( 1, 2, - pytest.raises( + partial( + pytest.raises, InvalidOperationError, match="`min_samples` must be less or equal than `window_size`", ), @@ -250,12 +269,20 @@ def test_rolling_sum_expr_invalid_params( ( 4.2, None, - pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+window_size="), + partial( + pytest.raises, + TypeError, + match=r"Expected '.+?', got: '.+?'\s+window_size=", + ), ), ( 2, 4.2, - pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+min_samples="), + partial( + pytest.raises, + TypeError, + match=r"Expected '.+?', got: '.+?'\s+min_samples=", + ), ), ], ) @@ -267,7 +294,7 @@ def test_rolling_sum_series_invalid_params( ) -> None: df = nw.from_native(constructor_eager(data)) - with context: + with context(): df["a"].rolling_sum(window_size=window_size, min_samples=min_samples) diff --git a/tests/expr_and_series/shift_test.py b/tests/expr_and_series/shift_test.py index a06ff4a872..f15b2616ea 100644 --- a/tests/expr_and_series/shift_test.py +++ b/tests/expr_and_series/shift_test.py @@ -1,6 +1,7 @@ from __future__ import annotations from contextlib import nullcontext +from functools import partial from typing import Any import pytest @@ -21,6 +22,8 @@ "c": [5, 4, 3, 2, 1], } +TYPE_ERROR_MSG = r"Expected '.+?', got: '.+?'\s+n=" + def test_shift(constructor_eager: ConstructorEager) -> None: df = nw.from_native(constructor_eager(data)) @@ -103,31 +106,31 @@ def test_shift_multi_chunk_pyarrow() -> None: @pytest.mark.parametrize( ("n", "context"), [ - (1.0, pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+n=")), - ("1", pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+n=")), - (None, pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+n=")), - (1, nullcontext()), - (0, nullcontext()), + (1.0, partial(pytest.raises, TypeError, match=TYPE_ERROR_MSG)), + ("1", partial(pytest.raises, TypeError, match=TYPE_ERROR_MSG)), + (None, partial(pytest.raises, TypeError, match=TYPE_ERROR_MSG)), + (1, nullcontext), + (0, nullcontext), ], ) def test_shift_expr_invalid_params(n: Any, context: Any) -> None: - with context: + with context(): nw.col("a").shift(n) @pytest.mark.parametrize( ("n", "context"), [ - (1.0, pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+n=")), - ("1", pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+n=")), - (None, pytest.raises(TypeError, match=r"Expected '.+?', got: '.+?'\s+n=")), - (1, nullcontext()), - (0, nullcontext()), + (1.0, partial(pytest.raises, TypeError, match=TYPE_ERROR_MSG)), + ("1", partial(pytest.raises, TypeError, match=TYPE_ERROR_MSG)), + (None, partial(pytest.raises, TypeError, match=TYPE_ERROR_MSG)), + (1, nullcontext), + (0, nullcontext), ], ) def test_shift_series_invalid_params( constructor_eager: ConstructorEager, n: Any, context: Any ) -> None: df = nw.from_native(constructor_eager(data), eager_only=True) - with context: + with context(): df["a"].shift(n) diff --git a/tests/frame/interchange_native_namespace_test.py b/tests/frame/interchange_native_namespace_test.py index 79a92ef6c9..a2e0961e03 100644 --- a/tests/frame/interchange_native_namespace_test.py +++ b/tests/frame/interchange_native_namespace_test.py @@ -60,9 +60,10 @@ def test_duckdb() -> None: pytest.importorskip("duckdb") import duckdb - df_pl = pl.DataFrame(data) # noqa: F841 + _df = pl.DataFrame(data) - rel = duckdb.sql("select * from df_pl") + con = duckdb.connect() + rel = con.sql("select * from _df") df = nw_v1.from_native(rel, eager_or_interchange_only=True) series = df["a"] diff --git a/tests/frame/interchange_schema_test.py b/tests/frame/interchange_schema_test.py index a9177278fe..f87c8e43e5 100644 --- a/tests/frame/interchange_schema_test.py +++ b/tests/frame/interchange_schema_test.py @@ -168,7 +168,7 @@ def test_interchange_schema_duckdb() -> None: pytest.importorskip("duckdb") import duckdb - df_pl = pl.DataFrame( # noqa: F841 + _df = pl.DataFrame( { "a": [1, 1, 2], "b": [4, 5, 6], @@ -206,7 +206,8 @@ def test_interchange_schema_duckdb() -> None: "p": pl.Boolean, }, ) - rel = duckdb.sql("select * from df_pl") + con = duckdb.connect() + rel = con.sql("select * from _df") df = nw_v1.from_native(rel, eager_or_interchange_only=True) result = df.schema expected = { diff --git a/tests/frame/interchange_select_test.py b/tests/frame/interchange_select_test.py index a927ba18c6..0eabf1ec3a 100644 --- a/tests/frame/interchange_select_test.py +++ b/tests/frame/interchange_select_test.py @@ -84,8 +84,9 @@ def test_interchange_duckdb() -> None: import duckdb import polars as pl - df_pl = pl.DataFrame(data) # noqa: F841 - rel = duckdb.sql("select * from df_pl") + _df = pl.DataFrame(data) + con = duckdb.connect() + rel = con.sql("select * from _df") df = nw_v1.from_native(rel, eager_or_interchange_only=True) out_cols = df.select("a", "z").schema.names() diff --git a/tests/frame/interchange_to_arrow_test.py b/tests/frame/interchange_to_arrow_test.py index 2277d498ea..e97222afcc 100644 --- a/tests/frame/interchange_to_arrow_test.py +++ b/tests/frame/interchange_to_arrow_test.py @@ -58,8 +58,9 @@ def test_interchange_duckdb_to_arrow() -> None: import polars as pl import pyarrow as pa - df_pl = pl.DataFrame(data) # noqa: F841 - rel = duckdb.sql("select * from df_pl") + _df = pl.DataFrame(data) + con = duckdb.connect() + rel = con.sql("select * from _df") df = nw_v1.from_native(rel, eager_or_interchange_only=True) result = df.to_arrow() diff --git a/tests/frame/interchange_to_pandas_test.py b/tests/frame/interchange_to_pandas_test.py index 254e712a51..f4714460a7 100644 --- a/tests/frame/interchange_to_pandas_test.py +++ b/tests/frame/interchange_to_pandas_test.py @@ -51,7 +51,8 @@ def test_interchange_duckdb_to_pandas(request: pytest.FixtureRequest) -> None: request.applymarker(pytest.mark.xfail) df_raw = pd.DataFrame(data) - rel = duckdb.sql("select * from df_raw") + con = duckdb.connect() + rel = con.sql("select * from df_raw") df = nw_v1.from_native(rel, eager_or_interchange_only=True) assert df.to_pandas().equals(df_raw) diff --git a/tests/frame/pivot_test.py b/tests/frame/pivot_test.py index 260006555e..c78a9b5dd6 100644 --- a/tests/frame/pivot_test.py +++ b/tests/frame/pivot_test.py @@ -1,6 +1,7 @@ from __future__ import annotations from contextlib import nullcontext as does_not_raise +from functools import partial from typing import Any import pytest @@ -141,8 +142,8 @@ def test_pivot( @pytest.mark.parametrize( ("data_", "context"), [ - (data_no_dups, does_not_raise()), - (data, pytest.raises((ValueError, NarwhalsError))), + (data_no_dups, does_not_raise), + (data, partial(pytest.raises, (ValueError, NarwhalsError))), ], ) def test_pivot_no_agg( @@ -155,7 +156,7 @@ def test_pivot_no_agg( request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor_eager(data_), eager_only=True) - with context: + with context(): df.pivot("col", index="ix", aggregate_function=None) diff --git a/tests/frame/schema_test.py b/tests/frame/schema_test.py index 1555cc35a7..8dec9f48b7 100644 --- a/tests/frame/schema_test.py +++ b/tests/frame/schema_test.py @@ -255,7 +255,8 @@ def test_validate_not_duplicated_columns_duckdb() -> None: pytest.importorskip("duckdb") import duckdb - rel = duckdb.sql("SELECT 1 AS a, 2 AS a") + con = duckdb.connect() + rel = con.sql("SELECT 1 AS a, 2 AS a") with pytest.raises( ValueError, match="Expected unique column names, got:\n- 'a' 2 times" ): @@ -301,7 +302,8 @@ def test_nested_dtypes() -> None: "b": nw.Array(nw.Int64, 2), "c": nw.Struct({"a": nw.Int64, "b": nw.String, "c": nw.Float64}), } - rel = duckdb.sql("select * from df_pa") + con = duckdb.connect() + rel = con.sql("select * from df_pa") nwdf = nw.from_native(rel) assert nwdf.collect_schema() == { "a": nw.List(nw.Int64), diff --git a/tests/free_threading_test.py b/tests/free_threading_test.py new file mode 100644 index 0000000000..efd99cddc0 --- /dev/null +++ b/tests/free_threading_test.py @@ -0,0 +1,335 @@ +"""Concurrency stress tests for narwhals-owned shared state. + +Targets: caches populated on first use, lazily-materialized dtype metadata, +expression `over` push-down, state stashed on shared Narwhals objects, and +`narwhals.sql`'s shared DuckDB catalog. + +See [docs/concepts/thread_safety.md]. + +NOTE: The tests are valid on any build (races are bugs under the GIL too), but are +most effective on a free-threaded build (`PYTHON_GIL=0`), where threads run in parallel. +""" + +from __future__ import annotations + +import sys +import sysconfig +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Any + +import pytest + +import narwhals as nw +from tests.utils import POLARS_VERSION, PYARROW_VERSION, assert_equal_data + +if TYPE_CHECKING: + from collections.abc import Callable + from concurrent.futures import Future + + from tests.utils import ConstructorEager + +DATA: dict[str, Any] = { + "g": [1, 1, 2, 2, 3, 3], + "i": [6, 5, 4, 3, 2, 1], + "v": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], +} +"""`i` is a unique ordering column, reversing the row order.""" + +_SELECT_TZ = "select timestamptz '2024-01-01' as t" +"""Time-zone-aware: forces Narwhals to query the connection for its time zone.""" + + +def run_threaded( + func: Callable[..., None], + max_workers: int = 8, + *, + outer_iterations: int = 1, + prepare_args: Callable[[], list[Any]] | None = None, +) -> None: + """Run `func` in `max_workers` threads at once, `outer_iterations` times. + + Each thread receives a shared `threading.Barrier` as its final positional argument, + so callers can line every thread up on the racy window with `barrier.wait()` before proceeding. + `prepare_args`, when given, is called once per iteration to build the arguments that precede + the barrier (e.g. to reset caches). + + Adapted from NumPy's `run_threaded` test helper, the pattern recommended by + https://py-free-threading.github.io/testing/. + Source: https://github.com/numpy/numpy/blob/7e1f94485495485c5cc3a408ab9e945940f1b91f/numpy/testing/_private/utils.py#L2831 + Copyright (c) 2005-2025, NumPy Developers. License: BSD 3-Clause. + """ + for _ in range(outer_iterations): + with ThreadPoolExecutor(max_workers=max_workers) as tpe: + args = [] if prepare_args is None else prepare_args() + barrier = threading.Barrier(max_workers) + args.append(barrier) + futures: list[Future[None]] = [] + try: + futures.extend(tpe.submit(func, *args) for _ in range(max_workers)) + except RuntimeError as e: # pragma: no cover + # Release any threads already blocked on the barrier so the pool can + # shut down instead of deadlocking, then skip. + barrier.abort() + pytest.skip( + f"Spawning {max_workers} threads failed with error {e!r} " + "(likely due to resource limits on the system running the tests)" + ) + for f in futures: + f.result() + + +def test_gil_stays_disabled_on_free_threaded_build() -> None: # pragma: no cover + if not sysconfig.get_config_var("Py_GIL_DISABLED"): + pytest.skip("not a free-threaded build") + # NOTE: Only reached on a free-threaded build, never on the GIL-enabled coverage jobs. + # Because of this, the entire function is flagged as "pragma: no cover" + is_gil_enabled = getattr(sys, "_is_gil_enabled", lambda: True) + assert not is_gil_enabled(), ( + "The GIL was re-enabled, likely by importing an extension module " + "without free-threading support." + ) + + +def test_from_native_cold_caches() -> None: + pytest.importorskip("pyarrow") + import pyarrow as pa + + from narwhals import _utils + + tbl = pa.table({"a": [1, 2, 3]}) + + def clear_caches() -> list[Any]: + _utils.backend_version.cache_clear() + _utils._import_native_namespace.cache_clear() + _utils._version_namespace.cache_clear() + _utils._version_dtypes.cache_clear() + _utils._version_dataframe.cache_clear() + _utils._version_lazyframe.cache_clear() + _utils._version_series.cache_clear() + return [] + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + nw_df = nw.from_native(tbl, eager_only=True) + assert nw_df["a"].sum() == 6 + + run_threaded(check, outer_iterations=3, prepare_args=clear_caches) + + +def test_plugin_discovery_cold_cache() -> None: + from narwhals import plugins + + def clear_caches() -> list[Any]: + plugins._discover_entrypoints.cache_clear() + return [] + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + assert plugins._discover_entrypoints() is not None + + run_threaded(check, outer_iterations=3, prepare_args=clear_caches) + + +def test_shared_expr_over_push_down() -> None: + # `.over()` must never mutate nodes reachable from the original, + # potentially shared, expression. + def make_expr() -> nw.Expr: + return nw.col("a").cum_sum() + nw.col("b").cum_sum().abs() + + base = make_expr() + expected_base = repr(base) + expected_over = repr(make_expr().over("g", order_by="i")) + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + result = base.over("g", order_by="i") + assert repr(result) == expected_over + assert repr(base) == expected_base + + run_threaded(check, outer_iterations=5) + + +def test_shared_expr_evaluation(constructor_eager: ConstructorEager) -> None: + """Evaluating one shared `Expr` in several contexts at once must not mutate it.""" + if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 10): + pytest.skip("`over(order_by=...)` requires polars>=1.10") + + df = nw.from_native(constructor_eager(DATA), eager_only=True) + expr = nw.col("v").cum_sum() + expected_repr = repr(expr) + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + for _ in range(10): + # Row order, `i` order, and one appended node: three rewrites of `expr`. + assert_equal_data(df.select(expr), {"v": [1.0, 3.0, 6.0, 10.0, 15.0, 21.0]}) + assert_equal_data( + df.select(expr.over(order_by="i")), + {"v": [21.0, 20.0, 18.0, 15.0, 11.0, 6.0]}, + ) + assert_equal_data( + df.with_columns(out=expr.abs()).select("out"), + {"out": [1.0, 3.0, 6.0, 10.0, 15.0, 21.0]}, + ) + assert repr(expr) == expected_repr + + run_threaded(check, outer_iterations=2) + + +def test_shared_group_by_agg(constructor_eager: ConstructorEager) -> None: + """A `GroupBy` reused from several threads must not leak state between them. + + Regression test: `agg` used to stash the native `groupby` on `self`, so one thread + could read another's grouping and silently aggregate the wrong rows. + """ + if "pyarrow_table" in str(constructor_eager) and PYARROW_VERSION < (14, 0): + pytest.skip("https://github.com/apache/arrow/issues/36709") + + grouped = nw.from_native(constructor_eager(DATA), eager_only=True).group_by("g") + # `sum` groups as-is, `first(order_by="i")` groups a sorted copy: two distinct + # native groupings, so a leak between threads shows up in the values. + unordered = {"g": [1, 2, 3], "v": [3.0, 7.0, 11.0]} + ordered = {"g": [1, 2, 3], "v": [2.0, 4.0, 6.0]} + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + for _ in range(20): + res_ordered = grouped.agg(nw.col("v").first(order_by="i")).sort("g") + assert_equal_data(res_ordered, ordered) + + res_unordered = grouped.agg(nw.col("v").sum()).sort("g") + assert_equal_data(res_unordered, unordered) + + run_threaded(check, outer_iterations=3) + + +def test_shared_dataframe_read_only(constructor_eager: ConstructorEager) -> None: + """Reading from a shared `DataFrame` is safe: no method may mutate it, or its input.""" + native = constructor_eager(DATA) + df = nw.from_native(native, eager_only=True) + native_before = repr(native) + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + for _ in range(10): + assert df.columns == ["g", "i", "v"] + assert df.schema == {"g": nw.Int64(), "i": nw.Int64(), "v": nw.Float64()} + assert df.shape == (6, 3) + assert_equal_data( + df.select(nw.col("v") * 2), {"v": [2.0, 4.0, 6.0, 8.0, 10.0, 12.0]} + ) + assert_equal_data( + df.filter(nw.col("g") == 1), {"g": [1, 1], "i": [6, 5], "v": [1.0, 2.0]} + ) + assert_equal_data( + df.sort("i").select("v"), {"v": [6.0, 5.0, 4.0, 3.0, 2.0, 1.0]} + ) + assert_equal_data(df.unique("g").sort("g").select("g"), {"g": [1, 2, 3]}) + assert_equal_data(df.lazy().collect().select("g"), {"g": [1, 1, 2, 2, 3, 3]}) + # `scatter` is the one method that reads as in-place: it must not be. + assert df["v"].scatter(0, 99.0).to_list() == [99.0, 2.0, 3.0, 4.0, 5.0, 6.0] + assert df["v"].to_list() == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + + run_threaded(check, outer_iterations=2) + assert repr(native) == native_before, "narwhals mutated the native input" + + +def test_enum_deferred_categories() -> None: + pytest.importorskip("polars") + import polars as pl + + from narwhals._polars.utils import native_to_narwhals_dtype + + categories = ("ft_x", "ft_y", "ft_z") + df = pl.DataFrame({"a": pl.Series(["ft_x"], dtype=pl.Enum(categories))}) + + def clear_caches() -> list[Any]: + native_to_narwhals_dtype.cache_clear() + return [] + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + dtype = nw.from_native(df, eager_only=True).schema["a"] + assert isinstance(dtype, nw.Enum) + assert dtype.categories == categories + + run_threaded(check, outer_iterations=3, prepare_args=clear_caches) + + +def test_shared_lazyframe_schema() -> None: + pytest.importorskip("duckdb") + import duckdb + + con = duckdb.connect() + rel = con.sql("select 1::BIGINT as a, 'x' as b") + lf = nw.from_native(rel) + expected = {"a": nw.Int64(), "b": nw.String()} + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + assert lf.collect_schema() == expected + assert lf.columns == ["a", "b"] + + run_threaded(check, outer_iterations=5) + + +def test_shared_duckdb_connection_schema() -> None: + pytest.importorskip("duckdb") + import duckdb + + con = duckdb.connect() + con.sql("set timezone = 'UTC'") + # One frame per thread *and* one shared frame, all on the same connection. + frames = [nw.from_native(con.sql(_SELECT_TZ)) for _ in range(8)] + shared = nw.from_native(con.sql(_SELECT_TZ)) + expected = {"t": nw.Datetime(time_zone="UTC")} + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + for lf in (*frames, shared): + assert lf.collect_schema() == expected + + run_threaded(check, outer_iterations=3) + + +def test_duckdb_per_thread_cursor() -> None: + """The supported recipe for concurrent DuckDB use: one cursor per thread. + + NOTE: `TimeZone` is `LOCAL`-scoped, so a cursor does not inherit the parent's value. + """ + pytest.importorskip("duckdb") + import duckdb + + con = duckdb.connect() + + def check(barrier: threading.Barrier) -> None: + cursor = con.cursor() + cursor.sql("set timezone = 'UTC'") + barrier.wait() + for _ in range(5): + lf = nw.from_native(cursor.sql(f"{_SELECT_TZ}, 1 as idx, 2 as a, 3 as b")) + assert lf.collect_schema()["t"] == nw.Datetime(time_zone="UTC") + assert_equal_data(lf.select("a").collect(), {"a": [2]}) + assert_equal_data( + lf.unpivot(on=["a", "b"], index=["idx"]).sort("variable"), + {"idx": [1, 1], "variable": ["a", "b"], "value": [2, 3]}, + ) + + run_threaded(check, outer_iterations=3) + + +def test_sql_table_concurrent() -> None: + pytest.importorskip("duckdb", minversion="1.3.0") + from narwhals.sql import table + + def check(barrier: threading.Barrier) -> None: + barrier.wait() + name = f"tbl_{uuid.uuid4().hex}" + result = table(name, {"a": nw.Int64(), "b": nw.String()}) + assert result.collect_schema() == {"a": nw.Int64(), "b": nw.String()} + assert name in result.to_sql() + assert result.to_native().fetchall() == [] + + run_threaded(check, outer_iterations=5) diff --git a/tests/read_scan_test.py b/tests/read_scan_test.py index 532502691c..b867cf6b62 100644 --- a/tests/read_scan_test.py +++ b/tests/read_scan_test.py @@ -137,6 +137,10 @@ def test_scan_csv( kwargs = {"session": sqlframe_session(), "inferSchema": True, "header": True} elif "pyspark" in str(constructor): kwargs = {"session": pyspark_session(), "inferSchema": True, "header": True} + elif "duckdb" in str(constructor): + import duckdb + + kwargs = {"connection": duckdb.connect()} else: kwargs = {} backend = native_namespace(constructor) @@ -181,12 +185,28 @@ def test_scan_parquet(parquet_path: FileSource, constructor: Constructor) -> Non kwargs = {"session": sqlframe_session(), "inferSchema": True} elif "pyspark" in str(constructor): kwargs = {"session": pyspark_session(), "inferSchema": True, "header": True} + elif "duckdb" in str(constructor): + import duckdb + + kwargs = {"connection": duckdb.connect()} else: kwargs = {} backend = native_namespace(constructor) assert_equal_lazy(nw.scan_parquet(parquet_path, backend=backend, **kwargs)) +# NOTE: Marked thread_unsafe on purpose +@pytest.mark.thread_unsafe( + reason="reads through the process-global duckdb default connection" +) +def test_scan_duckdb_default_connection( + csv_path: FileSource, parquet_path: FileSource +) -> None: + pytest.importorskip("duckdb") + assert_equal_lazy(nw.scan_csv(csv_path, backend="duckdb")) + assert_equal_lazy(nw.scan_parquet(parquet_path, backend="duckdb")) + + @skipif_pandas_lt_1_5 def test_scan_parquet_kwargs(parquet_path: FileSource) -> None: pytest.importorskip("pandas") diff --git a/tests/repr_test.py b/tests/repr_test.py index 04880c4392..6cc82d248c 100644 --- a/tests/repr_test.py +++ b/tests/repr_test.py @@ -53,7 +53,8 @@ def test_repr(request: pytest.FixtureRequest) -> None: "└─────────────────────┘" ) assert result == expected - result = nw.from_native(duckdb.table("df")).__repr__() + con = duckdb.connect() + result = nw.from_native(con.table("df")).__repr__() expected = ( "┌───────────────────┐\n" "|Narwhals LazyFrame |\n" @@ -71,7 +72,8 @@ def test_repr(request: pytest.FixtureRequest) -> None: assert result == expected # Make something wider than the terminal size df = pd.DataFrame({"a": [1, 2, 3], "b": ["fdaf" * 100, "fda", "cf"]}) - result = nw.from_native(duckdb.table("df")).__repr__() + con = duckdb.connect() + result = nw.from_native(con.table("df")).__repr__() expected = ( "┌───────────────────────────────────────┐\n" "| Narwhals LazyFrame |\n" diff --git a/tests/sql_test.py b/tests/sql_test.py index 4f0b697929..d986b280c4 100644 --- a/tests/sql_test.py +++ b/tests/sql_test.py @@ -1,42 +1,62 @@ from __future__ import annotations +import uuid + import pytest import narwhals as nw -from tests.utils import DUCKDB_VERSION + +pytest.importorskip("duckdb", minversion="1.3.0") + +from narwhals.sql import table + + +def _unique_name(prefix: str) -> str: + # `narwhals.sql.table` creates the table in a process-wide catalog, so + # names must be unique for tests to be re-runnable and thread-safe. + return f"{prefix}_{uuid.uuid4().hex}" def test_sql() -> None: - pytest.importorskip("duckdb") pytest.importorskip("sqlparse") - if DUCKDB_VERSION < (1, 3): - pytest.skip() - from narwhals.sql import table + name = _unique_name("assets") schema = {"date": nw.Date(), "price": nw.Int64(), "symbol": nw.String()} - assets = table("assets", schema) + assets = table(name, schema) result = assets.with_columns( returns=(nw.col("price") / nw.col("price").shift(1)).over( "symbol", order_by="date" ) ) - expected = """SELECT date, price, symbol, (price / lag(price, 1) OVER (PARTITION BY symbol ORDER BY date ASC NULLS FIRST)) AS "returns" FROM main.assets""" + expected = f"""SELECT date, price, symbol, (price / lag(price, 1) OVER (PARTITION BY symbol ORDER BY date ASC NULLS FIRST)) AS "returns" FROM main.{name}""" # noqa: S608 assert result.to_sql() == expected expected = ( "SELECT date, price,\n" " symbol,\n" " (price / lag(price, 1) OVER (PARTITION BY symbol\n" ' ORDER BY date ASC NULLS FIRST)) AS "returns"\n' - "FROM main.assets" + f"FROM main.{name}" ) assert result.to_sql(pretty=True) == expected def test_sql_table_schema_pairs() -> None: - pytest.importorskip("duckdb") - if DUCKDB_VERSION < (1, 3): - pytest.skip() - from narwhals.sql import table - - result = table("assets_pairs", [("date", nw.Date), ("price", nw.Int64())]) + name = _unique_name("assets_pairs") + result = table(name, [("date", nw.Date), ("price", nw.Int64())]) assert result.collect_schema() == {"date": nw.Date(), "price": nw.Int64()} + + +def test_sql_table_combine() -> None: + """Tables created in the same thread must be combinable.""" + lhs_name, rhs_name = _unique_name("lhs"), _unique_name("rhs") + lhs = table(lhs_name, {"id": nw.Int64(), "x": nw.Int64()}) + rhs = table(rhs_name, {"id": nw.Int64(), "y": nw.Int64()}) + + joined = lhs.join(rhs, on="id", how="inner").to_sql() + assert f"main.{lhs_name}" in joined + assert f"main.{rhs_name}" in joined + + # `nw.concat` is typed as a plain `LazyFrame`, so go through the relation. + concatenated = nw.concat([lhs.select("id"), rhs.select("id")]).to_native().sql_query() + assert f"main.{lhs_name}" in concatenated + assert f"main.{rhs_name}" in concatenated diff --git a/tests/testing/assert_frame_equal_test.py b/tests/testing/assert_frame_equal_test.py index 2ca5e0f498..0392d4362c 100644 --- a/tests/testing/assert_frame_equal_test.py +++ b/tests/testing/assert_frame_equal_test.py @@ -2,6 +2,7 @@ import re from contextlib import AbstractContextManager, nullcontext as does_not_raise +from functools import partial from typing import TYPE_CHECKING, Any import pytest @@ -12,7 +13,7 @@ from tests.utils import PANDAS_VERSION if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Callable, Mapping from narwhals.typing import IntoDType from tests.conftest import Data @@ -71,7 +72,7 @@ def test_check_same_input_type(constructor_eager: ConstructorEager) -> None: {"a": nw.Int32(), "b": nw.Float32()}, True, True, - does_not_raise(), + does_not_raise, ), # Same order, different dtypes ( @@ -79,14 +80,14 @@ def test_check_same_input_type(constructor_eager: ConstructorEager) -> None: {"a": nw.Int32(), "b": nw.Float64()}, False, True, - does_not_raise(), + does_not_raise, ), ( {"a": nw.Int32(), "b": nw.Float32()}, {"a": nw.Int32(), "b": nw.Float64()}, True, True, - _assertion_error("dtypes do not match"), + partial(_assertion_error, "dtypes do not match"), ), # Different order, same dtype ( @@ -94,14 +95,14 @@ def test_check_same_input_type(constructor_eager: ConstructorEager) -> None: {"b": nw.Float32(), "a": nw.Int32()}, True, False, - does_not_raise(), + does_not_raise, ), ( {"a": nw.Int32(), "b": nw.Float32()}, {"b": nw.Float32(), "a": nw.Int32()}, True, True, - _assertion_error("columns are not in the same order"), + partial(_assertion_error, "columns are not in the same order"), ), # Different order, different dtype ( @@ -109,28 +110,28 @@ def test_check_same_input_type(constructor_eager: ConstructorEager) -> None: {"b": nw.Float64(), "a": nw.Int16()}, False, False, - does_not_raise(), + does_not_raise, ), ( {"a": nw.Int32(), "b": nw.Float32()}, {"b": nw.Float64(), "a": nw.Int16()}, True, False, - _assertion_error("dtypes do not match"), + partial(_assertion_error, "dtypes do not match"), ), ( {"a": nw.Int32(), "b": nw.Float32()}, {"b": nw.Float64(), "a": nw.Int16()}, False, True, - _assertion_error("columns are not in the same order"), + partial(_assertion_error, "columns are not in the same order"), ), ( {"a": nw.Int32(), "b": nw.Float32()}, {"b": nw.Float64(), "a": nw.Int16()}, True, True, - _assertion_error("columns are not in the same order"), + partial(_assertion_error, "columns are not in the same order"), ), # Different columns (left not in right) ( @@ -138,7 +139,7 @@ def test_check_same_input_type(constructor_eager: ConstructorEager) -> None: {"b": nw.Float64()}, True, True, - _assertion_error("['a', 'z'] in left, but not in right"), + partial(_assertion_error, "['a', 'z'] in left, but not in right"), ), # Different columns (right not in left) ( @@ -146,7 +147,7 @@ def test_check_same_input_type(constructor_eager: ConstructorEager) -> None: {"z": nw.String(), "b": nw.Float64()}, True, True, - _assertion_error("['b'] in right, but not in left"), + partial(_assertion_error, "['b'] in right, but not in left"), ), ], ) @@ -157,7 +158,7 @@ def test_check_schema_mismatch( *, check_dtypes: bool, check_column_order: bool, - context: AbstractContextManager[Any], + context: Callable[[], AbstractContextManager[Any]], ) -> None: data = {"a": [1, 2, 3], "b": [4.5, 6.7, 8.9], "z": ["foo", "bar", "baz"]} left = nw.from_native(constructor(data)).select( @@ -167,7 +168,7 @@ def test_check_schema_mismatch( nw.col(name).cast(dtype) for name, dtype in right_schema.items() ) - with context: + with context(): assert_frame_equal( left, right, check_column_order=check_column_order, check_dtypes=check_dtypes ) @@ -206,12 +207,12 @@ def test_check_row_order( ) context = ( - _assertion_error('value mismatch for column "a"') + partial(_assertion_error, 'value mismatch for column "a"') if check_row_order and left.implementation in GUARANTEES_ROW_ORDER - else does_not_raise() + else does_not_raise ) - with context: + with context(): assert_frame_equal(left, right, check_row_order=check_row_order) diff --git a/tests/testing/assert_series_equal_test.py b/tests/testing/assert_series_equal_test.py index 0517fa8aaf..7cf22913a6 100644 --- a/tests/testing/assert_series_equal_test.py +++ b/tests/testing/assert_series_equal_test.py @@ -2,6 +2,7 @@ import re from contextlib import AbstractContextManager, nullcontext as does_not_raise +from functools import partial from typing import TYPE_CHECKING, Any import pytest @@ -135,10 +136,10 @@ def test_metadata_checks_with_flags( @pytest.mark.parametrize( ("dtype", "check_order", "context"), [ - (nw.List(nw.Int32()), False, pytest.raises(NotImplementedError)), - (nw.List(nw.Int32()), True, does_not_raise()), - (nw.Int32(), False, does_not_raise()), - (nw.Int32(), True, does_not_raise()), + (nw.List(nw.Int32()), False, partial(pytest.raises, NotImplementedError)), + (nw.List(nw.Int32()), True, does_not_raise), + (nw.Int32(), False, does_not_raise), + (nw.Int32(), True, does_not_raise), ], ) def test_check_order( @@ -147,7 +148,7 @@ def test_check_order( dtype: nw.dtypes.DType, *, check_order: bool, - context: AbstractContextManager[Any], + context: Callable[[], AbstractContextManager[Any]], ) -> None: """Test check_order behavior with nested and simple data.""" if "cudf" in str(constructor_eager) and dtype.is_nested(): @@ -164,7 +165,7 @@ def test_check_order( frame = nw.from_native(constructor_eager({"a": data}), eager_only=True) left = right = frame["a"].cast(dtype) - with context: + with context(): assert_series_equal(left, right, check_order=check_order, check_names=False) @@ -186,9 +187,9 @@ def test_null_mismatch(constructor_eager: ConstructorEager, null_data: Data) -> @pytest.mark.parametrize( ("check_exact", "abs_tol", "rel_tol", "context"), [ - (True, 1e-3, 1e-3, _assertion_error("exact value mismatch")), - (False, 1e-3, 1e-3, _assertion_error("values not within tolerance")), - (False, 2e-1, 2e-1, does_not_raise()), + (True, 1e-3, 1e-3, partial(_assertion_error, "exact value mismatch")), + (False, 1e-3, 1e-3, partial(_assertion_error, "values not within tolerance")), + (False, 2e-1, 2e-1, does_not_raise), ], ) def test_numeric( @@ -197,7 +198,7 @@ def test_numeric( check_exact: bool, abs_tol: float, rel_tol: float, - context: AbstractContextManager[Any], + context: Callable[[], AbstractContextManager[Any]], ) -> None: data = { "left": [1.0, float("nan"), float("inf"), None, 1.1], @@ -206,7 +207,7 @@ def test_numeric( frame = nw.from_native(constructor_eager(data), eager_only=True) left, right = frame["left"], frame["right"] - with context: + with context(): assert_series_equal( left, right, @@ -224,36 +225,36 @@ def test_numeric( [["foo", "bar"]], [["foo", None]], True, - _assertion_error("nested value mismatch"), + partial(_assertion_error, "nested value mismatch"), nw.List(nw.String()), ), ( [["foo", "bar"]], [["foo", None]], True, - _assertion_error("nested value mismatch"), + partial(_assertion_error, "nested value mismatch"), nw.Array(nw.String(), 2), ), ( [[0.0, 0.1]], [[0.1, 0.1]], True, - _assertion_error("nested value mismatch"), + partial(_assertion_error, "nested value mismatch"), nw.List(nw.Float32()), ), ( [[0.0, 0.1]], [[0.1, 0.1]], True, - _assertion_error("nested value mismatch"), + partial(_assertion_error, "nested value mismatch"), nw.Array(nw.Float32(), 2), ), - ([[0.0, 1e-10]], [[1e-10, 0.0]], False, does_not_raise(), nw.List(nw.Float64())), + ([[0.0, 1e-10]], [[1e-10, 0.0]], False, does_not_raise, nw.List(nw.Float64())), ( [[0.0, 1e-10]], [[1e-10, 0.0]], False, - does_not_raise(), + does_not_raise, nw.Array(nw.Float64(), 2), ), ], @@ -265,7 +266,7 @@ def test_list_like( r_vals: list[list[Any]], *, check_exact: bool, - context: AbstractContextManager[Any], + context: Callable[[], AbstractContextManager[Any]], dtype: nw.dtypes.DType, ) -> None: if "cudf" in str(constructor_eager): @@ -292,7 +293,7 @@ def test_list_like( data = {"left": l_vals, "right": r_vals} frame = nw.from_native(constructor_eager(data), eager_only=True) left, right = frame["left"].cast(dtype), frame["right"].cast(dtype) - with context: + with context(): assert_series_equal(left, right, check_names=False, check_exact=check_exact) @@ -303,19 +304,19 @@ def test_list_like( [{"a": 0.0, "b": ["orca"]}, None], [{"a": 1e-10, "b": ["orca"]}, None], True, - _assertion_error("exact value mismatch"), + partial(_assertion_error, "exact value mismatch"), ), ( [{"a": 0.0, "b": ["beluga"]}, None], [{"a": 0.0, "b": ["orca"]}, None], False, - _assertion_error("exact value mismatch"), + partial(_assertion_error, "exact value mismatch"), ), ( [{"a": 0.0, "b": ["orca"]}, None], [{"a": 1e-10, "b": ["orca"]}, None], False, - does_not_raise(), + does_not_raise, ), ], ) @@ -326,7 +327,7 @@ def test_struct( r_vals: list[dict[str, Any]], *, check_exact: bool, - context: AbstractContextManager[Any], + context: Callable[[], AbstractContextManager[Any]], ) -> None: if "cudf" in str(constructor_eager): reason = "NotImplementedError" @@ -342,7 +343,7 @@ def test_struct( data = {"left": l_vals, "right": r_vals} frame = nw.from_native(constructor_eager(data), eager_only=True) left, right = frame["left"].cast(dtype), frame["right"].cast(dtype) - with context: + with context(): assert_series_equal(left, right, check_names=False, check_exact=check_exact) @@ -363,10 +364,11 @@ def test_non_nw_series() -> None: @pytest.mark.parametrize( ("categorical_as_str", "context"), [ - (True, does_not_raise()), + (True, does_not_raise), ( False, - pytest.raises( + partial( + pytest.raises, AssertionError, match="Cannot compare categoricals coming from different sources", ), @@ -378,7 +380,7 @@ def test_categorical_as_str( constructor_eager: ConstructorEager, *, categorical_as_str: bool, - context: AbstractContextManager[Any], + context: Callable[[], AbstractContextManager[Any]], ) -> None: if ( "polars" in str(constructor_eager) @@ -411,7 +413,7 @@ def test_categorical_as_str( left = frame["left"].cast(nw.Categorical())[2:] right = frame["right"].cast(nw.Categorical())[2:] - with context: + with context(): assert_series_equal( left, right, check_names=False, categorical_as_str=categorical_as_str ) diff --git a/tests/translate/from_native_test.py b/tests/translate/from_native_test.py index 9f87b5220b..5995331203 100644 --- a/tests/translate/from_native_test.py +++ b/tests/translate/from_native_test.py @@ -22,6 +22,7 @@ # Using pyright's assert type instead # mypy: disallow-any-generics=false, disable-error-code="assert-type" from contextlib import nullcontext as does_not_raise +from functools import partial from importlib.util import find_spec from itertools import chain from typing import TYPE_CHECKING, Any, Literal, cast @@ -124,12 +125,12 @@ def __narwhals_series__(self) -> Any: @pytest.mark.parametrize( ("eager_only", "context"), [ - (False, does_not_raise()), - (True, pytest.raises(TypeError, match="Cannot only use `eager_only`")), + (False, does_not_raise), + (True, partial(pytest.raises, TypeError, match="Cannot only use `eager_only`")), ], ) def test_eager_only_lazy(dframe: Any, eager_only: Any, context: Any) -> None: - with context: + with context(): res = nw.from_native(dframe, eager_only=eager_only) assert isinstance(res, nw.LazyFrame) if eager_only: @@ -147,14 +148,17 @@ def test_eager_only_eager(dframe: Any, eager_only: Any) -> None: ("obj", "context"), [ *[ - (frame, pytest.raises(TypeError, match="Cannot only use `series_only`")) + ( + frame, + partial(pytest.raises, TypeError, match="Cannot only use `series_only`"), + ) for frame in all_frames ], - *[(series, does_not_raise()) for series in all_series], + *[(series, does_not_raise) for series in all_series], ], ) def test_series_only(obj: Any, context: Any) -> None: - with context: + with context(): res = nw.from_native(obj, series_only=True) assert isinstance(res, nw.Series) assert nw.from_native(obj, series_only=True, pass_through=True) is obj or isinstance( @@ -166,17 +170,19 @@ def test_series_only(obj: Any, context: Any) -> None: @pytest.mark.parametrize( ("allow_series", "context"), [ - (True, does_not_raise()), + (True, does_not_raise), ( False, - pytest.raises( - TypeError, match="Please set `allow_series=True` or `series_only=True`" + partial( + pytest.raises, + TypeError, + match="Please set `allow_series=True` or `series_only=True`", ), ), ], ) def test_allow_series(series: Any, allow_series: Any, context: Any) -> None: - with context: + with context(): res = nw.from_native(series, allow_series=allow_series) assert isinstance(res, nw.Series) if not allow_series: @@ -275,8 +281,8 @@ def test_series_only_dask() -> None: @pytest.mark.parametrize( ("eager_only", "context"), [ - (False, does_not_raise()), - (True, pytest.raises(TypeError, match="Cannot only use `eager_only`")), + (False, does_not_raise), + (True, partial(pytest.raises, TypeError, match="Cannot only use `eager_only`")), ], ) def test_eager_only_lazy_dask(eager_only: Any, context: Any) -> None: @@ -285,7 +291,7 @@ def test_eager_only_lazy_dask(eager_only: Any, context: Any) -> None: dframe = dd.from_pandas(df_pd) - with context: + with context(): res = nw.from_native(dframe, eager_only=eager_only) assert isinstance(res, nw.LazyFrame) if eager_only: @@ -303,10 +309,11 @@ def test_series_only_sqlframe() -> None: # pragma: no cover @pytest.mark.parametrize( ("eager_only", "context"), [ - (False, does_not_raise()), + (False, does_not_raise), ( True, - pytest.raises( + partial( + pytest.raises, TypeError, match="Cannot only use `series_only`, `eager_only` or `eager_or_interchange_only` with sqlframe DataFrame", ), @@ -317,7 +324,7 @@ def test_eager_only_sqlframe(eager_only: Any, context: Any) -> None: # pragma: pytest.importorskip("sqlframe") df = sqlframe_pyspark_lazy_constructor(data) - with context: + with context(): res = nw.from_native(df, eager_only=eager_only) assert isinstance(res, nw.LazyFrame) diff --git a/tests/translate/to_native_test.py b/tests/translate/to_native_test.py index f84cf80dd6..263cdda468 100644 --- a/tests/translate/to_native_test.py +++ b/tests/translate/to_native_test.py @@ -1,6 +1,7 @@ from __future__ import annotations from contextlib import nullcontext as does_not_raise +from functools import partial from typing import TYPE_CHECKING, Any import pytest @@ -14,13 +15,13 @@ @pytest.mark.parametrize( ("method", "pass_through", "context"), [ - ("head", False, does_not_raise()), - ("head", True, does_not_raise()), - ("to_numpy", True, does_not_raise()), + ("head", False, does_not_raise), + ("head", True, does_not_raise), + ("to_numpy", True, does_not_raise), ( "to_numpy", False, - pytest.raises(TypeError, match="Expected Narwhals object, got"), + partial(pytest.raises, TypeError, match="Expected Narwhals object, got"), ), ], ) @@ -31,10 +32,10 @@ def test_to_native( pytest.importorskip("numpy") df = nw.from_native(constructor_eager({"a": [1, 2, 3]})) - with context: + with context(): nw.to_native(getattr(df, method)(), pass_through=pass_through) s = df["a"] - with context: + with context(): nw.to_native(getattr(s, method)(), pass_through=pass_through) diff --git a/tests/v1_test.py b/tests/v1_test.py index ac048112e5..56d0ac4bbe 100644 --- a/tests/v1_test.py +++ b/tests/v1_test.py @@ -5,6 +5,7 @@ from collections import deque from contextlib import nullcontext as does_not_raise from datetime import datetime, timedelta +from functools import partial from typing import TYPE_CHECKING, Any, cast import pytest @@ -369,12 +370,13 @@ def test_v1_enum_duckdb_2550() -> None: pytest.importorskip("duckdb") import duckdb + con = duckdb.connect() result_v1 = nw_v1.from_native( - duckdb.sql("select 'a'::enum('a', 'b', 'c') as a") + con.sql("select 'a'::enum('a', 'b', 'c') as a") ).collect_schema() assert result_v1 == {"a": nw_v1.Enum()} result = nw.from_native( - duckdb.sql("select 'a'::enum('a', 'b', 'c') as a") + con.sql("select 'a'::enum('a', 'b', 'c') as a") ).collect_schema() assert result == {"a": nw.Enum(("a", "b", "c"))} @@ -483,12 +485,12 @@ def test_with_row_index(constructor: Constructor) -> None: msg = "Cannot pass `order_by`" context = ( - pytest.raises(TypeError, match=msg) + partial(pytest.raises, TypeError, match=msg) if any(x in str(constructor) for x in ("duckdb", "pyspark")) - else does_not_raise() + else does_not_raise ) - with context: + with context(): result = frame.with_row_index() expected = {"index": [0, 1], **data} @@ -574,8 +576,8 @@ def test_dtypes() -> None: @pytest.mark.parametrize( ("strict", "context"), [ - (True, pytest.raises(TypeError, match="Unsupported dataframe type")), - (False, does_not_raise()), + (True, partial(pytest.raises, TypeError, match="Unsupported dataframe type")), + (False, does_not_raise), ], ) def test_strict(strict: Any, context: Any) -> None: @@ -584,7 +586,7 @@ def test_strict(strict: Any, context: Any) -> None: arr = np.array([1, 2, 3]) - with context: + with context(): res = nw_v1.from_native(arr, strict=strict) assert isinstance(res, np.ndarray) diff --git a/zensical.toml b/zensical.toml index b35b104847..774467f3ae 100644 --- a/zensical.toml +++ b/zensical.toml @@ -24,6 +24,7 @@ nav = [ "concepts/column_names.md", "concepts/boolean.md", "concepts/null_handling.md", + "concepts/thread_safety.md", ]}, {"Overhead" = "overhead.md"}, {"Perfect backwards compatibility policy" = "backcompat.md"},