davisusanibar commented on code in PR #34227:
URL: https://github.com/apache/arrow/pull/34227#discussion_r1154663621
##########
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:
This loop iterate over the different tables defined on the Substrait plan,
then for any Table that appear on the plan tray to assign the Table Data
mapping by the user.
Yes, this could cause an error, if the user map the same Table Name to
different Table Data then only one of the will be used.
--
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]
