From fb2774e7c73527d331026c48548ce5a400ae4068 Mon Sep 17 00:00:00 2001 From: Stephen Dodson Date: Thu, 14 Nov 2019 20:25:34 +0000 Subject: [PATCH] Adding Series.isin --- eland/series.py | 8 +++++++- eland/tests/dataframe/test_query_pytest.py | 13 +++++++++---- 2 files changed, 16 insertions(+), 5 deletions(-) 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