fix: limit before adding new data

Author: ChenZiHong-Gavin

graphgen/models/vis/__init__.py

@@ -121,16 +121,16 @@ def read(
             else:
                 combined_ds = read_tasks[0].union(*read_tasks[1:])
 
+            if read_nums is not None:
+                combined_ds = combined_ds.limit(read_nums)
+
             combined_ds = combined_ds.map(
                 lambda record: {
                     **record,
                     "_trace_id": compute_mm_hash(record, prefix="doc-"),
                 }
             )
 
-            if read_nums is not None:
-                combined_ds = combined_ds.limit(read_nums)
-
             # sample record
             for i, item in enumerate(combined_ds.take(1)):
                 logger.debug("[READ] Sample record %d: %s", i, item)