python - 使用 .loc 后 Dask categorize() 将不起作用

Question

我在使用 dask 时遇到严重问题(dask 版本:1.00，pandas 版本:0.23.3)。我正在尝试从 CSV 文件加载 dask 数据帧，将结果过滤到两个单独的数据帧中，并对两个数据帧执行操作。

但是，在分割数据帧并尝试将类别列设置为“已知”后，它们仍然是“未知”。因此我无法继续我的操作(这需要“已知”类别列。)

注意:我已经按照建议使用 pandas 而不是 read_csv() 创建了一个最小示例。

import pandas as pd
import dask.dataframe as dd

# Specify dtypes
b_dtypes = {
    'symbol': 'category',
    'price': 'float64',
}

i_dtypes = {
    'symbol': 'category',
    'price': 'object'
}

# Specify a function to quickly set dtypes
def to_dtypes(df, dtypes):
    for column, dtype in dtypes.items():
        if column in df.columns:
            df[column] = df.loc[:, column].astype(dtype)
    return df

# Set up our test data
data = [
    ['B', 'IBN', '9.9800'],
    ['B', 'PAY', '21.5000'],
    ['I', 'PAY', 'seventeen'],
    ['I', 'SPY', 'ten']
]

# Create pandas dataframe
pdf = pd.DataFrame(data, columns=['type', 'symbol', 'price'], dtype='object')

# Convert into dask
df = dd.from_pandas(pdf, npartitions=3)

#
## At this point 'df' simulates what I get when I read the mixed-type CSV file via dask
#

# Split the dataframe by the 'type' column
b_df = df.loc[df['type'] == 'B', :]
i_df = df.loc[df['type'] == 'I', :]

# Convert columns into our intended dtypes
b_df = to_dtypes(b_df, b_dtypes)
i_df = to_dtypes(i_df, i_dtypes)

# Let's convert our 'symbol' column to known categories
b_df = b_df.categorize(columns=['symbol'])
i_df['symbol'] = i_df['symbol'].cat.as_known()

# Is our symbol column known now?
print(b_df['symbol'].cat.known, flush=True)
print(i_df['symbol'].cat.known, flush=True)

#
## print() returns 'False' for both, this makes me want to kill myself.
## (Please help...)
#

更新:因此，如果我将“npartitions”参数移至 1，则 print() 在两种情况下都会返回 True。因此，这似乎是包含不同类别的分区的问题。然而，将两个数据帧仅加载到两个分区中是不可行的，所以有没有办法告诉 dask 进行某种重新排序以使分区之间的类别保持一致？

Answer 1

您的问题的答案基本上包含在 doc 中。我指的是 # categorize 需要计算，并导致已知分类的结果 注释的部分代码，我将在这里展开，因为在我看来，您滥用了 loc

import pandas as pd
import dask.dataframe as dd

# Set up our test data
data = [['B', 'IBN', '9.9800'],
        ['B', 'PAY', '21.5000'],
        ['I', 'PAY', 'seventeen'],
        ['I', 'SPY', 'ten']
       ]

# Create pandas dataframe
pdf = pd.DataFrame(data, columns=['type', 'symbol', 'price'], dtype='object')

# Convert into dask
ddf = dd.from_pandas(pdf, npartitions=3)

# Split the dataframe by the 'type' column
# reset_index is not necessary
b_df = ddf[ddf["type"] == "B"].reset_index(drop=True)
i_df = ddf[ddf["type"] == "I"].reset_index(drop=True)

# Convert columns into our intended dtypes
b_df = b_df.categorize(columns=['symbol'])
b_df["price"] = b_df["price"].astype('float64')
i_df = i_df.categorize(columns=['symbol'])

# Is our symbol column known now? YES
print(b_df['symbol'].cat.known, flush=True)
print(i_df['symbol'].cat.known, flush=True)