diff --git a/.github/release-drafter.yml b/.github/release-drafter.yml index b3ae667ca..8f4c98811 100644 --- a/.github/release-drafter.yml +++ b/.github/release-drafter.yml @@ -9,31 +9,31 @@ autolabeler: - label: breaking title: # Example: feat!: ... - - '/^(build|chore|ci|depr|docs|feat|fix|perf|refactor|release|test)(\(.*\))?\!\: /' + - '/^([Bb]uild|[Cc]hore|CI|ci|[Dd]epr|[Dd]oc|DOC|[Ff]eat|[Ff]ix|[Pp]erf|[Rr]efactor|[Rr]elease|[Tt]est)\! /' - label: build title: - - '/^build/' + - '/^([Bb]uild)/' - label: internal title: - - '/^(chore|ci|refactor|test|template|bench)/' + - '/^(chore|ci|refactor|test|template|bench|Chore|CI|Refactor|Test|Template|Bench)/' - label: deprecation title: - - '/^depr/' + - '/^([Dd]epr)/' - label: documentation title: - - '/(.*doc|.*docstring)/' + - '/^([Dd]oc|DOC)/' - label: enhancement title: - - '/^(.*feat|.*enh)/' + - '/^(feat|enh|Feat|ENH|Enh)/' - label: fix title: - - '/^fix/' + - '/^([Ff]ix)/' - label: performance title: - - '/^perf/' + - '/^([Pp]erf)/' - label: release title: - - '/^release/' + - '/^([Rr]elease)/' version-resolver: major: diff --git a/.github/workflows/check_docs_build.yml b/.github/workflows/check_docs_build.yml index 4553a3a3a..3b418baff 100644 --- a/.github/workflows/check_docs_build.yml +++ b/.github/workflows/check_docs_build.yml @@ -1,4 +1,4 @@ -name: ci +name: Check Docs Build on: pull_request: @@ -18,20 +18,14 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Cache multiple paths - uses: actions/cache@v4 - with: - path: | - ~/.cache/pip - $RUNNER_TOOL_CACHE/Python/* - ~\AppData\Local\pip\Cache - key: ${{ runner.os }}-build-${{ matrix.python-version }} + - name: Install uv (Unix) + run: curl -LsSf https://astral.sh/uv/install.sh | sh - name: install-reqs - run: python -m pip install --upgrade tox virtualenv setuptools pip -r requirements-dev.txt + run: uv pip install --upgrade tox virtualenv setuptools pip -r requirements-dev.txt --system - name: install-docs-reqs - run: python -m pip install --upgrade -r docs/requirements-docs.txt + run: uv pip install --upgrade -r docs/requirements-docs.txt --system - name: local-install - run: python -m pip install -e . + run: uv pip install -e . --system - name: check-no-errors run: python -m mkdocs build > output.txt 2>&1 - name: assert-no-errors diff --git a/.github/workflows/downstream_tests.yml b/.github/workflows/downstream_tests.yml index 5b8e58cfd..c733e348d 100644 --- a/.github/workflows/downstream_tests.yml +++ b/.github/workflows/downstream_tests.yml @@ -1,4 +1,4 @@ -name: ci +name: Test Downstream Libraries on: pull_request: @@ -18,31 +18,25 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Cache multiple paths - uses: actions/cache@v4 - with: - path: | - ~/.cache/pip - $RUNNER_TOOL_CACHE/Python/* - ~\AppData\Local\pip\Cache - key: ${{ runner.os }}-build-${{ matrix.python-version }} + - name: Install uv (Unix) + run: curl -LsSf https://astral.sh/uv/install.sh | sh - name: clone-altair run: | git clone https://github.com/vega/altair.git --depth=1 cd altair git log - name: install-basics - run: python -m pip install --upgrade tox virtualenv setuptools pip + run: uv pip install --upgrade tox virtualenv setuptools --system - name: install-altair-dev run: | cd altair - pip install -e ".[dev, all]" + uv pip install -e ".[dev, all]" --system - name: install-narwhals-dev run: | - pip uninstall narwhals -y - pip install -e . + uv pip uninstall narwhals --system + uv pip install -e . --system - name: show-deps - run: pip freeze + run: uv pip freeze - name: Run pytest run: | cd altair @@ -64,28 +58,22 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Cache multiple paths - uses: actions/cache@v4 - with: - path: | - ~/.cache/pip - $RUNNER_TOOL_CACHE/Python/* - ~\AppData\Local\pip\Cache - key: ${{ runner.os }}-build-${{ matrix.python-version }} + - name: Install uv (Unix) + run: curl -LsSf https://astral.sh/uv/install.sh | sh - name: clone-scikit-lego run: git clone https://github.com/koaning/scikit-lego.git --depth 1 - name: install-basics - run: python -m pip install --upgrade tox virtualenv setuptools pip + run: uv pip install --upgrade tox virtualenv setuptools --system - name: install-scikit-lego-dev run: | cd scikit-lego - pip install -e ".[test]" + uv pip install -e ".[test]" --system - name: install-narwhals-dev run: | - pip uninstall narwhals -y - pip install -e . + uv pip uninstall narwhals --system + uv pip install -e . --system - name: show-deps - run: pip freeze + run: uv pip freeze - name: Run pytest run: | cd scikit-lego diff --git a/.github/workflows/extremes.yml b/.github/workflows/extremes.yml index e689de046..ae9c79009 100644 --- a/.github/workflows/extremes.yml +++ b/.github/workflows/extremes.yml @@ -1,4 +1,4 @@ -name: ci +name: Min, old, and nightly versions on: pull_request: @@ -18,20 +18,14 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Cache multiple paths - uses: actions/cache@v4 - with: - path: | - ~/.cache/pip - $RUNNER_TOOL_CACHE/Python/* - ~\AppData\Local\pip\Cache - key: ${{ runner.os }}-build-${{ matrix.python-version }} - - name: install-minimu-versions - run: python -m pip install tox virtualenv setuptools pandas==0.25.3 polars==0.20.3 numpy==1.17.5 pyarrow==11.0.0 scipy==1.5.0 scikit-learn==1.1.0 tzdata + - name: Install uv (Unix) + run: curl -LsSf https://astral.sh/uv/install.sh | sh + - name: install-minimum-versions + run: uv pip install tox virtualenv setuptools pandas==0.25.3 polars==0.20.3 numpy==1.17.5 pyarrow==11.0.0 scipy==1.5.0 scikit-learn==1.1.0 tzdata --system - name: install-reqs - run: python -m pip install -r requirements-dev.txt + run: uv pip install -r requirements-dev.txt --system - name: show-deps - run: pip freeze + run: uv pip freeze - name: Run pytest run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=50 --runslow @@ -47,20 +41,39 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Cache multiple paths - uses: actions/cache@v4 + - name: Install uv (Unix) + run: curl -LsSf https://astral.sh/uv/install.sh | sh + - name: install-minimum-versions + run: uv pip install tox virtualenv setuptools pandas==1.1.5 polars==0.20.3 numpy==1.17.5 pyarrow==11.0.0 scipy==1.5.0 scikit-learn==1.1.0 tzdata --system + - name: install-reqs + run: uv pip install -r requirements-dev.txt --system + - name: show-deps + run: uv pip freeze + - name: Run pytest + run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=50 --runslow + - name: Run doctests + run: pytest narwhals --doctest-modules + + not_so_old_versions: + strategy: + matrix: + python-version: ["3.9"] + os: [ubuntu-latest] + + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 with: - path: | - ~/.cache/pip - $RUNNER_TOOL_CACHE/Python/* - ~\AppData\Local\pip\Cache - key: ${{ runner.os }}-build-${{ matrix.python-version }} - - name: install-minimu-versions - run: python -m pip install tox virtualenv setuptools pandas==1.1.5 polars==0.20.3 numpy==1.17.5 pyarrow==11.0.0 scipy==1.5.0 scikit-learn==1.1.0 tzdata + python-version: ${{ matrix.python-version }} + - name: Install uv (Unix) + run: curl -LsSf https://astral.sh/uv/install.sh | sh + - name: install-minimum-versions + run: uv pip install tox virtualenv setuptools pandas==2.0.3 polars==0.20.8 numpy==1.24.4 pyarrow==14.0.0 scipy==1.8.0 scikit-learn==1.3.0 dask[dataframe]==2024.7 tzdata --system - name: install-reqs - run: python -m pip install -r requirements-dev.txt + run: uv pip install -r requirements-dev.txt --system - name: show-deps - run: pip freeze + run: uv pip freeze - name: Run pytest run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=50 --runslow - name: Run doctests @@ -78,35 +91,29 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Cache multiple paths - uses: actions/cache@v4 - with: - path: | - ~/.cache/pip - $RUNNER_TOOL_CACHE/Python/* - ~\AppData\Local\pip\Cache - key: ${{ runner.os }}-build-${{ matrix.python-version }} + - name: Install uv (Unix) + run: curl -LsSf https://astral.sh/uv/install.sh | sh - name: install-polars - run: pip install polars + run: uv pip install polars --system - name: install-reqs - run: python -m pip install --upgrade tox virtualenv setuptools pip -r requirements-dev.txt + run: uv pip install --upgrade tox virtualenv setuptools pip -r requirements-dev.txt --system - name: uninstall pyarrow - run: python -m pip uninstall pyarrow -y - - name: install pyarrow nightly - run: python -m pip install --extra-index-url https://pypi.fury.io/arrow-nightlies/ --prefer-binary --pre pyarrow + run: uv pip uninstall pyarrow --system + # - name: install pyarrow nightly + # run: uv pip install --extra-index-url https://pypi.fury.io/arrow-nightlies/ --pre pyarrow --system - name: uninstall pandas - run: python -m pip uninstall pandas -y + run: uv pip uninstall pandas --system - name: install-pandas-nightly - run: python -m pip install --pre --extra-index https://pypi.anaconda.org/scientific-python-nightly-wheels/simple pandas + run: uv pip install --pre --extra-index-url https://pypi.anaconda.org/scientific-python-nightly-wheels/simple pandas --system - name: uninstall numpy - run: python -m pip uninstall numpy -y + run: uv pip uninstall numpy --system - name: install numpy nightly - run: python -m pip install --pre --extra-index-url https://pypi.anaconda.org/scientific-python-nightly-wheels/simple numpy + run: uv pip install --pre --extra-index-url https://pypi.anaconda.org/scientific-python-nightly-wheels/simple numpy --system - name: install dask run: | - pip install git+https://github.com/dask/distributed git+https://github.com/dask/dask git+https://github.com/dask/dask-expr + python -m pip install git+https://github.com/dask/distributed git+https://github.com/dask/dask git+https://github.com/dask/dask-expr - name: show-deps - run: pip freeze + run: uv pip freeze - name: Run pytest run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=50 --runslow - name: Run doctests diff --git a/.github/workflows/mkdocs.yml b/.github/workflows/mkdocs.yml index bded1b45d..37b188413 100644 --- a/.github/workflows/mkdocs.yml +++ b/.github/workflows/mkdocs.yml @@ -1,4 +1,4 @@ -name: mkdocs +name: Publish Docs on: push: diff --git a/.github/workflows/publish_to_pypi.yml b/.github/workflows/publish_to_pypi.yml index e2e2ed1e8..bc9003ce0 100644 --- a/.github/workflows/publish_to_pypi.yml +++ b/.github/workflows/publish_to_pypi.yml @@ -30,7 +30,7 @@ jobs: publish-to-pypi: name: >- Publish Python 🐍 distribution 📦 to PyPI - if: startsWith(github.ref, 'refs/tags/') # only publish to PyPI on tag pushes + if: startsWith(github.ref, 'refs/tags/v') # only publish to PyPI on tag pushes needs: - build runs-on: ubuntu-latest diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 5a213550a..0458a4729 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -1,4 +1,4 @@ -name: ci +name: PyTest on: pull_request: @@ -18,20 +18,18 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Cache multiple paths - uses: actions/cache@v4 - with: - path: | - ~/.cache/pip - $RUNNER_TOOL_CACHE/Python/* - ~\AppData\Local\pip\Cache - key: ${{ runner.os }}-build-${{ matrix.python-version }} + - name: Install uv (Unix) + if: runner.os != 'Windows' + run: curl -LsSf https://astral.sh/uv/install.sh | sh + - name: Install uv (Windows) + if: runner.os == 'Windows' + run: powershell -c "irm https://astral.sh/uv/install.ps1 | iex" - name: install-reqs - run: python -m pip install --upgrade tox virtualenv setuptools pip -r requirements-dev.txt + run: uv pip install --upgrade tox virtualenv setuptools -r requirements-dev.txt --system - name: show-deps - run: pip freeze + run: uv pip freeze - name: Run pytest - run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=90 + run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=85 - name: Run doctests run: pytest narwhals --doctest-modules @@ -47,20 +45,14 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Cache multiple paths - uses: actions/cache@v4 - with: - path: | - ~/.cache/pip - $RUNNER_TOOL_CACHE/Python/* - ~\AppData\Local\pip\Cache - key: ${{ runner.os }}-build-${{ matrix.python-version }} + - name: Install uv (Windows) + run: powershell -c "irm https://astral.sh/uv/install.ps1 | iex" - name: install-reqs - run: python -m pip install --upgrade tox virtualenv setuptools pip -r requirements-dev.txt + run: uv pip install --upgrade tox virtualenv setuptools -r requirements-dev.txt --system - name: install-modin - run: python -m pip install --upgrade modin[dask] + run: uv pip install --upgrade modin[dask] --system - name: show-deps - run: pip freeze + run: uv pip freeze - name: Run pytest run: pytest tests --cov=narwhals --cov=tests --runslow --cov-fail-under=95 - name: Run doctests @@ -78,20 +70,14 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Cache multiple paths - uses: actions/cache@v4 - with: - path: | - ~/.cache/pip - $RUNNER_TOOL_CACHE/Python/* - ~\AppData\Local\pip\Cache - key: ${{ runner.os }}-build-${{ matrix.python-version }} + - name: Install uv (Unix) + run: curl -LsSf https://astral.sh/uv/install.sh | sh - name: install-reqs - run: python -m pip install --upgrade tox virtualenv setuptools pip -r requirements-dev.txt + run: uv pip install --upgrade tox virtualenv setuptools -r requirements-dev.txt --system - name: install-modin - run: python -m pip install --upgrade modin[dask] + run: uv pip install --upgrade modin[dask] --system - name: show-deps - run: pip freeze + run: uv pip freeze - name: Run pytest run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=100 --runslow - name: Run doctests diff --git a/.github/workflows/random_ci_pytest.yml b/.github/workflows/random_ci_pytest.yml index 184355dfa..e25bdcb68 100644 --- a/.github/workflows/random_ci_pytest.yml +++ b/.github/workflows/random_ci_pytest.yml @@ -1,4 +1,4 @@ -name: random-versions +name: Random Versions on: pull_request: @@ -16,17 +16,19 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} + - name: Install uv (Unix) + run: curl -LsSf https://astral.sh/uv/install.sh | sh - name: install package - run: pip install -e . + run: uv pip install -e . --system - name: generate-random-versions run: python utils/generate_random_versions.py - name: install-reqs - run: python -m pip install --upgrade tox virtualenv setuptools pip && python -m pip install -r requirements-dev.txt + run: uv pip install --upgrade tox virtualenv setuptools --system && uv pip install -r requirements-dev.txt --system - name: uninstall scipy/sklearn - run: python -m pip uninstall -y scipy scikit-learn + run: uv pip uninstall scipy scikit-learn --system - name: install-random-verions - run: python -m pip install -r random-requirements.txt + run: uv pip install -r random-requirements.txt --system - name: show versions - run: python -m pip freeze + run: uv pip freeze - name: Run pytest run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=80 diff --git a/.gitignore b/.gitignore index 3911158a8..d6edf57e6 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,7 @@ todo.md site/ .coverage.* .nox -docs/api-completeness.md \ No newline at end of file +*.lock + +docs/api-completeness/*.md +!docs/api-completeness/index.md \ No newline at end of file diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2a02b6353..4914e7e16 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,7 +1,7 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: 'v0.5.1' + rev: 'v0.5.7' hooks: # Run the formatter. - id: ruff-format @@ -9,10 +9,10 @@ repos: - id: ruff args: [--fix] - repo: https://github.com/pre-commit/mirrors-mypy - rev: 'v1.10.1' + rev: 'v1.11.1' hooks: - id: mypy - additional_dependencies: ['polars==1.1.0', 'pytest==8.0.1'] + additional_dependencies: ['polars==1.4.1', 'pytest==8.3.2'] exclude: utils|tpch - repo: https://github.com/codespell-project/codespell rev: 'v2.3.0' @@ -31,11 +31,18 @@ repos: additional_dependencies: [polars] - id: imports-are-banned name: import are banned (use `get_pandas` instead of `import pandas`) - entry: (?>> )import (pandas|polars|modin|cudf|pyarrow) - language: pygrep + entry: python utils/import_check.py + language: python files: ^narwhals/ exclude: ^narwhals/dependencies\.py - repo: https://github.com/kynan/nbstripout rev: 0.7.1 hooks: - id: nbstripout +- repo: https://github.com/adamchainz/blacken-docs + rev: "1.18.0" # replace with latest tag on GitHub + hooks: + - id: blacken-docs + args: [--skip-errors] + additional_dependencies: + - black==22.12.0 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ae34067ec..d36d21a55 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -2,7 +2,50 @@ Thank you for your interest in contributing to Narwhals! Any kind of improvement is welcome! -## Setting up your environment +## Local development vs Codespaces +You can contribute to Narwhals in your local development environment, using python3, git and your editor of choice. +You can also contribute to Narwhals using [Github Codespaces](https://docs.github.com/en/codespaces/overview) - a development environment that's hosted in the cloud. +This way you can easily start to work from your browser without installing git and cloning the repo. +Scroll down for instructions on how to use [Codespaces](#working-with-codespaces). + +## Working with local development environment + +### 1. Make sure you have git on your machine and a GitHub account + +Open your terminal and run the following command: + +```bash +git --version +``` + +If the output looks like `git version 2.34.1` and you have a personal account on GitHub - you're good to go to the next step. +If the terminal output informs about `command not found` you need to [install git](https://docs.github.com/en/get-started/quickstart/set-up-git). + +If you're new to GitHub, you'll need to create an account on [GitHub.com](https://github.com/) and verify your email address. + +### 2. Fork the repository + +Go to the [main project page](https://github.com/narwhals-dev/narwhals). +Fork the repository by clicking on the fork button. You can find it in the right corner on the top of the page. + +### 3. Clone the repository + +Go to the forked repository on your GitHub account - you'll find it on your account in the tab Repositories. +Click on the green `Code` button and then click the `Copy url to clipboard` icon. +Open a terminal, choose the directory where you would like to have Narwhals repository and run the following git command: + +```bash +git clone +``` + +for example: + +```bash +git clone git@github.com:YOUR-USERNAME/narwhals.git +``` + + +### 4. Setting up your environment Here's how you can set up your local development environment to contribute: @@ -20,13 +63,13 @@ pre-commit install ``` This will automatically format and lint your code before each commit, and it will block the commit if any issues are found. -## Working on your issue +### 5. Working on your issue Create a new git branch from the `main` branch in your local repository. Note that your work cannot be merged if the test below fail. If you add code that should be tested, please add tests. -## Running tests +### 6. Running tests To run tests, run `pytest`. To check coverage: `pytest --cov=narwhals`. To run tests on the docset-module, use `pytest narwhals --doctest-modules`. @@ -42,27 +85,77 @@ nox Notice that nox will also require to have all the python versions that are defined in the `noxfile.py` installed in your system. -## Building docs +### 7. Building docs To build the docs, run `mkdocs serve`, and then open the link provided in a browser. The docs should refresh when you make changes. If they don't, press `ctrl+C`, and then do `mkdocs build` and then `mkdocs serve`. -## Pull requests +### 8. Pull requests When you have resolved your issue, [open a pull request](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request-from-a-fork) in the Narwhals repository. Please adhere to the following guidelines: -1. Start your pull request title with a [conventional commit](https://www.conventionalcommits.org/) tag. This helps us add your contribution to the right section of the changelog. We use "Type" from the [Angular convention](https://github.com/angular/angular/blob/22b96b9/CONTRIBUTING.md#type). -2. Use a descriptive title starting with an uppercase letter. This text will end up in the [changelog](https://github.com/narwhals-dev/narwhals/releases). +1. Start your pull request title with a [conventional commit](https://www.conventionalcommits.org/) tag. This helps us add your contribution to the right section of the changelog. We use "Type" from the [Angular convention](https://github.com/angular/angular/blob/22b96b9/CONTRIBUTING.md#type).
+ TLDR: + The PR title should start with any of these abbreviations: `build`, `chore`, `ci`, `depr`, + `docs`, `feat`, `fix`, `perf`, `refactor`, `release`, `test`. Add a `!`at the end, if it is a breaking change. For example `refactor!`. +
+2. This text will end up in the [changelog](https://github.com/narwhals-dev/narwhals/releases). 3. Please follow the instructions in the pull request form and submit. +## Working with Codespaces +Codespaces is a great way to work on Narwhals without the need of configuring your local development environment. +Every GitHub.com user has a monthly quota of free use of GitHub Codespaces, and you can start working in a codespace without providing any payment details. +You'll be informed per email if you'll be close to using 100% of included services. +To learn more about it visit [GitHub Docs](https://docs.github.com/en/codespaces/overview) + + +### 1. Make sure you have GitHub account + +If you're new to GitHub, you'll need to create an account on [GitHub.com](https://github.com/) and verify your email address. + +### 2. Fork the repository + +Go to the [main project page](https://github.com/narwhals-dev/narwhals). +Fork the repository by clicking on the fork button. You can find it in the right corner on the top of the page. + +### 3. Create codespace + +Go to the forked repository on your GitHub account - you'll find it on your account in the tab Repositories. +Click on the green `Code` button and navigate to the `Codespaces` tab. +Click on the green button `Create codespace on main` - it will open a browser version of VSCode, +with the complete repository and git installed. +You can now proceed with the steps [4. Setting up your environment](#4-setting-up-your-environment) up to [8. Pull request](#8-pull-requests) +listed above in [Working with local development environment](#working-with-local-development-environment). + + ## How it works If Narwhals looks like underwater unicorn magic to you, then please read [how it works](https://narwhals-dev.github.io/narwhals/how-it-works/). +## Imports + +In Narwhals, we are very particular about imports. When it comes to importing +heavy third-party libraries (pandas, NumPy, Polars, etc...) please follow these rules: + +- Never import anything to do `isinstance` checks. Instead, just use the functions + in `narwhals.dependencies` (such as `is_pandas_dataframe`); +- If you need to import anything, do it in a place where you know that the import + is definitely available. For example, NumPy is a required dependency of PyArrow, + so it's OK to import NumPy to implement a PyArrow function - however, NumPy + should never be imported to implement a Polars function. The only exception is + for when there's simply no way around it by definition - for example, `Series.to_numpy` + always requires NumPy to be installed. +- Don't place a third-party import at the top of a file. Instead, place it in the + function where it's used, so that we minimise the chances of it being imported + unnecessarily. + +We're trying to be really lightweight and minimal-overhead, and +unnecessary imports can slow things down. + ## Happy contributing! Please remember to abide by the code of conduct, else you'll be conducted away from this project. diff --git a/README.md b/README.md index a9954402f..d26107e67 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ [![PyPI version](https://badge.fury.io/py/narwhals.svg)](https://badge.fury.io/py/narwhals) +[![Downloads](https://static.pepy.tech/badge/narwhals/month)](https://pepy.tech/project/narwhals) Extremely lightweight and extensible compatibility layer between dataframe libraries! @@ -39,6 +40,7 @@ Get started! Join the party! +- [Altair](https://github.com/vega/altair/) - [Hamilton](https://github.com/DAGWorks-Inc/hamilton/tree/main/examples/narwhals) - [scikit-lego](https://github.com/koaning/scikit-lego) - [scikit-playtime](https://github.com/koaning/scikit-playtime) @@ -74,6 +76,14 @@ There are three steps to writing dataframe-agnostic code using Narwhals: - if you started with cuDF, you'll get cuDF back (and compute will happen on GPU) - if you started with PyArrow, you'll get PyArrow back +

+ narwhals_gif + +

+ ## Example See the [tutorial](https://narwhals-dev.github.io/narwhals/basics/dataframe/) for several examples! diff --git a/docs/api-completeness.md b/docs/api-completeness.md new file mode 100644 index 000000000..f24ccc4b9 --- /dev/null +++ b/docs/api-completeness.md @@ -0,0 +1,221 @@ +# API Completeness + +Narwhals has two different level of support for libraries: "full" and "interchange". + +Libraries for which we have full support we intend to support the whole Narwhals API, however this is a work in progress. + +In the following table it is possible to check which method is implemented for which backend. + +!!! info + + - "pandas-like" means pandas, cuDF and Modin + - Polars supports all the methods (by design) + +| Class | Method | pandas-like | arrow | +|-------------------------|--------------------|--------------------|--------------------| +| DataFrame | clone | :white_check_mark: | :white_check_mark: | +| DataFrame | collect_schema | :white_check_mark: | :white_check_mark: | +| DataFrame | columns | :white_check_mark: | :white_check_mark: | +| DataFrame | drop | :white_check_mark: | :white_check_mark: | +| DataFrame | drop_nulls | :white_check_mark: | :white_check_mark: | +| DataFrame | filter | :white_check_mark: | :white_check_mark: | +| DataFrame | gather_every | :white_check_mark: | :white_check_mark: | +| DataFrame | get_column | :white_check_mark: | :white_check_mark: | +| DataFrame | group_by | :white_check_mark: | :white_check_mark: | +| DataFrame | head | :white_check_mark: | :white_check_mark: | +| DataFrame | is_duplicated | :white_check_mark: | :white_check_mark: | +| DataFrame | is_empty | :white_check_mark: | :white_check_mark: | +| DataFrame | is_unique | :white_check_mark: | :white_check_mark: | +| DataFrame | item | :white_check_mark: | :white_check_mark: | +| DataFrame | iter_rows | :white_check_mark: | :white_check_mark: | +| DataFrame | join | :white_check_mark: | :white_check_mark: | +| DataFrame | lazy | :white_check_mark: | :white_check_mark: | +| DataFrame | null_count | :white_check_mark: | :white_check_mark: | +| DataFrame | pipe | :x: | :x: | +| DataFrame | rename | :white_check_mark: | :white_check_mark: | +| DataFrame | rows | :white_check_mark: | :white_check_mark: | +| DataFrame | schema | :white_check_mark: | :white_check_mark: | +| DataFrame | select | :white_check_mark: | :white_check_mark: | +| DataFrame | shape | :white_check_mark: | :white_check_mark: | +| DataFrame | sort | :white_check_mark: | :white_check_mark: | +| DataFrame | tail | :white_check_mark: | :white_check_mark: | +| DataFrame | to_dict | :white_check_mark: | :white_check_mark: | +| DataFrame | to_numpy | :white_check_mark: | :white_check_mark: | +| DataFrame | to_pandas | :white_check_mark: | :white_check_mark: | +| DataFrame | unique | :white_check_mark: | :white_check_mark: | +| DataFrame | with_columns | :white_check_mark: | :white_check_mark: | +| DataFrame | with_row_index | :white_check_mark: | :white_check_mark: | +| DataFrame | write_parquet | :white_check_mark: | :white_check_mark: | +| Expr | abs | :white_check_mark: | :white_check_mark: | +| Expr | alias | :white_check_mark: | :white_check_mark: | +| Expr | all | :white_check_mark: | :white_check_mark: | +| Expr | any | :white_check_mark: | :white_check_mark: | +| Expr | arg_true | :white_check_mark: | :white_check_mark: | +| Expr | cast | :white_check_mark: | :white_check_mark: | +| Expr | cat | :white_check_mark: | :white_check_mark: | +| Expr | clip | :white_check_mark: | :x: | +| Expr | count | :white_check_mark: | :white_check_mark: | +| Expr | cum_sum | :white_check_mark: | :white_check_mark: | +| Expr | diff | :white_check_mark: | :white_check_mark: | +| Expr | drop_nulls | :white_check_mark: | :white_check_mark: | +| Expr | dt | :white_check_mark: | :white_check_mark: | +| Expr | fill_null | :white_check_mark: | :white_check_mark: | +| Expr | filter | :white_check_mark: | :white_check_mark: | +| Expr | gather_every | :white_check_mark: | :white_check_mark: | +| Expr | head | :white_check_mark: | :white_check_mark: | +| Expr | is_between | :white_check_mark: | :white_check_mark: | +| Expr | is_duplicated | :white_check_mark: | :white_check_mark: | +| Expr | is_first_distinct | :white_check_mark: | :white_check_mark: | +| Expr | is_in | :white_check_mark: | :white_check_mark: | +| Expr | is_last_distinct | :white_check_mark: | :white_check_mark: | +| Expr | is_null | :white_check_mark: | :white_check_mark: | +| Expr | is_unique | :white_check_mark: | :white_check_mark: | +| Expr | len | :white_check_mark: | :white_check_mark: | +| Expr | max | :white_check_mark: | :white_check_mark: | +| Expr | mean | :white_check_mark: | :white_check_mark: | +| Expr | min | :white_check_mark: | :white_check_mark: | +| Expr | n_unique | :white_check_mark: | :white_check_mark: | +| Expr | name | :white_check_mark: | :white_check_mark: | +| Expr | null_count | :white_check_mark: | :white_check_mark: | +| Expr | over | :white_check_mark: | :x: | +| Expr | quantile | :white_check_mark: | :white_check_mark: | +| Expr | round | :white_check_mark: | :x: | +| Expr | sample | :white_check_mark: | :white_check_mark: | +| Expr | shift | :white_check_mark: | :x: | +| Expr | sort | :white_check_mark: | :white_check_mark: | +| Expr | std | :white_check_mark: | :white_check_mark: | +| Expr | str | :white_check_mark: | :white_check_mark: | +| Expr | sum | :white_check_mark: | :white_check_mark: | +| Expr | tail | :white_check_mark: | :white_check_mark: | +| Expr | unique | :white_check_mark: | :white_check_mark: | +| ExprCatNamespace | get_categories | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | day | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | hour | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | microsecond | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | millisecond | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | minute | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | month | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | nanosecond | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | ordinal_day | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | second | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | to_string | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | total_microseconds | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | total_milliseconds | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | total_minutes | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | total_nanoseconds | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | total_seconds | :white_check_mark: | :white_check_mark: | +| ExprDateTimeNamespace | year | :white_check_mark: | :white_check_mark: | +| ExprNameNamespace | keep | :white_check_mark: | :white_check_mark: | +| ExprNameNamespace | map | :white_check_mark: | :white_check_mark: | +| ExprNameNamespace | prefix | :white_check_mark: | :white_check_mark: | +| ExprNameNamespace | suffix | :white_check_mark: | :white_check_mark: | +| ExprNameNamespace | to_lowercase | :white_check_mark: | :white_check_mark: | +| ExprNameNamespace | to_uppercase | :white_check_mark: | :white_check_mark: | +| ExprStringNamespace | contains | :white_check_mark: | :white_check_mark: | +| ExprStringNamespace | ends_with | :white_check_mark: | :white_check_mark: | +| ExprStringNamespace | head | :x: | :x: | +| ExprStringNamespace | slice | :white_check_mark: | :white_check_mark: | +| ExprStringNamespace | starts_with | :white_check_mark: | :white_check_mark: | +| ExprStringNamespace | tail | :x: | :x: | +| ExprStringNamespace | to_datetime | :white_check_mark: | :white_check_mark: | +| ExprStringNamespace | to_lowercase | :white_check_mark: | :white_check_mark: | +| ExprStringNamespace | to_uppercase | :white_check_mark: | :white_check_mark: | +| LazyFrame | clone | :white_check_mark: | :white_check_mark: | +| LazyFrame | collect | :white_check_mark: | :white_check_mark: | +| LazyFrame | collect_schema | :white_check_mark: | :white_check_mark: | +| LazyFrame | columns | :white_check_mark: | :white_check_mark: | +| LazyFrame | drop | :white_check_mark: | :white_check_mark: | +| LazyFrame | drop_nulls | :white_check_mark: | :white_check_mark: | +| LazyFrame | filter | :white_check_mark: | :white_check_mark: | +| LazyFrame | gather_every | :white_check_mark: | :white_check_mark: | +| LazyFrame | group_by | :white_check_mark: | :white_check_mark: | +| LazyFrame | head | :white_check_mark: | :white_check_mark: | +| LazyFrame | join | :white_check_mark: | :white_check_mark: | +| LazyFrame | lazy | :white_check_mark: | :white_check_mark: | +| LazyFrame | pipe | :x: | :x: | +| LazyFrame | rename | :white_check_mark: | :white_check_mark: | +| LazyFrame | schema | :white_check_mark: | :white_check_mark: | +| LazyFrame | select | :white_check_mark: | :white_check_mark: | +| LazyFrame | sort | :white_check_mark: | :white_check_mark: | +| LazyFrame | tail | :white_check_mark: | :white_check_mark: | +| LazyFrame | unique | :white_check_mark: | :white_check_mark: | +| LazyFrame | with_columns | :white_check_mark: | :white_check_mark: | +| LazyFrame | with_row_index | :white_check_mark: | :white_check_mark: | +| Series | abs | :white_check_mark: | :white_check_mark: | +| Series | alias | :white_check_mark: | :white_check_mark: | +| Series | all | :white_check_mark: | :white_check_mark: | +| Series | any | :white_check_mark: | :white_check_mark: | +| Series | arg_true | :white_check_mark: | :white_check_mark: | +| Series | cast | :white_check_mark: | :white_check_mark: | +| Series | cat | :white_check_mark: | :white_check_mark: | +| Series | clip | :white_check_mark: | :x: | +| Series | count | :white_check_mark: | :white_check_mark: | +| Series | cum_sum | :white_check_mark: | :white_check_mark: | +| Series | diff | :white_check_mark: | :white_check_mark: | +| Series | drop_nulls | :white_check_mark: | :white_check_mark: | +| Series | dt | :white_check_mark: | :white_check_mark: | +| Series | dtype | :white_check_mark: | :white_check_mark: | +| Series | fill_null | :white_check_mark: | :white_check_mark: | +| Series | filter | :white_check_mark: | :white_check_mark: | +| Series | gather_every | :white_check_mark: | :white_check_mark: | +| Series | head | :white_check_mark: | :white_check_mark: | +| Series | is_between | :white_check_mark: | :white_check_mark: | +| Series | is_duplicated | :white_check_mark: | :white_check_mark: | +| Series | is_empty | :white_check_mark: | :white_check_mark: | +| Series | is_first_distinct | :white_check_mark: | :white_check_mark: | +| Series | is_in | :white_check_mark: | :white_check_mark: | +| Series | is_last_distinct | :white_check_mark: | :white_check_mark: | +| Series | is_null | :white_check_mark: | :white_check_mark: | +| Series | is_sorted | :white_check_mark: | :white_check_mark: | +| Series | is_unique | :white_check_mark: | :white_check_mark: | +| Series | item | :white_check_mark: | :white_check_mark: | +| Series | len | :white_check_mark: | :white_check_mark: | +| Series | max | :white_check_mark: | :white_check_mark: | +| Series | mean | :white_check_mark: | :white_check_mark: | +| Series | min | :white_check_mark: | :white_check_mark: | +| Series | n_unique | :white_check_mark: | :white_check_mark: | +| Series | name | :white_check_mark: | :white_check_mark: | +| Series | null_count | :white_check_mark: | :white_check_mark: | +| Series | quantile | :white_check_mark: | :white_check_mark: | +| Series | round | :white_check_mark: | :x: | +| Series | sample | :white_check_mark: | :white_check_mark: | +| Series | shape | :white_check_mark: | :white_check_mark: | +| Series | shift | :white_check_mark: | :x: | +| Series | sort | :white_check_mark: | :white_check_mark: | +| Series | std | :white_check_mark: | :white_check_mark: | +| Series | str | :white_check_mark: | :white_check_mark: | +| Series | sum | :white_check_mark: | :white_check_mark: | +| Series | tail | :white_check_mark: | :white_check_mark: | +| Series | to_dummies | :white_check_mark: | :white_check_mark: | +| Series | to_frame | :white_check_mark: | :white_check_mark: | +| Series | to_list | :white_check_mark: | :white_check_mark: | +| Series | to_numpy | :white_check_mark: | :white_check_mark: | +| Series | to_pandas | :white_check_mark: | :white_check_mark: | +| Series | unique | :white_check_mark: | :white_check_mark: | +| Series | value_counts | :white_check_mark: | :white_check_mark: | +| Series | zip_with | :white_check_mark: | :white_check_mark: | +| SeriesCatNamespace | get_categories | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | day | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | hour | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | microsecond | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | millisecond | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | minute | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | month | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | nanosecond | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | ordinal_day | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | second | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | to_string | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | total_microseconds | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | total_milliseconds | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | total_minutes | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | total_nanoseconds | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | total_seconds | :white_check_mark: | :white_check_mark: | +| SeriesDateTimeNamespace | year | :white_check_mark: | :white_check_mark: | +| SeriesStringNamespace | contains | :white_check_mark: | :white_check_mark: | +| SeriesStringNamespace | ends_with | :white_check_mark: | :white_check_mark: | +| SeriesStringNamespace | head | :x: | :x: | +| SeriesStringNamespace | slice | :white_check_mark: | :white_check_mark: | +| SeriesStringNamespace | starts_with | :white_check_mark: | :white_check_mark: | +| SeriesStringNamespace | tail | :x: | :x: | +| SeriesStringNamespace | to_lowercase | :white_check_mark: | :white_check_mark: | +| SeriesStringNamespace | to_uppercase | :white_check_mark: | :white_check_mark: | \ No newline at end of file diff --git a/docs/api-completeness/index.md b/docs/api-completeness/index.md new file mode 100644 index 000000000..3daba31ca --- /dev/null +++ b/docs/api-completeness/index.md @@ -0,0 +1,14 @@ +# API Completeness + +Narwhals has two different level of support for libraries: "full" and "interchange". + +Libraries for which we have full support we intend to support the whole Narwhals API, +however this is a continuous work in progress. + +In the following section it is possible to check which method is implemented for which +class and backend. + +!!! info + + - By design, Polars supports all the methods of the Narwhals API. + - "pandas-like" means pandas, cuDF and Modin. diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md index 7231b01d1..c144b4af0 100644 --- a/docs/api-reference/dataframe.md +++ b/docs/api-reference/dataframe.md @@ -4,6 +4,7 @@ handler: python options: members: + - __arrow_c_stream__ - __getitem__ - clone - collect_schema @@ -25,18 +26,21 @@ - null_count - pipe - rename + - row - rows - schema - select - shape - sort - tail + - to_arrow - to_dict - to_numpy - to_pandas - unique - with_columns - with_row_index + - write_csv - write_parquet show_source: false show_bases: false diff --git a/docs/api-reference/dependencies.md b/docs/api-reference/dependencies.md index 72541875f..6c1a93d91 100644 --- a/docs/api-reference/dependencies.md +++ b/docs/api-reference/dependencies.md @@ -4,11 +4,25 @@ handler: python options: members: + - get_cudf + - get_modin - get_pandas - get_polars - - get_modin - - get_cudf - get_pyarrow + - is_cudf_dataframe + - is_cudf_series + - is_dask_dataframe + - is_modin_dataframe + - is_modin_series + - is_numpy_array - is_pandas_dataframe + - is_pandas_like_dataframe + - is_pandas_like_series + - is_pandas_series + - is_polars_dataframe + - is_polars_lazyframe + - is_polars_series + - is_pyarrow_chunked_array + - is_pyarrow_table show_source: false show_bases: false diff --git a/docs/api-reference/expr.md b/docs/api-reference/expr.md index 492066aa0..cc1290a85 100644 --- a/docs/api-reference/expr.md +++ b/docs/api-reference/expr.md @@ -18,6 +18,7 @@ - filter - gather_every - head + - clip - is_between - is_duplicated - is_first_distinct @@ -32,6 +33,7 @@ - null_count - n_unique - over + - pipe - quantile - round - sample diff --git a/docs/api-reference/expr_dt.md b/docs/api-reference/expr_dt.md index af929486b..e04e92889 100644 --- a/docs/api-reference/expr_dt.md +++ b/docs/api-reference/expr_dt.md @@ -4,6 +4,7 @@ handler: python options: members: + - date - year - month - day diff --git a/docs/api-reference/expr_str.md b/docs/api-reference/expr_str.md index f64b360f4..8cb0dd9ed 100644 --- a/docs/api-reference/expr_str.md +++ b/docs/api-reference/expr_str.md @@ -8,7 +8,10 @@ - ends_with - head - slice + - replace + - replace_all - starts_with + - strip_chars - tail - to_datetime - to_lowercase diff --git a/docs/api-reference/index.md b/docs/api-reference/index.md index 97ed4ee2b..0c2c81fa3 100644 --- a/docs/api-reference/index.md +++ b/docs/api-reference/index.md @@ -8,8 +8,8 @@ For example: import narwhals as nw df.with_columns( - a_mean = nw.col('a').mean(), - a_std = nw.col('a').std(), + a_mean=nw.col("a").mean(), + a_std=nw.col("a").std(), ) ``` is supported, as `DataFrame.with_columns`, `narwhals.col`, `Expr.mean`, and `Expr.std` are @@ -20,7 +20,7 @@ However, import narwhals as nw df.with_columns( - a_ewm_mean = nw.col('a').ewm_mean(alpha=.7), + a_ewm_mean=nw.col("a").ewm_mean(alpha=0.7), ) ``` is not - `Expr.ewm_mean` only appears in the Polars API reference, but not in the Narwhals diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md index 42c1a2e44..9d7ac384e 100644 --- a/docs/api-reference/narwhals.md +++ b/docs/api-reference/narwhals.md @@ -17,14 +17,17 @@ Here are the top-level functions available in Narwhals. - get_native_namespace - is_ordered_categorical - len - - maybe_align_index - - maybe_set_index - - maybe_convert_dtypes - lit - max + - maybe_align_index + - maybe_convert_dtypes + - maybe_get_index + - maybe_set_index - mean + - mean_horizontal - min - narwhalify + - new_series - sum - sum_horizontal - when diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md index 89ed25c63..f9cc2e6bb 100644 --- a/docs/api-reference/series.md +++ b/docs/api-reference/series.md @@ -4,6 +4,8 @@ handler: python options: members: + - __arrow_c_stream__ + - __getitem__ - abs - alias - all @@ -20,6 +22,7 @@ - gather_every - head - is_between + - clip - is_duplicated - is_empty - is_first_distinct @@ -36,6 +39,7 @@ - name - null_count - n_unique + - pipe - quantile - round - sample @@ -45,6 +49,7 @@ - std - sum - tail + - to_arrow - to_dummies - to_frame - to_list diff --git a/docs/api-reference/series_cat.md b/docs/api-reference/series_cat.md index 55dc6e634..de44a36b1 100644 --- a/docs/api-reference/series_cat.md +++ b/docs/api-reference/series_cat.md @@ -1,6 +1,6 @@ # `narwhals.Series.cat` -::: narwhals.series.SeriesStringNamespace +::: narwhals.series.SeriesCatNamespace handler: python options: members: diff --git a/docs/api-reference/series_dt.md b/docs/api-reference/series_dt.md index ad33a083b..ba342ad30 100644 --- a/docs/api-reference/series_dt.md +++ b/docs/api-reference/series_dt.md @@ -4,6 +4,7 @@ handler: python options: members: + - date - year - month - day diff --git a/docs/api-reference/series_str.md b/docs/api-reference/series_str.md index 3405b8a43..af657deff 100644 --- a/docs/api-reference/series_str.md +++ b/docs/api-reference/series_str.md @@ -7,8 +7,11 @@ - contains - ends_with - head + - replace + - replace_all - slice - starts_with + - strip_chars - tail show_source: false show_bases: false diff --git a/docs/api-reference/typing.md b/docs/api-reference/typing.md index 8dcdfa725..b3861b4f9 100644 --- a/docs/api-reference/typing.md +++ b/docs/api-reference/typing.md @@ -11,9 +11,10 @@ accepts a `nw.DataFrame` and returns a `nw.DataFrame` backed by the same backend import narwhals as nw from narwhals.typing import DataFrameT + @nw.narwhalify def func(df: DataFrameT) -> DataFrameT: - return df.with_columns(c=df['a']+1) + return df.with_columns(c=df["a"] + 1) ``` ## `Frame` @@ -25,6 +26,7 @@ either and your function doesn't care about its backend, for example: import narwhals as nw from narwhals.typing import Frame + @nw.narwhalify def func(df: Frame) -> list[str]: return df.columns @@ -38,9 +40,10 @@ or `nw.LazyFrame` and returns an object backed by the same backend, for example: import narwhals as nw from narwhals.typing import FrameT + @nw.narwhalify def func(df: FrameT) -> FrameT: - return df.with_columns(c=nw.col('a')+1) + return df.with_columns(c=nw.col("a") + 1) ``` ## `IntoDataFrame` @@ -53,6 +56,7 @@ Use this if your function accepts a narwhalifiable object but doesn't care about import narwhals as nw from narwhals.typing import IntoDataFrame + def func(df_native: IntoDataFrame) -> tuple[int, int]: df = nw.from_native(df_native, eager_only=True) return df.shape @@ -67,9 +71,10 @@ class: import narwhals as nw from narwhals.typing import IntoDataFrameT + def func(df_native: IntoDataFrameT) -> IntoDataFrameT: df = nw.from_native(df_native, eager_only=True) - return nw.to_native(df.with_columns(c=df['a']+1)) + return nw.to_native(df.with_columns(c=df["a"] + 1)) ``` ## `IntoExpr` @@ -88,6 +93,7 @@ care about its backend: import narwhals as nw from narwhals.typing import IntoFrame + def func(df_native: IntoFrame) -> list[str]: df = nw.from_native(df_native) return df.columns @@ -102,9 +108,10 @@ of the same type: import narwhals as nw from narwhals.typing import IntoFrameT + def func(df_native: IntoFrameT) -> IntoFrameT: df = nw.from_native(df_native) - return nw.to_native(df.with_columns(c=nw.col('a')+1)) + return nw.to_native(df.with_columns(c=nw.col("a") + 1)) ``` ## `nw.narwhalify`, or `nw.from_native`? @@ -117,17 +124,21 @@ import polars as pl import narwhals as nw from narwhals.typing import IntoDataFrameT, DataFrameT -df = pl.DataFrame({'a': [1,2,3]}) +df = pl.DataFrame({"a": [1, 2, 3]}) + def func(df: IntoDataFrameT) -> IntoDataFrameT: df = nw.from_native(df, eager_only=True) - return nw.to_native(df.select(b=nw.col('a'))) + return nw.to_native(df.select(b=nw.col("a"))) + reveal_type(func(df)) + @nw.narwhalify(strict=True) def func_2(df: DataFrameT) -> DataFrameT: - return df.select(b=nw.col('a')) + return df.select(b=nw.col("a")) + reveal_type(func_2(df)) ``` diff --git a/docs/backcompat.md b/docs/backcompat.md index c8e50a8d7..3c4a19a7c 100644 --- a/docs/backcompat.md +++ b/docs/backcompat.md @@ -22,12 +22,12 @@ Ever upgraded a package, only to find that it breaks all your tests because of a API change? Did you end up having to litter your code with statements such as the following? ```python -if parse_version(pdx.__version__) < parse_version('1.3.0'): +if parse_version(pdx.__version__) < parse_version("1.3.0"): df = df.brewbeer() -elif parse_version('1.3.0') <= parse_version(pdx.__version__) < parse_version('1.5.0'): +elif parse_version("1.3.0") <= parse_version(pdx.__version__) < parse_version("1.5.0"): df = df.brew_beer() else: - df = df.brew_drink('beer') + df = df.brew_drink("beer") ``` Now imagine multiplying that complexity over all the dataframe libraries you want to support... @@ -51,9 +51,10 @@ it. That is to say, if you write your code like this: import narwhals.stable.v1 as nw from narwhals.typing import FrameT + @nw.narwhalify def func(df: FrameT) -> FrameT: - return df.with_columns(nw.col('a').cum_sum()) + return df.with_columns(nw.col("a").cum_sum()) ``` then we, in Narwhals, promise that your code will keep working, even in newer versions of Polars diff --git a/docs/basics/column.md b/docs/basics/column.md index f44217b26..62951cb4b 100644 --- a/docs/basics/column.md +++ b/docs/basics/column.md @@ -17,16 +17,17 @@ This can stay lazy, so we just use `nw.from_native` and expressions: import narwhals as nw from narwhals.typing import FrameT + @nw.narwhalify def my_func(df: FrameT) -> FrameT: - return df.filter(nw.col('a') > 0) + return df.filter(nw.col("a") > 0) ``` === "pandas" ```python exec="true" source="material-block" result="python" session="ex1" import pandas as pd - df = pd.DataFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pd.DataFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df)) ``` @@ -34,7 +35,7 @@ def my_func(df: FrameT) -> FrameT: ```python exec="true" source="material-block" result="python" session="ex1" import polars as pl - df = pl.DataFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pl.DataFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df)) ``` @@ -42,7 +43,7 @@ def my_func(df: FrameT) -> FrameT: ```python exec="true" source="material-block" result="python" session="ex1" import polars as pl - df = pl.LazyFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pl.LazyFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df).collect()) ``` @@ -55,16 +56,17 @@ Let's write a dataframe-agnostic function which multiplies the values in column import narwhals as nw from narwhals.typing import FrameT + @nw.narwhalify def my_func(df: FrameT) -> FrameT: - return df.with_columns(nw.col('a')*2) + return df.with_columns(nw.col("a") * 2) ``` === "pandas" ```python exec="true" source="material-block" result="python" session="ex2" import pandas as pd - df = pd.DataFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pd.DataFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df)) ``` @@ -72,7 +74,7 @@ def my_func(df: FrameT) -> FrameT: ```python exec="true" source="material-block" result="python" session="ex2" import polars as pl - df = pl.DataFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pl.DataFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df)) ``` @@ -80,7 +82,7 @@ def my_func(df: FrameT) -> FrameT: ```python exec="true" source="material-block" result="python" session="ex2" import polars as pl - df = pl.LazyFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pl.LazyFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df).collect()) ``` @@ -91,16 +93,17 @@ values multiplied by 2, we could have used `Expr.alias`: import narwhals as nw from narwhals.typing import FrameT + @nw.narwhalify def my_func(df: FrameT) -> FrameT: - return df.with_columns((nw.col('a')*2).alias('c')) + return df.with_columns((nw.col("a") * 2).alias("c")) ``` === "pandas" ```python exec="true" source="material-block" result="python" session="ex2.1" import pandas as pd - df = pd.DataFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pd.DataFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df)) ``` @@ -108,7 +111,7 @@ def my_func(df: FrameT) -> FrameT: ```python exec="true" source="material-block" result="python" session="ex2.1" import polars as pl - df = pl.DataFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pl.DataFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df)) ``` @@ -116,7 +119,7 @@ def my_func(df: FrameT) -> FrameT: ```python exec="true" source="material-block" result="python" session="ex2.1" import polars as pl - df = pl.LazyFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pl.LazyFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df).collect()) ``` @@ -131,16 +134,17 @@ of using expressions, we'll extract a `Series`. from __future__ import annotations import narwhals as nw + @nw.narwhalify(eager_only=True) def my_func(df: nw.DataFrame) -> float | None: - return df['a'].mean() + return df["a"].mean() ``` === "pandas" ```python exec="true" source="material-block" result="python" session="ex2" import pandas as pd - df = pd.DataFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pd.DataFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df)) ``` @@ -148,7 +152,7 @@ def my_func(df: nw.DataFrame) -> float | None: ```python exec="true" source="material-block" result="python" session="ex2" import polars as pl - df = pl.DataFrame({'a': [-1, 1, 3], 'b': [3, 5, -3]}) + df = pl.DataFrame({"a": [-1, 1, 3], "b": [3, 5, -3]}) print(my_func(df)) ``` diff --git a/docs/basics/complete_example.md b/docs/basics/complete_example.md index 343321bcb..1e2cfe30d 100644 --- a/docs/basics/complete_example.md +++ b/docs/basics/complete_example.md @@ -30,11 +30,13 @@ the argument will be propagated to `nw.from_native`. import narwhals as nw from typing import Any + class StandardScaler: @nw.narwhalify(eager_only=True) def fit(self, df: nw.DataFrame[Any]) -> None: self._means = {col: df[col].mean() for col in df.columns} self._std_devs = {col: df[col].std() for col in df.columns} + self._columns = df.columns ``` ## Transform method @@ -43,12 +45,11 @@ We're going to take in a dataframe, and return a dataframe of the same type. Therefore, we use `@nw.narwhalify`: ```python - @nw.narwhalify - def transform(self, df: FrameT) -> FrameT: - return df.with_columns( - (nw.col(col) - self._means[col]) / self._std_devs[col] - for col in df.columns - ) +@nw.narwhalify +def transform(self, df: FrameT) -> FrameT: + return df.with_columns( + (nw.col(col) - self._means[col]) / self._std_devs[col] for col in self._columns + ) ``` Note that all the calculations here can stay lazy if the underlying library permits it, @@ -64,17 +65,19 @@ from typing import Any import narwhals as nw from narwhals.typing import FrameT + class StandardScaler: @nw.narwhalify(eager_only=True) def fit(self, df: nw.DataFrame[Any]) -> None: self._means = {col: df[col].mean() for col in df.columns} self._std_devs = {col: df[col].std() for col in df.columns} + self._columns = df.columns @nw.narwhalify def transform(self, df: FrameT) -> FrameT: return df.with_columns( (nw.col(col) - self._means[col]) / self._std_devs[col] - for col in df.columns + for col in self._columns ) ``` @@ -86,8 +89,8 @@ stay lazy! ```python exec="true" source="material-block" result="python" session="tute-ex1" import pandas as pd - df_train = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 7]}) - df_test = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 7]}) + df_train = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 7]}) + df_test = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 7]}) scaler = StandardScaler() scaler.fit(df_train) print(scaler.transform(df_test)) @@ -97,8 +100,8 @@ stay lazy! ```python exec="true" source="material-block" result="python" session="tute-ex1" import polars as pl - df_train = pl.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 7]}) - df_test = pl.LazyFrame({'a': [1, 2, 3], 'b': [4, 5, 7]}) + df_train = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 7]}) + df_test = pl.LazyFrame({"a": [1, 2, 3], "b": [4, 5, 7]}) scaler = StandardScaler() scaler.fit(df_train) print(scaler.transform(df_test).collect()) diff --git a/docs/basics/dataframe.md b/docs/basics/dataframe.md index a47c5a530..3cd2c884c 100644 --- a/docs/basics/dataframe.md +++ b/docs/basics/dataframe.md @@ -25,12 +25,13 @@ Make a Python file with the following content: import narwhals as nw from narwhals.typing import FrameT + @nw.narwhalify def func(df: FrameT) -> FrameT: return df.select( - a_sum=nw.col('a').sum(), - a_mean=nw.col('a').mean(), - a_std=nw.col('a').std(), + a_sum=nw.col("a").sum(), + a_mean=nw.col("a").mean(), + a_std=nw.col("a").std(), ) ``` Let's try it out: @@ -39,7 +40,7 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex1" import pandas as pd - df = pd.DataFrame({'a': [1, 1, 2]}) + df = pd.DataFrame({"a": [1, 1, 2]}) print(func(df)) ``` @@ -47,7 +48,7 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex1" import polars as pl - df = pl.DataFrame({'a': [1, 1, 2]}) + df = pl.DataFrame({"a": [1, 1, 2]}) print(func(df)) ``` @@ -55,7 +56,7 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex1" import polars as pl - df = pl.LazyFrame({'a': [1, 1, 2]}) + df = pl.LazyFrame({"a": [1, 1, 2]}) print(func(df).collect()) ``` @@ -64,17 +65,18 @@ Alternatively, we could have opted for the more explicit version: import narwhals as nw from narwhals.typing import IntoFrameT + def func(df_native: IntoFrameT) -> IntoFrameT: df = nw.from_native(df_native) df = df.select( - a_sum=nw.col('a').sum(), - a_mean=nw.col('a').mean(), - a_std=nw.col('a').std(), + a_sum=nw.col("a").sum(), + a_mean=nw.col("a").mean(), + a_std=nw.col("a").std(), ) return nw.to_native(df) ``` Despite being more verbose, it has the advantage of preserving the type annotation of the native -object - see [typing](# todo) for more details. +object - see [typing](../api-reference/typing.md) for more details. In general, in this tutorial, we'll use the former. @@ -86,9 +88,10 @@ Make a Python file with the following content: import narwhals as nw from narwhals.typing import FrameT + @nw.narwhalify def func(df: FrameT) -> FrameT: - return df.group_by('a').agg(nw.col('b').mean()).sort('a') + return df.group_by("a").agg(nw.col("b").mean()).sort("a") ``` Let's try it out: @@ -96,7 +99,7 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex2" import pandas as pd - df = pd.DataFrame({'a': [1, 1, 2], 'b': [4, 5, 6]}) + df = pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]}) print(func(df)) ``` @@ -104,7 +107,7 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex2" import polars as pl - df = pl.DataFrame({'a': [1, 1, 2], 'b': [4, 5, 6]}) + df = pl.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]}) print(func(df)) ``` @@ -112,7 +115,7 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex2" import polars as pl - df = pl.LazyFrame({'a': [1, 1, 2], 'b': [4, 5, 6]}) + df = pl.LazyFrame({"a": [1, 1, 2], "b": [4, 5, 6]}) print(func(df).collect()) ``` @@ -126,9 +129,10 @@ Make a Python file with the following content: import narwhals as nw from narwhals.typing import FrameT + @nw.narwhalify def func(df: FrameT) -> FrameT: - return df.with_columns(a_plus_b=nw.sum_horizontal('a', 'b')) + return df.with_columns(a_plus_b=nw.sum_horizontal("a", "b")) ``` Let's try it out: @@ -136,7 +140,7 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex3" import pandas as pd - df = pd.DataFrame({'a': [1, 1, 2], 'b': [4, 5, 6]}) + df = pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]}) print(func(df)) ``` @@ -144,7 +148,7 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex3" import polars as pl - df = pl.DataFrame({'a': [1, 1, 2], 'b': [4, 5, 6]}) + df = pl.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]}) print(func(df)) ``` @@ -152,7 +156,7 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex3" import polars as pl - df = pl.LazyFrame({'a': [1, 1, 2], 'b': [4, 5, 6]}) + df = pl.LazyFrame({"a": [1, 1, 2], "b": [4, 5, 6]}) print(func(df).collect()) ``` @@ -170,6 +174,7 @@ from typing import Any import narwhals as nw + @nw.narwhalify(eager_only=True) def func(df: nw.DataFrame[Any], s: nw.Series, col_name: str) -> int: return df.filter(nw.col(col_name).is_in(s)).shape[0] @@ -183,16 +188,16 @@ Let's try it out: ```python exec="true" source="material-block" result="python" session="df_ex4" import pandas as pd - df = pd.DataFrame({'a': [1, 1, 2, 2, 3], 'b': [4, 5, 6, 7, 8]}) + df = pd.DataFrame({"a": [1, 1, 2, 2, 3], "b": [4, 5, 6, 7, 8]}) s = pd.Series([1, 3]) - print(func(df, s.to_numpy(), 'a')) + print(func(df, s.to_numpy(), "a")) ``` === "Polars (eager)" ```python exec="true" source="material-block" result="python" session="df_ex4" import polars as pl - df = pl.DataFrame({'a': [1, 1, 2, 2, 3], 'b': [4, 5, 6, 7, 8]}) + df = pl.DataFrame({"a": [1, 1, 2, 2, 3], "b": [4, 5, 6, 7, 8]}) s = pl.Series([1, 3]) - print(func(df, s.to_numpy(), 'a')) + print(func(df, s.to_numpy(), "a")) ``` diff --git a/docs/extending.md b/docs/extending.md index dc0d4671b..1a750431f 100644 --- a/docs/extending.md +++ b/docs/extending.md @@ -13,7 +13,7 @@ Alternatively, if you can't do that (for example, if you library is closed-sourc the next section for what else you can do. To check which methods are supported for which backend in depth, please refer to the -[API completeness page](api-completeness.md). +[API completeness page](api-completeness/index.md). ## Extending Narwhals @@ -26,14 +26,16 @@ Make sure that, in addition to the public Narwhals API, you also define: from `Narwhals.DataFrame` - `DataFrame.__narwhals_namespace__`: return an object which implements public top-level functions from `narwhals` (e.g. `narwhals.col`, `narwhals.concat`, ...) + - `DataFrame.__native_namespace__`: return a native namespace object which must have a + `from_dict` method - `LazyFrame.__narwhals_lazyframe__`: return an object which implements public methods from `Narwhals.LazyFrame` - `LazyFrame.__narwhals_namespace__`: return an object which implements public top-level functions from `narwhals` (e.g. `narwhals.col`, `narwhals.concat`, ...) + - `LazyFrame.__native_namespace__`: return a native namespace object which must have a + `from_dict` method - `Series.__narwhals_series__`: return an object which implements public methods from `Narwhals.Series` - - `Series.__narwhals_namespace__`: return an object which implements public top-level - functions from `narwhals` (e.g. `narwhals.col`, `narwhals.concat`, ...) If your library doesn't distinguish between lazy and eager, then it's OK for your dataframe object to implement both `__narwhals_dataframe__` and `__narwhals_lazyframe__`. In fact, diff --git a/docs/how_it_works.md b/docs/how_it_works.md index 4bba89e0e..cda98a2b6 100644 --- a/docs/how_it_works.md +++ b/docs/how_it_works.md @@ -13,33 +13,34 @@ Translating this to pandas syntax, we get: ```python exec="1" source="above" def col_a(df): - return [df.loc[:, 'a']] + return [df.loc[:, "a"]] ``` Let's step up the complexity. How about `nw.col('a')+1`? We already know what the `nw.col('a')` part looks like, so we just need to add `1` to each of its outputs: -```python exec="1" +```python exec="1" source="above" def col_a(df): - return [df.loc[:, 'a']] + return [df.loc[:, "a"]] + def col_a_plus_1(df): - return [x+1 for x in col_a(df)] + return [x + 1 for x in col_a(df)] ``` Expressions can return multiple Series - for example, `nw.col('a', 'b')` translates to: -```python exec="1" +```python exec="1" source="above" def col_a_b(df): - return [df.loc[:, 'a'], df.loc[:, 'b']] + return [df.loc[:, "a"], df.loc[:, "b"]] ``` Expressions can also take multiple columns as input - for example, `nw.sum_horizontal('a', 'b')` translates to: -```python exec="1" +```python exec="1" source="above" def sum_horizontal_a_b(df): - return [df.loc[:, 'a'] + df.loc[:, 'b']] + return [df.loc[:, "a"] + df.loc[:, "b"]] ``` Note that although an expression may have multiple columns as input, @@ -75,7 +76,7 @@ pn = PandasLikeNamespace( implementation=Implementation.PANDAS, backend_version=parse_version(pd.__version__), ) -print(nw.col('a')._call(pn)) +print(nw.col("a")._call(pn)) ``` The result from the last line above is the same as we'd get from `pn.col('a')`, and it's a `narwhals._pandas_like.expr.PandasLikeExpr` object, which we'll call `PandasLikeExpr` for @@ -87,7 +88,7 @@ The `_call` method gives us that function! Let's see it in action. Note: the following examples use `PandasLikeDataFrame` and `PandasLikeSeries`. These are backed by actual `pandas.DataFrame`s and `pandas.Series` respectively and are Narwhals-compliant. We can access the -underlying pandas objects via `PandasLikeDataFrame._native_dataframe` and `PandasLikeSeries._native_series`. +underlying pandas objects via `PandasLikeDataFrame._native_frame` and `PandasLikeSeries._native_series`. ```python exec="1" result="python" session="pandas_impl" source="above" import narwhals as nw @@ -102,16 +103,16 @@ pn = PandasLikeNamespace( backend_version=parse_version(pd.__version__), ) -df_pd = pd.DataFrame({'a': [1,2,3], 'b': [4,5,6]}) +df_pd = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) df = PandasLikeDataFrame( df_pd, implementation=Implementation.PANDAS, backend_version=parse_version(pd.__version__), ) -expression = pn.col('a') + 1 +expression = pn.col("a") + 1 result = expression._call(df) -print(f'length of result: {len(result)}\n') -print('native series of first value of result: ') +print(f"length of result: {len(result)}\n") +print("native series of first value of result: ") print([x._native_series for x in result][0]) ``` @@ -133,29 +134,112 @@ this section better are 110% welcome. ## Polars and other implementations -Other implementations are similar to the above: their define their own Narwhals-compliant +Other implementations are similar to the above: they define their own Narwhals-compliant objects. So, all-in-all, there are a couple of layers here: - `nw.DataFrame` is backed by a Narwhals-compliant Dataframe, such as: - - `narwhals._pandas_like.dataframe.PandasLikeDataFrame` - - `narwhals._arrow.dataframe.ArrowDataFrame` - - `narwhals._polars.dataframe.PolarsDataFrame` + - `narwhals._pandas_like.dataframe.PandasLikeDataFrame` + - `narwhals._arrow.dataframe.ArrowDataFrame` + - `narwhals._polars.dataframe.PolarsDataFrame` - each Narwhals-compliant DataFrame is backed by a native Dataframe, for example: - - `narwhals._pandas_like.dataframe.PandasLikeDataFrame` is backed by a pandas DataFrame - - `narwhals._arrow.dataframe.ArrowDataFrame` is backed by a PyArrow Table - - `narwhals._polars.dataframe.PolarsDataFrame` is backed by a Polars DataFrame + - `narwhals._pandas_like.dataframe.PandasLikeDataFrame` is backed by a pandas DataFrame + - `narwhals._arrow.dataframe.ArrowDataFrame` is backed by a PyArrow Table + - `narwhals._polars.dataframe.PolarsDataFrame` is backed by a Polars DataFrame Each implementation defines its own objects in subfolders such as `narwhals._pandas_like`, `narwhals._arrow`, `narwhals._polars`, whereas the top-level modules such as `narwhals.dataframe` and `narwhals.series` coordinate how to dispatch the Narwhals API to each backend. +## Mapping from API to implementations + +If an end user executes some Narwhals code, such as + +```python +df.select(nw.col("a") + 1) +``` +then how does that get mapped to the underlying dataframe's native API? Let's walk through +this example to see. + +Things generally go through a couple of layers: + +- The user calls some top-level Narwhals API. +- The Narwhals API forwards the call to a Narwhals-compliant dataframe wrapper, such as + - `PandasLikeDataFrame` / `ArrowDataFrame` / `PolarsDataFrame` / ... + - `PandasLikeSeries` / `ArrowSeries` / `PolarsSeries` / ... + - `PandasLikeExpr` / `ArrowExpr` / `PolarsExpr` / ... +- The dataframe wrapper forwards the call to the underlying library, e.g.: + - `PandasLikeDataFrame` forwards the call to the underlying pandas/Modin/cuDF dataframe. + - `ArrowDataFrame` forwards the call to the underlying PyArrow table. + - `PolarsDataFrame` forwards the call to the underlying Polars DataFrame. + +The way you access the Narwhals-compliant wrapper depends on the object: + +- `narwhals.DataFrame` and `narwhals.LazyFrame`: use the `._compliant_frame` attribute. +- `narwhals.Series`: use the `._compliant_series` attribute. +- `narwhals.Expr`: call the `._call` method, and pass to it the Narwhals-compliant namespace associated with + the given backend. + +🛑 BUT WAIT! What's a Narwhals-compliant namespace? + +Each backend is expected to implement a Narwhals-compliant +namespace (`PandasLikeNamespace`, `ArrowNamespace`, `PolarsNamespace`). These can be used to interact with the Narwhals-compliant +Dataframe and Series objects described above - let's work through the motivating example to see how. + +```python exec="1" session="pandas_api_mapping" source="above" +import narwhals as nw +from narwhals._pandas_like.namespace import PandasLikeNamespace +from narwhals._pandas_like.utils import Implementation +from narwhals._pandas_like.dataframe import PandasLikeDataFrame +from narwhals.utils import parse_version +import pandas as pd + +pn = PandasLikeNamespace( + implementation=Implementation.PANDAS, + backend_version=parse_version(pd.__version__), +) + +df_pd = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) +df = nw.from_native(df_pd) +df.select(nw.col("a") + 1) +``` + +The first thing `narwhals.DataFrame.select` does is to parse each input expression to end up with a compliant expression for the given +backend, and it does so by passing a Narwhals-compliant namespace to `nw.Expr._call`: + +```python exec="1" result="python" session="pandas_api_mapping" source="above" +pn = PandasLikeNamespace( + implementation=Implementation.PANDAS, + backend_version=parse_version(pd.__version__), +) +expr = (nw.col("a") + 1)._call(pn) +print(expr) +``` +If we then extract a Narwhals-compliant dataframe from `df` by +calling `._compliant_frame`, we get a `PandasLikeDataFrame` - and that's an object which we can pass `expr` to! + +```python exec="1" session="pandas_api_mapping" source="above" +df_compliant = df._compliant_frame +result = df_compliant.select(expr) +``` + +We can then view the underlying pandas Dataframe which was produced by calling `._native_frame`: + +```python exec="1" result="python" session="pandas_api_mapping" source="above" +print(result._native_frame) +``` +which is the same as we'd have obtained by just using the Narwhals API directly: + +```python exec="1" result="python" session="pandas_api_mapping" source="above" +print(nw.to_native(df.select(nw.col("a") + 1))) +``` + ## Group-by Group-by is probably one of Polars' most significant innovations (on the syntax side) with respect to pandas. We can write something like ```python df: pl.DataFrame -df.group_by('a').agg((pl.col('c') > pl.col('b').mean()).max()) +df.group_by("a").agg((pl.col("c") > pl.col("b").mean()).max()) ``` To do this in pandas, we need to either use `GroupBy.apply` (sloooow), or do some crazy manual optimisations to get it to work. @@ -165,9 +249,8 @@ In Narwhals, here's what we do: - if somebody uses a simple group-by aggregation (e.g. `df.group_by('a').agg(nw.col('b').mean())`), then on the pandas side we translate it to ```python - df: pd.DataFrame - df.groupby('a').agg({'b': ['mean']}) + df.groupby("a").agg({"b": ["mean"]}) ``` - if somebody passes a complex group-by aggregation, then we use `apply` and raise a `UserWarning`, warning users of the performance penalty and advising them to refactor their code so that the aggregation they perform @@ -176,9 +259,9 @@ In Narwhals, here's what we do: In order to tell whether an aggregation is simple, Narwhals uses the private `_depth` attribute of `PandasLikeExpr`: ```python exec="1" result="python" session="pandas_impl" source="above" -print(pn.col('a').mean()) -print((pn.col('a')+1).mean()) -print(pn.mean('a')) +print(pn.col("a").mean()) +print((pn.col("a") + 1).mean()) +print(pn.mean("a")) ``` For simple aggregations, Narwhals can just look at `_depth` and `function_name` and figure out diff --git a/docs/installation.md b/docs/installation.md index 12d8257b1..04625c686 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -11,6 +11,6 @@ Then, if you start the Python REPL and see the following: ```python >>> import narwhals >>> narwhals.__version__ -'1.1.8' +'1.5.5' ``` then installation worked correctly! diff --git a/docs/levels.md b/docs/levels.md index afe0ea2eb..743334663 100644 --- a/docs/levels.md +++ b/docs/levels.md @@ -11,11 +11,12 @@ For example: import narwhals as nw from narwhals.typing import FrameT + @nw.narwhalify def func(df: FrameT) -> FrameT: - return df.group_by('a').agg( - b_mean=nw.col('b').mean(), - b_std=nw.col('b').std(), + return df.group_by("a").agg( + b_mean=nw.col("b").mean(), + b_std=nw.col("b").std(), ) ``` will work for any of pandas, Polars, cuDF, Modin, and PyArrow. diff --git a/docs/other/column_names.md b/docs/other/column_names.md new file mode 100644 index 000000000..14346a303 --- /dev/null +++ b/docs/other/column_names.md @@ -0,0 +1,24 @@ +# Column names + +Polars and PyArrow only allow for string column names. What about pandas? + +```python +>>> import pandas as pd +>>> pd.concat([pd.Series([1, 2], name=0), pd.Series([1, 3], name=0)], axis=1) + 0 0 +0 1 1 +1 2 3 +``` + +Oh...not only does it let us create a dataframe with a column named `0` - it lets us +create one with _two_ such columns! + +What does Narwhals do about this? + +- In general, non-string column names are supported. In some places where this might + create ambiguity (such as `DataFrame.__getitem__` or `DataFrame.select`) we may be strict and only + allow passing in column names if they're strings. +- If you have a use-case that's + failing for non-string column names, please report it to [https://github.com/narwhals-dev/narwhals/issues](https://github.com/narwhals-dev/narwhals/issues) + and we'll see if we can support it. +- Duplicate column names are 🚫 banned 🚫. diff --git a/docs/other/pandas_index.md b/docs/other/pandas_index.md index 8217d4153..ccc38d490 100644 --- a/docs/other/pandas_index.md +++ b/docs/other/pandas_index.md @@ -21,9 +21,10 @@ Let's learn about what Narwhals promises. ```python exec="1" source="above" session="ex1" import narwhals as nw + def my_func(df): df = nw.from_native(df) - df = df.with_columns(a_plus_one=nw.col('a')+1) + df = df.with_columns(a_plus_one=nw.col("a") + 1) return nw.to_native(df) ``` @@ -32,7 +33,7 @@ Let's start with a dataframe with an Index with values `[7, 8, 9]`. ```python exec="true" source="material-block" result="python" session="ex1" import pandas as pd -df = pd.DataFrame({'a': [2, 1, 3], 'b': [3, 5, -3]}, index=[7, 8, 9]) +df = pd.DataFrame({"a": [2, 1, 3], "b": [3, 5, -3]}, index=[7, 8, 9]) print(my_func(df)) ``` @@ -46,9 +47,9 @@ pandas automatically aligns indices for users. For example: ```python exec="1" source="above" session="ex2" import pandas as pd -df_pd = pd.DataFrame({'a': [2, 1, 3], 'b': [4, 5, 6]}) -s_pd = df_pd['a'].sort_values() -df_pd['a_sorted'] = s_pd +df_pd = pd.DataFrame({"a": [2, 1, 3], "b": [4, 5, 6]}) +s_pd = df_pd["a"].sort_values() +df_pd["a_sorted"] = s_pd ``` Reading the code, you might expect that `'a_sorted'` will contain the values `[1, 2, 3]`. diff --git a/docs/other/user_warning.md b/docs/other/user_warning.md new file mode 100644 index 000000000..60e20d4d6 --- /dev/null +++ b/docs/other/user_warning.md @@ -0,0 +1,103 @@ +# Avoiding the `UserWarning` error while using Pandas `group_by` + +## Introduction + +If you have ever experienced the + +> UserWarning: Found complex group-by expression, which can't be expressed efficiently with the pandas API. If you can, please rewrite your query such that group-by aggregations are simple (e.g. mean, std, min, max, ...) + +message while using the narwhals `group_by()` method, this is for you. If you haven't, this is also for you as you might experience it and you need to know how to avoid it. + +The pandas API most likely cannot efficiently handle the complexity of the aggregation operations you are trying to run. Take the following two codes as an example. + +=== "Approach 1" + ```python exec="true" source="above" result="python" session="df_ex1" + import narwhals as nw + import pandas as pd + + data = {"a": [1, 2, 3, 4, 5], "b": [5, 4, 3, 2, 1], "c": [10, 20, 30, 40, 50]} + + df_pd = pd.DataFrame(data) + + + @nw.narwhalify + def approach_1(df): + + # Pay attention to this next line + df = df.group_by("a").agg(d=(nw.col("b") + nw.col("c")).sum()) + + return df + + + print(approach_1(df_pd)) + ``` + +=== "Approach 2" + ```python exec="true" source="above" result="python" session="df_ex2" + import narwhals as nw + import pandas as pd + + data = {"a": [1, 2, 3, 4, 5], "b": [5, 4, 3, 2, 1], "c": [10, 20, 30, 40, 50]} + + df_pd = pd.DataFrame(data) + + + @nw.narwhalify + def approach_2(df): + + # Pay attention to this next line + df = df.with_columns(d=nw.col("b") + nw.col("c")).group_by("a").agg(nw.sum("d")) + + return df + + + print(approach_2(df_pd)) + ``` + + +Both Approaches shown above return the exact same result, but Approach 1 is inefficient and returns the warning message +we showed at the top. + +What makes the first approach inefficient and the second approach efficient? It comes down to what the +pandas API lets us express. + +## Approach 1 +```python +# From line 11 + +return df.group_by("a").agg((nw.col("b") + nw.col("c")).sum().alias("d")) +``` + +To translate this to pandas, we would do: +```python +df.groupby("a").apply( + lambda df: pd.Series([(df["b"] + df["c"]).sum()], index=["d"]), include_groups=False +) +``` +Any time you use `apply` in pandas, that's a performance footgun - best to avoid it and use vectorised operations instead. +Let's take a look at how "approach 2" gets translated to pandas to see the difference. + +## Approach 2 +```python +# Line 11 in Approach 2 + +return df.with_columns(d=nw.col("b") + nw.col("c")).group_by("a").agg({"d": "sum"}) +``` + +This gets roughly translated to: +```python +df.assign(d=lambda df: df["b"] + df["c"]).groupby("a").agg({"d": "sum"}) +``` +Because we're using pandas' own API, as opposed to `apply` and a custom `lambda` function, then this is going to be much more efficient. + +## Tips for Avoiding the `UserWarning` + +To ensure efficiency and avoid warnings similar to those seen in Approach 1, we recommend that you follow these practices: + +1. Decompose complex operations: break down complex transformations into simpler steps. In this case, keep the `.agg` method simple. Compute new columns first, then use these columns in aggregation or other operations. +2. Avoid redundant computations: if an operation (like addition) is used multiple times, compute it once and store the result in a new column. +3. Leverage built-in functions: use built-in functions provided by the DataFrame library. In this case, using the `with_columns()` method allows you to pre-compute before grouping and aggregation. + +By following these guidelines, you can are sure to avoid the aforementioned warning. + +**_Happy grouping!_** 🫡 diff --git a/docs/quick_start.md b/docs/quick_start.md index e106ac024..f3ff8c05a 100644 --- a/docs/quick_start.md +++ b/docs/quick_start.md @@ -29,12 +29,12 @@ def my_function(df_native: IntoFrame) -> list[str]: return column_names -df_pandas = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]}) -df_polars = pl.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]}) +df_pandas = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) +df_polars = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) -print('pandas output') +print("pandas output") print(my_function(df_pandas)) -print('Polars output') +print("Polars output") print(my_function(df_polars)) ``` diff --git a/docs/roadmap.md b/docs/roadmap.md index f0397ca34..87b224bf9 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -1,8 +1,11 @@ # Roadmap -Priorities, as of July 2024, are: +Priorities, as of August 2024, are: -- Works towards supporting projects which have shown interest in Narwhals: scikit-learn, shiny, tubular -- Implement when/then/otherwise so that Narwhals is API-complete enough to complete all the TPC-H queries -- Add support for extra backends such as Dask -- Add extra docs and tutorials to make the project more accessible and easy to get started with +- Works towards supporting projects which have shown interest in Narwhals. +- Implement when/then/otherwise so that Narwhals is API-complete enough to complete all the TPC-H queries. +- Make Dask support complete-enough, at least to the point that it can execute TPC-H queries. +- Improve support for cuDF, which we can't currently test in CI (unless NVIDIA helps us out :wink:) but + which we can and do test manually in Kaggle notebooks. +- Add extra docs and tutorials to make the project more accessible and easy to get started with. +- Look into extra backends, such as DuckDB and Ibis. diff --git a/docs/why.md b/docs/why.md index 3c6f4825c..adf8f39b4 100644 --- a/docs/why.md +++ b/docs/why.md @@ -18,13 +18,13 @@ Polars checks if it's in the values. For another example, try running the code below - note how the outputs have different column names after the join! ```python -pd_df_left = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]}) -pd_df_right = pd.DataFrame({'a': [1, 2, 3], 'c': [4, 5, 6]}) -pd_left_merge = pd_df_left.merge(pd_df_right, left_on='b', right_on='c', how='left') +pd_df_left = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) +pd_df_right = pd.DataFrame({"a": [1, 2, 3], "c": [4, 5, 6]}) +pd_left_merge = pd_df_left.merge(pd_df_right, left_on="b", right_on="c", how="left") -pl_df_left = pl.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]}) -pl_df_right = pl.DataFrame({'a': [1, 2, 3], 'c': [4, 5, 6]}) -pl_left_merge = pl_df_left.join(pl_df_right, left_on='b', right_on='c', how='left') +pl_df_left = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) +pl_df_right = pl.DataFrame({"a": [1, 2, 3], "c": [4, 5, 6]}) +pl_left_merge = pl_df_left.join(pl_df_right, left_on="b", right_on="c", how="left") print(pd_left_merge.columns) print(pl_df_right.columns) diff --git a/mkdocs.yml b/mkdocs.yml index c27bb4048..8b635f78d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -11,8 +11,10 @@ nav: - basics/dataframe.md - basics/column.md - basics/complete_example.md - - Other concepts: + - Pandas-like concepts: - other/pandas_index.md + - other/user_warning.md + - other/column_names.md - levels.md - overhead.md - backcompat.md @@ -20,7 +22,11 @@ nav: - how_it_works.md - Roadmap: roadmap.md - Related projects: related.md - - API Completeness: api-completeness.md + - API Completeness: + - api-completeness/index.md + - api-completeness/dataframe.md + - api-completeness/expr.md + - api-completeness/series.md - API Reference: - api-reference/narwhals.md - api-reference/dataframe.md @@ -44,12 +50,15 @@ nav: theme: name: material font: false + favicon: assets/image.png + logo: assets/image.png features: - content.code.copy - content.code.annotate - navigation.footer - navigation.indexes - palette: + - navigation.top + palette: # Palette toggle for automatic mode - media: "(prefers-color-scheme)" toggle: diff --git a/narwhals/__init__.py b/narwhals/__init__.py index c31459447..0977e716a 100644 --- a/narwhals/__init__.py +++ b/narwhals/__init__.py @@ -1,3 +1,4 @@ +from narwhals import dependencies from narwhals import selectors from narwhals import stable from narwhals.dataframe import DataFrame @@ -22,14 +23,15 @@ from narwhals.dtypes import UInt64 from narwhals.dtypes import Unknown from narwhals.expr import Expr -from narwhals.expr import all +from narwhals.expr import all_ as all from narwhals.expr import all_horizontal from narwhals.expr import any_horizontal from narwhals.expr import col -from narwhals.expr import len +from narwhals.expr import len_ as len from narwhals.expr import lit from narwhals.expr import max from narwhals.expr import mean +from narwhals.expr import mean_horizontal from narwhals.expr import min from narwhals.expr import sum from narwhals.expr import sum_horizontal @@ -37,6 +39,7 @@ from narwhals.functions import concat from narwhals.functions import from_dict from narwhals.functions import get_level +from narwhals.functions import new_series from narwhals.functions import show_versions from narwhals.schema import Schema from narwhals.series import Series @@ -47,20 +50,24 @@ from narwhals.utils import is_ordered_categorical from narwhals.utils import maybe_align_index from narwhals.utils import maybe_convert_dtypes +from narwhals.utils import maybe_get_index from narwhals.utils import maybe_set_index -__version__ = "1.1.8" +__version__ = "1.5.5" __all__ = [ + "dependencies", "selectors", "concat", "from_dict", "get_level", + "new_series", "to_native", "from_native", "is_ordered_categorical", "maybe_align_index", "maybe_convert_dtypes", + "maybe_get_index", "maybe_set_index", "get_native_namespace", "all", @@ -72,6 +79,7 @@ "min", "max", "mean", + "mean_horizontal", "sum", "sum_horizontal", "when", diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py index 4aee5c2f0..20e507166 100644 --- a/narwhals/_arrow/dataframe.py +++ b/narwhals/_arrow/dataframe.py @@ -8,18 +8,19 @@ from typing import Sequence from typing import overload +from narwhals._arrow.utils import broadcast_series from narwhals._arrow.utils import translate_dtype from narwhals._arrow.utils import validate_dataframe_comparand from narwhals._expression_parsing import evaluate_into_exprs -from narwhals.dependencies import get_numpy from narwhals.dependencies import get_pyarrow -from narwhals.dependencies import get_pyarrow_compute -from narwhals.dependencies import get_pyarrow_parquet +from narwhals.dependencies import is_numpy_array from narwhals.utils import Implementation from narwhals.utils import flatten from narwhals.utils import generate_unique_token +from narwhals.utils import parse_columns_to_drop if TYPE_CHECKING: + import numpy as np from typing_extensions import Self from narwhals._arrow.group_by import ArrowGroupBy @@ -34,7 +35,7 @@ class ArrowDataFrame: def __init__( self, native_dataframe: Any, *, backend_version: tuple[int, ...] ) -> None: - self._native_dataframe = native_dataframe + self._native_frame = native_dataframe self._implementation = Implementation.PYARROW self._backend_version = backend_version @@ -52,15 +53,18 @@ def __narwhals_dataframe__(self) -> Self: def __narwhals_lazyframe__(self) -> Self: return self - def _from_native_dataframe(self, df: Any) -> Self: + def _from_native_frame(self, df: Any) -> Self: return self.__class__(df, backend_version=self._backend_version) @property def shape(self) -> tuple[int, int]: - return self._native_dataframe.shape # type: ignore[no-any-return] + return self._native_frame.shape # type: ignore[no-any-return] def __len__(self) -> int: - return len(self._native_dataframe) + return len(self._native_frame) + + def row(self, index: int) -> tuple[Any, ...]: + return tuple(col[index] for col in self._native_frame) def rows( self, *, named: bool = False @@ -68,7 +72,7 @@ def rows( if not named: msg = "Unnamed rows are not yet supported on PyArrow tables" raise NotImplementedError(msg) - return self._native_dataframe.to_pylist() # type: ignore[no-any-return] + return self._native_frame.to_pylist() # type: ignore[no-any-return] def iter_rows( self, @@ -76,7 +80,7 @@ def iter_rows( named: bool = False, buffer_size: int = 512, ) -> Iterator[tuple[Any, ...]] | Iterator[dict[str, Any]]: - df = self._native_dataframe + df = self._native_frame num_rows = df.num_rows if not named: @@ -95,11 +99,14 @@ def get_column(self, name: str) -> ArrowSeries: raise TypeError(msg) return ArrowSeries( - self._native_dataframe[name], + self._native_frame[name], name=name, backend_version=self._backend_version, ) + def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray: + return self._native_frame.__array__(dtype, copy=copy) + @overload def __getitem__(self, item: tuple[Sequence[int], str | int]) -> ArrowSeries: ... # type: ignore[overload-overlap] @@ -119,7 +126,7 @@ def __getitem__( from narwhals._arrow.series import ArrowSeries return ArrowSeries( - self._native_dataframe[item], + self._native_frame[item], name=item, backend_version=self._backend_version, ) @@ -128,17 +135,43 @@ def __getitem__( and len(item) == 2 and isinstance(item[1], (list, tuple)) ): - return self._from_native_dataframe( - self._native_dataframe.take(item[0]).select(item[1]) + return self._from_native_frame( + self._native_frame.take(item[0]).select(item[1]) ) elif isinstance(item, tuple) and len(item) == 2: + if isinstance(item[1], slice): + columns = self.columns + if isinstance(item[1].start, str) or isinstance(item[1].stop, str): + start = ( + columns.index(item[1].start) + if item[1].start is not None + else None + ) + stop = ( + columns.index(item[1].stop) + 1 + if item[1].stop is not None + else None + ) + step = item[1].step + return self._from_native_frame( + self._native_frame.take(item[0]).select(columns[start:stop:step]) + ) + if isinstance(item[1].start, int) or isinstance(item[1].stop, int): + return self._from_native_frame( + self._native_frame.take(item[0]).select( + columns[item[1].start : item[1].stop : item[1].step] + ) + ) + msg = f"Expected slice of integers or strings, got: {type(item[1])}" # pragma: no cover + raise TypeError(msg) # pragma: no cover + from narwhals._arrow.series import ArrowSeries # PyArrow columns are always strings col_name = item[1] if isinstance(item[1], str) else self.columns[item[1]] return ArrowSeries( - self._native_dataframe[col_name].take(item[0]), + self._native_frame[col_name].take(item[0]), name=col_name, backend_version=self._backend_version, ) @@ -148,17 +181,13 @@ def __getitem__( msg = "Slicing with step is not supported on PyArrow tables" raise NotImplementedError(msg) start = item.start or 0 - stop = item.stop or len(self._native_dataframe) - return self._from_native_dataframe( - self._native_dataframe.slice(item.start, stop - start), + stop = item.stop or len(self._native_frame) + return self._from_native_frame( + self._native_frame.slice(item.start, stop - start), ) - elif isinstance(item, Sequence) or ( - (np := get_numpy()) is not None - and isinstance(item, np.ndarray) - and item.ndim == 1 - ): - return self._from_native_dataframe(self._native_dataframe.take(item)) + elif isinstance(item, Sequence) or (is_numpy_array(item) and item.ndim == 1): + return self._from_native_frame(self._native_frame.take(item)) else: # pragma: no cover msg = f"Expected str or slice, got: {type(item)}" @@ -166,7 +195,7 @@ def __getitem__( @property def schema(self) -> dict[str, DType]: - schema = self._native_dataframe.schema + schema = self._native_frame.schema return { name: translate_dtype(dtype) for name, dtype in zip(schema.names, schema.types) @@ -177,23 +206,25 @@ def collect_schema(self) -> dict[str, DType]: @property def columns(self) -> list[str]: - return self._native_dataframe.schema.names # type: ignore[no-any-return] + return self._native_frame.schema.names # type: ignore[no-any-return] def select( self, *exprs: IntoArrowExpr, **named_exprs: IntoArrowExpr, ) -> Self: + import pyarrow as pa # ignore-banned-import() + new_series = evaluate_into_exprs(self, *exprs, **named_exprs) if not new_series: # return empty dataframe, like Polars does - return self._from_native_dataframe( - self._native_dataframe.__class__.from_arrays([]) - ) + return self._from_native_frame(self._native_frame.__class__.from_arrays([])) names = [s.name for s in new_series] - pa = get_pyarrow() - df = pa.Table.from_arrays([s._native_series for s in new_series], names=names) - return self._from_native_dataframe(df) + df = pa.Table.from_arrays( + broadcast_series(new_series), + names=names, + ) + return self._from_native_frame(df) def with_columns( self, @@ -216,7 +247,7 @@ def with_columns( ) ) else: - to_concat.append(self._native_dataframe[name]) + to_concat.append(self._native_frame[name]) output_names.append(name) for s in new_column_name_to_new_column_map: to_concat.append( @@ -227,8 +258,8 @@ def with_columns( ) ) output_names.append(s) - df = self._native_dataframe.__class__.from_arrays(to_concat, names=output_names) - return self._from_native_dataframe(df) + df = self._native_frame.__class__.from_arrays(to_concat, names=output_names) + return self._from_native_frame(df) def group_by(self, *keys: str) -> ArrowGroupBy: from narwhals._arrow.group_by import ArrowGroupBy @@ -261,19 +292,21 @@ def join( n_bytes=8, columns=[*self.columns, *other.columns] ) - return self._from_native_dataframe( - self.with_columns(**{key_token: plx.lit(0, None)})._native_dataframe.join( - other.with_columns(**{key_token: plx.lit(0, None)})._native_dataframe, + return self._from_native_frame( + self.with_columns(**{key_token: plx.lit(0, None)}) + ._native_frame.join( + other.with_columns(**{key_token: plx.lit(0, None)})._native_frame, keys=key_token, right_keys=key_token, join_type="inner", right_suffix="_right", - ), - ).drop(key_token) + ) + .drop([key_token]), + ) - return self._from_native_dataframe( - self._native_dataframe.join( - other._native_dataframe, + return self._from_native_frame( + self._native_frame.join( + other._native_frame, keys=left_on, right_keys=right_on, join_type=how_to_join_map[how], @@ -281,11 +314,18 @@ def join( ), ) - def drop(self, *columns: str) -> Self: - return self._from_native_dataframe(self._native_dataframe.drop(list(columns))) + def drop(self: Self, columns: list[str], strict: bool) -> Self: # noqa: FBT001 + to_drop = parse_columns_to_drop( + compliant_frame=self, columns=columns, strict=strict + ) + return self._from_native_frame(self._native_frame.drop(to_drop)) - def drop_nulls(self) -> Self: - return self._from_native_dataframe(self._native_dataframe.drop_null()) + def drop_nulls(self: Self, subset: str | list[str] | None) -> Self: + if subset is None: + return self._from_native_frame(self._native_frame.drop_null()) + subset = [subset] if isinstance(subset, str) else subset + plx = self.__narwhals_namespace__() + return self.filter(~plx.any_horizontal(plx.col(*subset).is_null())) def sort( self, @@ -294,7 +334,7 @@ def sort( descending: bool | Sequence[bool] = False, ) -> Self: flat_keys = flatten([*flatten([by]), *more_by]) - df = self._native_dataframe + df = self._native_frame if isinstance(descending, bool): order = "descending" if descending else "ascending" @@ -304,18 +344,18 @@ def sort( (key, "descending" if is_descending else "ascending") for key, is_descending in zip(flat_keys, descending) ] - return self._from_native_dataframe(df.sort_by(sorting=sorting)) + return self._from_native_frame(df.sort_by(sorting=sorting)) def to_pandas(self) -> Any: - return self._native_dataframe.to_pandas() + return self._native_frame.to_pandas() def to_numpy(self) -> Any: - import numpy as np + import numpy as np # ignore-banned-import - return np.column_stack([col.to_numpy() for col in self._native_dataframe.columns]) + return np.column_stack([col.to_numpy() for col in self._native_frame.columns]) def to_dict(self, *, as_series: bool) -> Any: - df = self._native_dataframe + df = self._native_frame names_and_values = zip(df.column_names, df.columns) if as_series: @@ -329,58 +369,61 @@ def to_dict(self, *, as_series: bool) -> Any: return {name: col.to_pylist() for name, col in names_and_values} def with_row_index(self, name: str) -> Self: - pa = get_pyarrow() - df = self._native_dataframe + import pyarrow as pa # ignore-banned-import() + + df = self._native_frame row_indices = pa.array(range(df.num_rows)) - return self._from_native_dataframe(df.append_column(name, row_indices)) + return self._from_native_frame(df.append_column(name, row_indices)) def filter( self, *predicates: IntoArrowExpr, ) -> Self: - from narwhals._arrow.namespace import ArrowNamespace - - plx = ArrowNamespace(backend_version=self._backend_version) - expr = plx.all_horizontal(*predicates) - # Safety: all_horizontal's expression only returns a single column. - mask = expr._call(self)[0] - return self._from_native_dataframe( - self._native_dataframe.filter(mask._native_series) - ) + if ( + len(predicates) == 1 + and isinstance(predicates[0], list) + and all(isinstance(x, bool) for x in predicates[0]) + ): + mask = predicates[0] + else: + plx = self.__narwhals_namespace__() + expr = plx.all_horizontal(*predicates) + # Safety: all_horizontal's expression only returns a single column. + mask = expr._call(self)[0]._native_series + return self._from_native_frame(self._native_frame.filter(mask)) def null_count(self) -> Self: - pa = get_pyarrow() - df = self._native_dataframe + import pyarrow as pa # ignore-banned-import() + + df = self._native_frame names_and_values = zip(df.column_names, df.columns) - return self._from_native_dataframe( + return self._from_native_frame( pa.table({name: [col.null_count] for name, col in names_and_values}) ) def head(self, n: int) -> Self: - df = self._native_dataframe + df = self._native_frame if n >= 0: - return self._from_native_dataframe(df.slice(0, n)) + return self._from_native_frame(df.slice(0, n)) else: num_rows = df.num_rows - return self._from_native_dataframe(df.slice(0, max(0, num_rows + n))) + return self._from_native_frame(df.slice(0, max(0, num_rows + n))) def tail(self, n: int) -> Self: - df = self._native_dataframe + df = self._native_frame if n >= 0: num_rows = df.num_rows - return self._from_native_dataframe(df.slice(max(0, num_rows - n))) + return self._from_native_frame(df.slice(max(0, num_rows - n))) else: - return self._from_native_dataframe(df.slice(abs(n))) + return self._from_native_frame(df.slice(abs(n))) def lazy(self) -> Self: return self def collect(self) -> ArrowDataFrame: - return ArrowDataFrame( - self._native_dataframe, backend_version=self._backend_version - ) + return ArrowDataFrame(self._native_frame, backend_version=self._backend_version) def clone(self) -> Self: msg = "clone is not yet supported on PyArrow tables" @@ -398,31 +441,44 @@ def item(self: Self, row: int | None = None, column: int | str | None = None) -> f" frame has shape {self.shape!r}" ) raise ValueError(msg) - return self._native_dataframe[0][0] + return self._native_frame[0][0] elif row is None or column is None: msg = "cannot call `.item()` with only one of `row` or `column`" raise ValueError(msg) _col = self.columns.index(column) if isinstance(column, str) else column - return self._native_dataframe[_col][row] + return self._native_frame[_col][row] def rename(self, mapping: dict[str, str]) -> Self: - df = self._native_dataframe + df = self._native_frame new_cols = [mapping.get(c, c) for c in df.column_names] - return self._from_native_dataframe(df.rename_columns(new_cols)) + return self._from_native_frame(df.rename_columns(new_cols)) def write_parquet(self, file: Any) -> Any: - pp = get_pyarrow_parquet() - pp.write_table(self._native_dataframe, file) + import pyarrow.parquet as pp # ignore-banned-import + + pp.write_table(self._native_frame, file) + + def write_csv(self, file: Any) -> Any: + import pyarrow as pa # ignore-banned-import + import pyarrow.csv as pa_csv # ignore-banned-import + + pa_table = self._native_frame + if file is None: + csv_buffer = pa.BufferOutputStream() + pa_csv.write_csv(pa_table, csv_buffer) + return csv_buffer.getvalue().to_pybytes().decode() + return pa_csv.write_csv(pa_table, file) def is_duplicated(self: Self) -> ArrowSeries: + import numpy as np # ignore-banned-import + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + from narwhals._arrow.series import ArrowSeries - np = get_numpy() - pa = get_pyarrow() - pc = get_pyarrow_compute() - df = self._native_dataframe + df = self._native_frame columns = self.columns col_token = generate_unique_token(n_bytes=8, columns=columns) @@ -440,9 +496,10 @@ def is_duplicated(self: Self) -> ArrowSeries: return ArrowSeries(is_duplicated, name="", backend_version=self._backend_version) def is_unique(self: Self) -> ArrowSeries: + import pyarrow.compute as pc # ignore-banned-import() + from narwhals._arrow.series import ArrowSeries - pc = get_pyarrow_compute() is_duplicated = self.is_duplicated()._native_series return ArrowSeries( @@ -461,12 +518,11 @@ def unique( The param `maintain_order` is only here for compatibility with the polars API and has no effect on the output. """ + import numpy as np # ignore-banned-import + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() - np = get_numpy() - pa = get_pyarrow() - pc = get_pyarrow_compute() - - df = self._native_dataframe + df = self._native_frame if isinstance(subset, str): subset = [subset] @@ -484,10 +540,13 @@ def unique( .column(f"{col_token}_{agg_func}") ) - return self._from_native_dataframe(pc.take(df, keep_idx)) + return self._from_native_frame(pc.take(df, keep_idx)) keep_idx = self.select(*subset).is_unique() return self.filter(keep_idx) def gather_every(self: Self, n: int, offset: int = 0) -> Self: - return self._from_native_dataframe(self._native_dataframe[offset::n]) + return self._from_native_frame(self._native_frame[offset::n]) + + def to_arrow(self: Self) -> Any: + return self._native_frame diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py index 7e7a0d200..ca9293b8b 100644 --- a/narwhals/_arrow/expr.py +++ b/narwhals/_arrow/expr.py @@ -56,7 +56,7 @@ def from_column_names( def func(df: ArrowDataFrame) -> list[ArrowSeries]: return [ ArrowSeries( - df._native_dataframe[column_name], + df._native_frame[column_name], name=column_name, backend_version=df._backend_version, ) @@ -145,6 +145,12 @@ def __truediv__(self, other: ArrowExpr | Any) -> Self: def __rtruediv__(self, other: ArrowExpr | Any) -> Self: return reuse_series_implementation(self, "__rtruediv__", other) + def __mod__(self, other: ArrowExpr | Any) -> Self: + return reuse_series_implementation(self, "__mod__", other) + + def __rmod__(self, other: ArrowExpr | Any) -> Self: + return reuse_series_implementation(self, "__rmod__", other) + def __invert__(self) -> Self: return reuse_series_implementation(self, "__invert__") @@ -152,9 +158,7 @@ def len(self) -> Self: return reuse_series_implementation(self, "len", returns_scalar=True) def filter(self, *predicates: Any) -> Self: - from narwhals._arrow.namespace import ArrowNamespace - - plx = ArrowNamespace(backend_version=self._backend_version) + plx = self.__narwhals_namespace__() expr = plx.all_horizontal(*predicates) return reuse_series_implementation(self, "filter", other=expr) @@ -182,6 +186,9 @@ def diff(self) -> Self: def cum_sum(self) -> Self: return reuse_series_implementation(self, "cum_sum") + def round(self, decimals: int) -> Self: + return reuse_series_implementation(self, "round", decimals) + def any(self) -> Self: return reuse_series_implementation(self, "any", returns_scalar=True) @@ -200,6 +207,9 @@ def sum(self) -> Self: def drop_nulls(self) -> Self: return reuse_series_implementation(self, "drop_nulls") + def shift(self, n: int) -> Self: + return reuse_series_implementation(self, "shift", n) + def alias(self, name: str) -> Self: # Define this one manually, so that we can # override `output_names` and not increase depth @@ -281,6 +291,35 @@ def quantile( def gather_every(self: Self, n: int, offset: int = 0) -> Self: return reuse_series_implementation(self, "gather_every", n=n, offset=offset) + def clip( + self: Self, lower_bound: Any | None = None, upper_bound: Any | None = None + ) -> Self: + return reuse_series_implementation( + self, "clip", lower_bound=lower_bound, upper_bound=upper_bound + ) + + def over(self: Self, keys: list[str]) -> Self: + def func(df: ArrowDataFrame) -> list[ArrowSeries]: + if self._output_names is None: + msg = ( + "Anonymous expressions are not supported in over.\n" + "Instead of `nw.all()`, try using a named expression, such as " + "`nw.col('a', 'b')`\n" + ) + raise ValueError(msg) + tmp = df.group_by(*keys).agg(self) + tmp = df.select(*keys).join(tmp, how="left", left_on=keys, right_on=keys) + return [tmp[name] for name in self._output_names] + + return self.__class__( + func, + depth=self._depth + 1, + function_name=self._function_name + "->over", + root_names=self._root_names, + output_names=self._output_names, + backend_version=self._backend_version, + ) + @property def dt(self: Self) -> ArrowExprDateTimeNamespace: return ArrowExprDateTimeNamespace(self) @@ -319,6 +358,9 @@ def to_string(self: Self, format: str) -> ArrowExpr: # noqa: A002 self._expr, "dt", "to_string", format ) + def date(self: Self) -> ArrowExpr: + return reuse_series_namespace_implementation(self._expr, "dt", "date") + def year(self: Self) -> ArrowExpr: return reuse_series_namespace_implementation(self._expr, "dt", "year") @@ -375,6 +417,48 @@ class ArrowExprStringNamespace: def __init__(self, expr: ArrowExpr) -> None: self._expr = expr + def replace( + self, + pattern: str, + value: str, + *, + literal: bool = False, + n: int = 1, + ) -> ArrowExpr: + return reuse_series_namespace_implementation( + self._expr, + "str", + "replace", + pattern, + value, + literal=literal, + n=n, + ) + + def replace_all( + self, + pattern: str, + value: str, + *, + literal: bool = False, + ) -> ArrowExpr: + return reuse_series_namespace_implementation( + self._expr, + "str", + "replace_all", + pattern, + value, + literal=literal, + ) + + def strip_chars(self, characters: str | None = None) -> ArrowExpr: + return reuse_series_namespace_implementation( + self._expr, + "str", + "strip_chars", + characters, + ) + def starts_with(self, prefix: str) -> ArrowExpr: return reuse_series_namespace_implementation( self._expr, diff --git a/narwhals/_arrow/group_by.py b/narwhals/_arrow/group_by.py index 4da0356a5..27c7ff368 100644 --- a/narwhals/_arrow/group_by.py +++ b/narwhals/_arrow/group_by.py @@ -4,11 +4,10 @@ from typing import TYPE_CHECKING from typing import Any from typing import Callable +from typing import Iterator from narwhals._expression_parsing import is_simple_aggregation from narwhals._expression_parsing import parse_into_exprs -from narwhals.dependencies import get_pyarrow -from narwhals.dependencies import get_pyarrow_compute from narwhals.utils import remove_prefix if TYPE_CHECKING: @@ -19,10 +18,11 @@ class ArrowGroupBy: def __init__(self, df: ArrowDataFrame, keys: list[str]) -> None: - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + self._df = df self._keys = list(keys) - self._grouped = pa.TableGroupBy(self._df._native_dataframe, list(self._keys)) + self._grouped = pa.TableGroupBy(self._df._native_frame, list(self._keys)) def agg( self, @@ -50,7 +50,24 @@ def agg( exprs, self._keys, output_names, - self._df._from_native_dataframe, + self._df._from_native_frame, + ) + + def __iter__(self) -> Iterator[tuple[Any, ArrowDataFrame]]: + key_values = ( + self._df.select(*self._keys) + .unique(subset=self._keys, keep="first") + .iter_rows() + ) + nw_namespace = self._df.__narwhals_namespace__() + yield from ( + ( + key_value, + self._df.filter( + *[nw_namespace.col(k) == v for k, v in zip(self._keys, key_value)] + ), + ) + for key_value in key_values ) @@ -61,7 +78,8 @@ def agg_arrow( output_names: list[str], from_dataframe: Callable[[Any], ArrowDataFrame], ) -> ArrowDataFrame: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + all_simple_aggs = True for expr in exprs: if not is_simple_aggregation(expr): diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py index fbb285b50..ed14dd335 100644 --- a/narwhals/_arrow/namespace.py +++ b/narwhals/_arrow/namespace.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING from typing import Any from typing import Iterable +from typing import cast from narwhals import dtypes from narwhals._arrow.dataframe import ArrowDataFrame @@ -13,7 +14,6 @@ from narwhals._arrow.utils import horizontal_concat from narwhals._arrow.utils import vertical_concat from narwhals._expression_parsing import parse_into_exprs -from narwhals.dependencies import get_pyarrow from narwhals.utils import Implementation if TYPE_CHECKING: @@ -87,9 +87,10 @@ def _create_series_from_scalar(self, value: Any, series: ArrowSeries) -> ArrowSe ) def _create_compliant_series(self, value: Any) -> ArrowSeries: + import pyarrow as pa # ignore-banned-import() + from narwhals._arrow.series import ArrowSeries - pa = get_pyarrow() return ArrowSeries( native_series=pa.chunked_array([value]), name="", @@ -114,7 +115,7 @@ def len(self) -> ArrowExpr: return ArrowExpr( # pragma: no cover lambda df: [ ArrowSeries._from_iterable( - [len(df._native_dataframe)], + [len(df._native_frame)], name="len", backend_version=self._backend_version, ) @@ -133,7 +134,7 @@ def all(self) -> ArrowExpr: return ArrowExpr( lambda df: [ ArrowSeries( - df._native_dataframe[column_name], + df._native_frame[column_name], name=column_name, backend_version=df._backend_version, ) @@ -167,33 +168,33 @@ def _lit_arrow_series(_: ArrowDataFrame) -> ArrowSeries: ) def all_horizontal(self, *exprs: IntoArrowExpr) -> ArrowExpr: - return reduce( - lambda x, y: x & y, - parse_into_exprs(*exprs, namespace=self), - ) + return reduce(lambda x, y: x & y, parse_into_exprs(*exprs, namespace=self)) def any_horizontal(self, *exprs: IntoArrowExpr) -> ArrowExpr: - return reduce( - lambda x, y: x | y, - parse_into_exprs(*exprs, namespace=self), - ) + return reduce(lambda x, y: x | y, parse_into_exprs(*exprs, namespace=self)) def sum_horizontal(self, *exprs: IntoArrowExpr) -> ArrowExpr: return reduce( lambda x, y: x + y, - parse_into_exprs( - *exprs, - namespace=self, - ), + [expr.fill_null(0) for expr in parse_into_exprs(*exprs, namespace=self)], ) + def mean_horizontal(self, *exprs: IntoArrowExpr) -> IntoArrowExpr: + arrow_exprs = parse_into_exprs(*exprs, namespace=self) + total = reduce(lambda x, y: x + y, (e.fill_null(0.0) for e in arrow_exprs)) + n_non_zero = reduce( + lambda x, y: x + y, + ((1 - e.is_null().cast(self.Int64())) for e in arrow_exprs), + ) + return total / n_non_zero + def concat( self, items: Iterable[ArrowDataFrame], *, how: str = "vertical", ) -> ArrowDataFrame: - dfs: list[Any] = [item._native_dataframe for item in items] + dfs: list[Any] = [item._native_frame for item in items] if how == "horizontal": return ArrowDataFrame( @@ -234,3 +235,121 @@ def min(self, *column_names: str) -> ArrowExpr: @property def selectors(self) -> ArrowSelectorNamespace: return ArrowSelectorNamespace(backend_version=self._backend_version) + + def when( + self, + *predicates: IntoArrowExpr, + ) -> ArrowWhen: + plx = self.__class__(backend_version=self._backend_version) + if predicates: + condition = plx.all_horizontal(*predicates) + else: + msg = "at least one predicate needs to be provided" + raise TypeError(msg) + + return ArrowWhen(condition, self._backend_version) + + +class ArrowWhen: + def __init__( + self, + condition: ArrowExpr, + backend_version: tuple[int, ...], + then_value: Any = None, + otherwise_value: Any = None, + ) -> None: + self._backend_version = backend_version + self._condition = condition + self._then_value = then_value + self._otherwise_value = otherwise_value + + def __call__(self, df: ArrowDataFrame) -> list[ArrowSeries]: + import pyarrow as pa # ignore-banned-import + import pyarrow.compute as pc # ignore-banned-import + + from narwhals._arrow.namespace import ArrowNamespace + from narwhals._expression_parsing import parse_into_expr + + plx = ArrowNamespace(backend_version=self._backend_version) + + condition = parse_into_expr(self._condition, namespace=plx)._call(df)[0] # type: ignore[arg-type] + try: + value_series = parse_into_expr(self._then_value, namespace=plx)._call(df)[0] # type: ignore[arg-type] + except TypeError: + # `self._otherwise_value` is a scalar and can't be converted to an expression + value_series = condition.__class__._from_iterable( # type: ignore[call-arg] + [self._then_value] * len(condition), + name="literal", + backend_version=self._backend_version, + ) + value_series = cast(ArrowSeries, value_series) + + value_series_native = value_series._native_series + condition_native = condition._native_series.combine_chunks() + + if self._otherwise_value is None: + otherwise_native = pa.array( + [None] * len(condition_native), type=value_series_native.type + ) + return [ + value_series._from_native_series( + pc.if_else(condition_native, value_series_native, otherwise_native) + ) + ] + try: + otherwise_series = parse_into_expr( + self._otherwise_value, namespace=plx + )._call(df)[0] # type: ignore[arg-type] + except TypeError: + # `self._otherwise_value` is a scalar and can't be converted to an expression + return [ + value_series._from_native_series( + pc.if_else( + condition_native, value_series_native, self._otherwise_value + ) + ) + ] + else: + otherwise_series = cast(ArrowSeries, otherwise_series) + condition = cast(ArrowSeries, condition) + return [value_series.zip_with(condition, otherwise_series)] + + def then(self, value: ArrowExpr | ArrowSeries | Any) -> ArrowThen: + self._then_value = value + + return ArrowThen( + self, + depth=0, + function_name="whenthen", + root_names=None, + output_names=None, + backend_version=self._backend_version, + ) + + +class ArrowThen(ArrowExpr): + def __init__( + self, + call: ArrowWhen, + *, + depth: int, + function_name: str, + root_names: list[str] | None, + output_names: list[str] | None, + backend_version: tuple[int, ...], + ) -> None: + self._backend_version = backend_version + + self._call = call + self._depth = depth + self._function_name = function_name + self._root_names = root_names + self._output_names = output_names + + def otherwise(self, value: ArrowExpr | ArrowSeries | Any) -> ArrowExpr: + # type ignore because we are setting the `_call` attribute to a + # callable object of type `PandasWhen`, base class has the attribute as + # only a `Callable` + self._call._otherwise_value = value # type: ignore[attr-defined] + self._function_name = "whenotherwise" + return self diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py index 479108672..1e9e4a08c 100644 --- a/narwhals/_arrow/series.py +++ b/narwhals/_arrow/series.py @@ -9,13 +9,11 @@ from narwhals._arrow.utils import cast_for_truediv from narwhals._arrow.utils import floordiv_compat -from narwhals._arrow.utils import reverse_translate_dtype +from narwhals._arrow.utils import narwhals_to_native_dtype from narwhals._arrow.utils import translate_dtype from narwhals._arrow.utils import validate_column_comparand -from narwhals.dependencies import get_numpy from narwhals.dependencies import get_pandas from narwhals.dependencies import get_pyarrow -from narwhals.dependencies import get_pyarrow_compute from narwhals.utils import Implementation from narwhals.utils import generate_unique_token @@ -23,7 +21,6 @@ from typing_extensions import Self from narwhals._arrow.dataframe import ArrowDataFrame - from narwhals._arrow.namespace import ArrowNamespace from narwhals.dtypes import DType @@ -37,7 +34,8 @@ def __init__( self._backend_version = backend_version def _from_native_series(self, series: Any) -> Self: - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + if isinstance(series, pa.Array): series = pa.chunked_array([series]) return self.__class__( @@ -54,7 +52,8 @@ def _from_iterable( *, backend_version: tuple[int, ...], ) -> Self: - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + return cls( pa.chunked_array([data]), name=name, @@ -65,67 +64,78 @@ def __len__(self) -> int: return len(self._native_series) def __eq__(self, other: object) -> Self: # type: ignore[override] - pc = get_pyarrow_compute() + import pyarrow.compute as pc + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.equal(ser, other)) def __ne__(self, other: object) -> Self: # type: ignore[override] - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.not_equal(ser, other)) def __ge__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.greater_equal(ser, other)) def __gt__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.greater(ser, other)) def __le__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.less_equal(ser, other)) def __lt__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.less(ser, other)) def __and__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.and_kleene(ser, other)) def __rand__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.and_kleene(other, ser)) def __or__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.or_kleene(ser, other)) def __ror__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.or_kleene(other, ser)) def __add__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + other = validate_column_comparand(other) return self._from_native_series(pc.add(self._native_series, other)) @@ -133,7 +143,8 @@ def __radd__(self, other: Any) -> Self: return self + other # type: ignore[no-any-return] def __sub__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + other = validate_column_comparand(other) return self._from_native_series(pc.subtract(self._native_series, other)) @@ -141,7 +152,8 @@ def __rsub__(self, other: Any) -> Self: return (self - other) * (-1) # type: ignore[no-any-return] def __mul__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + other = validate_column_comparand(other) return self._from_native_series(pc.multiply(self._native_series, other)) @@ -149,13 +161,15 @@ def __rmul__(self, other: Any) -> Self: return self * other # type: ignore[no-any-return] def __pow__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.power(ser, other)) def __rpow__(self, other: Any) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) return self._from_native_series(pc.power(other, ser)) @@ -171,8 +185,9 @@ def __rfloordiv__(self, other: Any) -> Self: return self._from_native_series(floordiv_compat(other, ser)) def __truediv__(self, other: Any) -> Self: - pa = get_pyarrow() - pc = get_pyarrow_compute() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) if not isinstance(other, (pa.Array, pa.ChunkedArray)): @@ -181,8 +196,9 @@ def __truediv__(self, other: Any) -> Self: return self._from_native_series(pc.divide(*cast_for_truediv(ser, other))) def __rtruediv__(self, other: Any) -> Self: - pa = get_pyarrow() - pc = get_pyarrow_compute() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series other = validate_column_comparand(other) if not isinstance(other, (pa.Array, pa.ChunkedArray)): @@ -190,58 +206,94 @@ def __rtruediv__(self, other: Any) -> Self: other = pa.scalar(other) return self._from_native_series(pc.divide(*cast_for_truediv(other, ser))) + def __mod__(self, other: Any) -> Self: + import pyarrow.compute as pc # ignore-banned-import() + + ser = self._native_series + other = validate_column_comparand(other) + floor_div = (self // other)._native_series + res = pc.subtract(ser, pc.multiply(floor_div, other)) + return self._from_native_series(res) + + def __rmod__(self, other: Any) -> Self: + import pyarrow.compute as pc # ignore-banned-import() + + ser = self._native_series + other = validate_column_comparand(other) + floor_div = (other // self)._native_series + res = pc.subtract(other, pc.multiply(floor_div, ser)) + return self._from_native_series(res) + def __invert__(self) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._from_native_series(pc.invert(self._native_series)) def len(self) -> int: return len(self._native_series) def filter(self, other: Any) -> Self: - other = validate_column_comparand(other) + if not (isinstance(other, list) and all(isinstance(x, bool) for x in other)): + other = validate_column_comparand(other) return self._from_native_series(self._native_series.filter(other)) def mean(self) -> int: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return pc.mean(self._native_series) # type: ignore[no-any-return] def min(self) -> int: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return pc.min(self._native_series) # type: ignore[no-any-return] def max(self) -> int: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return pc.max(self._native_series) # type: ignore[no-any-return] def sum(self) -> int: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return pc.sum(self._native_series) # type: ignore[no-any-return] def drop_nulls(self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._from_native_series(pc.drop_null(self._native_series)) + def shift(self, n: int) -> Self: + import pyarrow as pa # ignore-banned-import() + + ca = self._native_series + + if n > 0: + result = pa.concat_arrays([pa.nulls(n, ca.type), *ca[:-n].chunks]) + elif n < 0: + result = pa.concat_arrays([*ca[-n:].chunks, pa.nulls(-n, ca.type)]) + else: + result = ca + return self._from_native_series(result) + def std(self, ddof: int = 1) -> int: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return pc.stddev(self._native_series, ddof=ddof) # type: ignore[no-any-return] def count(self) -> int: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return pc.count(self._native_series) # type: ignore[no-any-return] def n_unique(self) -> int: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + unique_values = pc.unique(self._native_series) return pc.count(unique_values, mode="all") # type: ignore[no-any-return] def __native_namespace__(self) -> Any: # pragma: no cover return get_pyarrow() - def __narwhals_namespace__(self) -> ArrowNamespace: - from narwhals._arrow.namespace import ArrowNamespace - - return ArrowNamespace(backend_version=self._backend_version) - @property def name(self) -> str: return self._name @@ -281,29 +333,42 @@ def dtype(self) -> DType: return translate_dtype(self._native_series.type) def abs(self) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._from_native_series(pc.abs(self._native_series)) def cum_sum(self) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._from_native_series(pc.cumulative_sum(self._native_series)) + def round(self, decimals: int) -> Self: + import pyarrow.compute as pc # ignore-banned-import() + + return self._from_native_series( + pc.round(self._native_series, decimals, round_mode="half_towards_infinity") + ) + def diff(self) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._from_native_series( pc.pairwise_diff(self._native_series.combine_chunks()) ) def any(self) -> bool: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return pc.any(self._native_series) # type: ignore[no-any-return] def all(self) -> bool: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return pc.all(self._native_series) # type: ignore[no-any-return] def is_between(self, lower_bound: Any, upper_bound: Any, closed: str = "both") -> Any: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series if closed == "left": ge = pc.greater_equal(ser, lower_bound) @@ -333,9 +398,10 @@ def is_null(self) -> Self: return self._from_native_series(ser.is_null()) def cast(self, dtype: DType) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series - dtype = reverse_translate_dtype(dtype) + dtype = narwhals_to_native_dtype(dtype) return self._from_native_series(pc.cast(ser, dtype)) def null_count(self: Self) -> int: @@ -358,14 +424,16 @@ def tail(self, n: int) -> Self: return self._from_native_series(ser.slice(abs(n))) def is_in(self, other: Any) -> Self: - pc = get_pyarrow_compute() - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + value_set = pa.array(other) ser = self._native_series return self._from_native_series(pc.is_in(ser, value_set=value_set)) def arg_true(self) -> Self: - np = get_numpy() + import numpy as np # ignore-banned-import + ser = self._native_series res = np.flatnonzero(ser) return self._from_iterable( @@ -392,10 +460,10 @@ def value_counts( normalize: bool = False, ) -> ArrowDataFrame: """Parallel is unused, exists for compatibility""" - from narwhals._arrow.dataframe import ArrowDataFrame + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() - pc = get_pyarrow_compute() - pa = get_pyarrow() + from narwhals._arrow.dataframe import ArrowDataFrame index_name_ = "index" if self._name is None else self._name value_name_ = name or ("proportion" if normalize else "count") @@ -420,13 +488,14 @@ def value_counts( ) def zip_with(self: Self, mask: Self, other: Self) -> Self: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + mask = mask._native_series.combine_chunks() return self._from_native_series( - pc.replace_with_mask( - self._native_series.combine_chunks(), - pc.invert(mask._native_series.combine_chunks()), - other._native_series.combine_chunks(), + pc.if_else( + mask, + self._native_series, + other._native_series, ) ) @@ -437,8 +506,9 @@ def sample( *, with_replacement: bool = False, ) -> Self: - np = get_numpy() - pc = get_pyarrow_compute() + import numpy as np # ignore-banned-import + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series num_rows = len(self) @@ -450,17 +520,19 @@ def sample( return self._from_native_series(pc.take(ser, mask)) def fill_null(self: Self, value: Any) -> Self: - pa = get_pyarrow() - pc = get_pyarrow_compute() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series dtype = ser.type return self._from_native_series(pc.fill_null(ser, pa.scalar(value, dtype))) def to_frame(self: Self) -> ArrowDataFrame: + import pyarrow as pa # ignore-banned-import() + from narwhals._arrow.dataframe import ArrowDataFrame - pa = get_pyarrow() df = pa.Table.from_arrays([self._native_series], names=[self.name]) return ArrowDataFrame(df, backend_version=self._backend_version) @@ -475,9 +547,9 @@ def is_unique(self: Self) -> ArrowSeries: return self.to_frame().is_unique().alias(self.name) def is_first_distinct(self: Self) -> Self: - np = get_numpy() - pa = get_pyarrow() - pc = get_pyarrow_compute() + import numpy as np # ignore-banned-import + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() row_number = pa.array(np.arange(len(self))) col_token = generate_unique_token(n_bytes=8, columns=[self.name]) @@ -492,9 +564,9 @@ def is_first_distinct(self: Self) -> Self: return self._from_native_series(pc.is_in(row_number, first_distinct_index)) def is_last_distinct(self: Self) -> Self: - np = get_numpy() - pa = get_pyarrow() - pc = get_pyarrow_compute() + import numpy as np # ignore-banned-import + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() row_number = pa.array(np.arange(len(self))) col_token = generate_unique_token(n_bytes=8, columns=[self.name]) @@ -512,7 +584,8 @@ def is_sorted(self: Self, *, descending: bool = False) -> bool: if not isinstance(descending, bool): msg = f"argument 'descending' should be boolean, found {type(descending)}" raise TypeError(msg) - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + ser = self._native_series if descending: return pc.all(pc.greater_equal(ser[:-1], ser[1:])) # type: ignore[no-any-return] @@ -520,13 +593,15 @@ def is_sorted(self: Self, *, descending: bool = False) -> bool: return pc.all(pc.less_equal(ser[:-1], ser[1:])) # type: ignore[no-any-return] def unique(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._from_native_series(pc.unique(self._native_series)) def sort( self: Self, *, descending: bool = False, nulls_last: bool = False ) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + series = self._native_series order = "descending" if descending else "ascending" null_placement = "at_end" if nulls_last else "at_start" @@ -539,10 +614,10 @@ def sort( def to_dummies( self: Self, *, separator: str = "_", drop_first: bool = False ) -> ArrowDataFrame: - from narwhals._arrow.dataframe import ArrowDataFrame + import numpy as np # ignore-banned-import + import pyarrow as pa # ignore-banned-import() - np = get_numpy() - pa = get_pyarrow() + from narwhals._arrow.dataframe import ArrowDataFrame series = self._native_series da = series.dictionary_encode().combine_chunks() @@ -561,7 +636,8 @@ def quantile( quantile: float, interpolation: Literal["nearest", "higher", "lower", "midpoint", "linear"], ) -> Any: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return pc.quantile(self._native_series, q=quantile, interpolation=interpolation)[ 0 ] @@ -569,6 +645,21 @@ def quantile( def gather_every(self: Self, n: int, offset: int = 0) -> Self: return self._from_native_series(self._native_series[offset::n]) + def clip( + self: Self, lower_bound: Any | None = None, upper_bound: Any | None = None + ) -> Self: + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + + arr = self._native_series + arr = pc.max_element_wise(arr, pa.scalar(lower_bound, type=arr.type)) + arr = pc.min_element_wise(arr, pa.scalar(upper_bound, type=arr.type)) + + return self._from_native_series(arr) + + def to_arrow(self: Self) -> Any: + return self._native_series.combine_chunks() + @property def shape(self) -> tuple[int]: return (len(self._native_series),) @@ -591,7 +682,8 @@ def __init__(self: Self, series: ArrowSeries) -> None: self._arrow_series = series def to_string(self: Self, format: str) -> ArrowSeries: # noqa: A002 - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + # PyArrow differs from other libraries in that %S also prints out # the fractional part of the second...:'( # https://arrow.apache.org/docs/python/generated/pyarrow.compute.strftime.html @@ -600,58 +692,73 @@ def to_string(self: Self, format: str) -> ArrowSeries: # noqa: A002 pc.strftime(self._arrow_series._native_series, format) ) + def date(self: Self) -> ArrowSeries: + import pyarrow as pa # ignore-banned-import() + + return self._arrow_series._from_native_series( + self._arrow_series._native_series.cast(pa.date64()) + ) + def year(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.year(self._arrow_series._native_series) ) def month(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.month(self._arrow_series._native_series) ) def day(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.day(self._arrow_series._native_series) ) def hour(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.hour(self._arrow_series._native_series) ) def minute(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.minute(self._arrow_series._native_series) ) def second(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.second(self._arrow_series._native_series) ) def millisecond(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.millisecond(self._arrow_series._native_series) ) def microsecond(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + arr = self._arrow_series._native_series result = pc.add(pc.multiply(pc.millisecond(arr), 1000), pc.microsecond(arr)) return self._arrow_series._from_native_series(result) def nanosecond(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + arr = self._arrow_series._native_series result = pc.add( pc.multiply(self.microsecond()._native_series, 1000), pc.nanosecond(arr) @@ -659,14 +766,16 @@ def nanosecond(self: Self) -> ArrowSeries: return self._arrow_series._from_native_series(result) def ordinal_day(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.day_of_year(self._arrow_series._native_series) ) def total_minutes(self: Self) -> ArrowSeries: - pa = get_pyarrow() - pc = get_pyarrow_compute() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + arr = self._arrow_series._native_series unit = arr.type.unit @@ -683,8 +792,9 @@ def total_minutes(self: Self) -> ArrowSeries: ) def total_seconds(self: Self) -> ArrowSeries: - pa = get_pyarrow() - pc = get_pyarrow_compute() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + arr = self._arrow_series._native_series unit = arr.type.unit @@ -701,8 +811,9 @@ def total_seconds(self: Self) -> ArrowSeries: ) def total_milliseconds(self: Self) -> ArrowSeries: - pa = get_pyarrow() - pc = get_pyarrow_compute() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + arr = self._arrow_series._native_series unit = arr.type.unit @@ -725,8 +836,9 @@ def total_milliseconds(self: Self) -> ArrowSeries: ) def total_microseconds(self: Self) -> ArrowSeries: - pa = get_pyarrow() - pc = get_pyarrow_compute() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + arr = self._arrow_series._native_series unit = arr.type.unit @@ -748,8 +860,9 @@ def total_microseconds(self: Self) -> ArrowSeries: ) def total_nanoseconds(self: Self) -> ArrowSeries: - pa = get_pyarrow() - pc = get_pyarrow_compute() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() + arr = self._arrow_series._native_series unit = arr.type.unit @@ -772,7 +885,8 @@ def __init__(self, series: ArrowSeries) -> None: self._arrow_series = series def get_categories(self) -> ArrowSeries: - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + ca = self._arrow_series._native_series # TODO(Unassigned): this looks potentially expensive - is there no better way? out = pa.chunked_array( @@ -785,27 +899,62 @@ class ArrowSeriesStringNamespace: def __init__(self: Self, series: ArrowSeries) -> None: self._arrow_series = series + def replace( + self, pattern: str, value: str, *, literal: bool = False, n: int = 1 + ) -> ArrowSeries: + import pyarrow.compute as pc # ignore-banned-import() + + method = "replace_substring" if literal else "replace_substring_regex" + return self._arrow_series._from_native_series( + getattr(pc, method)( + self._arrow_series._native_series, + pattern=pattern, + replacement=value, + max_replacements=n, + ) + ) + + def replace_all( + self, pattern: str, value: str, *, literal: bool = False + ) -> ArrowSeries: + return self.replace(pattern, value, literal=literal, n=-1) + + def strip_chars(self: Self, characters: str | None = None) -> ArrowSeries: + import pyarrow.compute as pc # ignore-banned-import() + + whitespace = " \t\n\r\v\f" + return self._arrow_series._from_native_series( + pc.utf8_trim( + self._arrow_series._native_series, + characters or whitespace, + ) + ) + def starts_with(self: Self, prefix: str) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.equal(self.slice(0, len(prefix))._native_series, prefix) ) def ends_with(self: Self, suffix: str) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.equal(self.slice(-len(suffix))._native_series, suffix) ) def contains(self: Self, pattern: str, *, literal: bool = False) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + check_func = pc.match_substring if literal else pc.match_substring_regex return self._arrow_series._from_native_series( check_func(self._arrow_series._native_series, pattern) ) def slice(self: Self, offset: int, length: int | None = None) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + stop = offset + length if length else None return self._arrow_series._from_native_series( pc.utf8_slice_codeunits( @@ -814,19 +963,22 @@ def slice(self: Self, offset: int, length: int | None = None) -> ArrowSeries: ) def to_datetime(self: Self, format: str | None = None) -> ArrowSeries: # noqa: A002 - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.strptime(self._arrow_series._native_series, format=format, unit="us") ) def to_uppercase(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.utf8_upper(self._arrow_series._native_series), ) def to_lowercase(self: Self) -> ArrowSeries: - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + return self._arrow_series._from_native_series( pc.utf8_lower(self._arrow_series._native_series), ) diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py index f98eb2e5b..6f7517aeb 100644 --- a/narwhals/_arrow/utils.py +++ b/narwhals/_arrow/utils.py @@ -1,15 +1,18 @@ from __future__ import annotations +from typing import TYPE_CHECKING from typing import Any from narwhals import dtypes -from narwhals.dependencies import get_pyarrow -from narwhals.dependencies import get_pyarrow_compute from narwhals.utils import isinstance_or_issubclass +if TYPE_CHECKING: + from narwhals._arrow.series import ArrowSeries + def translate_dtype(dtype: Any) -> dtypes.DType: - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + if pa.types.is_int64(dtype): return dtypes.Int64() if pa.types.is_int32(dtype): @@ -51,10 +54,10 @@ def translate_dtype(dtype: Any) -> dtypes.DType: raise AssertionError -def reverse_translate_dtype(dtype: dtypes.DType | type[dtypes.DType]) -> Any: - from narwhals import dtypes +def narwhals_to_native_dtype(dtype: dtypes.DType | type[dtypes.DType]) -> Any: + import pyarrow as pa # ignore-banned-import() - pa = get_pyarrow() + from narwhals import dtypes if isinstance_or_issubclass(dtype, dtypes.Float64): return pa.float64() @@ -139,7 +142,8 @@ def validate_dataframe_comparand( return NotImplemented if isinstance(other, ArrowSeries): if len(other) == 1: - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + value = other.item() if backend_version < (13,) and hasattr(value, "as_py"): # pragma: no cover value = value.as_py() @@ -155,7 +159,8 @@ def horizontal_concat(dfs: list[Any]) -> Any: Should be in namespace. """ - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + if not dfs: msg = "No dataframes to concatenate" # pragma: no cover raise AssertionError(msg) @@ -187,15 +192,16 @@ def vertical_concat(dfs: list[Any]) -> Any: msg = "unable to vstack, column names don't match" raise TypeError(msg) - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + return pa.concat_tables(dfs).combine_chunks() def floordiv_compat(left: Any, right: Any) -> Any: # The following lines are adapted from pandas' pyarrow implementation. # Ref: https://github.com/pandas-dev/pandas/blob/262fcfbffcee5c3116e86a951d8b693f90411e68/pandas/core/arrays/arrow/array.py#L124-L154 - pc = get_pyarrow_compute() - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() if isinstance(left, (int, float)): left = pa.scalar(left) @@ -233,8 +239,8 @@ def floordiv_compat(left: Any, right: Any) -> Any: def cast_for_truediv(arrow_array: Any, pa_object: Any) -> tuple[Any, Any]: # Lifted from: # https://github.com/pandas-dev/pandas/blob/262fcfbffcee5c3116e86a951d8b693f90411e68/pandas/core/arrays/arrow/array.py#L108-L122 - pc = get_pyarrow_compute() - pa = get_pyarrow() + import pyarrow as pa # ignore-banned-import() + import pyarrow.compute as pc # ignore-banned-import() # Ensure int / int -> float mirroring Python/Numpy behavior # as pc.divide_checked(int, int) -> int @@ -246,3 +252,27 @@ def cast_for_truediv(arrow_array: Any, pa_object: Any) -> tuple[Any, Any]: ) return arrow_array, pa_object + + +def broadcast_series(series: list[ArrowSeries]) -> list[Any]: + lengths = [len(s) for s in series] + max_length = max(lengths) + fast_path = all(_len == max_length for _len in lengths) + + if fast_path: + return [s._native_series for s in series] + + import pyarrow as pa # ignore-banned-import() + + reshaped = [] + for s, length in zip(series, lengths): + s_native = s._native_series + if max_length > 1 and length == 1: + value = s_native[0] + if s._backend_version < (13,) and hasattr(value, "as_py"): # pragma: no cover + value = value.as_py() + reshaped.append(pa.array([value] * max_length, type=s_native.type)) + else: + reshaped.append(s_native) + + return reshaped diff --git a/narwhals/_dask/dataframe.py b/narwhals/_dask/dataframe.py index 0861c122f..9774d6c8e 100644 --- a/narwhals/_dask/dataframe.py +++ b/narwhals/_dask/dataframe.py @@ -2,12 +2,19 @@ from typing import TYPE_CHECKING from typing import Any +from typing import Iterable +from typing import Literal +from typing import Sequence +from narwhals._dask.utils import add_row_index from narwhals._dask.utils import parse_exprs_and_named_exprs -from narwhals._expression_parsing import evaluate_into_exprs +from narwhals._pandas_like.utils import translate_dtype from narwhals.dependencies import get_dask_dataframe from narwhals.dependencies import get_pandas from narwhals.utils import Implementation +from narwhals.utils import flatten +from narwhals.utils import generate_unique_token +from narwhals.utils import parse_columns_to_drop from narwhals.utils import parse_version if TYPE_CHECKING: @@ -16,14 +23,16 @@ from narwhals._dask.expr import DaskExpr from narwhals._dask.namespace import DaskNamespace from narwhals._dask.typing import IntoDaskExpr + from narwhals.dtypes import DType class DaskLazyFrame: def __init__( self, native_dataframe: Any, *, backend_version: tuple[int, ...] ) -> None: - self._native_dataframe = native_dataframe + self._native_frame = native_dataframe self._backend_version = backend_version + self._implementation = Implementation.DASK def __native_namespace__(self) -> Any: # pragma: no cover return get_dask_dataframe() @@ -36,19 +45,19 @@ def __narwhals_namespace__(self) -> DaskNamespace: def __narwhals_lazyframe__(self) -> Self: return self - def _from_native_dataframe(self, df: Any) -> Self: + def _from_native_frame(self, df: Any) -> Self: return self.__class__(df, backend_version=self._backend_version) def with_columns(self, *exprs: DaskExpr, **named_exprs: DaskExpr) -> Self: - df = self._native_dataframe + df = self._native_frame new_series = parse_exprs_and_named_exprs(self, *exprs, **named_exprs) df = df.assign(**new_series) - return self._from_native_dataframe(df) + return self._from_native_frame(df) def collect(self) -> Any: from narwhals._pandas_like.dataframe import PandasLikeDataFrame - result = self._native_dataframe.compute() + result = self._native_frame.compute() return PandasLikeDataFrame( result, implementation=Implementation.PANDAS, @@ -57,7 +66,26 @@ def collect(self) -> Any: @property def columns(self) -> list[str]: - return self._native_dataframe.columns.tolist() # type: ignore[no-any-return] + return self._native_frame.columns.tolist() # type: ignore[no-any-return] + + def filter( + self, + *predicates: DaskExpr, + ) -> Self: + if ( + len(predicates) == 1 + and isinstance(predicates[0], list) + and all(isinstance(x, bool) for x in predicates[0]) + ): + mask = predicates[0] + else: + from narwhals._dask.namespace import DaskNamespace + + plx = DaskNamespace(backend_version=self._backend_version) + expr = plx.all_horizontal(*predicates) + # Safety: all_horizontal's expression only returns a single column. + mask = expr._call(self)[0] + return self._from_native_frame(self._native_frame.loc[mask]) def lazy(self) -> Self: return self @@ -67,16 +95,223 @@ def select( *exprs: IntoDaskExpr, **named_exprs: IntoDaskExpr, ) -> Self: - dd = get_dask_dataframe() + import dask.dataframe as dd # ignore-banned-import if exprs and all(isinstance(x, str) for x in exprs) and not named_exprs: # This is a simple slice => fastpath! - return self._from_native_dataframe(self._native_dataframe.loc[:, exprs]) + return self._from_native_frame(self._native_frame.loc[:, exprs]) + + new_series = parse_exprs_and_named_exprs(self, *exprs, **named_exprs) - new_series = evaluate_into_exprs(self, *exprs, **named_exprs) if not new_series: # return empty dataframe, like Polars does - pd = get_pandas() - return self._from_native_dataframe(dd.from_pandas(pd.DataFrame())) - df = dd.concat(new_series, axis=1) - return self._from_native_dataframe(df) + import pandas as pd # ignore-banned-import + + return self._from_native_frame( + dd.from_pandas(pd.DataFrame(), npartitions=self._native_frame.npartitions) + ) + + if all(getattr(expr, "_returns_scalar", False) for expr in exprs) and all( + getattr(val, "_returns_scalar", False) for val in named_exprs.values() + ): + df = dd.concat( + [val.to_series().rename(name) for name, val in new_series.items()], axis=1 + ) + return self._from_native_frame(df) + + df = self._native_frame.assign(**new_series).loc[:, list(new_series.keys())] + return self._from_native_frame(df) + + def drop_nulls(self: Self, subset: str | list[str] | None) -> Self: + if subset is None: + return self._from_native_frame(self._native_frame.dropna()) + subset = [subset] if isinstance(subset, str) else subset + plx = self.__narwhals_namespace__() + return self.filter(~plx.any_horizontal(plx.col(*subset).is_null())) + + @property + def schema(self) -> dict[str, DType]: + return { + col: translate_dtype(self._native_frame.loc[:, col]) + for col in self._native_frame.columns + } + + def collect_schema(self) -> dict[str, DType]: + return self.schema + + def drop(self: Self, columns: list[str], strict: bool) -> Self: # noqa: FBT001 + to_drop = parse_columns_to_drop( + compliant_frame=self, columns=columns, strict=strict + ) + + return self._from_native_frame(self._native_frame.drop(columns=to_drop)) + + def with_row_index(self: Self, name: str) -> Self: + # Implementation is based on the following StackOverflow reply: + # https://stackoverflow.com/questions/60831518/in-dask-how-does-one-add-a-range-of-integersauto-increment-to-a-new-column/60852409#60852409 + return self._from_native_frame(add_row_index(self._native_frame, name)) + + def rename(self: Self, mapping: dict[str, str]) -> Self: + return self._from_native_frame(self._native_frame.rename(columns=mapping)) + + def head(self: Self, n: int) -> Self: + return self._from_native_frame( + self._native_frame.head(n=n, compute=False, npartitions=-1) + ) + + def unique( + self: Self, + subset: str | list[str] | None, + *, + keep: Literal["any", "first", "last", "none"] = "any", + maintain_order: bool = False, + ) -> Self: + """ + NOTE: + The param `maintain_order` is only here for compatibility with the polars API + and has no effect on the output. + """ + subset = flatten(subset) if subset else None + native_frame = self._native_frame + if keep == "none": + subset = subset or self.columns + token = generate_unique_token(n_bytes=8, columns=subset) + ser = native_frame.groupby(subset).size().rename(token) + ser = ser.loc[ser == 1] + unique = ser.reset_index().drop(columns=token) + result = native_frame.merge(unique, on=subset, how="inner") + else: + mapped_keep = {"any": "first"}.get(keep, keep) + result = native_frame.drop_duplicates(subset=subset, keep=mapped_keep) + return self._from_native_frame(result) + + def sort( + self: Self, + by: str | Iterable[str], + *more_by: str, + descending: bool | Sequence[bool] = False, + ) -> Self: + flat_keys = flatten([*flatten([by]), *more_by]) + df = self._native_frame + if isinstance(descending, bool): + ascending: bool | list[bool] = not descending + else: + ascending = [not d for d in descending] + return self._from_native_frame(df.sort_values(flat_keys, ascending=ascending)) + + def join( + self: Self, + other: Self, + *, + how: Literal["left", "inner", "outer", "cross", "anti", "semi"] = "inner", + left_on: str | list[str] | None, + right_on: str | list[str] | None, + ) -> Self: + if isinstance(left_on, str): + left_on = [left_on] + if isinstance(right_on, str): + right_on = [right_on] + + if how == "cross": + key_token = generate_unique_token( + n_bytes=8, columns=[*self.columns, *other.columns] + ) + + return self._from_native_frame( + self._native_frame.assign(**{key_token: 0}) + .merge( + other._native_frame.assign(**{key_token: 0}), + how="inner", + left_on=key_token, + right_on=key_token, + suffixes=("", "_right"), + ) + .drop(columns=key_token), + ) + + if how == "anti": + indicator_token = generate_unique_token( + n_bytes=8, columns=[*self.columns, *other.columns] + ) + + other_native = ( + other._native_frame.loc[:, right_on] + .rename( # rename to avoid creating extra columns in join + columns=dict(zip(right_on, left_on)) # type: ignore[arg-type] + ) + .drop_duplicates() + ) + df = self._native_frame.merge( + other_native, + how="outer", + indicator=indicator_token, + left_on=left_on, + right_on=left_on, + ) + return self._from_native_frame( + df.loc[df[indicator_token] == "left_only"].drop(columns=[indicator_token]) + ) + + if how == "semi": + other_native = ( + other._native_frame.loc[:, right_on] + .rename( # rename to avoid creating extra columns in join + columns=dict(zip(right_on, left_on)) # type: ignore[arg-type] + ) + .drop_duplicates() # avoids potential rows duplication from inner join + ) + return self._from_native_frame( + self._native_frame.merge( + other_native, + how="inner", + left_on=left_on, + right_on=left_on, + ) + ) + + if how == "left": + other_native = other._native_frame + result_native = self._native_frame.merge( + other_native, + how="left", + left_on=left_on, + right_on=right_on, + suffixes=("", "_right"), + ) + extra = [] + for left_key, right_key in zip(left_on, right_on): # type: ignore[arg-type] + if right_key != left_key and right_key not in self.columns: + extra.append(right_key) + elif right_key != left_key: + extra.append(f"{right_key}_right") + return self._from_native_frame(result_native.drop(columns=extra)) + + return self._from_native_frame( + self._native_frame.merge( + other._native_frame, + left_on=left_on, + right_on=right_on, + how=how, + suffixes=("", "_right"), + ), + ) + + def group_by(self, *by: str) -> Any: + from narwhals._dask.group_by import DaskLazyGroupBy + + return DaskLazyGroupBy(self, list(by)) + + def tail(self: Self, n: int) -> Self: + return self._from_native_frame(self._native_frame.tail(n=n, compute=False)) + + def gather_every(self: Self, n: int, offset: int) -> Self: + row_index_token = generate_unique_token(n_bytes=8, columns=self.columns) + pln = self.__narwhals_namespace__() + return ( + self.with_row_index(name=row_index_token) + .filter( + pln.col(row_index_token) >= offset, # type: ignore[operator] + (pln.col(row_index_token) - offset) % n == 0, # type: ignore[arg-type] + ) + .drop([row_index_token], strict=False) + ) diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py index 36033afd9..62aaa85e6 100644 --- a/narwhals/_dask/expr.py +++ b/narwhals/_dask/expr.py @@ -4,17 +4,21 @@ from typing import TYPE_CHECKING from typing import Any from typing import Callable +from typing import Literal +from typing import NoReturn +from narwhals._dask.utils import add_row_index +from narwhals._dask.utils import maybe_evaluate +from narwhals._dask.utils import reverse_translate_dtype from narwhals.dependencies import get_dask -from narwhals.dependencies import get_dask_expr +from narwhals.utils import generate_unique_token if TYPE_CHECKING: from typing_extensions import Self from narwhals._dask.dataframe import DaskLazyFrame from narwhals._dask.namespace import DaskNamespace - -from narwhals._dask.utils import maybe_evaluate + from narwhals.dtypes import DType class DaskExpr: @@ -27,6 +31,9 @@ def __init__( function_name: str, root_names: list[str] | None, output_names: list[str] | None, + # Whether the expression is a length-1 Series resulting from + # a reduction, such as `nw.col('a').sum()` + returns_scalar: bool, backend_version: tuple[int, ...], ) -> None: self._call = call @@ -34,15 +41,17 @@ def __init__( self._function_name = function_name self._root_names = root_names self._output_names = output_names + self._returns_scalar = returns_scalar self._backend_version = backend_version + def __narwhals_expr__(self) -> None: ... + def __narwhals_namespace__(self) -> DaskNamespace: # pragma: no cover + # Unused, just for compatibility with PandasLikeExpr from narwhals._dask.namespace import DaskNamespace return DaskNamespace(backend_version=self._backend_version) - def __narwhals_expr__(self) -> None: ... - @classmethod def from_column_names( cls: type[Self], @@ -50,9 +59,7 @@ def from_column_names( backend_version: tuple[int, ...], ) -> Self: def func(df: DaskLazyFrame) -> list[Any]: - return [ - df._native_dataframe.loc[:, column_name] for column_name in column_names - ] + return [df._native_frame.loc[:, column_name] for column_name in column_names] return cls( func, @@ -60,6 +67,7 @@ def func(df: DaskLazyFrame) -> list[Any]: function_name="col", root_names=list(column_names), output_names=list(column_names), + returns_scalar=False, backend_version=backend_version, ) @@ -69,19 +77,19 @@ def _from_call( call: Any, expr_name: str, *args: Any, + returns_scalar: bool, **kwargs: Any, ) -> Self: def func(df: DaskLazyFrame) -> list[Any]: results = [] inputs = self._call(df) + _args = [maybe_evaluate(df, x) for x in args] + _kwargs = {key: maybe_evaluate(df, value) for key, value in kwargs.items()} for _input in inputs: - _args = [maybe_evaluate(df, x) for x in args] - _kwargs = { - key: maybe_evaluate(df, value) for key, value in kwargs.items() - } result = call(_input, *_args, **_kwargs) - if isinstance(result, get_dask_expr()._collection.Series): - result = result.rename(_input.name) + if returns_scalar: + result = result.to_series() + result = result.rename(_input.name) results.append(result) return results @@ -118,17 +126,14 @@ def func(df: DaskLazyFrame) -> list[Any]: function_name=f"{self._function_name}->{expr_name}", root_names=root_names, output_names=output_names, + returns_scalar=self._returns_scalar or returns_scalar, backend_version=self._backend_version, ) def alias(self, name: str) -> Self: def func(df: DaskLazyFrame) -> list[Any]: - results = [] inputs = self._call(df) - for _input in inputs: - result = _input.rename(name) - results.append(result) - return results + return [_input.rename(name) for _input in inputs] return self.__class__( func, @@ -136,6 +141,7 @@ def func(df: DaskLazyFrame) -> list[Any]: function_name=self._function_name, root_names=self._root_names, output_names=[name], + returns_scalar=self._returns_scalar, backend_version=self._backend_version, ) @@ -144,6 +150,15 @@ def __add__(self, other: Any) -> Self: lambda _input, other: _input.__add__(other), "__add__", other, + returns_scalar=False, + ) + + def __radd__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__radd__(other), + "__radd__", + other, + returns_scalar=False, ) def __sub__(self, other: Any) -> Self: @@ -151,6 +166,15 @@ def __sub__(self, other: Any) -> Self: lambda _input, other: _input.__sub__(other), "__sub__", other, + returns_scalar=False, + ) + + def __rsub__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__rsub__(other), + "__rsub__", + other, + returns_scalar=False, ) def __mul__(self, other: Any) -> Self: @@ -158,12 +182,195 @@ def __mul__(self, other: Any) -> Self: lambda _input, other: _input.__mul__(other), "__mul__", other, + returns_scalar=False, + ) + + def __rmul__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__rmul__(other), + "__rmul__", + other, + returns_scalar=False, + ) + + def __truediv__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__truediv__(other), + "__truediv__", + other, + returns_scalar=False, + ) + + def __rtruediv__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__rtruediv__(other), + "__rtruediv__", + other, + returns_scalar=False, + ) + + def __floordiv__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__floordiv__(other), + "__floordiv__", + other, + returns_scalar=False, + ) + + def __rfloordiv__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__rfloordiv__(other), + "__rfloordiv__", + other, + returns_scalar=False, + ) + + def __pow__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__pow__(other), + "__pow__", + other, + returns_scalar=False, + ) + + def __rpow__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__rpow__(other), + "__rpow__", + other, + returns_scalar=False, + ) + + def __mod__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__mod__(other), + "__mod__", + other, + returns_scalar=False, + ) + + def __rmod__(self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.__rmod__(other), + "__rmod__", + other, + returns_scalar=False, + ) + + def __eq__(self, other: DaskExpr) -> Self: # type: ignore[override] + return self._from_call( + lambda _input, other: _input.__eq__(other), + "__eq__", + other, + returns_scalar=False, + ) + + def __ne__(self, other: DaskExpr) -> Self: # type: ignore[override] + return self._from_call( + lambda _input, other: _input.__ne__(other), + "__ne__", + other, + returns_scalar=False, + ) + + def __ge__(self, other: DaskExpr) -> Self: + return self._from_call( + lambda _input, other: _input.__ge__(other), + "__ge__", + other, + returns_scalar=False, + ) + + def __gt__(self, other: DaskExpr) -> Self: + return self._from_call( + lambda _input, other: _input.__gt__(other), + "__gt__", + other, + returns_scalar=False, + ) + + def __le__(self, other: DaskExpr) -> Self: + return self._from_call( + lambda _input, other: _input.__le__(other), + "__le__", + other, + returns_scalar=False, + ) + + def __lt__(self, other: DaskExpr) -> Self: + return self._from_call( + lambda _input, other: _input.__lt__(other), + "__lt__", + other, + returns_scalar=False, + ) + + def __and__(self, other: DaskExpr) -> Self: + return self._from_call( + lambda _input, other: _input.__and__(other), + "__and__", + other, + returns_scalar=False, + ) + + def __rand__(self, other: DaskExpr) -> Self: # pragma: no cover + return self._from_call( + lambda _input, other: _input.__rand__(other), + "__rand__", + other, + returns_scalar=False, + ) + + def __or__(self, other: DaskExpr) -> Self: + return self._from_call( + lambda _input, other: _input.__or__(other), + "__or__", + other, + returns_scalar=False, + ) + + def __ror__(self, other: DaskExpr) -> Self: # pragma: no cover + return self._from_call( + lambda _input, other: _input.__ror__(other), + "__ror__", + other, + returns_scalar=False, + ) + + def __invert__(self: Self) -> Self: + return self._from_call( + lambda _input: _input.__invert__(), + "__invert__", + returns_scalar=False, ) def mean(self) -> Self: return self._from_call( lambda _input: _input.mean(), "mean", + returns_scalar=True, + ) + + def min(self) -> Self: + return self._from_call( + lambda _input: _input.min(), + "min", + returns_scalar=True, + ) + + def max(self) -> Self: + return self._from_call( + lambda _input: _input.max(), + "max", + returns_scalar=True, + ) + + def std(self, ddof: int = 1) -> Self: + return self._from_call( + lambda _input, ddof: _input.std(ddof=ddof), + "std", + ddof, + returns_scalar=True, ) def shift(self, n: int) -> Self: @@ -171,12 +378,14 @@ def shift(self, n: int) -> Self: lambda _input, n: _input.shift(n), "shift", n, + returns_scalar=False, ) def cum_sum(self) -> Self: return self._from_call( lambda _input: _input.cumsum(), "cum_sum", + returns_scalar=False, ) def is_between( @@ -185,6 +394,8 @@ def is_between( upper_bound: Any, closed: str = "both", ) -> Self: + if closed == "none": + closed = "neither" return self._from_call( lambda _input, lower_bound, upper_bound, closed: _input.between( lower_bound, @@ -195,12 +406,242 @@ def is_between( lower_bound, upper_bound, closed, + returns_scalar=False, ) def sum(self) -> Self: return self._from_call( lambda _input: _input.sum(), "sum", + returns_scalar=True, + ) + + def count(self) -> Self: + return self._from_call( + lambda _input: _input.count(), + "count", + returns_scalar=True, + ) + + def round(self, decimals: int) -> Self: + return self._from_call( + lambda _input, decimals: _input.round(decimals), + "round", + decimals, + returns_scalar=False, + ) + + def drop_nulls(self) -> NoReturn: + # We can't (yet?) allow methods which modify the index + msg = "`Expr.drop_nulls` is not supported for the Dask backend. Please use `LazyFrame.drop_nulls` instead." + raise NotImplementedError(msg) + + def head(self) -> NoReturn: + # We can't (yet?) allow methods which modify the index + msg = "`Expr.head` is not supported for the Dask backend. Please use `LazyFrame.head` instead." + raise NotImplementedError(msg) + + def sort(self, *, descending: bool = False, nulls_last: bool = False) -> NoReturn: + # We can't (yet?) allow methods which modify the index + msg = "`Expr.sort` is not supported for the Dask backend. Please use `LazyFrame.sort` instead." + raise NotImplementedError(msg) + + def abs(self) -> Self: + return self._from_call( + lambda _input: _input.abs(), + "abs", + returns_scalar=False, + ) + + def all(self) -> Self: + return self._from_call( + lambda _input: _input.all( + axis=None, skipna=True, split_every=False, out=None + ), + "all", + returns_scalar=True, + ) + + def any(self) -> Self: + return self._from_call( + lambda _input: _input.any(axis=0, skipna=True, split_every=False), + "any", + returns_scalar=True, + ) + + def fill_null(self, value: Any) -> DaskExpr: + return self._from_call( + lambda _input, _val: _input.fillna(_val), + "fillna", + value, + returns_scalar=False, + ) + + def clip( + self: Self, + lower_bound: Any | None = None, + upper_bound: Any | None = None, + ) -> Self: + return self._from_call( + lambda _input, _lower, _upper: _input.clip(lower=_lower, upper=_upper), + "clip", + lower_bound, + upper_bound, + returns_scalar=False, + ) + + def diff(self: Self) -> Self: + return self._from_call( + lambda _input: _input.diff(), + "diff", + returns_scalar=False, + ) + + def n_unique(self: Self) -> Self: + return self._from_call( + lambda _input: _input.nunique(dropna=False), + "n_unique", + returns_scalar=True, + ) + + def is_null(self: Self) -> Self: + return self._from_call( + lambda _input: _input.isna(), + "is_null", + returns_scalar=False, + ) + + def len(self: Self) -> Self: + return self._from_call( + lambda _input: _input.size, + "len", + returns_scalar=True, + ) + + def quantile( + self: Self, + quantile: float, + interpolation: Literal["nearest", "higher", "lower", "midpoint", "linear"], + ) -> Self: + if interpolation == "linear": + return self._from_call( + lambda _input, quantile: _input.quantile(q=quantile, method="dask"), + "quantile", + quantile, + returns_scalar=True, + ) + else: + msg = "`higher`, `lower`, `midpoint`, `nearest` - interpolation methods are not supported by Dask. Please use `linear` instead." + raise NotImplementedError(msg) + + def is_first_distinct(self: Self) -> Self: + def func(_input: Any) -> Any: + _name = _input.name + col_token = generate_unique_token(n_bytes=8, columns=[_name]) + _input = add_row_index(_input.to_frame(), col_token) + first_distinct_index = _input.groupby(_name).agg({col_token: "min"})[ + col_token + ] + + return _input[col_token].isin(first_distinct_index) + + return self._from_call( + func, + "is_first_distinct", + returns_scalar=False, + ) + + def is_last_distinct(self: Self) -> Self: + def func(_input: Any) -> Any: + _name = _input.name + col_token = generate_unique_token(n_bytes=8, columns=[_name]) + _input = add_row_index(_input.to_frame(), col_token) + last_distinct_index = _input.groupby(_name).agg({col_token: "max"})[col_token] + + return _input[col_token].isin(last_distinct_index) + + return self._from_call( + func, + "is_last_distinct", + returns_scalar=False, + ) + + def is_duplicated(self: Self) -> Self: + def func(_input: Any) -> Any: + _name = _input.name + return ( + _input.to_frame().groupby(_name).transform("size", meta=(_name, int)) > 1 + ) + + return self._from_call( + func, + "is_duplicated", + returns_scalar=False, + ) + + def is_unique(self: Self) -> Self: + def func(_input: Any) -> Any: + _name = _input.name + return ( + _input.to_frame().groupby(_name).transform("size", meta=(_name, int)) == 1 + ) + + return self._from_call( + func, + "is_unique", + returns_scalar=False, + ) + + def is_in(self: Self, other: Any) -> Self: + return self._from_call( + lambda _input, other: _input.isin(other), + "is_in", + other, + returns_scalar=False, + ) + + def null_count(self: Self) -> Self: + return self._from_call( + lambda _input: _input.isna().sum(), + "null_count", + returns_scalar=True, + ) + + def tail(self: Self) -> NoReturn: + # We can't (yet?) allow methods which modify the index + msg = "`Expr.tail` is not supported for the Dask backend. Please use `LazyFrame.tail` instead." + raise NotImplementedError(msg) + + def gather_every(self: Self, n: int, offset: int = 0) -> NoReturn: + # We can't (yet?) allow methods which modify the index + msg = "`Expr.gather_every` is not supported for the Dask backend. Please use `LazyFrame.gather_every` instead." + raise NotImplementedError(msg) + + def over(self: Self, keys: list[str]) -> Self: + def func(df: DaskLazyFrame) -> list[Any]: + if self._output_names is None: + msg = ( + "Anonymous expressions are not supported in over.\n" + "Instead of `nw.all()`, try using a named expression, such as " + "`nw.col('a', 'b')`\n" + ) + raise ValueError(msg) + tmp = df.group_by(*keys).agg(self) + tmp = ( + df.select(*keys) + .join(tmp, how="left", left_on=keys, right_on=keys) + ._native_frame + ) + return [tmp[name] for name in self._output_names] + + return self.__class__( + func, + depth=self._depth + 1, + function_name=self._function_name + "->over", + root_names=self._root_names, + output_names=self._output_names, + returns_scalar=False, + backend_version=self._backend_version, ) @property @@ -211,19 +652,90 @@ def str(self: Self) -> DaskExprStringNamespace: def dt(self: Self) -> DaskExprDateTimeNamespace: return DaskExprDateTimeNamespace(self) + @property + def name(self: Self) -> DaskExprNameNamespace: + return DaskExprNameNamespace(self) + + def cast( + self: Self, + dtype: DType | type[DType], + ) -> Self: + def func(_input: Any, dtype: DType | type[DType]) -> Any: + dtype = reverse_translate_dtype(dtype) + return _input.astype(dtype) + + return self._from_call( + func, + "cast", + dtype, + returns_scalar=False, + ) + class DaskExprStringNamespace: def __init__(self, expr: DaskExpr) -> None: self._expr = expr + def replace( + self, + pattern: str, + value: str, + *, + literal: bool = False, + n: int = 1, + ) -> DaskExpr: + return self._expr._from_call( + lambda _input, _pattern, _value, _literal, _n: _input.str.replace( + _pattern, _value, regex=not _literal, n=_n + ), + "replace", + pattern, + value, + literal, + n, + returns_scalar=False, + ) + + def replace_all( + self, + pattern: str, + value: str, + *, + literal: bool = False, + ) -> DaskExpr: + return self._expr._from_call( + lambda _input, _pattern, _value, _literal: _input.str.replace( + _pattern, _value, n=-1, regex=not _literal + ), + "replace", + pattern, + value, + literal, + returns_scalar=False, + ) + + def strip_chars(self, characters: str | None = None) -> DaskExpr: + return self._expr._from_call( + lambda _input, characters: _input.str.strip(characters), + "strip", + characters, + returns_scalar=False, + ) + def starts_with(self, prefix: str) -> DaskExpr: return self._expr._from_call( - lambda _input, prefix: _input.str.startswith(prefix), "starts_with", prefix + lambda _input, prefix: _input.str.startswith(prefix), + "starts_with", + prefix, + returns_scalar=False, ) def ends_with(self, suffix: str) -> DaskExpr: return self._expr._from_call( - lambda _input, suffix: _input.str.endswith(suffix), "ends_with", suffix + lambda _input, suffix: _input.str.endswith(suffix), + "ends_with", + suffix, + returns_scalar=False, ) def contains(self, pattern: str, *, literal: bool = False) -> DaskExpr: @@ -232,6 +744,7 @@ def contains(self, pattern: str, *, literal: bool = False) -> DaskExpr: "contains", pattern, not literal, + returns_scalar=False, ) def slice(self, offset: int, length: int | None = None) -> DaskExpr: @@ -241,6 +754,7 @@ def slice(self, offset: int, length: int | None = None) -> DaskExpr: "slice", offset, stop, + returns_scalar=False, ) def to_datetime(self, format: str | None = None) -> DaskExpr: # noqa: A002 @@ -248,18 +762,21 @@ def to_datetime(self, format: str | None = None) -> DaskExpr: # noqa: A002 lambda _input, fmt: get_dask().dataframe.to_datetime(_input, format=fmt), "to_datetime", format, + returns_scalar=False, ) def to_uppercase(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.str.upper(), "to_uppercase", + returns_scalar=False, ) def to_lowercase(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.str.lower(), "to_lowercase", + returns_scalar=False, ) @@ -267,62 +784,276 @@ class DaskExprDateTimeNamespace: def __init__(self, expr: DaskExpr) -> None: self._expr = expr + def date(self) -> DaskExpr: + return self._expr._from_call( + lambda _input: _input.dt.date, + "date", + returns_scalar=False, + ) + def year(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.year, "year", + returns_scalar=False, ) def month(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.month, "month", + returns_scalar=False, ) def day(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.day, "day", + returns_scalar=False, ) def hour(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.hour, "hour", + returns_scalar=False, ) def minute(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.minute, "minute", + returns_scalar=False, ) def second(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.second, "second", + returns_scalar=False, ) def millisecond(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.microsecond // 1000, "millisecond", + returns_scalar=False, ) def microsecond(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.microsecond, "microsecond", + returns_scalar=False, ) def nanosecond(self) -> DaskExpr: return self._expr._from_call( - lambda _input: _input.dt.microsecond * 1000, + lambda _input: _input.dt.microsecond * 1000 + _input.dt.nanosecond, "nanosecond", + returns_scalar=False, ) def ordinal_day(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.dayofyear, "ordinal_day", + returns_scalar=False, + ) + + def to_string(self, format: str) -> DaskExpr: # noqa: A002 + return self._expr._from_call( + lambda _input, _format: _input.dt.strftime(_format), + "strftime", + format.replace("%.f", ".%f"), + returns_scalar=False, + ) + + def total_minutes(self) -> DaskExpr: + return self._expr._from_call( + lambda _input: _input.dt.total_seconds() // 60, + "total_minutes", + returns_scalar=False, + ) + + def total_seconds(self) -> DaskExpr: + return self._expr._from_call( + lambda _input: _input.dt.total_seconds() // 1, + "total_seconds", + returns_scalar=False, + ) + + def total_milliseconds(self) -> DaskExpr: + return self._expr._from_call( + lambda _input: _input.dt.total_seconds() * 1000 // 1, + "total_milliseconds", + returns_scalar=False, + ) + + def total_microseconds(self) -> DaskExpr: + return self._expr._from_call( + lambda _input: _input.dt.total_seconds() * 1_000_000 // 1, + "total_microseconds", + returns_scalar=False, + ) + + def total_nanoseconds(self) -> DaskExpr: + return self._expr._from_call( + lambda _input: _input.dt.total_seconds() * 1_000_000_000 // 1, + "total_nanoseconds", + returns_scalar=False, + ) + + +class DaskExprNameNamespace: + def __init__(self: Self, expr: DaskExpr) -> None: + self._expr = expr + + def keep(self: Self) -> DaskExpr: + root_names = self._expr._root_names + + if root_names is None: + msg = ( + "Anonymous expressions are not supported in `.name.keep`.\n" + "Instead of `nw.all()`, try using a named expression, such as " + "`nw.col('a', 'b')`\n" + ) + raise ValueError(msg) + + return self._expr.__class__( + lambda df: [ + series.rename(name) + for series, name in zip(self._expr._call(df), root_names) + ], + depth=self._expr._depth, + function_name=self._expr._function_name, + root_names=root_names, + output_names=root_names, + returns_scalar=self._expr._returns_scalar, + backend_version=self._expr._backend_version, + ) + + def map(self: Self, function: Callable[[str], str]) -> DaskExpr: + root_names = self._expr._root_names + + if root_names is None: + msg = ( + "Anonymous expressions are not supported in `.name.map`.\n" + "Instead of `nw.all()`, try using a named expression, such as " + "`nw.col('a', 'b')`\n" + ) + raise ValueError(msg) + + output_names = [function(str(name)) for name in root_names] + + return self._expr.__class__( + lambda df: [ + series.rename(name) + for series, name in zip(self._expr._call(df), output_names) + ], + depth=self._expr._depth, + function_name=self._expr._function_name, + root_names=root_names, + output_names=output_names, + returns_scalar=self._expr._returns_scalar, + backend_version=self._expr._backend_version, + ) + + def prefix(self: Self, prefix: str) -> DaskExpr: + root_names = self._expr._root_names + if root_names is None: + msg = ( + "Anonymous expressions are not supported in `.name.prefix`.\n" + "Instead of `nw.all()`, try using a named expression, such as " + "`nw.col('a', 'b')`\n" + ) + raise ValueError(msg) + + output_names = [prefix + str(name) for name in root_names] + return self._expr.__class__( + lambda df: [ + series.rename(name) + for series, name in zip(self._expr._call(df), output_names) + ], + depth=self._expr._depth, + function_name=self._expr._function_name, + root_names=root_names, + output_names=output_names, + returns_scalar=self._expr._returns_scalar, + backend_version=self._expr._backend_version, + ) + + def suffix(self: Self, suffix: str) -> DaskExpr: + root_names = self._expr._root_names + if root_names is None: + msg = ( + "Anonymous expressions are not supported in `.name.suffix`.\n" + "Instead of `nw.all()`, try using a named expression, such as " + "`nw.col('a', 'b')`\n" + ) + raise ValueError(msg) + + output_names = [str(name) + suffix for name in root_names] + + return self._expr.__class__( + lambda df: [ + series.rename(name) + for series, name in zip(self._expr._call(df), output_names) + ], + depth=self._expr._depth, + function_name=self._expr._function_name, + root_names=root_names, + output_names=output_names, + returns_scalar=self._expr._returns_scalar, + backend_version=self._expr._backend_version, + ) + + def to_lowercase(self: Self) -> DaskExpr: + root_names = self._expr._root_names + + if root_names is None: + msg = ( + "Anonymous expressions are not supported in `.name.to_lowercase`.\n" + "Instead of `nw.all()`, try using a named expression, such as " + "`nw.col('a', 'b')`\n" + ) + raise ValueError(msg) + output_names = [str(name).lower() for name in root_names] + + return self._expr.__class__( + lambda df: [ + series.rename(name) + for series, name in zip(self._expr._call(df), output_names) + ], + depth=self._expr._depth, + function_name=self._expr._function_name, + root_names=root_names, + output_names=output_names, + returns_scalar=self._expr._returns_scalar, + backend_version=self._expr._backend_version, + ) + + def to_uppercase(self: Self) -> DaskExpr: + root_names = self._expr._root_names + + if root_names is None: + msg = ( + "Anonymous expressions are not supported in `.name.to_uppercase`.\n" + "Instead of `nw.all()`, try using a named expression, such as " + "`nw.col('a', 'b')`\n" + ) + raise ValueError(msg) + output_names = [str(name).upper() for name in root_names] + + return self._expr.__class__( + lambda df: [ + series.rename(name) + for series, name in zip(self._expr._call(df), output_names) + ], + depth=self._expr._depth, + function_name=self._expr._function_name, + root_names=root_names, + output_names=output_names, + returns_scalar=self._expr._returns_scalar, + backend_version=self._expr._backend_version, ) diff --git a/narwhals/_dask/group_by.py b/narwhals/_dask/group_by.py new file mode 100644 index 000000000..8538c62d2 --- /dev/null +++ b/narwhals/_dask/group_by.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from copy import copy +from typing import TYPE_CHECKING +from typing import Any +from typing import Callable + +from narwhals._expression_parsing import is_simple_aggregation +from narwhals._expression_parsing import parse_into_exprs +from narwhals.utils import remove_prefix + +if TYPE_CHECKING: + from narwhals._dask.dataframe import DaskLazyFrame + from narwhals._dask.expr import DaskExpr + from narwhals._dask.typing import IntoDaskExpr + +POLARS_TO_PANDAS_AGGREGATIONS = { + "len": "size", +} + + +class DaskLazyGroupBy: + def __init__(self, df: DaskLazyFrame, keys: list[str]) -> None: + self._df = df + self._keys = keys + self._grouped = self._df._native_frame.groupby( + list(self._keys), + dropna=False, + ) + + def agg( + self, + *aggs: IntoDaskExpr, + **named_aggs: IntoDaskExpr, + ) -> DaskLazyFrame: + exprs = parse_into_exprs( + *aggs, + namespace=self._df.__narwhals_namespace__(), + **named_aggs, + ) + output_names: list[str] = copy(self._keys) + for expr in exprs: + if expr._output_names is None: + msg = ( + "Anonymous expressions are not supported in group_by.agg.\n" + "Instead of `nw.all()`, try using a named expression, such as " + "`nw.col('a', 'b')`\n" + ) + raise ValueError(msg) + + output_names.extend(expr._output_names) + + return agg_dask( + self._grouped, + exprs, + self._keys, + self._from_native_frame, + ) + + def _from_native_frame(self, df: DaskLazyFrame) -> DaskLazyFrame: + from narwhals._dask.dataframe import DaskLazyFrame + + return DaskLazyFrame( + df, + backend_version=self._df._backend_version, + ) + + +def agg_dask( + grouped: Any, + exprs: list[DaskExpr], + keys: list[str], + from_dataframe: Callable[[Any], DaskLazyFrame], +) -> DaskLazyFrame: + """ + This should be the fastpath, but cuDF is too far behind to use it. + + - https://github.com/rapidsai/cudf/issues/15118 + - https://github.com/rapidsai/cudf/issues/15084 + """ + all_simple_aggs = True + for expr in exprs: + if not is_simple_aggregation(expr): + all_simple_aggs = False + break + + if all_simple_aggs: + simple_aggregations: dict[str, tuple[str, str]] = {} + for expr in exprs: + if expr._depth == 0: + # e.g. agg(nw.len()) # noqa: ERA001 + if expr._output_names is None: # pragma: no cover + msg = "Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues" + raise AssertionError(msg) + + function_name = POLARS_TO_PANDAS_AGGREGATIONS.get( + expr._function_name, expr._function_name + ) + for output_name in expr._output_names: + simple_aggregations[output_name] = (keys[0], function_name) + continue + + # e.g. agg(nw.mean('a')) # noqa: ERA001 + if ( + expr._depth != 1 or expr._root_names is None or expr._output_names is None + ): # pragma: no cover + msg = "Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues" + raise AssertionError(msg) + + function_name = remove_prefix(expr._function_name, "col->") + function_name = POLARS_TO_PANDAS_AGGREGATIONS.get( + function_name, function_name + ) + for root_name, output_name in zip(expr._root_names, expr._output_names): + simple_aggregations[output_name] = (root_name, function_name) + try: + result_simple = grouped.agg(**simple_aggregations) + except ValueError as exc: + msg = "Failed to aggregated - does your aggregation function return a scalar?" + raise RuntimeError(msg) from exc + return from_dataframe(result_simple.reset_index()) + + msg = ( + "Non-trivial complex found.\n\n" + "Hint: you were probably trying to apply a non-elementary aggregation with a " + "dask dataframe.\n" + "Please rewrite your query such that group-by aggregations " + "are elementary. For example, instead of:\n\n" + " df.group_by('a').agg(nw.col('b').round(2).mean())\n\n" + "use:\n\n" + " df.with_columns(nw.col('b').round(2)).group_by('a').agg(nw.col('b').mean())\n\n" + ) + raise ValueError(msg) diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py index 7ae42ea1a..e6019b509 100644 --- a/narwhals/_dask/namespace.py +++ b/narwhals/_dask/namespace.py @@ -1,16 +1,23 @@ from __future__ import annotations +from functools import reduce from typing import TYPE_CHECKING from typing import Any +from typing import Callable from typing import NoReturn +from typing import cast from narwhals import dtypes from narwhals._dask.expr import DaskExpr +from narwhals._dask.selectors import DaskSelectorNamespace +from narwhals._dask.utils import validate_comparand +from narwhals._expression_parsing import parse_into_exprs if TYPE_CHECKING: - from typing import Callable + import dask_expr from narwhals._dask.dataframe import DaskLazyFrame + from narwhals._dask.typing import IntoDaskExpr class DaskNamespace: @@ -34,15 +41,113 @@ class DaskNamespace: Duration = dtypes.Duration Date = dtypes.Date + @property + def selectors(self) -> DaskSelectorNamespace: + return DaskSelectorNamespace(backend_version=self._backend_version) + def __init__(self, *, backend_version: tuple[int, ...]) -> None: self._backend_version = backend_version + def all(self) -> DaskExpr: + def func(df: DaskLazyFrame) -> list[Any]: + return [df._native_frame.loc[:, column_name] for column_name in df.columns] + + return DaskExpr( + func, + depth=0, + function_name="all", + root_names=None, + output_names=None, + returns_scalar=False, + backend_version=self._backend_version, + ) + def col(self, *column_names: str) -> DaskExpr: return DaskExpr.from_column_names( *column_names, backend_version=self._backend_version, ) + def lit(self, value: Any, dtype: dtypes.DType | None) -> DaskExpr: + # TODO @FBruzzesi: cast to dtype once `narwhals_to_native_dtype` is implemented. + # It should be enough to add `.astype(narwhals_to_native_dtype(dtype))` + return DaskExpr( + lambda df: [df._native_frame.assign(lit=value).loc[:, "lit"]], + depth=0, + function_name="lit", + root_names=None, + output_names=["lit"], + returns_scalar=False, + backend_version=self._backend_version, + ) + + def min(self, *column_names: str) -> DaskExpr: + return DaskExpr.from_column_names( + *column_names, + backend_version=self._backend_version, + ).min() + + def max(self, *column_names: str) -> DaskExpr: + return DaskExpr.from_column_names( + *column_names, + backend_version=self._backend_version, + ).max() + + def mean(self, *column_names: str) -> DaskExpr: + return DaskExpr.from_column_names( + *column_names, + backend_version=self._backend_version, + ).mean() + + def sum(self, *column_names: str) -> DaskExpr: + return DaskExpr.from_column_names( + *column_names, + backend_version=self._backend_version, + ).sum() + + def len(self) -> DaskExpr: + import dask.dataframe as dd # ignore-banned-import + import pandas as pd # ignore-banned-import + + def func(df: DaskLazyFrame) -> list[Any]: + if not df.columns: + return [ + dd.from_pandas( + pd.Series([0], name="len"), + npartitions=df._native_frame.npartitions, + ) + ] + return [df._native_frame.loc[:, df.columns[0]].size.to_series().rename("len")] + + # coverage bug? this is definitely hit + return DaskExpr( # pragma: no cover + func, + depth=0, + function_name="len", + root_names=None, + output_names=["len"], + returns_scalar=True, + backend_version=self._backend_version, + ) + + def all_horizontal(self, *exprs: IntoDaskExpr) -> DaskExpr: + return reduce(lambda x, y: x & y, parse_into_exprs(*exprs, namespace=self)) + + def any_horizontal(self, *exprs: IntoDaskExpr) -> DaskExpr: + return reduce(lambda x, y: x | y, parse_into_exprs(*exprs, namespace=self)) + + def sum_horizontal(self, *exprs: IntoDaskExpr) -> DaskExpr: + return reduce( + lambda x, y: x + y, + [expr.fill_null(0) for expr in parse_into_exprs(*exprs, namespace=self)], + ) + + def mean_horizontal(self, *exprs: IntoDaskExpr) -> IntoDaskExpr: + dask_exprs = parse_into_exprs(*exprs, namespace=self) + total = reduce(lambda x, y: x + y, (e.fill_null(0.0) for e in dask_exprs)) + n_non_zero = reduce(lambda x, y: x + y, ((1 - e.is_null()) for e in dask_exprs)) + return total / n_non_zero + def _create_expr_from_series(self, _: Any) -> NoReturn: msg = "`_create_expr_from_series` for DaskNamespace exists only for compatibility" raise NotImplementedError(msg) @@ -66,11 +171,110 @@ def _create_expr_from_callable( # pragma: no cover root_names: list[str] | None, output_names: list[str] | None, ) -> DaskExpr: - return DaskExpr( - call=func, - depth=depth, - function_name=function_name, - root_names=root_names, - output_names=output_names, + msg = ( + "`_create_expr_from_callable` for DaskNamespace exists only for compatibility" + ) + raise NotImplementedError(msg) + + def when( + self, + *predicates: IntoDaskExpr, + ) -> DaskWhen: + plx = self.__class__(backend_version=self._backend_version) + if predicates: + condition = plx.all_horizontal(*predicates) + else: + msg = "at least one predicate needs to be provided" + raise TypeError(msg) + + return DaskWhen(condition, self._backend_version, returns_scalar=False) + + +class DaskWhen: + def __init__( + self, + condition: DaskExpr, + backend_version: tuple[int, ...], + then_value: Any = None, + otherwise_value: Any = None, + *, + returns_scalar: bool, + ) -> None: + self._backend_version = backend_version + self._condition = condition + self._then_value = then_value + self._otherwise_value = otherwise_value + self._returns_scalar = returns_scalar + + def __call__(self, df: DaskLazyFrame) -> list[Any]: + from narwhals._dask.namespace import DaskNamespace + from narwhals._expression_parsing import parse_into_expr + + plx = DaskNamespace(backend_version=self._backend_version) + + condition = parse_into_expr(self._condition, namespace=plx)._call(df)[0] # type: ignore[arg-type] + condition = cast("dask_expr.Series", condition) + try: + value_series = parse_into_expr(self._then_value, namespace=plx)._call(df)[0] # type: ignore[arg-type] + except TypeError: + # `self._otherwise_value` is a scalar and can't be converted to an expression + _df = condition.to_frame("a") + _df["tmp"] = self._then_value + value_series = _df["tmp"] + value_series = cast("dask_expr.Series", value_series) + validate_comparand(condition, value_series) + + if self._otherwise_value is None: + return [value_series.where(condition)] + try: + otherwise_series = parse_into_expr( + self._otherwise_value, namespace=plx + )._call(df)[0] # type: ignore[arg-type] + except TypeError: + # `self._otherwise_value` is a scalar and can't be converted to an expression + return [value_series.where(condition, self._otherwise_value)] + validate_comparand(condition, otherwise_series) + return [value_series.zip_with(condition, otherwise_series)] + + def then(self, value: DaskExpr | Any) -> DaskThen: + self._then_value = value + + return DaskThen( + self, + depth=0, + function_name="whenthen", + root_names=None, + output_names=None, + returns_scalar=self._returns_scalar, backend_version=self._backend_version, ) + + +class DaskThen(DaskExpr): + def __init__( + self, + call: DaskWhen, + *, + depth: int, + function_name: str, + root_names: list[str] | None, + output_names: list[str] | None, + returns_scalar: bool, + backend_version: tuple[int, ...], + ) -> None: + self._backend_version = backend_version + + self._call = call + self._depth = depth + self._function_name = function_name + self._root_names = root_names + self._output_names = output_names + self._returns_scalar = returns_scalar + + def otherwise(self, value: DaskExpr | Any) -> DaskExpr: + # type ignore because we are setting the `_call` attribute to a + # callable object of type `DaskWhen`, base class has the attribute as + # only a `Callable` + self._call._otherwise_value = value # type: ignore[attr-defined] + self._function_name = "whenotherwise" + return self diff --git a/narwhals/_dask/selectors.py b/narwhals/_dask/selectors.py new file mode 100644 index 000000000..073b3abd8 --- /dev/null +++ b/narwhals/_dask/selectors.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import NoReturn + +from narwhals import dtypes +from narwhals._dask.expr import DaskExpr + +if TYPE_CHECKING: + from typing_extensions import Self + + from narwhals._dask.dataframe import DaskLazyFrame + from narwhals.dtypes import DType + + +class DaskSelectorNamespace: + def __init__(self: Self, *, backend_version: tuple[int, ...]) -> None: + self._backend_version = backend_version + + def by_dtype(self: Self, dtypes: list[DType | type[DType]]) -> DaskSelector: + def func(df: DaskLazyFrame) -> list[Any]: + return [ + df._native_frame[col] for col in df.columns if df.schema[col] in dtypes + ] + + return DaskSelector( + func, + depth=0, + function_name="type_selector", + root_names=None, + output_names=None, + backend_version=self._backend_version, + returns_scalar=False, + ) + + def numeric(self: Self) -> DaskSelector: + return self.by_dtype( + [ + dtypes.Int64, + dtypes.Int32, + dtypes.Int16, + dtypes.Int8, + dtypes.UInt64, + dtypes.UInt32, + dtypes.UInt16, + dtypes.UInt8, + dtypes.Float64, + dtypes.Float32, + ], + ) + + def categorical(self: Self) -> DaskSelector: + return self.by_dtype([dtypes.Categorical]) + + def string(self: Self) -> DaskSelector: + return self.by_dtype([dtypes.String]) + + def boolean(self: Self) -> DaskSelector: + return self.by_dtype([dtypes.Boolean]) + + def all(self: Self) -> DaskSelector: + def func(df: DaskLazyFrame) -> list[Any]: + return [df._native_frame[col] for col in df.columns] + + return DaskSelector( + func, + depth=0, + function_name="type_selector", + root_names=None, + output_names=None, + backend_version=self._backend_version, + returns_scalar=False, + ) + + +class DaskSelector(DaskExpr): + def __repr__(self: Self) -> str: # pragma: no cover + return ( + f"DaskSelector(" + f"depth={self._depth}, " + f"function_name={self._function_name}, " + f"root_names={self._root_names}, " + f"output_names={self._output_names}" + ) + + def _to_expr(self: Self) -> DaskExpr: + return DaskExpr( + self._call, + depth=self._depth, + function_name=self._function_name, + root_names=self._root_names, + output_names=self._output_names, + backend_version=self._backend_version, + returns_scalar=self._returns_scalar, + ) + + def __sub__(self: Self, other: DaskSelector | Any) -> DaskSelector | Any: + if isinstance(other, DaskSelector): + + def call(df: DaskLazyFrame) -> list[Any]: + lhs = self._call(df) + rhs = other._call(df) + return [x for x in lhs if x.name not in [x.name for x in rhs]] + + return DaskSelector( + call, + depth=0, + function_name="type_selector", + root_names=None, + output_names=None, + backend_version=self._backend_version, + returns_scalar=self._returns_scalar, + ) + else: + return self._to_expr() - other + + def __or__(self: Self, other: DaskSelector | Any) -> DaskSelector | Any: + if isinstance(other, DaskSelector): + + def call(df: DaskLazyFrame) -> list[Any]: + lhs = self._call(df) + rhs = other._call(df) + return [ # type: ignore[no-any-return] + x for x in lhs if x.name not in [x.name for x in rhs] + ] + rhs + + return DaskSelector( + call, + depth=0, + function_name="type_selector", + root_names=None, + output_names=None, + backend_version=self._backend_version, + returns_scalar=self._returns_scalar, + ) + else: + return self._to_expr() | other + + def __and__(self: Self, other: DaskSelector | Any) -> DaskSelector | Any: + if isinstance(other, DaskSelector): + + def call(df: DaskLazyFrame) -> list[Any]: + lhs = self._call(df) + rhs = other._call(df) + return [x for x in lhs if x.name in [x.name for x in rhs]] + + return DaskSelector( + call, + depth=0, + function_name="type_selector", + root_names=None, + output_names=None, + backend_version=self._backend_version, + returns_scalar=self._returns_scalar, + ) + else: + return self._to_expr() & other + + def __invert__(self: Self) -> DaskSelector: + return DaskSelectorNamespace(backend_version=self._backend_version).all() - self + + def __rsub__(self: Self, other: Any) -> NoReturn: + raise NotImplementedError + + def __rand__(self: Self, other: Any) -> NoReturn: + raise NotImplementedError + + def __ror__(self: Self, other: Any) -> NoReturn: + raise NotImplementedError diff --git a/narwhals/_dask/utils.py b/narwhals/_dask/utils.py index d6eeeb5ea..e7fb64d02 100644 --- a/narwhals/_dask/utils.py +++ b/narwhals/_dask/utils.py @@ -3,10 +3,16 @@ from typing import TYPE_CHECKING from typing import Any -from narwhals.dependencies import get_dask_expr +from narwhals.dependencies import get_pandas +from narwhals.dependencies import get_pyarrow +from narwhals.utils import isinstance_or_issubclass +from narwhals.utils import parse_version if TYPE_CHECKING: + import dask_expr + from narwhals._dask.dataframe import DaskLazyFrame + from narwhals.dtypes import DType def maybe_evaluate(df: DaskLazyFrame, obj: Any) -> Any: @@ -18,21 +24,10 @@ def maybe_evaluate(df: DaskLazyFrame, obj: Any) -> Any: msg = "Multi-output expressions not supported in this context" raise NotImplementedError(msg) result = results[0] - if not get_dask_expr()._expr.are_co_aligned( - df._native_dataframe._expr, result._expr - ): # pragma: no cover - # are_co_aligned is a method which cheaply checks if two Dask expressions - # have the same index, and therefore don't require index alignment. - # If someone only operates on a Dask DataFrame via expressions, then this - # should always be the case: expression outputs (by definition) all come from the - # same input dataframe, and Dask Series does not have any operations which - # change the index. Nonetheless, we perform this safety check anyway. - - # However, we still need to carefully vet which methods we support for Dask, to - # avoid issues where `are_co_aligned` doesn't do what we want it to do: - # https://github.com/dask/dask-expr/issues/1112. - msg = "Implicit index alignment is not support for Dask DataFrame in Narwhals" - raise NotImplementedError(msg) + validate_comparand(df._native_frame, result) + if obj._returns_scalar: + # Return scalar, let Dask do its broadcasting + return result[0] return result return obj @@ -42,13 +37,93 @@ def parse_exprs_and_named_exprs( ) -> dict[str, Any]: results = {} for expr in exprs: - _results = expr._call(df) + if hasattr(expr, "__narwhals_expr__"): + _results = expr._call(df) + elif isinstance(expr, str): + _results = [df._native_frame.loc[:, expr]] + else: # pragma: no cover + msg = f"Expected expression or column name, got: {expr}" + raise TypeError(msg) for _result in _results: - results[_result.name] = _result + if getattr(expr, "_returns_scalar", False): + results[_result.name] = _result[0] + else: + results[_result.name] = _result for name, value in named_exprs.items(): _results = value._call(df) if len(_results) != 1: # pragma: no cover msg = "Named expressions must return a single column" raise AssertionError(msg) - results[name] = _results[0] + for _result in _results: + if getattr(value, "_returns_scalar", False): + results[name] = _result[0] + else: + results[name] = _result return results + + +def add_row_index(frame: Any, name: str) -> Any: + frame = frame.assign(**{name: 1}) + return frame.assign(**{name: frame[name].cumsum(method="blelloch") - 1}) + + +def validate_comparand(lhs: dask_expr.Series, rhs: dask_expr.Series) -> None: + import dask_expr # ignore-banned-import + + if not dask_expr._expr.are_co_aligned(lhs._expr, rhs._expr): # pragma: no cover + # are_co_aligned is a method which cheaply checks if two Dask expressions + # have the same index, and therefore don't require index alignment. + # If someone only operates on a Dask DataFrame via expressions, then this + # should always be the case: expression outputs (by definition) all come from the + # same input dataframe, and Dask Series does not have any operations which + # change the index. Nonetheless, we perform this safety check anyway. + + # However, we still need to carefully vet which methods we support for Dask, to + # avoid issues where `are_co_aligned` doesn't do what we want it to do: + # https://github.com/dask/dask-expr/issues/1112. + msg = "Objects are not co-aligned, so this operation is not supported for Dask backend" + raise RuntimeError(msg) + + +def reverse_translate_dtype(dtype: DType | type[DType]) -> Any: + from narwhals import dtypes + + if isinstance_or_issubclass(dtype, dtypes.Float64): + return "float64" + if isinstance_or_issubclass(dtype, dtypes.Float32): + return "float32" + if isinstance_or_issubclass(dtype, dtypes.Int64): + return "int64" + if isinstance_or_issubclass(dtype, dtypes.Int32): + return "int32" + if isinstance_or_issubclass(dtype, dtypes.Int16): + return "int16" + if isinstance_or_issubclass(dtype, dtypes.Int8): + return "int8" + if isinstance_or_issubclass(dtype, dtypes.UInt64): + return "uint64" + if isinstance_or_issubclass(dtype, dtypes.UInt32): + return "uint32" + if isinstance_or_issubclass(dtype, dtypes.UInt16): + return "uint16" + if isinstance_or_issubclass(dtype, dtypes.UInt8): + return "uint8" + if isinstance_or_issubclass(dtype, dtypes.String): + if (pd := get_pandas()) is not None and parse_version( + pd.__version__ + ) >= parse_version("2.0.0"): + if get_pyarrow() is not None: + return "string[pyarrow]" + return "string[python]" # pragma: no cover + return "object" # pragma: no cover + if isinstance_or_issubclass(dtype, dtypes.Boolean): + return "bool" + if isinstance_or_issubclass(dtype, dtypes.Categorical): + return "category" + if isinstance_or_issubclass(dtype, dtypes.Datetime): + return "datetime64[us]" + if isinstance_or_issubclass(dtype, dtypes.Duration): + return "timedelta64[ns]" + + msg = f"Unknown dtype: {dtype}" # pragma: no cover + raise AssertionError(msg) diff --git a/narwhals/_exceptions.py b/narwhals/_exceptions.py new file mode 100644 index 000000000..189954516 --- /dev/null +++ b/narwhals/_exceptions.py @@ -0,0 +1,4 @@ +from __future__ import annotations + + +class ColumnNotFoundError(Exception): ... diff --git a/narwhals/_expression_parsing.py b/narwhals/_expression_parsing.py index 4c1bb3a22..a74ca3c63 100644 --- a/narwhals/_expression_parsing.py +++ b/narwhals/_expression_parsing.py @@ -11,7 +11,7 @@ from typing import cast from typing import overload -from narwhals.dependencies import get_numpy +from narwhals.dependencies import is_numpy_array from narwhals.utils import flatten if TYPE_CHECKING: @@ -29,17 +29,27 @@ from narwhals._pandas_like.namespace import PandasLikeNamespace from narwhals._pandas_like.series import PandasLikeSeries from narwhals._pandas_like.typing import IntoPandasLikeExpr + from narwhals._polars.expr import PolarsExpr + from narwhals._polars.namespace import PolarsNamespace + from narwhals._polars.series import PolarsSeries + from narwhals._polars.typing import IntoPolarsExpr - CompliantNamespace = Union[PandasLikeNamespace, ArrowNamespace, DaskNamespace] - CompliantExpr = Union[PandasLikeExpr, ArrowExpr, DaskExpr] - IntoCompliantExpr = Union[IntoPandasLikeExpr, IntoArrowExpr, IntoDaskExpr] + CompliantNamespace = Union[ + PandasLikeNamespace, ArrowNamespace, DaskNamespace, PolarsNamespace + ] + CompliantExpr = Union[PandasLikeExpr, ArrowExpr, DaskExpr, PolarsExpr] + IntoCompliantExpr = Union[ + IntoPandasLikeExpr, IntoArrowExpr, IntoDaskExpr, IntoPolarsExpr + ] IntoCompliantExprT = TypeVar("IntoCompliantExprT", bound=IntoCompliantExpr) CompliantExprT = TypeVar("CompliantExprT", bound=CompliantExpr) - CompliantSeries = Union[PandasLikeSeries, ArrowSeries] + CompliantSeries = Union[PandasLikeSeries, ArrowSeries, PolarsSeries] ListOfCompliantSeries = Union[ - list[PandasLikeSeries], list[ArrowSeries], list[DaskExpr] + list[PandasLikeSeries], list[ArrowSeries], list[DaskExpr], list[PolarsSeries] + ] + ListOfCompliantExpr = Union[ + list[PandasLikeExpr], list[ArrowExpr], list[DaskExpr], list[PolarsExpr] ] - ListOfCompliantExpr = Union[list[PandasLikeExpr], list[ArrowExpr], list[DaskExpr]] CompliantDataFrame = Union[PandasLikeDataFrame, ArrowDataFrame, DaskLazyFrame] T = TypeVar("T") @@ -93,7 +103,8 @@ def evaluate_into_exprs( if len(evaluated_expr) > 1: msg = "Named expressions must return a single column" # pragma: no cover raise AssertionError(msg) - series.append(evaluated_expr[0].alias(name)) # type: ignore[arg-type] + to_append = evaluated_expr[0].alias(name) + series.append(to_append) # type: ignore[arg-type] return series @@ -123,6 +134,22 @@ def parse_into_exprs( ) -> list[ArrowExpr]: ... +@overload +def parse_into_exprs( + *exprs: IntoDaskExpr, + namespace: DaskNamespace, + **named_exprs: IntoDaskExpr, +) -> list[DaskExpr]: ... + + +@overload +def parse_into_exprs( + *exprs: IntoPolarsExpr, + namespace: PolarsNamespace, + **named_exprs: IntoPolarsExpr, +) -> list[PolarsExpr]: ... + + def parse_into_exprs( *exprs: IntoCompliantExpr, namespace: CompliantNamespace, @@ -130,12 +157,12 @@ def parse_into_exprs( ) -> ListOfCompliantExpr: """Parse each input as an expression (if it's not already one). See `parse_into_expr` for more details.""" - out = [ + return [ parse_into_expr(into_expr, namespace=namespace) for into_expr in flatten(exprs) + ] + [ + parse_into_expr(expr, namespace=namespace).alias(name) + for name, expr in named_exprs.items() ] - for name, expr in named_exprs.items(): - out.append(parse_into_expr(expr, namespace=namespace).alias(name)) - return out # type: ignore[return-value] def parse_into_expr( @@ -161,11 +188,16 @@ def parse_into_expr( return namespace._create_expr_from_series(into_expr) # type: ignore[arg-type] if isinstance(into_expr, str): return namespace.col(into_expr) - if (np := get_numpy()) is not None and isinstance(into_expr, np.ndarray): + if is_numpy_array(into_expr): series = namespace._create_compliant_series(into_expr) return namespace._create_expr_from_series(series) # type: ignore[arg-type] - msg = f"Expected IntoExpr, got {type(into_expr)}" # pragma: no cover - raise AssertionError(msg) + msg = ( + f"Expected an object which can be converted into an expression, got {type(into_expr)}\n\n" # pragma: no cover + "Hint: if you were trying to select a column which does not have a string column name, then " + "you should explicitly use `nw.col`.\nFor example, `df.select(nw.col(0))` if you have a column " + "named `0`." + ) + raise TypeError(msg) def reuse_series_implementation( diff --git a/narwhals/_interchange/dataframe.py b/narwhals/_interchange/dataframe.py index c7fc5ea3d..bf1b17243 100644 --- a/narwhals/_interchange/dataframe.py +++ b/narwhals/_interchange/dataframe.py @@ -69,7 +69,7 @@ def map_interchange_dtype_to_narwhals_dtype( class InterchangeFrame: def __init__(self, df: Any) -> None: - self._native_dataframe = df + self._native_frame = df def __narwhals_dataframe__(self) -> Any: return self @@ -77,15 +77,15 @@ def __narwhals_dataframe__(self) -> Any: def __getitem__(self, item: str) -> InterchangeSeries: from narwhals._interchange.series import InterchangeSeries - return InterchangeSeries(self._native_dataframe.get_column_by_name(item)) + return InterchangeSeries(self._native_frame.get_column_by_name(item)) @property def schema(self) -> dict[str, dtypes.DType]: return { column_name: map_interchange_dtype_to_narwhals_dtype( - self._native_dataframe.get_column_by_name(column_name).dtype + self._native_frame.get_column_by_name(column_name).dtype ) - for column_name in self._native_dataframe.column_names() + for column_name in self._native_frame.column_names() } def __getattr__(self, attr: str) -> NoReturn: diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py index 215886160..193955cbd 100644 --- a/narwhals/_pandas_like/dataframe.py +++ b/narwhals/_pandas_like/dataframe.py @@ -1,6 +1,5 @@ from __future__ import annotations -import collections from typing import TYPE_CHECKING from typing import Any from typing import Iterable @@ -11,20 +10,23 @@ from narwhals._expression_parsing import evaluate_into_exprs from narwhals._pandas_like.expr import PandasLikeExpr +from narwhals._pandas_like.utils import broadcast_series from narwhals._pandas_like.utils import create_native_series from narwhals._pandas_like.utils import horizontal_concat from narwhals._pandas_like.utils import translate_dtype from narwhals._pandas_like.utils import validate_dataframe_comparand -from narwhals._pandas_like.utils import validate_indices from narwhals.dependencies import get_cudf from narwhals.dependencies import get_modin -from narwhals.dependencies import get_numpy from narwhals.dependencies import get_pandas +from narwhals.dependencies import is_numpy_array from narwhals.utils import Implementation from narwhals.utils import flatten from narwhals.utils import generate_unique_token +from narwhals.utils import parse_columns_to_drop if TYPE_CHECKING: + import numpy as np + import pandas as pd from typing_extensions import Self from narwhals._pandas_like.group_by import PandasLikeGroupBy @@ -44,7 +46,7 @@ def __init__( backend_version: tuple[int, ...], ) -> None: self._validate_columns(native_dataframe.columns) - self._native_dataframe = native_dataframe + self._native_frame = native_dataframe self._implementation = implementation self._backend_version = backend_version @@ -70,19 +72,20 @@ def __native_namespace__(self) -> Any: raise AssertionError(msg) def __len__(self) -> int: - return len(self._native_dataframe) - - def _validate_columns(self, columns: Sequence[str]) -> None: - if len(columns) != len(set(columns)): - counter = collections.Counter(columns) - for col, count in counter.items(): - if count > 1: - msg = f"Expected unique column names, got {col!r} {count} time(s)" - raise ValueError(msg) - msg = "Please report a bug" # pragma: no cover - raise AssertionError(msg) - - def _from_native_dataframe(self, df: Any) -> Self: + return len(self._native_frame) + + def _validate_columns(self, columns: pd.Index) -> None: + try: + len_unique_columns = len(columns.drop_duplicates()) + except Exception: # noqa: BLE001 # pragma: no cover + msg = f"Expected hashable (e.g. str or int) column names, got: {columns}" + raise ValueError(msg) from None + + if len(columns) != len_unique_columns: + msg = f"Expected unique column names, got: {columns}" + raise ValueError(msg) + + def _from_native_frame(self, df: Any) -> Self: return self.__class__( df, implementation=self._implementation, @@ -93,11 +96,14 @@ def get_column(self, name: str) -> PandasLikeSeries: from narwhals._pandas_like.series import PandasLikeSeries return PandasLikeSeries( - self._native_dataframe.loc[:, name], + self._native_frame.loc[:, name], implementation=self._implementation, backend_version=self._backend_version, ) + def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray: + return self.to_numpy(dtype=dtype, copy=copy) + @overload def __getitem__(self, item: tuple[Sequence[int], str | int]) -> PandasLikeSeries: ... # type: ignore[overload-overlap] @@ -117,7 +123,7 @@ def __getitem__( from narwhals._pandas_like.series import PandasLikeSeries return PandasLikeSeries( - self._native_dataframe.loc[:, item], + self._native_frame.loc[:, item], implementation=self._implementation, backend_version=self._backend_version, ) @@ -128,26 +134,50 @@ def __getitem__( and isinstance(item[1], (tuple, list)) ): if all(isinstance(x, int) for x in item[1]): - return self._from_native_dataframe(self._native_dataframe.iloc[item]) + return self._from_native_frame(self._native_frame.iloc[item]) if all(isinstance(x, str) for x in item[1]): item = ( item[0], - self._native_dataframe.columns.get_indexer(item[1]), + self._native_frame.columns.get_indexer(item[1]), ) - return self._from_native_dataframe(self._native_dataframe.iloc[item]) + return self._from_native_frame(self._native_frame.iloc[item]) msg = ( f"Expected sequence str or int, got: {type(item[1])}" # pragma: no cover ) raise TypeError(msg) # pragma: no cover + elif isinstance(item, tuple) and len(item) == 2 and isinstance(item[1], slice): + columns = self._native_frame.columns + if isinstance(item[1].start, str) or isinstance(item[1].stop, str): + start = ( + columns.get_loc(item[1].start) if item[1].start is not None else None + ) + stop = ( + columns.get_loc(item[1].stop) + 1 + if item[1].stop is not None + else None + ) + step = item[1].step + return self._from_native_frame( + self._native_frame.iloc[item[0], slice(start, stop, step)] + ) + if isinstance(item[1].start, int) or isinstance(item[1].stop, int): + return self._from_native_frame( + self._native_frame.iloc[ + item[0], slice(item[1].start, item[1].stop, item[1].step) + ] + ) + msg = f"Expected slice of integers or strings, got: {type(item[1])}" # pragma: no cover + raise TypeError(msg) # pragma: no cover + elif isinstance(item, tuple) and len(item) == 2: from narwhals._pandas_like.series import PandasLikeSeries if isinstance(item[1], str): - item = (item[0], self._native_dataframe.columns.get_loc(item[1])) - native_series = self._native_dataframe.iloc[item] + item = (item[0], self._native_frame.columns.get_loc(item[1])) + native_series = self._native_frame.iloc[item] elif isinstance(item[1], int): - native_series = self._native_dataframe.iloc[item] + native_series = self._native_frame.iloc[item] else: # pragma: no cover msg = f"Expected str or int, got: {type(item[1])}" raise TypeError(msg) @@ -159,11 +189,9 @@ def __getitem__( ) elif isinstance(item, (slice, Sequence)) or ( - (np := get_numpy()) is not None - and isinstance(item, np.ndarray) - and item.ndim == 1 + is_numpy_array(item) and item.ndim == 1 ): - return self._from_native_dataframe(self._native_dataframe.iloc[item]) + return self._from_native_frame(self._native_frame.iloc[item]) else: # pragma: no cover msg = f"Expected str or slice, got: {type(item)}" @@ -172,15 +200,15 @@ def __getitem__( # --- properties --- @property def columns(self) -> list[str]: - return self._native_dataframe.columns.tolist() # type: ignore[no-any-return] + return self._native_frame.columns.tolist() # type: ignore[no-any-return] def rows( self, *, named: bool = False ) -> list[tuple[Any, ...]] | list[dict[str, Any]]: if not named: - return list(self._native_dataframe.itertuples(index=False, name=None)) + return list(self._native_frame.itertuples(index=False, name=None)) - return self._native_dataframe.to_dict(orient="records") # type: ignore[no-any-return] + return self._native_frame.to_dict(orient="records") # type: ignore[no-any-return] def iter_rows( self, @@ -194,19 +222,19 @@ def iter_rows( and has no effect on the output. """ if not named: - yield from self._native_dataframe.itertuples(index=False, name=None) + yield from self._native_frame.itertuples(index=False, name=None) else: - col_names = self._native_dataframe.columns + col_names = self._native_frame.columns yield from ( dict(zip(col_names, row)) - for row in self._native_dataframe.itertuples(index=False) + for row in self._native_frame.itertuples(index=False) ) # type: ignore[misc] @property def schema(self) -> dict[str, DType]: return { - col: translate_dtype(self._native_dataframe.loc[:, col]) - for col in self._native_dataframe.columns + col: translate_dtype(self._native_frame.loc[:, col]) + for col in self._native_frame.columns } def collect_schema(self) -> dict[str, DType]: @@ -220,57 +248,73 @@ def select( ) -> Self: if exprs and all(isinstance(x, str) for x in exprs) and not named_exprs: # This is a simple slice => fastpath! - return self._from_native_dataframe(self._native_dataframe.loc[:, exprs]) + return self._from_native_frame(self._native_frame.loc[:, list(exprs)]) new_series = evaluate_into_exprs(self, *exprs, **named_exprs) if not new_series: # return empty dataframe, like Polars does - return self._from_native_dataframe(self._native_dataframe.__class__()) - new_series = validate_indices(new_series) + return self._from_native_frame(self._native_frame.__class__()) + new_series = broadcast_series(new_series) df = horizontal_concat( new_series, implementation=self._implementation, backend_version=self._backend_version, ) - return self._from_native_dataframe(df) + return self._from_native_frame(df) - def drop_nulls(self) -> Self: - return self._from_native_dataframe(self._native_dataframe.dropna(axis=0)) + def drop_nulls(self, subset: str | list[str] | None) -> Self: + if subset is None: + return self._from_native_frame(self._native_frame.dropna(axis=0)) + subset = [subset] if isinstance(subset, str) else subset + plx = self.__narwhals_namespace__() + return self.filter(~plx.any_horizontal(plx.col(*subset).is_null())) def with_row_index(self, name: str) -> Self: row_index = create_native_series( - range(len(self._native_dataframe)), - index=self._native_dataframe.index, + range(len(self._native_frame)), + index=self._native_frame.index, implementation=self._implementation, backend_version=self._backend_version, ).alias(name) - return self._from_native_dataframe( + return self._from_native_frame( horizontal_concat( - [row_index._native_series, self._native_dataframe], + [row_index._native_series, self._native_frame], implementation=self._implementation, backend_version=self._backend_version, ) ) + def row(self, row: int) -> tuple[Any, ...]: + return tuple(x for x in self._native_frame.iloc[row]) + def filter( self, *predicates: IntoPandasLikeExpr, ) -> Self: - from narwhals._pandas_like.namespace import PandasLikeNamespace - - plx = PandasLikeNamespace(self._implementation, self._backend_version) - expr = plx.all_horizontal(*predicates) - # Safety: all_horizontal's expression only returns a single column. - mask = expr._call(self)[0] - _mask = validate_dataframe_comparand(self._native_dataframe.index, mask) - return self._from_native_dataframe(self._native_dataframe.loc[_mask]) + plx = self.__narwhals_namespace__() + if ( + len(predicates) == 1 + and isinstance(predicates[0], list) + and all(isinstance(x, bool) for x in predicates[0]) + ): + _mask = predicates[0] + else: + expr = plx.all_horizontal(*predicates) + # Safety: all_horizontal's expression only returns a single column. + mask = expr._call(self)[0] + _mask = validate_dataframe_comparand(self._native_frame.index, mask) + return self._from_native_frame(self._native_frame.loc[_mask]) def with_columns( self, *exprs: IntoPandasLikeExpr, **named_exprs: IntoPandasLikeExpr, ) -> Self: - index = self._native_dataframe.index + index = self._native_frame.index new_columns = evaluate_into_exprs(self, *exprs, **named_exprs) + + if not new_columns and len(self) == 0: + return self + # If the inputs are all Expressions which return full columns # (as opposed to scalars), we can use a fast path (concat, instead of assign). # We can't use the fastpath if any input is not an expression (e.g. @@ -286,7 +330,7 @@ def with_columns( new_column_name_to_new_column_map = {s.name: s for s in new_columns} to_concat = [] # Make sure to preserve column order - for name in self._native_dataframe.columns: + for name in self._native_frame.columns: if name in new_column_name_to_new_column_map: to_concat.append( validate_dataframe_comparand( @@ -294,7 +338,7 @@ def with_columns( ) ) else: - to_concat.append(self._native_dataframe.loc[:, name]) + to_concat.append(self._native_frame.loc[:, name]) to_concat.extend( validate_dataframe_comparand(index, new_column_name_to_new_column_map[s]) for s in new_column_name_to_new_column_map @@ -306,18 +350,19 @@ def with_columns( backend_version=self._backend_version, ) else: - df = self._native_dataframe.copy(deep=False) + df = self._native_frame.copy(deep=False) for s in new_columns: df[s.name] = validate_dataframe_comparand(index, s) - return self._from_native_dataframe(df) + return self._from_native_frame(df) def rename(self, mapping: dict[str, str]) -> Self: - return self._from_native_dataframe(self._native_dataframe.rename(columns=mapping)) + return self._from_native_frame(self._native_frame.rename(columns=mapping)) - def drop(self, *columns: str) -> Self: - return self._from_native_dataframe( - self._native_dataframe.drop(columns=list(columns)) + def drop(self: Self, columns: list[str], strict: bool) -> Self: # noqa: FBT001 + to_drop = parse_columns_to_drop( + compliant_frame=self, columns=columns, strict=strict ) + return self._from_native_frame(self._native_frame.drop(columns=to_drop)) # --- transform --- def sort( @@ -327,17 +372,17 @@ def sort( descending: bool | Sequence[bool] = False, ) -> Self: flat_keys = flatten([*flatten([by]), *more_by]) - df = self._native_dataframe + df = self._native_frame if isinstance(descending, bool): ascending: bool | list[bool] = not descending else: ascending = [not d for d in descending] - return self._from_native_dataframe(df.sort_values(flat_keys, ascending=ascending)) + return self._from_native_frame(df.sort_values(flat_keys, ascending=ascending)) # --- convert --- def collect(self) -> PandasLikeDataFrame: return PandasLikeDataFrame( - self._native_dataframe, + self._native_frame, implementation=self._implementation, backend_version=self._backend_version, ) @@ -376,19 +421,21 @@ def join( n_bytes=8, columns=[*self.columns, *other.columns] ) - return self._from_native_dataframe( - self._native_dataframe.assign(**{key_token: 0}).merge( - other._native_dataframe.assign(**{key_token: 0}), + return self._from_native_frame( + self._native_frame.assign(**{key_token: 0}) + .merge( + other._native_frame.assign(**{key_token: 0}), how="inner", left_on=key_token, right_on=key_token, suffixes=("", "_right"), - ), - ).drop(key_token) + ) + .drop(columns=key_token), + ) else: - return self._from_native_dataframe( - self._native_dataframe.merge( - other._native_dataframe, + return self._from_native_frame( + self._native_frame.merge( + other._native_frame, how="cross", suffixes=("", "_right"), ), @@ -400,14 +447,14 @@ def join( ) other_native = ( - other._native_dataframe.loc[:, right_on] + other._native_frame.loc[:, right_on] .rename( # rename to avoid creating extra columns in join columns=dict(zip(right_on, left_on)) # type: ignore[arg-type] ) .drop_duplicates() ) - return self._from_native_dataframe( - self._native_dataframe.merge( + return self._from_native_frame( + self._native_frame.merge( other_native, how="outer", indicator=indicator_token, @@ -415,19 +462,19 @@ def join( right_on=left_on, ) .loc[lambda t: t[indicator_token] == "left_only"] - .drop(columns=[indicator_token]) + .drop(columns=indicator_token) ) if how == "semi": other_native = ( - other._native_dataframe.loc[:, right_on] + other._native_frame.loc[:, right_on] .rename( # rename to avoid creating extra columns in join columns=dict(zip(right_on, left_on)) # type: ignore[arg-type] ) .drop_duplicates() # avoids potential rows duplication from inner join ) - return self._from_native_dataframe( - self._native_dataframe.merge( + return self._from_native_frame( + self._native_frame.merge( other_native, how="inner", left_on=left_on, @@ -436,8 +483,8 @@ def join( ) if how == "left": - other_native = other._native_dataframe - result_native = self._native_dataframe.merge( + other_native = other._native_frame + result_native = self._native_frame.merge( other_native, how="left", left_on=left_on, @@ -450,11 +497,11 @@ def join( extra.append(right_key) elif right_key != left_key: extra.append(f"{right_key}_right") - return self._from_native_dataframe(result_native.drop(columns=extra)) + return self._from_native_frame(result_native.drop(columns=extra)) - return self._from_native_dataframe( - self._native_dataframe.merge( - other._native_dataframe, + return self._from_native_frame( + self._native_frame.merge( + other._native_frame, left_on=left_on, right_on=right_on, how=how, @@ -465,10 +512,10 @@ def join( # --- partial reduction --- def head(self, n: int) -> Self: - return self._from_native_dataframe(self._native_dataframe.head(n)) + return self._from_native_frame(self._native_frame.head(n)) def tail(self, n: int) -> Self: - return self._from_native_dataframe(self._native_dataframe.tail(n)) + return self._from_native_frame(self._native_frame.tail(n)) def unique( self: Self, @@ -482,11 +529,10 @@ def unique( The param `maintain_order` is only here for compatibility with the polars API and has no effect on the output. """ - mapped_keep = {"none": False, "any": "first"}.get(keep, keep) subset = flatten(subset) if subset else None - return self._from_native_dataframe( - self._native_dataframe.drop_duplicates(subset=subset, keep=mapped_keep) + return self._from_native_frame( + self._native_frame.drop_duplicates(subset=subset, keep=mapped_keep) ) # --- lazy-only --- @@ -495,7 +541,7 @@ def lazy(self) -> Self: @property def shape(self) -> tuple[int, int]: - return self._native_dataframe.shape # type: ignore[no-any-return] + return self._native_frame.shape # type: ignore[no-any-return] def to_dict(self, *, as_series: bool = False) -> dict[str, Any]: from narwhals._pandas_like.series import PandasLikeSeries @@ -504,63 +550,75 @@ def to_dict(self, *, as_series: bool = False) -> dict[str, Any]: # TODO(Unassigned): should this return narwhals series? return { col: PandasLikeSeries( - self._native_dataframe.loc[:, col], + self._native_frame.loc[:, col], implementation=self._implementation, backend_version=self._backend_version, ) for col in self.columns } - return self._native_dataframe.to_dict(orient="list") # type: ignore[no-any-return] + return self._native_frame.to_dict(orient="list") # type: ignore[no-any-return] - def to_numpy(self) -> Any: + def to_numpy(self, dtype: Any = None, copy: bool | None = None) -> Any: from narwhals._pandas_like.series import PANDAS_TO_NUMPY_DTYPE_MISSING - # pandas return `object` dtype for nullable dtypes, so we cast each - # Series to numpy and let numpy find a common dtype. + if copy is None: + # pandas default differs from Polars + copy = False + + if dtype is not None: + return self._native_frame.to_numpy(dtype=dtype, copy=copy) + + # pandas return `object` dtype for nullable dtypes if dtype=None, + # so we cast each Series to numpy and let numpy find a common dtype. # If there aren't any dtypes where `to_numpy()` is "broken" (i.e. it # returns Object) then we just call `to_numpy()` on the DataFrame. - for dtype in self._native_dataframe.dtypes: - if str(dtype) in PANDAS_TO_NUMPY_DTYPE_MISSING: - import numpy as np + for col_dtype in self._native_frame.dtypes: + if str(col_dtype) in PANDAS_TO_NUMPY_DTYPE_MISSING: + import numpy as np # ignore-banned-import - return np.hstack([self[col].to_numpy()[:, None] for col in self.columns]) - return self._native_dataframe.to_numpy() + return np.hstack( + [self[col].to_numpy(copy=copy)[:, None] for col in self.columns] + ) + return self._native_frame.to_numpy(copy=copy) def to_pandas(self) -> Any: if self._implementation is Implementation.PANDAS: - return self._native_dataframe + return self._native_frame if self._implementation is Implementation.MODIN: # pragma: no cover - return self._native_dataframe._to_pandas() - return self._native_dataframe.to_pandas() # pragma: no cover + return self._native_frame._to_pandas() + return self._native_frame.to_pandas() # pragma: no cover def write_parquet(self, file: Any) -> Any: - self._native_dataframe.to_parquet(file) + self._native_frame.to_parquet(file) + + def write_csv(self, file: Any = None) -> Any: + return self._native_frame.to_csv(file, index=False) # --- descriptive --- def is_duplicated(self: Self) -> PandasLikeSeries: from narwhals._pandas_like.series import PandasLikeSeries return PandasLikeSeries( - self._native_dataframe.duplicated(keep=False), + self._native_frame.duplicated(keep=False), implementation=self._implementation, backend_version=self._backend_version, ) def is_empty(self: Self) -> bool: - return self._native_dataframe.empty # type: ignore[no-any-return] + return self._native_frame.empty # type: ignore[no-any-return] def is_unique(self: Self) -> PandasLikeSeries: from narwhals._pandas_like.series import PandasLikeSeries return PandasLikeSeries( - ~self._native_dataframe.duplicated(keep=False), + ~self._native_frame.duplicated(keep=False), implementation=self._implementation, backend_version=self._backend_version, ) def null_count(self: Self) -> PandasLikeDataFrame: return PandasLikeDataFrame( - self._native_dataframe.isna().sum(axis=0).to_frame().transpose(), + self._native_frame.isna().sum(axis=0).to_frame().transpose(), implementation=self._implementation, backend_version=self._backend_version, ) @@ -574,17 +632,26 @@ def item(self: Self, row: int | None = None, column: int | str | None = None) -> f" frame has shape {self.shape!r}" ) raise ValueError(msg) - return self._native_dataframe.iloc[0, 0] + return self._native_frame.iloc[0, 0] elif row is None or column is None: msg = "cannot call `.item()` with only one of `row` or `column`" raise ValueError(msg) _col = self.columns.index(column) if isinstance(column, str) else column - return self._native_dataframe.iloc[row, _col] + return self._native_frame.iloc[row, _col] def clone(self: Self) -> Self: - return self._from_native_dataframe(self._native_dataframe.copy()) + return self._from_native_frame(self._native_frame.copy()) def gather_every(self: Self, n: int, offset: int = 0) -> Self: - return self._from_native_dataframe(self._native_dataframe.iloc[offset::n]) + return self._from_native_frame(self._native_frame.iloc[offset::n]) + + def to_arrow(self: Self) -> Any: + if self._implementation is Implementation.CUDF: # pragma: no cover + msg = "`to_arrow` is not implemented for CuDF backend." + raise NotImplementedError(msg) + + import pyarrow as pa # ignore-banned-import() + + return pa.Table.from_pandas(self._native_frame) diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py index cb6211dd0..44154453d 100644 --- a/narwhals/_pandas_like/expr.py +++ b/narwhals/_pandas_like/expr.py @@ -64,7 +64,7 @@ def from_column_names( def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]: return [ PandasLikeSeries( - df._native_dataframe.loc[:, column_name], + df._native_frame.loc[:, column_name], implementation=df._implementation, backend_version=df._backend_version, ) @@ -196,6 +196,12 @@ def min(self) -> Self: return reuse_series_implementation(self, "min", returns_scalar=True) # Other + + def clip(self, lower_bound: Any, upper_bound: Any) -> Self: + return reuse_series_implementation( + self, "clip", lower_bound=lower_bound, upper_bound=upper_bound + ) + def is_between( self, lower_bound: Any, upper_bound: Any, closed: str = "both" ) -> Self: @@ -220,11 +226,9 @@ def arg_true(self) -> Self: return reuse_series_implementation(self, "arg_true") def filter(self, *predicates: Any) -> Self: - from narwhals._pandas_like.namespace import PandasLikeNamespace - - plx = PandasLikeNamespace(self._implementation, self._backend_version) - expr = plx.all_horizontal(*predicates) - return reuse_series_implementation(self, "filter", other=expr) + plx = self.__narwhals_namespace__() + other = plx.all_horizontal(*predicates) + return reuse_series_implementation(self, "filter", other=other) def drop_nulls(self) -> Self: return reuse_series_implementation(self, "drop_nulls") @@ -365,6 +369,37 @@ class PandasLikeExprStringNamespace: def __init__(self, expr: PandasLikeExpr) -> None: self._expr = expr + def replace( + self, + pattern: str, + value: str, + *, + literal: bool = False, + n: int = 1, + ) -> PandasLikeExpr: + return reuse_series_namespace_implementation( + self._expr, "str", "replace", pattern, value, literal=literal, n=n + ) + + def replace_all( + self, + pattern: str, + value: str, + *, + literal: bool = False, + ) -> PandasLikeExpr: + return reuse_series_namespace_implementation( + self._expr, "str", "replace_all", pattern, value, literal=literal + ) + + def strip_chars(self, characters: str | None = None) -> PandasLikeExpr: + return reuse_series_namespace_implementation( + self._expr, + "str", + "strip_chars", + characters, + ) + def starts_with(self, prefix: str) -> PandasLikeExpr: return reuse_series_namespace_implementation( self._expr, @@ -422,6 +457,9 @@ class PandasLikeExprDateTimeNamespace: def __init__(self, expr: PandasLikeExpr) -> None: self._expr = expr + def date(self) -> PandasLikeExpr: + return reuse_series_namespace_implementation(self._expr, "dt", "date") + def year(self) -> PandasLikeExpr: return reuse_series_namespace_implementation(self._expr, "dt", "year") diff --git a/narwhals/_pandas_like/group_by.py b/narwhals/_pandas_like/group_by.py index 129b047b5..11abc85c8 100644 --- a/narwhals/_pandas_like/group_by.py +++ b/narwhals/_pandas_like/group_by.py @@ -32,16 +32,16 @@ def __init__(self, df: PandasLikeDataFrame, keys: list[str]) -> None: self._df._implementation is Implementation.PANDAS and self._df._backend_version < (1, 0) ): # pragma: no cover - if self._df._native_dataframe.loc[:, self._keys].isna().any().any(): + if self._df._native_frame.loc[:, self._keys].isna().any().any(): msg = "Grouping by null values is not supported in pandas < 1.0.0" raise NotImplementedError(msg) - self._grouped = self._df._native_dataframe.groupby( + self._grouped = self._df._native_frame.groupby( list(self._keys), sort=False, as_index=True, ) else: - self._grouped = self._df._native_dataframe.groupby( + self._grouped = self._df._native_frame.groupby( list(self._keys), sort=False, as_index=True, @@ -75,13 +75,13 @@ def agg( exprs, self._keys, output_names, - self._from_native_dataframe, - dataframe_is_empty=self._df._native_dataframe.empty, + self._from_native_frame, + dataframe_is_empty=self._df._native_frame.empty, implementation=implementation, backend_version=self._df._backend_version, ) - def _from_native_dataframe(self, df: PandasLikeDataFrame) -> PandasLikeDataFrame: + def _from_native_frame(self, df: PandasLikeDataFrame) -> PandasLikeDataFrame: from narwhals._pandas_like.dataframe import PandasLikeDataFrame return PandasLikeDataFrame( @@ -100,9 +100,7 @@ def __iter__(self) -> Iterator[tuple[Any, PandasLikeDataFrame]]: category=FutureWarning, ) iterator = self._grouped.__iter__() - yield from ( - (key, self._from_native_dataframe(sub_df)) for (key, sub_df) in iterator - ) + yield from ((key, self._from_native_frame(sub_df)) for (key, sub_df) in iterator) def agg_pandas( diff --git a/narwhals/_pandas_like/namespace.py b/narwhals/_pandas_like/namespace.py index 360c18a1d..753b49f69 100644 --- a/narwhals/_pandas_like/namespace.py +++ b/narwhals/_pandas_like/namespace.py @@ -5,6 +5,7 @@ from typing import Any from typing import Callable from typing import Iterable +from typing import cast from narwhals import dtypes from narwhals._expression_parsing import parse_into_exprs @@ -15,7 +16,6 @@ from narwhals._pandas_like.utils import create_native_series from narwhals._pandas_like.utils import horizontal_concat from narwhals._pandas_like.utils import vertical_concat -from narwhals.utils import flatten if TYPE_CHECKING: from narwhals._pandas_like.typing import IntoPandasLikeExpr @@ -127,7 +127,7 @@ def all(self) -> PandasLikeExpr: return PandasLikeExpr( lambda df: [ PandasLikeSeries( - df._native_dataframe.loc[:, column_name], + df._native_frame.loc[:, column_name], implementation=self._implementation, backend_version=self._backend_version, ) @@ -146,7 +146,7 @@ def _lit_pandas_series(df: PandasLikeDataFrame) -> PandasLikeSeries: pandas_series = PandasLikeSeries._from_iterable( data=[value], name="lit", - index=df._native_dataframe.index[0:1], + index=df._native_frame.index[0:1], implementation=self._implementation, backend_version=self._backend_version, ) @@ -197,7 +197,7 @@ def len(self) -> PandasLikeExpr: return PandasLikeExpr( lambda df: [ PandasLikeSeries._from_iterable( - [len(df._native_dataframe)], + [len(df._native_frame)], name="len", index=[0], implementation=self._implementation, @@ -216,23 +216,22 @@ def len(self) -> PandasLikeExpr: def sum_horizontal(self, *exprs: IntoPandasLikeExpr) -> PandasLikeExpr: return reduce( lambda x, y: x + y, - parse_into_exprs( - *exprs, - namespace=self, - ), + [expr.fill_null(0) for expr in parse_into_exprs(*exprs, namespace=self)], ) def all_horizontal(self, *exprs: IntoPandasLikeExpr) -> PandasLikeExpr: - return reduce( - lambda x, y: x & y, - parse_into_exprs(*exprs, namespace=self), - ) + return reduce(lambda x, y: x & y, parse_into_exprs(*exprs, namespace=self)) def any_horizontal(self, *exprs: IntoPandasLikeExpr) -> PandasLikeExpr: - return reduce( - lambda x, y: x | y, - parse_into_exprs(*exprs, namespace=self), + return reduce(lambda x, y: x | y, parse_into_exprs(*exprs, namespace=self)) + + def mean_horizontal(self, *exprs: IntoPandasLikeExpr) -> PandasLikeExpr: + pandas_like_exprs = parse_into_exprs(*exprs, namespace=self) + total = reduce(lambda x, y: x + y, (e.fill_null(0.0) for e in pandas_like_exprs)) + n_non_zero = reduce( + lambda x, y: x + y, ((1 - e.is_null()) for e in pandas_like_exprs) ) + return total / n_non_zero def concat( self, @@ -240,7 +239,7 @@ def concat( *, how: str = "vertical", ) -> PandasLikeDataFrame: - dfs: list[Any] = [item._native_dataframe for item in items] + dfs: list[Any] = [item._native_frame for item in items] if how == "horizontal": return PandasLikeDataFrame( horizontal_concat( @@ -265,26 +264,16 @@ def concat( def when( self, - *predicates: IntoPandasLikeExpr | Iterable[IntoPandasLikeExpr], + *predicates: IntoPandasLikeExpr, ) -> PandasWhen: - return PandasWhen( - when_processing(self, *predicates), - self._implementation, - self._backend_version, - ) - + plx = self.__class__(self._implementation, self._backend_version) + if predicates: + condition = plx.all_horizontal(*predicates) + else: + msg = "at least one predicate needs to be provided" + raise TypeError(msg) -def when_processing( - plx: PandasLikeNamespace, - *predicates: IntoPandasLikeExpr | Iterable[IntoPandasLikeExpr], -) -> PandasLikeExpr: - if predicates: - condition = plx.all_horizontal(*flatten(predicates)) - else: - msg = "at least one predicate needs to be provided" - raise TypeError(msg) - - return condition + return PandasWhen(condition, self._implementation, self._backend_version) class PandasWhen: @@ -294,33 +283,61 @@ def __init__( implementation: Implementation, backend_version: tuple[int, ...], then_value: Any = None, - otherise_value: Any = None, + otherwise_value: Any = None, ) -> None: self._implementation = implementation self._backend_version = backend_version self._condition = condition self._then_value = then_value - self._otherwise_value = otherise_value - self._already_set = self._condition + self._otherwise_value = otherwise_value def __call__(self, df: PandasLikeDataFrame) -> list[PandasLikeSeries]: + from narwhals._expression_parsing import parse_into_expr from narwhals._pandas_like.namespace import PandasLikeNamespace + from narwhals._pandas_like.utils import validate_column_comparand plx = PandasLikeNamespace( implementation=self._implementation, backend_version=self._backend_version ) - condition = self._condition._call(df)[0] + condition = parse_into_expr(self._condition, namespace=plx)._call(df)[0] # type: ignore[arg-type] + try: + value_series = parse_into_expr(self._then_value, namespace=plx)._call(df)[0] # type: ignore[arg-type] + except TypeError: + # `self._otherwise_value` is a scalar and can't be converted to an expression + value_series = condition.__class__._from_iterable( # type: ignore[call-arg] + [self._then_value] * len(condition), + name="literal", + index=condition._native_series.index, + implementation=self._implementation, + backend_version=self._backend_version, + ) + value_series = cast(PandasLikeSeries, value_series) - value_series = plx._create_broadcast_series_from_scalar( - self._then_value, condition - ) - otherwise_series = plx._create_broadcast_series_from_scalar( - self._otherwise_value, condition - ) - return [value_series.zip_with(condition, otherwise_series)] + value_series_native = value_series._native_series + condition_native = validate_column_comparand(value_series_native.index, condition) - def then(self, value: Any) -> PandasThen: + if self._otherwise_value is None: + return [ + value_series._from_native_series( + value_series_native.where(condition_native) + ) + ] + try: + otherwise_series = parse_into_expr( + self._otherwise_value, namespace=plx + )._call(df)[0] # type: ignore[arg-type] + except TypeError: + # `self._otherwise_value` is a scalar and can't be converted to an expression + return [ + value_series._from_native_series( + value_series_native.where(condition_native, self._otherwise_value) + ) + ] + else: + return [value_series.zip_with(condition, otherwise_series)] + + def then(self, value: PandasLikeExpr | PandasLikeSeries | Any) -> PandasThen: self._then_value = value return PandasThen( @@ -329,8 +346,8 @@ def then(self, value: Any) -> PandasThen: function_name="whenthen", root_names=None, output_names=None, - implementation=self._condition._implementation, - backend_version=self._condition._backend_version, + implementation=self._implementation, + backend_version=self._backend_version, ) @@ -355,22 +372,7 @@ def __init__( self._root_names = root_names self._output_names = output_names - def when( - self, - *predicates: IntoPandasLikeExpr | Iterable[IntoPandasLikeExpr], - ) -> PandasChainedWhen: - return PandasChainedWhen( - self._call, # type: ignore[arg-type] - when_processing( - PandasLikeNamespace(self._implementation, self._backend_version), - *predicates, - ), - depth=self._depth + 1, - implementation=self._implementation, - backend_version=self._backend_version, - ) - - def otherwise(self, value: Any) -> PandasLikeExpr: + def otherwise(self, value: PandasLikeExpr | PandasLikeSeries | Any) -> PandasLikeExpr: # type ignore because we are setting the `_call` attribute to a # callable object of type `PandasWhen`, base class has the attribute as # only a `Callable` @@ -379,105 +381,105 @@ def otherwise(self, value: Any) -> PandasLikeExpr: return self -class PandasChainedWhen: - def __init__( - self, - above_when: PandasWhen | PandasChainedWhen, - condition: PandasLikeExpr, - depth: int, - implementation: Implementation, - backend_version: tuple[int, ...], - then_value: Any = None, - otherise_value: Any = None, - ) -> None: - self._implementation = implementation - self._depth = depth - self._backend_version = backend_version - self._condition = condition - self._above_when = above_when - self._then_value = then_value - self._otherwise_value = otherise_value - - # TODO @aivanoved: this is way slow as during computation time this takes - # quadratic time need to improve this to linear time - self._condition = self._condition & (~self._above_when._already_set) # type: ignore[has-type] - self._already_set = self._above_when._already_set | self._condition # type: ignore[has-type] - - def __call__(self, df: PandasLikeDataFrame) -> list[PandasLikeSeries]: - from narwhals._pandas_like.namespace import PandasLikeNamespace - - plx = PandasLikeNamespace( - implementation=self._implementation, backend_version=self._backend_version - ) - - set_then = self._condition._call(df)[0] - already_set = self._already_set._call(df)[0] - - value_series = plx._create_broadcast_series_from_scalar( - self._then_value, set_then - ) - otherwise_series = plx._create_broadcast_series_from_scalar( - self._otherwise_value, set_then - ) - - above_result = self._above_when(df)[0] - - result = value_series.zip_with(set_then, above_result).zip_with( - already_set, otherwise_series - ) - - return [result] - - def then(self, value: Any) -> PandasChainedThen: - self._then_value = value - return PandasChainedThen( - self, - depth=self._depth, - implementation=self._implementation, - function_name="chainedwhen", - root_names=None, - output_names=None, - backend_version=self._backend_version, - ) - - -class PandasChainedThen(PandasLikeExpr): - def __init__( - self, - call: PandasChainedWhen, - *, - depth: int, - function_name: str, - root_names: list[str] | None, - output_names: list[str] | None, - implementation: Implementation, - backend_version: tuple[int, ...], - ) -> None: - self._implementation = implementation - self._backend_version = backend_version - - self._call = call - self._depth = depth - self._function_name = function_name - self._root_names = root_names - self._output_names = output_names - - def when( - self, - *predicates: IntoPandasLikeExpr | Iterable[IntoPandasLikeExpr], - ) -> PandasChainedWhen: - return PandasChainedWhen( - self._call, # type: ignore[arg-type] - when_processing( - PandasLikeNamespace(self._implementation, self._backend_version), - *predicates, - ), - depth=self._depth + 1, - implementation=self._implementation, - backend_version=self._backend_version, - ) - - def otherwise(self, value: Any) -> PandasLikeExpr: - self._call._otherwise_value = value # type: ignore[attr-defined] - self._function_name = "chainedwhenotherwise" - return self +# class PandasChainedWhen: +# def __init__( +# self, +# above_when: PandasWhen | PandasChainedWhen, +# condition: PandasLikeExpr, +# depth: int, +# implementation: Implementation, +# backend_version: tuple[int, ...], +# then_value: Any = None, +# otherise_value: Any = None, +# ) -> None: +# self._implementation = implementation +# self._depth = depth +# self._backend_version = backend_version +# self._condition = condition +# self._above_when = above_when +# self._then_value = then_value +# self._otherwise_value = otherise_value +# +# # TODO @aivanoved: this is way slow as during computation time this takes +# # quadratic time need to improve this to linear time +# self._condition = self._condition & (~self._above_when._already_set) # type: ignore[has-type] +# self._already_set = self._above_when._already_set | self._condition # type: ignore[has-type] +# +# def __call__(self, df: PandasLikeDataFrame) -> list[PandasLikeSeries]: +# from narwhals._pandas_like.namespace import PandasLikeNamespace +# +# plx = PandasLikeNamespace( +# implementation=self._implementation, backend_version=self._backend_version +# ) +# +# set_then = self._condition._call(df)[0] +# already_set = self._already_set._call(df)[0] +# +# value_series = plx._create_broadcast_series_from_scalar( +# self._then_value, set_then +# ) +# otherwise_series = plx._create_broadcast_series_from_scalar( +# self._otherwise_value, set_then +# ) +# +# above_result = self._above_when(df)[0] +# +# result = value_series.zip_with(set_then, above_result).zip_with( +# already_set, otherwise_series +# ) +# +# return [result] +# +# def then(self, value: Any) -> PandasChainedThen: +# self._then_value = value +# return PandasChainedThen( +# self, +# depth=self._depth, +# implementation=self._implementation, +# function_name="chainedwhen", +# root_names=None, +# output_names=None, +# backend_version=self._backend_version, +# ) +# +# +# class PandasChainedThen(PandasLikeExpr): +# def __init__( +# self, +# call: PandasChainedWhen, +# *, +# depth: int, +# function_name: str, +# root_names: list[str] | None, +# output_names: list[str] | None, +# implementation: Implementation, +# backend_version: tuple[int, ...], +# ) -> None: +# self._implementation = implementation +# self._backend_version = backend_version +# +# self._call = call +# self._depth = depth +# self._function_name = function_name +# self._root_names = root_names +# self._output_names = output_names +# +# def when( +# self, +# *predicates: IntoPandasLikeExpr | Iterable[IntoPandasLikeExpr], +# ) -> PandasChainedWhen: +# return PandasChainedWhen( +# self._call, # type: ignore[arg-type] +# when_processing( +# PandasLikeNamespace(self._implementation, self._backend_version), +# *predicates, +# ), +# depth=self._depth + 1, +# implementation=self._implementation, +# backend_version=self._backend_version, +# ) +# +# def otherwise(self, value: Any) -> PandasLikeExpr: +# self._call._otherwise_value = value # type: ignore[attr-defined] +# self._function_name = "chainedwhenotherwise" +# return self diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py index 8cb98a9d1..e94c95a8c 100644 --- a/narwhals/_pandas_like/series.py +++ b/narwhals/_pandas_like/series.py @@ -8,23 +8,20 @@ from typing import overload from narwhals._pandas_like.utils import int_dtype_mapper +from narwhals._pandas_like.utils import narwhals_to_native_dtype from narwhals._pandas_like.utils import native_series_from_iterable -from narwhals._pandas_like.utils import reverse_translate_dtype from narwhals._pandas_like.utils import to_datetime from narwhals._pandas_like.utils import translate_dtype from narwhals._pandas_like.utils import validate_column_comparand from narwhals.dependencies import get_cudf from narwhals.dependencies import get_modin -from narwhals.dependencies import get_numpy from narwhals.dependencies import get_pandas -from narwhals.dependencies import get_pyarrow_compute from narwhals.utils import Implementation if TYPE_CHECKING: from typing_extensions import Self from narwhals._pandas_like.dataframe import PandasLikeDataFrame - from narwhals._pandas_like.namespace import PandasLikeNamespace from narwhals.dtypes import DType PANDAS_TO_NUMPY_DTYPE_NO_MISSING = { @@ -98,11 +95,6 @@ def __init__( else: self._use_copy_false = False - def __narwhals_namespace__(self) -> PandasLikeNamespace: - from narwhals._pandas_like.namespace import PandasLikeNamespace - - return PandasLikeNamespace(self._implementation, self._backend_version) - def __native_namespace__(self) -> Any: if self._implementation is Implementation.PANDAS: return get_pandas() @@ -180,7 +172,7 @@ def cast( dtype: Any, ) -> Self: ser = self._native_series - dtype = reverse_translate_dtype(dtype, ser.dtype, self._implementation) + dtype = narwhals_to_native_dtype(dtype, ser.dtype, self._implementation) return self._from_native_series(ser.astype(dtype)) def item(self: Self, index: int | None = None) -> Any: @@ -229,23 +221,16 @@ def is_in(self, other: Any) -> PandasLikeSeries: return self._from_native_series(res) def arg_true(self) -> PandasLikeSeries: - np = get_numpy() ser = self._native_series - res = np.flatnonzero(ser) - return self._from_native_series( - native_series_from_iterable( - res, - name=self.name, - index=range(len(res)), - implementation=self._implementation, - ) - ) + result = ser.__class__(range(len(ser)), name=ser.name, index=ser.index).loc[ser] + return self._from_native_series(result) # Binary comparisons def filter(self, other: Any) -> PandasLikeSeries: ser = self._native_series - other = validate_column_comparand(self._native_series.index, other) + if not (isinstance(other, list) and all(isinstance(x, bool) for x in other)): + other = validate_column_comparand(self._native_series.index, other) return self._from_native_series(self._rename(ser.loc[other], ser.name)) def __eq__(self, other: object) -> PandasLikeSeries: # type: ignore[override] @@ -530,22 +515,30 @@ def to_pandas(self) -> Any: # --- descriptive --- def is_duplicated(self: Self) -> Self: - return self._from_native_series(self._native_series.duplicated(keep=False)) + res = self._native_series.duplicated(keep=False) + res = self._rename(res, self.name) + return self._from_native_series(res) def is_empty(self: Self) -> bool: return self._native_series.empty # type: ignore[no-any-return] def is_unique(self: Self) -> Self: - return self._from_native_series(~self._native_series.duplicated(keep=False)) + res = ~self._native_series.duplicated(keep=False) + res = self._rename(res, self.name) + return self._from_native_series(res) def null_count(self: Self) -> int: return self._native_series.isna().sum() # type: ignore[no-any-return] def is_first_distinct(self: Self) -> Self: - return self._from_native_series(~self._native_series.duplicated(keep="first")) + res = ~self._native_series.duplicated(keep="first") + res = self._rename(res, self.name) + return self._from_native_series(res) def is_last_distinct(self: Self) -> Self: - return self._from_native_series(~self._native_series.duplicated(keep="last")) + res = ~self._native_series.duplicated(keep="last") + res = self._rename(res, self.name) + return self._from_native_series(res) def is_sorted(self: Self, *, descending: bool = False) -> bool: if not isinstance(descending, bool): @@ -597,7 +590,9 @@ def quantile( def zip_with(self: Self, mask: Any, other: Any) -> PandasLikeSeries: ser = self._native_series - res = ser.where(mask._native_series, other._native_series) + mask = validate_column_comparand(ser.index, mask) + other = validate_column_comparand(ser.index, other) + res = ser.where(mask, other) return self._from_native_series(res) def head(self: Self, n: int) -> Self: @@ -631,6 +626,22 @@ def to_dummies( def gather_every(self: Self, n: int, offset: int = 0) -> Self: return self._from_native_series(self._native_series.iloc[offset::n]) + def clip( + self: Self, lower_bound: Any | None = None, upper_bound: Any | None = None + ) -> Self: + return self._from_native_series( + self._native_series.clip(lower_bound, upper_bound) + ) + + def to_arrow(self: Self) -> Any: + if self._implementation is Implementation.CUDF: # pragma: no cover + msg = "`to_arrow` is not implemented for CuDF backend." + raise NotImplementedError(msg) + + import pyarrow as pa # ignore-banned-import() + + return pa.Array.from_pandas(self._native_series) + @property def str(self) -> PandasLikeSeriesStringNamespace: return PandasLikeSeriesStringNamespace(self) @@ -659,6 +670,25 @@ class PandasLikeSeriesStringNamespace: def __init__(self, series: PandasLikeSeries) -> None: self._pandas_series = series + def replace( + self, pattern: str, value: str, *, literal: bool = False, n: int = 1 + ) -> PandasLikeSeries: + return self._pandas_series._from_native_series( + self._pandas_series._native_series.str.replace( + pat=pattern, repl=value, n=n, regex=not literal + ), + ) + + def replace_all( + self, pattern: str, value: str, *, literal: bool = False + ) -> PandasLikeSeries: + return self.replace(pattern, value, literal=literal, n=-1) + + def strip_chars(self, characters: str | None) -> PandasLikeSeries: + return self._pandas_series._from_native_series( + self._pandas_series._native_series.str.strip(characters), + ) + def starts_with(self, prefix: str) -> PandasLikeSeries: return self._pandas_series._from_native_series( self._pandas_series._native_series.str.startswith(prefix), @@ -704,6 +734,21 @@ class PandasLikeSeriesDateTimeNamespace: def __init__(self, series: PandasLikeSeries) -> None: self._pandas_series = series + def date(self) -> PandasLikeSeries: + result = self._pandas_series._from_native_series( + self._pandas_series._native_series.dt.date, + ) + if str(result.dtype).lower() == "object": + msg = ( + "Accessing `date` on the default pandas backend " + "will return a Series of type `object`." + "\nThis differs from polars API and will prevent `.dt` chaining. " + "Please switch to the `pyarrow` backend:" + '\ndf.convert_dtypes(dtype_backend="pyarrow")' + ) + raise NotImplementedError(msg) + return result + def year(self) -> PandasLikeSeries: return self._pandas_series._from_native_series( self._pandas_series._native_series.dt.year, @@ -742,7 +787,8 @@ def microsecond(self) -> PandasLikeSeries: self._pandas_series._native_series.dtype ): # crazy workaround for https://github.com/pandas-dev/pandas/issues/59154 - pc = get_pyarrow_compute() + import pyarrow.compute as pc # ignore-banned-import() + native_series = self._pandas_series._native_series arr = native_series.array.__arrow_array__() result_arr = pc.add( diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py index f58c1297a..9e1d79ce9 100644 --- a/narwhals/_pandas_like/utils.py +++ b/narwhals/_pandas_like/utils.py @@ -197,6 +197,10 @@ def set_axis( implementation: Implementation, backend_version: tuple[int, ...], ) -> T: + if implementation is Implementation.CUDF: # pragma: no cover + obj = obj.copy(deep=False) # type: ignore[attr-defined] + obj.index = index # type: ignore[attr-defined] + return obj if implementation is Implementation.PANDAS and backend_version < ( 1, ): # pragma: no cover @@ -210,7 +214,7 @@ def set_axis( kwargs["copy"] = False else: # pragma: no cover pass - return obj.set_axis(index, axis=0, **kwargs) # type: ignore[no-any-return, attr-defined] + return obj.set_axis(index, axis=0, **kwargs) # type: ignore[attr-defined, no-any-return] def translate_dtype(column: Any) -> DType: @@ -273,15 +277,29 @@ def translate_dtype(column: Any) -> DType: if str(dtype) == "date32[day][pyarrow]": return dtypes.Date() if str(dtype) == "object": - if (idx := column.first_valid_index()) is not None and isinstance( - column.loc[idx], str - ): + if ( # pragma: no cover TODO(unassigned): why does this show as uncovered? + idx := getattr(column, "first_valid_index", lambda: None)() + ) is not None and isinstance(column.loc[idx], str): # Infer based on first non-missing value. # For pandas pre 3.0, this isn't perfect. # After pandas 3.0, pandas has a dedicated string dtype # which is inferred by default. return dtypes.String() - return dtypes.Object() + else: + df = column.to_frame() + if hasattr(df, "__dataframe__"): + from narwhals._interchange.dataframe import ( + map_interchange_dtype_to_narwhals_dtype, + ) + + try: + return map_interchange_dtype_to_narwhals_dtype( + df.__dataframe__().get_column(0).dtype + ) + except Exception: # noqa: BLE001 + return dtypes.Object() + else: # pragma: no cover + return dtypes.Object() return dtypes.Unknown() @@ -302,11 +320,20 @@ def get_dtype_backend(dtype: Any, implementation: Implementation) -> str: return "numpy" -def reverse_translate_dtype( # noqa: PLR0915 +def narwhals_to_native_dtype( # noqa: PLR0915 dtype: DType | type[DType], starting_dtype: Any, implementation: Implementation ) -> Any: from narwhals import dtypes + if "polars" in str(type(dtype)): + msg = ( + f"Expected Narwhals object, got: {type(dtype)}.\n\n" + "Perhaps you:\n" + "- Forgot a `nw.from_native` somewhere?\n" + "- Used `pl.Int64` instead of `nw.Int64`?" + ) + raise TypeError(msg) + dtype_backend = get_dtype_backend(starting_dtype, implementation) if isinstance_or_issubclass(dtype, dtypes.Float64): if dtype_backend == "pyarrow-nullable": @@ -412,25 +439,45 @@ def reverse_translate_dtype( # noqa: PLR0915 return "date32[pyarrow]" msg = "Date dtype only supported for pyarrow-backed data types in pandas" raise NotImplementedError(msg) + if isinstance_or_issubclass(dtype, dtypes.Enum): + msg = "Converting to Enum is not (yet) supported" + raise NotImplementedError(msg) msg = f"Unknown dtype: {dtype}" # pragma: no cover raise AssertionError(msg) -def validate_indices(series: list[PandasLikeSeries]) -> list[Any]: - idx = series[0]._native_series.index - reindexed = [series[0]._native_series] - for s in series[1:]: - if s._native_series.index is not idx: +def broadcast_series(series: list[PandasLikeSeries]) -> list[Any]: + native_namespace = series[0].__native_namespace__() + + lengths = [len(s) for s in series] + max_length = max(lengths) + + idx = series[lengths.index(max_length)]._native_series.index + reindexed = [] + + for s, length in zip(series, lengths): + s_native = s._native_series + if max_length > 1 and length == 1: + reindexed.append( + native_namespace.Series( + [s_native.iloc[0]] * max_length, + index=idx, + name=s_native.name, + dtype=s_native.dtype, + ) + ) + + elif s_native.index is not idx: reindexed.append( set_axis( - s._native_series, + s_native, idx, implementation=s._implementation, backend_version=s._backend_version, ) ) else: - reindexed.append(s._native_series) + reindexed.append(s_native) return reindexed diff --git a/narwhals/_polars/dataframe.py b/narwhals/_polars/dataframe.py index f137739fe..c5164147a 100644 --- a/narwhals/_polars/dataframe.py +++ b/narwhals/_polars/dataframe.py @@ -8,14 +8,16 @@ from narwhals._polars.utils import translate_dtype from narwhals.dependencies import get_polars from narwhals.utils import Implementation +from narwhals.utils import parse_columns_to_drop if TYPE_CHECKING: + import numpy as np from typing_extensions import Self class PolarsDataFrame: def __init__(self, df: Any, *, backend_version: tuple[int, ...]) -> None: - self._native_dataframe = df + self._native_frame = df self._implementation = Implementation.POLARS self._backend_version = backend_version @@ -31,7 +33,7 @@ def __narwhals_namespace__(self) -> PolarsNamespace: def __native_namespace__(self) -> Any: return get_polars() - def _from_native_dataframe(self, df: Any) -> Self: + def _from_native_frame(self, df: Any) -> Self: return self.__class__(df, backend_version=self._backend_version) def _from_native_object(self, obj: Any) -> Any: @@ -41,7 +43,7 @@ def _from_native_object(self, obj: Any) -> Any: return PolarsSeries(obj, backend_version=self._backend_version) if isinstance(obj, pl.DataFrame): - return self._from_native_dataframe(obj) + return self._from_native_frame(obj) # scalar return obj @@ -52,30 +54,63 @@ def __getattr__(self, attr: str) -> Any: def func(*args: Any, **kwargs: Any) -> Any: args, kwargs = extract_args_kwargs(args, kwargs) # type: ignore[assignment] return self._from_native_object( - getattr(self._native_dataframe, attr)(*args, **kwargs) + getattr(self._native_frame, attr)(*args, **kwargs) ) return func + def __array__(self, dtype: Any | None = None, copy: bool | None = None) -> np.ndarray: + if self._backend_version < (0, 20, 28) and copy is not None: # pragma: no cover + msg = "`copy` in `__array__` is only supported for Polars>=0.20.28" + raise NotImplementedError(msg) + if self._backend_version < (0, 20, 28): # pragma: no cover + return self._native_frame.__array__(dtype) + return self._native_frame.__array__(dtype) + @property def schema(self) -> dict[str, Any]: - schema = self._native_dataframe.schema + schema = self._native_frame.schema return {name: translate_dtype(dtype) for name, dtype in schema.items()} def collect_schema(self) -> dict[str, Any]: if self._backend_version < (1,): # pragma: no cover - schema = self._native_dataframe.schema + schema = self._native_frame.schema else: - schema = dict(self._native_dataframe.collect_schema()) + schema = dict(self._native_frame.collect_schema()) return {name: translate_dtype(dtype) for name, dtype in schema.items()} @property def shape(self) -> tuple[int, int]: - return self._native_dataframe.shape # type: ignore[no-any-return] + return self._native_frame.shape # type: ignore[no-any-return] def __getitem__(self, item: Any) -> Any: + if isinstance(item, tuple) and len(item) == 2 and isinstance(item[1], slice): + # TODO(marco): we can delete this branch after Polars==0.20.30 becomes the minimum + # Polars version we support + columns = self.columns + if isinstance(item[1].start, str) or isinstance(item[1].stop, str): + start = ( + columns.index(item[1].start) if item[1].start is not None else None + ) + stop = ( + columns.index(item[1].stop) + 1 if item[1].stop is not None else None + ) + step = item[1].step + return self._from_native_frame( + self._native_frame.select(columns[start:stop:step]).__getitem__( + item[0] + ) + ) + if isinstance(item[1].start, int) or isinstance(item[1].stop, int): + return self._from_native_frame( + self._native_frame.select( + columns[item[1].start : item[1].stop : item[1].step] + ).__getitem__(item[0]) + ) + msg = f"Expected slice of integers or strings, got: {type(item[1])}" # pragma: no cover + raise TypeError(msg) # pragma: no cover pl = get_polars() - result = self._native_dataframe.__getitem__(item) + result = self._native_frame.__getitem__(item) if isinstance(result, pl.Series): from narwhals._polars.series import PolarsSeries @@ -86,23 +121,23 @@ def get_column(self, name: str) -> Any: from narwhals._polars.series import PolarsSeries return PolarsSeries( - self._native_dataframe.get_column(name), backend_version=self._backend_version + self._native_frame.get_column(name), backend_version=self._backend_version ) def is_empty(self) -> bool: - return len(self._native_dataframe) == 0 + return len(self._native_frame) == 0 @property def columns(self) -> list[str]: - return self._native_dataframe.columns # type: ignore[no-any-return] + return self._native_frame.columns # type: ignore[no-any-return] def lazy(self) -> PolarsLazyFrame: return PolarsLazyFrame( - self._native_dataframe.lazy(), backend_version=self._backend_version + self._native_frame.lazy(), backend_version=self._backend_version ) def to_dict(self, *, as_series: bool) -> Any: - df = self._native_dataframe + df = self._native_frame if as_series: from narwhals._polars.series import PolarsSeries @@ -121,15 +156,21 @@ def group_by(self, *by: str) -> Any: def with_row_index(self, name: str) -> Any: if self._backend_version < (0, 20, 4): # pragma: no cover - return self._from_native_dataframe( - self._native_dataframe.with_row_count(name) + return self._from_native_frame(self._native_frame.with_row_count(name)) + return self._from_native_frame(self._native_frame.with_row_index(name)) + + def drop(self: Self, columns: list[str], strict: bool) -> Self: # noqa: FBT001 + if self._backend_version < (1, 0, 0): # pragma: no cover + to_drop = parse_columns_to_drop( + compliant_frame=self, columns=columns, strict=strict ) - return self._from_native_dataframe(self._native_dataframe.with_row_index(name)) + return self._from_native_frame(self._native_frame.drop(to_drop)) + return self._from_native_frame(self._native_frame.drop(columns, strict=strict)) class PolarsLazyFrame: def __init__(self, df: Any, *, backend_version: tuple[int, ...]) -> None: - self._native_dataframe = df + self._native_frame = df self._backend_version = backend_version def __repr__(self) -> str: # pragma: no cover @@ -144,37 +185,37 @@ def __narwhals_namespace__(self) -> PolarsNamespace: def __native_namespace__(self) -> Any: # pragma: no cover return get_polars() - def _from_native_dataframe(self, df: Any) -> Self: + def _from_native_frame(self, df: Any) -> Self: return self.__class__(df, backend_version=self._backend_version) def __getattr__(self, attr: str) -> Any: def func(*args: Any, **kwargs: Any) -> Any: args, kwargs = extract_args_kwargs(args, kwargs) # type: ignore[assignment] - return self._from_native_dataframe( - getattr(self._native_dataframe, attr)(*args, **kwargs) + return self._from_native_frame( + getattr(self._native_frame, attr)(*args, **kwargs) ) return func @property def columns(self) -> list[str]: - return self._native_dataframe.columns # type: ignore[no-any-return] + return self._native_frame.columns # type: ignore[no-any-return] @property def schema(self) -> dict[str, Any]: - schema = self._native_dataframe.schema + schema = self._native_frame.schema return {name: translate_dtype(dtype) for name, dtype in schema.items()} def collect_schema(self) -> dict[str, Any]: if self._backend_version < (1,): # pragma: no cover - schema = self._native_dataframe.schema + schema = self._native_frame.schema else: - schema = dict(self._native_dataframe.collect_schema()) + schema = dict(self._native_frame.collect_schema()) return {name: translate_dtype(dtype) for name, dtype in schema.items()} def collect(self) -> PolarsDataFrame: return PolarsDataFrame( - self._native_dataframe.collect(), backend_version=self._backend_version + self._native_frame.collect(), backend_version=self._backend_version ) def group_by(self, *by: str) -> Any: @@ -184,7 +225,10 @@ def group_by(self, *by: str) -> Any: def with_row_index(self, name: str) -> Any: if self._backend_version < (0, 20, 4): # pragma: no cover - return self._from_native_dataframe( - self._native_dataframe.with_row_count(name) - ) - return self._from_native_dataframe(self._native_dataframe.with_row_index(name)) + return self._from_native_frame(self._native_frame.with_row_count(name)) + return self._from_native_frame(self._native_frame.with_row_index(name)) + + def drop(self: Self, columns: list[str], strict: bool) -> Self: # noqa: FBT001 + if self._backend_version < (1, 0, 0): # pragma: no cover + return self._from_native_frame(self._native_frame.drop(columns)) + return self._from_native_frame(self._native_frame.drop(columns, strict=strict)) diff --git a/narwhals/_polars/expr.py b/narwhals/_polars/expr.py index 98aac298e..4f1532823 100644 --- a/narwhals/_polars/expr.py +++ b/narwhals/_polars/expr.py @@ -3,10 +3,9 @@ from typing import TYPE_CHECKING from typing import Any -from narwhals._polars.namespace import PolarsNamespace from narwhals._polars.utils import extract_args_kwargs from narwhals._polars.utils import extract_native -from narwhals._polars.utils import reverse_translate_dtype +from narwhals._polars.utils import narwhals_to_native_dtype from narwhals.utils import Implementation if TYPE_CHECKING: @@ -23,12 +22,6 @@ def __init__(self, expr: Any) -> None: def __repr__(self) -> str: # pragma: no cover return "PolarsExpr" - def __narwhals_expr__(self) -> Self: # pragma: no cover - return self - - def __narwhals_namespace__(self) -> PolarsNamespace: # pragma: no cover - return PolarsNamespace(backend_version=self._backend_version) - def _from_native_expr(self, expr: Any) -> Self: return self.__class__(expr) @@ -43,7 +36,7 @@ def func(*args: Any, **kwargs: Any) -> Any: def cast(self, dtype: DType) -> Self: expr = self._native_expr - dtype = reverse_translate_dtype(dtype) + dtype = narwhals_to_native_dtype(dtype) return self._from_native_expr(expr.cast(dtype)) def __eq__(self, other: object) -> Self: # type: ignore[override] diff --git a/narwhals/_polars/group_by.py b/narwhals/_polars/group_by.py index c0de75736..f03da610e 100644 --- a/narwhals/_polars/group_by.py +++ b/narwhals/_polars/group_by.py @@ -14,27 +14,27 @@ class PolarsGroupBy: def __init__(self, df: Any, keys: list[str]) -> None: self._compliant_frame = df self.keys = keys - self._grouped = df._native_dataframe.group_by(keys) + self._grouped = df._native_frame.group_by(keys) def agg(self, *aggs: Any, **named_aggs: Any) -> PolarsDataFrame: aggs, named_aggs = extract_args_kwargs(aggs, named_aggs) # type: ignore[assignment] - return self._compliant_frame._from_native_dataframe( # type: ignore[no-any-return] + return self._compliant_frame._from_native_frame( # type: ignore[no-any-return] self._grouped.agg(*aggs, **named_aggs), ) def __iter__(self) -> Any: for key, df in self._grouped: - yield tuple(key), self._compliant_frame._from_native_dataframe(df) + yield tuple(key), self._compliant_frame._from_native_frame(df) class PolarsLazyGroupBy: def __init__(self, df: Any, keys: list[str]) -> None: self._compliant_frame = df self.keys = keys - self._grouped = df._native_dataframe.group_by(keys) + self._grouped = df._native_frame.group_by(keys) def agg(self, *aggs: Any, **named_aggs: Any) -> PolarsLazyFrame: aggs, named_aggs = extract_args_kwargs(aggs, named_aggs) # type: ignore[assignment] - return self._compliant_frame._from_native_dataframe( # type: ignore[no-any-return] + return self._compliant_frame._from_native_frame( # type: ignore[no-any-return] self._grouped.agg(*aggs, **named_aggs), ) diff --git a/narwhals/_polars/namespace.py b/narwhals/_polars/namespace.py index 683b83c12..48ee8ebc0 100644 --- a/narwhals/_polars/namespace.py +++ b/narwhals/_polars/namespace.py @@ -1,13 +1,15 @@ from __future__ import annotations +from functools import reduce from typing import TYPE_CHECKING from typing import Any from typing import Iterable from typing import Sequence from narwhals import dtypes +from narwhals._expression_parsing import parse_into_exprs from narwhals._polars.utils import extract_args_kwargs -from narwhals._polars.utils import reverse_translate_dtype +from narwhals._polars.utils import narwhals_to_native_dtype from narwhals.dependencies import get_polars from narwhals.utils import Implementation @@ -15,6 +17,7 @@ from narwhals._polars.dataframe import PolarsDataFrame from narwhals._polars.dataframe import PolarsLazyFrame from narwhals._polars.expr import PolarsExpr + from narwhals._polars.typing import IntoPolarsExpr class PolarsNamespace: @@ -71,7 +74,7 @@ def concat( from narwhals._polars.dataframe import PolarsLazyFrame pl = get_polars() - dfs: list[Any] = [item._native_dataframe for item in items] + dfs: list[Any] = [item._native_frame for item in items] result = pl.concat(dfs, how=how) if isinstance(result, pl.DataFrame): return PolarsDataFrame(result, backend_version=items[0]._backend_version) @@ -82,14 +85,31 @@ def lit(self, value: Any, dtype: dtypes.DType | None = None) -> PolarsExpr: pl = get_polars() if dtype is not None: - return PolarsExpr(pl.lit(value, dtype=reverse_translate_dtype(dtype))) + return PolarsExpr(pl.lit(value, dtype=narwhals_to_native_dtype(dtype))) return PolarsExpr(pl.lit(value)) - def mean(self, *column_names: str) -> Any: + def mean(self, *column_names: str) -> PolarsExpr: + from narwhals._polars.expr import PolarsExpr + pl = get_polars() if self._backend_version < (0, 20, 4): # pragma: no cover - return pl.mean([*column_names]) - return pl.mean(*column_names) + return PolarsExpr(pl.mean([*column_names])) + return PolarsExpr(pl.mean(*column_names)) + + def mean_horizontal(self, *exprs: IntoPolarsExpr) -> PolarsExpr: + from narwhals._polars.expr import PolarsExpr + + pl = get_polars() + polars_exprs = parse_into_exprs(*exprs, namespace=self) + + if self._backend_version < (0, 20, 8): # pragma: no cover + total = reduce(lambda x, y: x + y, (e.fill_null(0.0) for e in polars_exprs)) + n_non_zero = reduce( + lambda x, y: x + y, ((1 - e.is_null()) for e in polars_exprs) + ) + return PolarsExpr(total._native_expr / n_non_zero._native_expr) + + return PolarsExpr(pl.mean_horizontal([e._native_expr for e in polars_exprs])) @property def selectors(self) -> PolarsSelectors: @@ -102,7 +122,7 @@ def by_dtype(self, dtypes: Iterable[dtypes.DType]) -> PolarsExpr: pl = get_polars() return PolarsExpr( - pl.selectors.by_dtype([reverse_translate_dtype(dtype) for dtype in dtypes]) + pl.selectors.by_dtype([narwhals_to_native_dtype(dtype) for dtype in dtypes]) ) def numeric(self) -> PolarsExpr: diff --git a/narwhals/_polars/series.py b/narwhals/_polars/series.py index 4a5fd2b7b..e71520042 100644 --- a/narwhals/_polars/series.py +++ b/narwhals/_polars/series.py @@ -17,8 +17,7 @@ from narwhals._polars.dataframe import PolarsDataFrame from narwhals.dtypes import DType -from narwhals._polars.namespace import PolarsNamespace -from narwhals._polars.utils import reverse_translate_dtype +from narwhals._polars.utils import narwhals_to_native_dtype from narwhals._polars.utils import translate_dtype PL = get_polars() @@ -39,9 +38,6 @@ def __narwhals_series__(self) -> Self: def __native_namespace__(self) -> Any: return get_polars() - def __narwhals_namespace__(self) -> PolarsNamespace: - return PolarsNamespace(backend_version=self._backend_version) - def _from_native_series(self, series: Any) -> Self: return self.__class__(series, backend_version=self._backend_version) @@ -94,7 +90,7 @@ def __getitem__(self, item: int | slice | Sequence[int]) -> Any | Self: def cast(self, dtype: DType) -> Self: ser = self._native_series - dtype = reverse_translate_dtype(dtype) + dtype = narwhals_to_native_dtype(dtype) return self._from_native_series(ser.cast(dtype)) def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray: diff --git a/narwhals/_polars/typing.py b/narwhals/_polars/typing.py new file mode 100644 index 000000000..3aed28991 --- /dev/null +++ b/narwhals/_polars/typing.py @@ -0,0 +1,17 @@ +from __future__ import annotations # pragma: no cover + +from typing import TYPE_CHECKING # pragma: no cover +from typing import Union # pragma: no cover + +if TYPE_CHECKING: + import sys + + if sys.version_info >= (3, 10): + from typing import TypeAlias + else: + from typing_extensions import TypeAlias + + from narwhals._polars.expr import PolarsExpr + from narwhals._polars.series import PolarsSeries + + IntoPolarsExpr: TypeAlias = Union[PolarsExpr, str, PolarsSeries] diff --git a/narwhals/_polars/utils.py b/narwhals/_polars/utils.py index 2dc92b3ad..51f0b1898 100644 --- a/narwhals/_polars/utils.py +++ b/narwhals/_polars/utils.py @@ -13,7 +13,7 @@ def extract_native(obj: Any) -> Any: from narwhals._polars.series import PolarsSeries if isinstance(obj, (PolarsDataFrame, PolarsLazyFrame)): - return obj._native_dataframe + return obj._native_frame if isinstance(obj, PolarsSeries): return obj._native_series if isinstance(obj, PolarsExpr): @@ -68,7 +68,7 @@ def translate_dtype(dtype: Any) -> dtypes.DType: return dtypes.Unknown() -def reverse_translate_dtype(dtype: dtypes.DType | type[dtypes.DType]) -> Any: +def narwhals_to_native_dtype(dtype: dtypes.DType | type[dtypes.DType]) -> Any: pl = get_polars() from narwhals import dtypes @@ -100,8 +100,9 @@ def reverse_translate_dtype(dtype: dtypes.DType | type[dtypes.DType]) -> Any: return pl.Object() if dtype == dtypes.Categorical: return pl.Categorical() - if dtype == dtypes.Enum: # pragma: no cover - return pl.Enum() + if dtype == dtypes.Enum: + msg = "Converting to Enum is not (yet) supported" + raise NotImplementedError(msg) if dtype == dtypes.Datetime: return pl.Datetime() if dtype == dtypes.Duration: diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py index 644db7bc3..64df5913f 100644 --- a/narwhals/dataframe.py +++ b/narwhals/dataframe.py @@ -11,16 +11,19 @@ from typing import TypeVar from typing import overload -from narwhals.dependencies import get_numpy from narwhals.dependencies import get_polars +from narwhals.dependencies import is_numpy_array from narwhals.schema import Schema from narwhals.utils import flatten +from narwhals.utils import parse_version if TYPE_CHECKING: from io import BytesIO from pathlib import Path import numpy as np + import pandas as pd + import pyarrow as pa from typing_extensions import Self from narwhals.group_by import GroupBy @@ -95,9 +98,9 @@ def with_row_index(self, name: str = "index") -> Self: self._compliant_frame.with_row_index(name), ) - def drop_nulls(self) -> Self: + def drop_nulls(self: Self, subset: str | list[str] | None = None) -> Self: return self._from_compliant_dataframe( - self._compliant_frame.drop_nulls(), + self._compliant_frame.drop_nulls(subset=subset), ) @property @@ -137,9 +140,9 @@ def head(self, n: int) -> Self: def tail(self, n: int) -> Self: return self._from_compliant_dataframe(self._compliant_frame.tail(n)) - def drop(self, *columns: str | Iterable[str]) -> Self: + def drop(self, *columns: Iterable[str], strict: bool) -> Self: return self._from_compliant_dataframe( - self._compliant_frame.drop(*flatten(columns)) + self._compliant_frame.drop(columns, strict=strict) ) def unique( @@ -155,8 +158,13 @@ def unique( ) ) - def filter(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> Self: - predicates, _ = self._flatten_and_extract(*predicates) + def filter(self, *predicates: IntoExpr | Iterable[IntoExpr] | list[bool]) -> Self: + if not ( + len(predicates) == 1 + and isinstance(predicates[0], list) + and all(isinstance(x, bool) for x in predicates[0]) + ): + predicates, _ = self._flatten_and_extract(*predicates) return self._from_compliant_dataframe( self._compliant_frame.filter(*predicates), ) @@ -230,8 +238,8 @@ def __init__( msg = f"Expected an object which implements `__narwhals_dataframe__`, got: {type(df)}" raise AssertionError(msg) - def __array__(self) -> np.ndarray: - return self._compliant_frame.to_numpy() + def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray: + return self._compliant_frame.__array__(dtype, copy=copy) def __repr__(self) -> str: # pragma: no cover header = " Narwhals DataFrame " @@ -247,6 +255,30 @@ def __repr__(self) -> str: # pragma: no cover + "┘" ) + def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: + """ + Export a DataFrame via the Arrow PyCapsule Interface. + + - if the underlying dataframe implements the interface, it'll return that + - else, it'll call `to_arrow` and then defer to PyArrow's implementation + + See [PyCapsule Interface](https://arrow.apache.org/docs/dev/format/CDataInterface/PyCapsuleInterface.html) + for more. + """ + native_frame = self._compliant_frame._native_frame + if hasattr(native_frame, "__arrow_c_stream__"): + return native_frame.__arrow_c_stream__(requested_schema=requested_schema) + try: + import pyarrow as pa # ignore-banned-import + except ModuleNotFoundError as exc: # pragma: no cover + msg = f"PyArrow>=14.0.0 is required for `DataFrame.__arrow_c_stream__` for object of type {type(native_frame)}" + raise ModuleNotFoundError(msg) from exc + if parse_version(pa.__version__) < (14, 0): # pragma: no cover + msg = f"PyArrow>=14.0.0 is required for `DataFrame.__arrow_c_stream__` for object of type {type(native_frame)}" + raise ModuleNotFoundError(msg) from None + pa_table = self.to_arrow() + return pa_table.__arrow_c_stream__(requested_schema=requested_schema) + def lazy(self) -> LazyFrame[Any]: """ Lazify the DataFrame (if possible). @@ -281,7 +313,7 @@ def lazy(self) -> LazyFrame[Any]: """ return super().lazy() - def to_pandas(self) -> Any: + def to_pandas(self) -> pd.DataFrame: """ Convert this DataFrame to a pandas DataFrame. @@ -316,6 +348,38 @@ def to_pandas(self) -> Any: """ return self._compliant_frame.to_pandas() + def write_csv(self, file: str | Path | BytesIO | None = None) -> Any: + r""" + Write dataframe to parquet file. + + Examples: + Construct pandas and Polars DataFrames: + + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> df = {"foo": [1, 2, 3], "bar": [6.0, 7.0, 8.0], "ham": ["a", "b", "c"]} + >>> df_pd = pd.DataFrame(df) + >>> df_pl = pl.DataFrame(df) + + We define a library agnostic function: + + >>> def func(df): + ... df = nw.from_native(df) + ... return df.write_csv() + + We can then pass either pandas or Polars to `func`: + + >>> func(df_pd) # doctest: +SKIP + 'foo,bar,ham\n1,6.0,a\n2,7.0,b\n3,8.0,c\n' + >>> func(df_pl) # doctest: +SKIP + 'foo,bar,ham\n1,6.0,a\n2,7.0,b\n3,8.0,c\n' + + If we had passed a file name to `write_csv`, it would have been + written to that file. + """ + return self._compliant_frame.write_csv(file) + def write_parquet(self, file: str | Path | BytesIO) -> Any: """ Write dataframe to parquet file. @@ -343,7 +407,7 @@ def write_parquet(self, file: str | Path | BytesIO) -> Any: """ self._compliant_frame.write_parquet(file) - def to_numpy(self) -> Any: + def to_numpy(self) -> np.ndarray: """ Convert this DataFrame to a NumPy ndarray. @@ -453,6 +517,8 @@ def get_column(self, name: str) -> Series: level=self._level, ) + @overload + def __getitem__(self, item: tuple[Sequence[int], slice]) -> Self: ... @overload def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ... @overload @@ -477,20 +543,33 @@ def __getitem__( | slice | Sequence[int] | tuple[Sequence[int], str | int] - | tuple[Sequence[int], Sequence[int] | Sequence[str]], + | tuple[Sequence[int], Sequence[int] | Sequence[str] | slice], ) -> Series | Self: """ Extract column or slice of DataFrame. Arguments: - item: how to slice dataframe: + item: How to slice dataframe. What happens depends on what is passed. It's easiest + to explain by example. Suppose we have a Dataframe `df`: + + - `df['a']` extracts column `'a'` and returns a `Series`. + - `df[0:2]` extracts the first two rows and returns a `DataFrame`. + - `df[0:2, 'a']` extracts the first two rows from column `'a'` and returns + a `Series`. + - `df[0:2, 0]` extracts the first two rows from the first column and returns + a `Series`. + - `df[[0, 1], [0, 1, 2]]` extracts the first two rows and the first three columns + and returns a `DataFrame` + - `df[0: 2, ['a', 'c']]` extracts the first two rows and columns `'a'` and `'c'` and + returns a `DataFrame` + - `df[:, 0: 2]` extracts all rows from the first two columns and returns a `DataFrame` + - `df[:, 'a': 'c']` extracts all rows and all columns positioned between `'a'` and `'c'` + _inclusive_ and returns a `DataFrame`. For example, if the columns are + `'a', 'd', 'c', 'b'`, then that would extract columns `'a'`, `'d'`, and `'c'`. - - str: extract column - - slice or Sequence of integers: slice rows from dataframe. - - tuple of Sequence of integers and str or int: slice rows and extract column at the same time. - - tuple of Sequence of integers and Sequence of integers: slice rows and extract columns at the same time. Notes: - Integers are always interpreted as positions, and strings always as column names. + - Integers are always interpreted as positions + - Strings are always interpreted as column names. In contrast with Polars, pandas allows non-string column names. If you don't know whether the column name you're trying to extract @@ -525,6 +604,8 @@ def __getitem__( 2 ] """ + if isinstance(item, int): + item = [item] if ( isinstance(item, tuple) and len(item) == 2 @@ -539,7 +620,7 @@ def __getitem__( if ( isinstance(item, tuple) and len(item) == 2 - and isinstance(item[1], (list, tuple)) + and isinstance(item[1], (list, tuple, slice)) ): return self._from_compliant_dataframe(self._compliant_frame[item]) if isinstance(item, str) or (isinstance(item, tuple) and len(item) == 2): @@ -551,9 +632,7 @@ def __getitem__( ) elif isinstance(item, (Sequence, slice)) or ( - (np := get_numpy()) is not None - and isinstance(item, np.ndarray) - and item.ndim == 1 + is_numpy_array(item) and item.ndim == 1 ): return self._from_compliant_dataframe(self._compliant_frame[item]) @@ -588,7 +667,7 @@ def to_dict( ... "A": [1, 2, 3, 4, 5], ... "fruits": ["banana", "banana", "apple", "apple", "banana"], ... "B": [5, 4, 3, 2, 1], - ... "cars": ["beetle", "audi", "beetle", "beetle", "beetle"], + ... "animals": ["beetle", "fly", "beetle", "beetle", "beetle"], ... "optional": [28, 300, None, 2, -30], ... } >>> df_pd = pd.DataFrame(df) @@ -603,9 +682,9 @@ def to_dict( We can then pass either pandas or Polars to `func`: >>> func(df_pd) - {'A': [1, 2, 3, 4, 5], 'fruits': ['banana', 'banana', 'apple', 'apple', 'banana'], 'B': [5, 4, 3, 2, 1], 'cars': ['beetle', 'audi', 'beetle', 'beetle', 'beetle'], 'optional': [28.0, 300.0, nan, 2.0, -30.0]} + {'A': [1, 2, 3, 4, 5], 'fruits': ['banana', 'banana', 'apple', 'apple', 'banana'], 'B': [5, 4, 3, 2, 1], 'animals': ['beetle', 'fly', 'beetle', 'beetle', 'beetle'], 'optional': [28.0, 300.0, nan, 2.0, -30.0]} >>> func(df_pl) - {'A': [1, 2, 3, 4, 5], 'fruits': ['banana', 'banana', 'apple', 'apple', 'banana'], 'B': [5, 4, 3, 2, 1], 'cars': ['beetle', 'audi', 'beetle', 'beetle', 'beetle'], 'optional': [28, 300, None, 2, -30]} + {'A': [1, 2, 3, 4, 5], 'fruits': ['banana', 'banana', 'apple', 'apple', 'banana'], 'B': [5, 4, 3, 2, 1], 'animals': ['beetle', 'fly', 'beetle', 'beetle', 'beetle'], 'optional': [28, 300, None, 2, -30]} """ from narwhals.series import Series @@ -621,6 +700,40 @@ def to_dict( } return self._compliant_frame.to_dict(as_series=as_series) # type: ignore[no-any-return] + def row(self, index: int) -> tuple[Any, ...]: + """ + Get values at given row. + + !!!note + You should NEVER use this method to iterate over a DataFrame; + if you require row-iteration you should strongly prefer use of iter_rows() instead. + + Arguments: + index: Row number. + + Examples: + >>> import narwhals as nw + >>> import pandas as pd + >>> import polars as pl + >>> data = {"a": [1, 2, 3], "b": [4, 5, 6]} + >>> df_pd = pd.DataFrame(data) + >>> df_pl = pl.DataFrame(data) + + Let's define a library-agnostic function to get the second row. + + >>> @nw.narwhalify + ... def func(df): + ... return df.row(1) + + We can then pass pandas / Polars / any other supported library: + + >>> func(df_pd) + (2, 5) + >>> func(df_pl) + (2, 5) + """ + return self._compliant_frame.row(index) # type: ignore[no-any-return] + # inherited def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self: """ @@ -663,10 +776,14 @@ def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Se """ return super().pipe(function, *args, **kwargs) - def drop_nulls(self) -> Self: + def drop_nulls(self: Self, subset: str | list[str] | None = None) -> Self: """ Drop null values. + Arguments: + subset: Column name(s) for which null values are considered. If set to None + (default), use all columns. + Notes: pandas and Polars handle null values differently. Polars distinguishes between NaN and Null, whereas pandas doesn't. @@ -700,7 +817,7 @@ def drop_nulls(self) -> Self: │ 1.0 ┆ 1.0 │ └─────┴─────┘ """ - return super().drop_nulls() + return super().drop_nulls(subset=subset) def with_row_index(self, name: str = "index") -> Self: """ @@ -1232,7 +1349,6 @@ def head(self, n: int = 5) -> Self: │ 3 ┆ 8 ┆ c │ └─────┴─────┴─────┘ """ - return super().head(n) def tail(self, n: int = 5) -> Self: @@ -1282,12 +1398,14 @@ def tail(self, n: int = 5) -> Self: """ return super().tail(n) - def drop(self, *columns: str | Iterable[str]) -> Self: + def drop(self, *columns: str | Iterable[str], strict: bool = True) -> Self: """ Remove columns from the dataframe. Arguments: *columns: Names of the columns that should be removed from the dataframe. + strict: Validate that all column names exist in the schema and throw an + exception if a column name does not exist in the schema. Examples: >>> import pandas as pd @@ -1345,7 +1463,7 @@ def drop(self, *columns: str | Iterable[str]) -> Self: │ 8.0 │ └─────┘ """ - return super().drop(*columns) + return super().drop(*flatten(columns), strict=strict) def unique( self, @@ -1406,14 +1524,15 @@ def unique( """ return super().unique(subset, keep=keep, maintain_order=maintain_order) - def filter(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> Self: + def filter(self, *predicates: IntoExpr | Iterable[IntoExpr] | list[bool]) -> Self: r""" Filter the rows in the DataFrame based on one or more predicate expressions. The original order of the remaining rows is preserved. Arguments: - predicates: Expression(s) that evaluates to a boolean Series. + *predicates: Expression(s) that evaluates to a boolean Series. Can + also be a (single!) boolean list. Examples: >>> import pandas as pd @@ -1802,7 +1921,6 @@ def is_empty(self: Self) -> bool: >>> func(df_pd), func(df_pl) (False, False) """ - return self._compliant_frame.is_empty() # type: ignore[no-any-return] def is_unique(self: Self) -> Series: @@ -1908,7 +2026,6 @@ def null_count(self: Self) -> Self: │ 1 ┆ 1 ┆ 0 │ └─────┴─────┴─────┘ """ - return self._from_compliant_dataframe(self._compliant_frame.null_count()) def item(self: Self, row: int | None = None, column: int | str | None = None) -> Any: @@ -1999,8 +2116,8 @@ def gather_every(self: Self, n: int, offset: int = 0) -> Self: starting from a offset of 1: >>> @nw.narwhalify - ... def func(df_any): - ... return df_any.gather_every(n=2, offset=1) + ... def func(df): + ... return df.gather_every(n=2, offset=1) >>> func(df_pd) a b @@ -2020,6 +2137,42 @@ def gather_every(self: Self, n: int, offset: int = 0) -> Self: """ return super().gather_every(n=n, offset=offset) + def to_arrow(self: Self) -> pa.Table: + r""" + Convert to arrow table. + + Examples: + >>> import narwhals as nw + >>> import pandas as pd + >>> import polars as pl + >>> data = {"foo": [1, 2, 3], "bar": ["a", "b", "c"]} + >>> df_pd = pd.DataFrame(data) + >>> df_pl = pl.DataFrame(data) + + Let's define a dataframe-agnostic function that converts to arrow table: + + >>> @nw.narwhalify + ... def func(df): + ... return df.to_arrow() + + >>> func(df_pd) # doctest:+SKIP + pyarrow.Table + foo: int64 + bar: string + ---- + foo: [[1,2,3]] + bar: [["a","b","c"]] + + >>> func(df_pl) # doctest:+NORMALIZE_WHITESPACE + pyarrow.Table + foo: int64 + bar: large_string + ---- + foo: [[1,2,3]] + bar: [["a","b","c"]] + """ + return self._compliant_frame.to_arrow() + class LazyFrame(BaseFrame[FrameT]): """ @@ -2143,10 +2296,14 @@ def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Se """ return super().pipe(function, *args, **kwargs) - def drop_nulls(self) -> Self: + def drop_nulls(self: Self, subset: str | list[str] | None = None) -> Self: """ Drop null values. + Arguments: + subset: Column name(s) for which null values are considered. If set to None + (default), use all columns. + Notes: pandas and Polars handle null values differently. Polars distinguishes between NaN and Null, whereas pandas doesn't. @@ -2180,7 +2337,7 @@ def drop_nulls(self) -> Self: │ 1.0 ┆ 1.0 │ └─────┴─────┘ """ - return super().drop_nulls() + return super().drop_nulls(subset=subset) def with_row_index(self, name: str = "index") -> Self: """ @@ -2378,11 +2535,15 @@ def select( Arguments: *exprs: Column(s) to select, specified as positional arguments. - Accepts expression input. Strings are parsed as column names, - other non-expression inputs are parsed as literals. - + Accepts expression input. Strings are parsed as column names. **named_exprs: Additional columns to select, specified as keyword arguments. - The columns will be renamed to the keyword used. + The columns will be renamed to the keyword used. + + Notes: + If you'd like to select a column whose name isn't a string (for example, + if you're working with pandas) then you should explicitly use `nw.col` instead + of just passing the column name. For example, to select a column named + `0` use `df.select(nw.col(0))`, not `df.select(0)`. Examples: >>> import pandas as pd @@ -2699,13 +2860,19 @@ def tail(self, n: int = 5) -> Self: """ return super().tail(n) - def drop(self, *columns: str | Iterable[str]) -> Self: + def drop(self, *columns: str | Iterable[str], strict: bool = True) -> Self: r""" Remove columns from the LazyFrame. Arguments: - *columns: Names of the columns that should be removed from the - dataframe. Accepts column selector input. + *columns: Names of the columns that should be removed from the dataframe. + strict: Validate that all column names exist in the schema and throw an + exception if a column name does not exist in the schema. + + Warning: + `strict` argument is ignored for `polars<1.0.0`. + + Please consider upgrading to a newer version or pass to eager mode. Examples: >>> import pandas as pd @@ -2763,7 +2930,7 @@ def drop(self, *columns: str | Iterable[str]) -> Self: │ 8.0 │ └─────┘ """ - return super().drop(*flatten(columns)) + return super().drop(*flatten(columns), strict=strict) def unique( self, @@ -2828,14 +2995,15 @@ def unique( """ return super().unique(subset, keep=keep, maintain_order=maintain_order) - def filter(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> Self: + def filter(self, *predicates: IntoExpr | Iterable[IntoExpr] | list[bool]) -> Self: r""" Filter the rows in the LazyFrame based on a predicate expression. The original order of the remaining rows is preserved. Arguments: - *predicates: Expression that evaluates to a boolean Series. + *predicates: Expression that evaluates to a boolean Series. Can + also be a (single!) boolean list. Examples: >>> import pandas as pd @@ -3299,8 +3467,8 @@ def gather_every(self: Self, n: int, offset: int = 0) -> Self: starting from a offset of 1: >>> @nw.narwhalify - ... def func(df_any): - ... return df_any.gather_every(n=2, offset=1) + ... def func(df): + ... return df.gather_every(n=2, offset=1) >>> func(df_pd) a b diff --git a/narwhals/dependencies.py b/narwhals/dependencies.py index 9eebdb703..66516eac9 100644 --- a/narwhals/dependencies.py +++ b/narwhals/dependencies.py @@ -8,11 +8,18 @@ from typing import Any if TYPE_CHECKING: + import numpy as np + if sys.version_info >= (3, 10): from typing import TypeGuard else: from typing_extensions import TypeGuard + import cudf + import dask.dataframe as dd + import modin.pandas as mpd import pandas as pd + import polars as pl + import pyarrow as pa def get_polars() -> Any: @@ -42,24 +49,6 @@ def get_pyarrow() -> Any: # pragma: no cover return sys.modules.get("pyarrow", None) -def get_pyarrow_compute() -> Any: # pragma: no cover - """Get pyarrow.compute module (if pyarrow has already been imported - else return None).""" - if "pyarrow" in sys.modules: - import pyarrow.compute as pc - - return pc - return None - - -def get_pyarrow_parquet() -> Any: # pragma: no cover - """Get pyarrow.parquet module (if pyarrow has already been imported - else return None).""" - if "pyarrow" in sys.modules: - import pyarrow.parquet as pp - - return pp - return None - - def get_numpy() -> Any: """Get numpy module (if already imported - else return None).""" return sys.modules.get("numpy", None) @@ -85,13 +74,104 @@ def is_pandas_dataframe(df: Any) -> TypeGuard[pd.DataFrame]: return bool((pd := get_pandas()) is not None and isinstance(df, pd.DataFrame)) +def is_pandas_series(ser: Any) -> TypeGuard[pd.Series[Any]]: + """Check whether `ser` is a pandas Series without importing pandas.""" + return bool((pd := get_pandas()) is not None and isinstance(ser, pd.Series)) + + +def is_modin_dataframe(df: Any) -> TypeGuard[mpd.DataFrame]: + """Check whether `df` is a modin DataFrame without importing modin.""" + return bool((pd := get_modin()) is not None and isinstance(df, pd.DataFrame)) + + +def is_modin_series(ser: Any) -> TypeGuard[mpd.Series]: + """Check whether `ser` is a modin Series without importing modin.""" + return bool((pd := get_modin()) is not None and isinstance(ser, pd.Series)) + + +def is_cudf_dataframe(df: Any) -> TypeGuard[cudf.DataFrame]: + """Check whether `df` is a cudf DataFrame without importing cudf.""" + return bool((pd := get_cudf()) is not None and isinstance(df, pd.DataFrame)) + + +def is_cudf_series(ser: Any) -> TypeGuard[pd.Series[Any]]: + """Check whether `ser` is a cudf Series without importing cudf.""" + return bool((pd := get_cudf()) is not None and isinstance(ser, pd.Series)) + + +def is_dask_dataframe(df: Any) -> TypeGuard[dd.DataFrame]: + """Check whether `df` is a Dask DataFrame without importing Dask.""" + return bool((dd := get_dask_dataframe()) is not None and isinstance(df, dd.DataFrame)) + + +def is_polars_dataframe(df: Any) -> TypeGuard[pl.DataFrame]: + """Check whether `df` is a Polars DataFrame without importing Polars.""" + return bool((pl := get_polars()) is not None and isinstance(df, pl.DataFrame)) + + +def is_polars_lazyframe(df: Any) -> TypeGuard[pl.LazyFrame]: + """Check whether `df` is a Polars LazyFrame without importing Polars.""" + return bool((pl := get_polars()) is not None and isinstance(df, pl.LazyFrame)) + + +def is_polars_series(ser: Any) -> TypeGuard[pl.Series]: + """Check whether `ser` is a Polars Series without importing Polars.""" + return bool((pl := get_polars()) is not None and isinstance(ser, pl.Series)) + + +def is_pyarrow_chunked_array(ser: Any) -> TypeGuard[pa.ChunkedArray]: + """Check whether `ser` is a PyArrow ChunkedArray without importing PyArrow.""" + return bool((pa := get_pyarrow()) is not None and isinstance(ser, pa.ChunkedArray)) + + +def is_pyarrow_table(df: Any) -> TypeGuard[pa.Table]: + """Check whether `df` is a PyArrow Table without importing PyArrow.""" + return bool((pa := get_pyarrow()) is not None and isinstance(df, pa.Table)) + + +def is_numpy_array(arr: Any) -> TypeGuard[np.ndarray]: + """Check whether `arr` is a NumPy Array without importing NumPy.""" + return bool((np := get_numpy()) is not None and isinstance(arr, np.ndarray)) + + +def is_pandas_like_dataframe(df: Any) -> bool: + """ + Check whether `df` is a pandas-like DataFrame without doing any imports + + By "pandas-like", we mean: pandas, Modin, cuDF. + """ + return is_pandas_dataframe(df) or is_modin_dataframe(df) or is_cudf_dataframe(df) + + +def is_pandas_like_series(arr: Any) -> bool: + """ + Check whether `arr` is a pandas-like Series without doing any imports + + By "pandas-like", we mean: pandas, Modin, cuDF. + """ + return is_pandas_series(arr) or is_modin_series(arr) or is_cudf_series(arr) + + __all__ = [ "get_polars", "get_pandas", "get_modin", "get_cudf", "get_pyarrow", - "get_pyarrow_compute", "get_numpy", "is_pandas_dataframe", + "is_pandas_series", + "is_polars_dataframe", + "is_polars_lazyframe", + "is_polars_series", + "is_modin_dataframe", + "is_modin_series", + "is_cudf_dataframe", + "is_cudf_series", + "is_pyarrow_table", + "is_pyarrow_chunked_array", + "is_numpy_array", + "is_dask_dataframe", + "is_pandas_like_dataframe", + "is_pandas_like_series", ] diff --git a/narwhals/dtypes.py b/narwhals/dtypes.py index 9b56d6141..4d8da4293 100644 --- a/narwhals/dtypes.py +++ b/narwhals/dtypes.py @@ -1,9 +1,6 @@ from __future__ import annotations from typing import TYPE_CHECKING -from typing import Any - -from narwhals.utils import isinstance_or_issubclass if TYPE_CHECKING: from typing_extensions import Self @@ -18,6 +15,8 @@ def is_numeric(cls: type[Self]) -> bool: return issubclass(cls, NumericType) def __eq__(self, other: DType | type[DType]) -> bool: # type: ignore[override] + from narwhals.utils import isinstance_or_issubclass + return isinstance_or_issubclass(other, type(self)) def __hash__(self) -> int: @@ -85,51 +84,3 @@ class Enum(DType): ... class Date(TemporalType): ... - - -def translate_dtype(plx: Any, dtype: DType) -> Any: - if "polars" in str(type(dtype)): - msg = ( - f"Expected Narwhals object, got: {type(dtype)}.\n\n" - "Perhaps you:\n" - "- Forgot a `nw.from_native` somewhere?\n" - "- Used `pl.Int64` instead of `nw.Int64`?" - ) - raise TypeError(msg) - if dtype == Float64: - return plx.Float64 - if dtype == Float32: - return plx.Float32 - if dtype == Int64: - return plx.Int64 - if dtype == Int32: - return plx.Int32 - if dtype == Int16: - return plx.Int16 - if dtype == Int8: - return plx.Int8 - if dtype == UInt64: - return plx.UInt64 - if dtype == UInt32: - return plx.UInt32 - if dtype == UInt16: - return plx.UInt16 - if dtype == UInt8: - return plx.UInt8 - if dtype == String: - return plx.String - if dtype == Boolean: - return plx.Boolean - if dtype == Categorical: - return plx.Categorical - if dtype == Enum: - msg = "Converting to Enum is not (yet) supported" - raise NotImplementedError(msg) - if dtype == Datetime: - return plx.Datetime - if dtype == Duration: - return plx.Duration - if dtype == Date: - return plx.Date - msg = f"Unknown dtype: {dtype}" # pragma: no cover - raise AssertionError(msg) diff --git a/narwhals/expr.py b/narwhals/expr.py index b4547883b..d8acadd20 100644 --- a/narwhals/expr.py +++ b/narwhals/expr.py @@ -6,14 +6,13 @@ from typing import Iterable from typing import Literal -from narwhals.dependencies import get_numpy -from narwhals.dtypes import DType -from narwhals.dtypes import translate_dtype +from narwhals.dependencies import is_numpy_array from narwhals.utils import flatten if TYPE_CHECKING: from typing_extensions import Self + from narwhals.dtypes import DType from narwhals.typing import IntoExpr @@ -77,6 +76,47 @@ def alias(self, name: str) -> Self: """ return self.__class__(lambda plx: self._call(plx).alias(name)) + def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self: + """ + Pipe function call. + + Examples: + >>> import polars as pl + >>> import pandas as pd + >>> import narwhals as nw + >>> data = {"a": [1, 2, 3, 4]} + >>> df_pd = pd.DataFrame(data) + >>> df_pl = pl.DataFrame(data) + + Lets define a library-agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... return df.select(nw.col("a").pipe(lambda x: x + 1)) + + We can then pass any supported library: + + >>> func(df_pd) + a + 0 2 + 1 3 + 2 4 + 3 5 + >>> func(df_pl) + shape: (4, 1) + ┌─────┐ + │ a │ + │ --- │ + │ i64 │ + ╞═════╡ + │ 2 │ + │ 3 │ + │ 4 │ + │ 5 │ + └─────┘ + """ + return function(self, *args, **kwargs) + def cast( self, dtype: Any, @@ -122,9 +162,8 @@ def cast( │ 3.0 ┆ 8 │ └─────┴─────┘ """ - return self.__class__( - lambda plx: self._call(plx).cast(translate_dtype(plx, dtype)), + lambda plx: self._call(plx).cast(dtype), ) # --- binary --- @@ -459,7 +498,6 @@ def min(self) -> Self: │ 1 ┆ 3 │ └─────┴─────┘ """ - return self.__class__(lambda plx: self._call(plx).min()) def max(self) -> Self: @@ -1076,8 +1114,8 @@ def arg_true(self) -> Self: >>> func(df_pd) a - 0 1 - 1 2 + 1 1 + 2 2 >>> func(df_pl) shape: (2, 1) ┌─────┐ @@ -1309,13 +1347,13 @@ def is_duplicated(self) -> Self: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pd) a b 0 True True 1 False True 2 False False 3 True False - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (4, 2) ┌───────┬───────┐ │ a ┆ b │ @@ -1350,13 +1388,13 @@ def is_unique(self) -> Self: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pd) a b 0 False False 1 True False 2 True True 3 False True - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (4, 2) ┌───────┬───────┐ │ a ┆ b │ @@ -1369,7 +1407,6 @@ def is_unique(self) -> Self: │ false ┆ true │ └───────┴───────┘ """ - return self.__class__(lambda plx: self._call(plx).is_unique()) def null_count(self) -> Self: @@ -1431,13 +1468,13 @@ def is_first_distinct(self) -> Self: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pd) a b 0 True True 1 True False 2 True True 3 False True - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (4, 2) ┌───────┬───────┐ │ a ┆ b │ @@ -1471,13 +1508,13 @@ def is_last_distinct(self) -> Self: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pd) a b 0 False False 1 True True 2 True True 3 True True - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (4, 2) ┌───────┬───────┐ │ a ┆ b │ @@ -1500,7 +1537,9 @@ def quantile( r"""Get quantile value. Note: - pandas and Polars may have implementation differences for a given interpolation method. + * pandas and Polars may have implementation differences for a given interpolation method. + * [dask](https://docs.dask.org/en/stable/generated/dask.dataframe.Series.quantile.html) has its own method to approximate quantile and it doesn't implement 'nearest', 'higher', 'lower', 'midpoint' + as interpolation method - use 'linear' which is closest to the native 'dask' - method. Arguments: quantile : float @@ -1524,11 +1563,11 @@ def quantile( We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE - a b + >>> func(df_pd) + a b 0 24.5 74.5 - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (1, 2) ┌──────┬──────┐ │ a ┆ b │ @@ -1566,12 +1605,12 @@ def head(self, n: int = 10) -> Self: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pd) a 0 0 1 1 2 2 - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (3, 1) ┌─────┐ │ a │ @@ -1583,7 +1622,6 @@ def head(self, n: int = 10) -> Self: │ 2 │ └─────┘ """ - return self.__class__(lambda plx: self._call(plx).head(n)) def tail(self, n: int = 10) -> Self: @@ -1610,12 +1648,12 @@ def tail(self, n: int = 10) -> Self: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE - a + >>> func(df_pd) + a 7 7 8 8 9 9 - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (3, 1) ┌─────┐ │ a │ @@ -1627,7 +1665,6 @@ def tail(self, n: int = 10) -> Self: │ 9 │ └─────┘ """ - return self.__class__(lambda plx: self._call(plx).tail(n)) def round(self, decimals: int = 0) -> Self: @@ -1638,12 +1675,12 @@ def round(self, decimals: int = 0) -> Self: decimals: Number of decimals to round by. Notes: - For values exactly halfway between rounded decimal values pandas and Polars behave differently. + For values exactly halfway between rounded decimal values pandas behaves differently than Polars and Arrow. pandas rounds to the nearest even value (e.g. -0.5 and 0.5 round to 0.0, 1.5 and 2.5 round to 2.0, 3.5 and 4.5 to 4.0, etc..). - Polars rounds away from 0 (e.g. -0.5 to -1.0, 0.5 to 1.0, 1.5 to 2.0, 2.5 to 3.0, etc..). + Polars and Arrow round away from 0 (e.g. -0.5 to -1.0, 0.5 to 1.0, 1.5 to 2.0, 2.5 to 3.0, etc..). Examples: @@ -1662,12 +1699,12 @@ def round(self, decimals: int = 0) -> Self: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pd) a 0 1.1 1 2.6 2 3.9 - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (3, 1) ┌─────┐ │ a │ @@ -1679,7 +1716,6 @@ def round(self, decimals: int = 0) -> Self: │ 3.9 │ └─────┘ """ - return self.__class__(lambda plx: self._call(plx).round(decimals)) def len(self) -> Self: @@ -1707,10 +1743,10 @@ def len(self) -> Self: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE - a1 a2 - 0 2 1 - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pd) + a1 a2 + 0 2 1 + >>> func(df_pl) shape: (1, 2) ┌─────┬─────┐ │ a1 ┆ a2 │ @@ -1765,6 +1801,118 @@ def gather_every(self: Self, n: int, offset: int = 0) -> Self: lambda plx: self._call(plx).gather_every(n=n, offset=offset) ) + # need to allow numeric typing + # TODO @aivanoved: make type alias for numeric type + def clip( + self, + lower_bound: Any | None = None, + upper_bound: Any | None = None, + ) -> Self: + r""" + Clip values in the Series. + + Arguments: + lower_bound: Lower bound value. + upper_bound: Upper bound value. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> + >>> s = [1, 2, 3] + >>> df_pd = pd.DataFrame({"s": s}) + >>> df_pl = pl.DataFrame({"s": s}) + + We define a library agnostic function: + + >>> @nw.narwhalify + ... def func_lower(df): + ... return df.select(nw.col("s").clip(2)) + + We can then pass either pandas or Polars to `func_lower`: + + >>> func_lower(df_pd) + s + 0 2 + 1 2 + 2 3 + >>> func_lower(df_pl) + shape: (3, 1) + ┌─────┐ + │ s │ + │ --- │ + │ i64 │ + ╞═════╡ + │ 2 │ + │ 2 │ + │ 3 │ + └─────┘ + + We define another library agnostic function: + + >>> @nw.narwhalify + ... def func_upper(df): + ... return df.select(nw.col("s").clip(upper_bound=2)) + + We can then pass either pandas or Polars to `func_upper`: + + >>> func_upper(df_pd) + s + 0 1 + 1 2 + 2 2 + >>> func_upper(df_pl) + shape: (3, 1) + ┌─────┐ + │ s │ + │ --- │ + │ i64 │ + ╞═════╡ + │ 1 │ + │ 2 │ + │ 2 │ + └─────┘ + + We can have both at the same time + + >>> s = [-1, 1, -3, 3, -5, 5] + >>> df_pd = pd.DataFrame({"s": s}) + >>> df_pl = pl.DataFrame({"s": s}) + + We define a library agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... return df.select(nw.col("s").clip(-1, 3)) + + We can pass either pandas or Polars to `func`: + + >>> func(df_pd) + s + 0 -1 + 1 1 + 2 -1 + 3 3 + 4 -1 + 5 3 + >>> func(df_pl) + shape: (6, 1) + ┌─────┐ + │ s │ + │ --- │ + │ i64 │ + ╞═════╡ + │ -1 │ + │ 1 │ + │ -1 │ + │ 3 │ + │ -1 │ + │ 3 │ + └─────┘ + """ + return self.__class__(lambda plx: self._call(plx).clip(lower_bound, upper_bound)) + @property def str(self: Self) -> ExprStringNamespace: return ExprStringNamespace(self) @@ -1833,6 +1981,121 @@ class ExprStringNamespace: def __init__(self, expr: Expr) -> None: self._expr = expr + def replace( + self, pattern: str, value: str, *, literal: bool = False, n: int = 1 + ) -> Expr: + r""" + Replace first matching regex/literal substring with a new string value. + + Arguments: + pattern: A valid regular expression pattern. + value: String that will replace the matched substring. + literal: Treat `pattern` as a literal string. + n: Number of matches to replace. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> data = {"foo": ["123abc", "abc abc123"]} + >>> df_pd = pd.DataFrame(data) + >>> df_pl = pl.DataFrame(data) + + We define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... df = df.with_columns(replaced=nw.col("foo").str.replace("abc", "")) + ... return df.to_dict(as_series=False) + + We can then pass either pandas or Polars to `func`: + + >>> func(df_pd) + {'foo': ['123abc', 'abc abc123'], 'replaced': ['123', ' abc123']} + + >>> func(df_pl) + {'foo': ['123abc', 'abc abc123'], 'replaced': ['123', ' abc123']} + + """ + return self._expr.__class__( + lambda plx: self._expr._call(plx).str.replace( + pattern, value, literal=literal, n=n + ) + ) + + def replace_all(self, pattern: str, value: str, *, literal: bool = False) -> Expr: + r""" + Replace all matching regex/literal substring with a new string value. + + Arguments: + pattern: A valid regular expression pattern. + value: String that will replace the matched substring. + literal: Treat `pattern` as a literal string. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> data = {"foo": ["123abc", "abc abc123"]} + >>> df_pd = pd.DataFrame(data) + >>> df_pl = pl.DataFrame(data) + + We define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... df = df.with_columns(replaced=nw.col("foo").str.replace_all("abc", "")) + ... return df.to_dict(as_series=False) + + We can then pass either pandas or Polars to `func`: + + >>> func(df_pd) + {'foo': ['123abc', 'abc abc123'], 'replaced': ['123', ' 123']} + + >>> func(df_pl) + {'foo': ['123abc', 'abc abc123'], 'replaced': ['123', ' 123']} + + """ + return self._expr.__class__( + lambda plx: self._expr._call(plx).str.replace_all( + pattern, value, literal=literal + ) + ) + + def strip_chars(self, characters: str | None = None) -> Expr: + r""" + Remove leading and trailing characters. + + Arguments: + characters: The set of characters to be removed. All combinations of this set of characters will be stripped from the start and end of the string. If set to None (default), all leading and trailing whitespace is removed instead. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> data = {"fruits": ["apple", "\nmango"]} + >>> df_pd = pd.DataFrame(data) + >>> df_pl = pl.DataFrame(data) + + We define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... df = df.with_columns(stripped=nw.col("fruits").str.strip_chars()) + ... return df.to_dict(as_series=False) + + We can then pass either pandas or Polars to `func`: + + >>> func(df_pd) + {'fruits': ['apple', '\nmango'], 'stripped': ['apple', 'mango']} + + >>> func(df_pl) + {'fruits': ['apple', '\nmango'], 'stripped': ['apple', 'mango']} + """ + return self._expr.__class__( + lambda plx: self._expr._call(plx).str.strip_chars(characters) + ) + def starts_with(self, prefix: str) -> Expr: r""" Check if string values start with a substring. @@ -1975,7 +2238,6 @@ def contains(self, pattern: str, *, literal: bool = False) -> Expr: │ null ┆ null ┆ null ┆ null │ └───────────────────┴───────────────┴────────────────────────┴───────────────┘ """ - return self._expr.__class__( lambda plx: self._expr._call(plx).str.contains(pattern, literal=literal) ) @@ -2012,7 +2274,7 @@ def slice(self, offset: int, length: int | None = None) -> Expr: 2 papaya ya 3 dragonfruit onf - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (4, 2) ┌─────────────┬──────────┐ │ s ┆ s_sliced │ @@ -2031,14 +2293,14 @@ def slice(self, offset: int, length: int | None = None) -> Expr: ... def func(df): ... return df.with_columns(s_sliced=nw.col("s").str.slice(-3)) - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pd) s s_sliced 0 pear ear 1 None None 2 papaya aya 3 dragonfruit uit - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (4, 2) ┌─────────────┬──────────┐ │ s ┆ s_sliced │ @@ -2225,8 +2487,8 @@ def to_uppercase(self) -> Expr: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE - fruits upper_col + >>> func(df_pd) + fruits upper_col 0 apple APPLE 1 mango MANGO 2 None None @@ -2266,13 +2528,13 @@ def to_lowercase(self) -> Expr: We can then pass either pandas or Polars to `func`: - >>> func(df_pd) # doctest: +NORMALIZE_WHITESPACE - fruits lower_col - 0 APPLE apple - 1 MANGO mango - 2 None None + >>> func(df_pd) + fruits lower_col + 0 APPLE apple + 1 MANGO mango + 2 None None - >>> func(df_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(df_pl) shape: (3, 2) ┌────────┬───────────┐ │ fruits ┆ lower_col │ @@ -2291,6 +2553,50 @@ class ExprDateTimeNamespace: def __init__(self, expr: Expr) -> None: self._expr = expr + def date(self) -> Expr: + """ + Extract the date from underlying DateTime representation. + + Raises: + NotImplementedError: If pandas default backend is being used. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> from datetime import datetime + >>> import narwhals as nw + >>> data = {"a": [datetime(2012, 1, 7, 10, 20), datetime(2023, 3, 10, 11, 32)]} + >>> df_pd = pd.DataFrame(data).convert_dtypes( + ... dtype_backend="pyarrow" + ... ) # doctest:+SKIP + >>> df_pl = pl.DataFrame(data) + + We define a library agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... return df.select(nw.col("a").dt.date()) + + We can then pass either pandas or Polars to `func`: + + >>> func(df_pd) # doctest:+SKIP + a + 0 2012-01-07 + 1 2023-03-10 + + >>> func(df_pl) # docetst + shape: (2, 1) + ┌────────────┐ + │ a │ + │ --- │ + │ date │ + ╞════════════╡ + │ 2012-01-07 │ + │ 2023-03-10 │ + └────────────┘ + """ + return self._expr.__class__(lambda plx: self._expr._call(plx).dt.date()) + def year(self) -> Expr: """ Extract year from underlying DateTime representation. @@ -3134,7 +3440,6 @@ def keep(self: Self) -> Expr: >>> func(df_pl).columns ['foo'] """ - return self._expr.__class__(lambda plx: self._expr._call(plx).name.keep()) def map(self: Self, function: Callable[[str], str]) -> Expr: @@ -3171,7 +3476,6 @@ def map(self: Self, function: Callable[[str], str]) -> Expr: >>> func(df_pl).columns ['oof', 'RAB'] """ - return self._expr.__class__(lambda plx: self._expr._call(plx).name.map(function)) def prefix(self: Self, prefix: str) -> Expr: @@ -3353,7 +3657,8 @@ def func(plx: Any) -> Any: return Expr(func) -def all() -> Expr: +# Add underscore so it doesn't conflict with builtin `all` +def all_() -> Expr: """ Instantiate an expression representing all columns. @@ -3392,7 +3697,8 @@ def all() -> Expr: return Expr(lambda plx: plx.all()) -def len() -> Expr: +# Add underscore so it doesn't conflict with builtin `len` +def len_() -> Expr: """ Return the number of rows. @@ -3599,17 +3905,22 @@ def max(*columns: str) -> Expr: def sum_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: """ - Sum all values horizontally across columns + Sum all values horizontally across columns. + + Warning: + Unlike Polars, we support horizontal sum over numeric columns only. Arguments: - exprs: Name(s) of the columns to use in the aggregation function. Accepts expression input. + exprs: Name(s) of the columns to use in the aggregation function. Accepts + expression input. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals as nw - >>> df_pl = pl.DataFrame({"a": [1, 2, 3], "b": [5, 10, 15]}) - >>> df_pd = pd.DataFrame({"a": [1, 2, 3], "b": [5, 10, 15]}) + >>> data = {"a": [1, 2, 3], "b": [5, 10, None]} + >>> df_pl = pl.DataFrame(data) + >>> df_pd = pd.DataFrame(data) We define a dataframe-agnostic function: @@ -3620,10 +3931,10 @@ def sum_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: We can then pass either pandas or polars to `func`: >>> func(df_pd) - a - 0 6 - 1 12 - 2 18 + a + 0 6.0 + 1 12.0 + 2 3.0 >>> func(df_pl) shape: (3, 1) ┌─────┐ @@ -3633,7 +3944,7 @@ def sum_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: ╞═════╡ │ 6 │ │ 12 │ - │ 18 │ + │ 3 │ └─────┘ """ return Expr( @@ -3643,17 +3954,18 @@ def sum_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: ) -def _extract_predicates(plx: Any, predicates: Iterable[IntoExpr]) -> Any: - return [extract_compliant(plx, v) for v in predicates] - - class When: def __init__(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> None: self._predicates = flatten([predicates]) + def _extract_predicates(self, plx: Any) -> Any: + return [extract_compliant(plx, v) for v in self._predicates] + def then(self, value: Any) -> Then: return Then( - lambda plx: plx.when(*_extract_predicates(plx, self._predicates)).then(value) + lambda plx: plx.when(*self._extract_predicates(plx)).then( + extract_compliant(plx, value) + ) ) @@ -3662,56 +3974,50 @@ def __init__(self, call: Callable[[Any], Any]) -> None: self._call = call def otherwise(self, value: Any) -> Expr: - return Expr(lambda plx: self._call(plx).otherwise(value)) + return Expr(lambda plx: self._call(plx).otherwise(extract_compliant(plx, value))) - def when( - self, - *predicates: IntoExpr | Iterable[IntoExpr], - ) -> ChainedWhen: - return ChainedWhen(self, *predicates) +# class ChainedWhen: +# def __init__( +# self, +# above_then: Then | ChainedThen, +# *predicates: IntoExpr | Iterable[IntoExpr], +# ) -> None: +# self._above_then = above_then +# self._predicates = predicates -class ChainedWhen: - def __init__( - self, - above_then: Then | ChainedThen, - *predicates: IntoExpr | Iterable[IntoExpr], - ) -> None: - self._above_then = above_then - self._predicates = predicates - - def then(self, value: Any) -> ChainedThen: - return ChainedThen( - lambda plx: self._above_then._call(plx) - .when(*_extract_predicates(plx, flatten([self._predicates]))) - .then(value) - ) +# def then(self, value: Any) -> ChainedThen: +# return ChainedThen( +# lambda plx: self._above_then._call(plx) +# .when(*_extract_predicates(plx, flatten([self._predicates]))) +# .then(value) +# ) -class ChainedThen(Expr): - def __init__(self, call: Callable[[Any], Any]) -> None: - self._call = call +# class ChainedThen(Expr): +# def __init__(self, call: Callable[[Any], Any]) -> None: +# self._call = call - def when( - self, - *predicates: IntoExpr | Iterable[IntoExpr], - ) -> ChainedWhen: - return ChainedWhen(self, *predicates) +# def when( +# self, +# *predicates: IntoExpr | Iterable[IntoExpr], +# ) -> ChainedWhen: +# return ChainedWhen(self, *predicates) - def otherwise(self, value: Any) -> Expr: - return Expr(lambda plx: self._call(plx).otherwise(value)) +# def otherwise(self, value: Any) -> Expr: +# return Expr(lambda plx: self._call(plx).otherwise(value)) def when(*predicates: IntoExpr | Iterable[IntoExpr]) -> When: """ Start a `when-then-otherwise` expression. + Expression similar to an `if-else` statement in Python. Always initiated by a `pl.when().then()`., and optionally followed by chaining one or more `.when().then()` statements. Chained when-then operations should be read as Python `if, elif, ... elif` blocks, not as `if, if, ... if`, i.e. the first condition that evaluates to `True` will be picked. If none of the conditions are `True`, an optional `.otherwise()` can be appended at the end. If not appended, and none of the conditions are `True`, `None` will be returned. - Parameters: - predicates - Condition(s) that must be met in order to apply the subsequent statement. Accepts one or more boolean expressions, which are implicitly combined with `&`. String input is parsed as a column name. + Arguments: + predicates: Condition(s) that must be met in order to apply the subsequent statement. Accepts one or more boolean expressions, which are implicitly combined with `&`. String input is parsed as a column name. Examples: >>> import pandas as pd @@ -3849,7 +4155,7 @@ def lit(value: Any, dtype: DType | None = None) -> Expr: └─────┴─────┘ """ - if (np := get_numpy()) is not None and isinstance(value, np.ndarray): + if is_numpy_array(value): msg = ( "numpy arrays are not supported as literal values. " "Consider using `with_columns` to create a new column from the array." @@ -3860,9 +4166,7 @@ def lit(value: Any, dtype: DType | None = None) -> Expr: msg = f"Nested datatypes are not supported yet. Got {value}" raise NotImplementedError(msg) - if dtype is None: - return Expr(lambda plx: plx.lit(value, dtype)) - return Expr(lambda plx: plx.lit(value, translate_dtype(plx, dtype))) + return Expr(lambda plx: plx.lit(value, dtype)) def any_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: @@ -3925,6 +4229,59 @@ def any_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: ) +def mean_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: + """ + Compute the mean of all values horizontally across columns. + + Arguments: + exprs: Name(s) of the columns to use in the aggregation function. Accepts + expression input. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> data = { + ... "a": [1, 8, 3], + ... "b": [4, 5, None], + ... "c": ["x", "y", "z"], + ... } + >>> df_pl = pl.DataFrame(data) + >>> df_pd = pd.DataFrame(data) + + We define a dataframe-agnostic function that computes the horizontal mean of "a" + and "b" columns: + + >>> @nw.narwhalify + ... def func(df): + ... return df.select(nw.mean_horizontal("a", "b")) + + We can then pass either pandas or polars to `func`: + + >>> func(df_pd) + a + 0 2.5 + 1 6.5 + 2 3.0 + >>> func(df_pl) + shape: (3, 1) + ┌─────┐ + │ a │ + │ --- │ + │ f64 │ + ╞═════╡ + │ 2.5 │ + │ 6.5 │ + │ 3.0 │ + └─────┘ + """ + return Expr( + lambda plx: plx.mean_horizontal( + *[extract_compliant(plx, v) for v in flatten(exprs)] + ) + ) + + __all__ = [ "Expr", ] diff --git a/narwhals/functions.py b/narwhals/functions.py index de7d45a15..51193c6c0 100644 --- a/narwhals/functions.py +++ b/narwhals/functions.py @@ -48,11 +48,105 @@ def concat( ) +def new_series( + name: str, + values: Any, + dtype: DType | type[DType] | None = None, + *, + native_namespace: ModuleType, +) -> Series: + """ + Instantiate Narwhals Series from raw data. + + Arguments: + name: Name of resulting Series. + values: Values of make Series from. + dtype: (Narwhals) dtype. If not provided, the native library + may auto-infer it from `values`. + native_namespace: The native library to use for DataFrame creation. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> data = {"a": [1, 2, 3], "b": [4, 5, 6]} + + Let's define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... values = [4, 1, 2] + ... native_namespace = nw.get_native_namespace(df) + ... return nw.new_series("c", values, nw.Int32, native_namespace=native_namespace) + + Let's see what happens when passing pandas / Polars input: + + >>> func(pd.DataFrame(data)) + 0 4 + 1 1 + 2 2 + Name: c, dtype: int32 + >>> func(pl.DataFrame(data)) # doctest: +NORMALIZE_WHITESPACE + shape: (3,) + Series: 'c' [i32] + [ + 4 + 1 + 2 + ] + """ + implementation = Implementation.from_native_namespace(native_namespace) + + if implementation is Implementation.POLARS: + if dtype: + from narwhals._polars.utils import ( + narwhals_to_native_dtype as polars_narwhals_to_native_dtype, + ) + + dtype = polars_narwhals_to_native_dtype(dtype) + + native_series = native_namespace.Series(name=name, values=values, dtype=dtype) + elif implementation in { + Implementation.PANDAS, + Implementation.MODIN, + Implementation.CUDF, + }: + if dtype: + from narwhals._pandas_like.utils import ( + narwhals_to_native_dtype as pandas_like_narwhals_to_native_dtype, + ) + + dtype = pandas_like_narwhals_to_native_dtype(dtype, None, implementation) + native_series = native_namespace.Series(values, name=name, dtype=dtype) + + elif implementation is Implementation.PYARROW: + if dtype: + from narwhals._arrow.utils import ( + narwhals_to_native_dtype as arrow_narwhals_to_native_dtype, + ) + + dtype = arrow_narwhals_to_native_dtype(dtype) + native_series = native_namespace.chunked_array([values], type=dtype) + + elif implementation is Implementation.DASK: + msg = "Dask support in Narwhals is lazy-only, so `new_series` is " "not supported" + raise NotImplementedError(msg) + else: # pragma: no cover + try: + # implementation is UNKNOWN, Narwhals extension using this feature should + # implement `from_dict` function in the top-level namespace. + native_series = native_namespace.new_series(name, values, dtype) + except AttributeError as e: + msg = "Unknown namespace is expected to implement `Series` constructor." + raise AttributeError(msg) from e + return from_native(native_series, series_only=True).alias(name) + + def from_dict( data: dict[str, Any], schema: dict[str, DType] | Schema | None = None, *, - native_namespace: ModuleType, + native_namespace: ModuleType | None = None, ) -> DataFrame[Any]: """ Instantiate DataFrame from dictionary. @@ -64,7 +158,8 @@ def from_dict( Arguments: data: Dictionary to create DataFrame from. schema: The DataFrame schema as Schema or dict of {name: type}. - native_namespace: The native library to use for DataFrame creation. + native_namespace: The native library to use for DataFrame creation. Only + necessary if inputs are not Narwhals Series. Examples: >>> import pandas as pd @@ -97,16 +192,31 @@ def from_dict( │ 2 ┆ 4 │ └─────┴─────┘ """ + from narwhals.series import Series + from narwhals.translate import to_native + + if not data: + msg = "from_dict cannot be called with empty dictionary" + raise ValueError(msg) + if native_namespace is None: + for val in data.values(): + if isinstance(val, Series): + native_namespace = val.__native_namespace__() + break + else: + msg = "Calling `from_dict` without `native_namespace` is only supported if all input values are already Narwhals Series" + raise TypeError(msg) + data = {key: to_native(value, strict=False) for key, value in data.items()} implementation = Implementation.from_native_namespace(native_namespace) if implementation is Implementation.POLARS: if schema: from narwhals._polars.utils import ( - reverse_translate_dtype as polars_reverse_translate_dtype, + narwhals_to_native_dtype as polars_narwhals_to_native_dtype, ) schema = { - name: polars_reverse_translate_dtype(dtype) + name: polars_narwhals_to_native_dtype(dtype) for name, dtype in schema.items() } @@ -120,11 +230,11 @@ def from_dict( if schema: from narwhals._pandas_like.utils import ( - reverse_translate_dtype as pandas_like_reverse_translate_dtype, + narwhals_to_native_dtype as pandas_like_narwhals_to_native_dtype, ) schema = { - name: pandas_like_reverse_translate_dtype( + name: pandas_like_narwhals_to_native_dtype( schema[name], native_type, implementation ) for name, native_type in native_frame.dtypes.items() @@ -134,19 +244,19 @@ def from_dict( elif implementation is Implementation.PYARROW: if schema: from narwhals._arrow.utils import ( - reverse_translate_dtype as arrow_reverse_translate_dtype, + narwhals_to_native_dtype as arrow_narwhals_to_native_dtype, ) schema = native_namespace.schema( [ - (name, arrow_reverse_translate_dtype(dtype)) + (name, arrow_narwhals_to_native_dtype(dtype)) for name, dtype in schema.items() ] ) native_frame = native_namespace.table(data, schema=schema) else: # pragma: no cover try: - # implementation is UNKNOWN, Narhwals extension using this feature should + # implementation is UNKNOWN, Narwhals extension using this feature should # implement `from_dict` function in the top-level namespace. native_frame = native_namespace.from_dict(data) except AttributeError as e: @@ -165,11 +275,11 @@ def _get_sys_info() -> dict[str, str]: """ python = sys.version.replace("\n", " ") - blob = [ + blob = ( ("python", python), ("executable", sys.executable), ("machine", platform.platform()), - ] + ) return dict(blob) @@ -185,14 +295,14 @@ def _get_deps_info() -> dict[str, str]: This function and show_versions were copied from sklearn and adapted """ - deps = [ + deps = ( "pandas", "polars", "cudf", "modin", "pyarrow", "numpy", - ] + ) from . import __version__ diff --git a/narwhals/selectors.py b/narwhals/selectors.py index fca1c4cdf..7c06a79c9 100644 --- a/narwhals/selectors.py +++ b/narwhals/selectors.py @@ -2,7 +2,6 @@ from typing import Any -from narwhals.dtypes import translate_dtype from narwhals.expr import Expr from narwhals.utils import flatten @@ -51,11 +50,7 @@ def by_dtype(*dtypes: Any) -> Expr: │ 4 ┆ 4.6 │ └─────┴─────┘ """ - return Selector( - lambda plx: plx.selectors.by_dtype( - [translate_dtype(plx, dtype) for dtype in flatten(dtypes)] - ) - ) + return Selector(lambda plx: plx.selectors.by_dtype(flatten(dtypes))) def numeric() -> Expr: diff --git a/narwhals/series.py b/narwhals/series.py index 27a594723..d80564d22 100644 --- a/narwhals/series.py +++ b/narwhals/series.py @@ -2,14 +2,17 @@ from typing import TYPE_CHECKING from typing import Any +from typing import Callable from typing import Literal from typing import Sequence from typing import overload -from narwhals.dtypes import translate_dtype +from narwhals.utils import parse_version if TYPE_CHECKING: import numpy as np + import pandas as pd + import pyarrow as pa from typing_extensions import Self from narwhals.dataframe import DataFrame @@ -56,8 +59,31 @@ def __getitem__(self, idx: int | slice | Sequence[int]) -> Any | Self: def __native_namespace__(self) -> Any: return self._compliant_series.__native_namespace__() - def __narwhals_namespace__(self) -> Any: - return self._compliant_series.__narwhals_namespace__() + def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: + """ + Export a Series via the Arrow PyCapsule Interface. + + Narwhals doesn't implement anything itself here: + + - if the underlying series implements the interface, it'll return that + - else, it'll call `to_arrow` and then defer to PyArrow's implementation + + See [PyCapsule Interface](https://arrow.apache.org/docs/dev/format/CDataInterface/PyCapsuleInterface.html) + for more. + """ + native_series = self._compliant_series._native_series + if hasattr(native_series, "__arrow_c_stream__"): + return native_series.__arrow_c_stream__(requested_schema=requested_schema) + try: + import pyarrow as pa # ignore-banned-import + except ModuleNotFoundError as exc: # pragma: no cover + msg = f"PyArrow>=16.0.0 is required for `Series.__arrow_c_stream__` for object of type {type(native_series)}" + raise ModuleNotFoundError(msg) from exc + if parse_version(pa.__version__) < (16, 0): # pragma: no cover + msg = f"PyArrow>=16.0.0 is required for `Series.__arrow_c_stream__` for object of type {type(native_series)}" + raise ModuleNotFoundError(msg) + ca = pa.chunked_array([self.to_arrow()]) + return ca.__arrow_c_stream__(requested_schema=requested_schema) @property def shape(self) -> tuple[int]: @@ -100,6 +126,44 @@ def _from_compliant_series(self, series: Any) -> Self: level=self._level, ) + def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self: + """ + Pipe function call. + + Examples: + >>> import polars as pl + >>> import pandas as pd + >>> import narwhals as nw + >>> s_pd = pd.Series([1, 2, 3, 4]) + >>> s_pl = pl.Series([1, 2, 3, 4]) + + Lets define a function to pipe into + >>> @nw.narwhalify + ... def func(s): + ... return s.pipe(lambda x: x + 2) + + Now apply it to the series + + >>> func(s_pd) + 0 3 + 1 4 + 2 5 + 3 6 + dtype: int64 + >>> func(s_pl) # doctest: +NORMALIZE_WHITESPACE + shape: (4,) + Series: '' [i64] + [ + 3 + 4 + 5 + 6 + ] + + + """ + return function(self, *args, **kwargs) + def __repr__(self) -> str: # pragma: no cover header = " Narwhals Series " length = len(header) @@ -139,9 +203,9 @@ def len(self) -> int: We can then pass either pandas or Polars to `func`: - >>> func(s_pd) # doctest: +NORMALIZE_WHITESPACE + >>> func(s_pd) 3 - >>> func(s_pl) # doctest: +NORMALIZE_WHITESPACE + >>> func(s_pl) 3 """ return len(self._compliant_series) @@ -242,11 +306,7 @@ def cast( 1 ] """ - return self._from_compliant_series( - self._compliant_series.cast( - translate_dtype(self.__narwhals_namespace__(), dtype) - ) - ) + return self._from_compliant_series(self._compliant_series.cast(dtype)) def to_frame(self) -> DataFrame[Any]: """ @@ -544,6 +604,107 @@ def std(self, *, ddof: int = 1) -> Any: """ return self._compliant_series.std(ddof=ddof) + def clip( + self, lower_bound: Any | None = None, upper_bound: Any | None = None + ) -> Self: + r""" + Clip values in the Series. + + Arguments: + lower_bound: Lower bound value. + upper_bound: Upper bound value. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> + >>> s = [1, 2, 3] + >>> s_pd = pd.Series(s) + >>> s_pl = pl.Series(s) + + We define a library agnostic function: + + >>> @nw.narwhalify + ... def func_lower(s): + ... return s.clip(2) + + We can then pass either pandas or Polars to `func_lower`: + + >>> func_lower(s_pd) + 0 2 + 1 2 + 2 3 + dtype: int64 + >>> func_lower(s_pl) # doctest: +NORMALIZE_WHITESPACE + shape: (3,) + Series: '' [i64] + [ + 2 + 2 + 3 + ] + + We define another library agnostic function: + + >>> @nw.narwhalify + ... def func_upper(s): + ... return s.clip(upper_bound=2) + + We can then pass either pandas or Polars to `func_upper`: + + >>> func_upper(s_pd) + 0 1 + 1 2 + 2 2 + dtype: int64 + >>> func_upper(s_pl) # doctest: +NORMALIZE_WHITESPACE + shape: (3,) + Series: '' [i64] + [ + 1 + 2 + 2 + ] + + We can have both at the same time + + >>> s = [-1, 1, -3, 3, -5, 5] + >>> s_pd = pd.Series(s) + >>> s_pl = pl.Series(s) + + We define a library agnostic function: + + >>> @nw.narwhalify + ... def func(s): + ... return s.clip(-1, 3) + + We can pass either pandas or Polars to `func`: + + >>> func(s_pd) + 0 -1 + 1 1 + 2 -1 + 3 3 + 4 -1 + 5 3 + dtype: int64 + >>> func(s_pl) # doctest: +NORMALIZE_WHITESPACE + shape: (6,) + Series: '' [i64] + [ + -1 + 1 + -1 + 3 + -1 + 3 + ] + """ + return self._from_compliant_series( + self._compliant_series.clip(lower_bound=lower_bound, upper_bound=upper_bound) + ) + def is_in(self, other: Any) -> Self: """ Check if the elements of this Series are in the other sequence. @@ -605,8 +766,8 @@ def arg_true(self) -> Self: We can then pass either pandas or Polars to `func`: >>> func(s_pd) - 0 1 - 1 2 + 1 1 + 2 2 Name: a, dtype: int64 >>> func(s_pl) # doctest: +NORMALIZE_WHITESPACE shape: (2,) @@ -1191,7 +1352,7 @@ def n_unique(self) -> int: """ return self._compliant_series.n_unique() # type: ignore[no-any-return] - def to_numpy(self) -> Any: + def to_numpy(self) -> np.ndarray: """ Convert to numpy. @@ -1218,7 +1379,7 @@ def to_numpy(self) -> Any: """ return self._compliant_series.to_numpy() - def to_pandas(self) -> Any: + def to_pandas(self) -> pd.Series: """ Convert to pandas. @@ -1537,7 +1698,6 @@ def null_count(self: Self) -> int: >>> func(s_pl) 2 """ - return self._compliant_series.null_count() # type: ignore[no-any-return] def is_first_distinct(self: Self) -> Self: @@ -1808,7 +1968,6 @@ def zip_with(self: Self, mask: Self, other: Self) -> Self: 4 5 dtype: int64 """ - return self._from_compliant_series( self._compliant_series.zip_with( self._extract_native(mask), self._extract_native(other) @@ -1882,7 +2041,6 @@ def head(self: Self, n: int = 10) -> Self: 2 ] """ - return self._from_compliant_series(self._compliant_series.head(n)) def tail(self: Self, n: int = 10) -> Self: @@ -1923,7 +2081,6 @@ def tail(self: Self, n: int = 10) -> Self: 9 ] """ - return self._from_compliant_series(self._compliant_series.tail(n)) def round(self: Self, decimals: int = 0) -> Self: @@ -1934,12 +2091,12 @@ def round(self: Self, decimals: int = 0) -> Self: decimals: Number of decimals to round by. Notes: - For values exactly halfway between rounded decimal values pandas and Polars behave differently. + For values exactly halfway between rounded decimal values pandas behaves differently than Polars and Arrow. pandas rounds to the nearest even value (e.g. -0.5 and 0.5 round to 0.0, 1.5 and 2.5 round to 2.0, 3.5 and 4.5 to 4.0, etc..). - Polars rounds away from 0 (e.g. -0.5 to -1.0, 0.5 to 1.0, 1.5 to 2.0, 2.5 to 3.0, etc..). + Polars and Arrow round away from 0 (e.g. -0.5 to -1.0, 0.5 to 1.0, 1.5 to 2.0, 2.5 to 3.0, etc..). Examples: >>> import narwhals as nw @@ -2039,7 +2196,6 @@ def to_dummies( │ 0 ┆ 1 │ └─────┴─────┘ """ - from narwhals.dataframe import DataFrame return DataFrame( @@ -2087,6 +2243,44 @@ def gather_every(self: Self, n: int, offset: int = 0) -> Self: self._compliant_series.gather_every(n=n, offset=offset) ) + def to_arrow(self: Self) -> pa.Array: + r""" + Convert to arrow. + + Examples: + >>> import narwhals as nw + >>> import pandas as pd + >>> import polars as pl + >>> data = [1, 2, 3, 4] + >>> s_pd = pd.Series(name="a", data=data) + >>> s_pl = pl.Series(name="a", values=data) + + Let's define a dataframe-agnostic function that converts to arrow: + + >>> @nw.narwhalify + ... def func(s): + ... return s.to_arrow() + + >>> func(s_pd) # doctest:+NORMALIZE_WHITESPACE + + [ + 1, + 2, + 3, + 4 + ] + + >>> func(s_pl) # doctest:+NORMALIZE_WHITESPACE + + [ + 1, + 2, + 3, + 4 + ] + """ + return self._compliant_series.to_arrow() + @property def str(self) -> SeriesStringNamespace: return SeriesStringNamespace(self) @@ -2148,6 +2342,119 @@ class SeriesStringNamespace: def __init__(self, series: Series) -> None: self._narwhals_series = series + def replace( + self, pattern: str, value: str, *, literal: bool = False, n: int = 1 + ) -> Series: + r""" + Replace first matching regex/literal substring with a new string value. + + Arguments: + pattern: A valid regular expression pattern. + value: String that will replace the matched substring. + literal: Treat `pattern` as a literal string. + n: Number of matches to replace. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> data = ["123abc", "abc abc123"] + >>> s_pd = pd.Series(data) + >>> s_pl = pl.Series(data) + + We define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(s): + ... s = s.str.replace("abc", "") + ... return s.to_list() + + We can then pass either pandas or Polars to `func`: + + >>> func(s_pd) + ['123', ' abc123'] + + >>> func(s_pl) + ['123', ' abc123'] + """ + return self._narwhals_series._from_compliant_series( + self._narwhals_series._compliant_series.str.replace( + pattern, value, literal=literal, n=n + ) + ) + + def replace_all(self, pattern: str, value: str, *, literal: bool = False) -> Series: + r""" + Replace all matching regex/literal substring with a new string value. + + Arguments: + pattern: A valid regular expression pattern. + value: String that will replace the matched substring. + literal: Treat `pattern` as a literal string. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> data = ["123abc", "abc abc123"] + >>> s_pd = pd.Series(data) + >>> s_pl = pl.Series(data) + + We define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(s): + ... s = s.str.replace_all("abc", "") + ... return s.to_list() + + We can then pass either pandas or Polars to `func`: + + >>> func(s_pd) + ['123', ' 123'] + + >>> func(s_pl) + ['123', ' 123'] + """ + return self._narwhals_series._from_compliant_series( + self._narwhals_series._compliant_series.str.replace_all( + pattern, value, literal=literal + ) + ) + + def strip_chars(self, characters: str | None = None) -> Series: + r""" + Remove leading and trailing characters. + + Arguments: + characters: The set of characters to be removed. All combinations of this set of characters will be stripped from the start and end of the string. If set to None (default), all leading and trailing whitespace is removed instead. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> data = ["apple", "\nmango"] + >>> s_pd = pd.Series(data) + >>> s_pl = pl.Series(data) + + We define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(s): + ... s = s.str.strip_chars() + ... return s.to_list() + + We can then pass either pandas or Polars to `func`: + + >>> func(s_pd) + ['apple', 'mango'] + + >>> func(s_pl) + ['apple', 'mango'] + """ + return self._narwhals_series._from_compliant_series( + self._narwhals_series._compliant_series.str.strip_chars(characters) + ) + def starts_with(self, prefix: str) -> Series: r""" Check if string values start with a substring. @@ -2543,6 +2850,49 @@ class SeriesDateTimeNamespace: def __init__(self, series: Series) -> None: self._narwhals_series = series + def date(self) -> Series: + """ + Get the date in a datetime series. + + Raises: + NotImplementedError: If pandas default backend is being used. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> from datetime import datetime + >>> import narwhals as nw + >>> dates = [datetime(2012, 1, 7, 10, 20), datetime(2023, 3, 10, 11, 32)] + >>> s_pd = pd.Series(dates).convert_dtypes( + ... dtype_backend="pyarrow" + ... ) # doctest:+SKIP + >>> s_pl = pl.Series(dates) + + We define a library agnostic function: + + >>> @nw.narwhalify + ... def func(s): + ... return s.dt.date() + + We can then pass either pandas or Polars to `func`: + + >>> func(s_pd) # doctest:+SKIP + 0 2012-01-07 + 1 2023-03-10 + dtype: date32[day][pyarrow] + + >>> func(s_pl) # doctest: +NORMALIZE_WHITESPACE + shape: (2,) + Series: '' [date] + [ + 2012-01-07 + 2023-03-10 + ] + """ + return self._narwhals_series._from_compliant_series( + self._narwhals_series._compliant_series.dt.date() + ) + def year(self) -> Series: """ Get the year in a datetime series. diff --git a/narwhals/stable/v1.py b/narwhals/stable/v1.py index aac1ba6b3..8363b36e9 100644 --- a/narwhals/stable/v1.py +++ b/narwhals/stable/v1.py @@ -1,6 +1,5 @@ from __future__ import annotations -from functools import wraps from typing import TYPE_CHECKING from typing import Any from typing import Callable @@ -11,6 +10,7 @@ from typing import overload import narwhals as nw +from narwhals import dependencies from narwhals import selectors from narwhals.dataframe import DataFrame as NwDataFrame from narwhals.dataframe import LazyFrame as NwLazyFrame @@ -33,8 +33,9 @@ from narwhals.dtypes import UInt32 from narwhals.dtypes import UInt64 from narwhals.dtypes import Unknown -from narwhals.expr import ChainedThen as NwChainedThen -from narwhals.expr import ChainedWhen as NwChainedWhen + +# from narwhals.expr import ChainedThen as NwChainedThen +# from narwhals.expr import ChainedWhen as NwChainedWhen from narwhals.expr import Expr as NwExpr from narwhals.expr import Then as NwThen from narwhals.expr import When as NwWhen @@ -44,12 +45,14 @@ from narwhals.schema import Schema as NwSchema from narwhals.series import Series as NwSeries from narwhals.translate import get_native_namespace as nw_get_native_namespace +from narwhals.translate import narwhalify as nw_narwhalify from narwhals.translate import to_native from narwhals.typing import IntoDataFrameT from narwhals.typing import IntoFrameT from narwhals.utils import is_ordered_categorical as nw_is_ordered_categorical from narwhals.utils import maybe_align_index as nw_maybe_align_index from narwhals.utils import maybe_convert_dtypes as nw_maybe_convert_dtypes +from narwhals.utils import maybe_get_index as nw_maybe_get_index from narwhals.utils import maybe_set_index as nw_maybe_set_index if TYPE_CHECKING: @@ -73,6 +76,8 @@ class DataFrame(NwDataFrame[IntoDataFrameT]): `narwhals.from_native`. """ + @overload + def __getitem__(self, item: tuple[Sequence[int], slice]) -> Self: ... @overload def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ... @@ -156,7 +161,7 @@ def to_dict( ... "A": [1, 2, 3, 4, 5], ... "fruits": ["banana", "banana", "apple", "apple", "banana"], ... "B": [5, 4, 3, 2, 1], - ... "cars": ["beetle", "audi", "beetle", "beetle", "beetle"], + ... "animals": ["beetle", "fly", "beetle", "beetle", "beetle"], ... "optional": [28, 300, None, 2, -30], ... } >>> df_pd = pd.DataFrame(df) @@ -171,9 +176,9 @@ def to_dict( We can then pass either pandas or Polars to `func`: >>> func(df_pd) - {'A': [1, 2, 3, 4, 5], 'fruits': ['banana', 'banana', 'apple', 'apple', 'banana'], 'B': [5, 4, 3, 2, 1], 'cars': ['beetle', 'audi', 'beetle', 'beetle', 'beetle'], 'optional': [28.0, 300.0, nan, 2.0, -30.0]} + {'A': [1, 2, 3, 4, 5], 'fruits': ['banana', 'banana', 'apple', 'apple', 'banana'], 'B': [5, 4, 3, 2, 1], 'animals': ['beetle', 'fly', 'beetle', 'beetle', 'beetle'], 'optional': [28.0, 300.0, nan, 2.0, -30.0]} >>> func(df_pl) - {'A': [1, 2, 3, 4, 5], 'fruits': ['banana', 'banana', 'apple', 'apple', 'banana'], 'B': [5, 4, 3, 2, 1], 'cars': ['beetle', 'audi', 'beetle', 'beetle', 'beetle'], 'optional': [28, 300, None, 2, -30]} + {'A': [1, 2, 3, 4, 5], 'fruits': ['banana', 'banana', 'apple', 'apple', 'banana'], 'B': [5, 4, 3, 2, 1], 'animals': ['beetle', 'fly', 'beetle', 'beetle', 'beetle'], 'optional': [28, 300, None, 2, -30]} """ if as_series: return {key: _stableify(value) for key, value in super().to_dict().items()} @@ -482,12 +487,12 @@ def _stableify(obj: NwSeries) -> Series: ... def _stableify(obj: NwExpr) -> Expr: ... @overload def _stableify(when_then: NwWhen) -> When: ... -@overload -def _stableify(when_then: NwChainedWhen) -> ChainedWhen: ... +# @overload +# def _stableify(when_then: NwChainedWhen) -> ChainedWhen: ... @overload def _stableify(when_then: NwThen) -> Then: ... -@overload -def _stableify(when_then: NwChainedThen) -> ChainedThen: ... +# @overload +# def _stableify(when_then: NwChainedThen) -> ChainedThen: ... @overload def _stableify(obj: Any) -> Any: ... @@ -498,9 +503,9 @@ def _stableify( | NwSeries | NwExpr | NwWhen - | NwChainedWhen + # | NwChainedWhen | NwThen - | NwChainedThen + # | NwChainedThen | Any, ) -> ( DataFrame[IntoFrameT] @@ -508,9 +513,9 @@ def _stableify( | Series | Expr | When - | ChainedWhen + # | ChainedWhen | Then - | ChainedThen + # | ChainedThen | Any ): if isinstance(obj, NwDataFrame): @@ -528,14 +533,14 @@ def _stableify( obj._compliant_series, level=obj._level, ) - elif isinstance(obj, NwChainedWhen): - return ChainedWhen.from_base(obj) + # elif isinstance(obj, NwChainedWhen): + # return ChainedWhen.from_base(obj) if isinstance(obj, NwWhen): - return When.from_base(obj) - elif isinstance(obj, NwChainedThen): - return ChainedThen.from_base(obj) + return When.from_when(obj) + # elif isinstance(obj, NwChainedThen): + # return ChainedThen.from_base(obj) elif isinstance(obj, NwThen): - return Then.from_base(obj) + return Then.from_then(obj) if isinstance(obj, NwExpr): return Expr(obj._call) return obj @@ -636,8 +641,8 @@ def from_native( allow_series: None = ..., ) -> DataFrame[IntoDataFrameT]: """ - from_native(df, strict=True, eager_or_interchange_only=True, allow_series=True) - from_native(df, eager_or_interchange_only=True, allow_series=True) + from_native(df, strict=True, eager_or_interchange_only=True) + from_native(df, eager_or_interchange_only=True) """ @@ -652,8 +657,8 @@ def from_native( allow_series: None = ..., ) -> DataFrame[IntoDataFrameT]: """ - from_native(df, strict=True, eager_only=True, allow_series=True) - from_native(df, eager_only=True, allow_series=True) + from_native(df, strict=True, eager_only=True) + from_native(df, eager_only=True) """ @@ -668,8 +673,8 @@ def from_native( allow_series: Literal[True], ) -> DataFrame[Any] | LazyFrame[Any] | Series: """ - from_native(df, strict=True, eager_only=True) - from_native(df, eager_only=True) + from_native(df, strict=True, allow_series=True) + from_native(df, allow_series=True) """ @@ -773,6 +778,7 @@ def narwhalify( *, strict: bool = False, eager_only: bool | None = False, + eager_or_interchange_only: bool | None = False, series_only: bool | None = False, allow_series: bool | None = True, ) -> Callable[..., Any]: @@ -815,7 +821,7 @@ def func(df): You can also pass in extra arguments, e.g. ```python - @nw.narhwalify(eager_only=True) + @nw.narwhalify(eager_only=True) ``` that will get passed down to `nw.from_native`. @@ -831,53 +837,14 @@ def func(df): allow_series: Whether to allow series (default is only dataframe / lazyframe). """ - # TODO(Unassigned): do we have a way to de-dupe this a bit? - def decorator(func: Callable[..., Any]) -> Callable[..., Any]: - @wraps(func) - def wrapper(*args: Any, **kwargs: Any) -> Any: - args = [ - from_native( - arg, - strict=strict, - eager_only=eager_only, - series_only=series_only, - allow_series=allow_series, - ) - for arg in args - ] # type: ignore[assignment] - - kwargs = { - name: from_native( - value, - strict=strict, - eager_only=eager_only, - series_only=series_only, - allow_series=allow_series, - ) - for name, value in kwargs.items() - } - - backends = { - b() - for v in [*args, *kwargs.values()] - if (b := getattr(v, "__native_namespace__", None)) - } - - if backends.__len__() > 1: - msg = "Found multiple backends. Make sure that all dataframe/series inputs come from the same backend." - raise ValueError(msg) - - result = func(*args, **kwargs) - - return to_native(result, strict=strict) - - return wrapper - - if func is None: - return decorator - else: - # If func is not None, it means the decorator is used without arguments - return decorator(func) + return nw_narwhalify( + func=func, + strict=strict, + eager_only=eager_only, + eager_or_interchange_only=eager_or_interchange_only, + series_only=series_only, + allow_series=allow_series, + ) def all() -> Expr: @@ -1202,17 +1169,22 @@ def sum(*columns: str) -> Expr: def sum_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: """ - Sum all values horizontally across columns + Sum all values horizontally across columns. + + Warning: + Unlike Polars, we support horizontal sum over numeric columns only. Arguments: - exprs: Name(s) of the columns to use in the aggregation function. Accepts expression input. + exprs: Name(s) of the columns to use in the aggregation function. Accepts + expression input. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals.stable.v1 as nw - >>> df_pl = pl.DataFrame({"a": [1, 2, 3], "b": [5, 10, 15]}) - >>> df_pd = pd.DataFrame({"a": [1, 2, 3], "b": [5, 10, 15]}) + >>> data = {"a": [1, 2, 3], "b": [5, 10, None]} + >>> df_pl = pl.DataFrame(data) + >>> df_pd = pd.DataFrame(data) We define a dataframe-agnostic function: @@ -1223,10 +1195,10 @@ def sum_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: We can then pass either pandas or polars to `func`: >>> func(df_pd) - a - 0 6 - 1 12 - 2 18 + a + 0 6.0 + 1 12.0 + 2 3.0 >>> func(df_pl) shape: (3, 1) ┌─────┐ @@ -1236,7 +1208,7 @@ def sum_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: ╞═════╡ │ 6 │ │ 12 │ - │ 18 │ + │ 3 │ └─────┘ """ return _stableify(nw.sum_horizontal(*exprs)) @@ -1354,6 +1326,55 @@ def any_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: return _stableify(nw.any_horizontal(*exprs)) +def mean_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: + """ + Compute the mean of all values horizontally across columns. + + Arguments: + exprs: Name(s) of the columns to use in the aggregation function. Accepts + expression input. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals.stable.v1 as nw + >>> data = { + ... "a": [1, 8, 3], + ... "b": [4, 5, None], + ... "c": ["x", "y", "z"], + ... } + >>> df_pl = pl.DataFrame(data) + >>> df_pd = pd.DataFrame(data) + + We define a dataframe-agnostic function that computes the horizontal mean of "a" + and "b" columns: + + >>> @nw.narwhalify + ... def func(df): + ... return df.select(nw.mean_horizontal("a", "b")) + + We can then pass either pandas or polars to `func`: + + >>> func(df_pd) + a + 0 2.5 + 1 6.5 + 2 3.0 + >>> func(df_pl) + shape: (3, 1) + ┌─────┐ + │ a │ + │ --- │ + │ f64 │ + ╞═════╡ + │ 2.5 │ + │ 6.5 │ + │ 3.0 │ + └─────┘ + """ + return _stableify(nw.mean_horizontal(*exprs)) + + def is_ordered_categorical(series: Series) -> bool: """ Return whether indices of categories are semantically meaningful. @@ -1424,7 +1445,12 @@ def maybe_align_index(lhs: T, rhs: Series | DataFrame[Any] | LazyFrame[Any]) -> def maybe_convert_dtypes(df: T, *args: bool, **kwargs: bool | str) -> T: """ - Convert columns to the best possible dtypes using dtypes supporting ``pd.NA``, if df is pandas-like. + Convert columns or series to the best possible dtypes using dtypes supporting ``pd.NA``, if df is pandas-like. + + Arguments: + obj: DataFrame or Series. + *args: Additional arguments which gets passed through. + **kwargs: Additional arguments which gets passed through. Notes: For non-pandas-like inputs, this is a no-op. @@ -1450,6 +1476,33 @@ def maybe_convert_dtypes(df: T, *args: bool, **kwargs: bool | str) -> T: return nw_maybe_convert_dtypes(df, *args, **kwargs) +def maybe_get_index(obj: T) -> Any | None: + """ + Get the index of a DataFrame or a Series, if it's pandas-like. + + Notes: + This is only really intended for backwards-compatibility purposes, + for example if your library already aligns indices for users. + If you're designing a new library, we highly encourage you to not + rely on the Index. + For non-pandas-like inputs, this returns `None`. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals.stable.v1 as nw + >>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}) + >>> df = nw.from_native(df_pd) + >>> nw.maybe_get_index(df) + RangeIndex(start=0, stop=2, step=1) + >>> series_pd = pd.Series([1, 2]) + >>> series = nw.from_native(series_pd, series_only=True) + >>> nw.maybe_get_index(series) + RangeIndex(start=0, stop=2, step=1) + """ + return nw_maybe_get_index(obj) + + def maybe_set_index(df: T, column_names: str | list[str]) -> T: """ Set columns `columns` to be the index of `df`, if `df` is pandas-like. @@ -1510,42 +1563,18 @@ def get_level( class When(NwWhen): @classmethod - def from_base(cls, when: NwWhen) -> Self: + def from_when(cls, when: NwWhen) -> Self: return cls(*when._predicates) def then(self, value: Any) -> Then: - return _stableify(super().then(value)) + return Then.from_then(super().then(value)) class Then(NwThen, Expr): @classmethod - def from_base(cls, then: NwThen) -> Self: + def from_then(cls, then: NwThen) -> Self: return cls(then._call) - def when(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> ChainedWhen: - return _stableify(super().when(*predicates)) - - def otherwise(self, value: Any) -> Expr: - return _stableify(super().otherwise(value)) - - -class ChainedWhen(NwChainedWhen): - @classmethod - def from_base(cls, chained_when: NwChainedWhen) -> Self: - return cls(_stableify(chained_when._above_then), *chained_when._predicates) - - def then(self, value: Any) -> ChainedThen: - return _stableify(super().then(value)) - - -class ChainedThen(NwChainedThen, Expr): - @classmethod - def from_base(cls, chained_then: NwChainedThen) -> Self: - return cls(chained_then._call) - - def when(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> ChainedWhen: - return _stableify(super().when(*predicates)) - def otherwise(self, value: Any) -> Expr: return _stableify(super().otherwise(value)) @@ -1553,13 +1582,13 @@ def otherwise(self, value: Any) -> Expr: def when(*predicates: IntoExpr | Iterable[IntoExpr]) -> When: """ Start a `when-then-otherwise` expression. + Expression similar to an `if-else` statement in Python. Always initiated by a `pl.when().then()`., and optionally followed by chaining one or more `.when().then()` statements. Chained when-then operations should be read as Python `if, elif, ... elif` blocks, not as `if, if, ... if`, i.e. the first condition that evaluates to `True` will be picked. If none of the conditions are `True`, an optional `.otherwise()` can be appended at the end. If not appended, and none of the conditions are `True`, `None` will be returned. - Parameters: - predicates - Condition(s) that must be met in order to apply the subsequent statement. Accepts one or more boolean expressions, which are implicitly combined with `&`. String input is parsed as a column name. + Arguments: + predicates: Condition(s) that must be met in order to apply the subsequent statement. Accepts one or more boolean expressions, which are implicitly combined with `&`. String input is parsed as a column name. Examples: >>> import pandas as pd @@ -1595,14 +1624,87 @@ def when(*predicates: IntoExpr | Iterable[IntoExpr]) -> When: │ 3 ┆ 15 ┆ 6 │ └─────┴─────┴────────┘ """ - return _stableify(nw_when(*predicates)) + return When.from_when(nw_when(*predicates)) + + +def new_series( + name: str, + values: Any, + dtype: DType | type[DType] | None = None, + *, + native_namespace: ModuleType, +) -> Series: + """ + Instantiate Narwhals Series from raw data. + + Arguments: + name: Name of resulting Series. + values: Values of make Series from. + dtype: (Narwhals) dtype. If not provided, the native library + may auto-infer it from `values`. + native_namespace: The native library to use for DataFrame creation. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals.stable.v1 as nw + >>> data = {"a": [1, 2, 3], "b": [4, 5, 6]} + + Let's define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... values = [4, 1, 2] + ... native_namespace = nw.get_native_namespace(df) + ... return nw.new_series("c", values, nw.Int32, native_namespace=native_namespace) + + Let's see what happens when passing pandas / Polars input: + + >>> func(pd.DataFrame(data)) + 0 4 + 1 1 + 2 2 + Name: c, dtype: int32 + >>> func(pl.DataFrame(data)) # doctest: +NORMALIZE_WHITESPACE + shape: (3,) + Series: 'c' [i32] + [ + 4 + 1 + 2 + ] + """ + return _stableify( + nw.new_series(name, values, dtype, native_namespace=native_namespace) + ) + + +# class ChainedWhen(NwChainedWhen): +# @classmethod +# def from_base(cls, chained_when: NwChainedWhen) -> Self: +# return cls(_stableify(chained_when._above_then), *chained_when._predicates) + +# def then(self, value: Any) -> ChainedThen: +# return _stableify(super().then(value)) + + +# class ChainedThen(NwChainedThen, Expr): +# @classmethod +# def from_base(cls, chained_then: NwChainedThen) -> Self: +# return cls(chained_then._call) + +# def when(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> ChainedWhen: +# return _stableify(super().when(*predicates)) + +# def otherwise(self, value: Any) -> Expr: +# return _stableify(super().otherwise(value)) def from_dict( data: dict[str, Any], schema: dict[str, DType] | Schema | None = None, *, - native_namespace: ModuleType, + native_namespace: ModuleType | None = None, ) -> DataFrame[Any]: """ Instantiate DataFrame from dictionary. @@ -1614,7 +1716,8 @@ def from_dict( Arguments: data: Dictionary to create DataFrame from. schema: The DataFrame schema as Schema or dict of {name: type}. - native_namespace: The native library to use for DataFrame creation. + native_namespace: The native library to use for DataFrame creation. Only + necessary if inputs are not Narwhals Series. Examples: >>> import pandas as pd @@ -1655,11 +1758,13 @@ def from_dict( __all__ = [ "selectors", "concat", + "dependencies", "to_native", "from_native", "is_ordered_categorical", "maybe_align_index", "maybe_convert_dtypes", + "maybe_get_index", "maybe_set_index", "get_native_namespace", "get_level", @@ -1672,6 +1777,7 @@ def from_dict( "min", "max", "mean", + "mean_horizontal", "sum", "sum_horizontal", "when", @@ -1702,4 +1808,5 @@ def from_dict( "show_versions", "Schema", "from_dict", + "new_series", ] diff --git a/narwhals/translate.py b/narwhals/translate.py index 3d5cc281b..c19ea7192 100644 --- a/narwhals/translate.py +++ b/narwhals/translate.py @@ -10,12 +10,23 @@ from narwhals.dependencies import get_cudf from narwhals.dependencies import get_dask -from narwhals.dependencies import get_dask_dataframe from narwhals.dependencies import get_dask_expr from narwhals.dependencies import get_modin from narwhals.dependencies import get_pandas from narwhals.dependencies import get_polars from narwhals.dependencies import get_pyarrow +from narwhals.dependencies import is_cudf_dataframe +from narwhals.dependencies import is_cudf_series +from narwhals.dependencies import is_dask_dataframe +from narwhals.dependencies import is_modin_dataframe +from narwhals.dependencies import is_modin_series +from narwhals.dependencies import is_pandas_dataframe +from narwhals.dependencies import is_pandas_series +from narwhals.dependencies import is_polars_dataframe +from narwhals.dependencies import is_polars_lazyframe +from narwhals.dependencies import is_polars_series +from narwhals.dependencies import is_pyarrow_chunked_array +from narwhals.dependencies import is_pyarrow_table if TYPE_CHECKING: from narwhals.dataframe import DataFrame @@ -60,7 +71,7 @@ def to_native( from narwhals.series import Series if isinstance(narwhals_object, BaseFrame): - return narwhals_object._compliant_frame._native_dataframe + return narwhals_object._compliant_frame._native_frame if isinstance(narwhals_object, Series): return narwhals_object._compliant_series._native_series @@ -340,27 +351,33 @@ def from_native( # noqa: PLR0915 level="full", ) + # TODO(marco): write all of these in terms of `is_` rather + # than `get_` + walrus + # Polars - elif (pl := get_polars()) is not None and isinstance(native_object, pl.DataFrame): + elif is_polars_dataframe(native_object): if series_only: msg = "Cannot only use `series_only` with polars.DataFrame" raise TypeError(msg) + pl = get_polars() return DataFrame( PolarsDataFrame(native_object, backend_version=parse_version(pl.__version__)), level="full", ) - elif (pl := get_polars()) is not None and isinstance(native_object, pl.LazyFrame): + elif is_polars_lazyframe(native_object): if series_only: msg = "Cannot only use `series_only` with polars.LazyFrame" raise TypeError(msg) if eager_only or eager_or_interchange_only: msg = "Cannot only use `eager_only` or `eager_or_interchange_only` with polars.LazyFrame" raise TypeError(msg) + pl = get_polars() return LazyFrame( PolarsLazyFrame(native_object, backend_version=parse_version(pl.__version__)), level="full", ) - elif (pl := get_polars()) is not None and isinstance(native_object, pl.Series): + elif is_polars_series(native_object): + pl = get_polars() if not allow_series: msg = "Please set `allow_series=True`" raise TypeError(msg) @@ -370,10 +387,11 @@ def from_native( # noqa: PLR0915 ) # pandas - elif (pd := get_pandas()) is not None and isinstance(native_object, pd.DataFrame): + elif is_pandas_dataframe(native_object): if series_only: msg = "Cannot only use `series_only` with dataframe" raise TypeError(msg) + pd = get_pandas() return DataFrame( PandasLikeDataFrame( native_object, @@ -382,10 +400,11 @@ def from_native( # noqa: PLR0915 ), level="full", ) - elif (pd := get_pandas()) is not None and isinstance(native_object, pd.Series): + elif is_pandas_series(native_object): if not allow_series: msg = "Please set `allow_series=True`" raise TypeError(msg) + pd = get_pandas() return Series( PandasLikeSeries( native_object, @@ -396,9 +415,8 @@ def from_native( # noqa: PLR0915 ) # Modin - elif (mpd := get_modin()) is not None and isinstance( - native_object, mpd.DataFrame - ): # pragma: no cover + elif is_modin_dataframe(native_object): # pragma: no cover + mpd = get_modin() if series_only: msg = "Cannot only use `series_only` with modin.DataFrame" raise TypeError(msg) @@ -410,9 +428,8 @@ def from_native( # noqa: PLR0915 ), level="full", ) - elif (mpd := get_modin()) is not None and isinstance( - native_object, mpd.Series - ): # pragma: no cover + elif is_modin_series(native_object): # pragma: no cover + mpd = get_modin() if not allow_series: msg = "Please set `allow_series=True`" raise TypeError(msg) @@ -426,9 +443,8 @@ def from_native( # noqa: PLR0915 ) # cuDF - elif (cudf := get_cudf()) is not None and isinstance( # pragma: no cover - native_object, cudf.DataFrame - ): + elif is_cudf_dataframe(native_object): # pragma: no cover + cudf = get_cudf() if series_only: msg = "Cannot only use `series_only` with cudf.DataFrame" raise TypeError(msg) @@ -440,9 +456,8 @@ def from_native( # noqa: PLR0915 ), level="full", ) - elif (cudf := get_cudf()) is not None and isinstance( - native_object, cudf.Series - ): # pragma: no cover + elif is_cudf_series(native_object): # pragma: no cover + cudf = get_cudf() if not allow_series: msg = "Please set `allow_series=True`" raise TypeError(msg) @@ -456,7 +471,8 @@ def from_native( # noqa: PLR0915 ) # PyArrow - elif (pa := get_pyarrow()) is not None and isinstance(native_object, pa.Table): + elif is_pyarrow_table(native_object): + pa = get_pyarrow() if series_only: msg = "Cannot only use `series_only` with arrow table" raise TypeError(msg) @@ -464,7 +480,8 @@ def from_native( # noqa: PLR0915 ArrowDataFrame(native_object, backend_version=parse_version(pa.__version__)), level="full", ) - elif (pa := get_pyarrow()) is not None and isinstance(native_object, pa.ChunkedArray): + elif is_pyarrow_chunked_array(native_object): + pa = get_pyarrow() if not allow_series: msg = "Please set `allow_series=True`" raise TypeError(msg) @@ -476,15 +493,11 @@ def from_native( # noqa: PLR0915 ) # Dask - elif (dd := get_dask_dataframe()) is not None and isinstance( - native_object, dd.DataFrame - ): - if series_only: # pragma: no cover - # TODO(unassigned): increase coverage + elif is_dask_dataframe(native_object): + if series_only: msg = "Cannot only use `series_only` with dask DataFrame" raise TypeError(msg) - if eager_only or eager_or_interchange_only: # pragma: no cover - # TODO(unassigned): increase coverage + if eager_only or eager_or_interchange_only: msg = "Cannot only use `eager_only` or `eager_or_interchange_only` with dask DataFrame" raise TypeError(msg) if get_dask_expr() is None: # pragma: no cover @@ -582,7 +595,7 @@ def func(df): You can also pass in extra arguments, e.g. ```python - @nw.narhwalify(eager_only=True) + @nw.narwhalify(eager_only=True) ``` that will get passed down to `nw.from_native`. diff --git a/narwhals/utils.py b/narwhals/utils.py index 2034c6feb..6c1b5c1b4 100644 --- a/narwhals/utils.py +++ b/narwhals/utils.py @@ -12,11 +12,21 @@ from typing import cast from narwhals import dtypes +from narwhals._exceptions import ColumnNotFoundError from narwhals.dependencies import get_cudf +from narwhals.dependencies import get_dask_dataframe from narwhals.dependencies import get_modin from narwhals.dependencies import get_pandas from narwhals.dependencies import get_polars from narwhals.dependencies import get_pyarrow +from narwhals.dependencies import is_cudf_series +from narwhals.dependencies import is_modin_series +from narwhals.dependencies import is_pandas_dataframe +from narwhals.dependencies import is_pandas_like_dataframe +from narwhals.dependencies import is_pandas_like_series +from narwhals.dependencies import is_pandas_series +from narwhals.dependencies import is_polars_series +from narwhals.dependencies import is_pyarrow_chunked_array from narwhals.translate import to_native if TYPE_CHECKING: @@ -36,6 +46,7 @@ class Implementation(Enum): CUDF = auto() PYARROW = auto() POLARS = auto() + DASK = auto() UNKNOWN = auto() @@ -50,6 +61,7 @@ def from_native_namespace( get_cudf(): Implementation.CUDF, get_pyarrow(): Implementation.PYARROW, get_polars(): Implementation.POLARS, + get_dask_dataframe(): Implementation.DASK, } return mapping.get(native_namespace, Implementation.UNKNOWN) @@ -94,7 +106,7 @@ def tupleify(arg: Any) -> Any: def _is_iterable(arg: Any | Iterable[Any]) -> bool: from narwhals.series import Series - if (pd := get_pandas()) is not None and isinstance(arg, (pd.Series, pd.DataFrame)): + if is_pandas_dataframe(arg) or is_pandas_series(arg): msg = f"Expected Narwhals class or scalar, got: {type(arg)}. Perhaps you forgot a `nw.from_native` somewhere?" raise TypeError(msg) if (pl := get_polars()) is not None and isinstance( @@ -177,23 +189,23 @@ def _validate_index(index: Any) -> None: if isinstance( getattr(lhs_any, "_compliant_frame", None), PandasLikeDataFrame ) and isinstance(getattr(rhs_any, "_compliant_frame", None), PandasLikeDataFrame): - _validate_index(lhs_any._compliant_frame._native_dataframe.index) - _validate_index(rhs_any._compliant_frame._native_dataframe.index) + _validate_index(lhs_any._compliant_frame._native_frame.index) + _validate_index(rhs_any._compliant_frame._native_frame.index) return lhs_any._from_compliant_dataframe( # type: ignore[no-any-return] - lhs_any._compliant_frame._from_native_dataframe( - lhs_any._compliant_frame._native_dataframe.loc[ - rhs_any._compliant_frame._native_dataframe.index + lhs_any._compliant_frame._from_native_frame( + lhs_any._compliant_frame._native_frame.loc[ + rhs_any._compliant_frame._native_frame.index ] ) ) if isinstance( getattr(lhs_any, "_compliant_frame", None), PandasLikeDataFrame ) and isinstance(getattr(rhs_any, "_compliant_series", None), PandasLikeSeries): - _validate_index(lhs_any._compliant_frame._native_dataframe.index) + _validate_index(lhs_any._compliant_frame._native_frame.index) _validate_index(rhs_any._compliant_series._native_series.index) return lhs_any._from_compliant_dataframe( # type: ignore[no-any-return] - lhs_any._compliant_frame._from_native_dataframe( - lhs_any._compliant_frame._native_dataframe.loc[ + lhs_any._compliant_frame._from_native_frame( + lhs_any._compliant_frame._native_frame.loc[ rhs_any._compliant_series._native_series.index ] ) @@ -202,11 +214,11 @@ def _validate_index(index: Any) -> None: getattr(lhs_any, "_compliant_series", None), PandasLikeSeries ) and isinstance(getattr(rhs_any, "_compliant_frame", None), PandasLikeDataFrame): _validate_index(lhs_any._compliant_series._native_series.index) - _validate_index(rhs_any._compliant_frame._native_dataframe.index) + _validate_index(rhs_any._compliant_frame._native_frame.index) return lhs_any._from_compliant_series( # type: ignore[no-any-return] lhs_any._compliant_series._from_native_series( lhs_any._compliant_series._native_series.loc[ - rhs_any._compliant_frame._native_dataframe.index + rhs_any._compliant_frame._native_frame.index ] ) ) @@ -228,6 +240,37 @@ def _validate_index(index: Any) -> None: return lhs +def maybe_get_index(obj: T) -> Any | None: + """ + Get the index of a DataFrame or a Series, if it's pandas-like. + + Notes: + This is only really intended for backwards-compatibility purposes, + for example if your library already aligns indices for users. + If you're designing a new library, we highly encourage you to not + rely on the Index. + For non-pandas-like inputs, this returns `None`. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}) + >>> df = nw.from_native(df_pd) + >>> nw.maybe_get_index(df) + RangeIndex(start=0, stop=2, step=1) + >>> series_pd = pd.Series([1, 2]) + >>> series = nw.from_native(series_pd, series_only=True) + >>> nw.maybe_get_index(series) + RangeIndex(start=0, stop=2, step=1) + """ + obj_any = cast(Any, obj) + native_obj = to_native(obj_any) + if is_pandas_like_dataframe(native_obj) or is_pandas_like_series(native_obj): + return native_obj.index + return None + + def maybe_set_index(df: T, column_names: str | list[str]) -> T: """ Set columns `columns` to be the index of `df`, if `df` is pandas-like. @@ -251,21 +294,25 @@ def maybe_set_index(df: T, column_names: str | list[str]) -> T: 4 1 5 2 """ - from narwhals._pandas_like.dataframe import PandasLikeDataFrame - df_any = cast(Any, df) - if isinstance(getattr(df_any, "_compliant_frame", None), PandasLikeDataFrame): + native_frame = to_native(df_any) + if is_pandas_like_dataframe(native_frame): return df_any._from_compliant_dataframe( # type: ignore[no-any-return] - df_any._compliant_frame._from_native_dataframe( - df_any._compliant_frame._native_dataframe.set_index(column_names) + df_any._compliant_frame._from_native_frame( + native_frame.set_index(column_names) ) ) - return df + return df_any # type: ignore[no-any-return] -def maybe_convert_dtypes(df: T, *args: bool, **kwargs: bool | str) -> T: +def maybe_convert_dtypes(obj: T, *args: bool, **kwargs: bool | str) -> T: """ - Convert columns to the best possible dtypes using dtypes supporting ``pd.NA``, if df is pandas-like. + Convert columns or series to the best possible dtypes using dtypes supporting ``pd.NA``, if df is pandas-like. + + Arguments: + obj: DataFrame or Series. + *args: Additional arguments which gets passed through. + **kwargs: Additional arguments which gets passed through. Notes: For non-pandas-like inputs, this is a no-op. @@ -288,16 +335,21 @@ def maybe_convert_dtypes(df: T, *args: bool, **kwargs: bool | str) -> T: b boolean dtype: object """ - from narwhals._pandas_like.dataframe import PandasLikeDataFrame - - df_any = cast(Any, df) - if isinstance(getattr(df_any, "_compliant_frame", None), PandasLikeDataFrame): - return df_any._from_compliant_dataframe( # type: ignore[no-any-return] - df_any._compliant_frame._from_native_dataframe( - df_any._compliant_frame._native_dataframe.convert_dtypes(*args, **kwargs) + obj_any = cast(Any, obj) + native_obj = to_native(obj_any) + if is_pandas_like_dataframe(native_obj): + return obj_any._from_compliant_dataframe( # type: ignore[no-any-return] + obj_any._compliant_frame._from_native_frame( + native_obj.convert_dtypes(*args, **kwargs) ) ) - return df + if is_pandas_like_series(native_obj): + return obj_any._from_compliant_series( # type: ignore[no-any-return] + obj_any._compliant_series._from_native_series( + native_obj.convert_dtypes(*args, **kwargs) + ) + ) + return obj_any # type: ignore[no-any-return] def is_ordered_categorical(series: Series) -> bool: @@ -351,19 +403,15 @@ def is_ordered_categorical(series: Series) -> bool: if series.dtype != dtypes.Categorical: return False native_series = to_native(series) - if (pl := get_polars()) is not None and isinstance(native_series, pl.Series): - return native_series.dtype.ordering == "physical" # type: ignore[no-any-return] - if (pd := get_pandas()) is not None and isinstance(native_series, pd.Series): + if is_polars_series(native_series): + return native_series.dtype.ordering == "physical" # type: ignore[attr-defined, no-any-return] + if is_pandas_series(native_series): return native_series.cat.ordered # type: ignore[no-any-return] - if (mpd := get_modin()) is not None and isinstance( - native_series, mpd.Series - ): # pragma: no cover + if is_modin_series(native_series): # pragma: no cover return native_series.cat.ordered # type: ignore[no-any-return] - if (cudf := get_cudf()) is not None and isinstance( - native_series, cudf.Series - ): # pragma: no cover + if is_cudf_series(native_series): # pragma: no cover return native_series.cat.ordered # type: ignore[no-any-return] - if (pa := get_pyarrow()) is not None and isinstance(native_series, pa.ChunkedArray): + if is_pyarrow_chunked_array(native_series): return native_series.type.ordered # type: ignore[no-any-return] # If it doesn't match any of the above, let's just play it safe and return False. return False # pragma: no cover @@ -395,3 +443,21 @@ def generate_unique_token(n_bytes: int, columns: list[str]) -> str: # pragma: n "join operation" ) raise AssertionError(msg) + + +def parse_columns_to_drop( + compliant_frame: Any, + columns: Iterable[str], + strict: bool, # noqa: FBT001 +) -> list[str]: + cols = set(compliant_frame.columns) + to_drop = list(columns) + + if strict: + for d in to_drop: + if d not in cols: + msg = f'"{d}" not found' + raise ColumnNotFoundError(msg) + else: + to_drop = list(cols.intersection(set(to_drop))) + return to_drop diff --git a/noxfile.py b/noxfile.py index 9e0508b7e..06cdc0284 100644 --- a/noxfile.py +++ b/noxfile.py @@ -23,17 +23,20 @@ def run_common(session: Session, coverage_threshold: float) -> None: @nox.session(python=PYTHON_VERSIONS) # type: ignore[misc] def pytest_coverage(session: Session) -> None: - coverage_threshold = 90 if session.python == "3.8" else 100 - - session.install("modin[dask]") + if session.python == "3.8": + coverage_threshold = 85 + else: + coverage_threshold = 100 + session.install("modin[dask]") run_common(session, coverage_threshold) @nox.session(python=PYTHON_VERSIONS[0]) # type: ignore[misc] -def minimum_versions(session: Session) -> None: +@nox.parametrize("pandas_version", ["0.25.3", "1.1.5"]) # type: ignore[misc] +def min_and_old_versions(session: Session, pandas_version: str) -> None: session.install( - "pandas==0.25.3", + f"pandas=={pandas_version}", "polars==0.20.3", "numpy==1.17.5", "pyarrow==11.0.0", @@ -46,12 +49,29 @@ def minimum_versions(session: Session) -> None: @nox.session(python=PYTHON_VERSIONS[-1]) # type: ignore[misc] def nightly_versions(session: Session) -> None: - session.install("modin[dask]", "polars") - session.install( + session.install("polars") + + session.install( # pandas nightly "--pre", "--extra-index-url", "https://pypi.anaconda.org/scientific-python-nightly-wheels/simple", "pandas", ) + session.install( # numpy nightly + "--pre", + "--extra-index-url", + "https://pypi.anaconda.org/scientific-python-nightly-wheels/simple", + "numpy", + ) + + session.run("uv", "pip", "install", "pip") + session.run( # dask nightly + "pip", + "install", + "git+https://github.com/dask/distributed", + "git+https://github.com/dask/dask", + "git+https://github.com/dask/dask-expr", + ) + run_common(session, coverage_threshold=50) diff --git a/pyproject.toml b/pyproject.toml index ddfbb4093..2aa6c8005 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "narwhals" -version = "1.1.8" +version = "1.5.5" authors = [ { name="Marco Gorelli", email="33491632+MarcoGorelli@users.noreply.github.com" }, ] @@ -26,9 +26,12 @@ exclude = [ ] [project.optional-dependencies] +cudf = ["cudf>=23.08.00"] +modin = ["modin"] pandas = ["pandas>=0.25.3"] polars = ["polars>=0.20.3"] -pyarrow = ['pyarrow>=11.0.0'] +pyarrow = ["pyarrow>=11.0.0"] +dask = ["dask[dataframe]>=2024.7"] [project.urls] "Homepage" = "https://github.com/narwhals-dev/narwhals" @@ -54,6 +57,7 @@ lint.ignore = [ "E501", "FIX", "ISC001", + "NPY002", "PD901", # This is a auxiliary library so dataframe variables have no concrete business meaning "PLR0911", "PLR0912", @@ -91,9 +95,14 @@ filterwarnings = [ 'ignore:.*implementation has mismatches with pandas', 'ignore:.*Do not use the `random` module inside strategies', 'ignore:.*You are using pyarrow version', + 'ignore:.*but when imported by', + 'ignore:Distributing .*This may take some time', ] xfail_strict = true markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"] +env = [ + "MODIN_ENGINE=python", +] [tool.coverage.run] plugins = ["covdefaults"] diff --git a/requirements-dev.txt b/requirements-dev.txt index 98bf9d8c7..2424d4ea1 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,10 +1,11 @@ covdefaults pandas -polars[timezones] +polars pre-commit pyarrow pytest pytest-cov +pytest-env hypothesis scikit-learn dask[dataframe]; python_version >= '3.9' diff --git a/tests/conftest.py b/tests/conftest.py index 93876de5f..cdf4e0be6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,6 +7,7 @@ import pyarrow as pa import pytest +from narwhals.dependencies import get_cudf from narwhals.dependencies import get_dask_dataframe from narwhals.dependencies import get_modin from narwhals.typing import IntoDataFrame @@ -17,6 +18,8 @@ import modin.pandas # noqa: F401 with contextlib.suppress(ImportError): import dask.dataframe # noqa: F401 +with contextlib.suppress(ImportError): + import cudf # noqa: F401 def pytest_addoption(parser: Any) -> None: @@ -56,6 +59,11 @@ def modin_constructor(obj: Any) -> IntoDataFrame: # pragma: no cover return mpd.DataFrame(pd.DataFrame(obj)).convert_dtypes(dtype_backend="pyarrow") # type: ignore[no-any-return] +def cudf_constructor(obj: Any) -> IntoDataFrame: # pragma: no cover + cudf = get_cudf() + return cudf.DataFrame(obj) # type: ignore[no-any-return] + + def polars_eager_constructor(obj: Any) -> IntoDataFrame: return pl.DataFrame(obj) @@ -87,9 +95,10 @@ def pyarrow_table_constructor(obj: Any) -> IntoDataFrame: if get_modin() is not None: # pragma: no cover eager_constructors.append(modin_constructor) -# TODO(unassigned): when Dask gets better support, remove the "False and" part -if False and get_dask_dataframe() is not None: # pragma: no cover # noqa: SIM223 - lazy_constructors.append(dask_lazy_constructor) +if get_cudf() is not None: + eager_constructors.append(cudf_constructor) # pragma: no cover +if get_dask_dataframe() is not None: # pragma: no cover + lazy_constructors.append(dask_lazy_constructor) # type: ignore # noqa: PGH003 @pytest.fixture(params=eager_constructors) diff --git a/tests/dask_test.py b/tests/dask_test.py deleted file mode 100644 index 110c88adb..000000000 --- a/tests/dask_test.py +++ /dev/null @@ -1,429 +0,0 @@ -""" -Dask support in Narwhals is still _very_ scant. - -Start with a simple test file whilst we develop the basics. -Once we're a bit further along (say, we can at least evaluate -TPC-H Q1 with Dask), then we can integrate dask tests into -the main test suite. -""" - -from __future__ import annotations - -import sys -import warnings -from datetime import datetime -from typing import Any - -import pandas as pd -import pytest - -import narwhals.stable.v1 as nw -from narwhals.utils import parse_version -from tests.utils import compare_dicts - -pytest.importorskip("dask") -with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - category=pytest.PytestDeprecationWarning, - ) - pytest.importorskip("dask_expr") - - -if sys.version_info < (3, 9): - pytest.skip("Dask tests require Python 3.9+", allow_module_level=True) - - -def test_with_columns() -> None: - import dask.dataframe as dd - - dfdd = dd.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})) - - df = nw.from_native(dfdd) - result = df.with_columns( - nw.col("a") + 1, - (nw.col("a") + nw.col("b").mean()).alias("c"), - d=nw.col("a"), - e=nw.col("a") + nw.col("b"), - f=nw.col("b") - 1, - g=nw.col("a") - nw.col("b"), - h=nw.col("a") * 3, - i=nw.col("a") * nw.col("b"), - ) - compare_dicts( - result, - { - "a": [2, 3, 4], - "b": [4, 5, 6], - "c": [6.0, 7.0, 8.0], - "d": [1, 2, 3], - "e": [5, 7, 9], - "f": [3, 4, 5], - "g": [-3, -3, -3], - "h": [3, 6, 9], - "i": [4, 10, 18], - }, - ) - - -def test_shift() -> None: - import dask.dataframe as dd - - dfdd = dd.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})) - df = nw.from_native(dfdd) - result = df.with_columns(nw.col("a").shift(1), nw.col("b").shift(-1)) - expected = {"a": [float("nan"), 1, 2], "b": [5, 6, float("nan")]} - compare_dicts(result, expected) - - -def test_cum_sum() -> None: - import dask.dataframe as dd - - dfdd = dd.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})) - df = nw.from_native(dfdd) - result = df.with_columns(nw.col("a", "b").cum_sum()) - expected = {"a": [1, 3, 6], "b": [4, 9, 15]} - compare_dicts(result, expected) - - -def test_sum() -> None: - import dask.dataframe as dd - - dfdd = dd.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})) - df = nw.from_native(dfdd) - result = df.with_columns((nw.col("a") + nw.col("b").sum()).alias("c")) - expected = {"a": [1, 2, 3], "b": [4, 5, 6], "c": [16, 17, 18]} - compare_dicts(result, expected) - - -@pytest.mark.parametrize( - ("closed", "expected"), - [ - ("left", [True, True, True, False]), - ("right", [False, True, True, True]), - ("both", [True, True, True, True]), - ("neither", [False, True, True, False]), - ], -) -def test_is_between(closed: str, expected: list[bool]) -> None: - import dask.dataframe as dd - - data = { - "a": [1, 4, 2, 5], - } - dfdd = dd.from_pandas(pd.DataFrame(data)) - - df = nw.from_native(dfdd) - result = df.with_columns(nw.col("a").is_between(1, 5, closed=closed)) - expected_dict = {"a": expected} - compare_dicts(result, expected_dict) - - -@pytest.mark.parametrize( - ("prefix", "expected"), - [ - ("fda", {"a": [True, False]}), - ("edf", {"a": [False, True]}), - ("asd", {"a": [False, False]}), - ], -) -def test_starts_with(prefix: str, expected: dict[str, list[bool]]) -> None: - import dask.dataframe as dd - - data = {"a": ["fdas", "edfas"]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(nw.col("a").str.starts_with(prefix)) - - compare_dicts(result, expected) - - -@pytest.mark.parametrize( - ("suffix", "expected"), - [ - ("das", {"a": [True, False]}), - ("fas", {"a": [False, True]}), - ("asd", {"a": [False, False]}), - ], -) -def test_ends_with(suffix: str, expected: dict[str, list[bool]]) -> None: - import dask.dataframe as dd - - data = {"a": ["fdas", "edfas"]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(nw.col("a").str.ends_with(suffix)) - - compare_dicts(result, expected) - - -def test_contains() -> None: - import dask.dataframe as dd - - data = {"pets": ["cat", "dog", "rabbit and parrot", "dove"]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - - result = df.with_columns( - case_insensitive_match=nw.col("pets").str.contains("(?i)parrot|Dove") - ) - expected = { - "pets": ["cat", "dog", "rabbit and parrot", "dove"], - "case_insensitive_match": [False, False, True, True], - } - compare_dicts(result, expected) - - -@pytest.mark.parametrize( - ("offset", "length", "expected"), - [(1, 2, {"a": ["da", "df"]}), (-2, None, {"a": ["as", "as"]})], -) -def test_str_slice(offset: int, length: int | None, expected: Any) -> None: - import dask.dataframe as dd - - data = {"a": ["fdas", "edfas"]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - - result_frame = df.with_columns(nw.col("a").str.slice(offset, length)) - compare_dicts(result_frame, expected) - - -def test_to_datetime() -> None: - import dask.dataframe as dd - - data = {"a": ["2020-01-01T12:34:56"]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - - format = "%Y-%m-%dT%H:%M:%S" - result = df.with_columns(b=nw.col("a").str.to_datetime(format=format)) - - expected = { - "a": ["2020-01-01T12:34:56"], - "b": [datetime.strptime("2020-01-01T12:34:56", format)], # noqa: DTZ007 - } - compare_dicts(result, expected) - - -@pytest.mark.parametrize( - ("data", "expected"), - [ - ({"a": ["foo", "bar"]}, {"a": ["FOO", "BAR"]}), - ( - { - "a": [ - "special case ß", - "ςpecial caσe", # noqa: RUF001 - ] - }, - {"a": ["SPECIAL CASE ẞ", "ΣPECIAL CAΣE"]}, - ), - ], -) -def test_str_to_uppercase( - request: pytest.FixtureRequest, - data: dict[str, list[str]], - expected: dict[str, list[str]], -) -> None: - import dask.dataframe as dd - import pyarrow as pa - - if (parse_version(pa.__version__) < (12, 0, 0)) and ("ß" in data["a"][0]): - # We are marking it xfail for these conditions above - # since the pyarrow backend will convert - # smaller cap 'ß' to upper cap 'ẞ' instead of 'SS' - request.applymarker(pytest.mark.xfail) - - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - - result_frame = df.with_columns(nw.col("a").str.to_uppercase()) - - compare_dicts(result_frame, expected) - - -@pytest.mark.parametrize( - ("data", "expected"), - [ - ({"a": ["FOO", "BAR"]}, {"a": ["foo", "bar"]}), - ( - {"a": ["SPECIAL CASE ß", "ΣPECIAL CAΣE"]}, - { - "a": [ - "special case ß", - "σpecial caσe", # noqa: RUF001 - ] - }, - ), - ], -) -def test_str_to_lowercase( - data: dict[str, list[str]], - expected: dict[str, list[str]], -) -> None: - import dask.dataframe as dd - - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - - result_frame = df.with_columns(nw.col("a").str.to_lowercase()) - compare_dicts(result_frame, expected) - - -def test_columns() -> None: - import dask.dataframe as dd - - dfdd = dd.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": ["cat", "bat", "mat"]})) - df = nw.from_native(dfdd) - - result = df.columns - assert result == ["a", "b"] - - -def test_select() -> None: - import dask.dataframe as dd - - data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} - df = nw.from_native(dd.from_pandas(pd.DataFrame(data))) - result = df.select("a", nw.col("b") + 1, (nw.col("z") * 2).alias("z*2")) - expected = {"a": [1, 3, 2], "b": [5, 5, 7], "z*2": [14.0, 16.0, 18.0]} - compare_dicts(result, expected) - - -def test_str_only_select() -> None: - import dask.dataframe as dd - - data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} - df = nw.from_native(dd.from_pandas(pd.DataFrame(data))) - result = df.select("a", "b") - expected = {"a": [1, 3, 2], "b": [4, 4, 6]} - compare_dicts(result, expected) - - -def test_empty_select() -> None: - import dask.dataframe as dd - - result = ( - nw.from_native(dd.from_pandas(pd.DataFrame({"a": [1, 2, 3]}))) - .lazy() - .select() - .collect() - ) - assert result.shape == (0, 0) - - -def test_dt_year() -> None: - import dask.dataframe as dd - - data = {"a": [datetime(2020, 1, 1), datetime(2021, 1, 1)]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(year=nw.col("a").dt.year()) - expected = {"a": data["a"], "year": [2020, 2021]} - compare_dicts(result, expected) - - -def test_dt_month() -> None: - import dask.dataframe as dd - - data = {"a": [datetime(2020, 1, 1), datetime(2021, 1, 1)]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(month=nw.col("a").dt.month()) - expected = {"a": data["a"], "month": [1, 1]} - compare_dicts(result, expected) - - -def test_dt_day() -> None: - import dask.dataframe as dd - - data = {"a": [datetime(2020, 1, 1), datetime(2021, 1, 1)]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(day=nw.col("a").dt.day()) - expected = {"a": data["a"], "day": [1, 1]} - compare_dicts(result, expected) - - -def test_dt_hour() -> None: - import dask.dataframe as dd - - data = {"a": [datetime(2020, 1, 1, 1), datetime(2021, 1, 1, 2)]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(hour=nw.col("a").dt.hour()) - expected = {"a": data["a"], "hour": [1, 2]} - compare_dicts(result, expected) - - -def test_dt_minute() -> None: - import dask.dataframe as dd - - data = {"a": [datetime(2020, 1, 1, 1, 1), datetime(2021, 1, 1, 2, 2)]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(minute=nw.col("a").dt.minute()) - expected = {"a": data["a"], "minute": [1, 2]} - compare_dicts(result, expected) - - -def test_dt_second() -> None: - import dask.dataframe as dd - - data = {"a": [datetime(2020, 1, 1, 1, 1, 1), datetime(2021, 1, 1, 2, 2, 2)]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(second=nw.col("a").dt.second()) - expected = {"a": data["a"], "second": [1, 2]} - compare_dicts(result, expected) - - -def test_dt_millisecond() -> None: - import dask.dataframe as dd - - data = { - "a": [datetime(2020, 1, 1, 1, 1, 1, 1000), datetime(2021, 1, 1, 2, 2, 2, 2000)] - } - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(millisecond=nw.col("a").dt.millisecond()) - expected = {"a": data["a"], "millisecond": [1, 2]} - compare_dicts(result, expected) - - -def test_dt_microsecond() -> None: - import dask.dataframe as dd - - data = { - "a": [datetime(2020, 1, 1, 1, 1, 1, 1000), datetime(2021, 1, 1, 2, 2, 2, 2000)] - } - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(microsecond=nw.col("a").dt.microsecond()) - expected = {"a": data["a"], "microsecond": [1000, 2000]} - compare_dicts(result, expected) - - -def test_dt_nanosecond() -> None: - import dask.dataframe as dd - - data = { - "a": [datetime(2020, 1, 1, 1, 1, 1, 1000), datetime(2021, 1, 1, 2, 2, 2, 2000)] - } - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(nanosecond=nw.col("a").dt.nanosecond()) - expected = {"a": data["a"], "nanosecond": [1000000, 2000000]} - compare_dicts(result, expected) - - -def test_dt_ordinal_day() -> None: - import dask.dataframe as dd - - data = {"a": [datetime(2020, 1, 7), datetime(2021, 2, 1)]} - dfdd = dd.from_pandas(pd.DataFrame(data)) - df = nw.from_native(dfdd) - result = df.with_columns(ordinal_day=nw.col("a").dt.ordinal_day()) - expected = {"a": data["a"], "ordinal_day": [7, 32]} - compare_dicts(result, expected) diff --git a/tests/expr_and_series/all_horizontal_test.py b/tests/expr_and_series/all_horizontal_test.py index eab76fe1b..256c45deb 100644 --- a/tests/expr_and_series/all_horizontal_test.py +++ b/tests/expr_and_series/all_horizontal_test.py @@ -1,16 +1,20 @@ from typing import Any +import pytest + import narwhals.stable.v1 as nw from tests.utils import compare_dicts -def test_allh(constructor: Any) -> None: +@pytest.mark.parametrize("expr1", ["a", nw.col("a")]) +@pytest.mark.parametrize("expr2", ["b", nw.col("b")]) +def test_allh(constructor: Any, expr1: Any, expr2: Any) -> None: data = { "a": [False, False, True], "b": [False, True, True], } df = nw.from_native(constructor(data)) - result = df.select(all=nw.all_horizontal(nw.col("a"), nw.col("b"))) + result = df.select(all=nw.all_horizontal(expr1, expr2)) expected = {"all": [False, False, True]} compare_dicts(result, expected) diff --git a/tests/expr_and_series/any_all_test.py b/tests/expr_and_series/any_all_test.py index cfe71a135..09cc8c9e3 100644 --- a/tests/expr_and_series/any_all_test.py +++ b/tests/expr_and_series/any_all_test.py @@ -14,7 +14,7 @@ def test_any_all(constructor: Any) -> None: } ) ) - result = df.select(nw.all().all()) + result = df.select(nw.col("a", "b", "c").all()) expected = {"a": [False], "b": [True], "c": [False]} compare_dicts(result, expected) result = df.select(nw.all().any()) diff --git a/tests/expr_and_series/any_horizontal_test.py b/tests/expr_and_series/any_horizontal_test.py index f845f6559..1f19aa304 100644 --- a/tests/expr_and_series/any_horizontal_test.py +++ b/tests/expr_and_series/any_horizontal_test.py @@ -1,16 +1,20 @@ from typing import Any +import pytest + import narwhals.stable.v1 as nw from tests.utils import compare_dicts -def test_anyh(constructor: Any) -> None: +@pytest.mark.parametrize("expr1", ["a", nw.col("a")]) +@pytest.mark.parametrize("expr2", ["b", nw.col("b")]) +def test_anyh(constructor: Any, expr1: Any, expr2: Any) -> None: data = { "a": [False, False, True], "b": [False, True, True], } df = nw.from_native(constructor(data)) - result = df.select(any=nw.any_horizontal(nw.col("a"), nw.col("b"))) + result = df.select(any=nw.any_horizontal(expr1, expr2)) expected = {"any": [False, True, True]} compare_dicts(result, expected) diff --git a/tests/expr_and_series/arg_true_test.py b/tests/expr_and_series/arg_true_test.py index 373ba41d8..eaa3d1ba6 100644 --- a/tests/expr_and_series/arg_true_test.py +++ b/tests/expr_and_series/arg_true_test.py @@ -1,10 +1,14 @@ from typing import Any -import narwhals as nw +import pytest + +import narwhals.stable.v1 as nw from tests.utils import compare_dicts -def test_arg_true(constructor: Any) -> None: +def test_arg_true(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor({"a": [1, None, None, 3]})) result = df.select(nw.col("a").is_null().arg_true()) expected = {"a": [1, 2]} diff --git a/tests/expr_and_series/arithmetic_test.py b/tests/expr_and_series/arithmetic_test.py index 599f4e28c..47d3e8ff0 100644 --- a/tests/expr_and_series/arithmetic_test.py +++ b/tests/expr_and_series/arithmetic_test.py @@ -2,9 +2,16 @@ from typing import Any +import hypothesis.strategies as st +import pandas as pd +import polars as pl +import pyarrow as pa import pytest +from hypothesis import assume +from hypothesis import given import narwhals.stable.v1 as nw +from narwhals.utils import parse_version from tests.utils import compare_dicts @@ -21,11 +28,15 @@ ("__pow__", 2, [1, 4, 9]), ], ) -def test_arithmetic( - attr: str, rhs: Any, expected: list[Any], constructor: Any, request: Any +def test_arithmetic_expr( + attr: str, + rhs: Any, + expected: list[Any], + constructor: Any, + request: Any, ) -> None: if attr == "__mod__" and any( - x in str(constructor) for x in ["pandas_pyarrow", "pyarrow_table", "modin"] + x in str(constructor) for x in ["pandas_pyarrow", "modin"] ): request.applymarker(pytest.mark.xfail) @@ -47,11 +58,15 @@ def test_arithmetic( ("__rpow__", 2, [2, 4, 8]), ], ) -def test_right_arithmetic( - attr: str, rhs: Any, expected: list[Any], constructor: Any, request: Any +def test_right_arithmetic_expr( + attr: str, + rhs: Any, + expected: list[Any], + constructor: Any, + request: Any, ) -> None: if attr == "__rmod__" and any( - x in str(constructor) for x in ["pandas_pyarrow", "pyarrow_table", "modin"] + x in str(constructor) for x in ["pandas_pyarrow", "modin"] ): request.applymarker(pytest.mark.xfail) @@ -75,10 +90,14 @@ def test_right_arithmetic( ], ) def test_arithmetic_series( - attr: str, rhs: Any, expected: list[Any], constructor_eager: Any, request: Any + attr: str, + rhs: Any, + expected: list[Any], + constructor_eager: Any, + request: Any, ) -> None: if attr == "__mod__" and any( - x in str(constructor_eager) for x in ["pandas_pyarrow", "pyarrow_table", "modin"] + x in str(constructor_eager) for x in ["pandas_pyarrow", "modin"] ): request.applymarker(pytest.mark.xfail) @@ -101,10 +120,14 @@ def test_arithmetic_series( ], ) def test_right_arithmetic_series( - attr: str, rhs: Any, expected: list[Any], constructor_eager: Any, request: Any + attr: str, + rhs: Any, + expected: list[Any], + constructor_eager: Any, + request: Any, ) -> None: if attr == "__rmod__" and any( - x in str(constructor_eager) for x in ["pandas_pyarrow", "pyarrow_table", "modin"] + x in str(constructor_eager) for x in ["pandas_pyarrow", "modin"] ): request.applymarker(pytest.mark.xfail) @@ -112,3 +135,86 @@ def test_right_arithmetic_series( df = nw.from_native(constructor_eager(data), eager_only=True) result = df.select(a=getattr(df["a"], attr)(rhs)) compare_dicts(result, {"a": expected}) + + +def test_truediv_same_dims(constructor_eager: Any, request: Any) -> None: + if "polars" in str(constructor_eager): + # https://github.com/pola-rs/polars/issues/17760 + request.applymarker(pytest.mark.xfail) + s_left = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"] + s_right = nw.from_native(constructor_eager({"a": [2, 2, 1]}), eager_only=True)["a"] + result = s_left / s_right + compare_dicts({"a": result}, {"a": [0.5, 1.0, 3.0]}) + result = s_left.__rtruediv__(s_right) + compare_dicts({"a": result}, {"a": [2, 1, 1 / 3]}) + + +@pytest.mark.slow() +@given( # type: ignore[misc] + left=st.integers(-100, 100), + right=st.integers(-100, 100), +) +@pytest.mark.skipif( + parse_version(pd.__version__) < (2, 0), reason="convert_dtypes not available" +) +def test_floordiv(left: int, right: int) -> None: + # hypothesis complains if we add `constructor` as an argument, so this + # test is a bit manual unfortunately + assume(right != 0) + expected = {"a": [left // right]} + result = nw.from_native(pd.DataFrame({"a": [left]}), eager_only=True).select( + nw.col("a") // right + ) + compare_dicts(result, expected) + if parse_version(pd.__version__) < (2, 2): # pragma: no cover + # Bug in old version of pandas + pass + else: + result = nw.from_native( + pd.DataFrame({"a": [left]}).convert_dtypes(dtype_backend="pyarrow"), + eager_only=True, + ).select(nw.col("a") // right) + compare_dicts(result, expected) + result = nw.from_native( + pd.DataFrame({"a": [left]}).convert_dtypes(), eager_only=True + ).select(nw.col("a") // right) + compare_dicts(result, expected) + result = nw.from_native(pl.DataFrame({"a": [left]}), eager_only=True).select( + nw.col("a") // right + ) + compare_dicts(result, expected) + result = nw.from_native(pa.table({"a": [left]}), eager_only=True).select( + nw.col("a") // right + ) + compare_dicts(result, expected) + + +@pytest.mark.slow() +@given( # type: ignore[misc] + left=st.integers(-100, 100), + right=st.integers(-100, 100), +) +@pytest.mark.skipif( + parse_version(pd.__version__) < (2, 0), reason="convert_dtypes not available" +) +def test_mod(left: int, right: int) -> None: + # hypothesis complains if we add `constructor` as an argument, so this + # test is a bit manual unfortunately + assume(right != 0) + expected = {"a": [left % right]} + result = nw.from_native(pd.DataFrame({"a": [left]}), eager_only=True).select( + nw.col("a") % right + ) + compare_dicts(result, expected) + result = nw.from_native( + pd.DataFrame({"a": [left]}).convert_dtypes(), eager_only=True + ).select(nw.col("a") % right) + compare_dicts(result, expected) + result = nw.from_native(pl.DataFrame({"a": [left]}), eager_only=True).select( + nw.col("a") % right + ) + compare_dicts(result, expected) + result = nw.from_native(pa.table({"a": [left]}), eager_only=True).select( + nw.col("a") % right + ) + compare_dicts(result, expected) diff --git a/tests/expr_and_series/binary_test.py b/tests/expr_and_series/binary_test.py new file mode 100644 index 000000000..2d55af228 --- /dev/null +++ b/tests/expr_and_series/binary_test.py @@ -0,0 +1,45 @@ +from typing import Any + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + + +def test_expr_binary(constructor: Any) -> None: + data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} + df_raw = constructor(data) + result = nw.from_native(df_raw).with_columns( + a=(1 + 3 * nw.col("a")) * (1 / nw.col("a")), + b=nw.col("z") / (2 - nw.col("b")), + c=nw.col("a") + nw.col("b") / 2, + d=nw.col("a") - nw.col("b"), + e=((nw.col("a") > nw.col("b")) & (nw.col("a") >= nw.col("z"))).cast(nw.Int64), + f=( + (nw.col("a") < nw.col("b")) + | (nw.col("a") <= nw.col("z")) + | (nw.col("a") == 1) + ).cast(nw.Int64), + g=nw.col("a") != 1, + h=(False & (nw.col("a") != 1)), + i=(False | (nw.col("a") != 1)), + j=2 ** nw.col("a"), + k=2 // nw.col("a"), + l=nw.col("a") // 2, + m=nw.col("a") ** 2, + ) + expected = { + "a": [4, 3.333333, 3.5], + "b": [-3.5, -4.0, -2.25], + "z": [7.0, 8.0, 9.0], + "c": [3, 5, 5], + "d": [-3, -1, -4], + "e": [0, 0, 0], + "f": [1, 1, 1], + "g": [False, True, True], + "h": [False, False, False], + "i": [False, True, True], + "j": [2, 8, 4], + "k": [2, 0, 1], + "l": [0, 1, 1], + "m": [1, 9, 4], + } + compare_dicts(result, expected) diff --git a/tests/expr_and_series/cast_test.py b/tests/expr_and_series/cast_test.py index eaf31f7be..0b496d7ae 100644 --- a/tests/expr_and_series/cast_test.py +++ b/tests/expr_and_series/cast_test.py @@ -1,5 +1,6 @@ from typing import Any +import pandas as pd import pyarrow as pa import pytest @@ -95,17 +96,21 @@ def test_cast(constructor: Any, request: Any) -> None: assert dict(result.collect_schema()) == expected -def test_cast_series(constructor_eager: Any, request: Any) -> None: - if "pyarrow_table_constructor" in str(constructor_eager) and parse_version( +def test_cast_series(constructor: Any, request: Any) -> None: + if "pyarrow_table_constructor" in str(constructor) and parse_version( pa.__version__ ) <= (15,): # pragma: no cover request.applymarker(pytest.mark.xfail) - if "modin" in str(constructor_eager): + if "modin" in str(constructor): # TODO(unassigned): in modin, we end up with `' None: df["p"].cast(nw.Duration), ) assert result.schema == expected + + +@pytest.mark.skipif( + parse_version(pd.__version__) < parse_version("1.0.0"), + reason="too old for convert_dtypes", +) +def test_cast_string() -> None: + s_pd = pd.Series([1, 2]).convert_dtypes() + s = nw.from_native(s_pd, series_only=True) + s = s.cast(nw.String) + result = nw.to_native(s) + assert str(result.dtype) in ("string", "object", "dtype('O')") + + +def test_cast_raises_for_unknown_dtype(constructor: Any, request: Any) -> None: + if "pyarrow_table_constructor" in str(constructor) and parse_version( + pa.__version__ + ) <= (15,): # pragma: no cover + request.applymarker(pytest.mark.xfail) + if "polars" in str(constructor): + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor(data)).select( + nw.col(key).cast(value) for key, value in schema.items() + ) + + class Banana: + pass + + with pytest.raises(AssertionError, match=r"Unknown dtype"): + df.select(nw.col("a").cast(Banana)) diff --git a/tests/expr_and_series/cat/get_categories_test.py b/tests/expr_and_series/cat/get_categories_test.py index 47145e4ae..6432826c2 100644 --- a/tests/expr_and_series/cat/get_categories_test.py +++ b/tests/expr_and_series/cat/get_categories_test.py @@ -25,8 +25,8 @@ def test_get_categories(request: Any, constructor_eager: Any) -> None: result_expr = df.select(nw.col("a").cat.get_categories()) compare_dicts(result_expr, expected) - result_series = df["a"].cat.get_categories().to_list() - assert set(result_series) == set(expected["a"]) + result_series = df["a"].cat.get_categories() + compare_dicts({"a": result_series}, expected) def test_get_categories_pyarrow() -> None: @@ -41,5 +41,5 @@ def test_get_categories_pyarrow() -> None: result_expr = df.select(nw.col("a").cat.get_categories()) compare_dicts(result_expr, expected) - result_series = df["a"].cat.get_categories().to_list() - assert result_series == expected["a"] + result_series = df["a"].cat.get_categories() + compare_dicts({"a": result_series}, expected) diff --git a/tests/expr_and_series/clip_test.py b/tests/expr_and_series/clip_test.py new file mode 100644 index 000000000..909b153b7 --- /dev/null +++ b/tests/expr_and_series/clip_test.py @@ -0,0 +1,35 @@ +from typing import Any + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + + +def test_clip(constructor: Any) -> None: + df = nw.from_native(constructor({"a": [1, 2, 3, -4, 5]})) + result = df.select( + lower_only=nw.col("a").clip(lower_bound=3), + upper_only=nw.col("a").clip(upper_bound=4), + both=nw.col("a").clip(3, 4), + ) + expected = { + "lower_only": [3, 3, 3, 3, 5], + "upper_only": [1, 2, 3, -4, 4], + "both": [3, 3, 3, 3, 4], + } + compare_dicts(result, expected) + + +def test_clip_series(constructor_eager: Any) -> None: + df = nw.from_native(constructor_eager({"a": [1, 2, 3, -4, 5]}), eager_only=True) + result = { + "lower_only": df["a"].clip(lower_bound=3), + "upper_only": df["a"].clip(upper_bound=4), + "both": df["a"].clip(3, 4), + } + + expected = { + "lower_only": [3, 3, 3, 3, 5], + "upper_only": [1, 2, 3, -4, 4], + "both": [3, 3, 3, 3, 4], + } + compare_dicts(result, expected) diff --git a/tests/expr_and_series/cum_sum_test.py b/tests/expr_and_series/cum_sum_test.py index 47274f03e..e169b28f9 100644 --- a/tests/expr_and_series/cum_sum_test.py +++ b/tests/expr_and_series/cum_sum_test.py @@ -12,7 +12,7 @@ def test_cum_sum_simple(constructor: Any) -> None: df = nw.from_native(constructor(data)) - result = df.select(nw.all().cum_sum()) + result = df.select(nw.col("a", "b", "c").cum_sum()) expected = { "a": [0, 1, 3, 6, 10], "b": [1, 3, 6, 11, 14], diff --git a/tests/expr_and_series/double_test.py b/tests/expr_and_series/double_test.py index c2114248b..3a6b622b8 100644 --- a/tests/expr_and_series/double_test.py +++ b/tests/expr_and_series/double_test.py @@ -16,5 +16,10 @@ def test_double_alias(constructor: Any) -> None: data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} df = nw.from_native(constructor(data)) result = df.with_columns(nw.col("a").alias("o"), nw.all() * 2) - expected = {"o": [1, 3, 2], "a": [2, 6, 4], "b": [8, 8, 12], "z": [14.0, 16.0, 18.0]} + expected = { + "o": [1, 3, 2], + "a": [2, 6, 4], + "b": [8, 8, 12], + "z": [14.0, 16.0, 18.0], + } compare_dicts(result, expected) diff --git a/tests/expr_and_series/drop_nulls_test.py b/tests/expr_and_series/drop_nulls_test.py index 5e54e8175..f4c8e2d7a 100644 --- a/tests/expr_and_series/drop_nulls_test.py +++ b/tests/expr_and_series/drop_nulls_test.py @@ -2,11 +2,15 @@ from typing import Any -import narwhals as nw +import pytest + +import narwhals.stable.v1 as nw from tests.utils import compare_dicts -def test_drop_nulls(constructor: Any) -> None: +def test_drop_nulls(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) data = { "A": [1, 2, None, 4], "B": [5, 6, 7, 8], diff --git a/tests/expr_and_series/dt/datetime_attributes_test.py b/tests/expr_and_series/dt/datetime_attributes_test.py index b0cfb3804..4d59567df 100644 --- a/tests/expr_and_series/dt/datetime_attributes_test.py +++ b/tests/expr_and_series/dt/datetime_attributes_test.py @@ -1,5 +1,6 @@ from __future__ import annotations +from datetime import date from datetime import datetime from typing import Any @@ -19,6 +20,7 @@ @pytest.mark.parametrize( ("attribute", "expected"), [ + ("date", [date(2021, 3, 1), date(2020, 1, 2)]), ("year", [2021, 2020]), ("month", [3, 1]), ("day", [1, 2]), @@ -32,11 +34,58 @@ ], ) def test_datetime_attributes( - constructor_eager: Any, attribute: str, expected: list[int] + request: Any, constructor: Any, attribute: str, expected: list[int] ) -> None: - df = nw.from_native(constructor_eager(data), eager_only=True) + if ( + attribute == "date" + and "pandas" in str(constructor) + and "pyarrow" not in str(constructor) + ): + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor(data)) result = df.select(getattr(nw.col("a").dt, attribute)()) compare_dicts(result, {"a": expected}) + +@pytest.mark.parametrize( + ("attribute", "expected"), + [ + ("date", [date(2021, 3, 1), date(2020, 1, 2)]), + ("year", [2021, 2020]), + ("month", [3, 1]), + ("day", [1, 2]), + ("hour", [12, 2]), + ("minute", [34, 4]), + ("second", [56, 14]), + ("millisecond", [49, 715]), + ("microsecond", [49000, 715000]), + ("nanosecond", [49000000, 715000000]), + ("ordinal_day", [60, 2]), + ], +) +def test_datetime_attributes_series( + request: Any, constructor_eager: Any, attribute: str, expected: list[int] +) -> None: + if ( + attribute == "date" + and "pandas" in str(constructor_eager) + and "pyarrow" not in str(constructor_eager) + ): + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor_eager(data), eager_only=True) result = df.select(getattr(df["a"].dt, attribute)()) compare_dicts(result, {"a": expected}) + + +def test_datetime_chained_attributes(request: Any, constructor_eager: Any) -> None: + if "pandas" in str(constructor_eager) and "pyarrow" not in str(constructor_eager): + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.select(df["a"].dt.date().dt.year()) + compare_dicts(result, {"a": [2021, 2020]}) + + result = df.select(nw.col("a").dt.date().dt.year()) + compare_dicts(result, {"a": [2021, 2020]}) diff --git a/tests/expr_and_series/dt/datetime_duration_test.py b/tests/expr_and_series/dt/datetime_duration_test.py index bc5c8f0e3..50d254ba3 100644 --- a/tests/expr_and_series/dt/datetime_duration_test.py +++ b/tests/expr_and_series/dt/datetime_duration_test.py @@ -23,58 +23,57 @@ timedelta(milliseconds=1, microseconds=300), ], "c": np.array([None, 20], dtype="timedelta64[ns]"), - "d": [ - timedelta(milliseconds=2), - timedelta(seconds=1), - ], } @pytest.mark.parametrize( - ("attribute", "expected_a", "expected_b"), + ("attribute", "expected_a", "expected_b", "expected_c"), [ - ("total_minutes", [0, 1], [0, 0]), - ("total_seconds", [0, 61], [0, 0]), - ("total_milliseconds", [0, 61001], [2, 1]), + ("total_minutes", [0, 1], [0, 0], [0, 0]), + ("total_seconds", [0, 61], [0, 0], [0, 0]), + ("total_milliseconds", [0, 61001], [2, 1], [0, 0]), + ("total_microseconds", [0, 61001001], [2000, 1300], [0, 0]), + ("total_nanoseconds", [0, 61001001000], [2000000, 1300000], [0, 20]), ], ) def test_duration_attributes( request: Any, - constructor_eager: Any, + constructor: Any, attribute: str, expected_a: list[int], expected_b: list[int], + expected_c: list[int], ) -> None: - if parse_version(pd.__version__) < (2, 2) and "pandas_pyarrow" in str( - constructor_eager - ): + if parse_version(pd.__version__) < (2, 2) and "pandas_pyarrow" in str(constructor): request.applymarker(pytest.mark.xfail) - df = nw.from_native(constructor_eager(data), eager_only=True) - result_a = df.select(getattr(nw.col("a").dt, attribute)().fill_null(0)) - compare_dicts(result_a, {"a": expected_a}) + df = nw.from_native(constructor(data)) - result_a = df.select(getattr(df["a"].dt, attribute)().fill_null(0)) + result_a = df.select(getattr(nw.col("a").dt, attribute)().fill_null(0)) compare_dicts(result_a, {"a": expected_a}) result_b = df.select(getattr(nw.col("b").dt, attribute)().fill_null(0)) compare_dicts(result_b, {"b": expected_b}) - result_b = df.select(getattr(df["b"].dt, attribute)().fill_null(0)) - compare_dicts(result_b, {"b": expected_b}) + result_c = df.select(getattr(nw.col("c").dt, attribute)().fill_null(0)) + compare_dicts(result_c, {"c": expected_c}) @pytest.mark.parametrize( - ("attribute", "expected_b", "expected_c"), + ("attribute", "expected_a", "expected_b", "expected_c"), [ - ("total_microseconds", [2000, 1300], [0, 0]), - ("total_nanoseconds", [2000000, 1300000], [0, 20]), + ("total_minutes", [0, 1], [0, 0], [0, 0]), + ("total_seconds", [0, 61], [0, 0], [0, 0]), + ("total_milliseconds", [0, 61001], [2, 1], [0, 0]), + ("total_microseconds", [0, 61001001], [2000, 1300], [0, 0]), + ("total_nanoseconds", [0, 61001001000], [2000000, 1300000], [0, 20]), ], ) -def test_duration_micro_nano( +def test_duration_attributes_series( request: Any, constructor_eager: Any, attribute: str, + expected_a: list[int], expected_b: list[int], expected_c: list[int], ) -> None: @@ -85,15 +84,12 @@ def test_duration_micro_nano( df = nw.from_native(constructor_eager(data), eager_only=True) - result_b = df.select(getattr(nw.col("b").dt, attribute)().fill_null(0)) - compare_dicts(result_b, {"b": expected_b}) + result_a = df.select(getattr(df["a"].dt, attribute)().fill_null(0)) + compare_dicts(result_a, {"a": expected_a}) result_b = df.select(getattr(df["b"].dt, attribute)().fill_null(0)) compare_dicts(result_b, {"b": expected_b}) - result_c = df.select(getattr(nw.col("c").dt, attribute)().fill_null(0)) - compare_dicts(result_c, {"c": expected_c}) - result_c = df.select(getattr(df["c"].dt, attribute)().fill_null(0)) compare_dicts(result_c, {"c": expected_c}) diff --git a/tests/expr_and_series/dt/to_string_test.py b/tests/expr_and_series/dt/to_string_test.py index b78f2d1f7..7cbbf72f2 100644 --- a/tests/expr_and_series/dt/to_string_test.py +++ b/tests/expr_and_series/dt/to_string_test.py @@ -6,6 +6,7 @@ import pytest import narwhals.stable.v1 as nw +from tests.utils import compare_dicts from tests.utils import is_windows data = { @@ -17,31 +18,71 @@ @pytest.mark.parametrize( - "fmt", ["%Y-%m-%d", "%Y-%m-%d %H:%M:%S", "%Y/%m/%d %H:%M:%S", "%G-W%V-%u", "%G-W%V"] + "fmt", + [ + "%Y-%m-%d", + "%Y-%m-%d %H:%M:%S", + "%Y/%m/%d %H:%M:%S", + "%G-W%V-%u", + "%G-W%V", + ], ) @pytest.mark.skipif(is_windows(), reason="pyarrow breaking on windows") -def test_dt_to_string(constructor_eager: Any, fmt: str) -> None: +def test_dt_to_string_series(constructor_eager: Any, fmt: str) -> None: input_frame = nw.from_native(constructor_eager(data), eager_only=True) input_series = input_frame["a"] expected_col = [datetime.strftime(d, fmt) for d in data["a"]] - result = input_series.dt.to_string(fmt).to_list() + result = {"a": input_series.dt.to_string(fmt)} + if any( x in str(constructor_eager) for x in ["pandas_pyarrow", "pyarrow_table", "modin"] ): # PyArrow differs from other libraries, in that %S also shows # the fraction of a second. - result = [x[: x.find(".")] if "." in x else x for x in result] - assert result == expected_col - result = input_frame.select(nw.col("a").dt.to_string(fmt))["a"].to_list() - if any( - x in str(constructor_eager) for x in ["pandas_pyarrow", "pyarrow_table", "modin"] - ): + result = {"a": input_series.dt.to_string(fmt).str.replace(r"\.\d+$", "")} + + compare_dicts(result, {"a": expected_col}) + + +@pytest.mark.parametrize( + "fmt", + [ + "%Y-%m-%d", + "%Y-%m-%d %H:%M:%S", + "%Y/%m/%d %H:%M:%S", + "%G-W%V-%u", + "%G-W%V", + ], +) +@pytest.mark.skipif(is_windows(), reason="pyarrow breaking on windows") +def test_dt_to_string_expr(constructor: Any, fmt: str) -> None: + input_frame = nw.from_native(constructor(data)) + + expected_col = [datetime.strftime(d, fmt) for d in data["a"]] + + result = input_frame.select(nw.col("a").dt.to_string(fmt).alias("b")) + if any(x in str(constructor) for x in ["pandas_pyarrow", "pyarrow_table", "modin"]): # PyArrow differs from other libraries, in that %S also shows # the fraction of a second. - result = [x[: x.find(".")] if "." in x else x for x in result] - assert result == expected_col + result = input_frame.select( + nw.col("a").dt.to_string(fmt).str.replace(r"\.\d+$", "").alias("b") + ) + compare_dicts(result, {"b": expected_col}) + + +def _clean_string(result: str) -> str: + # rstrip '0' to remove trailing zeros, as different libraries handle this differently + # if there's then a trailing `.`, remove that too. + if "." in result: + result = result.rstrip("0").rstrip(".") + return result + + +def _clean_string_expr(e: Any) -> Any: + # Same as `_clean_string` but for Expr + return e.str.replace_all(r"0+$", "").str.replace_all(r"\.$", "") @pytest.mark.parametrize( @@ -50,48 +91,60 @@ def test_dt_to_string(constructor_eager: Any, fmt: str) -> None: (datetime(2020, 1, 9), "2020-01-09T00:00:00.000000"), (datetime(2020, 1, 9, 12, 34, 56), "2020-01-09T12:34:56.000000"), (datetime(2020, 1, 9, 12, 34, 56, 123), "2020-01-09T12:34:56.000123"), - (datetime(2020, 1, 9, 12, 34, 56, 123456), "2020-01-09T12:34:56.123456"), + ( + datetime(2020, 1, 9, 12, 34, 56, 123456), + "2020-01-09T12:34:56.123456", + ), ], ) @pytest.mark.skipif(is_windows(), reason="pyarrow breaking on windows") -def test_dt_to_string_iso_local_datetime( +def test_dt_to_string_iso_local_datetime_series( constructor_eager: Any, data: datetime, expected: str ) -> None: - def _clean_string(result: str) -> str: - # rstrip '0' to remove trailing zeros, as different libraries handle this differently - # if there's then a trailing `.`, remove that too. - if "." in result: - result = result.rstrip("0").rstrip(".") - return result - df = constructor_eager({"a": [data]}) result = ( nw.from_native(df, eager_only=True)["a"] .dt.to_string("%Y-%m-%dT%H:%M:%S.%f") - .to_list()[0] - ) - assert _clean_string(result) == _clean_string(expected) - - result = ( - nw.from_native(df, eager_only=True) - .select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M:%S.%f"))["a"] - .to_list()[0] + .item(0) ) - assert _clean_string(result) == _clean_string(expected) + assert _clean_string(str(result)) == _clean_string(expected) result = ( nw.from_native(df, eager_only=True)["a"] .dt.to_string("%Y-%m-%dT%H:%M:%S%.f") - .to_list()[0] + .item(0) ) - assert _clean_string(result) == _clean_string(expected) + assert _clean_string(str(result)) == _clean_string(expected) - result = ( - nw.from_native(df, eager_only=True) - .select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M:%S%.f"))["a"] - .to_list()[0] + +@pytest.mark.parametrize( + ("data", "expected"), + [ + (datetime(2020, 1, 9, 12, 34, 56), "2020-01-09T12:34:56.000000"), + (datetime(2020, 1, 9, 12, 34, 56, 123), "2020-01-09T12:34:56.000123"), + ( + datetime(2020, 1, 9, 12, 34, 56, 123456), + "2020-01-09T12:34:56.123456", + ), + ], +) +@pytest.mark.skipif(is_windows(), reason="pyarrow breaking on windows") +def test_dt_to_string_iso_local_datetime_expr( + request: Any, constructor: Any, data: datetime, expected: str +) -> None: + if "modin" in str(constructor): + request.applymarker(pytest.mark.xfail) + df = constructor({"a": [data]}) + + result = nw.from_native(df).with_columns( + _clean_string_expr(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M:%S.%f")).alias("b") ) - assert _clean_string(result) == _clean_string(expected) + compare_dicts(result, {"a": [data], "b": [_clean_string(expected)]}) + + result = nw.from_native(df).with_columns( + _clean_string_expr(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M:%S%.f")).alias("b") + ) + compare_dicts(result, {"a": [data], "b": [_clean_string(expected)]}) @pytest.mark.parametrize( @@ -99,18 +152,27 @@ def _clean_string(result: str) -> str: [(datetime(2020, 1, 9), "2020-01-09")], ) @pytest.mark.skipif(is_windows(), reason="pyarrow breaking on windows") -def test_dt_to_string_iso_local_date( +def test_dt_to_string_iso_local_date_series( constructor_eager: Any, data: datetime, expected: str ) -> None: df = constructor_eager({"a": [data]}) - result = ( - nw.from_native(df, eager_only=True)["a"].dt.to_string("%Y-%m-%d").to_list()[0] - ) - assert result == expected + result = nw.from_native(df, eager_only=True)["a"].dt.to_string("%Y-%m-%d").item(0) + assert str(result) == expected - result = ( - nw.from_native(df, eager_only=True) - .select(b=nw.col("a").dt.to_string("%Y-%m-%d"))["b"] - .to_list()[0] + +@pytest.mark.parametrize( + ("data", "expected"), + [(datetime(2020, 1, 9), "2020-01-09")], +) +@pytest.mark.skipif(is_windows(), reason="pyarrow breaking on windows") +def test_dt_to_string_iso_local_date_expr( + request: Any, constructor: Any, data: datetime, expected: str +) -> None: + if "modin" in str(constructor): + request.applymarker(pytest.mark.xfail) + + df = constructor({"a": [data]}) + result = nw.from_native(df).with_columns( + nw.col("a").dt.to_string("%Y-%m-%d").alias("b") ) - assert result == expected + compare_dicts(result, {"a": [data], "b": [expected]}) diff --git a/tests/expr_and_series/filter_test.py b/tests/expr_and_series/filter_test.py index 3d3ca0ed7..b55a0368e 100644 --- a/tests/expr_and_series/filter_test.py +++ b/tests/expr_and_series/filter_test.py @@ -1,5 +1,7 @@ from typing import Any +import pytest + import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -11,7 +13,9 @@ } -def test_filter(constructor: Any) -> None: +def test_filter(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor(data)) result = df.select(nw.col("a").filter(nw.col("i") < 2, nw.col("c") == 5)) expected = {"a": [0]} @@ -23,3 +27,6 @@ def test_filter_series(constructor_eager: Any) -> None: result = df.select(df["a"].filter((df["i"] < 2) & (df["c"] == 5))) expected = {"a": [0]} compare_dicts(result, expected) + result_s = df["a"].filter([True, False, False, False, False]) + expected = {"a": [0]} + compare_dicts({"a": result_s}, expected) diff --git a/tests/expr_and_series/gather_every_test.py b/tests/expr_and_series/gather_every_test.py index a1a6f66ee..b00014f20 100644 --- a/tests/expr_and_series/gather_every_test.py +++ b/tests/expr_and_series/gather_every_test.py @@ -10,7 +10,9 @@ @pytest.mark.parametrize("n", [1, 2, 3]) @pytest.mark.parametrize("offset", [1, 2, 3]) -def test_gather_every_expr(constructor: Any, n: int, offset: int) -> None: +def test_gather_every_expr(constructor: Any, n: int, offset: int, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor(data)) result = df.select(nw.col("a").gather_every(n=n, offset=offset)) @@ -27,4 +29,4 @@ def test_gather_every_series(constructor_eager: Any, n: int, offset: int) -> Non result = series.gather_every(n=n, offset=offset) expected = data["a"][offset::n] - assert result.to_list() == expected + compare_dicts({"a": result}, {"a": expected}) diff --git a/tests/expr_and_series/head_test.py b/tests/expr_and_series/head_test.py new file mode 100644 index 000000000..ef2ed1bf1 --- /dev/null +++ b/tests/expr_and_series/head_test.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +import narwhals as nw +from tests.utils import compare_dicts + + +@pytest.mark.parametrize("n", [2, -1]) +def test_head(constructor: Any, n: int, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) + if "polars" in str(constructor) and n < 0: + request.applymarker(pytest.mark.xfail) + df = nw.from_native(constructor({"a": [1, 2, 3]})) + result = df.select(nw.col("a").head(n)) + expected = {"a": [1, 2]} + compare_dicts(result, expected) + + +@pytest.mark.parametrize("n", [2, -1]) +def test_head_series(constructor_eager: Any, n: int) -> None: + df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True) + result = df.select(df["a"].head(n)) + expected = {"a": [1, 2]} + compare_dicts(result, expected) diff --git a/tests/expr_and_series/is_duplicated_test.py b/tests/expr_and_series/is_duplicated_test.py index 1aaa7fa82..5fa060312 100644 --- a/tests/expr_and_series/is_duplicated_test.py +++ b/tests/expr_and_series/is_duplicated_test.py @@ -1,8 +1,5 @@ from typing import Any -import numpy as np -import pytest - import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -12,10 +9,7 @@ } -def test_is_duplicated_expr(constructor: Any, request: Any) -> None: - if "modin" in str(constructor): - # TODO(unassigned): why is Modin failing here? - request.applymarker(pytest.mark.xfail) +def test_is_duplicated_expr(constructor: Any) -> None: df = nw.from_native(constructor(data)) result = df.select(nw.all().is_duplicated()) expected = { @@ -28,5 +22,5 @@ def test_is_duplicated_expr(constructor: Any, request: Any) -> None: def test_is_duplicated_series(constructor_eager: Any) -> None: series = nw.from_native(constructor_eager(data), eager_only=True)["a"] result = series.is_duplicated() - expected = np.array([True, True, False]) - assert (result.to_numpy() == expected).all() + expected = {"a": [True, True, False]} + compare_dicts({"a": result}, expected) diff --git a/tests/expr_and_series/is_first_distinct_test.py b/tests/expr_and_series/is_first_distinct_test.py index e412f08cd..8521661d6 100644 --- a/tests/expr_and_series/is_first_distinct_test.py +++ b/tests/expr_and_series/is_first_distinct_test.py @@ -1,8 +1,5 @@ from typing import Any -import numpy as np -import pytest - import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -12,10 +9,7 @@ } -def test_is_first_distinct_expr(constructor: Any, request: Any) -> None: - if "modin" in str(constructor): - # TODO(unassigned): why is Modin failing here? - request.applymarker(pytest.mark.xfail) +def test_is_first_distinct_expr(constructor: Any) -> None: df = nw.from_native(constructor(data)) result = df.select(nw.all().is_first_distinct()) expected = { @@ -28,5 +22,7 @@ def test_is_first_distinct_expr(constructor: Any, request: Any) -> None: def test_is_first_distinct_series(constructor_eager: Any) -> None: series = nw.from_native(constructor_eager(data), eager_only=True)["a"] result = series.is_first_distinct() - expected = np.array([True, False, True, True, False]) - assert (result.to_numpy() == expected).all() + expected = { + "a": [True, False, True, True, False], + } + compare_dicts({"a": result}, expected) diff --git a/tests/expr_and_series/is_in_test.py b/tests/expr_and_series/is_in_test.py index 9579d3d15..40c7b2718 100644 --- a/tests/expr_and_series/is_in_test.py +++ b/tests/expr_and_series/is_in_test.py @@ -1,12 +1,11 @@ from typing import Any +import pytest + import narwhals.stable.v1 as nw from tests.utils import compare_dicts -series = [1, 4, 2, 5] -data = { - "a": series, -} +data = {"a": [1, 4, 2, 5]} def test_expr_is_in(constructor: Any) -> None: @@ -18,9 +17,19 @@ def test_expr_is_in(constructor: Any) -> None: def test_ser_is_in(constructor_eager: Any) -> None: - ser = nw.from_native(constructor_eager({"a": series}), eager_only=True)["a"] - result = ser.is_in([4, 5]).to_list() - assert not result[0] - assert result[1] - assert not result[2] - assert result[3] + ser = nw.from_native(constructor_eager(data), eager_only=True)["a"] + result = {"a": ser.is_in([4, 5])} + expected = {"a": [False, True, False, True]} + + compare_dicts(result, expected) + + +def test_is_in_other(constructor: Any) -> None: + df_raw = constructor(data) + with pytest.raises( + NotImplementedError, + match=( + "Narwhals `is_in` doesn't accept expressions as an argument, as opposed to Polars. You should provide an iterable instead." + ), + ): + nw.from_native(df_raw).with_columns(contains=nw.col("a").is_in("sets")) diff --git a/tests/expr_and_series/is_last_distinct_test.py b/tests/expr_and_series/is_last_distinct_test.py index 01a5eb6dd..2e4709efd 100644 --- a/tests/expr_and_series/is_last_distinct_test.py +++ b/tests/expr_and_series/is_last_distinct_test.py @@ -1,8 +1,5 @@ from typing import Any -import numpy as np -import pytest - import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -12,10 +9,7 @@ } -def test_is_last_distinct_expr(constructor: Any, request: Any) -> None: - if "modin" in str(constructor): - # TODO(unassigned): why is Modin failing here? - request.applymarker(pytest.mark.xfail) +def test_is_last_distinct_expr(constructor: Any) -> None: df = nw.from_native(constructor(data)) result = df.select(nw.all().is_last_distinct()) expected = { @@ -28,5 +22,7 @@ def test_is_last_distinct_expr(constructor: Any, request: Any) -> None: def test_is_last_distinct_series(constructor_eager: Any) -> None: series = nw.from_native(constructor_eager(data), eager_only=True)["a"] result = series.is_last_distinct() - expected = np.array([False, True, False, True, True]) - assert (result.to_numpy() == expected).all() + expected = { + "a": [False, True, False, True, True], + } + compare_dicts({"a": result}, expected) diff --git a/tests/expr_and_series/is_null_test.py b/tests/expr_and_series/is_null_test.py new file mode 100644 index 000000000..07465fd9b --- /dev/null +++ b/tests/expr_and_series/is_null_test.py @@ -0,0 +1,22 @@ +from typing import Any + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + + +def test_null(constructor: Any) -> None: + data_na = {"a": [None, 3, 2], "z": [7.0, None, None]} + expected = {"a": [True, False, False], "z": [True, False, False]} + df = nw.from_native(constructor(data_na)) + result = df.select(nw.col("a").is_null(), ~nw.col("z").is_null()) + + compare_dicts(result, expected) + + +def test_null_series(constructor_eager: Any) -> None: + data_na = {"a": [None, 3, 2], "z": [7.0, None, None]} + expected = {"a": [True, False, False], "z": [True, False, False]} + df = nw.from_native(constructor_eager(data_na), eager_only=True) + result = {"a": df["a"].is_null(), "z": ~df["z"].is_null()} + + compare_dicts(result, expected) diff --git a/tests/expr_and_series/is_unique_test.py b/tests/expr_and_series/is_unique_test.py index 6f803243c..8bddbb647 100644 --- a/tests/expr_and_series/is_unique_test.py +++ b/tests/expr_and_series/is_unique_test.py @@ -1,8 +1,5 @@ from typing import Any -import numpy as np -import pytest - import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -12,10 +9,7 @@ } -def test_is_unique_expr(constructor: Any, request: Any) -> None: - if "modin" in str(constructor): - # TODO(unassigned): why is Modin failing here? - request.applymarker(pytest.mark.xfail) +def test_is_unique_expr(constructor: Any) -> None: df = nw.from_native(constructor(data)) result = df.select(nw.all().is_unique()) expected = { @@ -28,5 +22,7 @@ def test_is_unique_expr(constructor: Any, request: Any) -> None: def test_is_unique_series(constructor_eager: Any) -> None: series = nw.from_native(constructor_eager(data), eager_only=True)["a"] result = series.is_unique() - expected = np.array([False, False, True]) - assert (result.to_numpy() == expected).all() + expected = { + "a": [False, False, True], + } + compare_dicts({"a": result}, expected) diff --git a/tests/expr_and_series/len_test.py b/tests/expr_and_series/len_test.py index 7ec1505cc..8a52dd327 100644 --- a/tests/expr_and_series/len_test.py +++ b/tests/expr_and_series/len_test.py @@ -1,17 +1,53 @@ from typing import Any +import pytest + import narwhals.stable.v1 as nw from tests.utils import compare_dicts -data = {"a": list("xyz"), "b": [1, 2, 1]} -expected = {"a1": [2], "a2": [1]} +def test_len_no_filter(constructor: Any) -> None: + data = {"a": list("xyz"), "b": [1, 2, 1]} + expected = {"l": [3], "l2": [6]} + df = nw.from_native(constructor(data)).select( + nw.col("a").len().alias("l"), + (nw.col("a").len() * 2).alias("l2"), + ) + + compare_dicts(df, expected) -def test_len(constructor: Any) -> None: - df_raw = constructor(data) - df = nw.from_native(df_raw).select( + +def test_len_chaining(constructor: Any, request: Any) -> None: + data = {"a": list("xyz"), "b": [1, 2, 1]} + expected = {"a1": [2], "a2": [1]} + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) + df = nw.from_native(constructor(data)).select( nw.col("a").filter(nw.col("b") == 1).len().alias("a1"), nw.col("a").filter(nw.col("b") == 2).len().alias("a2"), ) compare_dicts(df, expected) + + +def test_namespace_len(constructor: Any) -> None: + df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]})).select( + nw.len(), a=nw.len() + ) + expected = {"len": [3], "a": [3]} + compare_dicts(df, expected) + df = ( + nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]})) + .select() + .select(nw.len(), a=nw.len()) + ) + expected = {"len": [0], "a": [0]} + compare_dicts(df, expected) + + +def test_len_series(constructor_eager: Any) -> None: + data = {"a": [1, 2, 1]} + s = nw.from_native(constructor_eager(data), eager_only=True)["a"] + + assert s.len() == 3 + assert len(s) == 3 diff --git a/tests/expr_and_series/mean_horizontal_test.py b/tests/expr_and_series/mean_horizontal_test.py new file mode 100644 index 000000000..d42d5e324 --- /dev/null +++ b/tests/expr_and_series/mean_horizontal_test.py @@ -0,0 +1,15 @@ +from typing import Any + +import pytest + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + + +@pytest.mark.parametrize("col_expr", [nw.col("a"), "a"]) +def test_meanh(constructor: Any, col_expr: Any) -> None: + data = {"a": [1, 3, None, None], "b": [4, None, 6, None]} + df = nw.from_native(constructor(data)) + result = df.select(horizontal_mean=nw.mean_horizontal(col_expr, nw.col("b"))) + expected = {"horizontal_mean": [2.5, 3.0, 6.0, float("nan")]} + compare_dicts(result, expected) diff --git a/tests/expr_and_series/name/keep_test.py b/tests/expr_and_series/name/keep_test.py index beceb7b2a..0b43abe40 100644 --- a/tests/expr_and_series/name/keep_test.py +++ b/tests/expr_and_series/name/keep_test.py @@ -6,7 +6,7 @@ import polars as pl import pytest -import narwhals as nw +import narwhals.stable.v1 as nw from tests.utils import compare_dicts data = {"foo": [1, 2, 3], "BAR": [4, 5, 6]} diff --git a/tests/expr_and_series/name/map_test.py b/tests/expr_and_series/name/map_test.py index 9aaf284b7..ff039e30d 100644 --- a/tests/expr_and_series/name/map_test.py +++ b/tests/expr_and_series/name/map_test.py @@ -6,7 +6,7 @@ import polars as pl import pytest -import narwhals as nw +import narwhals.stable.v1 as nw from tests.utils import compare_dicts data = {"foo": [1, 2, 3], "BAR": [4, 5, 6]} diff --git a/tests/expr_and_series/name/prefix_test.py b/tests/expr_and_series/name/prefix_test.py index 1c9046606..f538d4136 100644 --- a/tests/expr_and_series/name/prefix_test.py +++ b/tests/expr_and_series/name/prefix_test.py @@ -6,7 +6,7 @@ import polars as pl import pytest -import narwhals as nw +import narwhals.stable.v1 as nw from tests.utils import compare_dicts data = {"foo": [1, 2, 3], "BAR": [4, 5, 6]} diff --git a/tests/expr_and_series/name/suffix_test.py b/tests/expr_and_series/name/suffix_test.py index f14326472..0e952449b 100644 --- a/tests/expr_and_series/name/suffix_test.py +++ b/tests/expr_and_series/name/suffix_test.py @@ -6,7 +6,7 @@ import polars as pl import pytest -import narwhals as nw +import narwhals.stable.v1 as nw from tests.utils import compare_dicts data = {"foo": [1, 2, 3], "BAR": [4, 5, 6]} diff --git a/tests/expr_and_series/name/to_lowercase_test.py b/tests/expr_and_series/name/to_lowercase_test.py index 3aa569f17..a9e8bfcfd 100644 --- a/tests/expr_and_series/name/to_lowercase_test.py +++ b/tests/expr_and_series/name/to_lowercase_test.py @@ -6,7 +6,7 @@ import polars as pl import pytest -import narwhals as nw +import narwhals.stable.v1 as nw from tests.utils import compare_dicts data = {"foo": [1, 2, 3], "BAR": [4, 5, 6]} diff --git a/tests/expr_and_series/name/to_uppercase_test.py b/tests/expr_and_series/name/to_uppercase_test.py index df3e89ee2..035dfeff2 100644 --- a/tests/expr_and_series/name/to_uppercase_test.py +++ b/tests/expr_and_series/name/to_uppercase_test.py @@ -6,7 +6,7 @@ import polars as pl import pytest -import narwhals as nw +import narwhals.stable.v1 as nw from tests.utils import compare_dicts data = {"foo": [1, 2, 3], "BAR": [4, 5, 6]} diff --git a/tests/expr_and_series/null_count_test.py b/tests/expr_and_series/null_count_test.py index 497e73a9f..a6cb58f71 100644 --- a/tests/expr_and_series/null_count_test.py +++ b/tests/expr_and_series/null_count_test.py @@ -9,7 +9,7 @@ } -def test_null_count(constructor: Any) -> None: +def test_null_count_expr(constructor: Any) -> None: df = nw.from_native(constructor(data)) result = df.select(nw.all().null_count()) expected = { @@ -17,3 +17,10 @@ def test_null_count(constructor: Any) -> None: "b": [1], } compare_dicts(result, expected) + + +def test_null_count_series(constructor_eager: Any) -> None: + data = [1, 2, None] + series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] + result = series.null_count() + assert result == 1 diff --git a/tests/expr_and_series/operators_test.py b/tests/expr_and_series/operators_test.py index f29c8e9eb..113824a94 100644 --- a/tests/expr_and_series/operators_test.py +++ b/tests/expr_and_series/operators_test.py @@ -19,7 +19,7 @@ ("__gt__", [False, False, True]), ], ) -def test_comparand_operators( +def test_comparand_operators_scalar_expr( constructor: Any, operator: str, expected: list[bool] ) -> None: data = {"a": [0, 1, 2]} @@ -28,6 +28,26 @@ def test_comparand_operators( compare_dicts(result, {"a": expected}) +@pytest.mark.parametrize( + ("operator", "expected"), + [ + ("__eq__", [True, False, False]), + ("__ne__", [False, True, True]), + ("__le__", [True, False, True]), + ("__lt__", [False, False, True]), + ("__ge__", [True, True, False]), + ("__gt__", [False, True, False]), + ], +) +def test_comparand_operators_expr( + constructor: Any, operator: str, expected: list[bool] +) -> None: + data = {"a": [0, 1, 1], "b": [0, 0, 2]} + df = nw.from_native(constructor(data)) + result = df.select(getattr(nw.col("a"), operator)(nw.col("b"))) + compare_dicts(result, {"a": expected}) + + @pytest.mark.parametrize( ("operator", "expected"), [ @@ -35,9 +55,69 @@ def test_comparand_operators( ("__or__", [True, True, True, False]), ], ) -def test_logic_operators(constructor: Any, operator: str, expected: list[bool]) -> None: +def test_logic_operators_expr( + constructor: Any, operator: str, expected: list[bool] +) -> None: data = {"a": [True, True, False, False], "b": [True, False, True, False]} df = nw.from_native(constructor(data)) result = df.select(getattr(nw.col("a"), operator)(nw.col("b"))) compare_dicts(result, {"a": expected}) + + +@pytest.mark.parametrize( + ("operator", "expected"), + [ + ("__eq__", [False, True, False]), + ("__ne__", [True, False, True]), + ("__le__", [True, True, False]), + ("__lt__", [True, False, False]), + ("__ge__", [False, True, True]), + ("__gt__", [False, False, True]), + ], +) +def test_comparand_operators_scalar_series( + constructor_eager: Any, operator: str, expected: list[bool] +) -> None: + data = {"a": [0, 1, 2]} + s = nw.from_native(constructor_eager(data), eager_only=True)["a"] + result = {"a": (getattr(s, operator)(1))} + compare_dicts(result, {"a": expected}) + + +@pytest.mark.parametrize( + ("operator", "expected"), + [ + ("__eq__", [True, False, False]), + ("__ne__", [False, True, True]), + ("__le__", [True, False, True]), + ("__lt__", [False, False, True]), + ("__ge__", [True, True, False]), + ("__gt__", [False, True, False]), + ], +) +def test_comparand_operators_series( + constructor_eager: Any, operator: str, expected: list[bool] +) -> None: + data = {"a": [0, 1, 1], "b": [0, 0, 2]} + df = nw.from_native(constructor_eager(data), eager_only=True) + series, other = df["a"], df["b"] + result = {"a": getattr(series, operator)(other)} + compare_dicts(result, {"a": expected}) + + +@pytest.mark.parametrize( + ("operator", "expected"), + [ + ("__and__", [True, False, False, False]), + ("__or__", [True, True, True, False]), + ], +) +def test_logic_operators_series( + constructor_eager: Any, operator: str, expected: list[bool] +) -> None: + data = {"a": [True, True, False, False], "b": [True, False, True, False]} + df = nw.from_native(constructor_eager(data), eager_only=True) + series, other = df["a"], df["b"] + result = {"a": getattr(series, operator)(other)} + compare_dicts(result, {"a": expected}) diff --git a/tests/expr_and_series/over_test.py b/tests/expr_and_series/over_test.py index b037d0386..fb01a3cfd 100644 --- a/tests/expr_and_series/over_test.py +++ b/tests/expr_and_series/over_test.py @@ -1,6 +1,5 @@ from typing import Any -import pandas as pd import pytest import narwhals.stable.v1 as nw @@ -13,10 +12,7 @@ } -def test_over_single(request: Any, constructor: Any) -> None: - if "pyarrow_table" in str(constructor): - request.applymarker(pytest.mark.xfail) - +def test_over_single(constructor: Any) -> None: df = nw.from_native(constructor(data)) result = df.with_columns(c_max=nw.col("c").max().over("a")) expected = { @@ -28,10 +24,7 @@ def test_over_single(request: Any, constructor: Any) -> None: compare_dicts(result, expected) -def test_over_multiple(request: Any, constructor: Any) -> None: - if "pyarrow_table" in str(constructor): - request.applymarker(pytest.mark.xfail) - +def test_over_multiple(constructor: Any) -> None: df = nw.from_native(constructor(data)) result = df.with_columns(c_min=nw.col("c").min().over("a", "b")) expected = { @@ -43,7 +36,10 @@ def test_over_multiple(request: Any, constructor: Any) -> None: compare_dicts(result, expected) -def test_over_invalid() -> None: - df = nw.from_native(pd.DataFrame(data)) +def test_over_invalid(request: Any, constructor: Any) -> None: + if "polars" in str(constructor): + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor(data)) with pytest.raises(ValueError, match="Anonymous expressions"): df.with_columns(c_min=nw.all().min().over("a", "b")) diff --git a/tests/expr_and_series/pipe_test.py b/tests/expr_and_series/pipe_test.py new file mode 100644 index 000000000..55de3548b --- /dev/null +++ b/tests/expr_and_series/pipe_test.py @@ -0,0 +1,21 @@ +from typing import Any + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + +input_list = {"a": [2, 4, 6, 8]} +expected = [4, 16, 36, 64] + + +def test_pipe_expr(constructor: Any) -> None: + df = nw.from_native(constructor(input_list)) + e = df.select(nw.col("a").pipe(lambda x: x**2)) + compare_dicts(e, {"a": expected}) + + +def test_pipe_series( + constructor_eager: Any, +) -> None: + s = nw.from_native(constructor_eager(input_list), eager_only=True)["a"] + result = s.pipe(lambda x: x**2) + compare_dicts({"a": result}, {"a": expected}) diff --git a/tests/expr_and_series/quantile_test.py b/tests/expr_and_series/quantile_test.py index 592a76857..d9064541f 100644 --- a/tests/expr_and_series/quantile_test.py +++ b/tests/expr_and_series/quantile_test.py @@ -24,7 +24,10 @@ def test_quantile_expr( constructor: Any, interpolation: Literal["nearest", "higher", "lower", "midpoint", "linear"], expected: dict[str, list[float]], + request: Any, ) -> None: + if "dask" in str(constructor) and interpolation != "linear": + request.applymarker(pytest.mark.xfail) q = 0.3 data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} df_raw = constructor(data) diff --git a/tests/expr_and_series/reduction_test.py b/tests/expr_and_series/reduction_test.py new file mode 100644 index 000000000..60750444e --- /dev/null +++ b/tests/expr_and_series/reduction_test.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + + +@pytest.mark.parametrize( + ("expr", "expected"), + [ + ( + [nw.col("a").min().alias("min"), nw.col("a", "b").mean()], + {"min": [1], "a": [2], "b": [5]}, + ), + ([(nw.col("a") + nw.col("b").max()).alias("x")], {"x": [7, 8, 9]}), + ([nw.col("a"), nw.col("b").min()], {"a": [1, 2, 3], "b": [4, 4, 4]}), + ([nw.col("a").max(), nw.col("b")], {"a": [3, 3, 3], "b": [4, 5, 6]}), + ( + [nw.col("a"), nw.col("b").min().alias("min")], + {"a": [1, 2, 3], "min": [4, 4, 4]}, + ), + ], + ids=range(5), +) +def test_scalar_reduction_select( + constructor: Any, expr: list[Any], expected: dict[str, list[Any]] +) -> None: + data = {"a": [1, 2, 3], "b": [4, 5, 6]} + df = nw.from_native(constructor(data)) + result = df.select(*expr) + compare_dicts(result, expected) + + +@pytest.mark.parametrize( + ("expr", "expected"), + [ + ( + [nw.col("a").min().alias("min"), nw.col("a", "b").mean()], + {"min": [1, 1, 1], "a": [2, 2, 2], "b": [5, 5, 5]}, + ), + ([(nw.col("a") + nw.col("b").max()).alias("x")], {"x": [7, 8, 9]}), + ([nw.col("a"), nw.col("b").min()], {"a": [1, 2, 3], "b": [4, 4, 4]}), + ([nw.col("a").max(), nw.col("b")], {"a": [3, 3, 3], "b": [4, 5, 6]}), + ( + [nw.col("a"), nw.col("b").min().alias("min")], + {"a": [1, 2, 3], "min": [4, 4, 4]}, + ), + ], + ids=range(5), +) +def test_scalar_reduction_with_columns( + constructor: Any, expr: list[Any], expected: dict[str, list[Any]] +) -> None: + data = {"a": [1, 2, 3], "b": [4, 5, 6]} + df = nw.from_native(constructor(data)) + result = df.with_columns(*expr).select(*expected.keys()) + compare_dicts(result, expected) diff --git a/tests/expr_and_series/round_test.py b/tests/expr_and_series/round_test.py index 0c021df33..769e4be11 100644 --- a/tests/expr_and_series/round_test.py +++ b/tests/expr_and_series/round_test.py @@ -9,10 +9,8 @@ @pytest.mark.parametrize("decimals", [0, 1, 2]) -def test_round(request: Any, constructor: Any, decimals: int) -> None: - if "pyarrow_table" in str(constructor): - request.applymarker(pytest.mark.xfail) - data = {"a": [1.12345, 2.56789, 3.901234]} +def test_round(constructor: Any, decimals: int) -> None: + data = {"a": [2.12345, 2.56789, 3.901234]} df_raw = constructor(data) df = nw.from_native(df_raw) @@ -22,9 +20,7 @@ def test_round(request: Any, constructor: Any, decimals: int) -> None: @pytest.mark.parametrize("decimals", [0, 1, 2]) -def test_round_series(request: Any, constructor_eager: Any, decimals: int) -> None: - if "pyarrow_table" in str(constructor_eager): - request.applymarker(pytest.mark.xfail) +def test_round_series(constructor_eager: Any, decimals: int) -> None: data = {"a": [1.12345, 2.56789, 3.901234]} df_raw = constructor_eager(data) df = nw.from_native(df_raw, eager_only=True) @@ -32,4 +28,4 @@ def test_round_series(request: Any, constructor_eager: Any, decimals: int) -> No expected_data = {k: [round(e, decimals) for e in v] for k, v in data.items()} result_series = df["a"].round(decimals) - assert result_series.to_list() == expected_data["a"] + compare_dicts({"a": result_series}, expected_data) diff --git a/tests/expr_and_series/sample_test.py b/tests/expr_and_series/sample_test.py index a19c686e6..c64703d3c 100644 --- a/tests/expr_and_series/sample_test.py +++ b/tests/expr_and_series/sample_test.py @@ -1,9 +1,13 @@ from typing import Any +import pytest + import narwhals.stable.v1 as nw -def test_expr_sample(constructor: Any) -> None: +def test_expr_sample(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]})).lazy() result_expr = df.select(nw.col("a").sample(n=2)).collect().shape @@ -15,7 +19,9 @@ def test_expr_sample(constructor: Any) -> None: assert result_series == expected_series -def test_expr_sample_fraction(constructor: Any) -> None: +def test_expr_sample_fraction(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor({"a": [1, 2, 3] * 10, "b": [4, 5, 6] * 10})).lazy() result_expr = df.select(nw.col("a").sample(fraction=0.1)).collect().shape diff --git a/tests/expr_and_series/shift_test.py b/tests/expr_and_series/shift_test.py index 9460357d6..02dbed6b0 100644 --- a/tests/expr_and_series/shift_test.py +++ b/tests/expr_and_series/shift_test.py @@ -1,6 +1,6 @@ from typing import Any -import pytest +import pyarrow as pa import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -13,10 +13,7 @@ } -def test_shift(request: Any, constructor: Any) -> None: - if "pyarrow_table" in str(constructor): - request.applymarker(pytest.mark.xfail) - +def test_shift(constructor: Any) -> None: df = nw.from_native(constructor(data)) result = df.with_columns(nw.col("a", "b", "c").shift(2)).filter(nw.col("i") > 1) expected = { @@ -28,21 +25,35 @@ def test_shift(request: Any, constructor: Any) -> None: compare_dicts(result, expected) -def test_shift_series(request: Any, constructor_eager: Any) -> None: - if "pyarrow_table" in str(constructor_eager): - request.applymarker(pytest.mark.xfail) - +def test_shift_series(constructor_eager: Any) -> None: df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.with_columns( + df["a"].shift(2), + df["b"].shift(2), + df["c"].shift(2), + ).filter(nw.col("i") > 1) expected = { "i": [2, 3, 4], "a": [0, 1, 2], "b": [1, 2, 3], "c": [5, 4, 3], } - result = df.select( - df["i"], - df["a"].shift(2), - df["b"].shift(2), - df["c"].shift(2), - ).filter(nw.col("i") > 1) + compare_dicts(result, expected) + + +def test_shift_multi_chunk_pyarrow() -> None: + tbl = pa.table({"a": [1, 2, 3]}) + tbl = pa.concat_tables([tbl, tbl, tbl]) + df = nw.from_native(tbl, eager_only=True) + + result = df.select(nw.col("a").shift(1)) + expected = {"a": [None, 1, 2, 3, 1, 2, 3, 1, 2]} + compare_dicts(result, expected) + + result = df.select(nw.col("a").shift(-1)) + expected = {"a": [2, 3, 1, 2, 3, 1, 2, 3, None]} + compare_dicts(result, expected) + + result = df.select(nw.col("a").shift(0)) + expected = {"a": [1, 2, 3, 1, 2, 3, 1, 2, 3]} compare_dicts(result, expected) diff --git a/tests/expr_and_series/str/contains_test.py b/tests/expr_and_series/str/contains_test.py index 7c69ee21c..5cc90f4ad 100644 --- a/tests/expr_and_series/str/contains_test.py +++ b/tests/expr_and_series/str/contains_test.py @@ -12,17 +12,25 @@ df_polars = pl.DataFrame(data) -def test_contains(constructor_eager: Any) -> None: - df = nw.from_native(constructor_eager(data), eager_only=True) +def test_contains(constructor: Any) -> None: + df = nw.from_native(constructor(data)) result = df.with_columns( - case_insensitive_match=nw.col("pets").str.contains("(?i)parrot|Dove") + nw.col("pets").str.contains("(?i)parrot|Dove").alias("result") ) expected = { "pets": ["cat", "dog", "rabbit and parrot", "dove"], - "case_insensitive_match": [False, False, True, True], + "result": [False, False, True, True], } compare_dicts(result, expected) + + +def test_contains_series(constructor_eager: Any) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) result = df.with_columns( case_insensitive_match=df["pets"].str.contains("(?i)parrot|Dove") ) + expected = { + "pets": ["cat", "dog", "rabbit and parrot", "dove"], + "case_insensitive_match": [False, False, True, True], + } compare_dicts(result, expected) diff --git a/tests/expr_and_series/str/head_test.py b/tests/expr_and_series/str/head_test.py index 013c31b54..1160920fd 100644 --- a/tests/expr_and_series/str/head_test.py +++ b/tests/expr_and_series/str/head_test.py @@ -6,12 +6,19 @@ data = {"a": ["foo", "bars"]} -def test_str_head(constructor_eager: Any) -> None: - df = nw.from_native(constructor_eager(data), eager_only=True) +def test_str_head(constructor: Any) -> None: + df = nw.from_native(constructor(data)) result = df.select(nw.col("a").str.head(3)) expected = { "a": ["foo", "bar"], } compare_dicts(result, expected) + + +def test_str_head_series(constructor_eager: Any) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) + expected = { + "a": ["foo", "bar"], + } result = df.select(df["a"].str.head(3)) compare_dicts(result, expected) diff --git a/tests/expr_and_series/str/replace_test.py b/tests/expr_and_series/str/replace_test.py new file mode 100644 index 000000000..95b5bd87c --- /dev/null +++ b/tests/expr_and_series/str/replace_test.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + +replace_data = [ + ( + {"a": ["123abc", "abc456"]}, + r"abc\b", + "ABC", + 1, + False, + {"a": ["123ABC", "abc456"]}, + ), + ({"a": ["abc abc", "abc456"]}, r"abc", "", 1, False, {"a": [" abc", "456"]}), + ({"a": ["abc abc abc", "456abc"]}, r"abc", "", -1, False, {"a": [" ", "456"]}), + ( + {"a": ["Dollar $ign", "literal"]}, + r"$", + "S", + -1, + True, + {"a": ["Dollar Sign", "literal"]}, + ), +] + +replace_all_data = [ + ( + {"a": ["123abc", "abc456"]}, + r"abc\b", + "ABC", + False, + {"a": ["123ABC", "abc456"]}, + ), + ({"a": ["abc abc", "abc456"]}, r"abc", "", False, {"a": [" ", "456"]}), + ({"a": ["abc abc abc", "456abc"]}, r"abc", "", False, {"a": [" ", "456"]}), + ( + {"a": ["Dollar $ign", "literal"]}, + r"$", + "S", + True, + {"a": ["Dollar Sign", "literal"]}, + ), +] + + +@pytest.mark.parametrize( + ("data", "pattern", "value", "n", "literal", "expected"), + replace_data, +) +def test_str_replace_series( + constructor_eager: Any, + data: dict[str, list[str]], + pattern: str, + value: str, + n: int, + literal: bool, # noqa: FBT001 + expected: dict[str, list[str]], +) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) + + result_series = df["a"].str.replace( + pattern=pattern, value=value, n=n, literal=literal + ) + compare_dicts({"a": result_series}, expected) + + +@pytest.mark.parametrize( + ("data", "pattern", "value", "literal", "expected"), + replace_all_data, +) +def test_str_replace_all_series( + constructor_eager: Any, + data: dict[str, list[str]], + pattern: str, + value: str, + literal: bool, # noqa: FBT001 + expected: dict[str, list[str]], +) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) + + result_series = df["a"].str.replace_all(pattern=pattern, value=value, literal=literal) + compare_dicts({"a": result_series}, expected) + + +@pytest.mark.parametrize( + ("data", "pattern", "value", "n", "literal", "expected"), + replace_data, +) +def test_str_replace_expr( + constructor: Any, + data: dict[str, list[str]], + pattern: str, + value: str, + n: int, + literal: bool, # noqa: FBT001 + expected: dict[str, list[str]], +) -> None: + df = nw.from_native(constructor(data)) + + result_df = df.select( + nw.col("a").str.replace(pattern=pattern, value=value, n=n, literal=literal) + ) + compare_dicts(result_df, expected) + + +@pytest.mark.parametrize( + ("data", "pattern", "value", "literal", "expected"), + replace_all_data, +) +def test_str_replace_all_expr( + constructor: Any, + data: dict[str, list[str]], + pattern: str, + value: str, + literal: bool, # noqa: FBT001 + expected: dict[str, list[str]], +) -> None: + df = nw.from_native(constructor(data)) + + result = df.select( + nw.col("a").str.replace_all(pattern=pattern, value=value, literal=literal) + ) + compare_dicts(result, expected) diff --git a/tests/expr_and_series/str/slice_test.py b/tests/expr_and_series/str/slice_test.py index 9974c5fa9..e4e7905f2 100644 --- a/tests/expr_and_series/str/slice_test.py +++ b/tests/expr_and_series/str/slice_test.py @@ -15,11 +15,21 @@ [(1, 2, {"a": ["da", "df"]}), (-2, None, {"a": ["as", "as"]})], ) def test_str_slice( - constructor_eager: Any, offset: int, length: int | None, expected: Any + constructor: Any, offset: int, length: int | None, expected: Any ) -> None: - df = nw.from_native(constructor_eager(data), eager_only=True) + df = nw.from_native(constructor(data)) result_frame = df.select(nw.col("a").str.slice(offset, length)) compare_dicts(result_frame, expected) + +@pytest.mark.parametrize( + ("offset", "length", "expected"), + [(1, 2, {"a": ["da", "df"]}), (-2, None, {"a": ["as", "as"]})], +) +def test_str_slice_series( + constructor_eager: Any, offset: int, length: int | None, expected: Any +) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) + result_series = df["a"].str.slice(offset, length) - assert result_series.to_numpy().tolist() == expected["a"] + compare_dicts({"a": result_series}, expected) diff --git a/tests/expr_and_series/str/starts_with_ends_with_test.py b/tests/expr_and_series/str/starts_with_ends_with_test.py index 72f13fda3..a5101edcb 100644 --- a/tests/expr_and_series/str/starts_with_ends_with_test.py +++ b/tests/expr_and_series/str/starts_with_ends_with_test.py @@ -12,14 +12,17 @@ def test_ends_with(constructor: Any) -> None: - df = nw.from_native(constructor(data)).lazy() + df = nw.from_native(constructor(data)) result = df.select(nw.col("a").str.ends_with("das")) expected = { "a": [True, False], } compare_dicts(result, expected) - result = df.select(df.collect()["a"].str.ends_with("das")) + +def test_ends_with_series(constructor_eager: Any) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.select(df["a"].str.ends_with("das")) expected = { "a": [True, False], } @@ -34,7 +37,10 @@ def test_starts_with(constructor: Any) -> None: } compare_dicts(result, expected) - result = df.select(df.collect()["a"].str.starts_with("fda")) + +def test_starts_with_series(constructor_eager: Any) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.select(df["a"].str.starts_with("fda")) expected = { "a": [True, False], } diff --git a/tests/expr_and_series/str/strip_chars_test.py b/tests/expr_and_series/str/strip_chars_test.py new file mode 100644 index 000000000..f6cbcc4fa --- /dev/null +++ b/tests/expr_and_series/str/strip_chars_test.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + +data = {"a": ["foobar", "bar\n", " baz"]} + + +@pytest.mark.parametrize( + ("characters", "expected"), + [ + (None, {"a": ["foobar", "bar", "baz"]}), + ("foo", {"a": ["bar", "bar\n", " baz"]}), + ], +) +def test_str_strip_chars(constructor: Any, characters: str | None, expected: Any) -> None: + df = nw.from_native(constructor(data)) + result_frame = df.select(nw.col("a").str.strip_chars(characters)) + compare_dicts(result_frame, expected) + + +@pytest.mark.parametrize( + ("characters", "expected"), + [ + (None, {"a": ["foobar", "bar", "baz"]}), + ("foo", {"a": ["bar", "bar\n", " baz"]}), + ], +) +def test_str_strip_chars_series( + constructor_eager: Any, characters: str | None, expected: Any +) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) + + result_series = df["a"].str.strip_chars(characters) + compare_dicts({"a": result_series}, expected) diff --git a/tests/expr_and_series/str/tail_test.py b/tests/expr_and_series/str/tail_test.py index d0698825f..c863cca0e 100644 --- a/tests/expr_and_series/str/tail_test.py +++ b/tests/expr_and_series/str/tail_test.py @@ -6,12 +6,17 @@ data = {"a": ["foo", "bars"]} -def test_str_tail(constructor_eager: Any) -> None: - df = nw.from_native(constructor_eager(data), eager_only=True) +def test_str_tail(constructor: Any) -> None: + df = nw.from_native(constructor(data)) expected = {"a": ["foo", "ars"]} result_frame = df.select(nw.col("a").str.tail(3)) compare_dicts(result_frame, expected) + +def test_str_tail_series(constructor_eager: Any) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) + expected = {"a": ["foo", "ars"]} + result_series = df["a"].str.tail(3) - assert result_series.to_numpy().tolist() == expected["a"] + compare_dicts({"a": result_series}, expected) diff --git a/tests/expr_and_series/str/to_datetime_test.py b/tests/expr_and_series/str/to_datetime_test.py index 2d003f1ca..8c3d1a51a 100644 --- a/tests/expr_and_series/str/to_datetime_test.py +++ b/tests/expr_and_series/str/to_datetime_test.py @@ -5,10 +5,12 @@ data = {"a": ["2020-01-01T12:34:56"]} -def test_to_datetime(constructor_eager: Any) -> None: +def test_to_datetime(constructor: Any) -> None: result = ( - nw.from_native(constructor_eager(data), eager_only=True) + nw.from_native(constructor(data)) + .lazy() .select(b=nw.col("a").str.to_datetime(format="%Y-%m-%dT%H:%M:%S")) + .collect() .item(row=0, column="b") ) assert str(result) == "2020-01-01 12:34:56" diff --git a/tests/expr_and_series/str/to_uppercase_to_lowercase_test.py b/tests/expr_and_series/str/to_uppercase_to_lowercase_test.py index ae28eaf42..4d2f2f745 100644 --- a/tests/expr_and_series/str/to_uppercase_to_lowercase_test.py +++ b/tests/expr_and_series/str/to_uppercase_to_lowercase_test.py @@ -2,9 +2,11 @@ from typing import Any +import pyarrow as pa import pytest import narwhals.stable.v1 as nw +from narwhals.utils import parse_version from tests.utils import compare_dicts @@ -24,13 +26,53 @@ ], ) def test_str_to_uppercase( + constructor: Any, + data: dict[str, list[str]], + expected: dict[str, list[str]], + request: Any, +) -> None: + df = nw.from_native(constructor(data)) + result_frame = df.select(nw.col("a").str.to_uppercase()) + + if any("ß" in s for value in data.values() for s in value) & ( + constructor.__name__ + in ( + "pandas_pyarrow_constructor", + "pyarrow_table_constructor", + "modin_constructor", + ) + or ("dask" in str(constructor) and parse_version(pa.__version__) >= (12,)) + ): + # We are marking it xfail for these conditions above + # since the pyarrow backend will convert + # smaller cap 'ß' to upper cap 'ẞ' instead of 'SS' + request.applymarker(pytest.mark.xfail) + + compare_dicts(result_frame, expected) + + +@pytest.mark.parametrize( + ("data", "expected"), + [ + ({"a": ["foo", "bar"]}, {"a": ["FOO", "BAR"]}), + ( + { + "a": [ + "special case ß", + "ςpecial caσe", # noqa: RUF001 + ] + }, + {"a": ["SPECIAL CASE SS", "ΣPECIAL CAΣE"]}, + ), + ], +) +def test_str_to_uppercase_series( constructor_eager: Any, data: dict[str, list[str]], expected: dict[str, list[str]], request: Any, ) -> None: df = nw.from_native(constructor_eager(data), eager_only=True) - result_frame = df.select(nw.col("a").str.to_uppercase()) if any("ß" in s for value in data.values() for s in value) & ( constructor_eager.__name__ @@ -45,10 +87,8 @@ def test_str_to_uppercase( # smaller cap 'ß' to upper cap 'ẞ' instead of 'SS' request.applymarker(pytest.mark.xfail) - compare_dicts(result_frame, expected) - result_series = df["a"].str.to_uppercase() - assert result_series.to_numpy().tolist() == expected["a"] + compare_dicts({"a": result_series}, expected) @pytest.mark.parametrize( @@ -67,13 +107,36 @@ def test_str_to_uppercase( ], ) def test_str_to_lowercase( - constructor_eager: Any, + constructor: Any, data: dict[str, list[str]], expected: dict[str, list[str]], ) -> None: - df = nw.from_native(constructor_eager(data), eager_only=True) + df = nw.from_native(constructor(data)) result_frame = df.select(nw.col("a").str.to_lowercase()) compare_dicts(result_frame, expected) + +@pytest.mark.parametrize( + ("data", "expected"), + [ + ({"a": ["FOO", "BAR"]}, {"a": ["foo", "bar"]}), + ( + {"a": ["SPECIAL CASE ß", "ΣPECIAL CAΣE"]}, + { + "a": [ + "special case ß", + "σpecial caσe", # noqa: RUF001 + ] + }, + ), + ], +) +def test_str_to_lowercase_series( + constructor_eager: Any, + data: dict[str, list[str]], + expected: dict[str, list[str]], +) -> None: + df = nw.from_native(constructor_eager(data), eager_only=True) + result_series = df["a"].str.to_lowercase() - assert result_series.to_numpy().tolist() == expected["a"] + compare_dicts({"a": result_series}, expected) diff --git a/tests/expr_and_series/sum_all_test.py b/tests/expr_and_series/sum_all_test.py deleted file mode 100644 index 11c202947..000000000 --- a/tests/expr_and_series/sum_all_test.py +++ /dev/null @@ -1,20 +0,0 @@ -from typing import Any - -import narwhals.stable.v1 as nw -from tests.utils import compare_dicts - -data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} - - -def test_sum_all_expr(constructor: Any) -> None: - df = nw.from_native(constructor(data)) - result = df.select(nw.all().sum()) - expected = {"a": [6], "b": [14], "z": [24.0]} - compare_dicts(result, expected) - - -def test_sum_all_namespace(constructor: Any) -> None: - df = nw.from_native(constructor(data)) - result = df.select(nw.sum("a", "b", "z")) - expected = {"a": [6], "b": [14], "z": [24.0]} - compare_dicts(result, expected) diff --git a/tests/expr_and_series/sum_horizontal_test.py b/tests/expr_and_series/sum_horizontal_test.py index 9411903cb..4c4ab924c 100644 --- a/tests/expr_and_series/sum_horizontal_test.py +++ b/tests/expr_and_series/sum_horizontal_test.py @@ -18,3 +18,12 @@ def test_sumh(constructor: Any, col_expr: Any) -> None: "horizontal_sum": [5, 7, 8], } compare_dicts(result, expected) + + +def test_sumh_nullable(constructor: Any) -> None: + data = {"a": [1, 8, 3], "b": [4, 5, None]} + expected = {"hsum": [5, 13, 3]} + + df = nw.from_native(constructor(data)) + result = df.select(hsum=nw.sum_horizontal("a", "b")) + compare_dicts(result, expected) diff --git a/tests/expr_and_series/tail_test.py b/tests/expr_and_series/tail_test.py new file mode 100644 index 000000000..be17ffb4e --- /dev/null +++ b/tests/expr_and_series/tail_test.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +import narwhals as nw +from tests.utils import compare_dicts + + +@pytest.mark.parametrize("n", [2, -1]) +def test_head(constructor: Any, n: int, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) + if "polars" in str(constructor) and n < 0: + request.applymarker(pytest.mark.xfail) + df = nw.from_native(constructor({"a": [1, 2, 3]})) + result = df.select(nw.col("a").tail(n)) + expected = {"a": [2, 3]} + compare_dicts(result, expected) + + +@pytest.mark.parametrize("n", [2, -1]) +def test_head_series(constructor_eager: Any, n: int) -> None: + df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True) + result = df.select(df["a"].tail(n)) + expected = {"a": [2, 3]} + compare_dicts(result, expected) diff --git a/tests/expr_and_series/unary_test.py b/tests/expr_and_series/unary_test.py index b5149b4da..7df0099dd 100644 --- a/tests/expr_and_series/unary_test.py +++ b/tests/expr_and_series/unary_test.py @@ -1,10 +1,14 @@ from typing import Any -import narwhals as nw +import pytest + +import narwhals.stable.v1 as nw from tests.utils import compare_dicts -def test_unary(constructor: Any) -> None: +def test_unary(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} result = ( nw.from_native(constructor(data)) diff --git a/tests/expr_and_series/unique_test.py b/tests/expr_and_series/unique_test.py index ad437f8cc..488d793cd 100644 --- a/tests/expr_and_series/unique_test.py +++ b/tests/expr_and_series/unique_test.py @@ -1,6 +1,6 @@ from typing import Any -import numpy as np +import pytest import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -8,7 +8,9 @@ data = {"a": [1, 1, 2]} -def test_unique_expr(constructor: Any) -> None: +def test_unique_expr(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor(data)) result = df.select(nw.col("a").unique()) expected = {"a": [1, 2]} @@ -18,5 +20,5 @@ def test_unique_expr(constructor: Any) -> None: def test_unique_series(constructor_eager: Any) -> None: series = nw.from_native(constructor_eager(data), eager_only=True)["a"] result = series.unique() - expected = np.array([1, 2]) - assert (result.to_numpy() == expected).all() + expected = {"a": [1, 2]} + compare_dicts({"a": result}, expected) diff --git a/tests/expr_and_series/when_test.py b/tests/expr_and_series/when_test.py index a18095743..50d69f5f5 100644 --- a/tests/expr_and_series/when_test.py +++ b/tests/expr_and_series/when_test.py @@ -2,122 +2,187 @@ from typing import Any +import numpy as np import pytest import narwhals.stable.v1 as nw from tests.utils import compare_dicts data = { - "a": [1, 2, 3, 4, 5], - "b": ["a", "b", "c", "d", "e"], - "c": [4.1, 5.0, 6.0, 7.0, 8.0], - "d": [True, False, True, False, True], + "a": [1, 2, 3], + "b": ["a", "b", "c"], + "c": [4.1, 5.0, 6.0], + "d": [True, False, True], + "e": [7.0, 2.0, 1.1], } -def test_when(request: Any, constructor: Any) -> None: - if "pyarrow_table" in str(constructor): - request.applymarker(pytest.mark.xfail) - +def test_when(constructor: Any) -> None: df = nw.from_native(constructor(data)) - result = df.with_columns(nw.when(nw.col("a") == 1).then(value=3).alias("a_when")) + result = df.select(nw.when(nw.col("a") == 1).then(value=3).alias("a_when")) expected = { - "a": [1, 2, 3, 4, 5], - "b": ["a", "b", "c", "d", "e"], - "c": [4.1, 5.0, 6.0, 7.0, 8.0], - "d": [True, False, True, False, True], - "a_when": [3, None, None, None, None], + "a_when": [3, np.nan, np.nan], } compare_dicts(result, expected) -def test_when_otherwise(request: Any, constructor: Any) -> None: - if "pyarrow_table" in str(constructor): - request.applymarker(pytest.mark.xfail) +def test_when_otherwise(constructor: Any) -> None: + df = nw.from_native(constructor(data)) + result = df.select(nw.when(nw.col("a") == 1).then(3).otherwise(6).alias("a_when")) + expected = { + "a_when": [3, 6, 6], + } + compare_dicts(result, expected) + +def test_multiple_conditions(constructor: Any) -> None: df = nw.from_native(constructor(data)) - result = df.with_columns( - nw.when(nw.col("a") == 1).then(3).otherwise(6).alias("a_when") + result = df.select( + nw.when(nw.col("a") < 3, nw.col("c") < 5.0).then(3).alias("a_when") ) expected = { - "a": [1, 2, 3, 4, 5], - "b": ["a", "b", "c", "d", "e"], - "c": [4.1, 5.0, 6.0, 7.0, 8.0], - "d": [True, False, True, False, True], - "a_when": [3, 6, 6, 6, 6], + "a_when": [3, np.nan, np.nan], } compare_dicts(result, expected) -def test_chained_when(request: Any, constructor: Any) -> None: - if "pyarrow_table" in str(constructor): +def test_no_arg_when_fail(constructor: Any) -> None: + df = nw.from_native(constructor(data)) + with pytest.raises((TypeError, ValueError)): + df.select(nw.when().then(value=3).alias("a_when")) + + +def test_value_numpy_array(request: Any, constructor: Any) -> None: + if "dask" in str(constructor): request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor(data)) - result = df.with_columns( - nw.when(nw.col("a") == 1) - .then(3) - .when(nw.col("a") == 2) - .then(5) - .otherwise(7) - .alias("a_when"), + import numpy as np + + result = df.select( + nw.when(nw.col("a") == 1).then(np.asanyarray([3, 4, 5])).alias("a_when") ) expected = { - "a": [1, 2, 3, 4, 5], - "b": ["a", "b", "c", "d", "e"], - "c": [4.1, 5.0, 6.0, 7.0, 8.0], - "d": [True, False, True, False, True], - "a_when": [3, 5, 7, 7, 7], + "a_when": [3, np.nan, np.nan], + } + compare_dicts(result, expected) + + +def test_value_series(constructor_eager: Any) -> None: + df = nw.from_native(constructor_eager(data)) + s_data = {"s": [3, 4, 5]} + s = nw.from_native(constructor_eager(s_data))["s"] + assert isinstance(s, nw.Series) + result = df.select(nw.when(nw.col("a") == 1).then(s).alias("a_when")) + expected = { + "a_when": [3, np.nan, np.nan], } compare_dicts(result, expected) -def test_when_with_multiple_conditions(request: Any, constructor: Any) -> None: - if "pyarrow_table" in str(constructor): +def test_value_expression(constructor: Any) -> None: + df = nw.from_native(constructor(data)) + result = df.select(nw.when(nw.col("a") == 1).then(nw.col("a") + 9).alias("a_when")) + expected = { + "a_when": [10, np.nan, np.nan], + } + compare_dicts(result, expected) + + +def test_otherwise_numpy_array(request: Any, constructor: Any) -> None: + if "dask" in str(constructor): request.applymarker(pytest.mark.xfail) + df = nw.from_native(constructor(data)) - result = df.with_columns( - nw.when(nw.col("a") == 1) - .then(3) - .when(nw.col("a") == 2) - .then(5) - .when(nw.col("a") == 3) - .then(7) - .otherwise(9) - .alias("a_when"), + import numpy as np + + result = df.select( + nw.when(nw.col("a") == 1).then(-1).otherwise(np.array([0, 9, 10])).alias("a_when") ) expected = { - "a": [1, 2, 3, 4, 5], - "b": ["a", "b", "c", "d", "e"], - "c": [4.1, 5.0, 6.0, 7.0, 8.0], - "d": [True, False, True, False, True], - "a_when": [3, 5, 7, 9, 9], + "a_when": [-1, 9, 10], } compare_dicts(result, expected) -def test_multiple_conditions(request: Any, constructor: Any) -> None: - if "pyarrow_table" in str(constructor): +def test_otherwise_series(constructor_eager: Any) -> None: + df = nw.from_native(constructor_eager(data)) + s_data = {"s": [0, 9, 10]} + s = nw.from_native(constructor_eager(s_data))["s"] + assert isinstance(s, nw.Series) + result = df.select(nw.when(nw.col("a") == 1).then(-1).otherwise(s).alias("a_when")) + expected = { + "a_when": [-1, 9, 10], + } + compare_dicts(result, expected) + + +def test_otherwise_expression(request: Any, constructor: Any) -> None: + if "dask" in str(constructor): request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor(data)) - result = df.with_columns( - nw.when(nw.col("a") < 3, nw.col("c") < 5.0).then(3).alias("a_when") + result = df.select( + nw.when(nw.col("a") == 1).then(-1).otherwise(nw.col("a") + 7).alias("a_when") ) expected = { - "a": [1, 2, 3, 4, 5], - "b": ["a", "b", "c", "d", "e"], - "c": [4.1, 5.0, 6.0, 7.0, 8.0], - "d": [True, False, True, False, True], - "a_when": [3, None, None, None, None], + "a_when": [-1, 9, 10], } compare_dicts(result, expected) -def test_no_arg_when_fail(request: Any, constructor: Any) -> None: - if "pyarrow_table" in str(constructor): +def test_when_then_otherwise_into_expr(request: Any, constructor: Any) -> None: + if "dask" in str(constructor): request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor(data)) - with pytest.raises(TypeError): - df.with_columns(nw.when().then(value=3).alias("a_when")) + result = df.select(nw.when(nw.col("a") > 1).then("c").otherwise("e")) + expected = {"c": [7, 5, 6]} + compare_dicts(result, expected) + + +# def test_chained_when(request: Any, constructor: Any) -> None: +# if "pyarrow_table" in str(constructor): +# request.applymarker(pytest.mark.xfail) + +# df = nw.from_native(constructor(data)) +# result = df.with_columns( +# nw.when(nw.col("a") == 1) +# .then(3) +# .when(nw.col("a") == 2) +# .then(5) +# .otherwise(7) +# .alias("a_when"), +# ) +# expected = { +# "a": [1, 2, 3, 4, 5], +# "b": ["a", "b", "c", "d", "e"], +# "c": [4.1, 5.0, 6.0, 7.0, 8.0], +# "d": [True, False, True, False, True], +# "a_when": [3, 5, 7, 7, 7], +# } +# compare_dicts(result, expected) + + +# def test_when_with_multiple_conditions(request: Any, constructor: Any) -> None: +# if "pyarrow_table" in str(constructor): +# request.applymarker(pytest.mark.xfail) +# df = nw.from_native(constructor(data)) +# result = df.with_columns( +# nw.when(nw.col("a") == 1) +# .then(3) +# .when(nw.col("a") == 2) +# .then(5) +# .when(nw.col("a") == 3) +# .then(7) +# .otherwise(9) +# .alias("a_when"), +# ) +# expected = { +# "a": [1, 2, 3, 4, 5], +# "b": ["a", "b", "c", "d", "e"], +# "c": [4.1, 5.0, 6.0, 7.0, 8.0], +# "d": [True, False, True, False, True], +# "a_when": [3, 5, 7, 9, 9], +# } +# compare_dicts(result, expected) diff --git a/tests/frame/array_dunder_test.py b/tests/frame/array_dunder_test.py new file mode 100644 index 000000000..8f9cbd16a --- /dev/null +++ b/tests/frame/array_dunder_test.py @@ -0,0 +1,61 @@ +from typing import Any + +import numpy as np +import pandas as pd +import polars as pl +import pyarrow as pa +import pytest + +import narwhals.stable.v1 as nw +from narwhals.utils import parse_version +from tests.utils import compare_dicts + + +def test_array_dunder(request: Any, constructor_eager: Any) -> None: + if "pyarrow_table" in str(constructor_eager) and parse_version( + pa.__version__ + ) < parse_version("16.0.0"): # pragma: no cover + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True) + result = df.__array__() + np.testing.assert_array_equal(result, np.array([[1], [2], [3]], dtype="int64")) + + +def test_array_dunder_with_dtype(request: Any, constructor_eager: Any) -> None: + if "pyarrow_table" in str(constructor_eager) and parse_version( + pa.__version__ + ) < parse_version("16.0.0"): # pragma: no cover + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True) + result = df.__array__(object) + np.testing.assert_array_equal(result, np.array([[1], [2], [3]], dtype=object)) + + +def test_array_dunder_with_copy(request: Any, constructor_eager: Any) -> None: + if "pyarrow_table" in str(constructor_eager) and parse_version(pa.__version__) < ( + 16, + 0, + 0, + ): # pragma: no cover + request.applymarker(pytest.mark.xfail) + if "polars" in str(constructor_eager) and parse_version(pl.__version__) < ( + 0, + 20, + 28, + ): # pragma: no cover + request.applymarker(pytest.mark.xfail) + + df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True) + result = df.__array__(copy=True) + np.testing.assert_array_equal(result, np.array([[1], [2], [3]], dtype="int64")) + if "pandas_constructor" in str(constructor_eager) and parse_version( + pd.__version__ + ) < (3,): + # If it's pandas, we know that `copy=False` definitely took effect. + # So, let's check it! + result = df.__array__(copy=False) + np.testing.assert_array_equal(result, np.array([[1], [2], [3]], dtype="int64")) + result[0, 0] = 999 + compare_dicts(df, {"a": [999, 2, 3]}) diff --git a/tests/frame/arrow_c_stream_test.py b/tests/frame/arrow_c_stream_test.py new file mode 100644 index 000000000..7a3403f69 --- /dev/null +++ b/tests/frame/arrow_c_stream_test.py @@ -0,0 +1,42 @@ +import polars as pl +import pyarrow as pa +import pyarrow.compute as pc +import pytest + +import narwhals.stable.v1 as nw +from narwhals.utils import parse_version + + +@pytest.mark.skipif( + parse_version(pl.__version__) < (1, 3), reason="too old for pycapsule in Polars" +) +def test_arrow_c_stream_test() -> None: + df = nw.from_native(pl.Series([1, 2, 3]).to_frame("a"), eager_only=True) + result = pa.table(df) + expected = pa.table({"a": [1, 2, 3]}) + assert pc.all(pc.equal(result["a"], expected["a"])).as_py() + + +@pytest.mark.skipif( + parse_version(pl.__version__) < (1, 3), reason="too old for pycapsule in Polars" +) +def test_arrow_c_stream_test_invalid(monkeypatch: pytest.MonkeyPatch) -> None: + # "poison" the dunder method to make sure it actually got called above + monkeypatch.setattr( + "narwhals.dataframe.DataFrame.__arrow_c_stream__", lambda *_: 1 / 0 + ) + df = nw.from_native(pl.Series([1, 2, 3]).to_frame("a"), eager_only=True) + with pytest.raises(ZeroDivisionError, match="division by zero"): + pa.table(df) + + +@pytest.mark.skipif( + parse_version(pl.__version__) < (1, 3), reason="too old for pycapsule in Polars" +) +def test_arrow_c_stream_test_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + # Check that fallback to PyArrow works + monkeypatch.delattr("polars.DataFrame.__arrow_c_stream__") + df = nw.from_native(pl.Series([1, 2, 3]).to_frame("a"), eager_only=True) + result = pa.table(df) + expected = pa.table({"a": [1, 2, 3]}) + assert pc.all(pc.equal(result["a"], expected["a"])).as_py() diff --git a/tests/frame/clone_test.py b/tests/frame/clone_test.py index 9d5b18caf..6e8b19beb 100644 --- a/tests/frame/clone_test.py +++ b/tests/frame/clone_test.py @@ -7,6 +7,8 @@ def test_clone(request: Any, constructor: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) if "pyarrow_table" in str(constructor): request.applymarker(pytest.mark.xfail) diff --git a/tests/frame/columns_test.py b/tests/frame/columns_test.py new file mode 100644 index 000000000..157051ba3 --- /dev/null +++ b/tests/frame/columns_test.py @@ -0,0 +1,14 @@ +from typing import Any + +import pytest + +import narwhals.stable.v1 as nw + + +@pytest.mark.filterwarnings("ignore:Determining|Resolving.*") +def test_columns(constructor: Any) -> None: + data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} + df = nw.from_native(constructor(data)) + result = df.columns + expected = ["a", "b", "z"] + assert result == expected diff --git a/tests/frame/concat_test.py b/tests/frame/concat_test.py index 94b8d7de1..970220bf2 100644 --- a/tests/frame/concat_test.py +++ b/tests/frame/concat_test.py @@ -6,7 +6,9 @@ from tests.utils import compare_dicts -def test_concat_horizontal(constructor: Any) -> None: +def test_concat_horizontal(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} df_left = nw.from_native(constructor(data)) @@ -27,7 +29,9 @@ def test_concat_horizontal(constructor: Any) -> None: nw.concat([]) -def test_concat_vertical(constructor: Any) -> None: +def test_concat_vertical(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} df_left = ( nw.from_native(constructor(data)).rename({"a": "c", "b": "d"}).drop("z").lazy() diff --git a/tests/frame/double_test.py b/tests/frame/double_test.py index a96f1d1be..5d52d0d26 100644 --- a/tests/frame/double_test.py +++ b/tests/frame/double_test.py @@ -13,5 +13,10 @@ def test_double(constructor: Any) -> None: compare_dicts(result, expected) result = df.with_columns(nw.col("a").alias("o"), nw.all() * 2) - expected = {"o": [1, 3, 2], "a": [2, 6, 4], "b": [8, 8, 12], "z": [14.0, 16.0, 18.0]} + expected = { + "o": [1, 3, 2], + "a": [2, 6, 4], + "b": [8, 8, 12], + "z": [14.0, 16.0, 18.0], + } compare_dicts(result, expected) diff --git a/tests/frame/drop_nulls_test.py b/tests/frame/drop_nulls_test.py index 53cc57bc2..58c9486ed 100644 --- a/tests/frame/drop_nulls_test.py +++ b/tests/frame/drop_nulls_test.py @@ -1,5 +1,9 @@ +from __future__ import annotations + from typing import Any +import pytest + import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -16,3 +20,13 @@ def test_drop_nulls(constructor: Any) -> None: "b": [3.0, 5.0], } compare_dicts(result, expected) + + +@pytest.mark.parametrize("subset", ["a", ["a"]]) +def test_drop_nulls_subset(constructor: Any, subset: str | list[str]) -> None: + result = nw.from_native(constructor(data)).drop_nulls(subset=subset) + expected = { + "a": [1, 2.0, 4.0], + "b": [float("nan"), 3.0, 5.0], + } + compare_dicts(result, expected) diff --git a/tests/frame/drop_test.py b/tests/frame/drop_test.py index f22d148b1..db039fcb2 100644 --- a/tests/frame/drop_test.py +++ b/tests/frame/drop_test.py @@ -1,21 +1,56 @@ from __future__ import annotations +from contextlib import nullcontext as does_not_raise from typing import Any +import polars as pl import pytest +from polars.exceptions import ColumnNotFoundError as PlColumnNotFoundError import narwhals.stable.v1 as nw +from narwhals._exceptions import ColumnNotFoundError +from narwhals.utils import parse_version @pytest.mark.parametrize( - ("drop", "left"), + ("to_drop", "expected"), [ - (["a"], ["b", "z"]), - (["a", "b"], ["z"]), + ("abc", ["b", "z"]), + (["abc"], ["b", "z"]), + (["abc", "b"], ["z"]), ], ) -def test_drop(constructor: Any, drop: list[str], left: list[str]) -> None: - data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} +def test_drop(constructor: Any, to_drop: list[str], expected: list[str]) -> None: + data = {"abc": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} df = nw.from_native(constructor(data)) - assert df.drop(drop).collect_schema().names() == left - assert df.drop(*drop).collect_schema().names() == left + assert df.drop(to_drop).collect_schema().names() == expected + if not isinstance(to_drop, str): + assert df.drop(*to_drop).collect_schema().names() == expected + + +@pytest.mark.parametrize( + ("strict", "context"), + [ + ( + True, + pytest.raises((ColumnNotFoundError, PlColumnNotFoundError), match="z"), + ), + (False, does_not_raise()), + ], +) +def test_drop_strict(request: Any, constructor: Any, strict: bool, context: Any) -> None: # noqa: FBT001 + if ( + "polars_lazy" in str(request) + and parse_version(pl.__version__) < (1, 0, 0) + and strict + ): + request.applymarker(pytest.mark.xfail) + + data = {"a": [1, 3, 2], "b": [4, 4, 6]} + to_drop = ["a", "z"] + + df = nw.from_native(constructor(data)) + + with context: + names_out = df.drop(to_drop, strict=strict).collect_schema().names() + assert names_out == ["b"] diff --git a/tests/frame/filter_test.py b/tests/frame/filter_test.py index 9b7ed45d2..a8d3144aa 100644 --- a/tests/frame/filter_test.py +++ b/tests/frame/filter_test.py @@ -12,11 +12,9 @@ def test_filter(constructor: Any) -> None: compare_dicts(result, expected) -def test_filter_series(constructor_eager: Any) -> None: +def test_filter_with_boolean_list(constructor: Any) -> None: data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} - df = nw.from_native(constructor_eager(data), eager_only=True).with_columns( - mask=nw.col("a") > 1 - ) - result = df.filter(df["mask"]).drop("mask") + df = nw.from_native(constructor(data)) + result = df.filter([False, True, True]) expected = {"a": [3, 2], "b": [4, 6], "z": [8.0, 9.0]} compare_dicts(result, expected) diff --git a/tests/frame/get_column_test.py b/tests/frame/get_column_test.py index 3478e156e..58766ac31 100644 --- a/tests/frame/get_column_test.py +++ b/tests/frame/get_column_test.py @@ -3,13 +3,14 @@ import pandas as pd import pytest -import narwhals as nw +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts def test_get_column(constructor_eager: Any) -> None: df = nw.from_native(constructor_eager({"a": [1, 2], "b": [3, 4]}), eager_only=True) result = df.get_column("a") - assert result.to_list() == [1, 2] + compare_dicts({"a": result}, {"a": [1, 2]}) assert result.name == "a" with pytest.raises( (KeyError, TypeError), match="Expected str|'int' object cannot be converted|0" @@ -21,8 +22,11 @@ def test_get_column(constructor_eager: Any) -> None: def test_non_string_name() -> None: df = pd.DataFrame({0: [1, 2]}) result = nw.from_native(df, eager_only=True).get_column(0) # type: ignore[arg-type] - assert result.to_list() == [1, 2] + compare_dicts({"a": result}, {"a": [1, 2]}) assert result.name == 0 # type: ignore[comparison-overlap] - with pytest.raises(TypeError, match="Expected str or slice"): - # Check that getitem would have raised - nw.from_native(df, eager_only=True)[0] # type: ignore[call-overload] + + +def test_get_single_row() -> None: + df = pd.DataFrame({"a": [1, 2], "b": [3, 4]}) + result = nw.from_native(df, eager_only=True)[0] # type: ignore[call-overload] + compare_dicts(result, {"a": [1], "b": [3]}) diff --git a/tests/frame/head_test.py b/tests/frame/head_test.py index f508a9c5e..e4b762f48 100644 --- a/tests/frame/head_test.py +++ b/tests/frame/head_test.py @@ -22,6 +22,3 @@ def test_head(constructor: Any) -> None: # negative indices not allowed for lazyframes result = df.lazy().collect().head(-1) compare_dicts(result, expected) - - result = df.select(nw.col("a").head(2)) - compare_dicts(result, {"a": expected["a"]}) diff --git a/tests/frame/invalid_test.py b/tests/frame/invalid_test.py new file mode 100644 index 000000000..41c780a7d --- /dev/null +++ b/tests/frame/invalid_test.py @@ -0,0 +1,20 @@ +import pandas as pd +import polars as pl +import pyarrow as pa +import pytest + +import narwhals.stable.v1 as nw + + +def test_invalid() -> None: + data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} + df = nw.from_native(pa.table({"a": [1, 2], "b": [3, 4]})) + with pytest.raises(ValueError, match="Multi-output"): + df.select(nw.all() + nw.all()) + df = nw.from_native(pd.DataFrame(data)) + with pytest.raises(ValueError, match="Multi-output"): + df.select(nw.all() + nw.all()) + with pytest.raises(TypeError, match="Perhaps you:"): + df.select([pl.col("a")]) # type: ignore[list-item] + with pytest.raises(TypeError, match="Perhaps you:"): + df.select([nw.col("a").cast(pl.Int64)]) diff --git a/tests/frame/is_duplicated_test.py b/tests/frame/is_duplicated_test.py index 7fa4bfb09..e1eb3f298 100644 --- a/tests/frame/is_duplicated_test.py +++ b/tests/frame/is_duplicated_test.py @@ -2,9 +2,8 @@ from typing import Any -import numpy as np - import narwhals.stable.v1 as nw +from tests.utils import compare_dicts def test_is_duplicated(constructor_eager: Any) -> None: @@ -12,5 +11,5 @@ def test_is_duplicated(constructor_eager: Any) -> None: df_raw = constructor_eager(data) df = nw.from_native(df_raw, eager_only=True) result = nw.concat([df, df.head(1)]).is_duplicated() - expected = np.array([True, False, False, True]) - assert (result.to_numpy() == expected).all() + expected = {"is_duplicated": [True, False, False, True]} + compare_dicts({"is_duplicated": result}, expected) diff --git a/tests/frame/is_unique_test.py b/tests/frame/is_unique_test.py index 7b36dce00..4259c8773 100644 --- a/tests/frame/is_unique_test.py +++ b/tests/frame/is_unique_test.py @@ -2,9 +2,8 @@ from typing import Any -import numpy as np - import narwhals.stable.v1 as nw +from tests.utils import compare_dicts def test_is_unique(constructor_eager: Any) -> None: @@ -12,5 +11,5 @@ def test_is_unique(constructor_eager: Any) -> None: df_raw = constructor_eager(data) df = nw.from_native(df_raw, eager_only=True) result = nw.concat([df, df.head(1)]).is_unique() - expected = np.array([False, True, True, False]) - assert (result.to_numpy() == expected).all() + expected = {"is_unique": [False, True, True, False]} + compare_dicts({"is_unique": result}, expected) diff --git a/tests/frame/join_test.py b/tests/frame/join_test.py index db4a25685..d5c88ee4c 100644 --- a/tests/frame/join_test.py +++ b/tests/frame/join_test.py @@ -15,9 +15,7 @@ def test_inner_join_two_keys(constructor: Any) -> None: data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} df = nw.from_native(constructor(data)) df_right = df - result = df.lazy().join( - df_right.lazy(), left_on=["a", "b"], right_on=["a", "b"], how="inner" - ) + result = df.join(df_right, left_on=["a", "b"], right_on=["a", "b"], how="inner") # type: ignore[arg-type] expected = { "a": [1, 3, 2], "b": [4, 4, 6], @@ -27,11 +25,11 @@ def test_inner_join_two_keys(constructor: Any) -> None: compare_dicts(result, expected) -def test_inner_join_single_key(constructor_eager: Any) -> None: +def test_inner_join_single_key(constructor: Any) -> None: data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} - df = nw.from_native(constructor_eager(data), eager_only=True) + df = nw.from_native(constructor(data)) df_right = df - result = df.join(df_right, left_on="a", right_on="a", how="inner") + result = df.join(df_right, left_on="a", right_on="a", how="inner") # type: ignore[arg-type] expected = { "a": [1, 3, 2], "b": [4, 4, 6], @@ -78,13 +76,13 @@ def test_cross_join_non_pandas() -> None: ], ) def test_anti_join( - constructor_eager: Any, + constructor: Any, join_key: list[str], filter_expr: nw.Expr, expected: dict[str, list[Any]], ) -> None: data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} - df = nw.from_native(constructor_eager(data)) + df = nw.from_native(constructor(data)) other = df.filter(filter_expr) result = df.join(other, how="anti", left_on=join_key, right_on=join_key) # type: ignore[arg-type] compare_dicts(result, expected) @@ -99,13 +97,13 @@ def test_anti_join( ], ) def test_semi_join( - constructor_eager: Any, + constructor: Any, join_key: list[str], filter_expr: nw.Expr, expected: dict[str, list[Any]], ) -> None: data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} - df = nw.from_native(constructor_eager(data)) + df = nw.from_native(constructor(data)) other = df.filter(filter_expr) result = df.join(other, how="semi", left_on=join_key, right_on=join_key) # type: ignore[arg-type] compare_dicts(result, expected) @@ -126,12 +124,12 @@ def test_join_not_implemented(constructor: Any, how: str) -> None: @pytest.mark.filterwarnings("ignore:the default coalesce behavior") -def test_left_join(constructor_eager: Any) -> None: +def test_left_join(constructor: Any) -> None: data_left = {"a": [1.0, 2, 3], "b": [4.0, 5, 6]} data_right = {"a": [1.0, 2, 3], "c": [4.0, 5, 7]} - df_left = nw.from_native(constructor_eager(data_left), eager_only=True) - df_right = nw.from_native(constructor_eager(data_right), eager_only=True) - result = df_left.join(df_right, left_on="b", right_on="c", how="left").select( + df_left = nw.from_native(constructor(data_left)) + df_right = nw.from_native(constructor(data_right)) + result = df_left.join(df_right, left_on="b", right_on="c", how="left").select( # type: ignore[arg-type] nw.all().fill_null(float("nan")) ) expected = {"a": [1, 2, 3], "b": [4, 5, 6], "a_right": [1, 2, float("nan")]} @@ -139,23 +137,23 @@ def test_left_join(constructor_eager: Any) -> None: @pytest.mark.filterwarnings("ignore: the default coalesce behavior") -def test_left_join_multiple_column(constructor_eager: Any) -> None: +def test_left_join_multiple_column(constructor: Any) -> None: data_left = {"a": [1, 2, 3], "b": [4, 5, 6]} data_right = {"a": [1, 2, 3], "c": [4, 5, 6]} - df_left = nw.from_native(constructor_eager(data_left), eager_only=True) - df_right = nw.from_native(constructor_eager(data_right), eager_only=True) - result = df_left.join(df_right, left_on=["a", "b"], right_on=["a", "c"], how="left") + df_left = nw.from_native(constructor(data_left)) + df_right = nw.from_native(constructor(data_right)) + result = df_left.join(df_right, left_on=["a", "b"], right_on=["a", "c"], how="left") # type: ignore[arg-type] expected = {"a": [1, 2, 3], "b": [4, 5, 6]} compare_dicts(result, expected) @pytest.mark.filterwarnings("ignore: the default coalesce behavior") -def test_left_join_overlapping_column(constructor_eager: Any) -> None: +def test_left_join_overlapping_column(constructor: Any) -> None: data_left = {"a": [1.0, 2, 3], "b": [4.0, 5, 6], "d": [1.0, 4, 2]} data_right = {"a": [1.0, 2, 3], "c": [4.0, 5, 6], "d": [1.0, 4, 2]} - df_left = nw.from_native(constructor_eager(data_left), eager_only=True) - df_right = nw.from_native(constructor_eager(data_right), eager_only=True) - result = df_left.join(df_right, left_on="b", right_on="c", how="left") + df_left = nw.from_native(constructor(data_left)) + df_right = nw.from_native(constructor(data_right)) + result = df_left.join(df_right, left_on="b", right_on="c", how="left") # type: ignore[arg-type] expected: dict[str, list[Any]] = { "a": [1, 2, 3], "b": [4, 5, 6], @@ -164,7 +162,7 @@ def test_left_join_overlapping_column(constructor_eager: Any) -> None: "d_right": [1, 4, 2], } compare_dicts(result, expected) - result = df_left.join(df_right, left_on="a", right_on="d", how="left").select( + result = df_left.join(df_right, left_on="a", right_on="d", how="left").select( # type: ignore[arg-type] nw.all().fill_null(float("nan")) ) expected = { diff --git a/tests/frame/lazy_test.py b/tests/frame/lazy_test.py new file mode 100644 index 000000000..14b9d1488 --- /dev/null +++ b/tests/frame/lazy_test.py @@ -0,0 +1,9 @@ +from typing import Any + +import narwhals.stable.v1 as nw + + +def test_lazy(constructor_eager: Any) -> None: + df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True) + result = df.lazy() + assert isinstance(result, nw.LazyFrame) diff --git a/tests/frame/lit_test.py b/tests/frame/lit_test.py index e5756e035..328e4d8e0 100644 --- a/tests/frame/lit_test.py +++ b/tests/frame/lit_test.py @@ -17,7 +17,11 @@ ("dtype", "expected_lit"), [(None, [2, 2, 2]), (nw.String, ["2", "2", "2"]), (nw.Float32, [2.0, 2.0, 2.0])], ) -def test_lit(constructor: Any, dtype: DType | None, expected_lit: list[Any]) -> None: +def test_lit( + constructor: Any, dtype: DType | None, expected_lit: list[Any], request: Any +) -> None: + if "dask" in str(constructor) and dtype == nw.String: + request.applymarker(pytest.mark.xfail) data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} df_raw = constructor(data) df = nw.from_native(df_raw).lazy() diff --git a/tests/frame/reindex_test.py b/tests/frame/reindex_test.py new file mode 100644 index 000000000..e21b31a8e --- /dev/null +++ b/tests/frame/reindex_test.py @@ -0,0 +1,30 @@ +from typing import Any + +import pandas as pd +import pytest + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + +data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} + + +@pytest.mark.parametrize("df_raw", [pd.DataFrame(data)]) +def test_reindex(df_raw: Any) -> None: + df = nw.from_native(df_raw, eager_only=True) + result = df.select("b", df["a"].sort(descending=True)) + expected = {"b": [4, 4, 6], "a": [3, 2, 1]} + compare_dicts(result, expected) + result = df.select("b", nw.col("a").sort(descending=True)) + compare_dicts(result, expected) + + s = df["a"] + result_s = s > s.sort() + assert not result_s[0] + assert result_s[1] + assert not result_s[2] + result = df.with_columns(s.sort()) + expected = {"a": [1, 2, 3], "b": [4, 4, 6], "z": [7.0, 8.0, 9.0]} # type: ignore[list-item] + compare_dicts(result, expected) + with pytest.raises(ValueError, match="Multi-output expressions are not supported"): + nw.to_native(df.with_columns(nw.all() + nw.all())) diff --git a/tests/frame/row_test.py b/tests/frame/row_test.py new file mode 100644 index 000000000..602c50f55 --- /dev/null +++ b/tests/frame/row_test.py @@ -0,0 +1,14 @@ +from typing import Any + +import narwhals.stable.v1 as nw + + +def test_row_column(constructor_eager: Any) -> None: + data = { + "a": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + "b": [11, 12, 13, 14, 15, 16], + } + result = nw.from_native(constructor_eager(data), eager_only=True).row(2) + if "pyarrow_table" in str(constructor_eager): + result = tuple(x.as_py() for x in result) + assert result == (3.0, 13) diff --git a/tests/frame/schema_test.py b/tests/frame/schema_test.py index 351a9b293..6e6b33aa1 100644 --- a/tests/frame/schema_test.py +++ b/tests/frame/schema_test.py @@ -45,7 +45,9 @@ def test_schema_comparison() -> None: def test_object() -> None: - df = pd.DataFrame({"a": [1, 2, 3]}).astype(object) + class Foo: ... + + df = pd.DataFrame({"a": [Foo()]}).astype(object) result = nw.from_native(df).schema assert result["a"] == nw.Object @@ -57,7 +59,7 @@ def test_string_disguised_as_object() -> None: def test_actual_object(request: Any, constructor_eager: Any) -> None: - if "pyarrow_table" in str(constructor_eager): + if any(x in str(constructor_eager) for x in ("modin", "pyarrow_table")): request.applymarker(pytest.mark.xfail) class Foo: ... @@ -168,3 +170,31 @@ def test_unknown_dtype_polars() -> None: def test_hash() -> None: assert nw.Int64() in {nw.Int64, nw.Int32} + + +@pytest.mark.parametrize( + ("method", "expected"), + [ + ("names", ["a", "b", "c"]), + ("dtypes", [nw.Int64(), nw.Float32(), nw.String()]), + ("len", 3), + ], +) +def test_schema_object(method: str, expected: Any) -> None: + data = {"a": nw.Int64(), "b": nw.Float32(), "c": nw.String()} + schema = nw.Schema(data) + assert getattr(schema, method)() == expected + + +@pytest.mark.skipif( + parse_version(pd.__version__) < (2,), + reason="Before 2.0, pandas would raise on `drop_duplicates`", +) +def test_from_non_hashable_column_name() -> None: + # This is technically super-illegal + # BUT, it shows up in a scikit-learn test, so... + df = pd.DataFrame([[1, 2], [3, 4]], columns=["pizza", ["a", "b"]]) + + df = nw.from_native(df, eager_only=True) + assert df.columns == ["pizza", ["a", "b"]] + assert df["pizza"].dtype == nw.Int64 diff --git a/tests/frame/select_test.py b/tests/frame/select_test.py index 6891ec4f1..450e91066 100644 --- a/tests/frame/select_test.py +++ b/tests/frame/select_test.py @@ -1,5 +1,8 @@ from typing import Any +import pandas as pd +import pytest + import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -12,6 +15,19 @@ def test_select(constructor: Any) -> None: compare_dicts(result, expected) -def test_empty_select(constructor_eager: Any) -> None: - result = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True).select() - assert result.shape == (0, 0) +def test_empty_select(constructor: Any) -> None: + result = nw.from_native(constructor({"a": [1, 2, 3]})).lazy().select() + assert result.collect().shape == (0, 0) + + +def test_non_string_select() -> None: + df = nw.from_native(pd.DataFrame({0: [1, 2], "b": [3, 4]})) + result = nw.to_native(df.select(nw.col(0))) # type: ignore[arg-type] + expected = pd.Series([1, 2], name=0).to_frame() + pd.testing.assert_frame_equal(result, expected) + + +def test_non_string_select_invalid() -> None: + df = nw.from_native(pd.DataFrame({0: [1, 2], "b": [3, 4]})) + with pytest.raises(TypeError, match="\n\nHint: if you were trying to select"): + nw.to_native(df.select(0)) # type: ignore[arg-type] diff --git a/tests/frame/series_sum_test.py b/tests/frame/series_sum_test.py deleted file mode 100644 index 2eb6cdb1b..000000000 --- a/tests/frame/series_sum_test.py +++ /dev/null @@ -1,21 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import narwhals.stable.v1 as nw -from tests.utils import compare_dicts - - -def test_series_sum(constructor: Any) -> None: - data = { - "a": [0, 1, 2, 3, 4], - "b": [1, 2, 3, 5, 3], - "c": [5, 4, None, 2, 1], - } - df = nw.from_native(constructor(data), strict=False, allow_series=True) - - result = df.select(nw.col("a", "b", "c").sum()) - - expected_sum = {"a": [10], "b": [14], "c": [12]} - - compare_dicts(result, expected_sum) diff --git a/tests/frame/slice_test.py b/tests/frame/slice_test.py index f381c174f..eea94d440 100644 --- a/tests/frame/slice_test.py +++ b/tests/frame/slice_test.py @@ -5,10 +5,8 @@ import polars as pl import pyarrow as pa import pytest -from pandas.testing import assert_series_equal import narwhals.stable.v1 as nw -from narwhals.utils import parse_version from tests.utils import compare_dicts data = { @@ -20,7 +18,7 @@ def test_slice_column(constructor_eager: Any) -> None: result = nw.from_native(constructor_eager(data))["a"] assert isinstance(result, nw.Series) - assert result.to_numpy().tolist() == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + compare_dicts({"a": result}, {"a": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]}) def test_slice_rows(constructor_eager: Any) -> None: @@ -51,9 +49,16 @@ def test_slice_lazy_fails() -> None: _ = nw.from_native(pl.LazyFrame(data))[1:] -def test_slice_int_fails(constructor_eager: Any) -> None: - with pytest.raises(TypeError, match="Expected str or slice, got: "): - _ = nw.from_native(constructor_eager(data))[1] # type: ignore[call-overload,index] +def test_slice_int(constructor_eager: Any) -> None: + result = nw.from_native(constructor_eager(data), eager_only=True)[1] # type: ignore[call-overload] + compare_dicts(result, {"a": [2], "b": [12]}) + + +def test_slice_fails(constructor_eager: Any) -> None: + class Foo: ... + + with pytest.raises(TypeError, match="Expected str or slice, got:"): + nw.from_native(constructor_eager(data), eager_only=True)[Foo()] # type: ignore[call-overload] def test_gather(constructor_eager: Any) -> None: @@ -83,21 +88,14 @@ def test_gather_pandas_index() -> None: def test_gather_rows_cols(constructor_eager: Any) -> None: native_df = constructor_eager(data) df = nw.from_native(native_df, eager_only=True) - is_pandas_wo_pyarrow = parse_version(pd.__version__) < parse_version("1.0.0") - if isinstance(native_df, pa.Table) or is_pandas_wo_pyarrow: - # PyArrowSeries do not have `to_pandas` - result = df[[0, 3, 1], 1].to_numpy() - expected = np.array([11, 14, 12]) - assert np.array_equal(result, expected) - result = df[np.array([0, 3, 1]), "b"].to_numpy() - assert np.array_equal(result, expected) - else: - result = df[[0, 3, 1], 1].to_pandas() - expected_index = range(3) if isinstance(native_df, pl.DataFrame) else [0, 3, 1] - expected = pd.Series([11, 14, 12], name="b", index=expected_index) - assert_series_equal(result, expected, check_dtype=False) - result = df[np.array([0, 3, 1]), "b"].to_pandas() - assert_series_equal(result, expected, check_dtype=False) + + expected = {"b": [11, 14, 12]} + + result = {"b": df[[0, 3, 1], 1]} + compare_dicts(result, expected) + + result = {"b": df[np.array([0, 3, 1]), "b"]} + compare_dicts(result, expected) def test_slice_both_tuples_of_ints(constructor_eager: Any) -> None: @@ -116,6 +114,35 @@ def test_slice_int_rows_str_columns(constructor_eager: Any) -> None: compare_dicts(result, expected) +def test_slice_slice_columns(constructor_eager: Any) -> None: + data = {"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9], "d": [1, 4, 2]} + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df[[0, 1], "b":"c"] # type: ignore[misc] + expected = {"b": [4, 5], "c": [7, 8]} + compare_dicts(result, expected) + result = df[[0, 1], :"c"] # type: ignore[misc] + expected = {"a": [1, 2], "b": [4, 5], "c": [7, 8]} + compare_dicts(result, expected) + result = df[[0, 1], "a":"d":2] # type: ignore[misc] + expected = {"a": [1, 2], "c": [7, 8]} + compare_dicts(result, expected) + result = df[[0, 1], "b":] # type: ignore[misc] + expected = {"b": [4, 5], "c": [7, 8], "d": [1, 4]} + compare_dicts(result, expected) + result = df[[0, 1], 1:3] + expected = {"b": [4, 5], "c": [7, 8]} + compare_dicts(result, expected) + result = df[[0, 1], :3] + expected = {"a": [1, 2], "b": [4, 5], "c": [7, 8]} + compare_dicts(result, expected) + result = df[[0, 1], 0:4:2] + expected = {"a": [1, 2], "c": [7, 8]} + compare_dicts(result, expected) + result = df[[0, 1], 1:] + expected = {"b": [4, 5], "c": [7, 8], "d": [1, 4]} + compare_dicts(result, expected) + + def test_slice_invalid(constructor_eager: Any) -> None: data = {"a": [1, 2], "b": [4, 5]} df = nw.from_native(constructor_eager(data), eager_only=True) diff --git a/tests/frame/test_common.py b/tests/frame/test_common.py deleted file mode 100644 index 4fe903b1f..000000000 --- a/tests/frame/test_common.py +++ /dev/null @@ -1,176 +0,0 @@ -from __future__ import annotations - -import warnings -from typing import Any - -import pandas as pd -import polars as pl -import pyarrow as pa -import pytest - -import narwhals.stable.v1 as nw -from narwhals.functions import _get_deps_info -from narwhals.functions import _get_sys_info -from narwhals.functions import show_versions -from tests.utils import compare_dicts - -data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} -data_na = {"a": [None, 3, 2], "b": [4, 4, 6], "z": [7.0, None, 9]} -data_right = {"c": [6, 12, -1], "d": [0, -4, 2]} - - -@pytest.mark.filterwarnings("ignore:Determining|Resolving.*") -def test_columns(constructor: Any) -> None: - df = nw.from_native(constructor(data)) - result = df.columns - expected = ["a", "b", "z"] - assert result == expected - - -def test_expr_binary(constructor: Any) -> None: - df_raw = constructor(data) - result = nw.from_native(df_raw).with_columns( - a=(1 + 3 * nw.col("a")) * (1 / nw.col("a")), - b=nw.col("z") / (2 - nw.col("b")), - c=nw.col("a") + nw.col("b") / 2, - d=nw.col("a") - nw.col("b"), - e=((nw.col("a") > nw.col("b")) & (nw.col("a") >= nw.col("z"))).cast(nw.Int64), - f=( - (nw.col("a") < nw.col("b")) - | (nw.col("a") <= nw.col("z")) - | (nw.col("a") == 1) - ).cast(nw.Int64), - g=nw.col("a") != 1, - h=(False & (nw.col("a") != 1)), - i=(False | (nw.col("a") != 1)), - j=2 ** nw.col("a"), - k=2 // nw.col("a"), - l=nw.col("a") // 2, - m=nw.col("a") ** 2, - ) - expected = { - "a": [4, 3.333333, 3.5], - "b": [-3.5, -4.0, -2.25], - "z": [7.0, 8.0, 9.0], - "c": [3, 5, 5], - "d": [-3, -1, -4], - "e": [0, 0, 0], - "f": [1, 1, 1], - "g": [False, True, True], - "h": [False, False, False], - "i": [False, True, True], - "j": [2, 8, 4], - "k": [2, 0, 1], - "l": [0, 1, 1], - "m": [1, 9, 4], - } - compare_dicts(result, expected) - - -def test_expr_transform(constructor: Any) -> None: - df = nw.from_native(constructor(data)) - result = df.with_columns(a=nw.col("a").is_between(-1, 1), b=nw.col("b").is_in([4, 5])) - expected = {"a": [True, False, False], "b": [True, True, False], "z": [7, 8, 9]} - compare_dicts(result, expected) - - -def test_expr_na(constructor: Any) -> None: - df = nw.from_native(constructor(data_na)).lazy() - result_nna = df.filter((~nw.col("a").is_null()) & (~df.collect()["z"].is_null())) - expected = {"a": [2], "b": [6], "z": [9]} - compare_dicts(result_nna, expected) - - -def test_lazy(constructor_eager: Any) -> None: - df = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True) - result = df.lazy() - assert isinstance(result, nw.LazyFrame) - - -def test_invalid() -> None: - df = nw.from_native(pa.table({"a": [1, 2], "b": [3, 4]})) - with pytest.raises(ValueError, match="Multi-output"): - df.select(nw.all() + nw.all()) - df = nw.from_native(pd.DataFrame(data)) - with pytest.raises(ValueError, match="Multi-output"): - df.select(nw.all() + nw.all()) - with pytest.raises(TypeError, match="Perhaps you:"): - df.select([pl.col("a")]) # type: ignore[list-item] - with pytest.raises(TypeError, match="Perhaps you:"): - df.select([nw.col("a").cast(pl.Int64)]) - - -@pytest.mark.parametrize("df_raw", [pd.DataFrame(data)]) -def test_reindex(df_raw: Any) -> None: - df = nw.from_native(df_raw, eager_only=True) - result = df.select("b", df["a"].sort(descending=True)) - expected = {"b": [4, 4, 6], "a": [3, 2, 1]} - compare_dicts(result, expected) - result = df.select("b", nw.col("a").sort(descending=True)) - compare_dicts(result, expected) - - s = df["a"] - result_s = s > s.sort() - assert not result_s[0] - assert result_s[1] - assert not result_s[2] - result = df.with_columns(s.sort()) - expected = {"a": [1, 2, 3], "b": [4, 4, 6], "z": [7.0, 8.0, 9.0]} # type: ignore[list-item] - compare_dicts(result, expected) - with pytest.raises(ValueError, match="Multi-output expressions are not supported"): - nw.to_native(df.with_columns(nw.all() + nw.all())) - - -def test_with_columns_order(constructor_eager: Any) -> None: - df = nw.from_native(constructor_eager(data)) - result = df.with_columns(nw.col("a") + 1, d=nw.col("a") - 1) - assert result.columns == ["a", "b", "z", "d"] - expected = {"a": [2, 4, 3], "b": [4, 4, 6], "z": [7.0, 8, 9], "d": [0, 2, 1]} - compare_dicts(result, expected) - - -def test_with_columns_order_single_row(constructor_eager: Any) -> None: - df = nw.from_native(constructor_eager(data)[:1]) - assert len(df) == 1 - result = df.with_columns(nw.col("a") + 1, d=nw.col("a") - 1) - assert result.columns == ["a", "b", "z", "d"] - expected = {"a": [2], "b": [4], "z": [7.0], "d": [0]} - compare_dicts(result, expected) - - -def test_get_sys_info() -> None: - with warnings.catch_warnings(): - warnings.filterwarnings("ignore") - show_versions() - sys_info = _get_sys_info() - - assert "python" in sys_info - assert "executable" in sys_info - assert "machine" in sys_info - - -def test_get_deps_info() -> None: - with warnings.catch_warnings(): - warnings.filterwarnings("ignore") - show_versions() - deps_info = _get_deps_info() - - assert "narwhals" in deps_info - assert "pandas" in deps_info - assert "polars" in deps_info - assert "cudf" in deps_info - assert "modin" in deps_info - assert "pyarrow" in deps_info - assert "numpy" in deps_info - - -def test_show_versions(capsys: Any) -> None: - with warnings.catch_warnings(): - warnings.filterwarnings("ignore") - show_versions() - out, _ = capsys.readouterr() - - assert "python" in out - assert "machine" in out - assert "pandas" in out - assert "polars" in out diff --git a/tests/frame/to_arrow_test.py b/tests/frame/to_arrow_test.py new file mode 100644 index 000000000..c1f395e59 --- /dev/null +++ b/tests/frame/to_arrow_test.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from typing import Any + +import pandas as pd +import pyarrow as pa +import pytest + +import narwhals.stable.v1 as nw +from narwhals.utils import parse_version + + +def test_to_arrow(request: Any, constructor_eager: Any) -> None: + if "pandas" in str(constructor_eager) and parse_version(pd.__version__) < (1, 0, 0): + # pyarrow requires pandas>=1.0.0 + request.applymarker(pytest.mark.xfail) + + data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.1, 8, 9]} + df_raw = constructor_eager(data) + result = nw.from_native(df_raw, eager_only=True).to_arrow() + + expected = pa.table(data) + assert result == expected diff --git a/tests/frame/to_dict_test.py b/tests/frame/to_dict_test.py index f9d246a6c..29c3d2270 100644 --- a/tests/frame/to_dict_test.py +++ b/tests/frame/to_dict_test.py @@ -3,6 +3,7 @@ import pytest import narwhals.stable.v1 as nw +from tests.utils import compare_dicts @pytest.mark.filterwarnings( @@ -22,4 +23,4 @@ def test_to_dict_as_series(constructor_eager: Any) -> None: assert isinstance(result["a"], nw.Series) assert isinstance(result["b"], nw.Series) assert isinstance(result["c"], nw.Series) - assert {key: value.to_list() for key, value in result.items()} == data + compare_dicts(result, data) diff --git a/tests/frame/to_numpy_test.py b/tests/frame/to_numpy_test.py index 4310f6760..d573f4322 100644 --- a/tests/frame/to_numpy_test.py +++ b/tests/frame/to_numpy_test.py @@ -7,7 +7,7 @@ import narwhals.stable.v1 as nw -def test_convert_numpy(constructor_eager: Any) -> None: +def test_to_numpy(constructor_eager: Any) -> None: data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.1, 8, 9]} df_raw = constructor_eager(data) result = nw.from_native(df_raw, eager_only=True).to_numpy() @@ -15,7 +15,3 @@ def test_convert_numpy(constructor_eager: Any) -> None: expected = np.array([[1, 3, 2], [4, 4, 6], [7.1, 8, 9]]).T np.testing.assert_array_equal(result, expected) assert result.dtype == "float64" - - result = nw.from_native(df_raw, eager_only=True).__array__() - np.testing.assert_array_equal(result, expected) - assert result.dtype == "float64" diff --git a/tests/frame/with_columns_sequence_test.py b/tests/frame/with_columns_sequence_test.py index 265e39e32..123425122 100644 --- a/tests/frame/with_columns_sequence_test.py +++ b/tests/frame/with_columns_sequence_test.py @@ -1,6 +1,7 @@ from typing import Any import numpy as np +import pytest import narwhals.stable.v1 as nw from tests.utils import compare_dicts @@ -11,7 +12,9 @@ } -def test_with_columns(constructor: Any) -> None: +def test_with_columns(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) result = ( nw.from_native(constructor(data)) .with_columns(d=np.array([4, 5])) diff --git a/tests/frame/with_columns_test.py b/tests/frame/with_columns_test.py index 605688fa2..864e689e8 100644 --- a/tests/frame/with_columns_test.py +++ b/tests/frame/with_columns_test.py @@ -1,7 +1,10 @@ +from typing import Any + import numpy as np import pandas as pd -import narwhals as nw +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts def test_with_columns_int_col_name_pandas() -> None: @@ -13,3 +16,28 @@ def test_with_columns_int_col_name_pandas() -> None: {0: [1, 4, 7], 1: [2, 5, 8], 2: [3, 6, 9], 4: [2, 5, 8]}, dtype="int64" ) pd.testing.assert_frame_equal(result, expected) + + +def test_with_columns_order(constructor: Any) -> None: + data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} + df = nw.from_native(constructor(data)) + result = df.with_columns(nw.col("a") + 1, d=nw.col("a") - 1) + assert result.collect_schema().names() == ["a", "b", "z", "d"] + expected = {"a": [2, 4, 3], "b": [4, 4, 6], "z": [7.0, 8, 9], "d": [0, 2, 1]} + compare_dicts(result, expected) + + +def test_with_columns_empty(constructor: Any) -> None: + data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]} + df = nw.from_native(constructor(data)) + result = df.select().with_columns() + compare_dicts(result, {}) + + +def test_with_columns_order_single_row(constructor: Any) -> None: + data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9], "i": [0, 1, 2]} + df = nw.from_native(constructor(data)).filter(nw.col("i") < 1).drop("i") + result = df.with_columns(nw.col("a") + 1, d=nw.col("a") - 1) + assert result.collect_schema().names() == ["a", "b", "z", "d"] + expected = {"a": [2], "b": [4], "z": [7.0], "d": [0]} + compare_dicts(result, expected) diff --git a/tests/frame/write_csv_test.py b/tests/frame/write_csv_test.py new file mode 100644 index 000000000..ed9303604 --- /dev/null +++ b/tests/frame/write_csv_test.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any + +import narwhals.stable.v1 as nw +from tests.utils import is_windows + +if TYPE_CHECKING: + import pytest + + +def test_write_csv(constructor_eager: Any, tmpdir: pytest.TempdirFactory) -> None: + data = {"a": [1, 2, 3]} + path = tmpdir / "foo.csv" # type: ignore[operator] + result = nw.from_native(constructor_eager(data), eager_only=True).write_csv(str(path)) + assert path.exists() + assert result is None + result = nw.from_native(constructor_eager(data), eager_only=True).write_csv() + if is_windows(): # pragma: no cover + result = result.replace("\r\n", "\n") + if "pyarrow_table" in str(constructor_eager): + assert result == '"a"\n1\n2\n3\n' + else: + assert result == "a\n1\n2\n3\n" diff --git a/tests/from_dict_test.py b/tests/from_dict_test.py index 17892d80c..cfaf99a7b 100644 --- a/tests/from_dict_test.py +++ b/tests/from_dict_test.py @@ -1,10 +1,14 @@ from typing import Any +import pytest + import narwhals.stable.v1 as nw from tests.utils import compare_dicts -def test_from_dict(constructor: Any) -> None: +def test_from_dict(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]})) native_namespace = nw.get_native_namespace(df) result = nw.from_dict({"c": [1, 2], "d": [5, 6]}, native_namespace=native_namespace) @@ -13,7 +17,9 @@ def test_from_dict(constructor: Any) -> None: assert isinstance(result, nw.DataFrame) -def test_from_dict_schema(constructor: Any) -> None: +def test_from_dict_schema(constructor: Any, request: Any) -> None: + if "dask" in str(constructor): + request.applymarker(pytest.mark.xfail) schema = {"c": nw.Int16(), "d": nw.Float32()} df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]})) native_namespace = nw.get_native_namespace(df) @@ -23,3 +29,27 @@ def test_from_dict_schema(constructor: Any) -> None: schema=schema, # type: ignore[arg-type] ) assert result.collect_schema() == schema + + +def test_from_dict_without_namespace(constructor: Any) -> None: + df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]})).lazy().collect() + result = nw.from_dict({"c": df["a"], "d": df["b"]}) + compare_dicts(result, {"c": [1, 2, 3], "d": [4, 5, 6]}) + + +def test_from_dict_without_namespace_invalid(constructor: Any) -> None: + df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]})).lazy().collect() + with pytest.raises(TypeError, match="namespace"): + nw.from_dict({"c": nw.to_native(df["a"]), "d": nw.to_native(df["b"])}) + + +def test_from_dict_one_native_one_narwhals(constructor: Any) -> None: + df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]})).lazy().collect() + result = nw.from_dict({"c": nw.to_native(df["a"]), "d": df["b"]}) + expected = {"c": [1, 2, 3], "d": [4, 5, 6]} + compare_dicts(result, expected) + + +def test_from_dict_empty() -> None: + with pytest.raises(ValueError, match="empty"): + nw.from_dict({}) diff --git a/tests/hypothesis/test_join.py b/tests/hypothesis/test_join.py index 1111959bb..ebdb88757 100644 --- a/tests/hypothesis/test_join.py +++ b/tests/hypothesis/test_join.py @@ -178,5 +178,6 @@ def test_left_join( # pragma: no cover .pipe(lambda df: df.sort(df.columns)) ) compare_dicts( - result_pa, result_pd.pipe(lambda df: df.sort(df.columns)).to_dict(as_series=False) + result_pa, + result_pd.pipe(lambda df: df.sort(df.columns)).to_dict(as_series=False), ) diff --git a/tests/new_series_test.py b/tests/new_series_test.py new file mode 100644 index 000000000..8ddcabd40 --- /dev/null +++ b/tests/new_series_test.py @@ -0,0 +1,36 @@ +from typing import Any + +import pandas as pd +import pytest + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + + +def test_new_series(constructor_eager: Any) -> None: + s = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"] + result = nw.new_series("b", [4, 1, 2], native_namespace=nw.get_native_namespace(s)) + expected = {"b": [4, 1, 2]} + # all supported libraries auto-infer this to be int64, we can always special-case + # something different if necessary + assert result.dtype == nw.Int64 + compare_dicts(result.to_frame(), expected) + + result = nw.new_series( + "b", [4, 1, 2], nw.Int32, native_namespace=nw.get_native_namespace(s) + ) + expected = {"b": [4, 1, 2]} + assert result.dtype == nw.Int32 + compare_dicts(result.to_frame(), expected) + + +def test_new_series_dask() -> None: + pytest.importorskip("dask") + pytest.importorskip("dask_expr", exc_type=ImportError) + import dask.dataframe as dd + + df = nw.from_native(dd.from_pandas(pd.DataFrame({"a": [1, 2, 3]}))) + with pytest.raises( + NotImplementedError, match="Dask support in Narwhals is lazy-only" + ): + nw.new_series("a", [1, 2, 3], native_namespace=nw.get_native_namespace(df)) diff --git a/tests/no_imports_test.py b/tests/no_imports_test.py new file mode 100644 index 000000000..a89ed0ed8 --- /dev/null +++ b/tests/no_imports_test.py @@ -0,0 +1,68 @@ +import sys + +import pandas as pd +import polars as pl +import pyarrow as pa +import pytest + +import narwhals.stable.v1 as nw + + +def test_polars(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delitem(sys.modules, "pandas") + monkeypatch.delitem(sys.modules, "numpy") + monkeypatch.delitem(sys.modules, "pyarrow") + monkeypatch.delitem(sys.modules, "dask", raising=False) + df = pl.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]}) + nw.from_native(df, eager_only=True).group_by("a").agg(nw.col("b").mean()).filter( + nw.col("a") > 1 + ) + assert "polars" in sys.modules + assert "pandas" not in sys.modules + assert "numpy" not in sys.modules + assert "pyarrow" not in sys.modules + assert "dask" not in sys.modules + + +def test_pandas(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delitem(sys.modules, "polars") + monkeypatch.delitem(sys.modules, "pyarrow") + monkeypatch.delitem(sys.modules, "dask", raising=False) + df = pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]}) + nw.from_native(df, eager_only=True).group_by("a").agg(nw.col("b").mean()).filter( + nw.col("a") > 1 + ) + assert "polars" not in sys.modules + assert "pandas" in sys.modules + assert "numpy" in sys.modules + assert "pyarrow" not in sys.modules + assert "dask" not in sys.modules + + +def test_dask(monkeypatch: pytest.MonkeyPatch) -> None: + pytest.importorskip("dask") + pytest.importorskip("dask_expr", exc_type=ImportError) + import dask.dataframe as dd + + monkeypatch.delitem(sys.modules, "polars") + monkeypatch.delitem(sys.modules, "pyarrow") + df = dd.from_pandas(pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]})) + nw.from_native(df).group_by("a").agg(nw.col("b").mean()).filter(nw.col("a") > 1) + assert "polars" not in sys.modules + assert "pandas" in sys.modules + assert "numpy" in sys.modules + assert "pyarrow" not in sys.modules + assert "dask" in sys.modules + + +def test_pyarrow(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delitem(sys.modules, "polars") + monkeypatch.delitem(sys.modules, "pandas") + monkeypatch.delitem(sys.modules, "dask", raising=False) + df = pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}) + nw.from_native(df).group_by("a").agg(nw.col("b").mean()).filter(nw.col("a") > 1) + assert "polars" not in sys.modules + assert "pandas" not in sys.modules + assert "numpy" in sys.modules + assert "pyarrow" in sys.modules + assert "dask" not in sys.modules diff --git a/tests/test_selectors.py b/tests/selectors_test.py similarity index 92% rename from tests/test_selectors.py rename to tests/selectors_test.py index 6edcb8061..ababee4a7 100644 --- a/tests/test_selectors.py +++ b/tests/selectors_test.py @@ -45,7 +45,10 @@ def test_boolean(constructor: Any) -> None: compare_dicts(result, expected) -def test_string(constructor: Any) -> None: +def test_string(constructor: Any, request: Any) -> None: + if "dask" in str(constructor) and parse_version(pa.__version__) < (12,): + # Dask doesn't infer `'b'` as String for old PyArrow versions + request.applymarker(pytest.mark.xfail) df = nw.from_native(constructor(data)) result = df.select(string()) expected = {"b": ["a", "b", "c"]} diff --git a/tests/series_only/arithmetic_test.py b/tests/series_only/arithmetic_test.py deleted file mode 100644 index b9eee6784..000000000 --- a/tests/series_only/arithmetic_test.py +++ /dev/null @@ -1,67 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import hypothesis.strategies as st -import pandas as pd -import polars as pl -import pyarrow as pa -import pytest -from hypothesis import assume -from hypothesis import given - -import narwhals.stable.v1 as nw -from narwhals.utils import parse_version -from tests.utils import compare_dicts - - -def test_truediv_same_dims(constructor_eager: Any, request: Any) -> None: - if "polars" in str(constructor_eager): - # https://github.com/pola-rs/polars/issues/17760 - request.applymarker(pytest.mark.xfail) - s_left = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"] - s_right = nw.from_native(constructor_eager({"a": [2, 2, 1]}), eager_only=True)["a"] - result = (s_left / s_right).to_list() - assert result == [0.5, 1.0, 3.0] - result = (s_left.__rtruediv__(s_right)).to_list() - assert result == [2, 1, 1 / 3] - - -@pytest.mark.slow() -@given( # type: ignore[misc] - left=st.integers(-100, 100), - right=st.integers(-100, 100), -) -@pytest.mark.skipif( - parse_version(pd.__version__) < (2, 0), reason="convert_dtypes not available" -) -def test_mod(left: int, right: int) -> None: - # hypothesis complains if we add `constructor` as an argument, so this - # test is a bit manual unfortunately - assume(right != 0) - expected = {"a": [left // right]} - result = nw.from_native(pd.DataFrame({"a": [left]}), eager_only=True).select( - nw.col("a") // right - ) - compare_dicts(result, expected) - if parse_version(pd.__version__) < (2, 2): # pragma: no cover - # Bug in old version of pandas - pass - else: - result = nw.from_native( - pd.DataFrame({"a": [left]}).convert_dtypes(dtype_backend="pyarrow"), - eager_only=True, - ).select(nw.col("a") // right) - compare_dicts(result, expected) - result = nw.from_native( - pd.DataFrame({"a": [left]}).convert_dtypes(), eager_only=True - ).select(nw.col("a") // right) - compare_dicts(result, expected) - result = nw.from_native(pl.DataFrame({"a": [left]}), eager_only=True).select( - nw.col("a") // right - ) - compare_dicts(result, expected) - result = nw.from_native(pa.table({"a": [left]}), eager_only=True).select( - nw.col("a") // right - ) - compare_dicts(result, expected) diff --git a/tests/series_only/array_dunder_test.py b/tests/series_only/array_dunder_test.py index 199059ae2..0449199ef 100644 --- a/tests/series_only/array_dunder_test.py +++ b/tests/series_only/array_dunder_test.py @@ -1,11 +1,13 @@ from typing import Any import numpy as np +import pandas as pd import pyarrow as pa import pytest import narwhals.stable.v1 as nw from narwhals.utils import parse_version +from tests.utils import compare_dicts def test_array_dunder(request: Any, constructor_eager: Any) -> None: @@ -14,6 +16,37 @@ def test_array_dunder(request: Any, constructor_eager: Any) -> None: ) < parse_version("16.0.0"): # pragma: no cover request.applymarker(pytest.mark.xfail) + s = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"] + result = s.__array__() + np.testing.assert_array_equal(result, np.array([1, 2, 3], dtype="int64")) + + +def test_array_dunder_with_dtype(request: Any, constructor_eager: Any) -> None: + if "pyarrow_table" in str(constructor_eager) and parse_version( + pa.__version__ + ) < parse_version("16.0.0"): # pragma: no cover + request.applymarker(pytest.mark.xfail) + s = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"] result = s.__array__(object) np.testing.assert_array_equal(result, np.array([1, 2, 3], dtype=object)) + + +def test_array_dunder_with_copy(request: Any, constructor_eager: Any) -> None: + if "pyarrow_table" in str(constructor_eager) and parse_version( + pa.__version__ + ) < parse_version("16.0.0"): # pragma: no cover + request.applymarker(pytest.mark.xfail) + + s = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"] + result = s.__array__(copy=True) + np.testing.assert_array_equal(result, np.array([1, 2, 3], dtype="int64")) + if "pandas_constructor" in str(constructor_eager) and parse_version( + pd.__version__ + ) < (3,): + # If it's pandas, we know that `copy=False` definitely took effect. + # So, let's check it! + result = s.__array__(copy=False) + np.testing.assert_array_equal(result, np.array([1, 2, 3], dtype="int64")) + result[0] = 999 + compare_dicts({"a": s}, {"a": [999, 2, 3]}) diff --git a/tests/series_only/arrow_c_stream_test.py b/tests/series_only/arrow_c_stream_test.py new file mode 100644 index 000000000..9964d7408 --- /dev/null +++ b/tests/series_only/arrow_c_stream_test.py @@ -0,0 +1,41 @@ +import polars as pl +import pyarrow as pa +import pyarrow.compute as pc +import pytest + +import narwhals.stable.v1 as nw +from narwhals.utils import parse_version + + +@pytest.mark.skipif( + parse_version(pl.__version__) < (1, 3), reason="too old for pycapsule in Polars" +) +def test_arrow_c_stream_test() -> None: + s = nw.from_native(pl.Series([1, 2, 3]), series_only=True) + result = pa.chunked_array(s) + expected = pa.chunked_array([[1, 2, 3]]) + assert pc.all(pc.equal(result, expected)).as_py() + + +@pytest.mark.skipif( + parse_version(pl.__version__) < (1, 3), reason="too old for pycapsule in Polars" +) +def test_arrow_c_stream_test_invalid(monkeypatch: pytest.MonkeyPatch) -> None: + # "poison" the dunder method to make sure it actually got called above + monkeypatch.setattr("narwhals.series.Series.__arrow_c_stream__", lambda *_: 1 / 0) + s = nw.from_native(pl.Series([1, 2, 3]), series_only=True) + with pytest.raises(ZeroDivisionError, match="division by zero"): + pa.chunked_array(s) + + +@pytest.mark.skipif( + parse_version(pl.__version__) < (1, 3), reason="too old for pycapsule in Polars" +) +def test_arrow_c_stream_test_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + # Check that fallback to PyArrow works + monkeypatch.delattr("polars.Series.__arrow_c_stream__") + s = nw.from_native(pl.Series([1, 2, 3]).to_frame("a"), eager_only=True)["a"] + s.__arrow_c_stream__() + result = pa.chunked_array(s) + expected = pa.chunked_array([[1, 2, 3]]) + assert pc.all(pc.equal(result, expected)).as_py() diff --git a/tests/series_only/dtype_test.py b/tests/series_only/dtype_test.py new file mode 100644 index 000000000..68d10fbca --- /dev/null +++ b/tests/series_only/dtype_test.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from typing import Any + +import narwhals.stable.v1 as nw + +data = {"a": [1, 3, 2]} + + +def test_dtype(constructor_eager: Any) -> None: + series = nw.from_native(constructor_eager(data), eager_only=True)["a"] + result = series.dtype + assert result == nw.Int64 + assert result.is_numeric() diff --git a/tests/series_only/head_test.py b/tests/series_only/head_test.py deleted file mode 100644 index bfc1c6cbf..000000000 --- a/tests/series_only/head_test.py +++ /dev/null @@ -1,14 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import pytest - -import narwhals as nw - - -@pytest.mark.parametrize("n", [2, -1]) -def test_head(constructor_eager: Any, n: int) -> None: - s = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"] - - assert s.head(n).to_list() == [1, 2] diff --git a/tests/series_only/is_between_test.py b/tests/series_only/is_between_test.py deleted file mode 100644 index 9f63613af..000000000 --- a/tests/series_only/is_between_test.py +++ /dev/null @@ -1,25 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import pytest - -import narwhals.stable.v1 as nw -from tests.utils import compare_dicts - -data = [1, 4, 2, 5] - - -@pytest.mark.parametrize( - ("closed", "expected"), - [ - ("left", [True, True, True, False]), - ("right", [False, True, True, True]), - ("both", [True, True, True, True]), - ("none", [False, True, True, False]), - ], -) -def test_is_between(constructor_eager: Any, closed: str, expected: list[bool]) -> None: - ser = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - result = ser.is_between(1, 5, closed=closed) - compare_dicts({"a": result}, {"a": expected}) diff --git a/tests/series_only/item_test.py b/tests/series_only/item_test.py new file mode 100644 index 000000000..869bd7c38 --- /dev/null +++ b/tests/series_only/item_test.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +import re +from typing import Any + +import pytest + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + +data = [1, 3, 2] + + +@pytest.mark.parametrize(("index", "expected"), [(0, 1), (1, 3)]) +def test_item(constructor_eager: Any, index: int, expected: int) -> None: + series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] + result = series.item(index) + compare_dicts({"a": [result]}, {"a": [expected]}) + compare_dicts({"a": [series.head(1).item()]}, {"a": [1]}) + + with pytest.raises( + ValueError, + match=re.escape("can only call '.item()' if the Series is of length 1,"), + ): + series.item(None) diff --git a/tests/series_only/null_count_test.py b/tests/series_only/null_count_test.py deleted file mode 100644 index 04300217a..000000000 --- a/tests/series_only/null_count_test.py +++ /dev/null @@ -1,12 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import narwhals as nw - - -def test_null_count(constructor_eager: Any) -> None: - data = [1, 2, None] - series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - result = series.null_count() - assert result == 1 diff --git a/tests/series_only/operators_test.py b/tests/series_only/operators_test.py deleted file mode 100644 index a47b34fa8..000000000 --- a/tests/series_only/operators_test.py +++ /dev/null @@ -1,45 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import pytest - -import narwhals.stable.v1 as nw - - -@pytest.mark.parametrize( - ("operator", "expected"), - [ - ("__eq__", [False, True, False]), - ("__ne__", [True, False, True]), - ("__le__", [True, True, False]), - ("__lt__", [True, False, False]), - ("__ge__", [False, True, True]), - ("__gt__", [False, False, True]), - ], -) -def test_comparand_operators( - constructor_eager: Any, operator: str, expected: list[bool] -) -> None: - data = [0, 1, 2] - s = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - result = getattr(s, operator)(1) - assert result.to_list() == expected - - -@pytest.mark.parametrize( - ("operator", "expected"), - [ - ("__and__", [True, False, False, False]), - ("__or__", [True, True, True, False]), - ], -) -def test_logic_operators( - constructor_eager: Any, operator: str, expected: list[bool] -) -> None: - data = [True, True, False, False] - other_data = [True, False, True, False] - series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - other = nw.from_native(constructor_eager({"a": other_data}), eager_only=True)["a"] - result = getattr(series, operator)(other) - assert result.to_list() == expected diff --git a/tests/series_only/tail_test.py b/tests/series_only/tail_test.py deleted file mode 100644 index f978c6e45..000000000 --- a/tests/series_only/tail_test.py +++ /dev/null @@ -1,14 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import pytest - -import narwhals as nw - - -@pytest.mark.parametrize("n", [2, -1]) -def test_tail(constructor_eager: Any, n: int) -> None: - s = nw.from_native(constructor_eager({"a": [1, 2, 3]}), eager_only=True)["a"] - - assert s.tail(n).to_list() == [2, 3] diff --git a/tests/series_only/test_common.py b/tests/series_only/test_common.py deleted file mode 100644 index 742258b73..000000000 --- a/tests/series_only/test_common.py +++ /dev/null @@ -1,122 +0,0 @@ -from __future__ import annotations - -import re -from typing import Any - -import numpy as np -import pandas as pd -import pytest -from numpy.testing import assert_array_equal -from pandas.testing import assert_series_equal - -import narwhals.stable.v1 as nw -from narwhals.utils import parse_version -from tests.utils import compare_dicts - -data = [1, 3, 2] -data_dups = [4, 4, 6] -data_sorted = [7.0, 8, 9] - - -def test_len(constructor_eager: Any) -> None: - series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - - result = len(series) - assert result == 3 - - result = series.len() - assert result == 3 - - -def test_is_in(constructor_eager: Any) -> None: - series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - - result = series.is_in([1, 2]).to_list() - assert result[0] - assert not result[1] - assert result[2] - - -def test_is_in_other(constructor: Any) -> None: - df_raw = constructor({"a": data}) - with pytest.raises( - NotImplementedError, - match=( - "Narwhals `is_in` doesn't accept expressions as an argument, as opposed to Polars. You should provide an iterable instead." - ), - ): - nw.from_native(df_raw).with_columns(contains=nw.col("a").is_in("sets")) - - -def test_dtype(constructor_eager: Any) -> None: - series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - result = series.dtype - assert result == nw.Int64 - assert result.is_numeric() - - -@pytest.mark.skipif( - parse_version(pd.__version__) < parse_version("2.0.0"), reason="too old for pyarrow" -) -def test_convert(request: Any, constructor_eager: Any) -> None: - if any( - cname in str(constructor_eager) - for cname in ("pandas_nullable", "pandas_pyarrow", "modin") - ): - request.applymarker(pytest.mark.xfail) - - series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"].alias( - "a" - ) - - result = series.to_numpy() - assert_array_equal(result, np.array([1, 3, 2])) - - result = series.to_pandas() - assert_series_equal(result, pd.Series([1, 3, 2], name="a")) - - -def test_to_numpy() -> None: - s = pd.Series([1, 2, None], dtype="Int64") - nw_series = nw.from_native(s, series_only=True) - assert nw_series.to_numpy().dtype == "float64" - assert nw_series.__array__().dtype == "float64" - assert nw_series.shape == (3,) - - -def test_zip_with(constructor_eager: Any) -> None: - series1 = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - series2 = nw.from_native(constructor_eager({"a": data_dups}), eager_only=True)["a"] - mask = nw.from_native(constructor_eager({"a": [True, False, True]}), eager_only=True)[ - "a" - ] - - result = series1.zip_with(mask, series2) - expected = [1, 4, 2] - assert result.to_list() == expected - - -@pytest.mark.skipif( - parse_version(pd.__version__) < parse_version("1.0.0"), - reason="too old for convert_dtypes", -) -def test_cast_string() -> None: - s_pd = pd.Series([1, 2]).convert_dtypes() - s = nw.from_native(s_pd, series_only=True) - s = s.cast(nw.String) - result = nw.to_native(s) - assert str(result.dtype) in ("string", "object", "dtype('O')") - - -@pytest.mark.parametrize(("index", "expected"), [(0, 1), (1, 3)]) -def test_item(constructor_eager: Any, index: int, expected: int) -> None: - series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - result = series.item(index) - compare_dicts({"a": [result]}, {"a": [expected]}) - compare_dicts({"a": [series.head(1).item()]}, {"a": [1]}) - - with pytest.raises( - ValueError, - match=re.escape("can only call '.item()' if the Series is of length 1,"), - ): - series.item(None) diff --git a/tests/series_only/to_arrow_test.py b/tests/series_only/to_arrow_test.py new file mode 100644 index 000000000..ebd90b7c2 --- /dev/null +++ b/tests/series_only/to_arrow_test.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from typing import Any + +import pyarrow as pa +import pyarrow.compute as pc +import pytest + +import narwhals.stable.v1 as nw + + +def test_to_arrow(constructor_eager: Any) -> None: + data = [1, 2, 3] + result = nw.from_native(constructor_eager({"a": data}), eager_only=True)[ + "a" + ].to_arrow() + + assert pa.types.is_int64(result.type) + assert pc.all(pc.equal(result, pa.array(data, type=pa.int64()))) + + +def test_to_arrow_with_nulls(constructor_eager: Any, request: Any) -> None: + if "pandas_constructor" in str(constructor_eager) or "modin_constructor" in str( + constructor_eager + ): + request.applymarker(pytest.mark.xfail) + + data = [1, 2, None] + result = ( + nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] + .cast(nw.Int64) + .to_arrow() + ) + + assert pa.types.is_int64(result.type) + assert pc.all(pc.equal(result, pa.array(data, type=pa.int64()))) diff --git a/tests/series_only/to_list_test.py b/tests/series_only/to_list_test.py index f4c7ec5e6..10e415916 100644 --- a/tests/series_only/to_list_test.py +++ b/tests/series_only/to_list_test.py @@ -1,10 +1,15 @@ from typing import Any +import pytest + import narwhals.stable.v1 as nw +from tests.utils import compare_dicts data = [1, 2, 3] -def test_to_list(constructor_eager: Any) -> None: +def test_to_list(constructor_eager: Any, request: Any) -> None: + if "cudf" in str(constructor_eager): # pragma: no cover + request.applymarker(pytest.mark.xfail) s = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"] - assert s.to_list() == [1, 2, 3] + compare_dicts({"a": s.to_list()}, {"a": [1, 2, 3]}) diff --git a/tests/series_only/to_numpy_test.py b/tests/series_only/to_numpy_test.py new file mode 100644 index 000000000..f5ed59fe1 --- /dev/null +++ b/tests/series_only/to_numpy_test.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest +from numpy.testing import assert_array_equal + +import narwhals.stable.v1 as nw + + +def test_to_numpy(constructor_eager: Any, request: Any) -> None: + if "pandas_constructor" in str(constructor_eager) or "modin_constructor" in str( + constructor_eager + ): + request.applymarker(pytest.mark.xfail) + + data = [1, 2, None] + s = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"].cast( + nw.Int64 + ) + assert s.to_numpy().dtype == "float64" + assert s.shape == (3,) + + assert_array_equal(s.to_numpy(), np.array(data, dtype=float)) diff --git a/tests/series_only/to_pandas_test.py b/tests/series_only/to_pandas_test.py new file mode 100644 index 000000000..747bed8b2 --- /dev/null +++ b/tests/series_only/to_pandas_test.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from typing import Any + +import pandas as pd +import pytest +from pandas.testing import assert_series_equal + +import narwhals.stable.v1 as nw +from narwhals.utils import parse_version + +data = [1, 3, 2] + + +@pytest.mark.skipif( + parse_version(pd.__version__) < parse_version("2.0.0"), reason="too old for pyarrow" +) +def test_convert(request: Any, constructor_eager: Any) -> None: + if any( + cname in str(constructor_eager) + for cname in ("pandas_nullable", "pandas_pyarrow", "modin") + ): + request.applymarker(pytest.mark.xfail) + + series = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"].alias( + "a" + ) + + result = series.to_pandas() + assert_series_equal(result, pd.Series([1, 3, 2], name="a")) diff --git a/tests/series_only/zip_with_test.py b/tests/series_only/zip_with_test.py new file mode 100644 index 000000000..0c068c386 --- /dev/null +++ b/tests/series_only/zip_with_test.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from typing import Any + +import narwhals.stable.v1 as nw +from tests.utils import compare_dicts + + +def test_zip_with(constructor_eager: Any) -> None: + series1 = nw.from_native(constructor_eager({"a": [1, 3, 2]}), eager_only=True)["a"] + series2 = nw.from_native(constructor_eager({"a": [4, 4, 6]}), eager_only=True)["a"] + mask = nw.from_native(constructor_eager({"a": [True, False, True]}), eager_only=True)[ + "a" + ] + + result = series1.zip_with(mask, series2) + expected = [1, 4, 2] + compare_dicts({"a": result}, {"a": expected}) diff --git a/tests/system_info_test.py b/tests/system_info_test.py new file mode 100644 index 000000000..30bb0c400 --- /dev/null +++ b/tests/system_info_test.py @@ -0,0 +1,44 @@ +import warnings +from typing import Any + +from narwhals.functions import _get_deps_info +from narwhals.functions import _get_sys_info +from narwhals.functions import show_versions + + +def test_get_sys_info() -> None: + with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + show_versions() + sys_info = _get_sys_info() + + assert "python" in sys_info + assert "executable" in sys_info + assert "machine" in sys_info + + +def test_get_deps_info() -> None: + with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + show_versions() + deps_info = _get_deps_info() + + assert "narwhals" in deps_info + assert "pandas" in deps_info + assert "polars" in deps_info + assert "cudf" in deps_info + assert "modin" in deps_info + assert "pyarrow" in deps_info + assert "numpy" in deps_info + + +def test_show_versions(capsys: Any) -> None: + with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + show_versions() + out, _ = capsys.readouterr() + + assert "python" in out + assert "machine" in out + assert "pandas" in out + assert "polars" in out diff --git a/tests/test_group_by.py b/tests/test_group_by.py index 6f16860ef..2bb8d435b 100644 --- a/tests/test_group_by.py +++ b/tests/test_group_by.py @@ -35,6 +35,25 @@ def test_group_by_complex() -> None: compare_dicts(result, expected) +def test_invalid_group_by_dask() -> None: + pytest.importorskip("dask") + pytest.importorskip("dask_expr", exc_type=ImportError) + import dask.dataframe as dd + + df_dask = dd.from_pandas(df_pandas) + + with pytest.raises(ValueError, match=r"Non-trivial complex found"): + nw.from_native(df_dask).group_by("a").agg(nw.col("b").mean().min()) + + with pytest.raises(RuntimeError, match="does your"): + nw.from_native(df_dask).group_by("a").agg(nw.col("b")) + + with pytest.raises( + ValueError, match=r"Anonymous expressions are not supported in group_by\.agg" + ): + nw.from_native(df_dask).group_by("a").agg(nw.all().mean()) + + def test_invalid_group_by() -> None: df = nw.from_native(df_pandas) with pytest.raises(RuntimeError, match="does your"): @@ -53,10 +72,7 @@ def test_invalid_group_by() -> None: ) -def test_group_by_iter(request: Any, constructor_eager: Any) -> None: - if "pyarrow_table" in str(constructor_eager): - request.applymarker(pytest.mark.xfail) - +def test_group_by_iter(constructor_eager: Any) -> None: df = nw.from_native(constructor_eager(data), eager_only=True) expected_keys = [(1,), (3,)] keys = [] @@ -69,11 +85,11 @@ def test_group_by_iter(request: Any, constructor_eager: Any) -> None: assert sorted(keys) == sorted(expected_keys) expected_keys = [(1, 4), (3, 6)] # type: ignore[list-item] keys = [] - for key, _df in df.group_by("a", "b"): + for key, _ in df.group_by("a", "b"): keys.append(key) assert sorted(keys) == sorted(expected_keys) keys = [] - for key, _df in df.group_by(["a", "b"]): + for key, _ in df.group_by(["a", "b"]): keys.append(key) assert sorted(keys) == sorted(expected_keys) @@ -97,13 +113,14 @@ def test_group_by_empty_result_pandas() -> None: def test_group_by_simple_named(constructor: Any) -> None: data = {"a": [1, 1, 2], "b": [4, 5, 6], "c": [7, 2, 1]} - df = nw.from_native(constructor(data)) + df = nw.from_native(constructor(data)).lazy() result = ( df.group_by("a") .agg( b_min=nw.col("b").min(), b_max=nw.col("b").max(), ) + .collect() .sort("a") ) expected = { @@ -116,13 +133,14 @@ def test_group_by_simple_named(constructor: Any) -> None: def test_group_by_simple_unnamed(constructor: Any) -> None: data = {"a": [1, 1, 2], "b": [4, 5, 6], "c": [7, 2, 1]} - df = nw.from_native(constructor(data)) + df = nw.from_native(constructor(data)).lazy() result = ( df.group_by("a") .agg( nw.col("b").min(), nw.col("c").max(), ) + .collect() .sort("a") ) expected = { @@ -135,13 +153,14 @@ def test_group_by_simple_unnamed(constructor: Any) -> None: def test_group_by_multiple_keys(constructor: Any) -> None: data = {"a": [1, 1, 2], "b": [4, 4, 6], "c": [7, 2, 1]} - df = nw.from_native(constructor(data)) + df = nw.from_native(constructor(data)).lazy() result = ( df.group_by("a", "b") .agg( c_min=nw.col("c").min(), c_max=nw.col("c").max(), ) + .collect() .sort("a") ) expected = { diff --git a/tests/test_schema.py b/tests/test_schema.py deleted file mode 100644 index f85fdd816..000000000 --- a/tests/test_schema.py +++ /dev/null @@ -1,22 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import pytest - -import narwhals.stable.v1 as nw - -data = {"a": nw.Int64(), "b": nw.Float32(), "c": nw.String()} - - -@pytest.mark.parametrize( - ("method", "expected"), - [ - ("names", ["a", "b", "c"]), - ("dtypes", [nw.Int64(), nw.Float32(), nw.String()]), - ("len", 3), - ], -) -def test_schema_object(method: str, expected: Any) -> None: - schema = nw.Schema(data) - assert getattr(schema, method)() == expected diff --git a/tests/test_utils.py b/tests/test_utils.py index f8b9b98ec..68dc90ed7 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -2,6 +2,7 @@ import polars as pl import pytest from pandas.testing import assert_frame_equal +from pandas.testing import assert_index_equal from pandas.testing import assert_series_equal import narwhals.stable.v1 as nw @@ -65,6 +66,24 @@ def test_maybe_set_index_polars() -> None: assert result is df +def test_maybe_get_index_pandas() -> None: + pandas_df = pd.DataFrame({"a": [1, 2, 3]}, index=[1, 2, 0]) + result = nw.maybe_get_index(nw.from_native(pandas_df)) + assert_index_equal(result, pandas_df.index) + pandas_series = pd.Series([1, 2, 3], index=[1, 2, 0]) + result_s = nw.maybe_get_index(nw.from_native(pandas_series, series_only=True)) + assert_index_equal(result_s, pandas_series.index) + + +def test_maybe_get_index_polars() -> None: + df = nw.from_native(pl.DataFrame({"a": [1, 2, 3]})) + result = nw.maybe_get_index(df) + assert result is None + series = nw.from_native(pl.Series([1, 2, 3]), series_only=True) + result = nw.maybe_get_index(series) + assert result is None + + @pytest.mark.skipif( parse_version(pd.__version__) < parse_version("1.0.0"), reason="too old for convert_dtypes", @@ -72,10 +91,15 @@ def test_maybe_set_index_polars() -> None: def test_maybe_convert_dtypes_pandas() -> None: import numpy as np - df = nw.from_native(pd.DataFrame({"a": [1, np.nan]}, dtype=np.dtype("float64"))) + df = nw.from_native( + pd.DataFrame({"a": [1, np.nan]}, dtype=np.dtype("float64")), eager_only=True + ) result = nw.to_native(nw.maybe_convert_dtypes(df)) expected = pd.DataFrame({"a": [1, pd.NA]}, dtype="Int64") pd.testing.assert_frame_equal(result, expected) + result_s = nw.to_native(nw.maybe_convert_dtypes(df["a"])) + expected_s = pd.Series([1, pd.NA], name="a", dtype="Int64") + pd.testing.assert_series_equal(result_s, expected_s) def test_maybe_convert_dtypes_polars() -> None: diff --git a/tests/tpch_q1_test.py b/tests/tpch_q1_test.py index da540e465..999f32f76 100644 --- a/tests/tpch_q1_test.py +++ b/tests/tpch_q1_test.py @@ -17,7 +17,7 @@ @pytest.mark.parametrize( "library", - ["pandas", "polars", "pyarrow"], + ["pandas", "polars", "pyarrow", "dask"], ) @pytest.mark.filterwarnings("ignore:.*Passing a BlockManager.*:DeprecationWarning") def test_q1(library: str, request: Any) -> None: @@ -28,6 +28,14 @@ def test_q1(library: str, request: Any) -> None: df_raw["l_shipdate"] = pd.to_datetime(df_raw["l_shipdate"]) elif library == "polars": df_raw = pl.scan_parquet("tests/data/lineitem.parquet") + elif library == "dask": + pytest.importorskip("dask") + pytest.importorskip("dask_expr", exc_type=ImportError) + import dask.dataframe as dd + + df_raw = dd.from_pandas( + pd.read_parquet("tests/data/lineitem.parquet", dtype_backend="pyarrow") + ) else: df_raw = pq.read_table("tests/data/lineitem.parquet") var_1 = datetime(1998, 9, 2) diff --git a/tests/translate/from_native_test.py b/tests/translate/from_native_test.py index 191ad49b5..af49c0226 100644 --- a/tests/translate/from_native_test.py +++ b/tests/translate/from_native_test.py @@ -169,3 +169,33 @@ def test_init_already_narwhals_unstable() -> None: s = df["a"] result_s = unstable_nw.from_native(s, allow_series=True) assert result_s is s + + +def test_series_only_dask() -> None: + pytest.importorskip("dask") + pytest.importorskip("dask_expr", exc_type=ImportError) + import dask.dataframe as dd + + dframe = dd.from_pandas(df_pd) + + with pytest.raises(TypeError, match="Cannot only use `series_only`"): + nw.from_native(dframe, series_only=True) + + +@pytest.mark.parametrize( + ("eager_only", "context"), + [ + (False, does_not_raise()), + (True, pytest.raises(TypeError, match="Cannot only use `eager_only`")), + ], +) +def test_eager_only_lazy_dask(eager_only: Any, context: Any) -> None: + pytest.importorskip("dask") + pytest.importorskip("dask_expr", exc_type=ImportError) + import dask.dataframe as dd + + dframe = dd.from_pandas(df_pd) + + with context: + res = nw.from_native(dframe, eager_only=eager_only) + assert isinstance(res, nw.LazyFrame) diff --git a/tpch/notebooks/gpu/execute.ipynb b/tpch/notebooks/gpu/execute.ipynb index c776cc1cf..a117c9187 100755 --- a/tpch/notebooks/gpu/execute.ipynb +++ b/tpch/notebooks/gpu/execute.ipynb @@ -452,7 +452,7 @@ "source": [ "import cudf\n", "fn = cudf.read_parquet\n", - "timings = %timeit -o q1(fn(lineitem))\n", + "timings = %timeit -o -q q1(fn(lineitem))\n", "results['q1'] = timings.all_runs" ] }, @@ -474,7 +474,7 @@ "source": [ "import cudf\n", "fn = cudf.read_parquet\n", - "timings = %timeit -o q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", + "timings = %timeit -o -q q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", "results['q2'] = timings.all_runs" ] }, @@ -496,7 +496,7 @@ "source": [ "import cudf\n", "fn = cudf.read_parquet\n", - "timings = %timeit -o q3(fn(customer), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q3(fn(customer), fn(lineitem), fn(orders))\n", "results['q3'] = timings.all_runs" ] }, @@ -518,7 +518,7 @@ "source": [ "import cudf\n", "fn = cudf.read_parquet\n", - "timings = %timeit -o q4(fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q4(fn(lineitem), fn(orders))\n", "results['q4'] = timings.all_runs" ] }, @@ -540,7 +540,7 @@ "source": [ "import cudf\n", "fn = cudf.read_parquet\n", - "timings = %timeit -o q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results['q5'] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q1/execute.ipynb b/tpch/notebooks/q1/execute.ipynb index 4760d2dfd..cc6dd4559 100755 --- a/tpch/notebooks/q1/execute.ipynb +++ b/tpch/notebooks/q1/execute.ipynb @@ -16,7 +16,7 @@ }, "outputs": [], "source": [ - "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow ibis-framework narwhals" + "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals dask[dataframe]" ] }, { @@ -46,50 +46,6 @@ "cell_type": "code", "execution_count": null, "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "from datetime import date\n", - "\n", - "def q1_pandas_native(lineitem):\n", - " VAR1 = date(1998, 9, 2)\n", - "\n", - " sel = lineitem.l_shipdate <= VAR1\n", - " lineitem_filtered = lineitem[sel]\n", - "\n", - " # This is lenient towards pandas as normally an optimizer should decide\n", - " # that this could be computed before the groupby aggregation.\n", - " # Other implementations don't enjoy this benefit.\n", - " lineitem_filtered[\"disc_price\"] = lineitem_filtered.l_extendedprice * (\n", - " 1 - lineitem_filtered.l_discount\n", - " )\n", - " lineitem_filtered[\"charge\"] = (\n", - " lineitem_filtered.l_extendedprice\n", - " * (1 - lineitem_filtered.l_discount)\n", - " * (1 + lineitem_filtered.l_tax)\n", - " )\n", - " gb = lineitem_filtered.groupby([\"l_returnflag\", \"l_linestatus\"], as_index=False)\n", - "\n", - " total = gb.agg(\n", - " sum_qty=pd.NamedAgg(column=\"l_quantity\", aggfunc=\"sum\"),\n", - " sum_base_price=pd.NamedAgg(column=\"l_extendedprice\", aggfunc=\"sum\"),\n", - " sum_disc_price=pd.NamedAgg(column=\"disc_price\", aggfunc=\"sum\"),\n", - " sum_charge=pd.NamedAgg(column=\"charge\", aggfunc=\"sum\"),\n", - " avg_qty=pd.NamedAgg(column=\"l_quantity\", aggfunc=\"mean\"),\n", - " avg_price=pd.NamedAgg(column=\"l_extendedprice\", aggfunc=\"mean\"),\n", - " avg_disc=pd.NamedAgg(column=\"l_discount\", aggfunc=\"mean\"),\n", - " count_order=pd.NamedAgg(column=\"l_orderkey\", aggfunc=\"size\"),\n", - " )\n", - "\n", - " result_df = total.sort_values([\"l_returnflag\", \"l_linestatus\"])\n", - "\n", - " return result_df # type: ignore[no-any-return]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", "metadata": { "papermill": { "duration": 0.021725, @@ -119,67 +75,25 @@ " * (1.0 + nw.col(\"l_tax\"))\n", " ),\n", " )\n", - " .group_by([\"l_returnflag\", \"l_linestatus\"])\n", + " .group_by(\"l_returnflag\", \"l_linestatus\")\n", " .agg(\n", - " [\n", - " nw.col(\"l_quantity\").sum().alias(\"sum_qty\"),\n", - " nw.col(\"l_extendedprice\").sum().alias(\"sum_base_price\"),\n", - " nw.col(\"disc_price\").sum().alias(\"sum_disc_price\"),\n", - " nw.col(\"charge\").sum().alias(\"sum_charge\"),\n", - " nw.col(\"l_quantity\").mean().alias(\"avg_qty\"),\n", - " nw.col(\"l_extendedprice\").mean().alias(\"avg_price\"),\n", - " nw.col(\"l_discount\").mean().alias(\"avg_disc\"),\n", - " nw.len().alias(\"count_order\"),\n", - " ],\n", + " nw.col(\"l_quantity\").sum().alias(\"sum_qty\"),\n", + " nw.col(\"l_extendedprice\").sum().alias(\"sum_base_price\"),\n", + " nw.col(\"disc_price\").sum().alias(\"sum_disc_price\"),\n", + " nw.col(\"charge\").sum().alias(\"sum_charge\"),\n", + " nw.col(\"l_quantity\").mean().alias(\"avg_qty\"),\n", + " nw.col(\"l_extendedprice\").mean().alias(\"avg_price\"),\n", + " nw.col(\"l_discount\").mean().alias(\"avg_disc\"),\n", + " nw.len().alias(\"count_order\"),\n", " )\n", - " .sort([\"l_returnflag\", \"l_linestatus\"])\n", + " .sort(\"l_returnflag\", \"l_linestatus\")\n", " )" ] }, { "cell_type": "code", "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [ - "def q1_ibis(lineitem: Any, *, tool):\n", - " var1 = datetime(1998, 9, 2)\n", - " lineitem = lineitem.filter(lineitem[\"l_shipdate\"] <= var1)\n", - " lineitem = lineitem.mutate(\n", - " disc_price=lineitem[\"l_extendedprice\"] * (1 - lineitem[\"l_discount\"]),\n", - " charge=(\n", - " lineitem[\"l_extendedprice\"]\n", - " * (1.0 - lineitem[\"l_discount\"])\n", - " * (1.0 + lineitem[\"l_tax\"])\n", - " ),\n", - " )\n", - " q_final = (\n", - " lineitem\n", - " .group_by([\"l_returnflag\", \"l_linestatus\"])\n", - " .aggregate(\n", - " sum_qty=lineitem[\"l_quantity\"].sum(),\n", - " sum_base_price=lineitem[\"l_extendedprice\"].sum(),\n", - " sum_disc_price=(lineitem['disc_price'].sum()),\n", - " sum_charge=(lineitem['charge'].sum()),\n", - " avg_qty=lineitem[\"l_quantity\"].mean(),\n", - " avg_price=lineitem[\"l_extendedprice\"].mean(),\n", - " avg_disc=lineitem[\"l_discount\"].mean(),\n", - " count_order=lambda lineitem: lineitem.count(),\n", - " )\n", - " .order_by([\"l_returnflag\", \"l_linestatus\"])\n", - " )\n", - " if tool == 'pandas':\n", - " return q_final.to_pandas()\n", - " if tool == 'polars':\n", - " return q_final.to_polars()\n", - " raise ValueError(\"expected pandas or polars\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", + "id": "3", "metadata": { "papermill": { "duration": 0.013325, @@ -206,7 +120,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "4", "metadata": { "papermill": { "duration": 0.014284, @@ -219,27 +133,23 @@ }, "outputs": [], "source": [ - "import ibis\n", "import pyarrow.parquet as pq\n", - "\n", - "con_pd = ibis.pandas.connect()\n", - "con_pl = ibis.polars.connect()\n", + "import dask.dataframe as dd\n", "\n", "IO_FUNCS = {\n", " 'pandas': lambda x: pd.read_parquet(x, engine='pyarrow'),\n", " 'pandas[pyarrow]': lambda x: pd.read_parquet(x, engine='pyarrow', dtype_backend='pyarrow'),\n", - " 'pandas[pyarrow][ibis]': lambda x: con_pd.read_parquet(x, engine='pyarrow', dtype_backend='pyarrow'),\n", " 'polars[eager]': lambda x: pl.read_parquet(x),\n", " 'polars[lazy]': lambda x: pl.scan_parquet(x),\n", - " 'polars[lazy][ibis]': lambda x: con_pl.read_parquet(x),\n", " 'pyarrow': lambda x: pq.read_table(x),\n", + " 'dask': lambda x: dd.read_parquet(x, engine='pyarrow', dtype_backend='pyarrow'),\n", "}" ] }, { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -248,7 +158,7 @@ }, { "cell_type": "markdown", - "id": "8", + "id": "6", "metadata": {}, "source": [ "## PyArrow.table" @@ -257,63 +167,19 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "7", "metadata": {}, "outputs": [], "source": [ "tool = 'pyarrow'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q1(fn(lineitem))\n", + "timings = %timeit -o -q q1(fn(lineitem))\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "10", - "metadata": {}, - "source": [ - "## pandas, pyarrow dtypes, ibis" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "11", - "metadata": {}, - "outputs": [], - "source": [ - "# very slow, comment-out at your own risk\n", - "\n", - "# tool = 'pandas[pyarrow][ibis]'\n", - "# fn = IO_FUNCS[tool]\n", - "# timings = %timeit -o q1_ibis(fn(lineitem), tool='pandas')\n", - "# results[tool] = timings.all_runs" - ] - }, - { - "cell_type": "markdown", - "id": "12", - "metadata": {}, - "source": [ - "## pandas, pyarrow dtypes, native" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "13", - "metadata": {}, - "outputs": [], - "source": [ - "tool = 'pandas[pyarrow]'\n", - "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q1_pandas_native(lineitem=fn(lineitem))\n", - "results[tool+'[native]'] = timings.all_runs" - ] - }, - { - "cell_type": "markdown", - "id": "14", + "id": "8", "metadata": { "papermill": { "duration": 0.005113, @@ -325,13 +191,13 @@ "tags": [] }, "source": [ - "## pandas via Narwhals" + "## pandas" ] }, { "cell_type": "code", "execution_count": null, - "id": "15", + "id": "9", "metadata": { "papermill": { "duration": 196.786925, @@ -346,13 +212,13 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q1(lineitem_ds=fn(lineitem))\n", + "timings = %timeit -o -q q1(lineitem_ds=fn(lineitem))\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "16", + "id": "10", "metadata": { "papermill": { "duration": 0.005184, @@ -364,13 +230,13 @@ "tags": [] }, "source": [ - "## pandas, pyarrow dtypes, via Narwhals" + "## pandas, pyarrow dtypes" ] }, { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "11", "metadata": { "papermill": { "duration": 158.748353, @@ -385,13 +251,13 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q1(fn(lineitem))\n", + "timings = %timeit -o -q q1(fn(lineitem))\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "18", + "id": "12", "metadata": { "papermill": { "duration": 0.005773, @@ -409,7 +275,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "13", "metadata": { "papermill": { "duration": 37.821116, @@ -424,13 +290,13 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q1(fn(lineitem))\n", + "timings = %timeit -o -q q1(fn(lineitem))\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "20", + "id": "14", "metadata": { "papermill": { "duration": 0.005515, @@ -448,7 +314,7 @@ { "cell_type": "code", "execution_count": null, - "id": "21", + "id": "15", "metadata": { "papermill": { "duration": 4.800698, @@ -463,34 +329,34 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q1(fn(lineitem)).collect()\n", + "timings = %timeit -o -q q1(fn(lineitem)).collect()\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "22", + "id": "16", "metadata": {}, "source": [ - "## Polars scan_parquet ibis" + "## Dask Dataframe" ] }, { "cell_type": "code", "execution_count": null, - "id": "23", + "id": "17", "metadata": {}, "outputs": [], "source": [ - "tool = 'polars[lazy][ibis]'\n", + "tool = 'dask'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q1_ibis(fn(lineitem), tool='polars')\n", + "timings = %timeit -o -q q1(fn(lineitem)).collect()\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "24", + "id": "18", "metadata": {}, "source": [ "## Save" @@ -499,7 +365,7 @@ { "cell_type": "code", "execution_count": null, - "id": "25", + "id": "19", "metadata": {}, "outputs": [], "source": [ diff --git a/tpch/notebooks/q10/execute.ipynb b/tpch/notebooks/q10/execute.ipynb index 307f69e7a..85ec0f14b 100644 --- a/tpch/notebooks/q10/execute.ipynb +++ b/tpch/notebooks/q10/execute.ipynb @@ -198,7 +198,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q10(fn(customer), fn(nation), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q10(fn(customer), fn(nation), fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -235,7 +235,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q10(fn(customer), fn(nation), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q10(fn(customer), fn(nation), fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -272,7 +272,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q10(fn(customer), fn(nation), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q10(fn(customer), fn(nation), fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -309,7 +309,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q10(fn(customer), fn(nation), fn(lineitem), fn(orders)).collect()\n", + "timings = %timeit -o -q q10(fn(customer), fn(nation), fn(lineitem), fn(orders)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q11/execute.ipynb b/tpch/notebooks/q11/execute.ipynb index fec9ee27e..33951d922 100644 --- a/tpch/notebooks/q11/execute.ipynb +++ b/tpch/notebooks/q11/execute.ipynb @@ -186,7 +186,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q11(fn(partsupp), fn(nation), fn(supplier))\n", + "timings = %timeit -o -q q11(fn(partsupp), fn(nation), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -223,7 +223,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q11(fn(partsupp), fn(nation), fn(supplier))\n", + "timings = %timeit -o -q q11(fn(partsupp), fn(nation), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -260,7 +260,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q11(fn(partsupp), fn(nation), fn(supplier))\n", + "timings = %timeit -o -q q11(fn(partsupp), fn(nation), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -297,7 +297,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q11(fn(partsupp), fn(nation), fn(supplier)).collect()\n", + "timings = %timeit -o -q q11(fn(partsupp), fn(nation), fn(supplier)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q15/execute.ipynb b/tpch/notebooks/q15/execute.ipynb index b487a9bf3..0baf11956 100644 --- a/tpch/notebooks/q15/execute.ipynb +++ b/tpch/notebooks/q15/execute.ipynb @@ -177,7 +177,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q15(fn(lineitem), fn(supplier))\n", + "timings = %timeit -o -q q15(fn(lineitem), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -214,7 +214,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q15(fn(lineitem), fn(supplier))\n", + "timings = %timeit -o -q q15(fn(lineitem), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -251,7 +251,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q15(fn(lineitem), fn(supplier))\n", + "timings = %timeit -o -q q15(fn(lineitem), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -288,7 +288,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q15(fn(lineitem), fn(supplier)).collect()\n", + "timings = %timeit -o -q q15(fn(lineitem), fn(supplier)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q17/execute.ipynb b/tpch/notebooks/q17/execute.ipynb index 958c7f5be..b13445d28 100644 --- a/tpch/notebooks/q17/execute.ipynb +++ b/tpch/notebooks/q17/execute.ipynb @@ -173,7 +173,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q17(fn(lineitem), fn(part))\n", + "timings = %timeit -o -q q17(fn(lineitem), fn(part))\n", "results[tool] = timings.all_runs" ] }, @@ -210,7 +210,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q17(fn(lineitem), fn(part))\n", + "timings = %timeit -o -q q17(fn(lineitem), fn(part))\n", "results[tool] = timings.all_runs" ] }, @@ -247,7 +247,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q17(fn(lineitem), fn(part))\n", + "timings = %timeit -o -q q17(fn(lineitem), fn(part))\n", "results[tool] = timings.all_runs" ] }, @@ -284,7 +284,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q17(fn(lineitem), fn(part)).collect()\n", + "timings = %timeit -o -q q17(fn(lineitem), fn(part)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q18/execute.ipynb b/tpch/notebooks/q18/execute.ipynb index 21557c957..c90629e0f 100644 --- a/tpch/notebooks/q18/execute.ipynb +++ b/tpch/notebooks/q18/execute.ipynb @@ -121,7 +121,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q19(fn(lineitem), fn(part))\n", + "timings = %timeit -o -q q19(fn(lineitem), fn(part))\n", "results[tool] = timings.all_runs" ] }, @@ -140,7 +140,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q18(fn(customer), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q18(fn(customer), fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -159,7 +159,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q18(fn(customer), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q18(fn(customer), fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -178,7 +178,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q18(fn(customer), fn(lineitem), fn(orders)).collect()\n", + "timings = %timeit -o -q q18(fn(customer), fn(lineitem), fn(orders)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q19/execute.ipynb b/tpch/notebooks/q19/execute.ipynb index a8cd3fea3..8483e06d5 100644 --- a/tpch/notebooks/q19/execute.ipynb +++ b/tpch/notebooks/q19/execute.ipynb @@ -194,7 +194,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q19(fn(lineitem), fn(part))\n", + "timings = %timeit -o -q q19(fn(lineitem), fn(part))\n", "results[tool] = timings.all_runs" ] }, @@ -231,7 +231,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q19(fn(lineitem), fn(part))\n", + "timings = %timeit -o -q q19(fn(lineitem), fn(part))\n", "results[tool] = timings.all_runs" ] }, @@ -268,7 +268,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q19(fn(lineitem), fn(part))\n", + "timings = %timeit -o -q q19(fn(lineitem), fn(part))\n", "results[tool] = timings.all_runs" ] }, @@ -305,7 +305,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q19(fn(lineitem), fn(part)).collect()\n", + "timings = %timeit -o -q q19(fn(lineitem), fn(part)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q2/execute.ipynb b/tpch/notebooks/q2/execute.ipynb index b4e59307b..c05345336 100755 --- a/tpch/notebooks/q2/execute.ipynb +++ b/tpch/notebooks/q2/execute.ipynb @@ -16,13 +16,23 @@ }, "outputs": [], "source": [ - "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework " + "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow dask[dataframe]" ] }, { "cell_type": "code", "execution_count": null, "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install git+https://github.com/MarcoGorelli/narwhals.git@more-dask-tpch" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", "metadata": { "papermill": { "duration": 0.907754, @@ -42,64 +52,6 @@ "pd.options.future.infer_string = True" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "from typing import Any\n", - "\n", - "def q2_pandas_native(\n", - " region_ds: Any,\n", - " nation_ds: Any,\n", - " supplier_ds: Any,\n", - " part_ds: Any,\n", - " part_supp_ds: Any,\n", - "):\n", - " var1 = 15\n", - " var2 = \"BRASS\"\n", - " var3 = \"EUROPE\"\n", - "\n", - " jn = (\n", - " part_ds.merge(part_supp_ds, left_on=\"p_partkey\", right_on=\"ps_partkey\")\n", - " .merge(supplier_ds, left_on=\"ps_suppkey\", right_on=\"s_suppkey\")\n", - " .merge(nation_ds, left_on=\"s_nationkey\", right_on=\"n_nationkey\")\n", - " .merge(region_ds, left_on=\"n_regionkey\", right_on=\"r_regionkey\")\n", - " )\n", - "\n", - " jn = jn[jn[\"p_size\"] == var1]\n", - " jn = jn[jn[\"p_type\"].str.endswith(var2)]\n", - " jn = jn[jn[\"r_name\"] == var3]\n", - "\n", - " gb = jn.groupby(\"p_partkey\", as_index=False)\n", - " agg = gb[\"ps_supplycost\"].min()\n", - " jn2 = agg.merge(jn, on=[\"p_partkey\", \"ps_supplycost\"])\n", - "\n", - " sel = jn2.loc[\n", - " :,\n", - " [\n", - " \"s_acctbal\",\n", - " \"s_name\",\n", - " \"n_name\",\n", - " \"p_partkey\",\n", - " \"p_mfgr\",\n", - " \"s_address\",\n", - " \"s_phone\",\n", - " \"s_comment\",\n", - " ],\n", - " ]\n", - "\n", - " sort = sel.sort_values(\n", - " by=[\"s_acctbal\", \"n_name\", \"s_name\", \"p_partkey\"],\n", - " ascending=[False, True, True, True],\n", - " )\n", - " result_df = sort.head(100)\n", - "\n", - " return result_df # type: ignore[no-any-return]" - ] - }, { "cell_type": "code", "execution_count": null, @@ -117,26 +69,20 @@ "outputs": [], "source": [ "from typing import Any\n", - "from datetime import datetime\n", "import narwhals as nw\n", "\n", + "@nw.narwhalify\n", "def q2(\n", - " region_ds_raw: Any,\n", - " nation_ds_raw: Any,\n", - " supplier_ds_raw: Any,\n", - " part_ds_raw: Any,\n", - " part_supp_ds_raw: Any,\n", + " region_ds: Any,\n", + " nation_ds: Any,\n", + " supplier_ds: Any,\n", + " part_ds: Any,\n", + " part_supp_ds: Any,\n", ") -> Any:\n", " var_1 = 15\n", " var_2 = \"BRASS\"\n", " var_3 = \"EUROPE\"\n", "\n", - " region_ds = nw.from_native(region_ds_raw)\n", - " nation_ds = nw.from_native(nation_ds_raw)\n", - " supplier_ds = nw.from_native(supplier_ds_raw)\n", - " part_ds = nw.from_native(part_ds_raw)\n", - " part_supp_ds = nw.from_native(part_supp_ds_raw)\n", - "\n", " result_q2 = (\n", " part_ds.join(part_supp_ds, left_on=\"p_partkey\", right_on=\"ps_partkey\")\n", " .join(supplier_ds, left_on=\"ps_suppkey\", right_on=\"s_suppkey\")\n", @@ -160,9 +106,9 @@ " \"s_comment\",\n", " ]\n", "\n", - " q_final = (\n", + " return (\n", " result_q2.group_by(\"p_partkey\")\n", - " .agg(nw.min(\"ps_supplycost\").alias(\"ps_supplycost\"))\n", + " .agg(nw.col(\"ps_supplycost\").min().alias(\"ps_supplycost\"))\n", " .join(\n", " result_q2,\n", " left_on=[\"p_partkey\", \"ps_supplycost\"],\n", @@ -170,77 +116,17 @@ " )\n", " .select(final_cols)\n", " .sort(\n", - " by=[\"s_acctbal\", \"n_name\", \"s_name\", \"p_partkey\"],\n", + " [\"s_acctbal\", \"n_name\", \"s_name\", \"p_partkey\"],\n", " descending=[True, False, False, False],\n", " )\n", " .head(100)\n", - " )\n", - "\n", - " return nw.to_native(q_final)" + " )" ] }, { "cell_type": "code", "execution_count": null, "id": "4", - "metadata": {}, - "outputs": [], - "source": [ - "from typing import Any\n", - "from datetime import datetime\n", - "import ibis\n", - "\n", - "def q2_ibis(\n", - " region: Any,\n", - " nation: Any,\n", - " supplier: Any,\n", - " part: Any,\n", - " partsupp: Any,\n", - " *,\n", - " tool: str,\n", - ") -> Any:\n", - " var1 = 15\n", - " var2 = \"BRASS\"\n", - " var3 = \"EUROPE\"\n", - "\n", - " q2 = (\n", - " part.join(partsupp, part[\"p_partkey\"] == partsupp[\"ps_partkey\"])\n", - " .join(supplier, partsupp[\"ps_suppkey\"] == supplier[\"s_suppkey\"])\n", - " .join(nation, supplier[\"s_nationkey\"] == nation[\"n_nationkey\"])\n", - " .join(region, nation[\"n_regionkey\"] == region[\"r_regionkey\"])\n", - " .filter(ibis._[\"p_size\"] == var1)\n", - " .filter(ibis._[\"p_type\"].endswith(var2))\n", - " .filter(ibis._[\"r_name\"] == var3)\n", - " )\n", - "\n", - " q_final = (\n", - " q2.group_by(\"p_partkey\")\n", - " .agg(ps_supplycost=ibis._[\"ps_supplycost\"].min())\n", - " .join(q2, [\"p_partkey\"])\n", - " .select(\n", - " \"s_acctbal\",\n", - " \"s_name\",\n", - " \"n_name\",\n", - " \"p_partkey\",\n", - " \"p_mfgr\",\n", - " \"s_address\",\n", - " \"s_phone\",\n", - " \"s_comment\",\n", - " )\n", - " .order_by(ibis.desc(\"s_acctbal\"), \"n_name\", \"s_name\", \"p_partkey\")\n", - " .limit(100)\n", - " )\n", - " if tool == 'pandas':\n", - " return q_final.to_pandas()\n", - " if tool == 'polars':\n", - " return q_final.to_polars()\n", - " raise ValueError(\"expected pandas or polars\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", "metadata": { "papermill": { "duration": 0.013325, @@ -267,7 +153,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "5", "metadata": { "papermill": { "duration": 0.014284, @@ -280,25 +166,23 @@ }, "outputs": [], "source": [ - "import ibis\n", - "\n", - "con_pd = ibis.pandas.connect()\n", - "con_pl = ibis.polars.connect()\n", + "import pyarrow.parquet as pq\n", + "import dask.dataframe as dd\n", "\n", "IO_FUNCS = {\n", " 'pandas': lambda x: pd.read_parquet(x, engine='pyarrow'),\n", " 'pandas[pyarrow]': lambda x: pd.read_parquet(x, engine='pyarrow', dtype_backend='pyarrow'),\n", - " 'pandas[pyarrow][ibis]': lambda x: con_pd.read_parquet(x, engine='pyarrow', dtype_backend='pyarrow'),\n", " 'polars[eager]': lambda x: pl.read_parquet(x),\n", " 'polars[lazy]': lambda x: pl.scan_parquet(x),\n", - " 'polars[lazy][ibis]': lambda x: con_pl.read_parquet(x),\n", + " 'pyarrow': lambda x: pq.read_table(x),\n", + " 'dask': lambda x: dd.read_parquet(x, engine='pyarrow', dtype_backend='pyarrow'),\n", "}" ] }, { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -307,70 +191,7 @@ }, { "cell_type": "markdown", - "id": "8", - "metadata": {}, - "source": [ - "## pandas, pyarrow dtypes, via ibis" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9", - "metadata": {}, - "outputs": [], - "source": [ - "tool = 'pandas[pyarrow][ibis]'\n", - "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q2_ibis(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp), tool='pandas')\n", - "results[tool] = timings.all_runs" - ] - }, - { - "cell_type": "markdown", - "id": "10", - "metadata": {}, - "source": [ - "## Polars scan_parquet via ibis" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "11", - "metadata": {}, - "outputs": [], - "source": [ - "tool = 'polars[lazy][ibis]'\n", - "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q2_ibis(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp), tool='polars')\n", - "results[tool] = timings.all_runs" - ] - }, - { - "cell_type": "markdown", - "id": "12", - "metadata": {}, - "source": [ - "## pandas, pyarrow dtypes, native" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "13", - "metadata": {}, - "outputs": [], - "source": [ - "tool = 'pandas[pyarrow]'\n", - "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q2_pandas_native(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", - "results[tool+'[native]'] = timings.all_runs" - ] - }, - { - "cell_type": "markdown", - "id": "14", + "id": "7", "metadata": { "papermill": { "duration": 0.005113, @@ -382,13 +203,13 @@ "tags": [] }, "source": [ - "## pandas via Narwhals" + "## pandas" ] }, { "cell_type": "code", "execution_count": null, - "id": "15", + "id": "8", "metadata": { "papermill": { "duration": 196.786925, @@ -403,13 +224,13 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", + "timings = %timeit -o -q q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "16", + "id": "9", "metadata": { "papermill": { "duration": 0.005184, @@ -421,13 +242,13 @@ "tags": [] }, "source": [ - "## pandas, pyarrow dtypes, via Narwhals" + "## pandas, pyarrow dtypes" ] }, { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "10", "metadata": { "papermill": { "duration": 158.748353, @@ -442,13 +263,13 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", + "timings = %timeit -o -q q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "18", + "id": "11", "metadata": { "papermill": { "duration": 0.005773, @@ -466,7 +287,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "12", "metadata": { "papermill": { "duration": 37.821116, @@ -481,13 +302,13 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", + "timings = %timeit -o -q q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "20", + "id": "13", "metadata": { "papermill": { "duration": 0.005515, @@ -505,7 +326,7 @@ { "cell_type": "code", "execution_count": null, - "id": "21", + "id": "14", "metadata": { "papermill": { "duration": 4.800698, @@ -520,13 +341,55 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp)).collect()\n", + "timings = %timeit -o -q q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp)).collect()\n", + "results[tool] = timings.all_runs" + ] + }, + { + "cell_type": "markdown", + "id": "15", + "metadata": {}, + "source": [ + "## PyArrow" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "tool = 'pyarrow'\n", + "fn = IO_FUNCS[tool]\n", + "timings = %timeit -o -q q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp))\n", "results[tool] = timings.all_runs" ] }, { "cell_type": "markdown", - "id": "22", + "id": "17", + "metadata": {}, + "source": [ + "## Dask" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "tool = 'dask'\n", + "fn = IO_FUNCS[tool]\n", + "timings = %timeit -o -q q2(fn(region), fn(nation), fn(supplier), fn(part), fn(partsupp)).compute()\n", + "results[tool] = timings.all_runs" + ] + }, + { + "cell_type": "markdown", + "id": "19", "metadata": {}, "source": [ "## Save" @@ -535,7 +398,7 @@ { "cell_type": "code", "execution_count": null, - "id": "23", + "id": "20", "metadata": {}, "outputs": [], "source": [ @@ -586,7 +449,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.13" + "version": "3.11.9" }, "papermill": { "default_parameters": {}, diff --git a/tpch/notebooks/q20/execute.ipynb b/tpch/notebooks/q20/execute.ipynb index f0719f317..aecb3a473 100644 --- a/tpch/notebooks/q20/execute.ipynb +++ b/tpch/notebooks/q20/execute.ipynb @@ -195,7 +195,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q20(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(supplier))\n", + "timings = %timeit -o -q q20(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -232,7 +232,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q20(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(supplier))\n", + "timings = %timeit -o -q q20(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -269,7 +269,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q20(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(supplier))\n", + "timings = %timeit -o -q q20(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -306,7 +306,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q20(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(supplier)).collect()\n", + "timings = %timeit -o -q q20(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(supplier)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q21/execute.ipynb b/tpch/notebooks/q21/execute.ipynb index dc5063f52..b51b15dce 100755 --- a/tpch/notebooks/q21/execute.ipynb +++ b/tpch/notebooks/q21/execute.ipynb @@ -218,7 +218,7 @@ "\n", "lineitem_raw, nation_raw, orders_raw, supplier_raw = fn(lineitem), fn(nation), fn(orders), fn(supplier)\n", "\n", - "timings = %timeit -o q21(lineitem_raw, nation_raw, orders_raw, supplier_raw)\n", + "timings = %timeit -o -q q21(lineitem_raw, nation_raw, orders_raw, supplier_raw)\n", "results[tool] = timings.all_runs" ] }, @@ -259,7 +259,7 @@ "fn = IO_FUNCS[tool]\n", "lineitem_raw, nation_raw, orders_raw, supplier_raw = fn(lineitem), fn(nation), fn(orders), fn(supplier)\n", "\n", - "timings = %timeit -o q21(lineitem_raw, nation_raw, orders_raw, supplier_raw)\n", + "timings = %timeit -o -q q21(lineitem_raw, nation_raw, orders_raw, supplier_raw)\n", "results[tool] = timings.all_runs" ] }, @@ -300,7 +300,7 @@ "fn = IO_FUNCS[tool]\n", "\n", "lineitem_raw, nation_raw, orders_raw, supplier_raw = fn(lineitem), fn(nation), fn(orders), fn(supplier)\n", - "timings = %timeit -o q21(lineitem_raw, nation_raw, orders_raw, supplier_raw)\n", + "timings = %timeit -o -q q21(lineitem_raw, nation_raw, orders_raw, supplier_raw)\n", "results[tool] = timings.all_runs" ] }, @@ -341,7 +341,7 @@ "fn = IO_FUNCS[tool]\n", "\n", "lineitem_raw, nation_raw, orders_raw, supplier_raw = fn(lineitem), fn(nation), fn(orders), fn(supplier)\n", - "timings = %timeit -o q21(lineitem_raw, nation_raw, orders_raw, supplier_raw).collect()\n", + "timings = %timeit -o -q q21(lineitem_raw, nation_raw, orders_raw, supplier_raw).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q3/execute.ipynb b/tpch/notebooks/q3/execute.ipynb index f289ea913..80178cae1 100755 --- a/tpch/notebooks/q3/execute.ipynb +++ b/tpch/notebooks/q3/execute.ipynb @@ -16,7 +16,7 @@ }, "outputs": [], "source": [ - "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework " + "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals" ] }, { @@ -278,7 +278,7 @@ "source": [ "tool = 'pandas[pyarrow][ibis]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q3_ibis(fn(customer), fn(lineitem), fn(orders), tool='pandas')\n", + "timings = %timeit -o -q q3_ibis(fn(customer), fn(lineitem), fn(orders), tool='pandas')\n", "results[tool] = timings.all_runs" ] }, @@ -299,7 +299,7 @@ "source": [ "tool = 'polars[lazy][ibis]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q3_ibis(fn(customer), fn(lineitem), fn(orders), tool='polars')\n", + "timings = %timeit -o -q q3_ibis(fn(customer), fn(lineitem), fn(orders), tool='polars')\n", "results[tool] = timings.all_runs" ] }, @@ -320,7 +320,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q3_pandas_native(fn(customer), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q3_pandas_native(fn(customer), fn(lineitem), fn(orders))\n", "results[tool+'[native]'] = timings.all_runs" ] }, @@ -359,7 +359,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q3(fn(customer), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q3(fn(customer), fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -398,7 +398,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q3(fn(customer), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q3(fn(customer), fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -437,7 +437,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q3(fn(customer), fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q3(fn(customer), fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -476,7 +476,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q3(fn(customer), fn(lineitem), fn(orders)).collect()\n", + "timings = %timeit -o -q q3(fn(customer), fn(lineitem), fn(orders)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q4/execute.ipynb b/tpch/notebooks/q4/execute.ipynb index f5d1b97bd..df07c9c5f 100755 --- a/tpch/notebooks/q4/execute.ipynb +++ b/tpch/notebooks/q4/execute.ipynb @@ -16,7 +16,7 @@ }, "outputs": [], "source": [ - "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework " + "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals" ] }, { @@ -243,7 +243,7 @@ "source": [ "tool = 'polars[lazy][ibis]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q4_ibis(fn(lineitem), fn(orders), tool='polars')\n", + "timings = %timeit -o -q q4_ibis(fn(lineitem), fn(orders), tool='polars')\n", "results[tool] = timings.all_runs" ] }, @@ -264,7 +264,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q4_pandas_native(fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q4_pandas_native(fn(lineitem), fn(orders))\n", "results[tool+'[native]'] = timings.all_runs" ] }, @@ -303,7 +303,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q4(fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q4(fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -342,7 +342,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q4(fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q4(fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -381,7 +381,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q4(fn(lineitem), fn(orders))\n", + "timings = %timeit -o -q q4(fn(lineitem), fn(orders))\n", "results[tool] = timings.all_runs" ] }, @@ -420,7 +420,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q4(fn(lineitem), fn(orders)).collect()\n", + "timings = %timeit -o -q q4(fn(lineitem), fn(orders)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q5/execute.ipynb b/tpch/notebooks/q5/execute.ipynb index a56ae03d1..5f6df9bbc 100755 --- a/tpch/notebooks/q5/execute.ipynb +++ b/tpch/notebooks/q5/execute.ipynb @@ -16,7 +16,7 @@ }, "outputs": [], "source": [ - "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework " + "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals" ] }, { @@ -275,7 +275,7 @@ "source": [ "tool = 'polars[lazy][ibis]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q5_ibis(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier), tool='polars')\n", + "timings = %timeit -o -q q5_ibis(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier), tool='polars')\n", "results[tool] = timings.all_runs" ] }, @@ -296,7 +296,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q5_pandas_native(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q5_pandas_native(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool+'[native]'] = timings.all_runs" ] }, @@ -335,7 +335,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -374,7 +374,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -413,7 +413,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -452,7 +452,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier)).collect()\n", + "timings = %timeit -o -q q5(fn(region), fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q6/execute.ipynb b/tpch/notebooks/q6/execute.ipynb index 0f8d6ce58..b101aa98d 100755 --- a/tpch/notebooks/q6/execute.ipynb +++ b/tpch/notebooks/q6/execute.ipynb @@ -16,7 +16,7 @@ }, "outputs": [], "source": [ - "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework " + "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals" ] }, { @@ -231,7 +231,7 @@ "source": [ "tool = 'pandas[pyarrow][ibis]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q6_ibis(fn(lineitem), tool='pandas')\n", + "timings = %timeit -o -q q6_ibis(fn(lineitem), tool='pandas')\n", "results[tool] = timings.all_runs" ] }, @@ -252,7 +252,7 @@ "source": [ "tool = 'polars[lazy][ibis]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q6_ibis(fn(lineitem), tool='polars')\n", + "timings = %timeit -o -q q6_ibis(fn(lineitem), tool='polars')\n", "results[tool] = timings.all_runs" ] }, @@ -273,7 +273,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q6_pandas_native(fn(lineitem))\n", + "timings = %timeit -o -q q6_pandas_native(fn(lineitem))\n", "results[tool+'[native]'] = timings.all_runs" ] }, @@ -312,7 +312,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q6(fn(lineitem))\n", + "timings = %timeit -o -q q6(fn(lineitem))\n", "results[tool] = timings.all_runs" ] }, @@ -351,7 +351,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q6(fn(lineitem))\n", + "timings = %timeit -o -q q6(fn(lineitem))\n", "results[tool] = timings.all_runs" ] }, @@ -390,7 +390,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q6(fn(lineitem))\n", + "timings = %timeit -o -q q6(fn(lineitem))\n", "results[tool] = timings.all_runs" ] }, @@ -429,7 +429,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q6(fn(lineitem)).collect()\n", + "timings = %timeit -o -q q6(fn(lineitem)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q7/execute.ipynb b/tpch/notebooks/q7/execute.ipynb index 3b64df2fc..1213043b0 100755 --- a/tpch/notebooks/q7/execute.ipynb +++ b/tpch/notebooks/q7/execute.ipynb @@ -326,7 +326,7 @@ "source": [ "tool = 'pandas[pyarrow][ibis]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7_ibis(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier), tool='pandas')\n", + "timings = %timeit -o -q q7_ibis(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier), tool='pandas')\n", "results[tool] = timings.all_runs" ] }, @@ -347,7 +347,7 @@ "source": [ "tool = 'polars[lazy][ibis]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7_ibis(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier), tool='polars')\n", + "timings = %timeit -o -q q7_ibis(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier), tool='polars')\n", "results[tool] = timings.all_runs" ] }, @@ -368,7 +368,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7_pandas_native(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q7_pandas_native(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool+'[native]'] = timings.all_runs" ] }, @@ -407,7 +407,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -446,7 +446,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -485,7 +485,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -524,7 +524,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier)).collect()\n", + "timings = %timeit -o -q q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q8/execute.ipynb b/tpch/notebooks/q8/execute.ipynb index 531cad195..b10b87907 100755 --- a/tpch/notebooks/q8/execute.ipynb +++ b/tpch/notebooks/q8/execute.ipynb @@ -260,7 +260,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7_pandas_native(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q7_pandas_native(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool+'[native]'] = timings.all_runs" ] }, @@ -299,7 +299,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -338,7 +338,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -377,7 +377,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -416,7 +416,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier)).collect()\n", + "timings = %timeit -o -q q7(fn(nation), fn(customer), fn(lineitem), fn(orders), fn(supplier)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/tpch/notebooks/q9/execute.ipynb b/tpch/notebooks/q9/execute.ipynb index d7412426c..86417e180 100644 --- a/tpch/notebooks/q9/execute.ipynb +++ b/tpch/notebooks/q9/execute.ipynb @@ -190,7 +190,7 @@ "source": [ "tool = 'pandas'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q9(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q9(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -227,7 +227,7 @@ "source": [ "tool = 'pandas[pyarrow]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q9(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q9(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -264,7 +264,7 @@ "source": [ "tool = 'polars[eager]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q9(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(orders), fn(supplier))\n", + "timings = %timeit -o -q q9(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(orders), fn(supplier))\n", "results[tool] = timings.all_runs" ] }, @@ -301,7 +301,7 @@ "source": [ "tool = 'polars[lazy]'\n", "fn = IO_FUNCS[tool]\n", - "timings = %timeit -o q9(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(orders), fn(supplier)).collect()\n", + "timings = %timeit -o -q q9(fn(part), fn(partsupp), fn(nation), fn(lineitem), fn(orders), fn(supplier)).collect()\n", "results[tool] = timings.all_runs" ] }, diff --git a/utils/api-completeness.md.jinja b/utils/api-completeness.md.jinja index 782edcd57..d79f157c6 100644 --- a/utils/api-completeness.md.jinja +++ b/utils/api-completeness.md.jinja @@ -1,14 +1,3 @@ -# API Completeness - -Narwhals has two different level of support for libraries: "full" and "interchange". - -Libraries for which we have full support we intend to support the whole Narwhals API, however this is a work in progress. - -In the following table it is possible to check which method is implemented for which backend. - -!!! info - - - "pandas-like" means pandas, cuDF and Modin - - Polars supports all the methods (by design) +# {{ title }} {{ backend_table }} diff --git a/utils/check_api_reference.py b/utils/check_api_reference.py index 80ee5d7aa..f6e5303c4 100644 --- a/utils/check_api_reference.py +++ b/utils/check_api_reference.py @@ -45,13 +45,13 @@ documented = [ remove_prefix(i, " - ") for i in content.splitlines() - if i.startswith(" - ") + if i.startswith(" - ") and not i.startswith(" - _") ] if missing := set(top_level_functions).difference(documented): print("DataFrame: not documented") # noqa: T201 print(missing) # noqa: T201 ret = 1 -if extra := set(documented).difference(top_level_functions).difference({"__getitem__"}): +if extra := set(documented).difference(top_level_functions): print("DataFrame: outdated") # noqa: T201 print(extra) # noqa: T201 ret = 1 @@ -87,7 +87,7 @@ documented = [ remove_prefix(i, " - ") for i in content.splitlines() - if i.startswith(" - ") + if i.startswith(" - ") and not i.startswith(" - _") ] if ( missing := set(top_level_functions) @@ -148,6 +148,7 @@ .difference(expr) .difference( { + "to_arrow", "to_dummies", "to_pandas", "to_list", diff --git a/utils/generate_backend_completeness.py b/utils/generate_backend_completeness.py index f7d2df50c..537701872 100644 --- a/utils/generate_backend_completeness.py +++ b/utils/generate_backend_completeness.py @@ -2,18 +2,40 @@ import importlib import inspect +from enum import Enum +from enum import auto from pathlib import Path from typing import Any from typing import Final +from typing import NamedTuple import polars as pl from jinja2 import Template TEMPLATE_PATH: Final[Path] = Path("utils") / "api-completeness.md.jinja" -DESTINATION_PATH: Final[Path] = Path("docs") / "api-completeness.md" +DESTINATION_PATH: Final[Path] = Path("docs") / "api-completeness" + + +class BackendType(Enum): + LAZY = auto() + EAGER = auto() + BOTH = auto() + + +class Backend(NamedTuple): + name: str + module: str + type_: BackendType MODULES = ["dataframe", "series", "expr"] + +BACKENDS = [ + Backend(name="pandas-like", module="_pandas_like", type_=BackendType.EAGER), + Backend(name="arrow", module="_arrow", type_=BackendType.EAGER), + Backend(name="dask", module="_dask", type_=BackendType.LAZY), +] + EXCLUDE_CLASSES = {"BaseFrame"} @@ -21,12 +43,26 @@ def get_class_methods(kls: type[Any]) -> list[str]: return [m[0] for m in inspect.getmembers(kls) if not m[0].startswith("_")] -def get_backend_completeness_table() -> str: - results = [] +def parse_module(module_name: str, backend: str, nw_class_name: str) -> list[str]: + try: + module_ = importlib.import_module(f"narwhals.{backend}.{module_name}") + class_ = inspect.getmembers( + module_, + predicate=lambda c: inspect.isclass(c) and c.__name__.endswith(nw_class_name), + ) + methods_ = get_class_methods(class_[0][1]) if class_ else [] + + except ModuleNotFoundError: + methods_ = [] + + return methods_ + +def get_backend_completeness_table() -> None: for module_name in MODULES: + results = [] + nw_namespace = f"narwhals.{module_name}" - sub_module_name = module_name narwhals_module_ = importlib.import_module(nw_namespace) classes_ = inspect.getmembers( @@ -37,71 +73,56 @@ def get_backend_completeness_table() -> str: for nw_class_name, nw_class in classes_: if nw_class_name in EXCLUDE_CLASSES: continue - if nw_class_name == "LazyFrame": - backend_class_name = "DataFrame" - else: - backend_class_name = nw_class_name - - arrow_class_name = f"Arrow{backend_class_name}" - arrow_module_ = importlib.import_module(f"narwhals._arrow.{sub_module_name}") - arrow_class = inspect.getmembers( - arrow_module_, - predicate=lambda c: inspect.isclass(c) and c.__name__ == arrow_class_name, # noqa: B023 - ) - - pandas_class_name = f"PandasLike{backend_class_name}" - pandas_module_ = importlib.import_module( - f"narwhals._pandas_like.{sub_module_name}" - ) - pandas_class = inspect.getmembers( - pandas_module_, - predicate=lambda c: inspect.isclass(c) - and c.__name__ == pandas_class_name, # noqa: B023 - ) nw_methods = get_class_methods(nw_class) - arrow_methods = get_class_methods(arrow_class[0][1]) if arrow_class else [] - pandas_methods = get_class_methods(pandas_class[0][1]) if pandas_class else [] narhwals = pl.DataFrame( {"Class": nw_class_name, "Backend": "narwhals", "Method": nw_methods} ) - arrow = pl.DataFrame( - {"Class": nw_class_name, "Backend": "arrow", "Method": arrow_methods} - ) - pandas = pl.DataFrame( - { - "Class": nw_class_name, - "Backend": "pandas-like", - "Method": pandas_methods, - } - ) - - results.extend([narhwals, pandas, arrow]) - - results = ( - pl.concat(results) # noqa: PD010 - .with_columns(supported=pl.lit(":white_check_mark:")) - .pivot(on="Backend", values="supported", index=["Class", "Method"]) - .filter(pl.col("narwhals").is_not_null()) - .drop("narwhals") - .fill_null(":x:") - .sort("Class", "Method") - ) - with pl.Config( - tbl_formatting="ASCII_MARKDOWN", - tbl_hide_column_data_types=True, - tbl_hide_dataframe_shape=True, - set_tbl_rows=results.shape[0], - ): - return str(results) + backend_methods = [ + pl.DataFrame( + { + "Class": nw_class_name, + "Backend": backend.name, + "Method": parse_module( + module_name, + backend=backend.module, + nw_class_name=nw_class_name, + ), + # "Type": backend.type_ + } + ) + for backend in BACKENDS + ] + + results.extend([narhwals, *backend_methods]) + + results = ( + pl.concat(results) # noqa: PD010 + .with_columns(supported=pl.lit(":white_check_mark:")) + .pivot(on="Backend", values="supported", index=["Class", "Method"]) + .filter(pl.col("narwhals").is_not_null()) + .drop("narwhals") + .fill_null(":x:") + .sort("Class", "Method") + ) + with pl.Config( + tbl_formatting="ASCII_MARKDOWN", + tbl_hide_column_data_types=True, + tbl_hide_dataframe_shape=True, + set_tbl_rows=results.shape[0], + ): + table = str(results) + + with TEMPLATE_PATH.open(mode="r") as stream: + new_content = Template(stream.read()).render( + {"backend_table": table, "title": module_name.capitalize()} + ) -backend_table = get_backend_completeness_table() + with (DESTINATION_PATH / f"{module_name}.md").open(mode="w") as destination: + destination.write(new_content) -with TEMPLATE_PATH.open(mode="r") as stream: - new_content = Template(stream.read()).render({"backend_table": backend_table}) -with DESTINATION_PATH.open(mode="w") as destination: - destination.write(new_content) +_ = get_backend_completeness_table() diff --git a/utils/import_check.py b/utils/import_check.py new file mode 100644 index 000000000..e8d776cde --- /dev/null +++ b/utils/import_check.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import ast +import sys + +BANNED_IMPORTS = { + "cudf", + "dask", + "dask.dataframe", + "dask_expr", + "duckdb", + "ibis", + "modin", + "numpy", + "pandas", + "polars", + "pyarrow", +} + + +class ImportPandasChecker(ast.NodeVisitor): + def __init__(self, file_name: str, lines: list[str]) -> None: + self.file_name = file_name + self.lines = lines + self.found_import = False + + def visit_If(self, node: ast.If) -> None: # noqa: N802 + # Check if the condition is `if TYPE_CHECKING` + if isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING": + # Skip the body of this if statement + return + self.generic_visit(node) + + def visit_Import(self, node: ast.Import) -> None: # noqa: N802 + for alias in node.names: + if ( + alias.name in BANNED_IMPORTS + and "# ignore-banned-import" not in self.lines[node.lineno - 1] + ): + print( # noqa: T201 + f"{self.file_name}:{node.lineno}:{node.col_offset}: found {alias.name} import" + ) + self.found_import = True + self.generic_visit(node) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: # noqa: N802 + if ( + node.module in BANNED_IMPORTS + and "# ignore-banned-import" not in self.lines[node.lineno - 1] + ): + print( # noqa: T201 + f"{self.file_name}:{node.lineno}:{node.col_offset}: found {node.module} import" + ) + self.found_import = True + self.generic_visit(node) + + +def check_import_pandas(filename: str) -> bool: + with open(filename) as file: + content = file.read() + tree = ast.parse(content, filename=filename) + + checker = ImportPandasChecker(filename, content.splitlines()) + checker.visit(tree) + + return checker.found_import + + +if __name__ == "__main__": + ret = 0 + for filename in sys.argv[1:]: + if not filename.endswith(".py"): + continue + ret |= check_import_pandas(filename) + sys.exit(ret)