diff --git a/poetry.lock b/poetry.lock index 53e3157..b84bf87 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1589,6 +1589,39 @@ files = [ dev = ["pre-commit", "tox"] testing = ["pytest", "pytest-benchmark"] +[[package]] +name = "polars" +version = "0.17.2" +description = "Blazingly fast DataFrame library" +category = "dev" +optional = false +python-versions = ">=3.7" +files = [ + {file = "polars-0.17.2-cp37-abi3-macosx_10_7_x86_64.whl", hash = "sha256:4631d29af25cebdec89d006afa9fee868138c1be91d6070dc7c8e1f3e9684287"}, + {file = "polars-0.17.2-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:cb69717eebd52d289eb5bf50dc1e90f380874280c4ef3c25b56b6fbc22b706bc"}, + {file = "polars-0.17.2-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e01317ea1902cc6ae327598b6998c6c6b8bcd6ff6dff15bad7a78c51231220fa"}, + {file = "polars-0.17.2-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:96c2e1c99738634eaeb510a724bf1aecc4ecb030cb45a28595d9ceac842300e9"}, + {file = "polars-0.17.2-cp37-abi3-win_amd64.whl", hash = "sha256:a6d1ba5396a241eabf5ab755c7fb9f18a92bebaa8be5327d7ee1f532823aaba7"}, + {file = "polars-0.17.2.tar.gz", hash = "sha256:f121ad3515dfc4c0276067870c20585f5e4fa8368efc9cf7ff8c56da4372e312"}, +] + +[package.dependencies] +typing_extensions = {version = ">=4.0.1", markers = "python_version < \"3.11\""} + +[package.extras] +all = ["polars[connectorx,deltalake,fsspec,matplotlib,numpy,pandas,pyarrow,sqlalchemy,timezone,xlsx2csv,xlsxwriter]"] +connectorx = ["connectorx"] +deltalake = ["deltalake (>=0.8.0)"] +fsspec = ["fsspec"] +matplotlib = ["matplotlib"] +numpy = ["numpy (>=1.16.0)"] +pandas = ["pandas", "pyarrow (>=7.0.0)"] +pyarrow = ["pyarrow (>=7.0.0)"] +sqlalchemy = ["pandas", "sqlalchemy"] +timezone = ["backports.zoneinfo", "tzdata"] +xlsx2csv = ["xlsx2csv (>=0.8.0)"] +xlsxwriter = ["xlsxwriter"] + [[package]] name = "posthog" version = "2.5.0" @@ -2892,4 +2925,4 @@ testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "cfdc44025f2b53a5d4d7eb2614f1ff4587daef1009b90295466ba289edf68276" +content-hash = "01c945cba3ebafdbf6775ce94108604b5dc9ee8926a7a1872173194671d0cbf1" diff --git a/pyproject.toml b/pyproject.toml index 9fcc808..af56aa8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ python = "^3.8" duckdb-engine = "^0.7" ipython = "^8" -jupysql = ">=0.6" +jupysql = ">=0.7" pandas = ">=1.5" prql-python = "^0.8" traitlets = "^5" @@ -54,6 +54,7 @@ ruff = "^0.0.225" safety = "^2" sphinx-rtd-theme = "^1.0.0" xdoctest = "^0.15.10" +polars = ">=0.16" [tool.poetry.scripts] pyprql = "pyprql.cli.__init__:main" diff --git a/pyprql/magic/README.md b/pyprql/magic/README.md index 92d89b6..76de5c3 100644 --- a/pyprql/magic/README.md +++ b/pyprql/magic/README.md @@ -159,6 +159,9 @@ PrqlMagic.autolimit= PrqlMagic.autopandas= Return Pandas DataFrames instead of regular result sets Current: True +PrqlMagic.autopolars= + Return Polars DataFrames instead of regular result sets + Current: False PrqlMagic.autoview= Display results Current: True @@ -183,6 +186,10 @@ PrqlMagic.dsn_filename= PrqlMagic.feedback= Print number of rows affected by DML Current: False +PrqlMagic.polars_dataframe_kwargs==... + Polars DataFrame constructor keyword arguments(e.g. infer_schema_length, + nan_to_null, schema_overrides, etc) + Current: {} PrqlMagic.short_errors= Don't display the full traceback on SQL Programming Error Current: True diff --git a/pyprql/magic/prql.py b/pyprql/magic/prql.py index e25a4ed..1cf44e9 100644 --- a/pyprql/magic/prql.py +++ b/pyprql/magic/prql.py @@ -39,6 +39,11 @@ class PrqlMagic(SqlMagic): config=True, help="Return Pandas DataFrames instead of regular result sets", ) + autopolars = Bool( + False, + config=True, + help="Return Polars DataFrames instead of regular result sets", + ) autoview = Bool(True, config=True, help="Display results") feedback = Bool(False, config=True, help="Print number of rows affected by DML") target = Unicode("sql.any", config=True, help="Compile target of prql-compiler") diff --git a/pyprql/tests/test_magic.py b/pyprql/tests/test_magic.py index 2b9204a..55ec122 100644 --- a/pyprql/tests/test_magic.py +++ b/pyprql/tests/test_magic.py @@ -14,6 +14,7 @@ import tempfile from textwrap import dedent +import polars as pl import pytest from sqlalchemy import create_engine @@ -262,13 +263,22 @@ def test_bind_vars(ip): def test_autopandas(ip): - ip.run_line_magic("config", "SqlMagic.autopandas = True") dframe = run_prql(ip, "from test") assert not dframe.empty assert dframe.ndim == 2 assert dframe.name[0] == "foo" +def test_autopolars(ip): + ip.run_line_magic("config", "PrqlMagic.autopolars = True") + dframe = run_prql(ip, "from test") + + assert type(dframe) == pl.DataFrame + assert not dframe.is_empty() + assert len(dframe.shape) == 2 + assert dframe["name"][0] == "foo" + + def test_target_dialect(ip): ip.run_line_magic("config", 'PrqlMagic.target = "sql.sqlite"') dframe = run_prql(