change dataset_utils

dmlc · Oct 15, 2024 · 13d2dac · 13d2dac
1 parent 7bf402f
commit 13d2dac
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tools/distpartitioning/dataset_utils.py b/tools/distpartitioning/dataset_utils.py
@@ -547,9 +547,11 @@ def get_dataset(
                     autogenerate_column_names=True,
                 )
                 parse_options = pyarrow.csv.ParseOptions(delimiter=" ")
+                if os.path.getsize(edge_file) == 0:
                 # if getsize() == 0, the file is empty, indicating that the partition doesn't have this attribute.
                 # The src_ids and dst_ids should remain empty.
-                if os.path.getsize(edge_file) != 0:
+                    continue
+                else:
                     with pyarrow.csv.open_csv(
                         edge_file,
                         read_options=read_options,