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)
}