///|
fn add_buffer(
body : Array[Byte],
buffers : Array[(Int, Int)],
bytes : Bytes,
) -> Unit {
align_bytes(body, 8)
buffers.push((body.length(), bytes.length()))
append_bytes(body, bytes)
}
///|
fn encode_column(
body : Array[Byte],
nodes : Array[(Int, Int)],
buffers : Array[(Int, Int)],
column : Column,
) -> Unit raise ArrowError {
let n = column.length()
let nulls = column.null_count()
nodes.push((n, nulls))
if column is Nulls(_) {
return
}
let bitmap : Array[Byte] = []
if nulls > 0 {
ignore(reserve(bitmap, n / 8 + (if n % 8 != 0 { 1 } else { 0 })))
for i in 0.. ()
Booleans(a) => {
ignore(reserve(values, n / 8 + (if n % 8 != 0 { 1 } else { 0 })))
for i in 0..
for value in a {
put_fixed(values, reserve(values, 1), value.unwrap_or(0).to_int64(), 1)
}
Int16s(a) =>
for value in a {
put_fixed(values, reserve(values, 2), value.unwrap_or(0).to_int64(), 2)
}
Int32s(a) =>
for value in a {
put_i32(values, reserve(values, 4), value.unwrap_or(0))
}
Int64s(a) =>
for value in a {
put_i64(values, reserve(values, 8), value.unwrap_or(0L))
}
UInt8s(a) =>
for value in a {
put_fixed(values, reserve(values, 1), value.unwrap_or(0).to_int64(), 1)
}
UInt16s(a) =>
for value in a {
put_fixed(values, reserve(values, 2), value.unwrap_or(0).to_int64(), 2)
}
UInt32s(a) =>
for value in a {
put_i32(
values,
reserve(values, 4),
value.unwrap_or(0U).reinterpret_as_int(),
)
}
UInt64s(a) =>
for value in a {
put_i64(
values,
reserve(values, 8),
value.unwrap_or(0UL).reinterpret_as_int64(),
)
}
Float32s(a) =>
for value in a {
put_i32(
values,
reserve(values, 4),
value.unwrap_or(0.0F).reinterpret_as_int(),
)
}
Float64s(a) =>
for value in a {
put_i64(
values,
reserve(values, 8),
value.unwrap_or(0.0).reinterpret_as_int64(),
)
}
Date32s(a) =>
for value in a {
put_i32(values, reserve(values, 4), value.unwrap_or(0))
}
Date64s(a) =>
for value in a {
put_i64(values, reserve(values, 8), value.unwrap_or(0L))
}
Timestamps(a, _, _) | Durations(a, _) =>
for value in a {
put_i64(values, reserve(values, 8), value.unwrap_or(0L))
}
Lists(offsets, _, child) => {
for offset in offsets {
put_i32(values, reserve(values, 4), offset)
}
add_buffer(body, buffers, Bytes::from_array(values))
encode_column(body, nodes, buffers, child)
return
}
LargeLists(offsets, _, child) => {
for offset in offsets {
put_i64(values, reserve(values, 8), offset.to_int64())
}
add_buffer(body, buffers, Bytes::from_array(values))
encode_column(body, nodes, buffers, child)
return
}
Maps(offsets, _, child, _) => {
for offset in offsets {
put_i32(values, reserve(values, 4), offset)
}
add_buffer(body, buffers, Bytes::from_array(values))
encode_column(body, nodes, buffers, child)
return
}
FixedLists(_, child, _) => {
encode_column(body, nodes, buffers, child)
return
}
Structs(_, children) => {
for child in children {
encode_column(body, nodes, buffers, child)
}
return
}
Dictionaries(indices, _, _, _) =>
for index in indices {
put_i32(values, reserve(values, 4), index.unwrap_or(0))
}
FixedBinaries(a, width) =>
for value in a {
match value {
Some(bytes) => append_bytes(values, bytes)
None => ignore(reserve(values, width))
}
}
Strings(_) | Binaries(_) | LargeStrings(_) | LargeBinaries(_) => {
let offsets : Array[Byte] = []
let large = column is LargeStrings(_) || column is LargeBinaries(_)
if large {
put_i64(offsets, reserve(offsets, 8), 0L)
} else {
put_i32(offsets, reserve(offsets, 4), 0)
}
for i in 0..
match a[i] {
Some(s) => @utf8.encode(s)
None => b""
}
Binaries(a) => a[i].unwrap_or(b"")
LargeStrings(a) =>
match a[i] {
Some(s) => @utf8.encode(s)
None => b""
}
LargeBinaries(a) => a[i].unwrap_or(b"")
_ => b""
}
if bytes.length() > 0x7fffffff - values.length() {
raise LimitExceeded("32-bit binary offsets")
}
append_bytes(values, bytes)
if large {
put_i64(offsets, reserve(offsets, 8), values.length().to_int64())
} else {
put_i32(offsets, reserve(offsets, 4), values.length())
}
}
add_buffer(body, buffers, Bytes::from_array(offsets))
}
}
add_buffer(body, buffers, Bytes::from_array(values))
}
///|
fn encode_body(batch : RecordBatch) -> BodyLayout raise ArrowError {
batch.validate()
let body : Array[Byte] = []
let nodes = []
let buffers = []
for column in batch.columns {
encode_column(body, nodes, buffers, column)
}
align_bytes(body, 8)
{ body: Bytes::from_array(body), nodes, buffers, }
}
///|
fn body_buffer(frame : Frame, base : Int, i : Int) -> Bytes raise ArrowError {
let offset = bounded_i64(
read_i64(frame.metadata, base + i * 16),
frame.body.length(),
"buffer offset",
)
let size = bounded_i64(
read_i64(frame.metadata, base + i * 16 + 8),
frame.body.length(),
"buffer size",
)
slice_bytes(frame.body, offset, size)
}
///|
priv struct BodyCursor {
nodes : Int
node_count : Int
mut node : Int
buffers : Int
buffer_count : Int
mut buffer : Int
mut remaining_values : Int
}
///|
fn field_layout(field : Field) -> (Int, Int) {
let mut nodes = 1
let mut buffers = match field.data_type {
Null => 0
Utf8 | Binary | LargeUtf8 | LargeBinary => 3
List | LargeList | Map(_) => 2
Struct | FixedSizeList(_) => 1
Dictionary(_, _, _) | FixedSizeBinary(_) => 2
_ => 2
}
for child in field.children {
let (n, b) = field_layout(child)
nodes += n
buffers += b
}
(nodes, buffers)
}
///|
fn next_body_buffer(
frame : Frame,
cursor : BodyCursor,
) -> Bytes raise ArrowError {
if cursor.buffer >= cursor.buffer_count {
raise Invalid("missing body buffer")
}
let result = body_buffer(frame, cursor.buffers, cursor.buffer)
cursor.buffer += 1
result
}
///|
fn decode_column(
frame : Frame,
field : Field,
expected : Int,
cursor : BodyCursor,
limits : ReadLimits,
depth : Int,
dictionaries : Array[DictionaryEntry],
) -> Column raise ArrowError {
if depth > limits.max_depth {
raise LimitExceeded("nested column depth")
}
if cursor.node >= cursor.node_count {
raise Invalid("missing field node")
}
let metadata = frame.metadata
let base = cursor.nodes + cursor.node * 16
cursor.node += 1
let rows = bounded_i64(
read_i64(metadata, base),
limits.max_values_per_batch,
"field node length",
)
let nulls = bounded_i64(read_i64(metadata, base + 8), rows, "null count")
if expected >= 0 && rows != expected {
raise Invalid("field node length differs from parent")
}
if rows > cursor.remaining_values {
raise LimitExceeded("values per batch")
}
cursor.remaining_values -= rows
if field.data_type == Null {
if nulls != rows {
raise Invalid("Null column null count differs from length")
}
return Nulls(rows)
}
let validity = next_body_buffer(frame, cursor)
let bitmap_size = rows / 8 + (if rows % 8 != 0 { 1 } else { 0 })
if (!validity.is_empty() || nulls > 0) && validity.length() < bitmap_size {
raise Invalid("short or missing validity bitmap")
}
let valid : Array[Bool] = []
let mut actual_nulls = 0
for j in 0..= offsets_bytes.length() / width {
raise Invalid("short List offsets")
}
let offsets = []
let mut previous = if width == 8 {
bounded_i64(
read_i64(offsets_bytes, 0),
limits.max_values_per_batch,
"LargeList offset",
)
} else {
read_i32(offsets_bytes, 0)
}
if previous < 0 {
raise Invalid("negative List offset")
}
offsets.push(previous)
for i in 0.. child.length() {
raise Invalid("List offset exceeds child length")
}
if field.data_type == LargeList {
return LargeLists(offsets, valid, child)
}
if field.data_type is Map(sorted) {
return Maps(offsets, valid, child, sorted)
}
return Lists(offsets, valid, child)
}
if field.data_type == Struct {
let children = []
for child_field in field.children {
children.push(
decode_column(
frame,
child_field,
rows,
cursor,
limits,
depth + 1,
dictionaries,
),
)
}
return Structs(valid, children)
}
if field.data_type is FixedSizeList(size) {
if field.children.length() != 1 || rows > limits.max_values_per_batch / size {
raise LimitExceeded("FixedSizeList child values")
}
let child = decode_column(
frame,
field.children[0],
rows * size,
cursor,
limits,
depth + 1,
dictionaries,
)
return FixedLists(valid, child, size)
}
let values = next_body_buffer(frame, cursor)
if field.data_type is Dictionary(id, _, ordered) {
if rows > values.length() / 4 {
raise Invalid("short dictionary index buffer")
}
let indices = Array::makei(rows, j => {
if valid[j] {
Some(read_i32(values, j * 4))
} else {
None
}
})
return Dictionaries(
indices,
dictionary_value(dictionaries, id),
id,
ordered,
)
}
match field.data_type {
Null => Column::Nulls(rows)
Boolean => {
if values.length() < bitmap_size {
raise Invalid("short boolean bitmap")
}
Booleans(
Array::makei(rows, j => {
if valid[j] {
Some((values[j / 8].to_int() & (1 << (j % 8))) != 0)
} else {
None
}
}),
)
}
Int8 | Int16 | UInt8 | UInt16 | UInt32 | UInt64 | Float32 => {
let width = match field.data_type {
Int8 | UInt8 => 1
Int16 | UInt16 => 2
UInt32 | Float32 => 4
_ => 8
}
if rows > values.length() / width {
raise Invalid("short fixed-width buffer")
}
match field.data_type {
Int8 =>
Int8s(
Array::makei(rows, j => {
if valid[j] {
Some(signed_fixed(read_fixed(values, j * width, width), width))
} else {
None
}
}),
)
Int16 =>
Int16s(
Array::makei(rows, j => {
if valid[j] {
Some(signed_fixed(read_fixed(values, j * width, width), width))
} else {
None
}
}),
)
UInt8 =>
UInt8s(
Array::makei(rows, j => {
if valid[j] {
Some(read_fixed(values, j * width, width).to_int())
} else {
None
}
}),
)
UInt16 =>
UInt16s(
Array::makei(rows, j => {
if valid[j] {
Some(read_fixed(values, j * width, width).to_int())
} else {
None
}
}),
)
UInt32 =>
UInt32s(
Array::makei(rows, j => {
if valid[j] {
Some(read_i32(values, j * width).reinterpret_as_uint())
} else {
None
}
}),
)
UInt64 =>
UInt64s(
Array::makei(rows, j => {
if valid[j] {
Some(read_i64(values, j * width).reinterpret_as_uint64())
} else {
None
}
}),
)
Float32 =>
Float32s(
Array::makei(rows, j => {
if valid[j] {
Some(Float::reinterpret_from_int(read_i32(values, j * width)))
} else {
None
}
}),
)
_ => raise Invalid("internal fixed-width type")
}
}
Int32 => {
if rows > values.length() / 4 {
raise Invalid("short Int32 buffer")
}
let result = []
for j in 0.. {
if rows > values.length() / 8 {
raise Invalid("short Int64 buffer")
}
let result = []
for j in 0.. {
if rows > values.length() / 8 {
raise Invalid("short Float64 buffer")
}
let result = []
for j in 0.. {
if rows > values.length() / 4 {
raise Invalid("short Date32 buffer")
}
Date32s(
Array::makei(rows, j => {
if valid[j] {
Some(read_i32(values, j * 4))
} else {
None
}
}),
)
}
Date64 | Timestamp(_, _) | Duration(_) => {
if rows > values.length() / 8 {
raise Invalid("short temporal buffer")
}
let result = Array::makei(rows, j => {
if valid[j] {
Some(read_i64(values, j * 8))
} else {
None
}
})
match field.data_type {
Date64 => Date64s(result)
Timestamp(unit, timezone) => Timestamps(result, unit, timezone)
Duration(unit) => Durations(result, unit)
_ => raise Invalid("internal temporal type")
}
}
FixedSizeBinary(width) => {
if rows > values.length() / width {
raise Invalid("short fixed binary buffer")
}
FixedBinaries(
Array::makei(rows, j => {
if valid[j] {
Some(slice_bytes(values, j * width, width))
} else {
None
}
}),
width,
)
}
Utf8 | Binary | LargeUtf8 | LargeBinary => {
let bytes = next_body_buffer(frame, cursor)
let large = field.data_type == LargeUtf8 || field.data_type == LargeBinary
let offset_width = if large { 8 } else { 4 }
if rows >= values.length() / offset_width {
raise Invalid("short offsets buffer")
}
let strings = []
let binaries = []
let mut start = if large {
bounded_i64(read_i64(values, 0), bytes.length(), "large binary offset")
} else {
read_i32(values, 0)
}
if start < 0 || start > bytes.length() {
raise Invalid("invalid first binary offset")
}
for j in 0.. bytes.length() {
raise Invalid("non-monotonic or out-of-range offset")
}
if field.data_type == Utf8 || field.data_type == LargeUtf8 {
strings.push(
if valid[j] {
Some(decode_utf8(bytes, start, end - start))
} else {
None
},
)
} else {
binaries.push(
if valid[j] {
Some(slice_bytes(bytes, start, end - start))
} else {
None
},
)
}
start = end
}
match field.data_type {
Utf8 => Strings(strings)
LargeUtf8 => LargeStrings(strings)
Binary => Binaries(binaries)
LargeBinary => LargeBinaries(binaries)
_ => raise Invalid("internal variable binary type")
}
}
List | LargeList | FixedSizeList(_) | Map(_) | Struct =>
raise Invalid("internal nested dispatch")
Dictionary(_, _, _) => raise Invalid("internal dictionary dispatch")
}
}
///|
fn decode_body(
frame : Frame,
schema : Schema,
limits : ReadLimits,
dictionaries : Array[DictionaryEntry],
) -> RecordBatch raise ArrowError {
if frame.kind != 3 {
raise Unsupported("expected RecordBatch message")
}
let metadata = frame.metadata
let table = frame.header
if field_ref(metadata, table, 3) is Some(_) {
raise Unsupported("compressed IPC body")
}
let (_, variadic) = vector_info(metadata, table, 4, 8, limits.max_fields)
if variadic != 0 {
raise Unsupported("variadic buffers")
}
let rows = bounded_i64(
field_long(metadata, table, 0),
limits.max_rows_per_batch,
"rows per batch",
)
let mut expected_nodes = 0
let mut expected_buffers = 0
for field in schema.fields {
let (n, b) = field_layout(field)
expected_nodes += n
expected_buffers += b
}
if expected_nodes > limits.max_fields {
raise LimitExceeded("field node count")
}
let (nodes, node_count) = vector_info(metadata, table, 1, 16, expected_nodes)
if node_count != expected_nodes {
raise Invalid("field node count differs from schema")
}
let (buffers, buffer_count) = vector_info(
metadata, table, 2, 16, expected_buffers,
)
if buffer_count != expected_buffers {
raise Invalid("buffer count differs from schema")
}
let mut remaining_bytes = limits.max_body_bytes
for i in 0..