Added vaex extract elsewhere

askap-vast · Dec 18, 2024 · 800f76a · 800f76a
1 parent ae7711a
commit 800f76a
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/vasttools/pipeline.py b/vasttools/pipeline.py
@@ -553,7 +553,7 @@ def load_two_epoch_metrics(self) -> None:
                 measurement_pairs_df.pair_epoch_key, agg='count'
             )
 
-            pair_counts = pair_counts.to_pandas_df().rename(
+            pair_counts = pair_counts.extract().to_pandas_df().rename(
                 columns={'count': 'total_pairs'}
             ).set_index('pair_epoch_key')
         else:
@@ -976,7 +976,7 @@ def _filter_meas_pairs_df(
         )
 
         if not self._vaex_meas_pairs:
-            new_measurement_pairs = new_measurement_pairs.to_pandas_df()
+            new_measurement_pairs = new_measurement_pairs.extract().to_pandas_df()
 
         return new_measurement_pairs
 
@@ -1025,7 +1025,7 @@ def recalc_measurement_pairs_df(
 
         # convert a vaex measurements df to panads so an index can be set
         if isinstance(measurements_df, vaex.dataframe.DataFrame):
-            measurements_df = measurements_df[flux_cols].to_pandas_df()
+            measurements_df = measurements_df[flux_cols].extract().to_pandas_df()
         else:
             measurements_df = measurements_df.loc[:, flux_cols].copy()
 
@@ -1211,7 +1211,7 @@ def recalc_sources_df(
         )
 
         # Switch to pandas at this point to perform join
-        sources_df = sources_df.to_pandas_df().set_index('source')
+        sources_df = sources_df.extract().to_pandas_df().set_index('source')
 
         sources_df = sources_df.join(sources_df_fluxes)
 
@@ -1847,7 +1847,7 @@ def run_two_epoch_analysis(
                 (pairs_df[vs_label] > vs) & (pairs_df[m_abs_label] > m)
             ]
 
-            candidate_pairs = candidate_pairs.to_pandas_df()
+            candidate_pairs = candidate_pairs.extract().to_pandas_df()
 
         else:
             candidate_pairs = pairs_df.loc[