从共享内存进行惰性扫描

编程语言 2026-07-12

我在尝试建立一个系统,通过共享内存使用Arrow IPC读写Polars。我已经有一份能够运行的代码(read_ipc 可以工作),但 scan_ipc 函数对我来说并不起作用。以下脚本演示了这个问题:

# /// script
# requires-python = "~=3.12.0"
# dependencies = [
#     "polars[pyarrow]==1.38.1",
# ]
# ///
"""Minimal working example for scan IPC bug."""

import logging
from multiprocessing.shared_memory import SharedMemory

import polars as pl
import pyarrow as pa

LOG = logging.getLogger(__name__)

pl.show_versions()

if __name__ == "__main__":
    # Create a toy dataset
    data = pl.DataFrame({"a": [1, 2, 3]})
    with pa.MockOutputStream() as sink:
        data.write_ipc(sink)

    shm = SharedMemory(name="test-ipc-polars", create=True, size=sink.size())
    with pa.FixedSizeBufferWriter(pa.py_buffer(shm.buf)) as stream:
        data.write_ipc(stream)
    del stream
    shm.close()

    # Read the dataset back
    existing_ = SharedMemory(shm.name, create=False)
    buf_ = pa.py_buffer(existing_.buf)
    # out = pl.read_ipc(buf_, use_pyarrow=True)  # Works
    # print(out)

    try:
        res = pl.scan_ipc(buf_)
    except Exception as exc:
        LOG.error(exc, stack_info=True, exc_info=True)
    finally:
        del buf_
        existing_.close()
        shm.unlink()

以详细日志运行这段脚本会得到如下输出:

_init_credential_provider_builder(): credential_provider_init = None
async thread count: 4
polars-stream: updating graph state
polars-stream: running in-memory-source in subgraph
polars-stream: running io-sink[single-file[ipc]] in subgraph
async upload_chunk_size: 67108864
async upload_concurrency: 8
io-sink[single-file[ipc]]: start_single_file_sink_pipeline: file_writer_starter: ipc, takeable_rows_provider: TakeableRowsProvider { max_size: NonZeroRowCountAndSize { num_rows: 122880, num_bytes: 18446744073709551615 }, byte_size_min_rows: 16384, allow_non_max_size: false }, upload_chunk_size: 67108864
polars-stream: done running graph phase
polars-stream: updating graph state
io-sink[single-file[ipc]]: Join on task_handle (recv PortState::Done)
io-sink[single-file[ipc]]: Statistics: total_size: RowCountAndSize { num_rows: 3, num_bytes: 24 }
_init_credential_provider_builder(): credential_provider_init = None
polars-stream: updating graph state
polars-stream: running io-sink[single-file[ipc]] in subgraph
polars-stream: running in-memory-source in subgraph
io-sink[single-file[ipc]]: start_single_file_sink_pipeline: file_writer_starter: ipc, takeable_rows_provider: TakeableRowsProvider { max_size: NonZeroRowCountAndSize { num_rows: 122880, num_bytes: 18446744073709551615 }, byte_size_min_rows: 16384, allow_non_max_size: false }, upload_chunk_size: 67108864
polars-stream: done running graph phase
polars-stream: updating graph state
io-sink[single-file[ipc]]: Join on task_handle (recv PortState::Done)
io-sink[single-file[ipc]]: Statistics: total_size: RowCountAndSize { num_rows: 3, num_bytes: 24 }
_init_credential_provider_builder(): credential_provider_init = None
argument 'sources': Object does not have a .read() method.
Traceback (most recent call last):
  File "/workspaces/github/mwe.py", line 36, in <module>
    res = pl.scan_ipc(buf_)
          ^^^^^^^^^^^^^^^^^
  File "/app/.cache/uv/environments-v2/mwe-63942327ce1a05a8/lib/python3.12/site-packages/polars/_utils/deprecation.py", line 128, in wrapper
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/.cache/uv/environments-v2/mwe-63942327ce1a05a8/lib/python3.12/site-packages/polars/_utils/deprecation.py", line 128, in wrapper
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/.cache/uv/environments-v2/mwe-63942327ce1a05a8/lib/python3.12/site-packages/polars/io/ipc/functions.py", line 506, in scan_ipc
    pylf = PyLazyFrame.new_from_ipc(
           ^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: argument 'sources': Object does not have a .read() method.
Stack (most recent call last):
  File "/workspaces/github/mwe.py", line 38, in <module>
    LOG.error(exc, stack_info=True, exc_info=True)

我很想能够对缓冲区进行惰性扫描!

解决方案

argument 'sources': Object does not have a .read() method. 行表明传入的对象没有实现 read() 方法,而 scan_ipc() 期望该方法。

文档 py_buffer 用来构造一个 Buffer,但并未实现 read() 方法。我们可以改用 BufferReader,它允许我们无问题地调用 scan_ipc()

# Old implementation
# buf_ = pa.py_buffer(existing_.buf)
buf_ = pa.BufferReader(existing_.buf)
站内所有文章版权归属LeftHeroAI导航站,无授权禁止任何主体转载、抄袭、复制内容,亦不得私自架设镜像站点。一经侵权,本站将通过法律途径追责。

相关文章