Something went wrong. Try again.
An experimental library which implements a subset of the apache arrow format in ocaml.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466(* See arrow_array.mli for the contract, the per-[dtype] physical layout, and the "no validity buffer" optimization.
Representation note: [buffers] holds the *raw* byte regions in Arrow's canonical order (validity first when [null_count > 0], then the data/offset buffers). The typed builders [Validity]/[Offsets] are used at construction time to produce those raw buffers; the accessors below reinterpret the bytes directly. Layout is a single point of dispatch on [dtype], so adding a new DataType is one new constructor + one arm per accessor it supports. *)
type t = { dtype : Datatype.t; length : int; null_count : int; buffers : Buffer.t list; children : t list; dictionary : t option;}
(* {1 Little-endian readers over a raw data Buffer} *)
let read_u8 buf off = Buffer.get buf offlet read_u16 buf off = read_u8 buf off lor (read_u8 buf (off + 1) lsl 8)
let read_u32 buf off = read_u8 buf off lor (read_u8 buf (off + 1) lsl 8) lor (read_u8 buf (off + 2) lsl 16) lor (read_u8 buf (off + 3) lsl 24)
(* Reinterpret the 32 low bits as a signed int32 (wraps the sign bit). *)let read_i32 buf off = Int32.of_int (read_u32 buf off)
let read_i64 buf off = let b k = Int64.of_int (Buffer.get buf (off + k)) in Int64.logor (b 0) (Int64.logor (Int64.shift_left (b 1) 8) (Int64.logor (Int64.shift_left (b 2) 16) (Int64.logor (Int64.shift_left (b 3) 24) (Int64.logor (Int64.shift_left (b 4) 32) (Int64.logor (Int64.shift_left (b 5) 40) (Int64.logor (Int64.shift_left (b 6) 48) (Int64.shift_left (b 7) 56)))))))
let read_f64 buf off = Int64.float_of_bits (read_i64 buf off)
let read_bytes buf off len = let b = Bytes.create len in Buffer.blit_to_bytes buf ~src_pos:off ~dst:b ~dst_pos:0 ~len; Bytes.unsafe_to_string b
(* {1 Dictionary index layout} *)
let index_byte_width = function | Datatype.UInt8 -> 1 | Datatype.UInt16 -> 2 | Datatype.UInt32 | Datatype.Int32 -> 4 | Datatype.Int64 -> 8 | other -> invalid_arg (Printf.sprintf "Arrow.Array: %s is not a valid dictionary index type" (Datatype.to_string other))
let read_index buf off width = match width with | 1 -> read_u8 buf off | 2 -> read_u16 buf off | 4 -> read_u32 buf off | 8 -> Int64.to_int (read_i64 buf off) | _ -> invalid_arg "Arrow.Array: unsupported dictionary index width"
(* {1 Layout accessors} *)
let dtype t = t.dtypelet length t = t.lengthlet null_count t = t.null_countlet buffers t = t.bufferslet children t = t.childrenlet dictionary_values t = t.dictionary
let child t i = match List.nth_opt t.children i with | Some c -> c | None -> invalid_arg (Printf.sprintf "Arrow.Array.child: index %d out of bounds [0, %d)" i (List.length t.children))
(* The validity buffer is present iff there are nulls; otherwise the data buffers start at [buffers] head. *)let validity_buffer t = if t.null_count > 0 then match t.buffers with b :: _ -> Some b | [] -> None else None
let data_buffers t = match validity_buffer t with Some _ -> List.tl t.buffers | None -> t.buffers
let data_buffer t k = List.nth (data_buffers t) k
(* {1 Validity} *)
let check_index t i = if i < 0 || i >= t.length then invalid_arg (Printf.sprintf "Arrow.Array: index %d out of bounds [0, %d)" i t.length)
let is_valid t i = check_index t i; match validity_buffer t with | None -> true | Some vb -> Buffer.get vb (i lsr 3) land (1 lsl (i land 7)) <> 0
let is_null t i = not (is_valid t i)
(* {1 Typed accessors} *)
let ensure_dtype t expect = if not (Datatype.equal t.dtype expect) then invalid_arg (Printf.sprintf "Arrow.Array: accessor expects a %s array but got %s" (Datatype.to_string expect) (Datatype.to_string t.dtype))
(* Shared shape for a fixed-width primitive accessor: check the dtype, return [None] for a null slot, otherwise read one value of [width] bytes. *)let get_fixed t expect width read i = ensure_dtype t expect; if is_valid t i then Some (read (data_buffer t 0) (i * width)) else None
let get_bool t i = ensure_dtype t Datatype.Boolean; if is_valid t i then let db = data_buffer t 0 in Some (Buffer.get db (i lsr 3) land (1 lsl (i land 7)) <> 0) else None
let get_int32 t i = get_fixed t Datatype.Int32 4 read_i32 ilet get_int64 t i = get_fixed t Datatype.Int64 8 read_i64 ilet get_uint8 t i = get_fixed t Datatype.UInt8 1 read_u8 ilet get_uint16 t i = get_fixed t Datatype.UInt16 2 read_u16 ilet get_uint32 t i = get_fixed t Datatype.UInt32 4 read_u32 ilet get_float64 t i = get_fixed t Datatype.Float64 8 read_f64 ilet get_timestamp t i = get_fixed t Datatype.Timestamp 8 read_i64 i
let get_timestamp_tz t i = (match t.dtype with | Datatype.Timestamp_tz _ -> () | other -> invalid_arg (Printf.sprintf "Arrow.Array: accessor expects a Timestamp_tz array but got %s" (Datatype.to_string other))); if is_valid t i then Some (read_i64 (data_buffer t 0) (i * 8)) else Nonelet get_uint64 t i = get_fixed t Datatype.UInt64 8 read_i64 i
(* Utf8 and Binary share the offsets + data layout. *)let get_var t expect i = ensure_dtype t expect; if is_valid t i then begin let offsets = data_buffer t 0 and data = data_buffer t 1 in let start = read_u32 offsets (i * 4) in let stop = read_u32 offsets ((i + 1) * 4) in Some (read_bytes data start (stop - start)) end else None
let get_utf8 t i = get_var t Datatype.Utf8 ilet get_binary t i = get_var t Datatype.Binary i
let get_fixed_size_binary t i = match t.dtype with | Datatype.Fixed_size_binary w -> if is_valid t i then Some (read_bytes (data_buffer t 0) (i * w) w) else None | other -> invalid_arg (Printf.sprintf "Arrow.Array.get_fixed_size_binary: array is %s, not FixedSizeBinary" (Datatype.to_string other))
let get_list t i = match t.dtype with | Datatype.List _ -> if is_valid t i then begin let offsets = data_buffer t 0 in Some (read_u32 offsets (i * 4), read_u32 offsets ((i + 1) * 4)) end else None | other -> invalid_arg (Printf.sprintf "Arrow.Array.get_list: array is %s, not List" (Datatype.to_string other))
let get_dictionary_index t i = match t.dtype with | Datatype.Dictionary { index; _ } -> if is_valid t i then Some (read_index (data_buffer t 0) (i * index_byte_width index) (index_byte_width index)) else None | other -> invalid_arg (Printf.sprintf "Arrow.Array.get_dictionary_index: array is %s, not Dictionary" (Datatype.to_string other))
(* {1 Construction} *)
(* Resolve the validity argument into a [(null_count, buffers-prefix)] pair, applying the no-validity optimization: an all-valid (or absent) bitmap yields [(0, [])] and no validity buffer. *)let validity_part validity ~length = match validity with | None -> (0, []) | Some v -> if Validity.length v <> length then invalid_arg (Printf.sprintf "Arrow.Array: validity describes %d slots but array length is %d" (Validity.length v) length); let nc = Validity.null_count v in if nc = 0 then (0, []) else (nc, [ Validity.buffer v ])
let make_fixed dtype width ?validity ~length ~data () = if length < 0 then invalid_arg "Arrow.Array: negative length"; if Buffer.length data < length * width then invalid_arg (Printf.sprintf "Arrow.Array: %s data holds %d bytes, need %d for %d slots" (Datatype.to_string dtype) (Buffer.length data) (length * width) length); let null_count, vbufs = validity_part validity ~length in { dtype; length; null_count; buffers = vbufs @ [ data ]; children = []; dictionary = None; }
let boolean ?validity ~length ~data () = if length < 0 then invalid_arg "Arrow.Array: negative length"; let need = (length + 7) / 8 in if Buffer.length data < need then invalid_arg (Printf.sprintf "Arrow.Array: boolean data holds %d bytes, need %d" (Buffer.length data) need); let null_count, vbufs = validity_part validity ~length in { dtype = Datatype.Boolean; length; null_count; buffers = vbufs @ [ data ]; children = []; dictionary = None; }
let int32 ?validity ~length ~data () = make_fixed Datatype.Int32 4 ?validity ~length ~data ()
let int64 ?validity ~length ~data () = make_fixed Datatype.Int64 8 ?validity ~length ~data ()
let uint8 ?validity ~length ~data () = make_fixed Datatype.UInt8 1 ?validity ~length ~data ()
let uint16 ?validity ~length ~data () = make_fixed Datatype.UInt16 2 ?validity ~length ~data ()
let uint32 ?validity ~length ~data () = make_fixed Datatype.UInt32 4 ?validity ~length ~data ()
let uint64 ?validity ~length ~data () = make_fixed Datatype.UInt64 8 ?validity ~length ~data ()
let float64 ?validity ~length ~data () = make_fixed Datatype.Float64 8 ?validity ~length ~data ()
let timestamp ?validity ~length ~data () = make_fixed Datatype.Timestamp 8 ?validity ~length ~data ()
let timestamp_tz ~unit ?tz ?validity ~length ~data () = make_fixed (Datatype.Timestamp_tz { unit; tz }) 8 ?validity ~length ~data ()
(* Utf8 and Binary: an offsets buffer then a data buffer. *)let make_var dtype ?validity ~length ~offsets ~data () = let name = String.lowercase_ascii (Datatype.to_string dtype) in if length < 0 then invalid_arg "Arrow.Array: negative length"; if Offsets.length offsets <> length then invalid_arg (Printf.sprintf "Arrow.Array: %s offsets describe %d slots but length is %d" name (Offsets.length offsets) length); if Buffer.length data < Offsets.last_offset offsets then invalid_arg (Printf.sprintf "Arrow.Array: %s data holds %d bytes, need %d" name (Buffer.length data) (Offsets.last_offset offsets)); let null_count, vbufs = validity_part validity ~length in { dtype; length; null_count; buffers = vbufs @ [ Offsets.buffer offsets; data ]; children = []; dictionary = None; }
let utf8 ?validity ~length ~offsets ~data () = make_var Datatype.Utf8 ?validity ~length ~offsets ~data ()
let binary ?validity ~length ~offsets ~data () = make_var Datatype.Binary ?validity ~length ~offsets ~data ()
let fixed_size_binary ?validity ~length ~byte_width ~data () = if byte_width < 0 then invalid_arg "Arrow.Array: negative byte_width"; make_fixed (Datatype.Fixed_size_binary byte_width) byte_width ?validity ~length ~data ()
let list ?validity ~length ~offsets ~values () = if length < 0 then invalid_arg "Arrow.Array: negative length"; if Offsets.length offsets <> length then invalid_arg (Printf.sprintf "Arrow.Array: list offsets describe %d slots but length is %d" (Offsets.length offsets) length); if Offsets.last_offset offsets > values.length then invalid_arg (Printf.sprintf "Arrow.Array: list offsets reach %d but element array has %d slots" (Offsets.last_offset offsets) values.length); let null_count, vbufs = validity_part validity ~length in { dtype = Datatype.List values.dtype; length; null_count; buffers = vbufs @ [ Offsets.buffer offsets ]; children = [ values ]; dictionary = None; }
let struct_ ?validity ~length ~children () = if length < 0 then invalid_arg "Arrow.Array: negative length"; List.iter (fun c -> if c.length <> length then invalid_arg (Printf.sprintf "Arrow.Array: struct child has %d slots but struct length is %d" c.length length)) children; let null_count, vbufs = validity_part validity ~length in { dtype = Datatype.Struct (List.map (fun c -> c.dtype) children); length; null_count; buffers = vbufs; children; dictionary = None; }
(* Every non-null index must point at a real value-array slot. Reads each live index after the array is assembled so it can reuse [is_valid]/[data_buffer]. *)let validate_dictionary_indices t ~dict_length ~width = let idx = data_buffer t 0 in for i = 0 to t.length - 1 do if is_valid t i then begin let v = read_index idx (i * width) width in if v < 0 || v >= dict_length then invalid_arg (Printf.sprintf "Arrow.Array: dictionary index %d at slot %d out of range [0, %d)" v i dict_length) end done
let dictionary ?validity ~length ~index ~indices ~values () = if length < 0 then invalid_arg "Arrow.Array: negative length"; let width = index_byte_width index in if Buffer.length indices < length * width then invalid_arg (Printf.sprintf "Arrow.Array: dictionary indices hold %d bytes, need %d" (Buffer.length indices) (length * width)); let null_count, vbufs = validity_part validity ~length in let t = { dtype = Datatype.Dictionary { index; values = values.dtype }; length; null_count; buffers = vbufs @ [ indices ]; children = []; dictionary = Some values; } in validate_dictionary_indices t ~dict_length:values.length ~width; t
(* {1 Logical equality} *)
(* Slot [i] of [a] against slot [j] of [b], for arrays of equal dtype: both null, or both valid with equal values. Reads only through the typed accessors, so bytes that carry no meaning (validity padding bits, the values under null slots, buffer slack) are never compared. *)let rec slot_equal a i b j = match (is_valid a i, is_valid b j) with | false, false -> true | true, true -> ( match a.dtype with | Datatype.Boolean -> get_bool a i = get_bool b j | Datatype.Int32 -> get_int32 a i = get_int32 b j | Datatype.Int64 -> get_int64 a i = get_int64 b j | Datatype.UInt8 -> get_uint8 a i = get_uint8 b j | Datatype.UInt16 -> get_uint16 a i = get_uint16 b j | Datatype.UInt32 -> get_uint32 a i = get_uint32 b j | Datatype.UInt64 -> get_uint64 a i = get_uint64 b j | Datatype.Float64 -> Option.equal Float.equal (get_float64 a i) (get_float64 b j) | Datatype.Timestamp -> get_timestamp a i = get_timestamp b j | Datatype.Timestamp_tz _ -> get_timestamp_tz a i = get_timestamp_tz b j | Datatype.Utf8 -> get_utf8 a i = get_utf8 b j | Datatype.Binary -> get_binary a i = get_binary b j | Datatype.Fixed_size_binary _ -> get_fixed_size_binary a i = get_fixed_size_binary b j | Datatype.List _ -> ( match (get_list a i, get_list b j) with | Some (sa, ea), Some (sb, eb) -> ea - sa = eb - sb && let va = child a 0 and vb = child b 0 in let rec loop k = k >= ea - sa || (slot_equal va (sa + k) vb (sb + k) && loop (k + 1)) in loop 0 | _ -> false) | Datatype.Struct _ -> List.for_all2 (fun ca cb -> slot_equal ca i cb j) a.children b.children | Datatype.Dictionary _ -> ( match ( get_dictionary_index a i, get_dictionary_index b j, a.dictionary, b.dictionary ) with | Some ia, Some ib, Some da, Some db -> slot_equal da ia db ib | _ -> false)) | _ -> false
let equal_values a b = Datatype.equal a.dtype b.dtype && a.length = b.length && let rec loop i = i >= a.length || (slot_equal a i b i && loop (i + 1)) in loop 0