Add MNIST inference unit test and enhance ONNX model handling

CMakeLists.txt

@@ -2,6 +2,7 @@
 
 # Set extension name here
 set(TARGET_NAME onnx)
+add_subdirectory(unit_test)
 
 # DuckDB's extension distribution supports vcpkg. As such, dependencies can be added in ./vcpkg.json and then
 # used in cmake with find_package. Feel free to remove or replace with other dependencies.

src/onnx_extension.cpp

@@ -1,5 +1,6 @@
 #define DUCKDB_EXTENSION_MAIN
 
+#include <filesystem>
 #include "onnx_extension.hpp"
 #include "duckdb.hpp"
 #include "duckdb/common/exception.hpp"
@@ -21,6 +22,11 @@ struct Tensor {
 void run_onnx_model_and_extract_results(const string &path,
                                         vector<Tensor> input_tensors,
                                         Value &struct_val_list) {
+  // check path exist
+  if (!std::__fs::filesystem::exists(path)) {
+    throw std::runtime_error(
+        std::string("ONNX model file not found: ") + path);
+  };
   Ort::Env ort_env;
   Ort::Session session{ort_env, (path.c_str()), Ort::SessionOptions{nullptr}};
   auto memory_info =
@@ -50,11 +56,23 @@ void run_onnx_model_and_extract_results(const string &path,
       memory_info, output_data.data(), output_size, output_shape.data(),
       output_shape.size());
 
-  const char *input_names[] = {"X"};
-  const char *output_names[] = {"Y"};
-
-  session.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1,
-              output_names, &output_tensor, 1);
+  Ort::AllocatorWithDefaultOptions allocator;
+  const std::string input_name =
+      session.GetInputNameAllocated(0, allocator).get();
+  const std::string output_name =
+      session.GetOutputNameAllocated(0, allocator).get();
+
+  const char *input_names[] = {input_name.c_str()};
+  const char *output_names[] = {output_name.c_str()};
+
+  try {
+    session.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1,
+                output_names, &output_tensor, 1);
+  } catch (const Ort::Exception &e) {
+    std::string persistent_message = e.what();
+    throw std::runtime_error(std::string("ONNXRuntime error: ") +
+                             persistent_message);
+  }
 
   std::vector<Value> shape_vl_list;
   shape_vl_list.reserve(output_shape.size());
@@ -138,9 +156,14 @@ inline void OnnxScalarFun(DataChunk &args, ExpressionState &state,
     Value struct_val_list;
     try {
       run_onnx_model_and_extract_results(path, onnx_inputs, struct_val_list);
+    } catch (Ort::Exception &e) {
+      throw std::runtime_error(std::string("ONNXRuntime error: ") + e.what());
+    } catch (std::exception &e) {
+      throw std::runtime_error(std::string("General error: ") + e.what());
     } catch (...) {
-      throw std::runtime_error("Onnxruntime error");
+      throw std::runtime_error("Unknown error occurred.");
     }
+
     result.SetValue(row, struct_val_list);
     result.SetVectorType(VectorType::CONSTANT_VECTOR);
   }

unit_test/CMakeLists.txt

@@ -0,0 +1,31 @@
+cmake_minimum_required(VERSION 3.10)
+
+project(MNIST_Inference)
+
+
+set(INPUT_DIR "${CMAKE_SOURCE_DIR}/../unit_test/mnist")
+set(OUTPUT_DIR "${CMAKE_BINARY_DIR}/extension/onnx/unit_test/mnist")
+message(STATUS "Input Directory: ${INPUT_DIR}")
+message(STATUS "Output Directory: ${OUTPUT_DIR}")
+file(MAKE_DIRECTORY ${OUTPUT_DIR})
+
+set(CMAKE_CXX_STANDARD 11)
+
+include_directories(${OpenCV_INCLUDE_DIRS})
+find_package(OpenCV REQUIRED)
+
+add_executable(mnist_inference mnist_inference.cpp)
+add_custom_command(TARGET mnist_inference POST_BUILD
+        COMMAND ${CMAKE_COMMAND} -E copy_directory ${INPUT_DIR} ${OUTPUT_DIR}
+        COMMENT "Copying files from ${INPUT_DIR} to ${OUTPUT_DIR}"
+)
+
+set_target_properties(mnist_inference PROPERTIES
+        BUILD_RPATH ${DUCKDB_LIB_DIR}
+        INSTALL_RPATH ${DUCKDB_LIB_DIR}
+)
+
+
+target_link_libraries(mnist_inference PRIVATE duckdb ${OpenCV_LIBS})
+
+

unit_test/mnist_inference.cpp

@@ -0,0 +1,153 @@
+#include <duckdb.h>
+#include <iostream>
+#include <opencv2/opencv.hpp>
+#include <vector>
+
+cv::Mat preprocess_image(const std::string &image_path) {
+  cv::Mat image = cv::imread(image_path, cv::IMREAD_GRAYSCALE);
+  if (image.empty()) {
+    std::cerr << "Error: Unable to load image!" << std::endl;
+    exit(1);
+  }
+  cv::resize(image, image, cv::Size(28, 28));
+  image.convertTo(image, CV_32F, 1.0 / 255.0);
+  return image;
+}
+
+int main() {
+  cv::Mat image = preprocess_image("mnist/images/7_12.png");
+  if (!image.isContinuous()) {
+    image = image.clone();
+  }
+  std::vector<float> input_data(image.ptr<float>(),
+                                image.ptr<float>() + image.total());
+
+  duckdb_database db;
+  duckdb_connection con;
+
+  if (duckdb_open(NULL, &db) == DuckDBError) {
+    // handle error
+  }
+  if (duckdb_connect(db, &con) == DuckDBError) {
+    // handle error
+  }
+
+  // run queries...
+  duckdb_result res;
+
+  std::string tensor_value = "[";
+  for (int i = 0; i < input_data.size(); ++i) {
+    tensor_value += std::to_string(input_data[i]);
+    if (i != input_data.size() - 1) {
+      tensor_value += ",";
+    }
+  }
+
+  std::string sql = "SELECT "
+                    "onnx('mnist/onnx/mnist-8.onnx',"
+                    "{'shape':[1,1,28,28],'value':" +
+                    tensor_value + "]}) as result";
+
+  duckdb_state state = duckdb_query(con, sql.c_str(), &res);
+  if (state == DuckDBError) {
+    // handle error
+    std::cerr << "Error" << std::endl;
+    exit(1);
+  }
+
+  std::vector<float> output_data;
+  output_data.reserve(10);
+  while (true) {
+    duckdb_data_chunk result = duckdb_fetch_chunk(res);
+    if (!result) {
+      // result is exhausted
+      break;
+    }
+    // get the number of rows from the data chunk
+    idx_t row_count = duckdb_data_chunk_get_size(result);
+    assert(row_count == 1);
+    // get the first column
+    duckdb_vector struct_col = duckdb_data_chunk_get_vector(result, 0);
+    uint64_t *struct_validity = duckdb_vector_get_validity(struct_col);
+
+    duckdb_vector col1_vector = duckdb_struct_vector_get_child(struct_col, 0);
+    duckdb_vector col2_vector = duckdb_struct_vector_get_child(struct_col, 1);
+
+    duckdb_list_entry *list_data_1 =
+        (duckdb_list_entry *)duckdb_vector_get_data(col1_vector);
+    uint64_t *list_validity_1 = duckdb_vector_get_validity(col1_vector);
+    // get the child column of the list
+    duckdb_vector list_child_1 = duckdb_list_vector_get_child(col1_vector);
+    int32_t *child_data_1 = (int32_t *)duckdb_vector_get_data(list_child_1);
+    uint64_t *child_validity_1 = duckdb_vector_get_validity(list_child_1);
+
+    duckdb_list_entry *list_data_2 =
+        (duckdb_list_entry *)duckdb_vector_get_data(col2_vector);
+    uint64_t *list_validity_2 = duckdb_vector_get_validity(col2_vector);
+    // get the child column of the list
+    duckdb_vector list_child_2 = duckdb_list_vector_get_child(col2_vector);
+    float *child_data_2 = (float *)duckdb_vector_get_data(list_child_2);
+    uint64_t *child_validity_2 = duckdb_vector_get_validity(list_child_2);
+
+    for (idx_t row = 0; row < row_count; row++) {
+      if (!duckdb_validity_row_is_valid(list_validity_1, row)) {
+        // entire list is NULL
+        printf("NULL\n");
+        continue;
+      }
+      // read the list offsets for this row
+      duckdb_list_entry list = list_data_1[row];
+      printf("[");
+      for (idx_t child_idx = list.offset; child_idx < list.offset + list.length;
+           child_idx++) {
+        if (child_idx > list.offset) {
+          printf(", ");
+        }
+        if (!duckdb_validity_row_is_valid(child_validity_1, child_idx)) {
+          // col1 is NULL
+          printf("NULL");
+        } else {
+          printf("%lld", child_data_1[child_idx]);
+        }
+      }
+      printf("]\n");
+
+      duckdb_list_entry list2 = list_data_2[row];
+      printf("[");
+      for (idx_t child_idx = list2.offset;
+           child_idx < list2.offset + list2.length; child_idx++) {
+        if (child_idx > list2.offset) {
+          printf(", ");
+        }
+        if (!duckdb_validity_row_is_valid(child_validity_2, child_idx)) {
+          // col2 is NULL
+          printf("NULL");
+        } else {
+          printf("%f", child_data_2[child_idx]);
+          output_data.push_back(child_data_2[child_idx]);
+        }
+      }
+      printf("]\n");
+    }
+    duckdb_destroy_data_chunk(&result);
+  }
+
+  // clean-up
+  duckdb_destroy_result(&res);
+  duckdb_disconnect(&con);
+  duckdb_close(&db);
+
+  std::cout << "Output probabilities: ";
+  for (int i = 0; i < 10; ++i) {
+    std::cout << output_data[i] << " ";
+  }
+  std::cout << std::endl;
+
+  //
+  int predicted_class =
+      std::distance(output_data.begin(),
+                    std::max_element(output_data.begin(), output_data.end()));
+  std::cout << "Predicted Class: " << predicted_class << std::endl;
+
+  return 0;
+}

vcpkg.json

@@ -1,5 +1,6 @@
 {
   "dependencies": [
-    "openssl"
+    "openssl",
+    "opencv"
   ]
 }