POC: pass date_unit to values_for_json

pandas/_libs/src/vendored/ujson/python/objToJSON.c

@@ -720,7 +720,25 @@ void PdBlock_iterBegin(JSOBJ _obj, JSONTypeContext *tc) {
         return;
     }
 
-    arrays = get_sub_attr(obj, "_mgr", "column_arrays");
+    NPY_DATETIMEUNIT dunit = ((PyObjectEncoder *)tc)->datetimeUnit;
+    char *date_unit;
+    if (dunit == NPY_FR_s) {
+        date_unit = "s";
+    } else if (dunit == NPY_FR_ms) {
+        date_unit = "ms";
+    } else if (dunit == NPY_FR_us) {
+        date_unit = "us";
+    } else if (dunit == NPY_FR_ns) {
+        date_unit = "ns";
+    }
+
+    PyObject *mgr = PyObject_GetAttrString(obj, "_mgr");
+    if (mgr == NULL) {
+        return;
+    }
+    arrays = PyObject_CallMethod(mgr, "column_arrays", "%s", date_unit);
+    Py_DECREF(mgr);
+
     if (!arrays) {
         GET_TC(tc)->iterNext = NpyArr_iterNextNone;
         return;

pandas/core/arrays/base.py

@@ -1736,7 +1736,7 @@ def _reduce(
     # Non-Optimized Default Methods; in the case of the private methods here,
     #  these are not guaranteed to be stable across pandas versions.
 
-    def _values_for_json(self) -> np.ndarray:
+    def _values_for_json(self, date_unit: str) -> np.ndarray:
         """
         Specify how to render our entries in to_json.
 

pandas/core/arrays/datetimelike.py

@@ -2202,11 +2202,25 @@ def _with_freq(self, freq) -> Self:
     # --------------------------------------------------------------
     # ExtensionArray Interface
 
-    def _values_for_json(self) -> np.ndarray:
+    def _values_for_json(self, date_unit: str) -> np.ndarray:
         # Small performance bump vs the base class which calls np.asarray(self)
-        if isinstance(self.dtype, np.dtype):
-            return self._ndarray
-        return super()._values_for_json()
+        if self.dtype.kind == "M":
+            if date_unit == "s":
+                return self.strftime("%Y-%M-%D %h-%m-%s")
+            elif date_unit == "ms":
+                raise NotImplementedError
+            elif date_unit == "us":
+                return self.strftime("%Y-%M-%D %h-%m-%s.%f")
+            elif date_unit == "ns":
+                return self.astype(str)
+            else:
+                raise NotImplementedError
+        elif self.dtype.kind == "m":
+            raise NotImplementedError
+        else:
+            return super()._values_for_json(date_unit)
+
+        return self._ndarray
 
     def factorize(
         self,

pandas/core/internals/array_manager.py

@@ -691,8 +691,7 @@ def iget_values(self, i: int) -> ArrayLike:
         """
         return self.arrays[i]
 
-    @property
-    def column_arrays(self) -> list[ArrayLike]:
+    def column_arrays(self, date_unit) -> list[np.ndarray]:
         """
         Used in the JSON C code to access column arrays.
         """

pandas/core/internals/managers.py

@@ -993,8 +993,7 @@ def iget_values(self, i: int) -> ArrayLike:
         values = block.iget(self.blklocs[i])
         return values
 
-    @property
-    def column_arrays(self) -> list[np.ndarray]:
+    def column_arrays(self, date_unit) -> list[np.ndarray]:
         """
         Used in the JSON C code to access column arrays.
         This optimizes compared to using `iget_values` by converting each
@@ -1008,7 +1007,7 @@ def column_arrays(self) -> list[np.ndarray]:
 
         for blk in self.blocks:
             mgr_locs = blk._mgr_locs
-            values = blk.array_values._values_for_json()
+            values = blk.array_values._values_for_json(date_unit)
             if values.ndim == 1:
                 # TODO(EA2D): special casing not needed with 2D EAs
                 result[mgr_locs[0]] = values