adding static postfix to partitioned exports (#77)

* adding static postfix to partitioned exports

* bump version

* address review comments

Author: stanislav-bedunkevich

cli-e2e-test/rel/config/config1.rel

@@ -14,6 +14,8 @@ module export_config
         def syntax:header = {
             1, :name
         }
+
+        def partition_size = 512  // defines the upper bound (MB) for the size of a partition
     end
 
     module stores_csv
@@ -22,6 +24,8 @@ module export_config
         def syntax:header = {
             1, :name
         }
+
+        def partition_size = 512  // defines the upper bound (MB) for the size of a partition
     end
 
     module products_csv

workflow/README.md

@@ -311,7 +311,12 @@ workflow:
 
 * `data` relation contains the well-defined (:col, key..., val) data to be exported;
 * `syntax:header` relation contains the header of the CSV file (Int, Symbol);
-* `partition_size` defines max size of the exported files in MB, if above threshold RAI partitions it in multiple files `filename_{part_nr}.csv`
+* `partition_size` defines max size of the exported files in MB, if above threshold RAI partitions it in multiple files
+  `filename_0_{part_nr}.csv`.
+
+Note: RWM puts a static `_0` postfix by default, as otherwise RAI would not add it if there's only one partition.
+Therefore, file names will look like `filename_0.csv` in case of a single partition and `filename_0_1.csv`, etc. in
+case of multiple partitions. *This only applies to CSV exports to Azure, as local exports can not be partitioned.*
 
 #### Basic exports
 The former needs to be used when a simple export config is passed as shown below:

workflow/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version_info__ = (0, 0, 43)
+__version_info__ = (0, 0, 44)
 __version__ = ".".join(map(str, __version_info__))

workflow/common.py

@@ -218,6 +218,7 @@ class Export:
     snapshot_binding: str
     container: Container
     offset_by_number_of_days: int = 0
+    is_partitioned: bool = False
 
 
 @dataclasses.dataclass

workflow/constants.py

@@ -23,6 +23,9 @@
 
 BLOB_PAGE_SIZE = 500
 
+# query constants
+PARTITIONED_EXPORT_POSTFIX = "_0"
+
 # Step types
 CONFIGURE_SOURCES = 'ConfigureSources'
 INSTALL_MODELS = 'InstallModels'

workflow/executor.py

@@ -488,6 +488,7 @@ def __init__(self, idt, name, type_value, state, timing, engine_size, exports, e
 
     def _execute(self, logger: logging.Logger, env_config: EnvConfig, rai_config: RaiConfig):
         exports = list(filter(lambda e: self._should_export(logger, rai_config, env_config, e), self.exports))
+        self._discover_partitioned(logger, rai_config, env_config, exports)
         if self.export_jointly:
             exports.sort(key=lambda e: e.container.name)
             container_groups = {container_name: list(group) for container_name, group in
@@ -527,6 +528,18 @@ def _should_export(self, logger: logging.Logger, rai_config: RaiConfig, env_conf
                 f"Skipping export of {export.relation}: defined as a snapshot and the current one is still valid")
         return should_export
 
+    @staticmethod
+    def _discover_partitioned(logger: logging.Logger, rai_config: RaiConfig, env_config: EnvConfig,
+                              exports: List[Export]):
+        logger.info("Identifying partitioned exports...")
+        rez = rai.execute_query_take_tuples(logger, rai_config, env_config, q.discover_partitioned_exports(exports))
+        if not rez:
+            return
+        # now get the keys of the output dict `rez` and drop the `:` from the start of the key
+        partitioned_export_names = [k[1:] for k in rez.keys()]
+        for export in exports:
+            export.is_partitioned = export.relation in partitioned_export_names
+
 
 class ExportWorkflowStepFactory(WorkflowStepFactory):
 

workflow/query.py

@@ -6,7 +6,7 @@
 
 from workflow import utils
 from workflow.common import FileType, Export, Source, ContainerType, AzureConfig
-from workflow.constants import IMPORT_CONFIG_REL, FILE_LOAD_RELATION
+from workflow.constants import IMPORT_CONFIG_REL, FILE_LOAD_RELATION, PARTITIONED_EXPORT_POSTFIX
 
 # Static queries
 DISABLE_IVM = "def insert:relconfig:disable_ivm = true"
@@ -171,6 +171,15 @@ def output(valid_until) {{
     """
 
 
+def discover_partitioned_exports(exports: List[Export]) -> str:
+    query = ""
+    for export in exports:
+        query += f"""
+        def output:{export.relation} = export_config:{export.relation}:partition_size = _
+        """
+    return query
+
+
 def export_relations_local(logger: logging.Logger, exports: List[Export]) -> str:
     query = ""
     for export in exports:
@@ -414,7 +423,9 @@ def output:{rel_name}[{key_str}] = csv_string[_export_csv_config:{rel_name}[{key
 
 def _export_relation_as_csv_to_azure(config: AzureConfig, export: Export, end_date: str, date_format: str) -> str:
     rel_name = export.relation
-    export_path = f"{_compose_export_path(config, export, end_date, date_format)}/{rel_name}.csv"
+    # for partitioned exports we append "_0" to the file name postfix so that we always have a partition number
+    postfix = PARTITIONED_EXPORT_POSTFIX if export.is_partitioned else ""
+    export_path = f"{_compose_export_path(config, export, end_date, date_format)}/{rel_name}{postfix}.csv"
     return f"""
     module _export_csv_config
         def {rel_name} = export_config:{rel_name}
@@ -428,6 +439,9 @@ def export:{rel_name} = export_csv[_export_csv_config:{rel_name}]
 def _export_meta_relation_as_csv_to_azure(config: AzureConfig, export: Export, end_date: str, date_format: str) -> str:
     rel_name = export.relation
     postfix = _to_rel_meta_key_as_str(export)
+    # for partitioned exports we append "_0" to the file name postfix so that we always have a partition number
+    if export.is_partitioned:
+        postfix += PARTITIONED_EXPORT_POSTFIX
     base_path = _compose_export_path(config, export, end_date, date_format)
     export_path = f"{base_path}/{rel_name}_{postfix}.csv"
     key_str = _to_rel_meta_key_as_seq(export)

workflow/rai.py

@@ -264,6 +264,24 @@ def execute_query_take_single(logger: logging.Logger, rai_config: RaiConfig, env
     return rsp.results[0]['table'].to_pydict()["v1"][0]
 
 
+def execute_query_take_tuples(logger: logging.Logger, rai_config: RaiConfig, env_config: EnvConfig, query: str,
+                              readonly: bool = True, ignore_problems: bool = False) -> any:
+    """
+    Execute query and take the results as an array of tuples.
+    :param logger:          logger
+    :param rai_config:      RAI config
+    :param env_config:      Env config
+    :param query:           Rel query
+    :param readonly:        Parameter to specify transaction type: Read/Write
+    :param ignore_problems: Ignore SDK problems if any
+    """
+    rsp = execute_query(logger, rai_config, env_config, query, readonly=readonly, ignore_problems=ignore_problems)
+    if not rsp.results:
+        logger.debug(f"Query returned no results: {query}")
+        return {}
+    return _parse_as_dict(rsp)
+
+
 def list_transactions(logger: logging.Logger, rai_config: RaiConfig) -> List:
     """
     List all transactions for the engine
@@ -322,13 +340,29 @@ def _parse_csv_string(rsp: api.TransactionAsyncResponse) -> Dict:
     Parse the output for a csv_string query ie
         def output:{rel_name} = csv_string[...]
     """
-    resp = {}
     rel_pattern = r'^/:output/:(.*)/String$'
+    return _parse_matching_pattern(rsp, rel_pattern)
+
+
+def _parse_as_dict(rsp: api.TransactionAsyncResponse) -> Dict:
+    """
+    Parse the output as dictionary, everything starting with :output
+    """
+    rel_pattern = r'^/:output/(.*)$'
+    return _parse_matching_pattern(rsp, rel_pattern)
+
+
+def _parse_matching_pattern(rsp: api.TransactionAsyncResponse, pattern: str) -> Dict:
+    resp = {}
     for result in rsp.results:
-        match = re.search(rel_pattern, result['relationId'])
+        match = re.search(pattern, result['relationId'])
         if match:
             relation_id = match.group(1)
-            data = result['table'].to_pydict()["v1"][0]
+            result_dict = result['table'].to_pydict()
+            if "v1" in result_dict:
+                data = result_dict["v1"][0]
+            else:
+                data = True  # if only relation label present, the value is True
             resp[relation_id] = data
     return resp