magmalake parquet

parquet · API

parquet.writer

ParquetWriter — write Arrow arrays back out as a Parquet file.

The writer takes exactly what the reader produces: an ArrayArena and the indices of the top-level arrays in it. That makes a round trip one line, and it means anything the reader can decode, the writer can write back.

from parquet import ParquetReader, ParquetWriter, WriterOptions

var r = ParquetReader.open("in.parquet")
var t = r.read_table()

var opts = WriterOptions()
opts.codec = CompressionCodec.SNAPPY.value
var w = ParquetWriter(opts^)
for b in t.batches:
    w.write_batch(b.arena, b.roots)
var bytes = w^.finish()

What it writes: Parquet 2.6 metadata with v1 data pages, PLAIN and RLE_DICTIONARY values, RLE levels, three-level lists, key_value maps, structs, field ids, per-chunk statistics with the sort order each logical type requires, an OffsetIndex and a ColumnIndex per chunk, and key/value metadata. Decimals go out as FIXED_LEN_BYTE_ARRAY(16), which is legal for every precision Arrow can hold.

Structs

WriterOptions

struct WriterOptions

How to write. Every field has a sensible default.

Implements Copyable, Defaultable, Movable

Fields
codec: Int32
row_group_size: Int
data_page_size: Int: Target uncompressed bytes of values per data page.
use_dictionary: Bool
write_statistics: Bool
write_page_index: Bool
created_by: String

__init__

def __init__(out self)

WNode

struct WNode

One node of the tree the writer walks, mirroring the Arrow field tree.

Implements Copyable, Defaultable, Movable

Fields
name: String
kind: Int
arrow: ArrowType
nullable: Bool
field_id: Int32
children: List[Int]
def_level: Int
rep_level: Int
elem_def_level: Int
elem_rep_level: Int
leaf: Int

__init__

def __init__(out self)

WLeaf

struct WLeaf

A column the writer will produce chunks for.

Implements Copyable, Defaultable, Movable

Fields
path: List[String]
arrow: ArrowType
physical: Int32
type_length: Int
max_def: Int
max_rep: Int
field_id: Int32

__init__

def __init__(out self)

WriteSchema

struct WriteSchema

The writer's mirror of the Arrow field tree, plus the Parquet schema.

Implements Copyable, Defaultable, Movable

Fields
nodes: List[WNode]
roots: List[Int]
leaves: List[WLeaf]
elements: List[SchemaElement]

__init__

def __init__(out self)

LeafBuffer

struct LeafBuffer

Implements Copyable, Defaultable, Movable

Fields
defs: List[UInt16]
reps: List[UInt16]
values: PhysBuffer

__init__

def __init__(out self)

DictBuilder

struct DictBuilder

Distinct values in first-seen order, with an open-addressed index.

The index is a flat power-of-two table of Int32 dictionary indices with linear probing — one array, no per-bucket allocation, no Dict and no List of collisions.

A 4- or 8-byte fixed-width column sets key_width and is indexed by the integer value itself: the hash is a multiply-shift of the key and the comparison is an integer compare, so a value costs no span construction and no byte loop at all. Everything else hashes and compares bytes.

Implements Defaultable, Movable

Fields
values: PhysBuffer
slots: List[Int32]
mask: UInt64
used: Int
key_width: Int: 4 or 8 for the integer-keyed mode, 0 for the byte-keyed one.

__init__

def __init__(out self)
def __init__(out self, kind: Int, width: Int, expect: Int = Int(0))

index_of_key

def index_of_key(mut self, key: UInt64) -> Int

The dictionary index of an integer-keyed value, adding it if new.

index_of

def index_of(mut self, v: Span[UInt8]) -> Int

ParquetWriter

struct ParquetWriter[Codecs: CodecSet = DefaultCodecs]

Writes Arrow arrays out as Parquet. One write_batch per record batch.

Implements Movable

Fields
options: WriterOptions
schema: WriteSchema
out: List[UInt8]
row_groups: List[RowGroup]
offset_indexes: List[List[OffsetIndex]]
column_indexes: List[List[ColumnIndex]]
num_rows: Int64
kv: List[KeyValue]

__init__

def __init__(out self, var options: WriterOptions)

add_metadata

def add_metadata(mut self, key: StringSpan, value: StringSpan)

write_batch

def write_batch(mut self, arena: ArrayArena, roots: List[Int])

Write one record batch, split into row groups of row_group_size.

finish

def finish(deinit self) -> List[UInt8]

Close the file and return its bytes.

Functions

physical_for

def physical_for(t: ArrowType) -> Tuple[Int32, Int]

The Parquet physical type and type_length for an Arrow type.

annotate

def annotate(mut el: SchemaElement, t: ArrowType)

Give a SchemaElement the LogicalType and ConvertedType of t.

w_def

def w_def(s: WriteSchema, node: Int) -> Int

build_write_schema

def build_write_schema(arena: ArrayArena, roots: List[Int]) -> WriteSchema

Derive the Parquet schema from the Arrow arrays that will be written.

append_physical

def append_physical(mut out: PhysBuffer, leaf: WLeaf, a: ArrayData, i: Int)

Append value i of a in the leaf's physical Parquet representation.

shred_flat

def shred_flat(s: WriteSchema, w: Int, arena: ArrayArena, node: Int, r0: Int, r1: Int, mut cols: List[LeafBuffer]) -> Bool

Shred a whole run of rows of an unnested column in one pass per buffer.

A top-level primitive column has one slot per row and no repetition, so its definition levels are a fill (with the null slots knocked back to zero) and its values are a copy. Returns False for anything that needs the recursive path — a nested column, or an Arrow type whose bytes are not the Parquet bytes.

shred

def shred(s: WriteSchema, w: Int, arena: ArrayArena, node: Int, i: Int, parent_def: Int, rep: Int, mut cols: List[LeafBuffer])

Write one value of one Arrow array into the leaf buffers under it.

min_max

def min_max(leaf: WLeaf, values: PhysBuffer, start: Int, count: Int) -> Tuple[List[UInt8], List[UInt8], Bool]

The bounds of count values from start, in this column's sort order.

column_statistics

def column_statistics(leaf: WLeaf, buf: LeafBuffer, null_count: Int) -> Statistics

dict_value_cap

def dict_value_cap(n_values: Int) -> Int

The largest dictionary index n_values values may produce, or abandon.

Half the chunk is the older, weaker bound and still applies to a short chunk, where 8192 distinct values are not reachable anyway; the floor of 64 keeps a tiny chunk — which cannot pay for a dictionary page either way — from abandoning on its second value.

Aliases

comptime WN_PRIM = 0
comptime WN_LIST = 1
comptime WN_MAP = 2
comptime WN_STRUCT = 3
comptime CREATED_BY = "parquet.mojo version 0.3.1"
comptime MM_GENERIC = 0

This type needs the generic byte comparator.

comptime MM_BOUNDS = 1

Bounds found and written.

comptime MM_NONE = 2

Handled, but there is nothing to report: every value was NaN.

comptime DICT_MAX_VALUES = 8192

How many distinct values a column chunk's dictionary may hold.

Past this the build is abandoned and the chunk is written PLAIN, so this is both the size rule and the ceiling on the work a chunk that will never be dictionary-encoded can cost: DICT_MAX_VALUES inserts, not one per value.

8192 is where a dictionary stops paying on compressible data. Writing 65,536 values with SNAPPY, dictionary against PLAIN, on values drawn cyclically from a pool of that size:

distinct INT64 DOUBLE BYTE_ARRAY
512 0.23× 0.23× 0.14×
2048 0.47× 0.47× 0.29×
8192 0.74× 0.74× 0.53×
16384 1.02× 1.02× 0.79×
32768 1.28× 1.31× 1.10×

At 32768 — which is what a cap of n_values / 2 allows — the dictionary makes the fixed-width chunks 30% bigger. On values that do not compress the dictionary keeps winning further out (0.82× at 32768 distinct INT64), so a chunk with 8k–32k distinct incompressible values is the case this cap gives up on; a mid-chunk fallback that kept what it had already encoded, the way parquet-cpp does, would not have to choose.

The dictionary indices are UInt16, so the cap can never exceed 65535.