Fix schema initialization in process_dataset_chunks

Author: alexcombessie

python-lib/dku_io_utils.py

@@ -18,28 +18,27 @@ def count_records(dataset: dataiku.Dataset) -> int:
 
     Returns:
         Number of records
+
     """
     metric_id = "records:COUNT_RECORDS"
     partitions = dataset.read_partitions
     client = dataiku.api_client()
     project = client.get_project(dataset.project_key)
     record_count = 0
-    logging.info("Counting records of dataset: {}...".format(dataset.name))
+    logging.info(f"Counting records of dataset: {dataset.name}...")
     if partitions is None or len(partitions) == 0:
         project.get_dataset(dataset.short_name).compute_metrics(metric_ids=[metric_id])
         metric = dataset.get_last_metric_values()
         record_count = dataiku.ComputedMetrics.get_value_from_data(metric.get_global_data(metric_id=metric_id))
-        logging.info("Dataset {} contains {:d} records and is not partitioned".format(dataset.name, record_count))
+        logging.info(f"Dataset {dataset.name} contains {record_count:d} records and is not partitioned")
     else:
         for partition in partitions:
             project.get_dataset(dataset.short_name).compute_metrics(partition=partition, metric_ids=[metric_id])
             metric = dataset.get_last_metric_values()
             record_count += dataiku.ComputedMetrics.get_value_from_data(
                 metric.get_partition_data(partition=partition, metric_id=metric_id)
             )
-        logging.info(
-            "Dataset {} contains {:d} records in partition(s) {}".format(dataset.name, record_count, partitions)
-        )
+        logging.info(f"Dataset {dataset.name} contains {record_count:d} records in partition(s) {partitions}")
     return record_count
 
 
@@ -48,8 +47,8 @@ def process_dataset_chunks(
 ) -> None:
     """Read a dataset by chunks, process each dataframe chunk with a function and write back to another dataset.
 
-    Passes keyword arguments to the function, adds a tqdm progress bar and generic logging.
-    Directly writes chunks to the output_dataset, so that only one chunk needs to be processed in-memory at a time.
+    Pass keyword arguments to the function, adds a tqdm progress bar and generic logging.
+    Directly write chunks to the output_dataset, so that only one chunk needs to be processed in-memory at a time.
 
     Args:
         input_dataset: Input dataiku.Dataset instance
@@ -59,45 +58,49 @@ def process_dataset_chunks(
             and output another pandas.DataFrame
         chunksize: Number of rows of each chunk of pandas.DataFrame fed to `func`
         **kwargs: Optional keyword arguments fed to `func`
+
+    Raises:
+        ValueError: If the input dataset is empty or if pandas cannot read it without type inference
+
     """
     input_count_records = count_records(input_dataset)
     if input_count_records == 0:
         raise ValueError("Input dataset has no records")
-    logging.info(
-        "Processing dataset {} of {:d} rows by chunks of {:d}...".format(
-            input_dataset.name, input_count_records, chunksize
-        )
-    )
+    logging.info(f"Processing dataset {input_dataset.name} of {input_count_records} rows by chunks of {chunksize}...")
     start = time()
+    # First, initialize output schema if not present. Required to show the real error if `iter_dataframes` fails.
+    if not output_dataset.read_schema(raise_if_empty=False):
+        df = input_dataset.get_dataframe(limit=5, infer_with_pandas=False)
+        output_df = func(df=df, **kwargs)
+        output_dataset.write_schema_from_dataframe(output_df)
     with output_dataset.get_writer() as writer:
         df_iterator = input_dataset.iter_dataframes(chunksize=chunksize, infer_with_pandas=False)
         len_iterator = math.ceil(input_count_records / chunksize)
-        for i, df in tqdm(enumerate(df_iterator), total=len_iterator):
+        for i, df in tqdm(enumerate(df_iterator), total=len_iterator, unit="chunk", mininterval=1.0):
             output_df = func(df=df, **kwargs)
             if i == 0:
                 output_dataset.write_schema_from_dataframe(
                     output_df, dropAndCreate=bool(not output_dataset.writePartition)
                 )
             writer.write_dataframe(output_df)
     logging.info(
-        "Processing dataset {} of {:d} rows: Done in {:.2f} seconds.".format(
-            input_dataset.name, input_count_records, time() - start
-        )
+        f"Processing dataset {input_dataset.name} of {input_count_records} rows: Done in {time() - start:.2f} seconds."
     )
 
 
-def set_column_description(
-    output_dataset: dataiku.Dataset, column_description_dict: Dict, input_dataset: dataiku.Dataset = None
+def set_column_descriptions(
+    output_dataset: dataiku.Dataset, column_descriptions: Dict, input_dataset: dataiku.Dataset = None
 ) -> None:
     """Set column descriptions of the output dataset based on a dictionary of column descriptions
 
-    Retains the column descriptions from the input dataset if the column name matches.
+    Retain the column descriptions from the input dataset if the column name matches.
 
     Args:
         output_dataset: Output dataiku.Dataset instance
-        column_description_dict: Dictionary holding column descriptions (value) by column name (key)
+        column_descriptions: Dictionary holding column descriptions (value) by column name (key)
         input_dataset: Optional input dataiku.Dataset instance
             in case you want to retain input column descriptions
+
     """
     output_dataset_schema = output_dataset.read_schema()
     input_dataset_schema = []
@@ -107,7 +110,7 @@ def set_column_description(
         input_columns_names = [col["name"] for col in input_dataset_schema]
     for output_col_info in output_dataset_schema:
         output_col_name = output_col_info.get("name", "")
-        output_col_info["comment"] = column_description_dict.get(output_col_name)
+        output_col_info["comment"] = column_descriptions.get(output_col_name)
         if output_col_name in input_columns_names:
             matched_comment = [
                 input_col_info.get("comment", "")