Merge pull request #1 from Keunyoung-Jung/dev

✅  Update : version 0.0.4

Author: keyog0

ParquetLoader/__init__.py

@@ -1,6 +1,6 @@
 __name__ = 'parquet-loader'
 __description__ = 'Parquet file Load and Read from minio & S3'
-__version__ = '0.0.3'
+__version__ = '0.0.4'
 __url__ = 'https://github.com/Keunyoung-Jung/ParquetLoader'
 __download_url__ = 'https://github.com/Keunyoung-Jung/ParquetLoader'
 __install_requires__ = [

ParquetLoader/loader.py

@@ -14,7 +14,8 @@ def __init__(self,
                  random_seed : int = int((time() - int(time()))*100000),
                  columns : list = None,
                  depth : int = 0,
-                 std_out: bool = True
+                 std_out: bool = True,
+                 filters: list = None
                  ):
         self.chunk_size = chunk_size
         self.cache = None
@@ -30,6 +31,8 @@ def __init__(self,
         self.fp_obj = None
         self.depth = depth
         self.random_seed = random_seed
+        self.filters = filters
+        self.check_filter()
 
         if root_path == '.' :
             self.root_path = os.getcwd()
@@ -84,7 +87,9 @@ def close(self):
         else:
             raise RuntimeError("generator ignored GeneratorExit")
     def generator(self):
-        for df in self.fp_obj.iter_row_groups(columns=self.select_columns):
+        for df in self.fp_obj.iter_row_groups(filters=self.filters,columns=self.select_columns):
+            if self.filters != None :
+                df = self.filtering(df)
             if self.dataset is None :
                 self.dataset = df
             else :
@@ -110,4 +115,51 @@ def generator(self):
         yield self.dataset
         self.dataset = None
         if self.std_out :
-            print(self.counter,'data loaded complete!',end='\n')
+            print(self.counter,'data loaded complete!',end='\n')
+            
+    def filtering(self,df) :
+        op = ''
+        df_store = []
+        for or_part in self.filters:
+            tmp_df = df.copy()
+            for and_part in or_part :
+                col = and_part[0]
+                op = and_part[1]
+                val = and_part[2]
+                if op == '==' or op == '=' :
+                    tmp_df = tmp_df[tmp_df[col] == val]
+                elif op == '>' :
+                    tmp_df = tmp_df[tmp_df[col] > val]
+                elif op == '>=' :
+                    tmp_df = tmp_df[tmp_df[col] >= val]
+                elif op == '<' :
+                    tmp_df = tmp_df[tmp_df[col] < val]
+                elif op == '<=' :
+                    tmp_df = tmp_df[tmp_df[col] <= val]
+                elif op == '!=' :
+                    tmp_df = tmp_df[tmp_df[col] != val]
+                elif op == 'in' :
+                    tmp_df = tmp_df[tmp_df[col] in val]
+                elif op == 'not in' :
+                    tmp_df = tmp_df[tmp_df[col] not in val]
+            df_store.append(tmp_df)
+        concat_df = pd.concat(df_store)
+        return concat_df.drop_duplicates()
+    
+    def check_filter(self) :
+        try :
+            op = ''
+            if self.filters != None :
+                for or_part in self.filters :
+                    for and_part in or_part :
+                        op = and_part[1]
+                        if len(and_part) != 3 :
+                            raise  IndexError
+                        if op not in ['==','=','>','>=','<','<=','!=','in','not in'] :
+                            raise ValueError
+        except ValueError :
+            print("ValueError :",f'"{op}" is wrong operator')
+            exit()
+        except IndexError :
+            print("IndexError :",f'{self.filters} invalid filter')
+            exit()

ParquetLoader/s3.py

@@ -16,7 +16,8 @@ def __init__(self,
                  random_seed : int = int((time() - int(time()))*100000),
                  columns : list = None,
                  depth : int = 0,
-                 std_out: bool = True
+                 std_out: bool = True,
+                 filters: list = None
                  ):
         self.client_dict = None
         if s3_endpoint_url != '' or \
@@ -33,7 +34,8 @@ def __init__(self,
             random_seed=random_seed,
             columns=columns,
             depth=depth,
-            std_out=std_out
+            std_out=std_out,
+            filters=filters
             )
 
     def initialize(self) :

README.md

@@ -67,7 +67,8 @@ dl = DataLoader(
     random_seed : int = int((time() - int(time()))*100000),
     columns : list = None,
     depth : int = 0,
-    std_out: bool = True
+    std_out: bool = True,
+    filters: list = None
     )
 ```
 * `chunk_size`
@@ -94,6 +95,9 @@ dl = DataLoader(
 * `std_out`
     * default : True
     * You can turn off output.
+* `filters`
+    * It is used when you want get filtered dataframe, It must use 2 dim list
+    * example : `[[("column","==",10)]]`
 
 ### 4.1. Select Columns
 `columns` param is taken as a list.
@@ -149,3 +153,37 @@ dl = S3Loader(
 * `s3_access_key` and `s3_secret_key`
     * you can set s3_access_key and s3_secret_key, but I don't recommend using it
     * it is recommended to use environment variables.
+
+## 7. Get Filtered Dataframe
+It is used when you want get filtered dataframe, It must use 2 dim list
+It is built with a two-dimensional list construction. (Equal fastparquet filter)
+```python
+dl = S3Loader(
+    bucket = 'test',
+    folder = 'data',
+    filters = [[[("col1",">",10)]]]
+    )
+```
+The first list consists of an OR operation.
+```python
+# col > 10 or col2 in ["children","kids"]
+filters = [
+    [("col1",">",10)],
+    ["col2","in",["children","kids"]]
+    ] 
+```
+The second list consists of an AND operation.
+```python
+# col > 10 and col2 == "male"
+filters = [
+    [("col1",">",10),("col2","==","male")]
+    ] 
+```
+You can also mix the two to make a filter.
+```python
+# (col > 10 and col2 == "male") or col3 in ["children","kids"]
+filters = [
+    [("col1",">",10),("col2","==","male")],
+    ["col3","in",["children","kids"]]
+    ]
+```

tests/filters.py

@@ -0,0 +1,70 @@
+from ParquetLoader import S3Loader
+
+def s3_data_loading():
+    sl = S3Loader(
+        chunk_size=10000,
+        s3_endpoint_url='http://localhost:9000',
+        s3_access_key='minio',
+        s3_secret_key='minio123',
+        bucket='test',
+        folder="data",
+        shuffle=False)
+    print(sl.schema)
+    for df in sl :
+        print(df.tail(10))
+        break
+    
+def filter_data():
+    sl = S3Loader(
+        chunk_size=10000,
+        s3_endpoint_url='http://localhost:9000',
+        s3_access_key='minio',
+        s3_secret_key='minio123',
+        bucket='test',
+        folder="data",
+        filters=[
+            [("tf",">",10),("tf","<",13)],
+            [("tf","==",1)]
+            ],
+        shuffle=False)
+    print(sl.schema)
+    for df in sl :
+        print(df.tail(10))
+        break
+
+def wrong_op():
+    sl = S3Loader(
+        chunk_size=10000,
+        s3_endpoint_url='http://localhost:9000',
+        s3_access_key='minio',
+        s3_secret_key='minio123',
+        bucket='test',
+        folder="data",
+        filters=[[("tf","%",10)]],
+        shuffle=False)
+    print(sl.schema)
+    for df in sl :
+        print(df.tail(10))
+        break
+
+def wrong_filter():
+    sl = S3Loader(
+        chunk_size=10000,
+        s3_endpoint_url='http://localhost:9000',
+        s3_access_key='minio',
+        s3_secret_key='minio123',
+        bucket='test',
+        folder="data",
+        filters=[
+            [("tf",">")]],
+        shuffle=False)
+    print(sl.schema)
+    for df in sl :
+        print(df.tail(10))
+        break
+        
+if __name__ == '__main__':
+    s3_data_loading()
+    filter_data()
+    wrong_op()
+    wrong_filter()