diff --git a/eland/series.py b/eland/series.py index 27c8fc5..d7a3b3b 100644 --- a/eland/series.py +++ b/eland/series.py @@ -20,7 +20,7 @@ import warnings import pandas as pd from eland import NDFrame -from eland.filter import NotFilter, Equal, Greater, Less, GreaterEqual, LessEqual, ScriptFilter +from eland.filter import NotFilter, Equal, Greater, Less, GreaterEqual, LessEqual, ScriptFilter, IsIn class Series(NDFrame): @@ -216,6 +216,12 @@ class Series(NDFrame): else: raise NotImplementedError(other, type(other)) + def isin(self, other): + if isinstance(other, list): + return IsIn(field=self.name, value=other) + else: + raise NotImplementedError(other, type(other)) + @property def ndim(self): """ diff --git a/eland/tests/dataframe/test_query_pytest.py b/eland/tests/dataframe/test_query_pytest.py index 8d0b06e..25cff71 100644 --- a/eland/tests/dataframe/test_query_pytest.py +++ b/eland/tests/dataframe/test_query_pytest.py @@ -43,11 +43,16 @@ class TestDataFrameQuery(TestData): assert_pandas_eland_frame_equal(pd_q4, ed_q4) - def test_query(self): + def test_simple_query(self): ed_flights = self.ed_flights() pd_flights = self.pd_flights() - #print(ed_flights.query('FlightDelayMin > 60').info_es()) + assert pd_flights.query('FlightDelayMin > 60').shape == \ + ed_flights.query('FlightDelayMin > 60').shape - print(pd_flights.query('FlightDelayMin > 60').shape) - print(ed_flights.query('FlightDelayMin > 60').shape) + def test_isin_query(self): + ed_flights = self.ed_flights() + pd_flights = self.pd_flights() + + assert pd_flights[pd_flights.OriginAirportID.isin(['LHR','SYD'])].shape == \ + ed_flights[ed_flights.OriginAirportID.isin(['LHR','SYD'])].shape