change initialization

Signed-off-by: Nitish Bharambe <[email protected]>

Author: nitbharambe

docs/examples/arrow_example.ipynb

@@ -183,7 +183,8 @@
     "\n",
     "The [power-grid-model documentation on Components](https://power-grid-model.readthedocs.io/en/stable/user_manual/components.html) provides documentation on which components are required and which ones are optional.\n",
     "\n",
-    "Construct the Arrow data as a table with the correct headers and data types."
+    "Construct the Arrow data as a table with the correct headers and data types. \n",
+    "The creation of arrays and combining it in a RecordBatch as well as the method of initializing that RecordBatch is up to the user."
    ]
   },
   {
@@ -208,41 +209,54 @@
     }
    ],
    "source": [
-    "nodes_dict = {\"id\": [1, 2, 3], \"u_rated\": [10500.0, 10500.0, 10500.0]}\n",
-    "\n",
-    "\n",
-    "lines_dict = {\n",
-    "    \"id\": [4, 5],\n",
-    "    \"from_node\": [1, 2],\n",
-    "    \"to_node\": [2, 3],\n",
-    "    \"from_status\": [1, 1],\n",
-    "    \"to_status\": [1, 1],\n",
-    "    \"r1\": [0.11, 0.15],\n",
-    "    \"x1\": [0.12, 0.16],\n",
-    "    \"c1\": [4.1e-05, 5.4e-05],\n",
-    "    \"tan1\": [0.1, 0.1],\n",
-    "    \"r0\": [0.01, 0.05],\n",
-    "    \"x0\": [0.22, 0.06],\n",
-    "    \"c0\": [4.1e-05, 5.4e-05],\n",
-    "    \"tan0\": [0.4, 0.1],\n",
-    "}\n",
-    "\n",
-    "sources_dict = {\"id\": [6], \"node\": [1], \"status\": [1], \"u_ref\": [1.0]}\n",
+    "nodes_schema = pgm_schema(DatasetType.input, ComponentType.node)\n",
+    "nodes = pa.record_batch(\n",
+    "    [\n",
+    "        pa.array([1, 2, 3], type=nodes_schema.field(\"id\").type),\n",
+    "        pa.array([10500.0, 10500.0, 10500.0], type=nodes_schema.field(\"u_rated\").type),\n",
+    "    ],\n",
+    "    names=(\"id\", \"u_rated\"),\n",
+    ")\n",
     "\n",
-    "sym_loads_dict = {\n",
-    "    \"id\": [7, 8],\n",
-    "    \"node\": [2, 3],\n",
-    "    \"status\": [1, 1],\n",
-    "    \"type\": [0, 0],\n",
-    "    \"p_specified\": [1.0, 2.0],\n",
-    "    \"q_specified\": [0.5, 1.5],\n",
-    "}\n",
+    "lines = pa.record_batch(\n",
+    "    {\n",
+    "        \"id\": [4, 5],\n",
+    "        \"from_node\": [1, 2],\n",
+    "        \"to_node\": [2, 3],\n",
+    "        \"from_status\": [1, 1],\n",
+    "        \"to_status\": [1, 1],\n",
+    "        \"r1\": [0.11, 0.15],\n",
+    "        \"x1\": [0.12, 0.16],\n",
+    "        \"c1\": [4.1e-05, 5.4e-05],\n",
+    "        \"tan1\": [0.1, 0.1],\n",
+    "        \"r0\": [0.01, 0.05],\n",
+    "        \"x0\": [0.22, 0.06],\n",
+    "        \"c0\": [4.1e-05, 5.4e-05],\n",
+    "        \"tan0\": [0.4, 0.1],\n",
+    "    },\n",
+    "    schema=pgm_schema(\n",
+    "        DatasetType.input,\n",
+    "        ComponentType.line,\n",
+    "        [\"id\", \"from_node\", \"to_node\", \"from_status\", \"to_status\", \"r1\", \"x1\", \"c1\", \"tan1\", \"r0\", \"x0\", \"c0\", \"tan0\"],\n",
+    "    ),\n",
+    ")\n",
     "\n",
-    "nodes = pa.record_batch(nodes_dict, schema=pgm_schema(DatasetType.input, ComponentType.node, nodes_dict.keys()))\n",
-    "lines = pa.record_batch(lines_dict, schema=pgm_schema(DatasetType.input, ComponentType.line, lines_dict.keys()))\n",
-    "sources = pa.record_batch(sources_dict, schema=pgm_schema(DatasetType.input, ComponentType.source, sources_dict.keys()))\n",
+    "sources = pa.record_batch(\n",
+    "    {\"id\": [6], \"node\": [1], \"status\": [1], \"u_ref\": [1.0]},\n",
+    "    schema=pgm_schema(DatasetType.input, ComponentType.source, [\"id\", \"node\", \"status\", \"u_ref\"]),\n",
+    ")\n",
     "sym_loads = pa.record_batch(\n",
-    "    sym_loads_dict, schema=pgm_schema(DatasetType.input, ComponentType.sym_load, sym_loads_dict.keys())\n",
+    "    {\n",
+    "        \"id\": [7, 8],\n",
+    "        \"node\": [2, 3],\n",
+    "        \"status\": [1, 1],\n",
+    "        \"type\": [0, 0],\n",
+    "        \"p_specified\": [1.0, 2.0],\n",
+    "        \"q_specified\": [0.5, 1.5],\n",
+    "    },\n",
+    "    schema=pgm_schema(\n",
+    "        DatasetType.input, ComponentType.sym_load, [\"id\", \"node\", \"status\", \"type\", \"p_specified\", \"q_specified\"]\n",
+    "    ),\n",
     ")\n",
     "\n",
     "nodes\n",
@@ -349,7 +363,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
@@ -564,7 +578,17 @@
     "}\n",
     "\n",
     "asym_loads = pa.record_batch(\n",
-    "    asym_loads_dict, schema=pgm_schema(DatasetType.input, ComponentType.asym_load, asym_loads_dict.keys())\n",
+    "    {\n",
+    "        \"id\": [7, 8],\n",
+    "        \"node\": [2, 3],\n",
+    "        \"status\": [1, 1],\n",
+    "        \"type\": [0, 0],\n",
+    "        \"p_specified\": [[1.0, 1.0e-2, 1.1e-2], [2.0, 2.5, 4.5e2]],\n",
+    "        \"q_specified\": [[0.5, 1.5e3, 0.1], [1.5, 2.5, 1.5e3]],\n",
+    "    },\n",
+    "    schema=pgm_schema(\n",
+    "        DatasetType.input, ComponentType.asym_load, [\"id\", \"node\", \"status\", \"type\", \"p_specified\", \"q_specified\"]\n",
+    "    ),\n",
     ")\n",
     "\n",
     "asym_loads"
@@ -782,6 +806,34 @@
     "pa_asym_node_result"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<pyarrow.lib.DoubleArray object at 0x0000020996813E20>\n",
+       "[\n",
+       "  1,\n",
+       "  0.01,\n",
+       "  0.011,\n",
+       "  2,\n",
+       "  2.5,\n",
+       "  450\n",
+       "]"
+      ]
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pa.array(asym_load_input[\"p_specified\"].flatten(), type=pa.float64())"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},