implement partial encode/decode for sharding codec.

lib/codecs/array_to_bytes.mli

@@ -15,7 +15,29 @@ module ArrayToBytes : sig
     array_tobytes ->
     ('a, 'b) Util.array_repr ->
     string ->
-    (('a, 'b, Bigarray.c_layout) Bigarray.Genarray.t, [> error]) result
+    (('a, 'b, Bigarray.c_layout) Bigarray.Genarray.t
+    ,[> `Store_read of string | error]) result
   val of_yojson : Yojson.Safe.t -> (array_tobytes, string) result
   val to_yojson : array_tobytes -> Yojson.Safe.t
 end
+
+module ShardingIndexedCodec : sig
+  type t = internal_shard_config
+  val partial_encode :
+    t ->
+    ((int * int option) list ->
+      (string list, [> `Store_read of string | error ] as 'c) result) ->
+    partial_setter ->
+    int ->
+    ('a, 'b) Util.array_repr ->
+    (int array * 'a) list ->
+    (unit, 'c) result
+  val partial_decode :
+    t ->
+    ((int * int option) list ->
+      (string list, [> `Store_read of string | error ] as 'c) result) ->
+    int ->
+    ('a, 'b) Util.array_repr ->
+    (int * int array) list ->
+    ((int * 'a) list, 'c) result
+end

lib/codecs/codecs.ml

@@ -93,11 +93,45 @@
       (fun acc c -> acc >>= BytesToBytes.encode c)
       (ArrayToBytes.encode t.a2b y) t.b2b
 
+  let is_just_sharding : t -> bool = function
+    | {a2a = []; a2b = `ShardingIndexed _; b2b = []} -> true
+    | _ -> false
+
+  let partial_encode :
+    t ->
+    ((int * int option) list ->
+      (string list, [> `Store_read of string | error] as 'c) result) ->
+    partial_setter ->
+    int ->
+    ('a, 'b) Util.array_repr ->
+    (int array * 'a) list ->
+    (unit, 'c) result
+    = fun t f g bsize repr pairs ->
+    match t.a2b with
+    | `ShardingIndexed c ->
+      ShardingIndexedCodec.partial_encode c f g bsize repr pairs
+    | `Bytes _ -> failwith "bytes codec does not support partial encoding." 
+
+  let partial_decode :
+    t ->
+    ((int * int option) list ->
+      (string list, [> `Store_read of string | error ] as 'c) result) ->
+    int ->
+    ('a, 'b) Util.array_repr ->
+    (int * int array) list ->
+    ((int * 'a) list, 'c) result
+    = fun t f s repr pairs ->
+    match t.a2b with
+    | `ShardingIndexed c ->
+      ShardingIndexedCodec.partial_decode c f s repr pairs
+    | `Bytes _ -> failwith "bytes codec does not support partial decoding."
+
   let decode :
     t ->
     ('a, 'b) Util.array_repr ->
     string ->
-    (('a, 'b, Bigarray.c_layout) Bigarray.Genarray.t, [> error]) result
+    (('a, 'b, Bigarray.c_layout) Bigarray.Genarray.t
+    ,[> `Store_read of string | error]) result
     = fun t repr x ->
     List.fold_right
       (fun c acc -> acc >>= BytesToBytes.decode c) t.b2b (Ok x)

lib/codecs/codecs.mli

@@ -21,6 +21,10 @@ module Chain : sig
       the required codecs as defined in the Zarr Version 3 specification. *)
   val default : t
 
+  (** [is_just_sharding t] is [true] if the codec chain [t] contains only
+      the [sharding_indexed] codec. *)
+  val is_just_sharding : t -> bool
+
   (** [encode t x] computes the encoded byte string representation of
       array chunk [x]. Returns an error upon failure. *)
   val encode :
@@ -34,7 +38,27 @@ module Chain : sig
     t ->
     ('a, 'b) Util.array_repr ->
     string ->
-    (('a, 'b, Bigarray.c_layout) Bigarray.Genarray.t, [> error ]) result
+    (('a, 'b, Bigarray.c_layout) Bigarray.Genarray.t
+    ,[> `Store_read of string | error ]) result
+
+  val partial_encode :
+    t ->
+    ((int * int option) list ->
+      (string list, [> `Store_read of string | error ] as 'c) result) ->
+    partial_setter ->
+    int ->
+    ('a, 'b) Util.array_repr ->
+    (int array * 'a) list ->
+    (unit, 'c) result
+
+  val partial_decode :
+    t ->
+    ((int * int option) list ->
+      (string list, [> `Store_read of string | error ] as 'c) result) ->
+    int ->
+    ('a, 'b) Util.array_repr ->
+    (int * int array) list ->
+    ((int * 'a) list, 'c) result
 
   (** [x = y] returns true if chain [x] is equal to chain [y],
       and false otherwise. *)

lib/codecs/codecs_intf.ml

@@ -39,6 +39,7 @@ type error =
   | `CodecChain of string
   | `Sharding of int array * int array * string ]
 
+type partial_setter = ?append:bool -> (int * string) list -> unit
 
 module type Interface = sig
   (** The type of [array -> array] codecs. *)
@@ -102,4 +103,6 @@ module type Interface = sig
     | `Transpose_order of int array * string
     | `CodecChain of string
     | `Sharding of int array * int array * string ]
+
+  type partial_setter = ?append:bool -> (int * string) list -> unit
 end

lib/storage/filesystem.ml

@@ -28,14 +28,49 @@
     with
     | Sys_error _ -> Error (`Store_read fpath)
 
+  let get_partial_values t key ranges =
+    let open Util.Result_syntax in
+    In_channel.with_open_gen
+      In_channel.[Open_rdonly]
+      t.file_perm
+      (key_to_fspath t key)
+      (fun ic ->
+        let size = In_channel.length ic |> Int64.to_int in
+        List.fold_right
+          (fun (rs, len) acc ->
+            acc >>= fun xs ->
+            let len' =
+              match len with
+              | Some l -> l
+              | None -> size - rs
+            in
+            In_channel.seek ic @@ Int64.of_int rs;
+            match In_channel.really_input_string ic len' with
+            | Some s -> Ok (s :: xs)
+            | None ->
+              Error (`Store_read "EOF reached before all bytes are read."))
+          ranges (Ok []))
+
   let set t key value =
     let filename = key_to_fspath t key in
     create_parent_dir filename t.file_perm;
     Out_channel.with_open_gen
       Out_channel.[Open_wronly; Open_trunc; Open_creat]
       t.file_perm
       filename
-      (fun oc -> Out_channel.output_string oc value)
+      (fun oc -> Out_channel.output_string oc value; Out_channel.flush oc)
+
+  let set_partial_values t key ?(append=false) rvs =
+    let open Out_channel in
+    Out_channel.with_open_gen
+      [if append then Open_append else Open_wronly]
+      t.file_perm
+      (key_to_fspath t key)
+      (fun oc ->
+        List.iter
+          (fun (rs, value) ->
+            Out_channel.seek oc @@ Int64.of_int rs;
+            Out_channel.output_string oc value) rvs; Out_channel.flush oc)
 
   let list t =
     let module StrSet = Storage_intf.Base.StrSet in
@@ -68,23 +103,18 @@
     try Sys.remove @@ key_to_fspath t key with
     | Sys_error _ -> ()
 
-  let get_partial_values t kr_pairs =
-    Storage_intf.Base.get_partial_values
-      ~get_fn:get t kr_pairs
-
-  let set_partial_values t krv_triplet =
-    Storage_intf.Base.set_partial_values
-      ~set_fn:set ~get_fn:get t krv_triplet
-
-  let erase_values t keys =
-    Storage_intf.Base.erase_values
-      ~erase_fn:erase t keys
+  let size t key =
+    In_channel.with_open_gen
+      In_channel.[Open_rdonly]
+      t.file_perm
+      (key_to_fspath t key)
+      (fun ic -> In_channel.length ic |> Int64.to_int)
 
   let erase_prefix t pre =
     Storage_intf.Base.erase_prefix
       ~list_fn:list ~erase_fn:erase t pre
 
-  let list_prefix pre t =
+  let list_prefix t pre =
     Storage_intf.Base.list_prefix ~list_fn:list t pre
 
   let list_dir t pre =

lib/storage/memory.ml

@@ -20,23 +20,23 @@ module Impl = struct
 
   let erase = StrMap.remove
 
+  let size t key =
+    get t key |> Result.get_ok |> String.length
+
   let erase_prefix t pre =
     StrMap.filter_map_inplace
       (fun k v ->
         if String.starts_with ~prefix:pre k then None else Some v) t
 
-  let get_partial_values t kr_pairs =
+  let get_partial_values t key ranges =
     Storage_intf.Base.get_partial_values
-      ~get_fn:get t kr_pairs
+      ~get_fn:get t key ranges
 
-  let set_partial_values t krv_triplet =
+  let set_partial_values t key ?(append=false) rv =
     Storage_intf.Base.set_partial_values
-      ~set_fn:set ~get_fn:get t krv_triplet
-
-  let erase_values t keys =
-    Storage_intf.Base.erase_values ~erase_fn:erase t keys
+      ~set_fn:set ~get_fn:get t key append rv
 
-  let list_prefix pre t =
+  let list_prefix t pre =
     Storage_intf.Base.list_prefix ~list_fn:list t pre
 
   let list_dir t pre =