diff --git a/src/quak/_util.py b/src/quak/_util.py index 427c2bb..049fc97 100644 --- a/src/quak/_util.py +++ b/src/quak/_util.py @@ -70,12 +70,22 @@ def arrow_table_from_ipc(data: bytes | memoryview) -> pa.lib.Table: return feather.read_table(io.BytesIO(data)) -def table_to_ipc(table: pa.lib.Table) -> memoryview: - """Convert a pyarrow Table to an Arrow IPC message.""" +def table_to_ipc(table: pa.lib.Table | pa.lib.RecordBatch | pa.lib.RecordBatchReader) -> memoryview: + """Convert Arrow tabular data to an Arrow IPC message.""" import io + import pyarrow as pa import pyarrow.feather as feather + if isinstance(table, pa.RecordBatchReader): + table = table.read_all() + elif isinstance(table, pa.RecordBatch): + table = pa.Table.from_batches([table], schema=table.schema) + elif not isinstance(table, pa.Table): + raise TypeError( + f"Expected a pyarrow Table, RecordBatch, or RecordBatchReader, got {type(table)!r}" + ) + sink = io.BytesIO() feather.write_feather(table, sink, compression="uncompressed") return sink.getbuffer() diff --git a/src/quak/_version.py b/src/quak/_version.py index 8f05f93..6a4982b 100644 --- a/src/quak/_version.py +++ b/src/quak/_version.py @@ -2,4 +2,7 @@ import importlib.metadata -__version__ = importlib.metadata.version("quak") +try: + __version__ = importlib.metadata.version("quak") +except importlib.metadata.PackageNotFoundError: # pragma: no cover - local dev fallback + __version__ = "0+local"