Add persistence methods

Author: ghanse

dbldatagen/config.py

@@ -0,0 +1,21 @@
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+This module implements configuration classes for writing generated data.
+"""
+
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen=True, slots=True)
+class OutputConfig:
+    """
+    This class implements an output sink configuration used to write generated data. Output sinks must extend from the
+    `OutputConfig` base class.
+    """
+    location: str
+    output_mode: str
+    format: str = "delta"
+    options: dict = field(default_factory=dict)

dbldatagen/data_generator.py

@@ -20,6 +20,7 @@
 from dbldatagen import datagen_constants
 from dbldatagen._version import _get_spark_version
 from dbldatagen.column_generation_spec import ColumnGenerationSpec
+from dbldatagen.config import OutputConfig
 from dbldatagen.constraints import Constraint, SqlExpr
 from dbldatagen.datarange import DataRange
 from dbldatagen.distributions import DataDistribution
@@ -28,7 +29,14 @@
 from dbldatagen.serialization import SerializableToDict
 from dbldatagen.spark_singleton import SparkSingleton
 from dbldatagen.text_generators import TextGenerator
-from dbldatagen.utils import DataGenError, deprecated, ensure, split_list_matching_condition, topologicalSort
+from dbldatagen.utils import (
+    DataGenError,
+    deprecated,
+    ensure,
+    split_list_matching_condition,
+    topologicalSort,
+    write_data_to_output,
+)
 
 
 _OLD_MIN_OPTION: str = "min"
@@ -1909,6 +1917,18 @@ def scriptMerge(
 
         return result
 
+    def writeGeneratedData(self, config: OutputConfig, is_streaming: bool = False) -> None:
+        """
+        Builds a `DataFrame` from the `DataGenerator` and writes the data to a target table.
+
+        :param config: Output configuration for writing generated data
+        :param is_streaming: Whether to write data with Structured Streaming (default `False`)
+        """
+        if is_streaming:
+            write_data_to_output(self.build(withStreaming=True), config=config, is_streaming=is_streaming)
+
+        write_data_to_output(self.build(), config=config)
+
     @staticmethod
     def loadFromJson(options: str) -> "DataGenerator":
         """

dbldatagen/utils.py

@@ -18,6 +18,9 @@
 from typing import Any
 
 import jmespath
+from pyspark.sql import DataFrame
+
+from dbldatagen.config import OutputConfig
 
 
 def deprecated(message: str = "") -> Callable[[Callable[..., Any]], Callable[..., Any]]:
@@ -360,3 +363,32 @@ def system_time_millis() -> int:
     """
     curr_time: int = round(time.time() / 1000)
     return curr_time
+
+
+def write_data_to_output(df: DataFrame, config: OutputConfig, is_streaming: bool = False) -> None:
+    """
+    Writes a DataFrame to the sink configured in the output configuration.
+
+    :param df: Spark DataFrame to write
+    :param config: Output configuration passed as an `OutputConfig`
+    :param is_streaming: Whether to write the data with Structured Streaming (default `False`)
+    """
+
+    if is_streaming:
+        (
+            df
+            .writeStream
+            .format(config.format)
+            .outputMode(config.output_mode)
+            .options(**config.options)
+            .start(config.location)
+        )
+
+    (
+        df
+        .write
+        .format(config.format)
+        .mode(config.output_mode)
+        .options(**config.options)
+        .save(config.location)
+    )

docs/source/index.rst

@@ -37,6 +37,7 @@ As it is installable via `%pip install`, it can also be incorporated in environm
    Using multiple tables <multi_table_data>
    Extending text generation <extending_text_generation>
    Use with Delta Live Tables <using_delta_live_tables>
+   Writing Generated Data <writing_generated_data>
    Troubleshooting data generation <troubleshooting>
 
 .. toctree::

docs/source/writing_generated_data.rst

@@ -0,0 +1,33 @@
+.. Databricks Labs Data Generator documentation master file, created by
+   sphinx-quickstart on Sun Jun 21 10:54:30 2020.
+
+Writing Generated Data to Tables or Files
+===========================================================
+
+Generated data can be written directly to output tables or files using the ``OutputConfig`` class.
+
+Writing Generated Data to a Table
+--------------------------------------------
+
+Once you've defined a ``DataGenerator``, call the ``writeGeneratedData`` method to write data to a target table.
+
+.. code-block:: python
+
+   from pyspark.sql.types import StringType
+   import dbldatagen as dg
+   from dbldatagen.config import OutputConfig
+
+   # Create a sample data generator with a few columns:
+   testDataSpec = (
+       dg.DataGenerator(spark, name="users_dataset", rows=1000)
+       .withColumn("user_name", StringType(), expr="concat('user_', id)")
+       .withColumn("email_address", StringType(), expr="concat(user_name, '@email.com')")
+       .withColumn("phone_number", StringType(), template="555-DDD-DDDD")
+   )
+
+   # Define an output configuration:
+   outputConfig = OutputConfig(location="main.demo.table", output_mode="overwrite")
+
+   # Get the data generation options as a Python dictionary:
+   testDataSpec.writeGeneratedData(config=outputConfig)
+