DataFrames and Arrow#

agensgraph.columnar turns a result into an Apache Arrow table, a pandas DataFrame or a polars DataFrame. Each column is built as a column, with its type taken from what the server said the column is wherever it said, rather than by turning rows into dicts and letting the library guess, which is also faster.

Function	Gives	Needs
`to_arrow(source)`	a `pyarrow.Table`	`arrow` extra
`to_pandas(source)`	a `pandas.DataFrame`	`arrow` and `pandas` extras
`to_polars(source)`	a `polars.DataFrame`	`arrow` and `polars` extras
`batches(source, size=8192)`	`pyarrow.RecordBatch`es, one chunk at a time	`arrow` extra
`reader(source, size=8192)`	a `pyarrow.RecordBatchReader`	`arrow` extra
`columns(records, keys)`	a `dict` of Python lists	nothing

Each library is imported only inside the function that uses it, so importing agensgraph.columnar loads none of them.

From a result#

>>> from agensgraph.columnar import to_arrow, to_pandas, to_polars
>>> result = conn.execute_query(
...     "MATCH (m:movie) RETURN m.title AS title, m.year AS year, m.genre AS genre"
...     " ORDER BY year LIMIT 4"
... )
>>> table = to_arrow(result)
>>> table.schema
title: string
year: int64
genre: string
>>> print(to_pandas(result))
                        title  year   genre
0                Blade Runner  1982  Sci-Fi
1              The Terminator  1984  Action
2                Total Recall  1990  Sci-Fi
3  Terminator 2: Judgment Day  1991  Action

to_polars(result) gives the same as a polars frame. to_pandas keeps Arrow’s types by default, so a struct or a vector column stays one; dtypes="numpy" asks for pandas’ classic types instead, with structs and lists as Python objects.

A source is one of three things:

a Result from execute_query(), which carries its column names and their types;
a psycopg cursor, including a named server-side one: read whole by to_arrow() and its siblings unless you pass size=, and a chunk at a time by batches() and reader();
any other sequence or iterator of rows, such as conn.stream(...), which carries no names, so pass them with keys=.

Rows must be tuples. A result read with a dict row factory cannot be transposed into columns correctly; use the default rows for export.

How values become columns#

Value	Arrow column
a vertex	`struct<id: uint64, label: dictionary<int32, string>, properties: string>`
an edge	the same, plus `start: uint64` and `end: uint64`
a path	`struct<vertices: list<vertex>, edges: list<edge>>`
a `GraphId`	`uint64`, the packed identity
a list of vertices or edges	a `list` of those structs
a dense vector	`fixed_size_list<float32>[n]`, straight from the wire bytes
a sparse vector	`struct<dimensions, indices, values>`
a property, or any other `jsonb`	inferred from the values: `int64`, `float64`, `string`, a list, or JSON text for a map
a SQL column of a known type	its Arrow equivalent: `int4` as `int32`, `timestamptz` as `timestamp[us, UTC]`, and so on

>>> to_arrow(conn.execute_query("MATCH (m:movie) RETURN m ORDER BY id(m) LIMIT 2")).schema
m: struct<id: uint64, label: dictionary<values=string, indices=int32, ordered=0>, properties: string>
  child 0, id: uint64
  child 1, label: dictionary<values=string, indices=int32, ordered=0>
  child 2, properties: string

A vertex’s property map goes in as JSON text, taken from the bytes it arrived in, so it is never decoded into Python on the way. An edge’s start and end are the same uint64 as a vertex’s id, so joining edges to vertices in Arrow, pandas or polars is an integer join.

Choosing another layout#

Layout decides the cases where more than one answer is reasonable:

import pyarrow
from agensgraph.columnar import Layout

Layout(elements="columns")      # a vertex as n.id, n.label, n.properties, not one struct
Layout(identity="text")         # identities as '3.1' rather than packed integers
Layout(properties="skip")       # leave the property map out
# named properties as typed fields
Layout(properties=pyarrow.struct([("title", pyarrow.string())]))
Layout(labels="text")           # labels as plain strings, not dictionary encoded
Layout(vectors="list")          # vectors as variable-length lists
Layout(sparse="dense")          # sparse vectors expanded to dense

>>> import pyarrow
>>> to_arrow(
...     conn.execute_query("MATCH (m:movie) RETURN m ORDER BY id(m) LIMIT 2"),
...     layout=Layout(
...         elements="columns",
...         properties=pyarrow.struct([("title", pyarrow.string()), ("year", pyarrow.int64())]),
...     ),
... ).to_pylist()
[{'m.id': 844424930131969, 'm.label': 'movie', 'm.properties': {'title': 'The Matrix', 'year': 1999}}, {'m.id': 844424930131970, 'm.label': 'movie', 'm.properties': {'title': 'The Matrix Reloaded', 'year': 2003}}]

Embeddings#

An embedding becomes a fixed-size list of 32-bit floats, built from the bytes on the wire with no Python float in between. Ask for the composite rendering when you read one, which sends the numbers themselves rather than their decimal text:

>>> conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
<agensgraph.connection.Cursor ...>
>>> conn.register_vectors()
('vector', 'halfvec', 'sparsevec')
>>> to_arrow(conn.execute_query(
...     "MATCH (m:movie) RETURN m.title AS title, m.embedding::vector(4) AS embedding", binary_=True
... )).schema
title: string
embedding: fixed_size_list<item: float>[4]
  child 0, item: float

Embedding vectors explains register_vectors().

Larger than memory#

to_arrow() and its siblings hold the whole result. For one too large for that, batches() and reader() read a chunk at a time and never hold more than one chunk of Python objects:

>>> from agensgraph import columnar
>>> with conn.transaction():
...     sizes = [batch.num_rows for batch in columnar.batches(
...         conn.stream("MATCH (m:movie) RETURN m.title AS title, m.year AS year"),
...         keys=["title", "year"], size=5,
...     )]
>>> sizes
[5, 5, 5, 3]

The first chunk settles the schema, and every later chunk is held to it, so the batches always concatenate into one table. When the first chunk cannot settle a column, for example because it is all nulls, or a number column holds integers in one chunk and decimals in the next, the export stops with a ValueError naming the column; pass schema= with a pyarrow.Schema to settle it in advance.

reader() returns a pyarrow.RecordBatchReader, which is what Arrow’s writers take. To write a label straight to Parquet through a server-side cursor:

import pyarrow.dataset

with conn.transaction(), conn.cursor(name="export") as cursor:
    cursor.execute("SELECT * FROM (MATCH (m:movie) RETURN m.title AS title, m.year AS year) AS t")
    pyarrow.dataset.write_dataset(columnar.reader(cursor, size=8192), "movies", format="parquet")

Passing size= to to_arrow() with a cursor or a stream also reads in chunks, but still builds one table at the end.

In awaiting code#

A cursor or a stream that has to be awaited needs the awaiting forms, async_batches, async_to_arrow, async_to_pandas and async_to_polars:

from agensgraph.columnar import async_batches, async_to_arrow

async with aconn.transaction():
    table = await async_to_arrow(
        aconn.stream("MATCH (m:movie) RETURN m.title AS title"), keys=["title"]
    )

    stream = aconn.stream("MATCH (m:movie) RETURN m")
    async for batch in async_batches(stream, keys=["m"], size=8192):
        writer.write_batch(batch)

There is no awaiting reader(): a RecordBatchReader is pulled from, and cannot pull from an awaiting source without blocking the event loop. Loop over async_batches() and hand each batch to the writer, as above.

Without a library#

columns() does the same transposition into plain Python lists, with identities as text and graph values as dicts. It takes the rows and the column names:

>>> columnar.columns(result.records, result.keys)
{'title': ['Blade Runner', 'The Terminator', 'Total Recall', 'Terminator 2: Judgment Day'], 'year': [1982, 1984, 1990, 1991], 'genre': ['Sci-Fi', 'Action', 'Sci-Fi', 'Action']}

Two columns with the same name are refused, since a dict cannot hold both; name them apart with AS, or use to_arrow(), which keeps both.

Loading a frame#

The other direction, writing a frame into the graph, is load_vertex_frame() and load_edge_frame().