migrate to latest API

oap-project · Dec 8, 2024 · b442a3a · b442a3a
1 parent 19b010b
commit b442a3a
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/python/raydp/spark/__init__.py b/python/raydp/spark/__init__.py
@@ -31,4 +31,4 @@
   "spark_dataframe_to_ray_dataset",
   "ray_dataset_to_spark_dataframe",
   "from_spark_recoverable"
-]
+]
diff --git a/python/raydp/spark/dataset.py b/python/raydp/spark/dataset.py
@@ -254,4 +254,4 @@ def ray_dataset_to_spark_dataframe(spark: sql.SparkSession,
     elif isinstance(sample, pa.Table):
         return _convert_by_udf(spark, blocks, locations, schema)
     else:
-        raise RuntimeError("ray.to_spark only supports arrow type blocks")
+        raise RuntimeError("ray.to_spark only supports arrow type blocks")
diff --git a/python/raydp/tf/estimator.py b/python/raydp/tf/estimator.py
@@ -226,8 +226,10 @@ def fit(self,
                     label_cols = self._label_columns
                     if not isinstance(label_cols, list):
                         label_cols = [label_cols]
-                    preprocessor = Concatenator(output_column_name="features",
-                                                exclude=label_cols)
+                    preprocessor = Concatenator(
+                        columns=self._feature_columns - label_cols,
+                        output_column_name="features",
+                        )
                     train_loop_config["feature_columns"] = "features"
                     train_ds = preprocessor.transform(train_ds)
                     if evaluate_ds is not None: