NDArray toNDArray()

in java/com_spotify_voyager_jni_Index.cpp [235:277]


NDArray<float, 2> toNDArray(JNIEnv *env, jobjectArray floatArrays) {
  jsize numElements = env->GetArrayLength(floatArrays);
  if (numElements == 0) {
    return NDArray<float, 2>({0, 0});
  }

  jobject firstElement = env->GetObjectArrayElement(floatArrays, 0);
  jsize numDimensions = env->GetArrayLength((jfloatArray)firstElement);
  env->DeleteLocalRef(firstElement);
  if (numDimensions == 0) {
    return NDArray<float, 2>({0, 0});
  }

  NDArray<float, 2> output = NDArray<float, 2>({numElements, numDimensions});

  float *outputPointer = output.data.data();
  for (int i = 0; i < numElements; i++) {
    jobject element = env->GetObjectArrayElement(floatArrays, i);

    // TODO: Ensure that each element is actually a float array!
    jfloatArray floatArray = (jfloatArray)element;

    jsize numVectorDimensions = env->GetArrayLength(floatArray);

    if (numVectorDimensions != numDimensions) {
      throw std::runtime_error("When passing an array of arrays, all "
                               "sub-arrays must be the same length.");
    }

    env->GetFloatArrayRegion(floatArray, 0, numDimensions, outputPointer);

    // Delete the local reference to the nested float array.
    // Note that this isn't necessary; it merely helps the Java GC
    // identify if/when these elements can be cleaned up earlier than when
    // this function returns.
    // Removing this call would not create a memory leak.
    env->DeleteLocalRef(element);

    outputPointer += numDimensions;
  }

  return output;
}