Skip to content

Commit

Permalink
Add support for storage transformers.
Browse files Browse the repository at this point in the history
  • Loading branch information
zoj613 committed Jul 9, 2024
1 parent aa31327 commit 912624e
Show file tree
Hide file tree
Showing 8 changed files with 142 additions and 25 deletions.
85 changes: 85 additions & 0 deletions lib/extensions.ml
Original file line number Diff line number Diff line change
Expand Up @@ -187,3 +187,88 @@ module Datatype = struct
| `String "nativeint" -> Ok Nativeint
| _ -> Error ("Unsupported metadata data_type")
end

type tf_error =
[ `Store_read of string
| `Store_write of string ]

module type STF = sig
type t
val get : t -> string -> (string, [> tf_error]) result
val set : t -> string -> string -> unit
val erase : t -> string -> unit
end

module StorageTransformers = struct
type transformer =
| Identity
type t = transformer list

let default = [Identity]

let deserialize x =
match
Util.get_name x,
Yojson.Safe.Util.(member "configuration" x)

Check warning on line 212 in lib/extensions.ml

View check run for this annotation

Codecov / codecov/patch

lib/extensions.ml#L210-L212

Added lines #L210 - L212 were not covered by tests
with
| "identity", `Null -> Ok Identity
| _ ->

Check warning on line 215 in lib/extensions.ml

View check run for this annotation

Codecov / codecov/patch

lib/extensions.ml#L214-L215

Added lines #L214 - L215 were not covered by tests
Error "Unsupported storage transformer name or configuration."

let of_yojson x =
let open Util.Result_syntax in

Check warning on line 219 in lib/extensions.ml

View check run for this annotation

Codecov / codecov/patch

lib/extensions.ml#L219

Added line #L219 was not covered by tests
List.fold_right
(fun x acc ->
acc >>= fun l ->
deserialize x >>| fun s ->
s :: l) (Yojson.Safe.Util.to_list x) (Ok [])

Check warning on line 224 in lib/extensions.ml

View check run for this annotation

Codecov / codecov/patch

lib/extensions.ml#L222-L224

Added lines #L222 - L224 were not covered by tests

let to_yojson x =
`List
(List.fold_right

Check warning on line 228 in lib/extensions.ml

View check run for this annotation

Codecov / codecov/patch

lib/extensions.ml#L227-L228

Added lines #L227 - L228 were not covered by tests
(fun x acc ->
match x with
| Identity -> acc) x [])

Check warning on line 231 in lib/extensions.ml

View check run for this annotation

Codecov / codecov/patch

lib/extensions.ml#L230-L231

Added lines #L230 - L231 were not covered by tests

let get
(type a)
(module M : STF with type t = a)
(store : a)
(transformers : t)
(key : string)
=
let open Util.Result_syntax in
M.get store key >>| fun raw ->
snd @@
List.fold_right
(fun x (k, v) ->
match x with
| Identity -> (k, v)) transformers (key, raw)

Check warning on line 246 in lib/extensions.ml

View check run for this annotation

Codecov / codecov/patch

lib/extensions.ml#L245-L246

Added lines #L245 - L246 were not covered by tests

let set
(type a)
(module M : STF with type t = a)
(store : a)
(transformers : t)
(key : string)
(value : string)
=
let k', v' =
List.fold_left
(fun (k, v) -> function
| Identity -> (k, v)) (key, value) transformers

Check warning on line 259 in lib/extensions.ml

View check run for this annotation

Codecov / codecov/patch

lib/extensions.ml#L259

Added line #L259 was not covered by tests
in
M.set store k' v'

let erase
(type a)
(module M : STF with type t = a)
(store : a)
(transformers : t)
(key : string)
=
M.erase store @@
List.fold_left
(fun k -> function
| Identity -> k) key transformers

Check warning on line 273 in lib/extensions.ml

View check run for this annotation

Codecov / codecov/patch

lib/extensions.ml#L273

Added line #L273 was not covered by tests
end
27 changes: 27 additions & 0 deletions lib/extensions.mli
Original file line number Diff line number Diff line change
Expand Up @@ -46,3 +46,30 @@ module Datatype : sig
val of_yojson : Yojson.Safe.t -> (t, string) result
val to_yojson : t -> Yojson.Safe.t
end

type tf_error =
[ `Store_read of string
| `Store_write of string ]

module type STF = sig
type t
val get : t -> string -> (string, [> tf_error]) result
val set : t -> string -> string -> unit
val erase : t -> string -> unit
end

module StorageTransformers : sig
type transformer =
| Identity
type t = transformer list

val default : t
val get :
(module STF with type t = 'a) -> 'a -> t -> string -> (string, [> tf_error ]) result
val set :
(module STF with type t = 'a) -> 'a -> t -> string -> string -> unit
val erase :
(module STF with type t = 'a) -> 'a -> t -> string -> unit
val to_yojson : t -> Yojson.Safe.t
val of_yojson : Yojson.Safe.t -> (t, string) result
end
18 changes: 14 additions & 4 deletions lib/metadata.ml
Original file line number Diff line number Diff line change
Expand Up @@ -112,13 +112,14 @@ module ArrayMetadata = struct
;chunk_key_encoding : ChunkKeyEncoding.t
;attributes : Yojson.Safe.t
;dimension_names : string option list
;storage_transformers : Yojson.Safe.t Util.ExtPoint.t list}
;storage_transformers : StorageTransformers.t}

let create
?(sep=`Slash)
?(codecs=Codecs.Chain.default)
?(dimension_names=[])
?(attributes=`Null)
?(storage_transformers=[])
~shape
kind
fv
Expand All @@ -134,7 +135,7 @@ module ArrayMetadata = struct
;dimension_names
;zarr_format = 3
;node_type = "array"
;storage_transformers = []
;storage_transformers
;fill_value = FillValue.of_kind kind fv
;data_type = Datatype.of_kind kind
;chunk_key_encoding = ChunkKeyEncoding.create sep}
Expand Down Expand Up @@ -168,6 +169,12 @@ module ArrayMetadata = struct
| None -> `Null) xs
in
l @ [("dimension_names", `List xs')]
in
let l =
match t.storage_transformers with
| [] | [StorageTransformers.Identity] -> l
| xs ->
l @ [("storage_transformers", StorageTransformers.to_yojson xs)]

Check warning on line 177 in lib/metadata.ml

View check run for this annotation

Codecov / codecov/patch

lib/metadata.ml#L175-L177

Added lines #L175 - L177 were not covered by tests
in `Assoc l

let of_yojson x =
Expand Down Expand Up @@ -256,8 +263,9 @@ module ArrayMetadata = struct
>>= fun dimension_names ->

(match member "storage_transformers" x with
| `Null -> Ok []
| _ -> Error "storage_transformers field is not yet supported.")
| `Null | `List [] -> Ok []

Check warning on line 266 in lib/metadata.ml

View check run for this annotation

Codecov / codecov/patch

lib/metadata.ml#L266

Added line #L266 was not covered by tests
| _ ->
Error "storage_transformers field is not yet supported.")
>>| fun storage_transformers ->

{zarr_format; shape; node_type; data_type; codecs; fill_value; chunk_grid
Expand Down Expand Up @@ -286,6 +294,8 @@ module ArrayMetadata = struct

let attributes t = t.attributes

let storage_transformers t = t.storage_transformers

let chunk_shape t =
RegularGrid.chunk_shape t.chunk_grid

Expand Down
5 changes: 5 additions & 0 deletions lib/metadata.mli
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ module ArrayMetadata : sig
?codecs:Codecs.Chain.t ->
?dimension_names:string option list ->
?attributes:Yojson.Safe.t ->
?storage_transformers:Extensions.StorageTransformers.t ->
shape:int array ->
('a, 'b) Bigarray.kind ->
'a ->
Expand Down Expand Up @@ -76,6 +77,10 @@ module ArrayMetadata : sig
(** [attributes t] Returns a Yojson type containing user attributes assigned
to the zarr array represented by [t]. *)

val storage_transformers : t -> Extensions.StorageTransformers.t
(** [storage_transformers t] Returns the storage transformers to be applied
to the keys and values of this store. *)

val dimension_names : t -> string option list
(** [dimension_name t] returns a list of dimension names. If none are
defined then an empty list is returned. *)
Expand Down
14 changes: 10 additions & 4 deletions lib/storage/storage.ml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ module ArraySet = Util.ArraySet
module Arraytbl = Util.Arraytbl
module AM = Metadata.ArrayMetadata
module GM = Metadata.GroupMetadata
module ST = Extensions.StorageTransformers

module Make (M : STORE) : S with type t = M.t = struct
include M
Expand Down Expand Up @@ -35,6 +36,7 @@ module Make (M : STORE) : S with type t = M.t = struct
?(sep=`Slash)
?(dimension_names=[])
?(attributes=`Null)
?(storage_transformers=[])
?codecs
~shape
~chunks
Expand All @@ -54,6 +56,7 @@ module Make (M : STORE) : S with type t = M.t = struct
~codecs
~dimension_names
~attributes
~storage_transformers
~shape
kind
fill_value
Expand Down Expand Up @@ -154,11 +157,12 @@ module Make (M : STORE) : S with type t = M.t = struct
in
let codecs = AM.codecs meta in
let prefix = ArrayNode.to_key node ^ "/" in
let tf = AM.storage_transformers meta in
let cindices = ArraySet.of_seq @@ Arraytbl.to_seq_keys tbl in
ArraySet.fold (fun idx acc ->
acc >>= fun () ->
let chunkkey = prefix ^ AM.chunk_key meta idx in
(match get t chunkkey with
(match ST.get (module M) t tf chunkkey with
| Ok b ->
Codecs.Chain.decode codecs repr b
| Error _ ->
Expand All @@ -173,7 +177,7 @@ module Make (M : STORE) : S with type t = M.t = struct
List.iter
(fun (c, v) -> Ndarray.set arr c v) @@ Arraytbl.find_all tbl idx;
Codecs.Chain.encode codecs arr >>| fun encoded ->
set t chunkkey encoded) cindices (Ok ())
ST.set (module M) t tf chunkkey encoded) cindices (Ok ())

let get_array
: type a b.
Expand Down Expand Up @@ -206,6 +210,7 @@ module Make (M : STORE) : S with type t = M.t = struct
let tbl = Arraytbl.create @@ Array.length pair in
let prefix = ArrayNode.to_key node ^ "/" in
let chain = AM.codecs meta in
let tf = AM.storage_transformers meta in
let repr =
{kind
;shape = AM.chunk_shape meta
Expand All @@ -217,7 +222,7 @@ module Make (M : STORE) : S with type t = M.t = struct
| Some arr ->
Ok (Ndarray.get arr coord :: l)
| None ->
(match get t @@ prefix ^ AM.chunk_key meta idx with
(match ST.get (module M) t tf @@ prefix ^ AM.chunk_key meta idx with
| Ok b ->
Codecs.Chain.decode chain repr b
| Error _ ->
Expand All @@ -243,8 +248,9 @@ module Make (M : STORE) : S with type t = M.t = struct
ArraySet.of_list @@ AM.chunk_indices meta @@ AM.shape meta in
let s' =
ArraySet.of_list @@ AM.chunk_indices meta shape in
let tf = AM.storage_transformers meta in
ArraySet.iter
(fun v -> erase t @@ pre ^ AM.chunk_key meta v)
(fun v -> ST.erase (module M) t tf @@ pre ^ AM.chunk_key meta v)
ArraySet.(diff s s');
Ok (set t mkey @@ AM.encode @@ AM.update_shape meta shape)
end
Expand Down
1 change: 1 addition & 0 deletions lib/storage/storage_intf.ml
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ module type S = sig
: ?sep:[< `Dot | `Slash > `Slash ] ->
?dimension_names:string option list ->
?attributes:Yojson.Safe.t ->
?storage_transformers:Extensions.StorageTransformers.t ->
?codecs:Codecs.chain ->
shape:int array ->
chunks:int array ->
Expand Down
10 changes: 0 additions & 10 deletions lib/util.ml
Original file line number Diff line number Diff line change
@@ -1,13 +1,3 @@
module ExtPoint = struct
type 'a t =
{name : string
;configuration : 'a}

let ( = ) cmp x y =
(x.name = y.name) &&
cmp x.configuration y.configuration
end

type ('a, 'b) array_repr =
{kind : ('a, 'b) Bigarray.kind
;shape : int array
Expand Down
7 changes: 0 additions & 7 deletions lib/util.mli
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,6 @@ type ('a, 'b) array_repr =
(** The type summarizing the decoded/encoded representation of a Zarr array
or chunk. *)

module ExtPoint : sig
(** The type representing a JSON extension point metadata configuration. *)

type 'a t = {name : string ; configuration : 'a}
val ( = ) : ('a -> 'a -> bool) -> 'a t -> 'a t -> bool
end

module StrMap : sig include Hashtbl.S with type key = string end
(** A hashtable with string keys. *)

Expand Down

0 comments on commit 912624e

Please sign in to comment.