///|
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_body(batch : RecordBatch) -> BodyLayout raise ArrowError {
batch.validate()
let body : Array[Byte] = []
let nodes = []
let buffers = []
for column in batch.columns {
let n = batch.num_rows
let nulls = column.null_count()
nodes.push((n, nulls))
if column is Nulls(_) {
continue
}
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_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))
}
Float64s(a) =>
for value in a {
put_i64(
values,
reserve(values, 8),
value.unwrap_or(0.0).reinterpret_as_int64(),
)
}
Strings(_) | Binaries(_) => {
let offsets : Array[Byte] = []
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"")
_ => b""
}
if bytes.length() > 0x7fffffff - values.length() {
raise LimitExceeded("32-bit binary offsets")
}
append_bytes(values, bytes)
put_i32(offsets, reserve(offsets, 4), values.length())
}
add_buffer(body, buffers, Bytes::from_array(offsets))
}
}
add_buffer(body, buffers, Bytes::from_array(values))
}
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)
}
///|
fn decode_body(
frame : Frame,
schema : Schema,
limits : ReadLimits,
) -> 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 count = schema.fields.length()
if count > 0 && rows > limits.max_values_per_batch / count {
raise LimitExceeded("values per batch")
}
let (nodes, node_count) = vector_info(
metadata,
table,
1,
16,
limits.max_fields,
)
if node_count != count {
raise Invalid("field node count differs from schema")
}
let mut expected_buffers = 0
for field in schema.fields {
expected_buffers += match field.data_type {
Null => 0
Utf8 | Binary => 3
_ => 2
}
}
let (buffers, buffer_count) = vector_info(
metadata, table, 2, 16, expected_buffers,
)
if buffer_count != expected_buffers {
raise Invalid("buffer count differs from schema")
}
// Bound copied buffers even if malformed metadata aliases the same body range
// many times. This is a byte budget, not an exact process-memory limit.
let mut remaining_bytes = limits.max_body_bytes
for i in 0.. 0 && validity.length() < bitmap_size {
raise Invalid("missing validity bitmap")
}
let valid : Array[Bool] = []
let mut actual_nulls = 0
for j in 0.. 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
}
}),
)
}
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.. {
let bytes = body_buffer(frame, buffers, cursor)
cursor += 1
if rows >= values.length() / 4 {
raise Invalid("short offsets buffer")
}
let strings = []
let binaries = []
let mut start = 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 {
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
}
if field.data_type == Utf8 {
Strings(strings)
} else {
Binaries(binaries)
}
}
}
columns.push(column)
}
RecordBatch::new(schema, columns, num_rows=rows)
}