Adds protobuf serialization to FeatureMap, Hasher, and VariableInfo implementations

Core/src/main/java/org/tribuo/CategoricalIDInfo.java

@@ -16,12 +16,30 @@
 
 package org.tribuo;
 
+import com.google.protobuf.Any;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.oracle.labs.mlrg.olcut.util.MutableLong;
+import org.tribuo.protos.ProtoSerializableClass;
+import org.tribuo.protos.ProtoSerializableField;
+import org.tribuo.protos.core.CategoricalIDInfoProto;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Objects;
+
 /**
  * Same as a {@link CategoricalInfo}, but with an additional int id field.
  */
+@ProtoSerializableClass(version = CategoricalIDInfo.CURRENT_VERSION, serializedDataClass = CategoricalIDInfoProto.class)
 public class CategoricalIDInfo extends CategoricalInfo implements VariableIDInfo {
     private static final long serialVersionUID = 2L;
 
+    /**
+     * Protobuf serialization version.
+     */
+    public static final int CURRENT_VERSION = 0;
+
+    @ProtoSerializableField
     private final int id;
 
     /**
@@ -31,6 +49,9 @@ public class CategoricalIDInfo extends CategoricalInfo implements VariableIDInfo
      */
     public CategoricalIDInfo(CategoricalInfo info, int id) {
         super(info);
+        if (id < 0) {
+            throw new IllegalArgumentException("Invalid id number, must be non-negative, found " + id);
+        }
         this.id = id;
     }
 
@@ -46,6 +67,61 @@ private CategoricalIDInfo(CategoricalIDInfo info, String newName) {
         this.id = info.id;
     }
 
+    /**
+     * Deserialization constructor.
+     * @param name The info name.
+     * @param id The info id.
+     */
+    private CategoricalIDInfo(String name, int id) {
+        super(name);
+        if (id < 0) {
+            throw new IllegalArgumentException("Invalid id number, must be non-negative, found " + id);
+        }
+        this.id = id;
+    }
+
+    /**
+     * Deserialization factory.
+     * @param version The serialized object version.
+     * @param className The class name.
+     * @param message The serialized data.
+     */
+    public static CategoricalIDInfo deserializeFromProto(int version, String className, Any message) throws InvalidProtocolBufferException {
+        if (version < 0 || version > CURRENT_VERSION) {
+            throw new IllegalArgumentException("Unknown version " + version + ", this class supports at most version " + CURRENT_VERSION);
+        }
+        CategoricalIDInfoProto proto = message.unpack(CategoricalIDInfoProto.class);
+        CategoricalIDInfo info = new CategoricalIDInfo(proto.getName(),proto.getId());
+        List<Double> keys = proto.getKeyList();
+        List<Long> values = proto.getValueList();
+        if (keys.size() != values.size()) {
+            throw new IllegalStateException("Invalid protobuf, keys and values don't match. keys.size() = " + keys.size() + ", values.size() = " + values.size());
+        }
+        int newCount = 0;
+        if (keys.size() > 1) {
+            info.valueCounts = new HashMap<>(keys.size());
+            for (int i = 0; i < keys.size(); i++) {
+                if (values.get(i) < 0) {
+                    throw new IllegalStateException("Invalid protobuf, counts must be positive, found " + values.get(i) + " for value " + keys.get(i));
+                }
+                info.valueCounts.put(keys.get(i),new MutableLong(values.get(i)));
+                newCount += values.get(i).intValue();
+            }
+        } else {
+            info.observedValue = proto.getObservedValue();
+            info.observedCount = proto.getObservedCount();
+            newCount = (int) proto.getObservedCount();
+            if (info.observedCount < 0) {
+                throw new IllegalStateException("Invalid protobuf, counts must be positive, found " + info.observedCount + " for value " + info.observedValue);
+            }
+        }
+        if (newCount != proto.getCount()) {
+            throw new IllegalStateException("Invalid protobuf, count " + newCount + " did not match expected value " + proto.getCount());
+        }
+        info.count = newCount;
+        return info;
+    }
+
     @Override
     public int getID() {
         return id;
@@ -84,4 +160,25 @@ public String toString() {
             return "CategoricalFeature(name=" + name + ",id=" + id + ",count=" + count + ",map={" +observedValue+","+observedCount+"})";
         }
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        if (!super.equals(o)) {
+            return false;
+        }
+        CategoricalIDInfo that = (CategoricalIDInfo) o;
+        return id == that.id;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), id);
+    }
+
 }

Core/src/main/java/org/tribuo/CategoricalInfo.java

@@ -16,15 +16,24 @@
 
 package org.tribuo;
 
+import com.google.protobuf.Any;
+import com.google.protobuf.InvalidProtocolBufferException;
 import com.oracle.labs.mlrg.olcut.util.MutableLong;
 import com.oracle.labs.mlrg.olcut.util.MutableNumber;
+import org.tribuo.protos.ProtoSerializableClass;
+import org.tribuo.protos.ProtoSerializableField;
+import org.tribuo.protos.ProtoSerializableKeysValuesField;
+import org.tribuo.protos.ProtoUtil;
+import org.tribuo.protos.core.CategoricalInfoProto;
+import org.tribuo.protos.core.VariableInfoProto;
 import org.tribuo.util.Util;
 
 import java.io.IOException;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Random;
 import java.util.SplittableRandom;
 import java.util.stream.Collectors;
@@ -47,9 +56,15 @@
  * are recomputed. Care should be taken if data is read while {@link #observe(double)} is called.
  * </p>
  */
+@ProtoSerializableClass(version = CategoricalInfo.CURRENT_VERSION, serializedDataClass = CategoricalInfoProto.class)
 public class CategoricalInfo extends SkeletalVariableInfo {
     private static final long serialVersionUID = 2L;
 
+    /**
+     * Protobuf serialization version.
+     */
+    public static final int CURRENT_VERSION = 0;
+
     private static final MutableLong ZERO = new MutableLong(0);
     /**
      * The default threshold for converting a categorical info into a {@link RealInfo}.
@@ -60,16 +75,19 @@ public class CategoricalInfo extends SkeletalVariableInfo {
     /**
      * The occurrence counts of each value.
      */
+    @ProtoSerializableKeysValuesField(keysName="key", valuesName="value")
     protected Map<Double,MutableLong> valueCounts = null;
 
     /**
      * The observed value if it's only seen a single one.
      */
+    @ProtoSerializableField
     protected double observedValue = Double.NaN;
 
     /**
      * The count of the observed value if it's only seen a single one.
      */
+    @ProtoSerializableField
     protected long observedCount = 0;
 
     // These variables are used in the sampling methods, and regenerated after serialization if a sample is required.
@@ -117,6 +135,50 @@ protected CategoricalInfo(CategoricalInfo info, String newName) {
         }
     }
 
+    /**
+     * Deserialization factory.
+     * @param version The serialized object version.
+     * @param className The class name.
+     * @param message The serialized data.
+     */
+    public static CategoricalInfo deserializeFromProto(int version, String className, Any message) throws InvalidProtocolBufferException {
+        if (version < 0 || version > CURRENT_VERSION) {
+            throw new IllegalArgumentException("Unknown version " + version + ", this class supports at most version " + CURRENT_VERSION);
+        }
+        CategoricalInfoProto proto = message.unpack(CategoricalInfoProto.class);
+        CategoricalInfo info = new CategoricalInfo(proto.getName());
+        List<Double> keys = proto.getKeyList();
+        List<Long> values = proto.getValueList();
+        if (keys.size() != values.size()) {
+            throw new IllegalStateException("Invalid protobuf, keys and values don't match. keys.size() = " + keys.size() + ", values.size() = " + values.size());
+        }
+        int newCount = 0;
+        if (keys.size() > 1) {
+            info.valueCounts = new HashMap<>(keys.size());
+            for (int i = 0; i < keys.size(); i++) {
+                if (values.get(i) < 0) {
+                    throw new IllegalStateException("Invalid protobuf, counts must be positive, found " + values.get(i) + " for value " + keys.get(i));
+                }
+                info.valueCounts.put(keys.get(i),new MutableLong(values.get(i)));
+                newCount += values.get(i).intValue();
+            }
+        } else {
+            info.observedValue = proto.getObservedValue();
+            info.observedCount = proto.getObservedCount();
+            newCount = (int) proto.getObservedCount();
+            if (info.observedCount < 0) {
+                throw new IllegalStateException("Invalid protobuf, counts must be positive, found " + info.observedCount + " for value " + info.observedValue);
+            }
+        }
+        info.count = newCount;
+        return info;
+    }
+
+    @Override
+    public VariableInfoProto serialize() {
+        return ProtoUtil.serialize(this);
+    }
+
     @Override
     protected void observe(double value) {
         if (value != 0.0) {
@@ -361,6 +423,43 @@ private synchronized void regenerateValues() {
         }
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        if (!super.equals(o)) {
+            return false;
+        }
+        CategoricalInfo that = (CategoricalInfo) o;
+        // MutableLong in OLCUT 5.2.0 doesn't implement equals,
+        // so we can't compare valueCounts with Objects.equals.
+        // That'll be fixed in the next OLCUT but for the time being we've got this workaround.
+        if (valueCounts != null ^ that.valueCounts != null) {
+            return false;
+        } else if (valueCounts != null && that.valueCounts != null) {
+            if (valueCounts.size() != that.valueCounts.size()) {
+                return false;
+            } else {
+                for (Map.Entry<Double, MutableLong> e : valueCounts.entrySet()) {
+                    MutableLong other = that.valueCounts.get(e.getKey());
+                    if ((other == null) || (e.getValue().longValue() != other.longValue())) {
+                        return false;
+                    }
+                }
+            }
+        }
+        return Double.compare(that.observedValue, observedValue) == 0 && observedCount == that.observedCount;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), valueCounts, observedValue, observedCount);
+    }
+
     @Override
     public String toString() {
         if (valueCounts != null) {

Core/src/main/java/org/tribuo/FeatureMap.java

@@ -16,23 +16,30 @@
 
 package org.tribuo;
 
+import org.tribuo.protos.core.FeatureDomainProto;
+import org.tribuo.protos.ProtoSerializable;
+import org.tribuo.protos.ProtoSerializableMapValuesField;
+import org.tribuo.protos.ProtoUtil;
+
 import java.io.Serializable;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.TreeMap;
 
 /**
  * A map from Strings to {@link VariableInfo} objects storing
  * information about a feature.
  */
-public abstract class FeatureMap implements Serializable, Iterable<VariableInfo> {
+public abstract class FeatureMap implements Serializable, ProtoSerializable<FeatureDomainProto>, Iterable<VariableInfo> {
     private static final long serialVersionUID = 1L;
 
     /**
      * Map from the feature names to their info.
      */
+    @ProtoSerializableMapValuesField(valuesName = "info")
     protected final Map<String, VariableInfo> m;
 
     /**
@@ -101,6 +108,23 @@ public Iterator<VariableInfo> iterator() {
         return m.values().iterator();
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        FeatureMap that = (FeatureMap) o;
+        return m.equals(that.m);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(m);
+    }
+
     /**
      * Same as the toString, but ordered by name, and with newlines.
      * @return A String representation of this FeatureMap.
@@ -136,4 +160,13 @@ public boolean domainEquals(FeatureMap other) {
         }
     }
 
+    /**
+     * Deserializes a {@link FeatureDomainProto} into a {@link FeatureMap} subclass.
+     * @param proto The proto to deserialize.
+     * @return The deserialized FeatureMap.
+     */
+    public static FeatureMap deserialize(FeatureDomainProto proto) {
+        return ProtoUtil.deserialize(proto);
+    }
+
 }