Auto-sync from demo at Fri Jan 30 05:51:20 UTC 2026

graphgen/bases/__init__.py

@@ -1,5 +1,6 @@
 from .base_evaluator import BaseKGEvaluator, BaseQAEvaluator, BaseTripleEvaluator
 from .base_extractor import BaseExtractor
+from .base_filter import BaseValueFilter
 from .base_generator import BaseGenerator
 from .base_kg_builder import BaseKGBuilder
 from .base_llm_wrapper import BaseLLMWrapper

graphgen/bases/base_filter.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+from typing import Any, Union
+
+import numpy as np
+
+
+class BaseFilter(ABC):
+    @abstractmethod
+    def filter(self, data: Any) -> bool:
+        """
+        Filter the data and return True if it passes the filter, False otherwise.
+        """
+        raise NotImplementedError
+
+
+class BaseValueFilter(BaseFilter, ABC):
+    @abstractmethod
+    def filter(self, data: Union[int, float, np.number]) -> bool:
+        """
+        Filter the numeric value and return True if it passes the filter, False otherwise.
+        """
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def filter_type(self) -> str:
+        """
+        Return the type of filter (e.g., "greater_than", "less_than", etc.)
+        """
+        raise NotImplementedError

graphgen/engine.py

@@ -2,7 +2,6 @@ import inspect
 import logging
 import os
 from collections import defaultdict, deque
-from functools import wraps
 from typing import Any, Callable, Dict, List, Set
 
 import ray
@@ -103,7 +102,6 @@ class Engine:
         kv_namespaces = set()
         graph_namespaces = set()
 
-        # TODO: Temporarily hard-coded; node storage will be centrally managed later.
         for node in self.config.nodes:
             op_name = node.op_name
             if self._function_needs_param(op_name, "kv_backend"):
@@ -232,62 +230,38 @@ class Engine:
 
         input_ds = self._get_input_dataset(node, initial_ds)
 
-        if inspect.isclass(op_handler):
-            execution_params = node.execution_params or {}
-            replicas = execution_params.get("replicas", 1)
-            batch_size = (
-                int(execution_params.get("batch_size"))
-                if "batch_size" in execution_params
-                else "default"
+        # if inspect.isclass(op_handler):
+        execution_params = node.execution_params or {}
+        replicas = execution_params.get("replicas", 1)
+        batch_size = (
+            int(execution_params.get("batch_size"))
+            if "batch_size" in execution_params
+            else "default"
+        )
+        compute_resources = execution_params.get("compute_resources", {})
+
+        if node.type == "aggregate":
+            self.datasets[node.id] = input_ds.repartition(1).map_batches(
+                op_handler,
+                compute=ray.data.ActorPoolStrategy(min_size=1, max_size=1),
+                batch_size=None,  # aggregate processes the whole dataset at once
+                num_gpus=compute_resources.get("num_gpus", 0)
+                if compute_resources
+                else 0,
+                fn_constructor_kwargs=node_params,
+                batch_format="pandas",
             )
-            compute_resources = execution_params.get("compute_resources", {})
-
-            if node.type == "aggregate":
-                self.datasets[node.id] = input_ds.repartition(1).map_batches(
-                    op_handler,
-                    compute=ray.data.ActorPoolStrategy(min_size=1, max_size=1),
-                    batch_size=None,  # aggregate processes the whole dataset at once
-                    num_gpus=compute_resources.get("num_gpus", 0)
-                    if compute_resources
-                    else 0,
-                    fn_constructor_kwargs=node_params,
-                    batch_format="pandas",
-                )
-            else:
-                # others like map, filter, flatmap, map_batch let actors process data inside batches
-                self.datasets[node.id] = input_ds.map_batches(
-                    op_handler,
-                    compute=ray.data.ActorPoolStrategy(min_size=1, max_size=replicas),
-                    batch_size=batch_size,
-                    num_gpus=compute_resources.get("num_gpus", 0)
-                    if compute_resources
-                    else 0,
-                    fn_constructor_kwargs=node_params,
-                    batch_format="pandas",
-                )
-
         else:
-
-            @wraps(op_handler)
-            def func_wrapper(row_or_batch: Dict[str, Any]) -> Dict[str, Any]:
-                return op_handler(row_or_batch, **node_params)
-
-            if node.type == "map":
-                self.datasets[node.id] = input_ds.map(func_wrapper)
-            elif node.type == "filter":
-                self.datasets[node.id] = input_ds.filter(func_wrapper)
-            elif node.type == "flatmap":
-                self.datasets[node.id] = input_ds.flat_map(func_wrapper)
-            elif node.type == "aggregate":
-                self.datasets[node.id] = input_ds.repartition(1).map_batches(
-                    func_wrapper, batch_format="default"
-                )
-            elif node.type == "map_batch":
-                self.datasets[node.id] = input_ds.map_batches(func_wrapper)
-            else:
-                raise ValueError(
-                    f"Unsupported node type {node.type} for node {node.id}"
-                )
+            self.datasets[node.id] = input_ds.map_batches(
+                op_handler,
+                compute=ray.data.ActorPoolStrategy(min_size=1, max_size=replicas),
+                batch_size=batch_size,
+                num_gpus=compute_resources.get("num_gpus", 0)
+                if compute_resources
+                else 0,
+                fn_constructor_kwargs=node_params,
+                batch_format="pandas",
+            )
 
     def execute(
         self, initial_ds: ray.data.Dataset, output_dir: str
@@ -315,6 +289,14 @@ class Engine:
                 logger.info("Node %s output saved to %s", node.id, node_output_path)
 
                 # ray will lazy read the dataset
-                self.datasets[node.id] = ray.data.read_json(node_output_path)
+                if os.path.exists(node_output_path) and os.listdir(node_output_path):
+                    self.datasets[node.id] = ray.data.read_json(node_output_path)
+                else:
+                    self.datasets[node.id] = ray.data.from_items([])
+                    logger.warning(
+                        "Node %s output path %s is empty. Created an empty dataset.",
+                        node.id,
+                        node_output_path,
+                    )
 
         return self.datasets

graphgen/models/__init__.py

@@ -6,6 +6,7 @@ from .evaluator import (
     StructureEvaluator,
     UniEvaluator,
 )
+from .filter import RangeFilter
 from .generator import (
     AggregatedGenerator,
     AtomicGenerator,

graphgen/models/filter/__init__.py

@@ -0,0 +1 @@
+from .range_filter import RangeFilter

graphgen/models/filter/range_filter.py

@@ -0,0 +1,40 @@
+from typing import Union
+
+import numpy as np
+
+from graphgen.bases import BaseValueFilter
+
+
+class RangeFilter(BaseValueFilter):
+    """
+    keeps values within a specified range [min_val, max_val] (inclusive or exclusive)
+    """
+
+    def __init__(
+        self,
+        min_val: float,
+        max_val: float,
+        left_inclusive: bool = True,
+        right_inclusive: bool = True,
+    ):
+        self.min_val = min_val
+        self.max_val = max_val
+        self.left_inclusive = left_inclusive
+        self.right_inclusive = right_inclusive
+
+    def filter(self, data: Union[int, float, np.number]) -> bool:
+        value = float(data)
+        if self.left_inclusive and self.right_inclusive:
+            return self.min_val <= value <= self.max_val
+        if self.left_inclusive and not self.right_inclusive:
+            return self.min_val <= value < self.max_val
+        if not self.left_inclusive and self.right_inclusive:
+            return self.min_val < value <= self.max_val
+        return self.min_val < value < self.max_val
+
+    @property
+    def filter_type(self) -> str:
+        return "range"
+
+    def __repr__(self) -> str:
+        return f"RangeFilter({self.min_val}, {self.max_val})"

graphgen/operators/__init__.py

@@ -2,6 +2,7 @@ from .build_kg import BuildKGService
 from .chunk import ChunkService
 from .evaluate import EvaluateService
 from .extract import ExtractService
+from .filter import FilterService
 from .generate import GenerateService
 from .judge import JudgeService
 from .partition import PartitionService
@@ -9,7 +10,6 @@ from .quiz import QuizService
 from .read import read
 from .search import SearchService
 
-
 operators = {
     "read": read,
     "chunk": ChunkService,
@@ -21,4 +21,5 @@ operators = {
     "partition": PartitionService,
     "generate": GenerateService,
     "evaluate": EvaluateService,
+    "filter": FilterService,
 }

graphgen/operators/filter/__init__.py

@@ -0,0 +1 @@
+from .filter_service import FilterService

graphgen/operators/filter/filter_service.py

@@ -0,0 +1,49 @@
+from typing import Tuple
+
+from graphgen.bases import BaseOperator
+from graphgen.utils import logger
+
+
+class FilterService(BaseOperator):
+    def __init__(
+        self, working_dir: str = "cache", kv_backend: str = "rocksdb", **filter_kwargs
+    ):
+        super().__init__(
+            working_dir=working_dir, kv_backend=kv_backend, op_name="filter"
+        )
+        method = filter_kwargs["method"]
+        method_params = filter_kwargs["method_params"]
+        self.metric = method_params["metric"]
+        if method == "range":
+            from graphgen.models import RangeFilter
+
+            self.filter_instance = RangeFilter(
+                min_val=method_params["min_val"],
+                max_val=method_params["max_val"],
+                left_inclusive=method_params.get("left_inclusive", True),
+                right_inclusive=method_params.get("right_inclusive", True),
+            )
+        else:
+            raise ValueError(f"Unsupported filter method: {method}")
+
+    def process(self, batch: list) -> Tuple[list, dict]:
+        """
+        Filter the items in the batch.
+        :return: A tuple of (results, meta_updates)
+            results: A list of filtered items.
+            meta_updates: empty as filtering does not create new items.
+        """
+        results = []
+        meta_updates = {}
+
+        for item in batch:
+            value = item["metrics"].get(self.metric)
+            if value is None:
+                logger.warning(
+                    f"Item {item} does not have metric {self.metric}. Skipping."
+                )
+                continue
+            if self.filter_instance.filter(value):
+                results.append(item)
+
+        return results, meta_updates