Skip to content

Commit

Permalink
feat(python!): Use Altair in DataFrame.plot
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Aug 1, 2024
1 parent 5134051 commit 872d784
Show file tree
Hide file tree
Showing 5 changed files with 155 additions and 71 deletions.
45 changes: 23 additions & 22 deletions py-polars/polars/dataframe/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,14 +83,14 @@
from polars.datatypes.group import INTEGER_DTYPES
from polars.dependencies import (
_GREAT_TABLES_AVAILABLE,
_HVPLOT_AVAILABLE,
_ALTAIR_AVAILABLE,
_PANDAS_AVAILABLE,
_PYARROW_AVAILABLE,
_check_for_numpy,
_check_for_pandas,
_check_for_pyarrow,
great_tables,
hvplot,
altair,
import_optional,
)
from polars.dependencies import numpy as np
Expand All @@ -106,6 +106,7 @@
from polars.interchange.protocol import CompatLevel
from polars.schema import Schema
from polars.selectors import _expand_selector_dicts, _expand_selectors
from polars.dataframe.plotting import Plot

with contextlib.suppress(ImportError): # Module not available when building docs
from polars.polars import PyDataFrame, PySeries
Expand All @@ -123,8 +124,8 @@
import numpy.typing as npt
import torch
from great_tables import GT
from hvplot.plotting.core import hvPlotTabularPolars
from xlsxwriter import Workbook
from polars.dataframe.plotting import Plot

from polars import DataType, Expr, LazyFrame, Series
from polars._typing import (
Expand Down Expand Up @@ -603,17 +604,25 @@ def _replace(self, column: str, new_column: Series) -> DataFrame:

@property
@unstable()
def plot(self) -> hvPlotTabularPolars:
def plot(self) -> Plot:
"""
Create a plot namespace.
.. warning::
This functionality is currently considered **unstable**. It may be
changed at any point without it being considered a breaking change.
.. versionchanged:: 1.4.0
In prior versions of Polars, HvPlot was the plotting backend. If you would like to restore
the previous plotting functionality, all you need to do add `import hvplot.polars` at the
top of your script and replace `df.plot` with `df.hvplot`.
Polars does not implement plotting logic itself, but instead defers to
hvplot. Please see the `hvplot reference gallery <https://hvplot.holoviz.org/reference/index.html>`_
for more information and documentation.
Altair:
- `df.plot.line(*args, **kwargs)` is shorthand for `alt.Chart(df).mark_line().encode(*args, **kwargs)`
- `df.plot.point(*args, **kwargs)` is shorthand for `alt.Chart(df).mark_point().encode(*args, **kwargs)`
- ...
Examples
--------
Expand All @@ -626,32 +635,24 @@ def plot(self) -> hvPlotTabularPolars:
... "species": ["setosa", "setosa", "versicolor"],
... }
... )
>>> df.plot.scatter(x="length", y="width", by="species") # doctest: +SKIP
>>> df.plot.point(x="length", y="width", color="species") # doctest: +SKIP
Line plot:
>>> from datetime import date
>>> df = pl.DataFrame(
... {
... "date": [date(2020, 1, 2), date(2020, 1, 3), date(2020, 1, 4)],
... "stock_1": [1, 4, 6],
... "stock_2": [1, 5, 2],
... "date": [date(2020, 1, 2), date(2020, 1, 3), date(2020, 1, 4)]*2,
... "price": [1, 4, 6, 1, 5, 2],
... "stock": ['a', 'a', 'a', 'b', 'b', 'b'],
... }
... )
>>> df.plot.line(x="date", y=["stock_1", "stock_2"]) # doctest: +SKIP
For more info on what you can pass, you can use ``hvplot.help``:
>>> import hvplot # doctest: +SKIP
>>> hvplot.help("scatter") # doctest: +SKIP
>>> df.plot.line(x="date", y="price", color="stock") # doctest: +SKIP
"""
if not _HVPLOT_AVAILABLE or parse_version(hvplot.__version__) < parse_version(
"0.9.1"
):
msg = "hvplot>=0.9.1 is required for `.plot`"
if not _ALTAIR_AVAILABLE or parse_version(altair.__version__) < (5, 3, 0):
msg = "altair>=5.3.0 is required for `.plot`"
raise ModuleUpgradeRequiredError(msg)
hvplot.post_patch()
return hvplot.plotting.core.hvPlotTabularPolars(self)
return Plot(self)

@property
@unstable()
Expand Down
123 changes: 123 additions & 0 deletions py-polars/polars/dataframe/plotting.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
from __future__ import annotations

from typing import TYPE_CHECKING
from typing import Any
from altair import Undefined


if TYPE_CHECKING:
from polars import DataFrame
import altair as alt


class Plot:
chart: alt.Chart

def __init__(self, df: DataFrame) -> None:
import altair as alt
self.chart = alt.Chart(df)

def line(
self,
x: str | Any,
y: str | Any,
color: str | Any,
order: str | Any,
tooltip: str | Any,
) -> alt.Chart:
"""
Draw line plot.
Polars does not implement plottinng logic itself but instead defers to Altair.
`df.plot.line(*args, **kwargs)` is shorthand for `alt.Chart(df).mark_line().encode(*args, **kwargs)`,
as is intended for convenience - for full customisatibility, use Altair directly.
.. versionchanged:: 1.4.0
In prior versions of Polars, HvPlot was the plotting backend. If you would like to restore
the previous plotting functionality, all you need to do add `import hvplot.polars` at the
top of your script and replace `df.plot` with `df.hvplot`.
Parameters
----------
x
Column with x-coordinates of lines.
y
Column with y-coordinates of lines.
color
Column to color lines by.
order
Column to use for order of data points in lines.
tooltip
Columns to show values of when hovering over points with pointer.
*args, **kwargs
Additional arguments and keyword arguments passed to Altair.
"""
encodings = {}
if x is not Undefined:
encodings['x'] = x
if y is not Undefined:
encodings['y'] = y
if color is not Undefined:
encodings['color'] = color
if order is not Undefined:
encodings['order'] = order
if tooltip is not Undefined:
encodings['tooltip'] = tooltip
return self.chart.mark_point().encode(*args, **{**encodings, **kwargs})

def point(
self,
x: str | Any=Undefined,
y: str | Any=Undefined,
color: str | Any=Undefined,
size: str | Any=Undefined,
tooltip: str | Any=Undefined,
*args,
**kwargs
) -> alt.Chart:
"""
Draw scatter plot.
Polars does not implement plottinng logic itself but instead defers to Altair.
`df.plot.point(*args, **kwargs)` is shorthand for `alt.Chart(df).mark_point().encode(*args, **kwargs)`,
as is intended for convenience - for full customisatibility, use Altair directly.
.. versionchanged:: 1.4.0
In prior versions of Polars, HvPlot was the plotting backend. If you would like to restore
the previous plotting functionality, all you need to do add `import hvplot.polars` at the
top of your script and replace `df.plot` with `df.hvplot`.
Parameters
----------
x
Column with x-coordinates of points.
y
Column with y-coordinates of points.
color
Column to color points by.
color
Column which determines points' sizes.
tooltip
Columns to show values of when hovering over points with pointer.
*args, **kwargs
Additional arguments and keyword arguments passed to Altair.
"""
encodings = {}
if x is not Undefined:
encodings['x'] = x
if y is not Undefined:
encodings['y'] = y
if color is not Undefined:
encodings['color'] = color
if size is not Undefined:
encodings['size'] = size
if tooltip is not Undefined:
encodings['tooltip'] = tooltip
return self.chart.mark_line().encode(*args, **{**encodings, **kwargs})

def __getattr__(self, attr: str, *args, **kwargs) -> alt.Chart:
method = self.chart.getattr(f'mark_{attr}', None)
if method is None:
msg = "Altair has no method 'mark_{attr}'"
raise AttributeError(msg)
return method().encode(*args, **kwargs)
10 changes: 5 additions & 5 deletions py-polars/polars/dependencies.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,11 @@
from types import ModuleType
from typing import TYPE_CHECKING, Any, ClassVar, Hashable, cast

_ALTAIR_AVAILABLE = True
_DELTALAKE_AVAILABLE = True
_FSSPEC_AVAILABLE = True
_GEVENT_AVAILABLE = True
_GREAT_TABLES_AVAILABLE = True
_HVPLOT_AVAILABLE = True
_HYPOTHESIS_AVAILABLE = True
_NUMPY_AVAILABLE = True
_PANDAS_AVAILABLE = True
Expand Down Expand Up @@ -154,7 +154,7 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]:
import fsspec
import gevent
import great_tables
import hvplot
import altair
import hypothesis
import numpy
import pandas
Expand All @@ -175,10 +175,10 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]:
subprocess, _ = _lazy_import("subprocess")

# heavy/optional third party libs
altair, _ALTAIR_AVAILABLE = _lazy_import("altair")
deltalake, _DELTALAKE_AVAILABLE = _lazy_import("deltalake")
fsspec, _FSSPEC_AVAILABLE = _lazy_import("fsspec")
great_tables, _GREAT_TABLES_AVAILABLE = _lazy_import("great_tables")
hvplot, _HVPLOT_AVAILABLE = _lazy_import("hvplot")
hypothesis, _HYPOTHESIS_AVAILABLE = _lazy_import("hypothesis")
numpy, _NUMPY_AVAILABLE = _lazy_import("numpy")
pandas, _PANDAS_AVAILABLE = _lazy_import("pandas")
Expand Down Expand Up @@ -301,11 +301,11 @@ def import_optional(
"pickle",
"subprocess",
# lazy-load third party libs
"altair",
"deltalake",
"fsspec",
"gevent",
"great_tables",
"hvplot",
"numpy",
"pandas",
"pydantic",
Expand All @@ -318,11 +318,11 @@ def import_optional(
"_check_for_pyarrow",
"_check_for_pydantic",
# exported flags/guards
"_ALTAIR_AVAILABLE",
"_DELTALAKE_AVAILABLE",
"_PYICEBERG_AVAILABLE",
"_FSSPEC_AVAILABLE",
"_GEVENT_AVAILABLE",
"_HVPLOT_AVAILABLE",
"_HYPOTHESIS_AVAILABLE",
"_NUMPY_AVAILABLE",
"_PANDAS_AVAILABLE",
Expand Down
4 changes: 2 additions & 2 deletions py-polars/polars/meta/versions.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,13 @@ def show_versions() -> None:
Python: 3.11.8 (main, Feb 6 2024, 21:21:21) [Clang 15.0.0 (clang-1500.1.0.2.5)]
----Optional dependencies----
adbc_driver_manager: 0.11.0
altair: 5.3.0
cloudpickle: 3.0.0
connectorx: 0.3.2
deltalake: 0.17.1
fastexcel: 0.10.4
fsspec: 2023.12.2
gevent: 24.2.1
hvplot: 0.9.2
matplotlib: 3.8.4
nest_asyncio: 1.6.0
numpy: 1.26.4
Expand Down Expand Up @@ -63,14 +63,14 @@ def _get_dependency_info() -> dict[str, str]:
# see the list of dependencies in pyproject.toml
opt_deps = [
"adbc_driver_manager",
"altair",
"cloudpickle",
"connectorx",
"deltalake",
"fastexcel",
"fsspec",
"gevent",
"great_tables",
"hvplot",
"matplotlib",
"nest_asyncio",
"numpy",
Expand Down
44 changes: 2 additions & 42 deletions py-polars/polars/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,12 +86,12 @@
)
from polars.datatypes._utils import dtype_to_init_repr
from polars.dependencies import (
_HVPLOT_AVAILABLE,
_ALTAIR_AVAILABLE,
_PYARROW_AVAILABLE,
_check_for_numpy,
_check_for_pandas,
_check_for_pyarrow,
hvplot,
altair,
import_optional,
)
from polars.dependencies import numpy as np
Expand All @@ -117,7 +117,6 @@
import jax
import numpy.typing as npt
import torch
from hvplot.plotting.core import hvPlotTabularPolars

from polars import DataFrame, DataType, Expr
from polars._typing import (
Expand Down Expand Up @@ -7378,45 +7377,6 @@ def struct(self) -> StructNameSpace:
"""Create an object namespace of all struct related methods."""
return StructNameSpace(self)

@property
@unstable()
def plot(self) -> hvPlotTabularPolars:
"""
Create a plot namespace.
.. warning::
This functionality is currently considered **unstable**. It may be
changed at any point without it being considered a breaking change.
Polars does not implement plotting logic itself, but instead defers to
hvplot. Please see the `hvplot reference gallery <https://hvplot.holoviz.org/reference/index.html>`_
for more information and documentation.
Examples
--------
Histogram:
>>> s = pl.Series("values", [1, 4, 2])
>>> s.plot.hist() # doctest: +SKIP
KDE plot (note: in addition to ``hvplot``, this one also requires ``scipy``):
>>> s.plot.kde() # doctest: +SKIP
For more info on what you can pass, you can use ``hvplot.help``:
>>> import hvplot # doctest: +SKIP
>>> hvplot.help("hist") # doctest: +SKIP
"""
if not _HVPLOT_AVAILABLE or parse_version(hvplot.__version__) < parse_version(
"0.9.1"
):
msg = "hvplot>=0.9.1 is required for `.plot`"
raise ModuleUpgradeRequiredError(msg)
hvplot.post_patch()
return hvplot.plotting.core.hvPlotTabularPolars(self)


def _resolve_temporal_dtype(
dtype: PolarsDataType | None,
ndtype: np.dtype[np.datetime64] | np.dtype[np.timedelta64],
Expand Down

0 comments on commit 872d784

Please sign in to comment.