Parquet: Handle reading data from multiple sqvers correctly

Reads that involved reading from multiple different sqvers versions was
not handled correctly causing such reads to return incorrect data. This
patch addresses that. Essentially, we need to read each sqvers folder
individually and then join all the dataframes together.

Signed-off-by: Dinesh Dutt <[email protected]>

Author: ddutt

suzieq/db/parquet/parquetdb.py

@@ -88,6 +88,9 @@ def read(self, table_name: str, data_format: str,
 
         folder = self._get_table_directory(table_name, False)
 
+        if not all(x in fields for x in key_fields):
+            raise ValueError('Key fields MUST be included in columns list')
+
         if addnl_filter:
             # This is for special cases that are specific to an object
             query_str = addnl_filter
@@ -108,7 +111,7 @@ def read(self, table_name: str, data_format: str,
 
         # If requesting a specific version of the data, handle that diff too
         sqvers = kwargs.pop('sqvers', None)
-        datasets = []
+        final_df = pd.DataFrame()
         try:
             dirs = Path(folder)
             try:
@@ -123,64 +126,43 @@ def read(self, table_name: str, data_format: str,
                         if vers > max_vers:
                             max_vers = vers
 
-                    datasets.append(ds.dataset(elem, format='parquet',
-                                               partitioning='hive'))
+                    dataset = ds.dataset(elem, format='parquet',
+                                         partitioning='hive')
+
+                    if not dataset:
+                        continue
+
+                    tmp_df = self._process_dataset(dataset, namespace, start,
+                                                   end, fields, merge_fields,
+                                                   query_str, **kwargs)
+                    if not tmp_df.empty:
+                        final_df = pd.concat([final_df, tmp_df])
             except FileNotFoundError:
                 pass
             except Exception as e:
                 raise e
 
-            # Now find the exact set of files we need to go over
+            # Now operate on the coalesced data set
             cp_dataset = self._get_cp_dataset(table_name, need_sqvers, sqvers,
                                               view, start, end)
             if cp_dataset:
-                datasets.append(cp_dataset)
-
-            if not datasets:
-                datasets = [ds.dataset(folder, format='parquet',
-                                       partitioning='hive')]
-
-            # Build the filters for predicate pushdown
-            master_schema = self._build_master_schema(datasets)
-
-            avail_fields = list(filter(lambda x: x in master_schema.names,
-                                       fields))
-            filters = self.build_ds_filters(
-                start, end, master_schema, merge_fields=merge_fields, **kwargs)
-
-            filtered_datasets = self._get_filtered_fileset(
-                datasets, namespace)
-
-            final_df = filtered_datasets \
-                .to_table(filter=filters, columns=avail_fields) \
-                .to_pandas(self_destruct=True) \
-                .query(query_str) \
-                .sort_values(by='timestamp')
-
-            if merge_fields:
-                # These are key fields that need to be set right before we do
-                # the drop duplicates to avoid missing out all the data
-                for field in merge_fields:
-                    newfld = merge_fields[field]
-                    if (field in final_df.columns and
-                            newfld in final_df.columns):
-                        final_df[newfld] = np.where(final_df[newfld],
-                                                    final_df[newfld],
-                                                    final_df[field])
-                    elif (field in final_df.columns and
-                          newfld not in final_df.columns):
-                        final_df = final_df.rename(columns={field: newfld})
+                tmp_df = self._process_dataset(cp_dataset, namespace, start,
+                                               end, fields, merge_fields,
+                                               query_str, **kwargs)
+                if not tmp_df.empty:
+                    final_df = pd.concat([final_df, tmp_df])
 
             # Because of how coalescing works, we can have multiple duplicated
             # entries with same timestamp. Remove them
-            dupts_keys = key_fields + ['timestamp']
-            final_df = final_df.set_index(dupts_keys) \
-                .query('~index.duplicated(keep="last")') \
-                .reset_index()
-            if (not final_df.empty and (view == 'latest') and
-                    all(x in final_df.columns for x in key_fields)):
-                final_df = final_df.set_index(key_fields) \
-                    .query('~index.duplicated(keep="last")')
+            if not final_df.empty:
+                final_df = final_df.sort_values(by=['timestamp'])
+                dupts_keys = key_fields + ['timestamp']
+                final_df = final_df.set_index(dupts_keys) \
+                    .query('~index.duplicated(keep="last")') \
+                    .reset_index()
+                if not final_df.empty and (view == 'latest'):
+                    final_df = final_df.set_index(key_fields) \
+                        .query('~index.duplicated(keep="last")')
         except (pa.lib.ArrowInvalid, OSError):
             return pd.DataFrame(columns=fields)
 
@@ -483,6 +465,46 @@ def _get_avail_sqvers(self, table_name: str, coalesced: bool) -> List[str]:
 
         return sqvers_list
 
+    def _process_dataset(self, dataset: ds.Dataset, namespace: List[str],
+                         start: str, end: str, fields: List[str],
+                         merge_fields: List[str], query_str: str,
+                         **kwargs) -> pd.DataFrame:
+        '''Process provided dataset and return a pandas DF'''
+
+        # Build the filters for predicate pushdown
+        master_schema = dataset.schema
+
+        avail_fields = [f for f in fields if f in master_schema.names]
+        filters = self.build_ds_filters(
+            start, end, master_schema, merge_fields=merge_fields,
+            **kwargs)
+
+        filtered_dataset = self._get_filtered_fileset(dataset, namespace)
+
+        if not filtered_dataset.files:
+            return pd.DataFrame()
+
+        tmp_df = filtered_dataset \
+            .to_table(filter=filters, columns=avail_fields) \
+            .to_pandas(self_destruct=True) \
+            .query(query_str)
+
+        if merge_fields and not tmp_df.empty:
+            # These are key fields that need to be set right before we do
+            # the drop duplicates to avoid missing out all the data
+            for field in merge_fields:
+                newfld = merge_fields[field]
+                if (field in tmp_df.columns and
+                        newfld in tmp_df.columns):
+                    tmp_df[newfld] = np.where(tmp_df[newfld],
+                                              tmp_df[newfld],
+                                              tmp_df[field])
+                elif (field in tmp_df.columns and
+                      newfld not in tmp_df.columns):
+                    tmp_df = tmp_df.rename(columns={field: newfld})
+
+        return tmp_df
+
     def _get_cp_dataset(self, table_name: str, need_sqvers: bool,
                         sqvers: str, view: str, start_time: float,
                         end_time: float) -> ds.dataset:
@@ -608,7 +630,7 @@ def _build_master_schema(self, datasets: list) -> pa.lib.Schema:
 
         return msch
 
-    def _get_filtered_fileset(self, datasets: list, namespace: list) -> ds:
+    def _get_filtered_fileset(self, dataset: ds, namespace: list) -> ds:
         """Filter the dataset based on the namespace
 
         We can use this method to filter out namespaces and hostnames based
@@ -620,48 +642,44 @@ def _get_filtered_fileset(self, datasets: list, namespace: list) -> ds:
         Returns:
             ds: pyarrow dataset of only the files that match filter
         """
-        filelist = []
-        for dataset in datasets:
-            if not namespace:
-                filelist.extend(dataset.files)
-                continue
+        if not namespace:
+            return dataset
 
-            notlist = [x for x in namespace
+        # Exclude not and regexp operators as they're handled elsewhere.
+        excluded_ns = [x for x in namespace
                        if x.startswith('!') or x.startswith('~!')]
-            if notlist != namespace:
-                newns = [x for x in namespace
-                         if not x.startswith('!') and not x.startswith('~!')]
-            else:
-                newns = namespace
-            ns_filelist = []
-            chklist = dataset.files
-            for ns in newns or []:
+        if excluded_ns != namespace:
+            ns_filters = [x for x in namespace
+                          if not x.startswith('!') and not x.startswith('~!')]
+        else:
+            ns_filters = namespace
+        ns_filelist = []
+        chklist = dataset.files
+        for ns in ns_filters or []:
+            if ns.startswith('!'):
+                ns = ns[1:]
+                ns_filelist = \
+                    [x for x in chklist
+                     if not re.search(f'namespace={ns}/', x)]
+                chklist = ns_filelist
+            elif ns.startswith('~'):
+                ns = ns[1:]
                 if ns.startswith('!'):
                     ns = ns[1:]
                     ns_filelist = \
                         [x for x in chklist
                          if not re.search(f'namespace={ns}/', x)]
                     chklist = ns_filelist
-                elif ns.startswith('~'):
-                    ns = ns[1:]
-                    if ns.startswith('!'):
-                        ns = ns[1:]
-                        ns_filelist = \
-                            [x for x in chklist
-                             if not re.search(f'namespace={ns}/', x)]
-                        chklist = ns_filelist
-                    else:
-                        ns_filelist.extend(
-                            [x for x in dataset.files
-                             if re.search(f'namespace={ns}/', x)])
                 else:
                     ns_filelist.extend(
                         [x for x in dataset.files
                          if re.search(f'namespace={ns}/', x)])
+            else:
+                ns_filelist.extend(
+                    [x for x in dataset.files
+                     if re.search(f'namespace={ns}/', x)])
 
-            filelist.extend(ns_filelist)
-
-        return ds.dataset(filelist, format='parquet', partitioning='hive')
+        return ds.dataset(ns_filelist, format='parquet', partitioning='hive')
 
     def _cons_int_filter(self, keyfld: str, filter_str: str) -> ds.Expression:
         '''Construct Integer filters with arithmetic operations'''