Java Code Examples for com.google.common.primitives.Longs#concat()

private void addDataBlocks() {

    try(RollbackCloseable rollbackable = new RollbackCloseable()) {

      // make sure can fit the next batch.
      listener.addBatch();

      {
        FixedBlockVector newFixed = new FixedBlockVector(allocator, pivot.getBlockWidth());
        rollbackable.add(newFixed);
        newFixed.ensureAvailableBlocks(MAX_VALUES_PER_BATCH);
        fixedBlocks = ObjectArrays.concat(fixedBlocks, newFixed);
        tableFixedAddresses = Longs.concat(tableFixedAddresses, new long[]{newFixed.getMemoryAddress()});
      }

      {
        VariableBlockVector newVariable = new VariableBlockVector(allocator, pivot.getVariableCount());
        rollbackable.add(newVariable);
        newVariable.ensureAvailableDataSpace(pivot.getVariableCount() == 0 ? 0 : MAX_VALUES_PER_BATCH * defaultVariableLengthSize);
        variableBlocks = ObjectArrays.concat(variableBlocks, newVariable);
        initVariableAddresses = Longs.concat(initVariableAddresses, new long[]{newVariable.getMemoryAddress()});
        openVariableAddresses = Longs.concat(openVariableAddresses, new long[]{newVariable.getMemoryAddress()});
        maxVariableAddresses = Longs.concat(maxVariableAddresses, new long[]{newVariable.getMaxMemoryAddress()});
      }
      rollbackable.commit();
    } catch (Exception e) {
      throw Throwables.propagate(e);
    }
  }
 

private void addDataBlocks(){

    // make sure can fit the next batch.
    listener.resized(currentOrdinal + MAX_VALUES_PER_BATCH);

    try(RollbackCloseable rollbackable = new RollbackCloseable()) {

      {
        FixedBlockVector newFixed = new FixedBlockVector(allocator, pivot.getBlockWidth());
        rollbackable.add(newFixed);
        newFixed.ensureAvailableBlocks(MAX_VALUES_PER_BATCH);
        fixedBlocks = ObjectArrays.concat(fixedBlocks, newFixed);
        tableFixedAddresses = Longs.concat(tableFixedAddresses, new long[]{newFixed.getMemoryAddress()});
      }

      {
        VariableBlockVector newVariable = new VariableBlockVector(allocator, pivot.getVariableCount());
        rollbackable.add(newVariable);
        newVariable.ensureAvailableDataSpace(pivot.getVariableCount() == 0 ? 0 : MAX_VALUES_PER_BATCH * defaultVariableLengthSize);
        variableBlocks = ObjectArrays.concat(variableBlocks, newVariable);
        initVariableAddresses = Longs.concat(initVariableAddresses, new long[]{newVariable.getMemoryAddress()});
        openVariableAddresses = Longs.concat(openVariableAddresses, new long[]{newVariable.getMemoryAddress()});
        maxVariableAddresses = Longs.concat(maxVariableAddresses, new long[]{newVariable.getMaxMemoryAddress()});
      }
      rollbackable.commit();
    } catch (Exception e) {
      throw Throwables.propagate(e);
    }
  }
 

@Override
public INDArray repeat(int dimension, long... repeats) {
    Nd4j.getCompressor().autoDecompress(this);


    if (dimension < 0)
        dimension += rank();

    if (repeats.length < rank()) {
        if (dimension > 0)
            repeats = Longs.concat(ArrayUtil.nTimes((long) rank() - repeats.length, 1), repeats);
            //append rather than prepend for dimension == 0
        else
            repeats = Longs.concat(repeats, ArrayUtil.nTimes((long) rank() - repeats.length, 1));

    }

    long[] newShape = new long[rank()];

    for (int i = 0; i < newShape.length; i++)
        newShape[i] = size(i) * repeats[i];

    INDArray ret = Nd4j.create(newShape);

    //number of times to repeat each value
    long repeatDelta = ArrayUtil.prod(newShape) / length();
    for (int i = 0; i < tensorssAlongDimension(dimension); i++) {
        INDArray thisTensor = tensorAlongDimension(i, dimension);
        INDArray retTensor = ret.tensorAlongDimension(i, dimension);
        int retIdx = 0;
        for (int k = 0; k < thisTensor.length(); k++) {
            for (int j = 0; j < repeatDelta; j++) {
                retTensor.putScalar(retIdx++, thisTensor.getDouble(k));
            }
        }
    }

    return ret;
}
 

/**
 * Add a new data block (batch) to the hashtable (and accumulator). Memory
 * allocation is needed for the following things:
 *
 * (1) Add new {@link FixedBlockVector} to array of fixed blocks.
 * (2) Add new {@link VariableBlockVector} to array of variable blocks.
 * (3) Add new {@link org.apache.arrow.vector.FieldVector} as a new target vector
 *     in {@link com.dremio.sabot.op.aggregate.vectorized.BaseSingleAccumulator}
 *     to store computed values. This is done for _each_ accumulator inside
 *     {@link com.dremio.sabot.op.aggregate.vectorized.AccumulatorSet}.
 *
 * All of the above operations have to be done in a single transaction
 * as one atomic unit of work. This allows us to handle OutOfMemory situations
 * without creating any inconsistent state of data structures.
 */
private void addDataBlocks(){
  final long currentAllocatedMemory = allocator.getAllocatedMemory();
  FixedBlockVector[] oldFixedBlocks = fixedBlocks;
  long[] oldTableFixedAddresses = tableFixedAddresses;
  try(RollbackCloseable rollbackable = new RollbackCloseable()) {
    FixedBlockVector newFixed;
    VariableBlockVector newVariable;
    {
      /* add new target accumulator vector to each accumulator.
       * since this is an array based allocation, we can fail in the middle
       * later on we revert on each accumulator and that might be a NO-OP
       * for some accumulators
       */
      listener.addBatch();
    }

    /* if the above operation was successful and we fail anywhere in the following
     * operations then we need to revert memory allocation on accumulator (all accumulators
     * in NestedAccumulator)
     */
    {
      newFixed = new FixedBlockVector(allocator, pivot.getBlockWidth());
      /* no need to rollback explicitly */
      rollbackable.add(newFixed);
      newFixed.ensureAvailableBlocks(MAX_VALUES_PER_BATCH);
      /* if we fail while allocating memory in above step, the state of fixed block array is still
       * consistent so we don't have to revert anything.
       */
      fixedBlocks = ObjectArrays.concat(fixedBlocks, newFixed);
      tableFixedAddresses = Longs.concat(tableFixedAddresses, new long[]{newFixed.getMemoryAddress()});
    }

    {
      newVariable = new VariableBlockVector(allocator, pivot.getVariableCount());
      /* no need to rollback explicitly */
      rollbackable.add(newVariable);
      newVariable.ensureAvailableDataSpace(variableBlockMaxLength);
      /* if we fail while allocating memory in above step, the state of variable block array is still consistent */
      variableBlocks = ObjectArrays.concat(variableBlocks, newVariable);
      initVariableAddresses = Longs.concat(initVariableAddresses, new long[]{newVariable.getMemoryAddress()});
      openVariableAddresses = Longs.concat(openVariableAddresses, new long[]{newVariable.getMemoryAddress()});
      maxVariableAddresses = Longs.concat(maxVariableAddresses, new long[]{newVariable.getMaxMemoryAddress()});
    }

    listener.commitResize();
    rollbackable.commit();
    /* bump these stats only after all new allocations have been successful as otherwise we would revert everything */
    allocatedForFixedBlocks += newFixed.getCapacity();
    allocatedForVarBlocks += newVariable.getCapacity();
  } catch (Exception e) {
    logger.debug("ERROR: failed to add data blocks, exception: ", e);
    /* explicitly rollback resizing operations on NestedAccumulator */
    listener.revertResize();
    fixedBlocks = oldFixedBlocks;
    tableFixedAddresses = oldTableFixedAddresses;
    /* do sanity checking on the state of all data structures after
     * memory allocation failed and we rollbacked. this helps in proactively detecting
     * potential IndexOutOfBoundsException and seg faults due to inconsistent state across
     * data structures.
     */
    Preconditions.checkArgument(fixedBlocks.length == variableBlocks.length,
                                "Error: detected inconsistent state in hashtable after memory allocation failed");
    listener.verifyBatchCount(fixedBlocks.length);
    /* at this point we are as good as no memory allocation was ever attempted */
    Preconditions.checkArgument(allocator.getAllocatedMemory() == currentAllocatedMemory,
                                "Error: detected inconsistent state of allocated memory");
    /* VectorizedHashAggOperator inserts data into hashtable and will handle (if OOM) this exception */
    throw Throwables.propagate(e);
  }
}
 

/**
 * Get the tensor matrix multiply shape
 * @param aShape the shape of the first array
 * @param bShape the shape of the second array
 * @param axes the axes to do the multiply
 * @return the shape for tensor matrix multiply
 */
public static long[] getTensorMmulShape(long[] aShape, long[] bShape, int[][] axes) {
    // FIXME: int cast


    int validationLength = Math.min(axes[0].length, axes[1].length);
    for (int i = 0; i < validationLength; i++) {
        if (aShape[axes[0][i]] != bShape[axes[1][i]])
            throw new IllegalArgumentException(
                            "Size of the given axes a" + " t each dimension must be the same size.");
        if (axes[0][i] < 0)
            axes[0][i] += aShape.length;
        if (axes[1][i] < 0)
            axes[1][i] += bShape.length;

    }

    List<Integer> listA = new ArrayList<>();
    for (int i = 0; i < aShape.length; i++) {
        if (!Ints.contains(axes[0], i))
            listA.add(i);
    }



    List<Integer> listB = new ArrayList<>();
    for (int i = 0; i < bShape.length; i++) {
        if (!Ints.contains(axes[1], i))
            listB.add(i);
    }


    int n2 = 1;
    int aLength = Math.min(aShape.length, axes[0].length);
    for (int i = 0; i < aLength; i++) {
        n2 *= aShape[axes[0][i]];
    }

    //if listA and listB are empty these donot initialize.
    //so initializing with {1} which will then get overriden if not empty
    long[] oldShapeA;
    if (listA.size() == 0) {
        oldShapeA = new long[] {1};
    } else {
        oldShapeA = Longs.toArray(listA);
        for (int i = 0; i < oldShapeA.length; i++)
            oldShapeA[i] = aShape[(int) oldShapeA[i]];
    }

    int n3 = 1;
    int bNax = Math.min(bShape.length, axes[1].length);
    for (int i = 0; i < bNax; i++) {
        n3 *= bShape[axes[1][i]];
    }


    long[] oldShapeB;
    if (listB.size() == 0) {
        oldShapeB = new long[] {1};
    } else {
        oldShapeB = Longs.toArray(listB);
        for (int i = 0; i < oldShapeB.length; i++)
            oldShapeB[i] = bShape[(int) oldShapeB[i]];
    }


    long[] aPlusB = Longs.concat(oldShapeA, oldShapeB);
    return aPlusB;
}
 

private SDVariable doTensorMmul(SDVariable a,
                                SDVariable b,
                                int[][] axes) {

    int validationLength = Math.min(axes[0].length, axes[1].length);
    for (int i = 0; i < validationLength; i++) {
        if (a.getShape()[axes[0][i]] != b.getShape()[axes[1][i]])
            throw new IllegalArgumentException("Size of the given axes at each dimension must be the same size.");
        if (axes[0][i] < 0)
            axes[0][i] += a.getShape().length;
        if (axes[1][i] < 0)
            axes[1][i] += b.getShape().length;

    }

    List<Integer> listA = new ArrayList<>();
    for (int i = 0; i < a.getShape().length; i++) {
        if (!Ints.contains(axes[0], i))
            listA.add(i);
    }

    int[] newAxesA = Ints.concat(Ints.toArray(listA), axes[0]);


    List<Integer> listB = new ArrayList<>();
    for (int i = 0; i < b.getShape().length; i++) {
        if (!Ints.contains(axes[1], i))
            listB.add(i);
    }

    int[] newAxesB = Ints.concat(axes[1], Ints.toArray(listB));

    int n2 = 1;
    int aLength = Math.min(a.getShape().length, axes[0].length);
    for (int i = 0; i < aLength; i++) {
        n2 *= a.getShape()[axes[0][i]];
    }

    //if listA and listB are empty these do not initialize.
    //so initializing with {1} which will then get overridden if not empty
    long[] newShapeA = {-1, n2};
    long[] oldShapeA;
    if (listA.size() == 0) {
        oldShapeA = new long[] {1};
    } else {
        oldShapeA = Longs.toArray(listA);
        for (int i = 0; i < oldShapeA.length; i++)
            oldShapeA[i] = a.getShape()[(int) oldShapeA[i]];
    }

    int n3 = 1;
    int bNax = Math.min(b.getShape().length, axes[1].length);
    for (int i = 0; i < bNax; i++) {
        n3 *= b.getShape()[axes[1][i]];
    }


    int[] newShapeB = {n3, -1};
    long[] oldShapeB;
    if (listB.size() == 0) {
        oldShapeB = new long[] {1};
    } else {
        oldShapeB = Longs.toArray(listB);
        for (int i = 0; i < oldShapeB.length; i++)
            oldShapeB[i] = b.getShape()[(int) oldShapeB[i]];
    }


    SDVariable at = f()
            .reshape(f().permute
                    (a,newAxesA),newShapeA);
    SDVariable bt = f()
            .reshape(f()
                    .permute(b,newAxesB),newShapeB);

    SDVariable ret = f().mmul(at,bt);
    long[] aPlusB = Longs.concat(oldShapeA, oldShapeB);
    return f().reshape(ret, aPlusB);
}