Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add sample JNI API #9728

Merged
merged 7 commits into from
Dec 2, 2021
Merged

Add sample JNI API #9728

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
3b37f06
Add sample JNI
Nov 17, 2021
c5ec858
Add comment
Nov 19, 2021
37e9b57
Update comment
Nov 19, 2021
6fe77c2
Add a test case for sample
Nov 19, 2021
6205baa
Update comments and test cases
Nov 22, 2021
3dbb139
Update comments
Nov 23, 2021
6282f4b
Format
Dec 2, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions java/src/main/java/ai/rapids/cudf/Table.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -656,6 +656,8 @@ private static native ContiguousTable[] contiguousSplitGroups(long inputTable,
boolean[] keysDescending,
boolean[] keysNullSmallest);

private static native long[] sample(long tableHandle, long n, boolean replacement, long seed);

/////////////////////////////////////////////////////////////////////////////
// TABLE CREATION APIs
/////////////////////////////////////////////////////////////////////////////
Expand Down Expand Up @@ -2743,6 +2745,19 @@ public static Table fromPackedTable(ByteBuffer metadata, DeviceMemoryBuffer data
return result;
}

/**
* Gather `n` samples from table randomly
* The output is not same with CPU Sample exec, but this is faster.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
* The output is not same with CPU Sample exec, but this is faster.
* The output is not the same as what CPU Sample Exec produces, but this is faster.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks

*
* @param n non-negative number of samples expected
* @param replacement Allow or disallow sampling of the same row more than once.
* @param seed Seed value to initiate random number generator.
* @return
*/
public Table sample(long n, boolean replacement, long seed) {
firestarman marked this conversation as resolved.
Show resolved Hide resolved
return new Table(sample(nativeHandle, n, replacement, seed));
}

/////////////////////////////////////////////////////////////////////////////
// HELPER CLASSES
/////////////////////////////////////////////////////////////////////////////
Expand Down
14 changes: 14 additions & 0 deletions java/src/main/native/src/TableJni.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
#include <cudf/sorting.hpp>
#include <cudf/stream_compaction.hpp>
#include <cudf/types.hpp>
#include <cudf/copying.hpp>
#include <rmm/cuda_stream_view.hpp>

#include "cudf_jni_apis.hpp"
Expand Down Expand Up @@ -3145,4 +3146,17 @@ JNIEXPORT jobjectArray JNICALL Java_ai_rapids_cudf_Table_contiguousSplitGroups(
CATCH_STD(env, NULL);
}

JNIEXPORT jlongArray JNICALL Java_ai_rapids_cudf_Table_sample(JNIEnv *env, jclass, jlong j_input,
jlong n, jboolean replacement, jlong seed) {
JNI_NULL_CHECK(env, j_input, "input table is null", 0);
try {
cudf::jni::auto_set_device(env);
cudf::table_view *input = reinterpret_cast<cudf::table_view *>(j_input);
auto sample_with_replacement =
replacement ? cudf::sample_with_replacement::TRUE : cudf::sample_with_replacement::FALSE;
std::unique_ptr<cudf::table> result = cudf::sample(*input, n, sample_with_replacement, seed);
return cudf::jni::convert_table_for_return(env, result);
}
CATCH_STD(env, 0);
}
} // extern "C"