Feature id fixes (#114)

* id column fixes

Doc updates 100522 (#119)

* fixed reference to dbx in pull_request_template

* reverted inadvertently changed file

* release 0.2.1

* doc updates

* doc updates

* updates for building docs

* updated public docs

* updated sphinx version

* updated docs

* doc updates

* removed generated docs

* removed changes to non-doc

* fixed typo

* reverted unrelated changes

* added further coverage tests and renamed option fromn 'seedColumn' to 'seedColumnName' for clarity

* added further coverage test for 'seedColumnName' property'

* updated build process to explicotly use python 3.8

* updated tests for alternative seed id names

* simplifying condition where warning is issued on column override

Author: ronanstokes-db

dbldatagen/column_generation_spec.py

@@ -16,7 +16,7 @@
     TimestampType, DataType, DateType
 
 from .column_spec_options import ColumnSpecOptions
-from .datagen_constants import RANDOM_SEED_FIXED, RANDOM_SEED_HASH_FIELD_NAME, RANDOM_SEED_RANDOM
+from .datagen_constants import RANDOM_SEED_FIXED, RANDOM_SEED_HASH_FIELD_NAME, RANDOM_SEED_RANDOM, DEFAULT_SEED_COLUMN
 from .daterange import DateRange
 from .distributions import Normal, DataDistribution
 from .nrange import NRange
@@ -68,13 +68,11 @@ class ColumnGenerationSpec(object):
     :param debug: If True, output debugging log statements. Defaults to False.
     :param verbose: If True, output logging statements at the info level. If False (the default),
                     only output warning and error logging statements.
+    :param seedColumnName: if supplied, specifies seed column name
 
     For full list of options, see :doc:`/reference/api/dbldatagen.column_spec_options`.
     """
 
-    #: row seed field for data set
-    SEED_COLUMN = "id"
-
     #: maxValue values for each column type, only if where value is intentionally restricted
     _max_type_range = {
         'byte': 256,
@@ -89,7 +87,9 @@ class ColumnGenerationSpec(object):
 
     def __init__(self, name, colType=None, minValue=0, maxValue=None, step=1, prefix='', random=False,
                  distribution=None, baseColumn=None, randomSeed=None, randomSeedMethod=None,
-                 implicit=False, omit=False, nullable=True, debug=False, verbose=False, **kwargs):
+                 implicit=False, omit=False, nullable=True, debug=False, verbose=False,
+                 seedColumnName=DEFAULT_SEED_COLUMN,
+                 **kwargs):
 
         # set up logging
         self.verbose = verbose
@@ -108,9 +108,11 @@ def __init__(self, name, colType=None, minValue=0, maxValue=None, step=1, prefix
         self._initialBuildPlan = []  # the build plan for the column - descriptive only
         self.executionHistory = []  # the execution history for the column
 
+        self._seedColumnName = seedColumnName
+
         # If no base column is specified, assume its dependent on the seed column
         if baseColumn is None:
-            baseColumn = self.SEED_COLUMN
+            baseColumn = self._seedColumnName
 
         # to allow for open ended extension of many column attributes, we use a few specific
         # parameters and pass the rest as keyword arguments
@@ -345,10 +347,10 @@ def _computeBasicDependencies(self):
 
         :return: base columns as list with dependency on seed column added
         """
-        if self.baseColumn != self.SEED_COLUMN:
-            return list(set(self.baseColumns + [self.SEED_COLUMN]))
+        if self.baseColumn != self._seedColumnName:
+            return list(set(self.baseColumns + [self._seedColumnName]))
         else:
-            return [self.SEED_COLUMN]
+            return [self._seedColumnName]
 
     def setBaseColumnDatatypes(self, columnDatatypes):
         """ Set the data types for the base columns

dbldatagen/data_generator.py

@@ -12,13 +12,16 @@
 from pyspark.sql.types import LongType, IntegerType, StringType, StructType, StructField, DataType
 from .spark_singleton import SparkSingleton
 from .column_generation_spec import ColumnGenerationSpec
-from .datagen_constants import DEFAULT_RANDOM_SEED, RANDOM_SEED_FIXED, RANDOM_SEED_HASH_FIELD_NAME, MIN_SPARK_VERSION
+from .datagen_constants import DEFAULT_RANDOM_SEED, RANDOM_SEED_FIXED, RANDOM_SEED_HASH_FIELD_NAME, \
+                               DEFAULT_SEED_COLUMN, SPARK_RANGE_COLUMN, MIN_SPARK_VERSION
 from .utils import ensure, topologicalSort, DataGenError, deprecated
 from . _version import _get_spark_version
 
 _OLD_MIN_OPTION = 'min'
 _OLD_MAX_OPTION = 'max'
 
+_STREAMING_TIMESTAMP_COLUMN = "_source_timestamp"
+
 
 class DataGenerator:
     """ Main Class for test data set generation
@@ -35,6 +38,7 @@ class DataGenerator:
     :param verbose: = if `True`, generate verbose output
     :param batchSize: = UDF batch number of rows to pass via Apache Arrow to Pandas UDFs
     :param debug: = if set to True, output debug level of information
+    :param seedColumnName: = if set, this should be the name of the `seed` or logical `id` column. Defaults to `id`
     """
 
     # class vars
@@ -52,14 +56,20 @@ class DataGenerator:
 
     def __init__(self, sparkSession=None, name=None, randomSeedMethod=None,
                  rows=1000000, startingId=0, randomSeed=None, partitions=None, verbose=False,
-                 batchSize=None, debug=False, **kwargs):
+                 batchSize=None, debug=False, seedColumnName=DEFAULT_SEED_COLUMN,
+                 **kwargs):
         """ Constructor for data generator object """
 
         # set up logging
         self.verbose = verbose
         self.debug = debug
 
         self._setupLogger()
+        self._seedColumnName = seedColumnName
+        self._outputStreamingFields = False
+
+        if seedColumnName != DEFAULT_SEED_COLUMN:
+            self.logger.info(f"Using '{self._seedColumnName}' for seed column in place of '{DEFAULT_SEED_COLUMN}")
 
         self.name = name if name is not None else self.generateName()
         self._rowCount = rows
@@ -129,7 +139,9 @@ def __init__(self, sparkSession=None, name=None, randomSeedMethod=None,
         self._buildOrder = []
         self._inferredSchemaFields = []
         self.buildPlanComputed = False
-        self.withColumn(ColumnGenerationSpec.SEED_COLUMN, LongType(), nullable=False, implicit=True, omit=True)
+
+        # lets add the seed column
+        self.withColumn(self._seedColumnName, LongType(), nullable=False, implicit=True, omit=True, noWarn=True)
         self._batchSize = batchSize
 
         # set up spark session
@@ -138,6 +150,11 @@ def __init__(self, sparkSession=None, name=None, randomSeedMethod=None,
         # set up use of pandas udfs
         self._setupPandas(batchSize)
 
+    @property
+    def seedColumnName(self):
+        """ return the name of data generation seed column"""
+        return self._seedColumnName
+
     @classmethod
     def _checkSparkVersion(cls, sparkVersion, minSparkVersion):
         """
@@ -291,9 +308,12 @@ def explain(self, suppressOutput=False):
         if not self.buildPlanComputed:
             self.computeBuildPlan()
 
+        rc = self._rowCount
+        tasks = self.partitions
         output = ["", "Data generation plan", "====================",
-                  f"spec=DateGenerator(name={self.name}, rows={self._rowCount}, startingId={self.starting_id}, partitions={self.partitions})"
-            , ")", "", f"column build order: {self._buildOrder}", "", "build plan:"]
+                  f"spec=DateGenerator(name={self.name}, rows={rc}, startingId={self.starting_id}, partitions={tasks})"
+                  , ")", "", f"seed column: {self._seedColumnName}", "",
+                  f"column build order: {self._buildOrder}", "", "build plan:"]
 
         for plan_action in self._buildPlan:
             output.append(" ==> " + plan_action)
@@ -345,14 +365,14 @@ def rowCount(self):
         return self._rowCount
 
     def withIdOutput(self):
-        """ output seed column field (defaults to `id`) as a column in the test data set if specified
+        """ output seed column field (defaults to `id`) as a column in the generated data set if specified
 
-        If this is not called, the seed column field is omitted from the final test data set
+        If this is not called, the seed column field is omitted from the final generated data set
 
         :returns: modified in-place instance of test data generator allowing for chaining of calls
                   following Builder pattern
         """
-        self._columnSpecsByName[ColumnGenerationSpec.SEED_COLUMN].omit = False
+        self._columnSpecsByName[self._seedColumnName].omit = False
         self._markForPlanRegen()
 
         return self
@@ -584,7 +604,7 @@ def withColumnSpecs(self, patterns=None, fields=None, matchTypes=None, **kwargs)
 
         all_fields = self.getInferredColumnNames()
         effective_fields = [x for x in all_fields if
-                            (fields is None or x in fields) and x != ColumnGenerationSpec.SEED_COLUMN]
+                            (fields is None or x in fields) and x != self._seedColumnName]
 
         if patterns is not None:
             effective_fields = [x for x in effective_fields for y in patterns if re.search(y, x) is not None]
@@ -605,7 +625,7 @@ def _checkColumnOrColumnList(self, columns, allowId=False):
         :returns: True if test passes
         """
         inferred_columns = self.getInferredColumnNames()
-        if allowId and columns == ColumnGenerationSpec.SEED_COLUMN:
+        if allowId and columns == self._seedColumnName:
             return True
 
         if type(columns) is list:
@@ -677,7 +697,7 @@ def hasColumnSpec(self, colName):
     def withColumn(self, colName, colType=StringType(), minValue=None, maxValue=None, step=1,
                    dataRange=None, prefix=None, random=False, distribution=None,
                    baseColumn=None, nullable=True,
-                   omit=False, implicit=False,
+                   omit=False, implicit=False, noWarn=False,
                    **kwargs):
         """ add a new column for specification
 
@@ -693,6 +713,10 @@ def withColumn(self, colName, colType=StringType(), minValue=None, maxValue=None
         if baseColumn is not None:
             self._checkColumnOrColumnList(baseColumn, allowId=True)
 
+        if not noWarn and colName == self._seedColumnName:
+            self.logger.warning(f"Adding a new column named '{colName}' overrides seed column '{self._seedColumnName}'")
+            self.logger.warning(f"Use `seedColumName` option on DataGenerator construction for different seed column")
+
         # handle migration of old `min` and `max` options
         if _OLD_MIN_OPTION in kwargs:
             assert minValue is None, \
@@ -774,6 +798,7 @@ def _generateColumnDefinition(self, colName, colType=None, baseColumn=None,
                                            nullable=nullable,
                                            verbose=self.verbose,
                                            debug=self.debug,
+                                           seedColumnName=self._seedColumnName,
                                            **new_props)
 
         self._columnSpecsByName[colName] = column_spec
@@ -812,8 +837,9 @@ def _getBaseDataFrame(self, startId=0, streaming=False, options=None):
                                           end=end_id,
                                           numPartitions=id_partitions)
 
-            if ColumnGenerationSpec.SEED_COLUMN != "id":
-                df1 = df1.withColumnRenamed("id", ColumnGenerationSpec.SEED_COLUMN)
+            # spark.range generates a dataframe with the column `id` so rename it if its not our seed column
+            if SPARK_RANGE_COLUMN != self._seedColumnName:
+                df1 = df1.withColumnRenamed(SPARK_RANGE_COLUMN, self._seedColumnName)
 
         else:
             status = (
@@ -831,12 +857,15 @@ def _getBaseDataFrame(self, startId=0, streaming=False, options=None):
 
                 for k, v in options.items():
                     df1 = df1.option(k, v)
-                df1 = df1.load().withColumnRenamed("value", ColumnGenerationSpec.SEED_COLUMN)
+                df1 = (df1.load()
+                       .withColumnRenamed("value", self._seedColumnName)
+                       )
+
             else:
                 df1 = (df1.option("rowsPerSecond", 1)
                        .option("numPartitions", id_partitions)
                        .load()
-                       .withColumnRenamed("value", ColumnGenerationSpec.SEED_COLUMN)
+                       .withColumnRenamed("value", self._seedColumnName)
                        )
 
         return df1
@@ -854,16 +883,16 @@ def _computeColumnBuildOrder(self):
 
         :returns: the build ordering
         """
-        dependency_ordering = [(x.name, set(x.dependencies)) if x.name != ColumnGenerationSpec.SEED_COLUMN else (
-            ColumnGenerationSpec.SEED_COLUMN, set())
+        dependency_ordering = [(x.name, set(x.dependencies)) if x.name != self._seedColumnName else (
+            self._seedColumnName, set())
                                for x in self._allColumnSpecs]
 
         # self.pp_list(dependency_ordering, msg="dependencies")
 
         self.logger.info("dependency list: %s", str(dependency_ordering))
 
         self._buildOrder = list(
-            topologicalSort(dependency_ordering, flatten=False, initial_columns=[ColumnGenerationSpec.SEED_COLUMN]))
+            topologicalSort(dependency_ordering, flatten=False, initial_columns=[self._seedColumnName]))
 
         self.logger.info("columnBuildOrder: %s", str(self._buildOrder))
 
@@ -876,7 +905,7 @@ def build_order(self):
 
         The build order will be a list of lists - each list specifying columns that can be built at the same time
         """
-        return [x for x in self._buildOrder if x != [ColumnGenerationSpec.SEED_COLUMN]]
+        return [x for x in self._buildOrder if x != [self._seedColumnName]]
 
     def _getColumnDataTypes(self, columns):
         """ Get data types for columns
@@ -897,7 +926,7 @@ def computeBuildPlan(self):
         self._buildPlan = []
         self.executionHistory = []
         self._processOptions()
-        self._buildPlan.append(f"Build Spark data frame with seed column: {ColumnGenerationSpec.SEED_COLUMN}")
+        self._buildPlan.append(f"Build Spark data frame with seed column: '{self._seedColumnName}'")
 
         # add temporary columns
         for cs in self._allColumnSpecs:

dbldatagen/datagen_constants.py

@@ -26,6 +26,13 @@
 RANDOM_SEED_FIXED = "fixed"
 RANDOM_SEED_HASH_FIELD_NAME = "hash_fieldname"
 
+# constants related to seed column
+DEFAULT_SEED_COLUMN = "id"
+
+# this is the column name produced by `spark.range`
+# dont change unless semantics of `spark.range` changes
+SPARK_RANGE_COLUMN = "id"
+
 # minimum versions for version checks
 MIN_PYTHON_VERSION = (3, 8)
 MIN_SPARK_VERSION = (3, 1, 2)

docs/source/generating_column_data.rst

@@ -62,7 +62,7 @@ added through this method by a subsequent call to ``withColumnSpec`` to change t
 should be generated
 See :data:`~dbldatagen.data_generator.DataGenerator.withColumnSpecs`.
 
-By default all columns are marked as being dependent on an internal ``id`` column.
+By default all columns are marked as being dependent on an internal ``id`` seed column.
 Use the ``baseColumn`` attribute to mark a column as being dependent on another column or set of columns.
 Use of the base column attribute has several effects:
 

docs/source/repeatable_data_generation.rst

@@ -9,7 +9,9 @@ Repeatable data generation
 One of the basic principles of the data generator is that all data can be generated multiple times and
 produce the same results unless a column is marked ``random`` and no random seed is being used.
 
-Data is generated by generating an internal autoincrementing field called ``id``.
+Data is generated by generating an internal autoincrementing field called ``id``. The name of this column may be
+overridden by use of an argument to the instantiation of the DataGenerator object.
+
 All additional columns are generated through some transformation of the ``id`` column or some other designated
 ``baseColumn``, either by using its value, a hash of its value or an array of values if you are using multiple
 ``baseColumn`` fields.