lidavidm commented on code in PR #34227:
URL: https://github.com/apache/arrow/pull/34227#discussion_r1154685338
##########
java/dataset/src/main/cpp/jni_wrapper.cc:
##########
@@ -578,3 +601,102 @@
Java_org_apache_arrow_dataset_file_JniWrapper_writeFromScannerToFile(
JniAssertOkOrThrow(arrow::dataset::FileSystemDataset::Write(options,
scanner));
JNI_METHOD_END()
}
+
+/*
+ * Class: org_apache_arrow_dataset_substrait_JniWrapper
+ * Method: executeSerializedPlanLocalFiles
+ * Signature: (Ljava/lang/String;J)V
+ */
+JNIEXPORT void JNICALL
Java_org_apache_arrow_dataset_substrait_JniWrapper_executeSerializedPlanLocalFiles__Ljava_lang_String_2J
(
+ JNIEnv* env, jobject, jstring plan, jlong
c_arrow_array_stream_address_out) {
+ JNI_METHOD_START
+ auto* arrow_stream =
reinterpret_cast<ArrowArrayStream*>(c_arrow_array_stream_address_out);
+ std::shared_ptr<arrow::Buffer> buffer =
JniGetOrThrow(arrow::engine::SerializeJsonPlan(JStringToCString(env, plan)));
+ std::shared_ptr<arrow::RecordBatchReader> reader =
JniGetOrThrow(arrow::engine::ExecuteSerializedPlan(*buffer));
+ JniAssertOkOrThrow(arrow::ExportRecordBatchReader(reader, arrow_stream));
+ JNI_METHOD_END()
+}
+
+/*
+ * Class: org_apache_arrow_dataset_substrait_JniWrapper
+ * Method: executeSerializedPlanLocalFiles
+ * Signature: (Ljava/nio/ByteBuffer;J)V
+ */
+JNIEXPORT void JNICALL
Java_org_apache_arrow_dataset_substrait_JniWrapper_executeSerializedPlanLocalFiles__Ljava_nio_ByteBuffer_2J
(
+ JNIEnv* env, jobject, jobject plan, jlong
c_arrow_array_stream_address_out) {
+ JNI_METHOD_START
+ auto* arrow_stream =
reinterpret_cast<ArrowArrayStream*>(c_arrow_array_stream_address_out);
+ // mapping arrow::Buffer
+ auto *buff = reinterpret_cast<jbyte*>(env->GetDirectBufferAddress(plan));
+ int length = env->GetDirectBufferCapacity(plan);
+ std::shared_ptr<arrow::Buffer> buffer =
JniGetOrThrow(arrow::AllocateBuffer(length));
+ std::memcpy(buffer->mutable_data(), buff, length);
+ // execute plan
+ std::shared_ptr<arrow::RecordBatchReader> reader =
JniGetOrThrow(arrow::engine::ExecuteSerializedPlan(*buffer));
+ JniAssertOkOrThrow(arrow::ExportRecordBatchReader(reader, arrow_stream));
+ JNI_METHOD_END()
+}
+
+/*
+ * Class: org_apache_arrow_dataset_substrait_JniWrapper
+ * Method: executeSerializedPlanNamedTables
+ * Signature: (Ljava/lang/String;[Ljava/lang/String;J)V
+ */
+JNIEXPORT void JNICALL
Java_org_apache_arrow_dataset_substrait_JniWrapper_executeSerializedPlanNamedTables__Ljava_lang_String_2_3Ljava_lang_String_2J
(
+ JNIEnv* env, jobject, jstring plan, jobjectArray
table_to_memory_address_input, jlong memory_address_output) {
+ JNI_METHOD_START
+ // get mapping of table name to memory address
+ std::unordered_map<std::string, std::shared_ptr<arrow::Table>>
map_table_to_reader = ToMapTableToArrowReader(env,
table_to_memory_address_input);
+ // create table provider
+ arrow::engine::NamedTableProvider table_provider =
[&map_table_to_reader](const std::vector<std::string>& names, const
arrow::Schema&) {
+ std::shared_ptr<arrow::Table> output_table;
+ for (const auto& name : names) {
+ output_table = map_table_to_reader[name];
Review Comment:
Also, what I mean is that if we have more than one entry in `names`, we're
just overwriting the same `output_table` multiple times. That seems wrong.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
