BUG: Fix StringAccessor __getitem__ error (#588)

xorbitsai · Jul 10, 2023 · b8a1b4d · b8a1b4d
1 parent 1aafbe0
commit b8a1b4d
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/python/xorbits/core/adapter.py b/python/xorbits/core/adapter.py
@@ -375,6 +375,9 @@ def __getattr__(self, item):
             # class variable
             return from_mars(attr)
 
+    def __getitem__(self, item):
+        return from_mars(self._mars_obj[to_mars(item)])
+
 
 def to_mars(inp: Union[DataRef, Tuple, List, Dict]):
     """

diff --git a/python/xorbits/core/tests/test_execution.py b/python/xorbits/core/tests/test_execution.py
@@ -324,3 +324,16 @@ def test_interactive_execution(setup, ip):
     ip.run_cell(raw_cell="df + 3")
     r = ip.run_cell(raw_cell="xorbits.core.execution.need_to_execute(_)")
     assert r.result
+
+
+def test_getitem(setup):
+    import pandas as pd
+
+    from ... import pandas as xpd
+
+    data = {"a": ["abc", "def"]}
+    df = pd.DataFrame(data)
+    xdf = xpd.DataFrame(data)
+    result = xdf[xdf["a"].str[0] == "a"]
+    expected = df[df["a"].str[0] == "a"]
+    pd.testing.assert_frame_equal(result.to_pandas(), expected)