util/sketch/BitArray.java

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one or more
0003  * contributor license agreements.  See the NOTICE file distributed with
0004  * this work for additional information regarding copyright ownership.
0005  * The ASF licenses this file to You under the Apache License, Version 2.0
0006  * (the "License"); you may not use this file except in compliance with
0007  * the License.  You may obtain a copy of the License at
0008  *
0009  *    http://www.apache.org/licenses/LICENSE-2.0
0010  *
0011  * Unless required by applicable law or agreed to in writing, software
0012  * distributed under the License is distributed on an "AS IS" BASIS,
0013  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0014  * See the License for the specific language governing permissions and
0015  * limitations under the License.
0016  */
0017
0018 package org.apache.spark.util.sketch;
0019
0020 import java.io.DataInputStream;
0021 import java.io.DataOutputStream;
0022 import java.io.IOException;
0023 import java.util.Arrays;
0024
0025 final class BitArray {
0026   private final long[] data;
0027   private long bitCount;
0028
0029   static int numWords(long numBits) {
0030     if (numBits <= 0) {
0031       throw new IllegalArgumentException("numBits must be positive, but got " + numBits);
0032     }
0033     long numWords = (long) Math.ceil(numBits / 64.0);
0034     if (numWords > Integer.MAX_VALUE) {
0035       throw new IllegalArgumentException("Can't allocate enough space for " + numBits + " bits");
0036     }
0037     return (int) numWords;
0038   }
0039
0040   BitArray(long numBits) {
0041     this(new long[numWords(numBits)]);
0042   }
0043
0044   private BitArray(long[] data) {
0045     this.data = data;
0046     long bitCount = 0;
0047     for (long word : data) {
0048       bitCount += Long.bitCount(word);
0049     }
0050     this.bitCount = bitCount;
0051   }
0052
0053   /** Returns true if the bit changed value. */
0054   boolean set(long index) {
0055     if (!get(index)) {
0056       data[(int) (index >>> 6)] |= (1L << index);
0057       bitCount++;
0058       return true;
0059     }
0060     return false;
0061   }
0062
0063   boolean get(long index) {
0064     return (data[(int) (index >>> 6)] & (1L << index)) != 0;
0065   }
0066
0067   /** Number of bits */
0068   long bitSize() {
0069     return (long) data.length * Long.SIZE;
0070   }
0071
0072   /** Number of set bits (1s) */
0073   long cardinality() {
0074     return bitCount;
0075   }
0076
0077   /** Combines the two BitArrays using bitwise OR. */
0078   void putAll(BitArray array) {
0079     assert data.length == array.data.length : "BitArrays must be of equal length when merging";
0080     long bitCount = 0;
0081     for (int i = 0; i < data.length; i++) {
0082       data[i] |= array.data[i];
0083       bitCount += Long.bitCount(data[i]);
0084     }
0085     this.bitCount = bitCount;
0086   }
0087
0088   void writeTo(DataOutputStream out) throws IOException {
0089     out.writeInt(data.length);
0090     for (long datum : data) {
0091       out.writeLong(datum);
0092     }
0093   }
0094
0095   static BitArray readFrom(DataInputStream in) throws IOException {
0096     int numWords = in.readInt();
0097     long[] data = new long[numWords];
0098     for (int i = 0; i < numWords; i++) {
0099       data[i] = in.readLong();
0100     }
0101     return new BitArray(data);
0102   }
0103
0104   @Override
0105   public boolean equals(Object other) {
0106     if (this == other) return true;
0107     if (other == null || !(other instanceof BitArray)) return false;
0108     BitArray that = (BitArray) other;
0109     return Arrays.equals(data, that.data);
0110   }
0111
0112   @Override
0113   public int hashCode() {
0114     return Arrays.hashCode(data);
0115   }
0116 }