SuperModelingFactory 0.5.4__tar.gz → 0.5.6__tar.gz

{supermodelingfactory-0.5.4 → supermodelingfactory-0.5.6}/Modeling_Tool/Core/ODPS_Tool.py

@@ -281,17 +281,31 @@ class ODPSRunner(object):
 
     @staticmethod
     def cre_table_schema(df, partition_name=None):
-
-        dtypes = df.dtypes
-        isnum_dtypes = [pd.api.types.is_numeric_dtype(x) for x in dtypes]
-        isint_dtypes = [pd.api.types.is_integer_dtype(x) for x in dtypes]
-
         table_columns = []
         table_partitions = []
-        for i in range(len(df.columns)):
-            col = df.columns[i]
-            col_type = "string" if not isnum_dtypes[i] else "float"
-            col_type = "bigint" if isint_dtypes[i] else col_type
+        for col, dtype in df.dtypes.items():
+            if pd.api.types.is_complex_dtype(dtype) or pd.api.types.is_timedelta64_dtype(dtype):
+                raise TypeError(
+                    f"Unsupported dtype for ODPS schema inference: column={col!r}, dtype={dtype!s}."
+                )
+            if pd.api.types.is_bool_dtype(dtype):
+                col_type = "boolean"
+            elif pd.api.types.is_integer_dtype(dtype):
+                col_type = "bigint"
+            elif pd.api.types.is_float_dtype(dtype):
+                col_type = "float" if str(dtype).lower() == "float32" else "double"
+            elif pd.api.types.is_datetime64_any_dtype(dtype):
+                col_type = "datetime"
+            elif (
+                pd.api.types.is_string_dtype(dtype)
+                or pd.api.types.is_object_dtype(dtype)
+                or isinstance(dtype, pd.CategoricalDtype)
+            ):
+                col_type = "string"
+            else:
+                raise TypeError(
+                    f"Unsupported dtype for ODPS schema inference: column={col!r}, dtype={dtype!s}."
+                )
             if col == partition_name:
                 table_partitions.append(Partition(name=col, type=col_type))
             else:
@@ -304,6 +318,28 @@ class ODPSRunner(object):
 
         return table_schema
 
+    @staticmethod
+    def _to_odps_records(df):
+        """Return Python records with pandas missing values converted to ODPS NULL."""
+        object_df = df.astype(object)
+        object_df = object_df.where(pd.notna(object_df), None)
+        return object_df.values.tolist()
+
+    def _execute_sql_blocking(self, sql):
+        """Execute DDL and wait until the remote ODPS instance has completed."""
+        execute_sql = getattr(self.o, "execute_sql", None)
+        if callable(execute_sql):
+            return execute_sql(sql)
+
+        run_sql = getattr(self.o, "run_sql", None)
+        if not callable(run_sql):
+            raise AttributeError("ODPS client must provide execute_sql() or run_sql().")
+        instance = run_sql(sql)
+        wait_for_success = getattr(instance, "wait_for_success", None)
+        if callable(wait_for_success):
+            wait_for_success()
+        return instance
+
     def upload_df(self, df, table_name, table_schema=None, partition=None, atomic=True):
         """上传数据集至mc中创建新表
 
@@ -328,9 +364,11 @@ class ODPSRunner(object):
             legacy behaviour, e.g. for tables whose downstream readers expect
             a specific object identity.
         """
+        upload_data = df.copy()
         if table_schema is None:
-            df.loc[:, "py_inserttime"] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
-            table_schema = self.cre_table_schema(df=df, partition_name=None)
+            upload_data.loc[:, "py_inserttime"] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+            table_schema = self.cre_table_schema(df=upload_data, partition_name=None)
+        records = self._to_odps_records(upload_data)
 
         if not atomic:
             # Legacy pre-0.4.2 non-atomic path — retained under an explicit
@@ -339,10 +377,10 @@ class ODPSRunner(object):
             t = self.o.create_table(table_name, table_schema)
             if bool(partition):
                 with t.open_writer(partition=partition, create_partition=True) as writer:
-                    writer.write(df.values.tolist())
+                    writer.write(records)
             else:
                 with t.open_writer() as writer:
-                    writer.write(df.values.tolist())
+                    writer.write(records)
             logger.info(f'<<<< 完成数据入表{table_name}: shape={df.shape} >>>>')
             return
 
@@ -360,10 +398,10 @@ class ODPSRunner(object):
         try:
             if bool(partition):
                 with tmp_table.open_writer(partition=partition, create_partition=True) as writer:
-                    writer.write(df.values.tolist())
+                    writer.write(records)
             else:
                 with tmp_table.open_writer() as writer:
-                    writer.write(df.values.tolist())
+                    writer.write(records)
         except Exception:
             try:
                 self.o.delete_table(tmp_name, if_exists=True)
@@ -379,7 +417,7 @@ class ODPSRunner(object):
         target_exists = self.o.exist_table(table_name)
         if target_exists:
             try:
-                self.o.run_sql(f"ALTER TABLE {table_name} RENAME TO {old_rename_name};")
+                self._execute_sql_blocking(f"ALTER TABLE {table_name} RENAME TO {old_rename_name};")
             except Exception:
                 # Rename of the live target failed — do not touch it. Drop the
                 # tmp table so we don't leak it, then re-raise.
@@ -392,13 +430,13 @@ class ODPSRunner(object):
                 raise
 
         try:
-            self.o.run_sql(f"ALTER TABLE {tmp_name} RENAME TO {target_rename_name};")
+            self._execute_sql_blocking(f"ALTER TABLE {tmp_name} RENAME TO {target_rename_name};")
         except Exception:
             # tmp -> target rename failed. If we already moved the original,
             # try to restore it so the caller is not left with a missing table.
             if target_exists:
                 try:
-                    self.o.run_sql(f"ALTER TABLE {old_name} RENAME TO {target_rename_name};")
+                    self._execute_sql_blocking(f"ALTER TABLE {old_name} RENAME TO {target_rename_name};")
                     logger.warning(
                         f"upload_df: swap of {tmp_name} -> {table_name} failed; restored original from {old_name}"
                     )
@@ -439,7 +477,9 @@ class ODPSRunner(object):
     def _rename_partition(self, table_name, source_partition, target_partition):
         src = _format_odps_partition_spec(_parse_odps_partition_spec(source_partition), quoted=True)
         dst = _format_odps_partition_spec(_parse_odps_partition_spec(target_partition), quoted=True)
-        self.o.run_sql(f"ALTER TABLE {table_name} PARTITION ({src}) RENAME TO PARTITION ({dst});")
+        self._execute_sql_blocking(
+            f"ALTER TABLE {table_name} PARTITION ({src}) RENAME TO PARTITION ({dst});"
+        )
 
     def insert_df(self, df, table_name, overwrite=True, partition=None, atomic=True):
         """将数据集插入至mc已存在的表中.
@@ -457,7 +497,9 @@ class ODPSRunner(object):
         """
         t = self.o.get_table(table_name)
         if "py_inserttime" in t.schema and "py_inserttime" not in df.columns:
+            df = df.copy()
             df.loc[:, "py_inserttime"] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+        records = self._to_odps_records(df)
 
         if bool(partition):
             if overwrite and atomic:
@@ -476,7 +518,7 @@ class ODPSRunner(object):
 
                 try:
                     with t.open_writer(partition=staging_partition, create_partition=True) as writer:
-                        writer.write(df.values.tolist())
+                        writer.write(records)
                     if target_exists:
                         self._rename_partition(table_name, partition, backup_partition)
                     self._rename_partition(table_name, staging_partition, partition)
@@ -510,10 +552,10 @@ class ODPSRunner(object):
             if overwrite:
                 t.delete_partition(partition, if_exists=True)
             with t.open_writer(partition=partition, create_partition=True) as writer:
-                writer.write(df.values.tolist())
+                writer.write(records)
         else:
             if overwrite:
                 t.truncate()
             with t.open_writer() as writer:
-                writer.write(df.values.tolist())
+                writer.write(records)
         logger.info('<<<< 完成数据入表: shape={0} >>>>'.format(df.shape))

{supermodelingfactory-0.5.4 → supermodelingfactory-0.5.6}/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",