bzm3r's questions

考虑以下小示例（基于此图库示例）：

import altair as alt
from vega_datasets import data
import polars as pl

# add a column indicating the year associated with each date
source = pl.from_pandas(data.stocks()).with_columns(year=pl.col.date.dt.year())

# an MSFT specific plot
msft_plot = (
    alt.Chart(source.filter(pl.col.symbol.eq("MSFT")))
    .mark_line()
    .encode(x="date:T", y="price:Q", color="year:O")
)

# the original plot: https://altair-viz.github.io/gallery/line_chart_with_points.html
all_plot = (
    alt.Chart(source)
    .mark_line()
    .encode(x="date:T", y="price:Q", color="symbol:N")
)

msft_plot & all_plot

这将产生以下输出：

另一方面，如果我仅绘制all_plot：

当我连接时，如何阻止图例合并msft_plot & all_plot？

请考虑以下示例：

import polars as pl

pl.DataFrame(pl.Series("x", ["1, 0", "2,3", "5 4"])).with_columns(
    pl.col("x").str.split(",").list.eval(pl.element().str.split(" "))
)

shape: (3, 1)
┌────────────────────┐
│ x                  │
│ ---                │
│ list[list[str]]    │
╞════════════════════╡
│ [["1"], ["", "0"]] │
│ [["2"], ["3"]]     │
│ [["5", "4"]]       │
└────────────────────┘

我想将列中的元素展平，这样这些元素就不再是嵌套列表，而是列表。我该怎么做？

以下是一个小例子：

import altair as alt
import polars as pl

source = pl.DataFrame(
    {
        "Category": list("AAABBBCCC"),
        "Value": [0.1, 0.6, 0.9, 0.7, 0.2, 1.1, 0.6, 0.1, 0.2],
        "Date": [f"2024-{m+1}-1" for m in range(3)] * 3,
    }
).with_columns(pl.col("Date").str.to_date())

bars = alt.Chart(source).mark_bar().encode(
    x=alt.X("Date:T"),
    xOffset="Category:N",
    y="Value:Q",
    color="Category:N",
)

bars

如果我设置x="Date:N"，则示例将按我希望的方式运行，但没有 x 轴时间格式的好处：

有什么方法可以帮助我解决xOffset这种情况吗x="Date:T"？

JoinStrategy是输入join：https：//docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.join.html

我的静态类型检查工具似乎能够掌握JoinStrategy，但我不知道如何/从哪里掌握。

通常，类型存根包在 PyPI 上可用，但在这种情况下没有什么明显的突出之处：https://pypi.org/user/ritchie46/

我如何导入JoinStrategy（或 Polars 提供的其他类型定义）供我自己使用？

作为一个简单的示例，考虑以下内容，使用groupby：

import polars as pl

df = pl.DataFrame(
    [pl.Series("id", ["a", "b", "a"]), pl.Series("x", [0, 1, 2])]
)
print(df.group_by("id").agg(pl.col("x")))

# shape: (2, 2)
# ┌─────┬───────────┐
# │ id  ┆ x         │
# │ --- ┆ ---       │
# │ str ┆ list[i64] │
# ╞═════╪═══════════╡
# │ b   ┆ [1]       │
# │ a   ┆ [0, 2]    │
# └─────┴───────────┘

但如果我们使用over，我们会得到：

import polars as pl

df = pl.DataFrame(
    [pl.Series("id", ["a", "b", "a"]), pl.Series("x", [0, 1, 2])]
)
print(df.with_columns(pl.col("x").over("id")))
# shape: (3, 2)
# ┌─────┬─────┐
# │ id  ┆ x   │
# │ --- ┆ --- │
# │ str ┆ i64 │
# ╞═════╪═════╡
# │ a   ┆ 0   │
# │ b   ┆ 1   │
# │ a   ┆ 2   │
# └─────┴─────┘

如何groupby使用 实现结果over？嗯，使用mapping_strategy="join"。

一个稍微复杂一点的例子，旨在展示为什么我们可能想要使用over而不是groupby：

import polars as pl

# the smallest value a Float32 can encode is 1e-38
# therefore, as far as we are concerned,
# 1e-41 and 1e-42 should be indistinguishable
# in other words, we do not want to use "other" as an id column
# but we do want to preserve other!
df = pl.DataFrame(
    [
        pl.Series("id", ["a", "b", "a"]),
        pl.Series("other", [1e-41, 1e-16, 1e-42], dtype=pl.Float32()),
        pl.Series("x", [0, 1, 2]),
    ]
)
print(df.group_by("id").agg(pl.col("x"), pl.col("other")).explode("other"))

# shape: (3, 3)
# ┌─────┬───────────┬────────────┐
# │ id  ┆ x         ┆ other      │
# │ --- ┆ ---       ┆ ---        │
# │ str ┆ list[i64] ┆ f32        │
# ╞═════╪═══════════╪════════════╡
# │ a   ┆ [0, 2]    ┆ 9.9997e-42 │
# │ a   ┆ [0, 2]    ┆ 1.0005e-42 │
# │ b   ┆ [1]       ┆ 1.0000e-16 │
# └─────┴───────────┴────────────┘

现在，使用over：

import polars as pl

# the smallest value a Float32 can encode is 1e-38
# therefore, as far as we are concerned,
# 1e-41 and 1e-42 should be indistinguishable
# in other words, we do not want to use "other" as an id column
# but we do want to preserve other!
df = pl.DataFrame(
    [
        pl.Series("id", ["a", "b", "a"]),
        pl.Series("other", [1e-41, 1e-16, 1e-42], dtype=pl.Float32()),
        pl.Series("x", [0, 1, 2]),
    ]
)
print(df.with_columns(pl.col("x").over(["id"], mapping_strategy="join")))

# shape: (3, 3)
# ┌─────┬────────────┬───────────┐
# │ id  ┆ other      ┆ x         │
# │ --- ┆ ---        ┆ ---       │
# │ str ┆ f32        ┆ list[i64] │
# ╞═════╪════════════╪═══════════╡
# │ a   ┆ 9.9997e-42 ┆ [0, 2]    │
# │ b   ┆ 1.0000e-16 ┆ [1]       │
# │ a   ┆ 1.0005e-42 ┆ [0, 2]    │
# └─────┴────────────┴───────────┘

问题mapping_strategy="join"是使用起来很慢。所以，这意味着我应该先执行 a，group_by然后执行 a join：

import polars as pl
import polars.selectors as cs

# the smallest value a Float32 can encode is 1e-38
# therefore, as far as we are concerned,
# 1e-41 and 1e-42 should be indistinguishable
# in other words, we do not want to use "other" as an id column
# but we do want to preserve other!
df = pl.DataFrame(
    [
        pl.Series("id", ["a", "b", "a"]),
        pl.Series("other", [1e-41, 1e-16, 1e-42], dtype=pl.Float32()),
        pl.Series("x", [0, 1, 2]),
    ]
)


print(
    df.select(cs.exclude("x")).join(
        df.group_by("id").agg("x"),
        on="id",
        # we expect there to be multiple "id"s on the left, matching
        # a single "id" on the right
        validate="m:1",
    )
)

# shape: (3, 3)
# ┌─────┬────────────┬───────────┐
# │ id  ┆ other      ┆ x         │
# │ --- ┆ ---        ┆ ---       │
# │ str ┆ f32        ┆ list[i64] │
# ╞═════╪════════════╪═══════════╡
# │ a   ┆ 9.9997e-42 ┆ [0, 2]    │
# │ b   ┆ 1.0000e-16 ┆ [1]       │
# │ a   ┆ 1.0005e-42 ┆ [0, 2]    │
# └─────┴────────────┴───────────┘

但也许我还遗漏了其他一些东西over？

当我输入类似的内容时，np.我认为这会触发 Visual Studio Code + Pylance（不确定）的自动导入完成，并暗示这import numpy as np可能是相关的。

我想创建类似的自定义自动导入/完成关联。例如：介于pl和之间polars，这样如果我输入类似的东西pl.，import polars as pl就会给出自动导入建议。

我该怎么做？这是我正在使用的 Pylance 扩展特有的，还是与 Visual Studio Code 有关？

请注意，自动导入/导入完成与自定义代码片段有很大不同，如如何在 VSCode 中添加自定义代码片段中所述。原因如下：

• VS Code 在文件顶部添加了一条新的导入语句（因此必须确定是否可以解析该导入）。它不会在光标所在的位置添加代码，而代码片段会这样做。
• 此功能依赖于某种语言服务器（因此我怀疑是 Pylance 提供了此功能）来解析导入，并在文件中的适当位置插入导入语句。

假设我有以下内容：

from collections import OrderedDict
from dataclasses import dataclass


@dataclass
class HelloWorld:
    x: OrderedDict[str, int]


a = OrderedDict([("a", 0), ("c", 2), ("b", 1)])
HelloWorld(a) <--- # type error here

产生的类型错误是：

Argument of type "OrderedDict[Literal['a', 'c', 'b'], Literal[0, 2, 1]]" cannot be assigned to parameter "x" of type "OrderedDict[str, int]" in function "__init__"
  "OrderedDict[Literal['a', 'c', 'b'], Literal[0, 2, 1]]" is incompatible with "OrderedDict[str, int]"
    Type parameter "_KT@OrderedDict" is invariant, but "Literal['a', 'c', 'b']" is not the same as "str"
    Type parameter "_VT@OrderedDict" is invariant, but "Literal[0, 2, 1]" is not the same as "int

奇怪的是，这个非常相似的代码片段不会产生错误：

from collections import OrderedDict
from dataclasses import dataclass


@dataclass
class HelloWorld:
    x: OrderedDict[str, int]


HelloWorld(OrderedDict([("a", 0), ("c", 2), ("b", 1)])) # <--- no error

考虑以下玩具示例：

import polars as pl

xs = pl.DataFrame(
    [
        pl.Series(
            "date",
            ["2024 Jan", "2024 Feb", "2024 Jan", "2024 Jan"],
            dtype=pl.String,
        )
    ]
)
ys = (
    xs.with_columns(
        pl.col("date").str.split(" ").list.to_struct(fields=["year", "month"]),
    )
    .with_columns(
        pl.col("date").struct.with_fields(pl.field("year").cast(pl.Int16()))
    )
    .unnest("date")
)
ys

shape: (4, 2)
┌──────┬───────┐
│ year ┆ month │
│ ---  ┆ ---   │
│ i16  ┆ str   │
╞══════╪═══════╡
│ 2024 ┆ Jan   │
│ 2024 ┆ Feb   │
│ 2024 ┆ Jan   │
│ 2024 ┆ Jan   │
└──────┴───────┘

我认为对一系列独特的日期数据进行操作会更有效率（我可以使用map_dict，但我选择这样做join没有什么好理由）：

unique_dates = (
    pl.DataFrame([xs["date"].unique()])
    .with_columns(
        pl.col("date")
        .str.split(" ")
        .list.to_struct(fields=["year", "month"])
        .alias("struct_date")
    )
    .with_columns(
        pl.col("struct_date").struct.with_fields(
            pl.field("year").cast(pl.Int16())
        )
    )
)
unique_dates

shape: (2, 2)
┌──────────┬──────────────┐
│ date     ┆ struct_date  │
│ ---      ┆ ---          │
│ str      ┆ struct[2]    │
╞══════════╪══════════════╡
│ 2024 Jan ┆ {2024,"Jan"} │
│ 2024 Feb ┆ {2024,"Feb"} │
└──────────┴──────────────┘

zs = (
    xs.join(unique_dates, on="date", left_on="date", right_on="struct_date")
    .drop("date")
    .rename({"struct_date": "date"})
    .unnest("date")
)

zs

shape: (4, 2)
┌──────┬───────┐
│ year ┆ month │
│ ---  ┆ ---   │
│ i16  ┆ str   │
╞══════╪═══════╡
│ 2024 ┆ Jan   │
│ 2024 ┆ Feb   │
│ 2024 ┆ Jan   │
│ 2024 ┆ Jan   │
└──────┴───────┘

我该怎么做才能进一步提高此操作的效率？我polars是否足够习惯使用？

import polars as pl

dates = ["2024 Jan", "2024 Feb"]
DateEnum = pl.Enum(dates)

date_table = pl.DataFrame(
    {"date": pl.Series(raw_dates, dtype=DateEnum)}
)
date_table.explode("date") # fails

    {
        "date": pl.Series(raw_dates, dtype=plDateEnum).cast(dtype=str),
        "index": pl.Series(raw_dates, dtype=plDateEnum).cast(dtype=int),
    }
)
date_table # fails

我希望看到像这样的数据框：

from typing import Enum

pyDateEnum = Enum("Dates", raw_dates)
date_table = pl.DataFrame({
  "date": pl.Series("dates", [x.name for x in pyDateEnum], dtype=str),
  "index": pl.Series("indices", [x.value - 1 for x in pyDateEnum], dtype=int)
})
date_table

shape: (2, 2)
┌──────────┬───────┐
│ date     ┆ index │
│ ---      ┆ ---   │
│ str      ┆ i64   │
╞══════════╪═══════╡
│ 2024 Jan ┆ 0     │
│ 2024 Feb ┆ 1     │
└──────────┴───────┘

如何使用 Polars 来实现这一点，而不使用 Python 枚举？

注意：本问题涉及 Python 3.12+

假设我有：

from typing import Any, TypeVar

import numpy as np


T = TypeVar("T")
U = TypeVar("U")
ListLike = T | list[T] | tuple[T, ...] | np.ndarray[Any, U]
ListLikeStr = ListLike[str, np.object_]
# ListLikeStr should be: str | list[str] | tuple[str, ...] | np.ndarray[Any, np.object_]

这有效，但幸运的是。我本可以这样写：ListLike[np.object_, str]，然后我会得到ListLikeStrbeing np.object_ | list[np.object_] | tuple[np.object_, ...] | np.ndarray[Any, str]，但这不是我想要的。

理想情况下，我可以这样做：ListLike[T=str, U=np.object_]，但这样做行不通。那么，当我在 中实例化类型变量时，是什么决定了顺序ListLike？当我写 时，ListLike“知道”T与str和U相对应？np.object_ListLike[str, np.object_]