在“pyarrow”测试中使用inmemory文件系统

2024-03-28 23:13:33 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一些pyarrow拼花数据集正在编写代码。我想有一个集成测试,以确保文件被正确写入。我想通过将一个小示例数据块写入内存文件系统来实现这一点。然而,我正在努力为Python找到一个pyarrow兼容的内存文件系统接口。你知道吗

您将在下面找到一段代码,其中包含一个filesystem变量。我想用内存中的文件系统替换filesystem变量,稍后我可以通过编程方式在集成测试中检查该文件系统。你知道吗

import pyarrow.parquet as pq
pq.write_to_dataset(
        score_table,
        root_path=AWS_ZEBRA_OUTPUT_S3_PREFIX,
        filesystem=filesystem,
        partition_cols=[
            EQF_SNAPSHOT_YEAR_PARTITION,
            EQF_SNAPSHOT_MONTH_PARTITION,
            EQF_SNAPSHOT_DAY_PARTITION,
            ZEBRA_COMPUTATION_TIMESTAMP
        ]
    )

Tags: 文件数据内存代码示例编程方式snapshot
2条回答

最后,我手动实现了pyarrow.FileSystemABC的一个实例。似乎使用mock进行测试是失败的,因为pyarrow(不是以最Pythonic的方式)检查传递给write_to_datasethttps://github.com/apache/arrow/blob/5e201fed061f2a95e66889fa527ae8ef547e9618/python/pyarrow/filesystem.py#L383filesystem参数的类型。我建议将此方法中的逻辑更改为不显式检查类型(甚至isinstance也更好!)以便于测试。你知道吗

如果filesystemNone,则可以将内存中的文件对象传递给write_to_dataset。你知道吗

所以你的电话可能会变成:

from io import BytesIO
import pyarrow.parquet as pq

with BytesIO() as f:
    pq.write_to_dataset(
        score_table,
        root_path=f,
        filesystem=None,
        partition_cols=[
            EQF_SNAPSHOT_YEAR_PARTITION,
            EQF_SNAPSHOT_MONTH_PARTITION,
            EQF_SNAPSHOT_DAY_PARTITION,
            ZEBRA_COMPUTATION_TIMESTAMP
        ]
    )

pyarrow来源的相关行:

def resolve_filesystem_and_path(where, filesystem=None):
    """
    Return filesystem from path which could be an HDFS URI, a local URI,
    or a plain filesystem path.
    """
    if not _is_path_like(where):
        if filesystem is not None:
            raise ValueError("filesystem passed but where is file-like, so"
                             " there is nothing to open with filesystem.")
        return filesystem, where

https://github.com/apache/arrow/blob/207b3507be82e92ebf29ec7d6d3b0bb86091c09a/python/pyarrow/filesystem.py#L402-L411

相关问题 更多 >