model_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\11918\\AppData\\Local\\Temp/ipykernel_20484/1676956454.py:95: TracerWarning: Iterating over a tensor might cause the trace to be incorrect. Passing a tensor of different shape won't change the number of iterations executed (and might lead to errors or silently give incorrect results).\n",
      "  for left_row in left_data:\n",
      "C:\\Users\\11918\\AppData\\Local\\Temp/ipykernel_20484/1676956454.py:97: TracerWarning: Iterating over a tensor might cause the trace to be incorrect. Passing a tensor of different shape won't change the number of iterations executed (and might lead to errors or silently give incorrect results).\n",
      "  for right_row in right_data:\n",
      "C:\\Users\\11918\\AppData\\Local\\Temp/ipykernel_20484/1676956454.py:99: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if left_key_value == right_key_value:\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============== Diagnostic Run torch.onnx.export version 2.0.1+cpu ==============\n",
      "verbose: False, log level: Level.ERROR\n",
      "======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================\n",
      "\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: numpy.ndarray",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp/ipykernel_20484/1676956454.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m    159\u001b[0m \u001b[1;31m# 执行查询计划\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    160\u001b[0m \u001b[1;32mwhile\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 161\u001b[1;33m     \u001b[0mtuple\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mjoin\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mnext\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    162\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0mtuple\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    163\u001b[0m         \u001b[1;32mbreak\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Local\\Temp/ipykernel_20484/1676956454.py\u001b[0m in \u001b[0;36mnext\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    128\u001b[0m         \u001b[0mdummy_left_input\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrandn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mleft_data\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mleft_data\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    129\u001b[0m         \u001b[0mdummy_right_input\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrandn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mright_data\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mright_data\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 130\u001b[1;33m         torch.onnx.export(model, (dummy_left_input, dummy_right_input), \"join_model.onnx\",\n\u001b[0m\u001b[0;32m    131\u001b[0m                           input_names=[\"left_input\", \"right_input\"], output_names=[\"output\"], verbose=False)\n\u001b[0;32m    132\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Python\\Python310\\lib\\site-packages\\torch\\onnx\\utils.py\u001b[0m in \u001b[0;36mexport\u001b[1;34m(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions)\u001b[0m\n\u001b[0;32m    504\u001b[0m     \"\"\"\n\u001b[0;32m    505\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 506\u001b[1;33m     _export(\n\u001b[0m\u001b[0;32m    507\u001b[0m         \u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    508\u001b[0m         \u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Python\\Python310\\lib\\site-packages\\torch\\onnx\\utils.py\u001b[0m in \u001b[0;36m_export\u001b[1;34m(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, onnx_shape_inference, export_modules_as_functions)\u001b[0m\n\u001b[0;32m   1546\u001b[0m             \u001b[0m_validate_dynamic_axes\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdynamic_axes\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput_names\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0moutput_names\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1547\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1548\u001b[1;33m             graph, params_dict, torch_out = _model_to_graph(\n\u001b[0m\u001b[0;32m   1549\u001b[0m                 \u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1550\u001b[0m                 \u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Python\\Python310\\lib\\site-packages\\torch\\onnx\\utils.py\u001b[0m in \u001b[0;36m_model_to_graph\u001b[1;34m(model, args, verbose, input_names, output_names, operator_export_type, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes)\u001b[0m\n\u001b[0;32m   1111\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1112\u001b[0m     \u001b[0mmodel\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_pre_trace_quant_model\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1113\u001b[1;33m     \u001b[0mgraph\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtorch_out\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmodule\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_create_jit_graph\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1114\u001b[0m     \u001b[0mparams_dict\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_get_named_param_dict\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mgraph\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1115\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Python\\Python310\\lib\\site-packages\\torch\\onnx\\utils.py\u001b[0m in \u001b[0;36m_create_jit_graph\u001b[1;34m(model, args)\u001b[0m\n\u001b[0;32m    987\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0mgraph\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtorch_out\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    988\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 989\u001b[1;33m     \u001b[0mgraph\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtorch_out\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_trace_and_get_graph_from_model\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    990\u001b[0m     \u001b[0m_C\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_jit_pass_onnx_lint\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mgraph\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    991\u001b[0m     \u001b[0mstate_dict\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_unique_state_dict\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Python\\Python310\\lib\\site-packages\\torch\\onnx\\utils.py\u001b[0m in \u001b[0;36m_trace_and_get_graph_from_model\u001b[1;34m(model, args)\u001b[0m\n\u001b[0;32m    891\u001b[0m     \u001b[0mprev_autocast_cache_enabled\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mis_autocast_cache_enabled\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    892\u001b[0m     \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mset_autocast_cache_enabled\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;32mFalse\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 893\u001b[1;33m     trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(\n\u001b[0m\u001b[0;32m    894\u001b[0m         \u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    895\u001b[0m         \u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Python\\Python310\\lib\\site-packages\\torch\\jit\\_trace.py\u001b[0m in \u001b[0;36m_get_trace_graph\u001b[1;34m(f, args, kwargs, strict, _force_outplace, return_inputs, _return_inputs_states)\u001b[0m\n\u001b[0;32m   1266\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1267\u001b[0m         \u001b[0margs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1268\u001b[1;33m     \u001b[0mouts\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mONNXTracedModule\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mf\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mstrict\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_force_outplace\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mreturn_inputs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_return_inputs_states\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1269\u001b[0m     \u001b[1;32mreturn\u001b[0m \u001b[0mouts\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1499\u001b[0m                 \u001b[1;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[1;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1500\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[1;32m-> 1501\u001b[1;33m             \u001b[1;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1502\u001b[0m         \u001b[1;31m# Do not call functions when jit is used\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1503\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Python\\Python310\\lib\\site-packages\\torch\\jit\\_trace.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, *args)\u001b[0m\n\u001b[0;32m    125\u001b[0m                 \u001b[1;32mreturn\u001b[0m \u001b[0mtuple\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout_vars\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    126\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 127\u001b[1;33m         graph, out = torch._C._create_graph_by_tracing(\n\u001b[0m\u001b[0;32m    128\u001b[0m             \u001b[0mwrapper\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    129\u001b[0m             \u001b[0min_vars\u001b[0m \u001b[1;33m+\u001b[0m \u001b[0mmodule_state\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Python\\Python310\\lib\\site-packages\\torch\\jit\\_trace.py\u001b[0m in \u001b[0;36mwrapper\u001b[1;34m(*args)\u001b[0m\n\u001b[0;32m    119\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_return_inputs_states\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    120\u001b[0m                 \u001b[0minputs_states\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0minputs_states\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtrace_inputs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 121\u001b[1;33m             \u001b[0mout_vars\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_flatten\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mouts\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    122\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout_vars\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m==\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    123\u001b[0m                 \u001b[1;32mreturn\u001b[0m \u001b[0mout_vars\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: numpy.ndarray"
     ]
    }
   ],
   "source": [
    "import pyarrow as pa\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import onnxruntime as ort\n",
    "import onnx\n",
    "\n",
    "class Scan():\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "        self.index = 0\n",
    "\n",
    "    def next(self):\n",
    "        if self.index < len(self.data):\n",
    "            result = self.data.slice(self.index, 1).to_pandas().iloc[0].to_numpy()\n",
    "            result = result[np.newaxis, :]\n",
    "            self.index += 1\n",
    "            return result\n",
    "        else:\n",
    "            return None\n",
    "\n",
    "class ProjectionModel(nn.Module):\n",
    "    def __init__(self, proj):\n",
    "        super(ProjectionModel, self).__init__()\n",
    "        self.proj = proj\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x[:, self.proj]\n",
    "\n",
    "class Projection():\n",
    "    def __init__(self, child, columns):\n",
    "        self.child = child\n",
    "        self.columns = columns\n",
    "\n",
    "    def next(self):\n",
    "        tensor = self.child.next()\n",
    "        if tensor is None:\n",
    "            return None\n",
    "        tensor = tensor.astype(np.float32)\n",
    "        model = ProjectionModel(self.columns)\n",
    "        # 将模型转换为ONNX格式\n",
    "        dummy_input = torch.randn(tensor.shape[0], tensor.shape[1])\n",
    "        torch.onnx.export(model, dummy_input, \"projection_model.onnx\", input_names=[\"input\"], output_names=[\"output\"], verbose=False)\n",
    "        # 使用ONNX Runtime执行模型\n",
    "        ort_session = ort.InferenceSession(\"projection_model.onnx\")\n",
    "        # 运行模型并获取输出\n",
    "        ort_inputs = {\"input\": tensor}\n",
    "        ort_outputs = ort_session.run(None, ort_inputs)\n",
    "        # 输出结果\n",
    "        result = ort_outputs[0]\n",
    "        return result\n",
    "\n",
    "class SelectionModel(nn.Module):\n",
    "    def __init__(self, predicate):\n",
    "        super(SelectionModel, self).__init__()\n",
    "        self.predicate = predicate\n",
    "\n",
    "    def forward(self, x):\n",
    "        mask = self.predicate(x)\n",
    "        # print(\"mask\")\n",
    "        return mask\n",
    "\n",
    "class Selection():\n",
    "    def __init__(self, child, predicate):\n",
    "        self.child = child\n",
    "        self.predicate = predicate\n",
    "\n",
    "    def next(self):\n",
    "        while True:\n",
    "            tensor = self.child.next()\n",
    "            if tensor is None:\n",
    "                return None\n",
    "            tensor = tensor.astype(np.float32)\n",
    "            model = SelectionModel(self.predicate)\n",
    "            dummy_input = torch.randn(tensor.shape[0], tensor.shape[1])\n",
    "            torch.onnx.export(model, dummy_input, \"selection_model.onnx\", input_names=[\"input\"], output_names=[\"output\"], verbose=False)\n",
    "            ort_session = ort.InferenceSession(\"selection_model.onnx\")\n",
    "            ort_inputs = {\"input\": tensor}\n",
    "            ort_outputs = ort_session.run(None, ort_inputs)\n",
    "            # print(ort_outputs[0][0])\n",
    "            if ort_outputs[0][0]:\n",
    "                return tensor\n",
    "\n",
    "class JoinModel(nn.Module):\n",
    "    def __init__(self, left_key, right_key):\n",
    "        super(JoinModel, self).__init__()\n",
    "        self.left_key = left_key\n",
    "        self.right_key = right_key\n",
    "\n",
    "    def forward(self, left_data, right_data):\n",
    "        left_key_data = left_data[:, self.left_key]\n",
    "        right_key_data = right_data[:, self.right_key]\n",
    "        \n",
    "        result = []\n",
    "        for left_row in left_data:\n",
    "            left_key_value = left_row[self.left_key]\n",
    "            for right_row in right_data:\n",
    "                right_key_value = right_row[self.right_key]\n",
    "                if left_key_value == right_key_value:\n",
    "                    joined_row = np.concatenate((left_row, right_row))\n",
    "                    result.append(joined_row)\n",
    "        \n",
    "        if len(result) > 0:\n",
    "            result = np.stack(result)\n",
    "        else:\n",
    "            result = np.empty((0, left_data.shape[1] + right_data.shape[1]))\n",
    "        \n",
    "        return result\n",
    "\n",
    "class Join():\n",
    "    def __init__(self, left_child, right_child, left_key, right_key):\n",
    "        self.left_child = left_child\n",
    "        self.right_child = right_child\n",
    "        self.left_key = left_key\n",
    "        self.right_key = right_key\n",
    "\n",
    "    def next(self):\n",
    "        left_data = self.left_child.next()\n",
    "        right_data = self.right_child.next()\n",
    "\n",
    "        if left_data is None or right_data is None:\n",
    "            return None\n",
    "\n",
    "        left_data = left_data.astype(np.float32)\n",
    "        right_data = right_data.astype(np.float32)\n",
    "\n",
    "        model = JoinModel(self.left_key, self.right_key)\n",
    "        dummy_left_input = torch.randn(left_data.shape[0], left_data.shape[1])\n",
    "        dummy_right_input = torch.randn(right_data.shape[0], right_data.shape[1])\n",
    "        torch.onnx.export(model, (dummy_left_input, dummy_right_input), \"join_model.onnx\",\n",
    "                          input_names=[\"left_input\", \"right_input\"], output_names=[\"output\"], verbose=False)\n",
    "\n",
    "        ort_session = ort.InferenceSession(\"join_model.onnx\")\n",
    "        ort_inputs = {\"left_input\": left_data, \"right_input\": right_data}\n",
    "        ort_outputs = ort_session.run(None, ort_inputs)\n",
    "\n",
    "        result = ort_outputs[0]\n",
    "        return result\n",
    "\n",
    "# 示例数据\n",
    "left_data = [\n",
    "    {'id': 1, 'gender': 0},\n",
    "    {'id': 2, 'gender': 1},\n",
    "    {'id': 3, 'gender': 1}\n",
    "]\n",
    "right_data = [\n",
    "    {'id': 1, 'age': 25},\n",
    "    {'id': 2, 'age': 28},\n",
    "    {'id': 4, 'age': 30}\n",
    "]\n",
    "left_table = pa.Table.from_pylist(left_data)\n",
    "right_table = pa.Table.from_pylist(right_data)\n",
    "\n",
    "# 构建查询计划\n",
    "left_scan = Scan(left_table)\n",
    "right_scan = Scan(right_table)\n",
    "join = Join(left_scan, right_scan, 0, 0)\n",
    "\n",
    "# 执行查询计划\n",
    "while True:\n",
    "    tuple = join.next()\n",
    "    if tuple is None:\n",
    "        break\n",
    "    print(tuple)\n",
    "\n",
    "\n",
    "\n",
    "# # 示例数据\n",
    "# data = [\n",
    "#     {'id': 1, 'gender': 0, 'age': 38, 'length': 177.2},\n",
    "#     {'id': 2, 'gender': 1, 'age': 40, 'length': 178.8},\n",
    "#     {'id': 3, 'gender': 0, 'age': 35, 'length': 175.5}\n",
    "# ]\n",
    "# data0 = pa.Table.from_pylist(data)\n",
    "# # 构建查询计划\n",
    "# scan0 = Scan(data0)\n",
    "# # print(scan0.next())\n",
    "# projection0 = Projection(scan0, [0, 2, 3])\n",
    "# selection0 = Selection(projection0, lambda t: t[:, 1] < 39)\n",
    "# # selection0 = Selection(projection0, lambda t: t[:, 2] == 177.2)\n",
    "\n",
    "# # 执行查询计划\n",
    "# while True:\n",
    "#     tuple = selection0.next()\n",
    "#     if tuple is None:\n",
    "#         break\n",
    "#     print(tuple)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}