Merge pull request #886 from haileyajohnson/zarr-filters

[5.x] add filters for zarr

Author: haileyajohnson

cdm/core/src/main/java/ucar/nc2/filter/Checksum32.java

@@ -33,8 +33,10 @@ public Checksum32(CType type, ByteOrder bo) {
     this.byteOrder = bo;
   }
 
+
   public Checksum32(CType type) {
-    this(type, ByteOrder.BIG_ENDIAN);
+    // TODO: can we do this better?
+    this(type, ByteOrder.LITTLE_ENDIAN);
   }
 
   @Override

cdm/core/src/main/java/ucar/nc2/iosp/hdf5/DataBTree.java

@@ -297,7 +297,7 @@ DataChunk next() throws IOException {
   public class DataChunk {
     public final int size; // size of chunk in bytes; need storage layout dimensions to interpret
     public final int filterMask; // bitfield indicating which filters have been skipped for this chunk
-    public final int[] offset; // offset index of this chunk, reletive to entire array
+    public final int[] offset; // offset index of this chunk, relative to entire array
     public final long filePos; // filePos of a single raw data chunk, already shifted by the offset if needed
 
     DataChunk(int ndim, boolean last) throws IOException {

cdm/zarr/src/main/java/ucar/nc2/iosp/zarr/ZArray.java

@@ -9,6 +9,9 @@
 import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import ucar.ma2.DataType;
+import ucar.nc2.filter.Filter;
+import ucar.nc2.filter.Filters;
+import ucar.nc2.filter.UnknownFilterException;
 
 import java.io.IOException;
 import java.nio.ByteOrder;
@@ -63,14 +66,14 @@ public enum Order {
   private final Object fillValue;
   private final DataType datatype;
   private final String dtype;
-  private final ZarrFilter compressor;
+  private final Filter compressor;
   private final ByteOrder byteOrder;
   private final Order order;
-  private final List<ZarrFilter> filters;
+  private final List<Filter> filters;
   private final String separator;
 
-  public ZArray(int[] shape, int[] chunks, Object fill_value, String dtype, ZarrFilter compressor, String order,
-      List<ZarrFilter> filters, String separator) throws ZarrFormatException {
+  public ZArray(int[] shape, int[] chunks, Object fill_value, String dtype, Filter compressor, String order,
+      List<Filter> filters, String separator) throws ZarrFormatException {
     this.shape = shape;
     this.chunks = chunks;
     this.fillValue = fill_value;
@@ -91,11 +94,11 @@ public int[] getChunks() {
     return this.chunks;
   }
 
-  public ZarrFilter getCompressor() {
+  public Filter getCompressor() {
     return this.compressor;
   }
 
-  public List<ZarrFilter> getFilters() {
+  public List<Filter> getFilters() {
     return this.filters;
   }
 
@@ -192,21 +195,24 @@ public ZArray deserialize(JsonParser p, DeserializationContext ctxt) throws IOEx
       TreeNode dim_sep = root.path(ZarrKeys.DIMENSION_SEPARATOR);
       String delimiter = dim_sep.isMissingNode() ? DEFAULT_SEPARATOR : ((JsonNode) dim_sep).asText();
 
-      Map<String, Object> compBean = codec.readValue(root.path(ZarrKeys.COMPRESSOR).traverse(codec), HashMap.class);
+      // Filters and compressor
+      try {
+        Map<String, Object> compBean = codec.readValue(root.path(ZarrKeys.COMPRESSOR).traverse(codec), HashMap.class);
 
-      ZarrFilter compressor;
-      if (compBean == null) {
-        compressor = null;
-      } else {
-        compressor = null; // TODO: implement compressors
-      }
+        Filter compressor = Filters.getFilterByName(compBean);
 
-      // TODO: filters
-      List<ZarrFilter> filters = null;
+        List<Filter> filters = new ArrayList<>();
 
-      try {
+        Map<String, Object>[] filtersBean =
+            codec.readValue(root.path(ZarrKeys.FILTERS).traverse(codec), HashMap[].class);
+
+        if (filtersBean != null) {
+          for (Map<String, Object> bean : filtersBean) {
+            filters.add(Filters.getFilterByName(bean));
+          }
+        }
         return new ZArray(shape, chunks, fill, dtype, compressor, order, filters, delimiter);
-      } catch (ZarrFormatException ex) {
+      } catch (UnknownFilterException | ZarrFormatException ex) {
         throw new IOException(ex.getMessage(), ex.getCause());
       }
     }

cdm/zarr/src/main/java/ucar/nc2/iosp/zarr/ZarrFilter.java

@@ -10,6 +10,7 @@
 import ucar.nc2.Dimension;
 import ucar.nc2.Group;
 import ucar.nc2.Variable;
+import ucar.nc2.filter.Filter;
 import ucar.unidata.io.RandomAccessFile;
 import ucar.unidata.io.zarr.RandomAccessDirectory;
 import ucar.unidata.io.zarr.RandomAccessDirectoryItem;
@@ -42,7 +43,7 @@ public ZarrHeader(RandomAccessDirectory raf, Group.Builder rootGroup) {
   private class DelayedVarMaker {
     private RandomAccessDirectoryItem var;
     private ZArray zarray;
-    private Set<Integer> initializedChunks; // track any uninitialized chunks for var
+    private Map<Integer, Long> initializedChunks; // track any uninitialized chunks for var
     private List<Attribute> attrs; // list of variable attributes
     private long dataOffset; // byte position where data starts
 
@@ -53,7 +54,7 @@ void setAttrs(List<Attribute> attrs) {
     void setVar(RandomAccessDirectoryItem var) {
       this.var = var;
       this.attrs = null;
-      this.initializedChunks = new HashSet<>();
+      this.initializedChunks = new HashMap<>();
       this.dataOffset = -1;
       if (var != null) {
         try {
@@ -90,7 +91,7 @@ void processItem(RandomAccessDirectoryItem item) {
         ZarrIosp.logger.error(new ZarrFormatException().getMessage());
         this.var = null; // skip rest of var is unrecognized files found
       }
-      this.initializedChunks.add(index);
+      this.initializedChunks.put(index, item.length());
       // if data offset is uninitialized, set here
       if (this.dataOffset < 0) {
         this.dataOffset = item.startIndex();
@@ -179,7 +180,7 @@ private void makeGroup(RandomAccessDirectoryItem item, List<Attribute> attrs) {
   }
 
   private void makeVariable(RandomAccessDirectoryItem item, long dataOffset, ZArray zarray,
-      Set<Integer> initializedChunks, List<Attribute> attrs) throws ZarrFormatException {
+      Map<Integer, Long> initializedChunks, List<Attribute> attrs) throws ZarrFormatException {
     // make new Variable
     Variable.Builder var = Variable.builder();
     String location = ZarrUtils.trimLocation(item.getLocation());
@@ -300,16 +301,16 @@ private Group.Builder findGroup(String location) throws ZarrFormatException {
   class VInfo {
     private final int[] chunks;
     private final Object fillValue;
-    private final ZarrFilter compressor;
+    private final Filter compressor;
     private final ByteOrder byteOrder;
     private final ZArray.Order order;
     private final String separator;
-    private final List<ZarrFilter> filters;
+    private final List<Filter> filters;
     private final long offset;
-    private final Set<Integer> initializedChunks;
+    private final Map<Integer, Long> initializedChunks;
 
-    VInfo(int[] chunks, Object fillValue, ZarrFilter compressor, ByteOrder byteOrder, ZArray.Order order,
-        String separator, List<ZarrFilter> filters, long offset, Set<Integer> initializedChunks) {
+    VInfo(int[] chunks, Object fillValue, Filter compressor, ByteOrder byteOrder, ZArray.Order order, String separator,
+        List<Filter> filters, long offset, Map<Integer, Long> initializedChunks) {
       this.chunks = chunks;
       this.fillValue = fillValue;
       this.byteOrder = byteOrder;
@@ -329,7 +330,7 @@ public Object getFillValue() {
       return this.fillValue;
     }
 
-    public ZarrFilter getCompressor() {
+    public Filter getCompressor() {
       return this.compressor;
     }
 
@@ -345,15 +346,15 @@ public String getSeparator() {
       return this.separator;
     }
 
-    public List<ZarrFilter> getFilters() {
+    public List<Filter> getFilters() {
       return this.filters;
     }
 
     public long getOffset() {
       return this.offset;
     }
 
-    public Set<Integer> getInitializedChunks() {
+    public Map<Integer, Long> getInitializedChunks() {
       return this.initializedChunks;
     }
 

cdm/zarr/src/main/java/ucar/nc2/iosp/zarr/ZarrHeader.java

@@ -4,6 +4,7 @@
 import ucar.ma2.Section;
 import ucar.nc2.Dimension;
 import ucar.nc2.Variable;
+import ucar.nc2.filter.Filter;
 import ucar.nc2.iosp.LayoutBB;
 import ucar.nc2.iosp.LayoutBBTiled;
 import ucar.unidata.io.RandomAccessFile;
@@ -12,10 +13,11 @@
 import java.nio.*;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
- * A tiled layout for Zarr formats that accommodates uncompressing anf filtering data before returning
+ * A tiled layout for Zarr formats that accommodates uncompressing and filtering data before returning
  */
 public class ZarrLayoutBB implements LayoutBB {
 
@@ -27,39 +29,35 @@ public class ZarrLayoutBB implements LayoutBB {
   private final Section want;
 
   private int[] chunkSize; // number of elements per chunks
-  private int elemSize; // size of eelements in bytes
+  private int elemSize; // size of elements in bytes
   private int nChunks[]; // number of chunks per dimension
-  private int nBytes; // number of bytes per chunk
   private int totalNChunks; // total number of chunks
-  private int totalChunkSize; // total number of elements per chunk
   private boolean F_order = false; // F order storage?
-  private Set<Integer> initializedChunks; // set of chunks that exist as files
-
-  // offset to start of data
-  private static final int ZARR_COMPRESSOR_OFFSET = 16;
-  private int data_bytes_offset;
+  private Map<Integer, Long> initializedChunks; // set of chunks that exist as files and their compressed size
+  private Filter compressor;
+  private List<Filter> filters;
 
   public ZarrLayoutBB(Variable v2, Section wantSection, RandomAccessFile raf) {
     // var data info
     this.raf = raf;
     ZarrHeader.VInfo vinfo = (ZarrHeader.VInfo) v2.getSPobject();
     this.byteOrder = vinfo.getByteOrder();
     this.varOffset = vinfo.getOffset();
-    this.data_bytes_offset = vinfo.getCompressor() == null ? 0 : ZARR_COMPRESSOR_OFFSET;
+    this.compressor = vinfo.getCompressor();
+    this.filters = vinfo.getFilters();
+
 
     // fill in chunk info
     this.chunkSize = vinfo.getChunks();
     int ndims = this.chunkSize.length;
     this.initializedChunks = vinfo.getInitializedChunks();
     this.nChunks = new int[ndims];
     this.totalNChunks = 1;
-    this.totalChunkSize = 1;
     for (int i = 0; i < ndims; i++) {
       Dimension dim = v2.getDimension(i);
       // round up nchunks if not evenly divisible by chunk size
       this.nChunks[i] = (int) Math.ceil(dim.getLength() / this.chunkSize[i]);
       this.totalNChunks *= nChunks[i];
-      this.totalChunkSize *= chunkSize[i];
     }
 
     // transpose wantsSection and chunk shape if F order
@@ -79,7 +77,6 @@ public ZarrLayoutBB(Variable v2, Section wantSection, RandomAccessFile raf) {
     }
 
     this.elemSize = v2.getDataType().getSize();
-    this.nBytes = totalChunkSize * elemSize;
 
     // create delegate and chunk iterator
     ZarrLayoutBB.DataChunkIterator iter = new ZarrLayoutBB.DataChunkIterator();
@@ -110,18 +107,21 @@ private class DataChunkIterator implements LayoutBBTiled.DataChunkIterator {
 
     private int[] currChunk; // current chunk in subscript coords
     private int chunkNum; // current chunk as flat index
+    private long currOffset; // byte position of current chunk
+
 
     DataChunkIterator() {
       this.currChunk = new int[chunkSize.length];
       this.chunkNum = 0;
+      this.currOffset = varOffset; // start at start of variable data
     }
 
     public boolean hasNext() {
       return this.chunkNum < totalNChunks;
     }
 
     public LayoutBBTiled.DataChunk next() {
-      DataChunk chunk = new ZarrLayoutBB.DataChunk(this.currChunk, this.chunkNum);
+      DataChunk chunk = new ZarrLayoutBB.DataChunk(this.currChunk, this.chunkNum, this.currOffset);
       incrementChunk();
       return chunk;
     }
@@ -134,6 +134,7 @@ private void incrementChunk() {
         i--;
       }
       this.currChunk[i]++;
+      this.currOffset += initializedChunks.getOrDefault(this.chunkNum, (long) 0);
       this.chunkNum = ZarrUtils.subscriptsToIndex(this.currChunk, nChunks);
     }
   }
@@ -144,9 +145,9 @@ private class DataChunk implements LayoutBBTiled.DataChunk {
     private long rafOffset; // start position of chunk in bytes
     private int chunkNum;
 
-    DataChunk(int[] index, int chunkNum) {
+    DataChunk(int[] index, int chunkNum, long rafOffset) {
+      this.rafOffset = rafOffset;
       this.offset = new int[index.length];
-      this.rafOffset = varOffset + (chunkNum * (nBytes + data_bytes_offset)) + data_bytes_offset;
       for (int i = 0; i < index.length; i++) {
         int j = F_order ? index.length - i - 1 : i;
         this.offset[i] = index[j] * chunkSize[i];
@@ -162,14 +163,23 @@ public ByteBuffer getByteBuffer() throws IOException {
       // read the data
       byte[] data;
       // if chunk does not exist as file, return empty buffer
-      if (!initializedChunks.contains(chunkNum)) {
-        data = new byte[0];
-      } else {
-        data = new byte[nBytes];
-        raf.seek(this.rafOffset);
-        raf.readFully(data);
-
-        // TODO: apply filters in reverse order
+      long dataLength = initializedChunks.getOrDefault(chunkNum, (long) 0);
+      if (dataLength == 0) {
+        ByteBuffer result = ByteBuffer.wrap(new byte[0]);
+        result.order(byteOrder);
+        return result;
+      }
+
+      data = new byte[(int) dataLength];
+      raf.seek(this.rafOffset);
+      // raf.read(data, 0, (int)dataLength);
+      raf.readFully(data);
+
+      // apply compressor
+      data = compressor.decode(data);
+      // apply filters in reverse order
+      for (int i = filters.size() - 1; i >= 0; i--) {
+        data = filters.get(i).decode(data);
       }
 
       ByteBuffer result = ByteBuffer.wrap(data);