dask.dataframe.DataFrame.set_index

dask.dataframe.DataFrame.set_index

DataFrame.set_index(other: str | dask.dataframe.core.Series, drop: bool = True, sorted: bool = False, npartitions: Optional[Union[int, Literal['auto']]] = None, divisions: collections.abc.Sequence | None = None, inplace: bool = False, sort: bool = True, **kwargs)[源代码]

使用现有列设置 DataFrame 索引(行标签)。

如果 sort=False,此函数操作与 pandas.set_index 完全相同,并在 DataFrame 上设置索引。如果 sort=True``(默认),此函数还会按新索引对 DataFrame 进行排序。这会对性能产生重大影响,因为在该列上的连接、分组、查找等操作都会快得多。然而,这种性能提升是有代价的,对并行数据集进行排序需要昂贵的洗牌操作。通常我们在数据摄取和过滤后直接 ``set_index 一次,然后对排序后的数据集执行许多廉价的计算。

使用 sort=True ,这个函数会变得更加昂贵。在正常操作下,这个函数首先遍历索引列以计算近似分位数,作为未来的分割点。然后,它再次遍历数据,将每个输入分区分割成若干部分,并将这些部分按排序顺序分配给所有的输出分区。

在某些情况下,我们可以减轻这些成本,例如,如果您的数据集已经排序,那么我们可以避免制作许多小块,或者如果您知道分割新索引列的良好值,那么我们可以避免初始的数据遍历。例如,如果您的新索引是日期时间索引,并且您的数据已经按天排序,那么整个操作可以免费完成。您可以使用以下参数控制这些选项。

参数
other: 字符串或Dask系列

用作索引的列。

drop: boolean, 默认 True

删除要作为新索引的列。

sorted: bool, 可选

如果索引列已经按升序排序。默认为 False

npartitions: int, None, 或 ‘auto’

理想的输出分区数量。如果为 None,则使用与输入相同的数量。如果为 ‘auto’,则根据内存使用情况决定。仅在未给出 divisions 时使用。如果给出了 divisions,则输出分区的数量将为 len(divisions) - 1

divisions: list, 可选

用于将新索引分割成分区的“分割线”。对于 divisions=[0, 10, 50, 100],将会有三个输出分区,新索引分别包含 [0, 10)、[10, 50) 和 [50, 100)。参见 http://www.aidoczh.com/dask/dataframe-design.html#partitions。如果没有指定(默认),将通过立即计算数据并查看其值的分布来计算好的分区。对于大型数据集,这可能会很昂贵。请注意,如果 sorted=True,指定的分区将被假定与数据中的现有分区匹配;如果这不是真的,你应该保持分区为空并在 set_index 之后调用 repartition

inplace: bool, 可选

Dask 不支持就地修改 DataFrame。默认为 False。

sort: bool, 可选

如果 True,则按新索引对 DataFrame 进行排序。否则在各个现有分区上设置索引。默认为 True

shuffle_method: {‘disk’, ‘tasks’, ‘p2p’}, 可选

可以是 'disk' 用于单节点操作,或者是 'tasks''p2p' 用于分布式操作。将由您当前的调度器推断。

compute: bool, 默认 False

是否触发立即计算。默认为 False。注意,即使你设置 compute=False,如果 divisionsNone,仍然会触发立即计算。

partition_size: int, 可选

每个分区期望的大小,以字节为单位。仅在 npartitions='auto' 时使用。

示例

>>> import dask
>>> ddf = dask.datasets.timeseries(start="2021-01-01", end="2021-01-07", freq="1h").reset_index()
>>> ddf2 = ddf.set_index("x")
>>> ddf2 = ddf.set_index(ddf.x)
>>> ddf2 = ddf.set_index(ddf.timestamp, sorted=True)

一个常见的情况是当我们有一个已知是按时间排序并且按天清晰分割的datetime列时。我们可以通过指定该列是预排序的以及它是沿着哪些特定分割的,来免费设置这个索引。

>>> import pandas as pd
>>> divisions = pd.date_range(start="2021-01-01", end="2021-01-07", freq='1D')
>>> divisions
DatetimeIndex(['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04',
               '2021-01-05', '2021-01-06', '2021-01-07'],
              dtype='datetime64[ns]', freq='D')

注意 len(divisions) 等于 npartitions + 1。这是因为 divisions 表示每个分区的上下界。第一个项目是第一个分区的下界,第二个项目是第二个分区的下界和第一个分区的上界,依此类推。倒数第二个项目是最后一个分区的下界,最后一个(额外的)项目是最后一个分区的上界。

>>> ddf2 = ddf.set_index("timestamp", sorted=True, divisions=divisions.tolist())

如果你将反复在相同(或相似)的数据集上运行 set_index,你可以通过让 Dask 一次性计算好的分区,然后复制粘贴它们以重复使用来节省时间。这在 Jupyter notebook 中运行时特别有用:

>>> ddf2 = ddf.set_index("name")  # slow, calculates data distribution
>>> ddf2.divisions  
["Alice", "Laura", "Ursula", "Zelda"]
>>> # ^ Now copy-paste this and edit the line above to:
>>> # ddf2 = ddf.set_index("name", divisions=["Alice", "Laura", "Ursula", "Zelda"])