Merge pull request #72 from wey-gu/fix_int_vid_field_ngload

fix: failed to ng_load for int col of vid
wey-gu · Aug 9, 2024 · 6d58ca6 · 6d58ca6
2 parents 2c9f116 + 34a60a2
commit 6d58ca6
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 5 deletions.
diff --git a/ngql/magic.py b/ngql/magic.py
@@ -148,7 +148,7 @@ def ngql(self, line, cell=None, local_ns={}):
             fancy_print("[ERROR] Connection is not ready", color="pink")
             return f"Connection State: { connection_state }"
         if connection_state == CONNECTION_POOL_CREATED:
-            fancy_print("Connection Pool Created", color="blue")
+            fancy_print("[OK] Connection Pool Created", color="green")
             if not cell:
                 return self._stylized(self._show_spaces())
             else:
@@ -192,13 +192,17 @@ def _init_connection_pool(self, args: Optional[Any] = None):
             except RuntimeError:
                 # When GraphD is over TLS
                 fancy_print(
-                    "[ERROR] Got RuntimeError, trying to connect assuming GraphD is over TLS",
+                    "[WARN] Got RuntimeError, trying to connect assuming NebulaGraph is over TLS",
                     color="pink",
                 )
                 ssl_config = SSL_config()
                 connect_init_result = connection_pool.init(
                     [(args.address, args.port)], config, ssl_config
                 )
+                fancy_print(
+                    f"[OK] Connection State: { connect_init_result }, TLS: True",
+                    color="blue",
+                )
             if not connect_init_result:
                 return CONNECTION_POOL_INIT_FAILURE
             else:

diff --git a/ngql/ng_load.py b/ngql/ng_load.py
@@ -243,7 +243,15 @@ def safe_load_file(source, file_type, header_option=None, limit=None):
             else:
                 query = f"INSERT VERTEX `{args.tag}` (`{'`, `'.join(prop_columns)}`) VALUES "
             for index, row in batch.iterrows():
-                raw_vid_str = row["___vid"].strip('"').replace('"', '\\"')
+                if pd.isna(row["___vid"]) or row["___vid"] == "":
+                    fancy_print(
+                        f"[WARNING] Skipping row with empty VID: {row}", "yellow"
+                    )
+                    continue
+                elif isinstance(row["___vid"], str):
+                    raw_vid_str = row["___vid"].strip('"').replace('"', '\\"')
+                else:
+                    raw_vid_str = str(row["___vid"])
                 vid_str = f"{QUOTE_VID}{raw_vid_str}{QUOTE_VID}"
 
                 prop_str = ""
@@ -304,9 +312,27 @@ def safe_load_file(source, file_type, header_option=None, limit=None):
                     f"INSERT EDGE `{args.edge}` (`{'`, `'.join(prop_columns)}`) VALUES "
                 )
             for index, row in batch.iterrows():
-                raw_src_str = row["___src"].strip('"').replace('"', '\\"')
+                if pd.isna(row["___src"]) or row["___src"] == "":
+                    fancy_print(
+                        f"[WARNING] Skipping row with empty source VID: {row}", "yellow"
+                    )
+                    continue
+                elif isinstance(row["___src"], str):
+                    raw_src_str = row["___src"].strip('"').replace('"', '\\"')
+                else:
+                    raw_src_str = str(row["___src"])
                 src_str = f"{QUOTE_VID}{raw_src_str}{QUOTE_VID}"
-                raw_dst_str = row["___dst"].strip('"').replace('"', '\\"')
+
+                if pd.isna(row["___dst"]) or row["___dst"] == "":
+                    fancy_print(
+                        f"[WARNING] Skipping row with empty destination VID: {row}",
+                        "yellow",
+                    )
+                    continue
+                elif isinstance(row["___dst"], str):
+                    raw_dst_str = row["___dst"].strip('"').replace('"', '\\"')
+                else:
+                    raw_dst_str = str(row["___dst"])
                 dst_str = f"{QUOTE_VID}{raw_dst_str}{QUOTE_VID}"
                 prop_str = ""
                 if with_props: