Fixed Object type for DataFrame, take 2

Author: ankane

CHANGELOG.md

@@ -4,7 +4,7 @@
 - Added `storage_options` and `retries` options to `sink_ipc` method
 - Added experimental support for Iceberg
 - Added experimental `cast_options` option to `scan_parquet` method
-- Fixed `Object` type for `Series`
+- Fixed `Object` type
 
 ## 0.21.1 (2025-08-18)
 

ext/polars/src/dataframe/mod.rs

@@ -4,10 +4,14 @@ mod general;
 mod io;
 mod serde;
 
+use magnus::{DataTypeFunctions, TypedData, gc};
 use polars::prelude::*;
 use std::cell::RefCell;
 
-#[magnus::wrap(class = "Polars::RbDataFrame")]
+use crate::series::mark_series;
+
+#[derive(TypedData)]
+#[magnus(class = "Polars::RbDataFrame", mark)]
 pub struct RbDataFrame {
     pub df: RefCell<DataFrame>,
 }
@@ -25,3 +29,23 @@ impl RbDataFrame {
         }
     }
 }
+
+impl DataTypeFunctions for RbDataFrame {
+    fn mark(&self, marker: &gc::Marker) {
+        // this is really, really not ideal, as objects will not be marked if unable to borrow
+        // currently, this should only happen for write_* methods,
+        // which should refuse to write Object datatype, and therefore be safe,
+        // since GC will not have a chance to run
+        if let Ok(df) = self.df.try_borrow() {
+            for column in df.get_columns() {
+                if let DataType::Object(_) = column.dtype() {
+                    match column {
+                        Column::Series(s) => mark_series(marker, s),
+                        Column::Partitioned(s) => mark_series(marker, s.partitions()),
+                        Column::Scalar(s) => mark_series(marker, &s.as_single_value_series()),
+                    }
+                }
+            }
+        }
+    }
+}

ext/polars/src/series/mod.rs

@@ -50,16 +50,19 @@ pub fn to_rbseries(s: Vec<Column>) -> RArray {
     )
 }
 
-impl DataTypeFunctions for RbSeries {
-    fn mark(&self, marker: &gc::Marker) {
-        let series = self.series.borrow();
-        if let DataType::Object(_) = series.dtype() {
-            for i in 0..series.len() {
-                let obj: Option<&ObjectValue> = series.get_object(i).map(|any| any.into());
-                if let Some(o) = obj {
-                    marker.mark(o.inner);
-                }
+pub fn mark_series(marker: &gc::Marker, series: &Series) {
+    if let DataType::Object(_) = series.dtype() {
+        for i in 0..series.len() {
+            let obj: Option<&ObjectValue> = series.get_object(i).map(|any| any.into());
+            if let Some(o) = obj {
+                marker.mark(o.inner);
             }
         }
     }
 }
+
+impl DataTypeFunctions for RbSeries {
+    fn mark(&self, marker: &gc::Marker) {
+        mark_series(marker, &self.series.borrow());
+    }
+}

test/types_test.rb

@@ -193,14 +193,35 @@ def test_object
 
     df = Polars::DataFrame.new({a: [Object.new]})
     GC.start
-    # TODO fix
-    # assert df.inspect
+    assert df.inspect
     assert df.to_a
 
     df = Polars::DataFrame.new({a: [Object.new, Object.new]})
     GC.start
-    # TODO fix
-    # assert df.inspect
+    assert df.inspect
+    assert df.to_a
+  end
+
+  # only fully tested in stress mode
+  def test_object_stress
+    df = Polars::DataFrame.new({a: [Object.new]})
+    io = StringIO.new
+    assert_raises(Polars::ComputeError) do
+      df.write_csv(io)
+    end
+    assert_raises(Polars::ComputeError) do
+      df.write_parquet(io)
+    end
+    assert_raises(Polars::ComputeError) do
+      df.write_json(io)
+    end
+    assert_raises(Polars::ComputeError) do
+      df.write_ndjson(io)
+    end
+    assert_raises(Polars::ComputeError) do
+      df.write_ipc(io)
+    end
+    assert df.inspect
     assert df.to_a
   end