private def setFeatures()

in sagemaker-spark-sdk/src/main/scala/com/amazonaws/services/sagemaker/sparksdk/protobuf/ProtobufConverter.scala [238:294]


  private def setFeatures(protobufBuilder: Record.Builder,
                          matrix: Matrix): Record.Builder = {
    val featuresTensorBuilder = Value.newBuilder().getFloat32TensorBuilder()

    // Matrix shape must be recorded for both dense/sparse matrices
    featuresTensorBuilder.addShape(matrix.numRows)
    featuresTensorBuilder.addShape(matrix.numCols)

    val featuresTensor = matrix match {
      case m: DenseMatrix =>
        if (!m.isTransposed) {
          // Convert to Row major order
          for (row <- (0 to m.numRows - 1)) {
            for (col <- (0 to m.numCols - 1)) {
              featuresTensorBuilder.addValues(m(row, col).toFloat)
            }
          }
        } else {
          // When transposed it is already in Row major order
          for (value <- m.values) {
            featuresTensorBuilder.addValues(value.toFloat)
          }
        }

        featuresTensorBuilder.build()

      case m: SparseMatrix =>
        // Construct the index for each value so that
        // row = floor(index / cols)
        // col = index % cols
        var rowIdx = 0
        var colIdx = 0
        for (colStart <- m.colPtrs.slice(1, m.colPtrs.size)) {
          while (rowIdx < colStart) {
            if (m.isTransposed) {
              // When transposed, rowIndices behave are colIndices, and colPtrs and rowPtrs
              // and rowIdx, colIdx are swapped
              featuresTensorBuilder.addKeys((colIdx * m.numCols) + m.rowIndices(rowIdx))
            } else {
              featuresTensorBuilder.addKeys((m.rowIndices(rowIdx) * m.numCols) + colIdx)
            }
            rowIdx += 1
          }
          colIdx += 1
        }

        for (value <- m.values) {
          featuresTensorBuilder.addValues(value.toFloat)
        }

        featuresTensorBuilder.build()
    }
    val featuresValue = Value.newBuilder().setFloat32Tensor(featuresTensor).build
    val mapEntry = MapEntry.newBuilder().setKey(ValuesIdentifierString)
      .setValue(featuresValue).build
    protobufBuilder.addFeatures(mapEntry)
  }