Merge branch 'main' into shuowei-time-series-bike

Author: shuoweil

bigframes/core/compile/polars/compiler.py

@@ -152,6 +152,11 @@ def _(
                 value = None
             if expression.dtype is None:
                 return pl.lit(None)
+
+            # Polars lit does not handle pandas timedelta well at v1.36
+            if isinstance(value, pd.Timedelta):
+                value = value.to_pytimedelta()
+
             return pl.lit(value, _bigframes_dtype_to_polars_dtype(expression.dtype))
 
         @compile_expression.register

bigframes/ml/base.py

@@ -248,6 +248,13 @@ def fit(
     ) -> _T:
         return self._fit(X, y)
 
+    def fit_predict(
+        self: _T,
+        X: utils.ArrayType,
+        y: Optional[utils.ArrayType] = None,
+    ) -> _T:
+        return self.fit(X).predict(X)
+
 
 class RetriableRemotePredictor(BaseEstimator):
     def _predict_and_retry(

notebooks/generative_ai/bq_dataframes_llm_kmeans.ipynb

@@ -1736,7 +1736,7 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "venv (3.10.14)",
    "language": "python",
    "name": "python3"
   },
@@ -1750,7 +1750,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.10.14"
   }
  },
  "nbformat": 4,

tests/unit/test_series_polars.py

@@ -5109,3 +5109,18 @@ def test_series_item_with_empty(session):
 
     with pytest.raises(ValueError, match=re.escape(expected_message)):
         bf_s_empty.item()
+
+
+def test_series_dt_total_seconds(scalars_df_index, scalars_pandas_df_index):
+    bf_result = scalars_df_index["duration_col"].dt.total_seconds().to_pandas()
+
+    pd_result = scalars_pandas_df_index["duration_col"].dt.total_seconds()
+
+    # Index will be object type in pandas, string type in bigframes, but same values
+    pd.testing.assert_series_equal(
+        bf_result,
+        pd_result,
+        check_index_type=False,
+        # bigframes uses Float64, newer pandas may use double[pyarrow]
+        check_dtype=False,
+    )

third_party/bigframes_vendored/sklearn/cluster/_kmeans.py

@@ -115,6 +115,26 @@ def predict(
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+    def fit_predict(
+        self,
+        X,
+        y=None,
+    ):
+        """Compute cluster centers and predict cluster index for each sample.
+
+        Convenience method; equivalent to calling fit(X) followed by predict(X).
+
+        Args:
+            X (bigframes.dataframe.DataFrame or bigframes.series.Series or pandas.core.frame.DataFrame or pandas.core.series.Series):
+                DataFrame of shape (n_samples, n_features). Training data.
+            y (default None):
+                Not used, present here for API consistency by convention.
+
+        Returns:
+            bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted labels.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
     def score(
         self,
         X,

third_party/bigframes_vendored/sklearn/decomposition/_mf.py

@@ -94,3 +94,23 @@ def predict(self, X):
         Returns:
             bigframes.dataframe.DataFrame: Predicted DataFrames."""
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+    def fit_predict(
+        self,
+        X,
+        y=None,
+    ):
+        """Fit the model with X and generate a predicted rating for every user-item row combination for a matrix factorization model. on X.
+
+        Convenience method; equivalent to calling fit(X) followed by predict(X).
+
+        Args:
+            X (bigframes.dataframe.DataFrame or bigframes.series.Series or pandas.core.frame.DataFrame or pandas.core.series.Series):
+                DataFrame of shape (n_samples, n_features). Training data.
+            y (default None):
+                Not used, present here for API consistency by convention.
+
+        Returns:
+            bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted labels.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)

third_party/bigframes_vendored/sklearn/decomposition/_pca.py

@@ -101,6 +101,26 @@ def predict(self, X):
             bigframes.dataframe.DataFrame: Predicted DataFrames."""
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+    def fit_predict(
+        self,
+        X,
+        y=None,
+    ):
+        """Fit the model with X and apply the dimensionality reduction on X.
+
+        Convenience method; equivalent to calling fit(X) followed by predict(X).
+
+        Args:
+            X (bigframes.dataframe.DataFrame or bigframes.series.Series or pandas.core.frame.DataFrame or pandas.core.series.Series):
+                DataFrame of shape (n_samples, n_features). Training data.
+            y (default None):
+                Not used, present here for API consistency by convention.
+
+        Returns:
+            bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted labels.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
     @property
     def components_(self):
         """Principal axes in feature space, representing the directions of maximum variance in the data.