在函数内使用日期时间索引对 DataFrame 进行适当切片

Question

这更多是一个与对象类型相关的问题，但我也收到了

mypy

错误。所以一定有更好的方法来做到这一点。

说到重点：

我正在使用以下函数加载时间序列：

def load_something_between_two_days(
    "some_path",
    start: Optional[str] = None,
    end: Optional[str] = None,
):
    df = pd.read_parquet(some_path) # index is assumed to be datetime, no shitty format
    if start is not None:
        df.index = df[start:]
    if end is not None:
        df.index = df[:end]
return df

我的错误：

error: Slice index must be an integer or None

我的第一个假设是宁愿使用

.loc

。但

Optional

标签仍然在我的定义中。可能还有

mypy´ is not happy about the weird 

DataFrame` 切片。那么处理这个问题的最佳方法是什么？我也不确定我是否没有解决任何基本的编码原则。

StackOverflow 适合提出此类问题吗？ ：D

干杯！

Answer 1

是的，你可以使用

loc

：

def load_something_between_two_days(
    some_path: str,
    start: Optional[str] = None,
    end: Optional[str] = None,
):

    df = pd.read_parquet(some_path).loc[start:end]

    return df

如果你有一个命名索引，你可以避免使用

filters

加载整个数据框：

def load_something_between_two_days(
    some_path: str,
    start: Optional[str] = None,
    end: Optional[str] = None,
):
    filters = []

    if start is not None:
        filters.append(("index", ">=", pd.Timestamp(start)))

    if end is not None:
        filters.append(("index", "<=", pd.Timestamp(end)))
        
    df = pd.read_parquet(some_path, filters=filters if filters else None)
    
    return df