fix: make _data code and unit tests pandas-3 compatible

mostlyai/sdk/_data/dtype.py

@@ -184,6 +184,16 @@ def coerce(self, data: pd.Series) -> pd.Series:
         return bool_coerce(data)
 
 
+def _to_numeric_or_same(data: pd.Series, downcast: str) -> pd.Series:
+    # pandas 3.0 dropped errors="ignore" for pd.to_numeric; emulate the
+    # previous "return input unchanged on failure" contract that our cast()
+    # callers relied on.
+    try:
+        return pd.to_numeric(data, errors="raise", downcast=downcast, dtype_backend="pyarrow")
+    except (ValueError, TypeError):
+        return data
+
+
 class VirtualInteger(VirtualDType):
     def get_encompass_data_fallback_dtypes(self) -> list["DType"]:
         return [VirtualFloat, VirtualVarchar]
@@ -192,15 +202,15 @@ def coerce(self, data: pd.Series) -> pd.Series:
         return int_coerce(data)
 
     def cast(self, data: pd.Series) -> pd.Series:
-        return pd.to_numeric(data, errors="ignore", downcast="integer", dtype_backend="pyarrow")
+        return _to_numeric_or_same(data, downcast="integer")
 
 
 class VirtualFloat(VirtualDType):
     def coerce(self, data: pd.Series) -> pd.Series:
         return float_coerce(data)
 
     def cast(self, data: pd.Series) -> pd.Series:
-        return pd.to_numeric(data, errors="ignore", downcast="float", dtype_backend="pyarrow")
+        return _to_numeric_or_same(data, downcast="float")
 
 
 class VirtualDate(VirtualDType):
@@ -468,8 +478,11 @@ def time_coerce(s: pd.Series) -> pd.Series:
     n_coerced = n_na_1 - n_na_0
     if n_coerced > 0:
         _LOG.warning(f"{n_coerced} values coerced during datetime_coerce")
-    # map to "object"
-    return s.astype("object")
+    # map to "object", and normalize pd.NaT to pd.NA so the result is uniform
+    # regardless of whether any row survived coercion (pandas 3 is stricter
+    # about distinguishing NaT vs NA under the object dtype)
+    s = s.astype("object")
+    return s.where(~s.isna(), pd.NA)
 
 
 class PandasDType(WrappedDType):

mostlyai/sdk/_data/file/table/json.py

@@ -125,13 +125,17 @@ def write_data(self, df: pd.DataFrame, if_exists: str = "append", **kwargs):
             if is_date_dtype(df[c]):
                 df[c] = df[c].dt.strftime("%Y-%m-%d")
             elif is_timestamp_dtype(df[c]):
-                # we need to strip off any milliseconds
+                # we need to strip off any microseconds (pyarrow-backed
+                # timestamps include them in strftime output regardless of
+                # format string; use str.slice rather than str[:-7] because
+                # pandas 3.0 removed StringMethods.__getitem__ for pyarrow
+                # extension arrays)
                 df[c] = (
                     df[c]
                     .dt.tz_localize(None)
                     .astype("timestamp[us][pyarrow]")
                     .dt.strftime("%Y-%m-%d %H:%M:%S")
-                    .str[:-7]
+                    .str.slice(stop=-7)
                 )
         mode = self.handle_if_exists(if_exists)
         df.to_json(self.container.path_str, orient="records", lines=True, mode=mode)

tests/_data/unit/test_finalize_generation.py

@@ -48,8 +48,10 @@ def test_finalize_table_generation(tmp_path, tgt_data):
     gen_data = tgt_table.read_data()
     gen_data_path = tmp_path / "gen"
     gen_data_path.mkdir(exist_ok=True, parents=True)
-    # make 3 partitions
-    get_data_partitioned = np.array_split(gen_data, 3)
+    # make 3 partitions (np.array_split returns ndarrays for DataFrames on
+    # numpy 2 / pandas 3, so split via positional indices to keep DataFrames)
+    split_points = np.array_split(np.arange(len(gen_data)), 3)
+    get_data_partitioned = [gen_data.iloc[idx] for idx in split_points]
     # save each part into a separate parquet file
     for i, df_part in enumerate(get_data_partitioned):
         part_fn = gen_data_path / f"part.{i:06}.parquet"
@@ -103,7 +105,7 @@ def test_export_data_to_excel(tmp_path):
     excel_file_path = tmp_path / "synthetic-samples.xlsx"
     xls = pd.ExcelFile(excel_file_path)
 
-    expected_values = pd.DataFrame({"A": [0, 1, 2, 3, 4, 5, 6, 7, 8], "X": ["x"] * 9, "N": [None] * 9})
+    expected_values = pd.DataFrame({"A": [0, 1, 2, 3, 4, 5, 6, 7, 8], "X": ["x"] * 9})
     # Iterate over each sheet
     for sheet_name in xls.sheet_names:
         df = pd.read_excel(xls, sheet_name=sheet_name)
@@ -115,7 +117,11 @@ def test_export_data_to_excel(tmp_path):
             # Check the expected values
             df.sort_values(by="A", ascending=True, inplace=True)
             df.reset_index(drop=True, inplace=True)
-            pd.testing.assert_frame_equal(df, expected_values, check_dtype=False)
+            # the "N" column is all-missing; compare via isna() to stay
+            # agnostic of how read_excel materializes nulls (NaN vs None vs NA)
+            # across pandas versions
+            assert df["N"].isna().all()
+            pd.testing.assert_frame_equal(df.drop(columns=["N"]), expected_values, check_dtype=False)
 
     # Check for the sheet names
     expected_sheets = ["_TOC_", "A", "B", "C"]

tests/_data/unit/test_push.py

@@ -55,7 +55,10 @@ def source_table_partitioned(tmp_path):
     source_data_path = tmp_path / "source"
     source_data_path.mkdir(exist_ok=True, parents=True)
     n_partitions = 10
-    get_data_partitioned = np.array_split(source_data, n_partitions)
+    # np.array_split returns ndarrays for DataFrames on numpy 2 / pandas 3,
+    # so split via positional indices to keep DataFrames
+    split_points = np.array_split(np.arange(len(source_data)), n_partitions)
+    get_data_partitioned = [source_data.iloc[idx] for idx in split_points]
     # save each part into a separate parquet file
     for i, df_part in enumerate(get_data_partitioned):
         part_fn = source_data_path / f"part.{i:06}.parquet"