diff --git a/cardinality.go b/cardinality.go index 1423724..2790c5c 100644 --- a/cardinality.go +++ b/cardinality.go @@ -1,6 +1,7 @@ package probably import ( + "encoding/json" "math" ) @@ -104,3 +105,36 @@ func (h *HyperLogLog) Merge(from *HyperLogLog) { } } } + +func (h *HyperLogLog) MarshalJSON() ([]byte, error) { + m := map[string]interface{}{ + "m": h.m, + "k": h.k, + "k_comp": h.k_comp, + "alpha_m": h.alpha_m, + "bits": h.bits, + } + return json.Marshal(m) +} + +func (h *HyperLogLog) UnmarshalJSON(by []byte) error { + + data := struct { + M uint `json:"m"` + K float64 `json:"k"` + KComp int `json:"k_comp"` + AlphaM float64 `json:"alpha_m"` + Bits []uint8 `json:"bits"` + }{} + + err := json.Unmarshal(by, &data) + if err != nil { + return err + } + h.m = data.M + h.k = data.K + h.k_comp = data.KComp + h.alpha_m = data.AlphaM + h.bits = data.Bits + return nil +} diff --git a/cardinality_test.go b/cardinality_test.go index d9f3548..880c69c 100644 --- a/cardinality_test.go +++ b/cardinality_test.go @@ -1,6 +1,7 @@ package probably import ( + "encoding/json" "hash/crc32" "testing" ) @@ -436,6 +437,37 @@ func TestCardinality(t *testing.T) { } } +func TestCardinalityJson(t *testing.T) { + hll := NewHyperLogLog(0.001) + for _, w := range words { + h := crc32.ChecksumIEEE([]byte(w)) + hll.Add(h) + } + t.Logf("Word list is %v words, estimate is %v", len(words), hll.Count()) + if hll.Count() != 2334 { + t.Fatalf("Expected estimate of 2,334, got %v", hll.Count()) + } + by, err := json.Marshal(hll) + if err != nil { + t.Fatalf("Must not error on json serialization: %v", err) + } + var llJson HyperLogLog + err = json.Unmarshal(by, &llJson) + if err != nil { + t.Fatalf("Must not error on json deserialization: %v", err) + } + if llJson.Count() != 2334 { + t.Fatalf("2 Expected estimate of 2,334, got %v", llJson.Count()) + } + for _, w := range words { + h := crc32.ChecksumIEEE([]byte(w)) + llJson.Add(h) + } + if llJson.Count() != 2334 { + t.Fatalf("3 Expected estimate of 2,334, got %v", llJson.Count()) + } +} + func TestCardinalityMerging(t *testing.T) { counters := make([]*HyperLogLog, 10) for i := 0; i < 10; i++ {