Skip to content

yggdrasil.arrow.tests

tests

Unittest base class for Apache Arrow tests.

Quick start

::

from yggdrasil.testing import ArrowTestCase
import pyarrow as pa

class TestMyCodec(ArrowTestCase):
    def test_roundtrip(self):
        tbl = pa.table({"id": [1, 2, 3], "val": ["a", "b", "c"]})
        out = self.tmp_path / "t.parquet"
        self.write_parquet(tbl, out)
        self.assertFrameEqual(self.read_parquet(out), tbl)

Auto-install

The class uses :func:yggdrasil.environ.runtime_import_module to load pyarrow. Auto-install is opt-in: set auto_install = True on the subclass or export YGG_TEST_AUTO_INSTALL=1. Otherwise a missing pyarrow skips the class with an install hint.

ArrowTestCase

Bases: TestCase

Base class for pyarrow integration tests.

Attributes

pa : module The imported pyarrow module. Populated by setUpClass. pq : module The imported pyarrow.parquet module (if require_parquet). tmp_path : pathlib.Path Per-test scratch directory.

Class attributes

auto_install : bool | None Override the global auto-install behaviour for this class. None means "follow YGG_TEST_AUTO_INSTALL". require_parquet : bool If True (default), also import pyarrow.parquet.

table

table(data: dict[str, Any], schema: Any = None) -> 'pa.Table'

Shorthand for pa.table(data, schema=schema).

record_batch

record_batch(data: dict[str, Any], schema: Any = None) -> 'pa.RecordBatch'

Shorthand for pa.record_batch(data, schema=schema).

write_parquet

write_parquet(table: 'pa.Table', path: Path | str) -> Path

Write table to path as Parquet and return the path.

read_parquet

read_parquet(path: Path | str) -> 'pa.Table'

Read a Parquet file as a pa.Table.

write_ipc

write_ipc(table: 'pa.Table', path: Path | str) -> Path

Write table as Arrow IPC (file format) and return the path.

read_ipc

read_ipc(path: Path | str) -> 'pa.Table'

Read an Arrow IPC file as a pa.Table.

assertFrameEqual

assertFrameEqual(
    actual: "pa.Table",
    expected: "pa.Table",
    *,
    check_schema: bool = True,
    check_metadata: bool = False
) -> None

Assert two Arrow tables are equal.

check_metadata=False ignores schema and field metadata, which is the usual "does the data match" check. Set it True for round-trip tests where metadata preservation matters.

assertSchemaEqual

assertSchemaEqual(
    actual: "pa.Table | pa.Schema", expected_fields: list[tuple[str, Any]]
) -> None

Assert a table/schema has exactly the given (name, type) fields.