SuperModelingFactory 0.5.3__tar.gz → 0.5.5__tar.gz

{supermodelingfactory-0.5.3 → supermodelingfactory-0.5.5}/Modeling_Tool/Core/Parallel_ODPS_Manager.py

@@ -17,6 +17,7 @@ from .utils import parse_sql_file
 
 Backend = Literal["thread", "process", "sequential"]
 WriteMode = Literal["overwrite", "append"]
+PullSplitStrategy = Literal["auto", "hash", "row_number"]
 
 
 @dataclass
@@ -26,6 +27,10 @@ class ParallelODPSConfig:
     n_chunks: int | None = None
     n_jobs: int = 3
     backend: Backend = "thread"
+    pull_split_strategy: PullSplitStrategy = "auto"
+    row_number_order_by: str | None = None
+    row_number_col: str = "__smf_parallel_odps_rn__"
+    validate_unique_key: bool = True
     chunk_filter_key: str = "chunk_filter"
     tmp_dir: Path = field(default_factory=lambda: Path("data/_chunks"))
     tmp_table_prefix: str = "tmp_parallel_odps"
@@ -55,22 +60,23 @@ def _delete_table(runner: ODPSRunner, table_name: str) -> None:
     runner.run_sql(f"DROP TABLE IF EXISTS {table_name};", to_df=False)
 
 
+def _strip_trailing_semicolon(sql: str) -> str:
+    return sql.strip().rstrip(";").strip()
+
+
 def _pull_one_chunk(
-    chunk_id: int,
+    chunk_spec: dict[str, Any],
     runner: ODPSRunner | None,
-    sql_path: str,
-    template_kwargs: dict[str, Any],
-    chunk_filter_key: str,
-    unique_key: str,
-    n_chunks: int,
     tmp_dir_str: str,
 ) -> dict[str, Any]:
     active_runner = _get_runner(runner)
-    kwargs = dict(template_kwargs)
-    kwargs[chunk_filter_key] = f"ABS(HASH({unique_key})) % {n_chunks} = {chunk_id}"
-    rendered_sql = parse_sql_file(sql_path=sql_path, **kwargs)
+    chunk_id = int(chunk_spec["chunk"])
+    rendered_sql = str(chunk_spec["sql"])
+    drop_cols = list(chunk_spec.get("drop_cols") or [])
 
     df_chunk = active_runner.run_sql(rendered_sql, to_df=True, n_process=1)
+    if drop_cols:
+        df_chunk = df_chunk.drop(columns=[col for col in drop_cols if col in df_chunk.columns])
     chunk_path = Path(tmp_dir_str) / f"_pull_chunk_{chunk_id:04d}.csv"
     df_chunk.to_csv(chunk_path, index=False)
     n_rows = len(df_chunk)
@@ -102,6 +108,7 @@ class ParallelODPSManager:
 
     _VALID_BACKENDS = {"thread", "process", "sequential"}
     _VALID_WRITE_MODES = {"overwrite", "append"}
+    _VALID_PULL_STRATEGIES = {"auto", "hash", "row_number"}
 
     def __init__(self, config: ParallelODPSConfig, odps_runner: ODPSRunner | None = None):
         self.config = config
@@ -112,6 +119,8 @@ class ParallelODPSManager:
         cfg = self.config
         if cfg.backend not in self._VALID_BACKENDS:
             raise ValueError(f"backend must be one of {sorted(self._VALID_BACKENDS)}")
+        if cfg.pull_split_strategy not in self._VALID_PULL_STRATEGIES:
+            raise ValueError(f"pull_split_strategy must be one of {sorted(self._VALID_PULL_STRATEGIES)}")
         if cfg.chunk_size is not None and cfg.chunk_size <= 0:
             raise ValueError("chunk_size must be a positive integer.")
         if cfg.n_chunks is not None and cfg.n_chunks <= 0:
@@ -122,6 +131,8 @@ class ParallelODPSManager:
             raise ValueError("n_jobs must be a positive integer.")
         if not cfg.chunk_filter_key:
             raise ValueError("chunk_filter_key cannot be empty.")
+        if not cfg.row_number_col:
+            raise ValueError("row_number_col cannot be empty.")
         if not cfg.tmp_table_prefix:
             raise ValueError("tmp_table_prefix cannot be empty.")
 
@@ -131,34 +142,151 @@ class ParallelODPSManager:
     def _auto_count_query(self, sql_path: str, template_kwargs: dict[str, Any]) -> str:
         kwargs = dict(template_kwargs)
         kwargs[self.config.chunk_filter_key] = "1=1"
-        rendered_sql = parse_sql_file(sql_path=sql_path, **kwargs).rstrip().rstrip(";")
+        rendered_sql = _strip_trailing_semicolon(parse_sql_file(sql_path=sql_path, **kwargs))
         return f"SELECT COUNT(1) FROM ({rendered_sql}) __count_src;"
 
-    def _resolve_pull_n_chunks(self, count_query: str | None, sql_path: str, template_kwargs: dict[str, Any]) -> int:
+    def _assert_chunk_filter_placeholder(self, sql_path: str, template_kwargs: dict[str, Any]) -> None:
+        sentinel = "__SMF_CHUNK_FILTER_SENTINEL__=1"
+        kwargs = dict(template_kwargs)
+        kwargs[self.config.chunk_filter_key] = sentinel
+        rendered_sql = parse_sql_file(sql_path=sql_path, **kwargs)
+        if sentinel not in rendered_sql:
+            raise ValueError(
+                f"pull SQL template must contain {{{self.config.chunk_filter_key}}}"
+            )
+
+    def _render_pull_sql(self, sql_path: str, template_kwargs: dict[str, Any], chunk_filter: str) -> str:
+        kwargs = dict(template_kwargs)
+        kwargs[self.config.chunk_filter_key] = chunk_filter
+        return parse_sql_file(sql_path=sql_path, **kwargs)
+
+    def _resolve_pull_strategy(self) -> str:
+        cfg = self.config
+        strategy = cfg.pull_split_strategy
+        if strategy == "auto":
+            return "hash" if cfg.unique_key else "row_number"
+        if strategy == "hash" and not cfg.unique_key:
+            raise ValueError("unique_key is required when pull_split_strategy='hash'.")
+        return strategy
+
+    def _resolve_pull_n_chunks(
+        self,
+        count_query: str | None,
+        sql_path: str,
+        template_kwargs: dict[str, Any],
+        strategy: str,
+        staging_table: str | None = None,
+    ) -> int:
         cfg = self.config
-        if not cfg.unique_key:
-            raise ValueError("unique_key is required for pull().")
         if cfg.n_chunks is not None:
             return cfg.n_chunks
         if cfg.chunk_size is None:
             raise ValueError("chunk_size or n_chunks is required for pull().")
         if count_query is None:
-            count_query = self._auto_count_query(sql_path, template_kwargs)
+            if strategy == "row_number":
+                if not staging_table:
+                    raise ValueError("staging_table is required to count row_number pull chunks.")
+                count_query = f"SELECT COUNT(1) FROM {staging_table};"
+            else:
+                count_query = self._auto_count_query(sql_path, template_kwargs)
         count_df = self.odps_runner.run_sql(count_query, to_df=True)
         total_rows = int(count_df.iloc[0, 0])
         return max(1, math.ceil(total_rows / cfg.chunk_size))
 
-    def pull(
+    def _hash_chunk_filter(self, chunk_id: int, n_chunks: int) -> str:
+        return f"ABS(HASH({self.config.unique_key})) % {n_chunks} = {chunk_id}"
+
+    def _validate_hash_pull_sql(
         self,
         sql_path: str,
-        out_path: str,
-        count_query: str | None = None,
-        **template_kwargs: Any,
-    ) -> dict[str, Any]:
+        template_kwargs: dict[str, Any],
+        n_chunks: int,
+    ) -> None:
+        if not self.config.validate_unique_key:
+            return
+        rendered_sql = self._render_pull_sql(
+            sql_path=sql_path,
+            template_kwargs=template_kwargs,
+            chunk_filter=self._hash_chunk_filter(0, n_chunks),
+        )
+        probe_sql = (
+            "SELECT * FROM (\n"
+            f"{_strip_trailing_semicolon(rendered_sql)}\n"
+            ") __smf_unique_key_probe LIMIT 1;"
+        )
+        try:
+            self.odps_runner.run_sql(probe_sql, to_df=True, n_process=1)
+        except Exception as exc:
+            raise ValueError(
+                f"unique_key validation failed for pull SQL. "
+                f"Check that unique_key={self.config.unique_key!r} is visible in the SQL scope."
+            ) from exc
+
+    def _build_hash_pull_chunks(
+        self,
+        sql_path: str,
+        template_kwargs: dict[str, Any],
+        n_chunks: int,
+    ) -> list[dict[str, Any]]:
+        return [
+            {
+                "chunk": chunk_id,
+                "sql": self._render_pull_sql(
+                    sql_path=sql_path,
+                    template_kwargs=template_kwargs,
+                    chunk_filter=self._hash_chunk_filter(chunk_id, n_chunks),
+                ),
+            }
+            for chunk_id in range(n_chunks)
+        ]
+
+    def _row_number_staging_table_name(self, run_id: str) -> str:
+        return f"{self.config.tmp_table_prefix}_{run_id}_pull_stage"
+
+    def _create_row_number_staging_table(
+        self,
+        sql_path: str,
+        template_kwargs: dict[str, Any],
+        staging_table: str,
+    ) -> str:
         cfg = self.config
-        n_chunks = self._resolve_pull_n_chunks(count_query, sql_path, template_kwargs)
-        cfg.tmp_dir.mkdir(parents=True, exist_ok=True)
+        base_sql = _strip_trailing_semicolon(
+            self._render_pull_sql(
+                sql_path=sql_path,
+                template_kwargs=template_kwargs,
+                chunk_filter="1=1",
+            )
+        )
+        order_by = cfg.row_number_order_by or "1"
+        create_sql = (
+            f"CREATE TABLE {staging_table} AS\n"
+            "SELECT\n"
+            f"  ROW_NUMBER() OVER (ORDER BY {order_by}) AS {cfg.row_number_col},\n"
+            "  __smf_base.*\n"
+            "FROM (\n"
+            f"{base_sql}\n"
+            ") __smf_base;"
+        )
+        self.odps_runner.run_sql(create_sql, to_df=False)
+        return create_sql
+
+    def _build_row_number_pull_chunks(self, staging_table: str, n_chunks: int) -> list[dict[str, Any]]:
+        row_number_col = self.config.row_number_col
+        return [
+            {
+                "chunk": chunk_id,
+                "sql": (
+                    f"SELECT *\n"
+                    f"FROM {staging_table}\n"
+                    f"WHERE ({row_number_col} - 1) % {n_chunks} = {chunk_id};"
+                ),
+                "drop_cols": [row_number_col],
+            }
+            for chunk_id in range(n_chunks)
+        ]
 
+    def _run_pull_chunks(self, chunk_specs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        cfg = self.config
         engine_cfg = ParallelApplyConfig(
             split_axis="chunk",
             backend=cfg.backend,
@@ -168,41 +296,83 @@ class ParallelODPSManager:
         )
         result = ParallelApplyEngine(engine_cfg).run(
             func=_pull_one_chunk,
-            chunks=list(range(n_chunks)),
+            chunks=chunk_specs,
             func_args=(
                 self._runner_for_backend(),
-                sql_path,
-                template_kwargs,
-                cfg.chunk_filter_key,
-                cfg.unique_key,
-                n_chunks,
                 str(cfg.tmp_dir),
             ),
         )
-
         if len(result.errors):
-            raise RuntimeError(f"{len(result.errors)}/{n_chunks} ODPS pull chunks failed:\n{result.errors}")
-
-        chunk_summaries = sorted(result.output, key=lambda item: item["chunk"])
-        final_path = Path(out_path)
-        final_path.parent.mkdir(parents=True, exist_ok=True)
-        final_path.unlink(missing_ok=True)
-
-        with open(final_path, "wb") as fout:
-            for idx, summary in enumerate(chunk_summaries):
-                chunk_file = Path(summary["path"])
-                with open(chunk_file, "rb") as fin:
-                    if idx > 0:
-                        fin.readline()
-                    shutil.copyfileobj(fin, fout)
-                chunk_file.unlink()
-
-        return {
-            "n_chunks": n_chunks,
-            "total_rows": sum(item["rows"] for item in chunk_summaries),
-            "out_path": str(final_path),
-            "per_chunk_rows": [item["rows"] for item in chunk_summaries],
-        }
+            raise RuntimeError(f"{len(result.errors)}/{len(chunk_specs)} ODPS pull chunks failed:\n{result.errors}")
+        return sorted(result.output, key=lambda item: item["chunk"])
+
+    def pull(
+        self,
+        sql_path: str,
+        out_path: str,
+        count_query: str | None = None,
+        **template_kwargs: Any,
+    ) -> dict[str, Any]:
+        cfg = self.config
+        self._assert_chunk_filter_placeholder(sql_path, template_kwargs)
+        cfg.tmp_dir.mkdir(parents=True, exist_ok=True)
+        strategy = self._resolve_pull_strategy()
+        staging_table: str | None = None
+        success = False
+
+        try:
+            if strategy == "row_number":
+                run_id = uuid.uuid4().hex[:12]
+                staging_table = self._row_number_staging_table_name(run_id)
+                self._create_row_number_staging_table(sql_path, template_kwargs, staging_table)
+                n_chunks = self._resolve_pull_n_chunks(
+                    count_query=count_query,
+                    sql_path=sql_path,
+                    template_kwargs=template_kwargs,
+                    strategy=strategy,
+                    staging_table=staging_table,
+                )
+                chunk_specs = self._build_row_number_pull_chunks(staging_table, n_chunks)
+            else:
+                n_chunks = self._resolve_pull_n_chunks(
+                    count_query=count_query,
+                    sql_path=sql_path,
+                    template_kwargs=template_kwargs,
+                    strategy=strategy,
+                )
+                self._validate_hash_pull_sql(sql_path, template_kwargs, n_chunks)
+                chunk_specs = self._build_hash_pull_chunks(sql_path, template_kwargs, n_chunks)
+
+            chunk_summaries = self._run_pull_chunks(chunk_specs)
+            final_path = Path(out_path)
+            final_path.parent.mkdir(parents=True, exist_ok=True)
+            final_path.unlink(missing_ok=True)
+
+            with open(final_path, "wb") as fout:
+                for idx, summary in enumerate(chunk_summaries):
+                    chunk_file = Path(summary["path"])
+                    with open(chunk_file, "rb") as fin:
+                        if idx > 0:
+                            fin.readline()
+                        shutil.copyfileobj(fin, fout)
+                    chunk_file.unlink()
+
+            success = True
+            return {
+                "pull_strategy": strategy,
+                "staging_table": staging_table,
+                "n_chunks": n_chunks,
+                "total_rows": sum(item["rows"] for item in chunk_summaries),
+                "out_path": str(final_path),
+                "per_chunk_rows": [item["rows"] for item in chunk_summaries],
+            }
+        finally:
+            if (
+                staging_table
+                and cfg.cleanup_tmp
+                and (success or not cfg.keep_tmp_on_error)
+            ):
+                _delete_table(self.odps_runner, staging_table)
 
     def push(
         self,
@@ -223,6 +393,7 @@ class ParallelODPSManager:
             raise ValueError("push data produced no chunks.")
 
         tmp_tables = [spec["tmp_table"] for spec in chunk_specs]
+        success = False
         try:
             engine_cfg = ParallelApplyConfig(
                 split_axis="chunk",

{supermodelingfactory-0.5.3 → supermodelingfactory-0.5.5}/Modeling_Tool/Pipeline/__init__.py

@@ -31,6 +31,23 @@ from .feature_validation import (
 )
 from .orchestrator import run_modeling_from_validation
 from .screening_artifact import FeatureScreeningArtifact, screen_result_to_summary
+from .field_meta import (
+    FieldMeta,
+    PipelineRegistryEntry,
+    PIPELINE_REGISTRY,
+    config_from_dict,
+    config_from_yaml,
+    config_to_dict,
+    config_to_yaml,
+    extract_config_schema,
+    extract_pipeline_schema,
+    extract_schema,
+    generate_pipeline_code,
+    get_config_field_meta,
+    get_pipeline_registry,
+    get_pipeline_registry_schema,
+    validate_pipeline_config,
+)
 
 __all__ = [
     "RejectInferencePipeline",
@@ -57,4 +74,19 @@ __all__ = [
     "FeatureScreeningArtifact",
     "screen_result_to_summary",
     "run_modeling_from_validation",
+    "FieldMeta",
+    "PipelineRegistryEntry",
+    "PIPELINE_REGISTRY",
+    "get_pipeline_registry",
+    "get_pipeline_registry_schema",
+    "get_config_field_meta",
+    "extract_config_schema",
+    "extract_pipeline_schema",
+    "extract_schema",
+    "config_to_dict",
+    "config_from_dict",
+    "config_to_yaml",
+    "config_from_yaml",
+    "validate_pipeline_config",
+    "generate_pipeline_code",
 ]

{supermodelingfactory-0.5.3 → supermodelingfactory-0.5.5}/Modeling_Tool/Pipeline/feature_validation.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import gc
 import logging
+import warnings
 from dataclasses import dataclass, field, replace
 from pathlib import Path
 from typing import Any, Literal
@@ -33,6 +34,7 @@ class FeatureValidationPipelineConfig:
 
     input_type: Literal["auto", "dataframe", "csv"] = "auto"
     csv_read_kwargs: dict[str, Any] = field(default_factory=dict)
+    enable_batch: bool = False
     feature_batch_size: int | None = None
     feature_batches: list[list[str]] | None = None
     batch_base_cols: list[str] | None = None
@@ -150,10 +152,19 @@ class FeatureValidationPipeline:
 
     def run(self, data: pd.DataFrame | str | Path) -> FeatureValidationPipelineResult:
         input_type = self._resolve_input_type(data)
+        if self.config.enable_batch and input_type != "csv":
+            raise ValueError("enable_batch=True currently requires CSV input.")
         if input_type == "csv":
             csv_path = Path(data)
             if self._is_batch_mode():
                 return self._run_csv_batches(csv_path)
+            if self._has_batch_config():
+                warnings.warn(
+                    "feature_batch_size/feature_batches are configured but enable_batch=False; "
+                    "FeatureValidationPipeline will read the CSV in full and ignore batch settings.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
             return self._run_dataframe(self._read_csv(csv_path))
         if not isinstance(data, pd.DataFrame):
             raise TypeError("data must be a pandas DataFrame when input_type='dataframe'.")
@@ -268,10 +279,18 @@ class FeatureValidationPipeline:
             return cfg.input_type
         return "dataframe" if isinstance(data, pd.DataFrame) else "csv"
 
-    def _is_batch_mode(self) -> bool:
+    def _has_batch_config(self) -> bool:
         cfg = self.config
         return bool(cfg.feature_batches) or cfg.feature_batch_size is not None
 
+    def _is_batch_mode(self) -> bool:
+        cfg = self.config
+        if not cfg.enable_batch:
+            return False
+        if not self._has_batch_config():
+            raise ValueError("enable_batch=True requires feature_batch_size or feature_batches.")
+        return True
+
     def _read_csv(self, csv_path: Path, usecols: list[str] | None = None, nrows: int | None = None) -> pd.DataFrame:
         if not csv_path.exists():
             raise FileNotFoundError(f"CSV file not found: {csv_path}")
@@ -325,6 +344,7 @@ class FeatureValidationPipeline:
                     cfg,
                     output_dir=str(batch_dir),
                     input_type="dataframe",
+                    enable_batch=False,
                     feature_batch_size=None,
                     feature_batches=None,
                     new_feature_cols=list(batch_features),