feat: add row index at file transform, if not using file transform add during dataa contract or business rules (but likely non-deterministic)

Author: stevenhsd

poetry.lock

@@ -42,6 +42,7 @@ include-groups = [
 [tool.poetry.group.dev.dependencies]
 commitizen = "4.9.1"
 pre-commit = "4.3.0"
+ipykernel = "^7.2.0"
 
 [tool.poetry.group.test]
 optional = true

pyproject.toml

@@ -163,7 +163,7 @@ def apply(
             return entities, get_parent(processing_errors_uri), successful
 
         for entity_name, entity in entities.items():
-            entities[entity_name] = self.step_implementations.add_row_id(entity)
+            entities[entity_name] = self.step_implementations.add_record_index(entity)
 
         # TODO: Handle entity manager creation errors.
         entity_manager = EntityManager(entities, reference_data)
@@ -172,9 +172,6 @@ def apply(
         # TODO: and return uri to errors
         _ = self.step_implementations.apply_rules(working_dir, entity_manager, rule_metadata)
 
-        for entity_name, entity in entity_manager.entities.items():
-            entity_manager.entities[entity_name] = self.step_implementations.drop_row_id(entity)
-
         return entity_manager.entities, get_parent(dc_feedback_errors_uri), True
 
     def process(

src/dve/core_engine/backends/base/backend.py

@@ -368,6 +368,14 @@ def read_raw_entities(
                 messages.extend(new_messages)
 
         return entities, dedup_messages(messages), successful
+    
+    def add_record_index(self, entity: EntityType, **kwargs) -> EntityType:
+        """Add a record index to the entity"""
+        raise NotImplementedError(f"add_record_index not implemented in {self.__class__}")
+    
+    def drop_record_index(self, entity: EntityType, **kwargs) -> EntityType:
+        """Drop a record index from the entity"""
+        raise NotImplementedError(f"drop_record_index not implemented in {self.__class__}")
 
     @abstractmethod
     def apply_data_contract(

src/dve/core_engine/backends/base/contract.py

@@ -126,6 +126,14 @@ def read_to_entity_type(
             raise ReaderLacksEntityTypeSupport(entity_type=entity_type) from err
 
         return reader_func(self, resource, entity_name, schema)
+    
+    def add_record_index(self, entity: EntityType, **kwargs) -> EntityType:
+        """Add a record index to the entity"""
+        raise NotImplementedError(f"add_record_index not implemented in {self.__class__}")
+    
+    def drop_record_index(self, entity: EntityType, **kwargs) -> EntityType:
+        """Drop a record index to the entity"""
+        raise NotImplementedError(f"drop_record_index not implemented in {self.__class__}")
 
     def write_parquet(
         self,

src/dve/core_engine/backends/base/reader.py

@@ -135,15 +135,13 @@ def register_udfs(cls, **kwargs):
         """Method to register all custom dve functions for use during business rules application"""
         raise NotImplementedError()
 
-    @staticmethod
-    def add_row_id(entity: EntityType) -> EntityType:
-        """Add a unique row id field to an entity"""
-        raise NotImplementedError()
+    def add_record_index(self, entity: EntityType, **kwargs) -> EntityType:
+        """Add a record index to the entity"""
+        raise NotImplementedError(f"add_record_index not implemented in {self.__class__}")
 
-    @staticmethod
-    def drop_row_id(entity: EntityType) -> EntityType:
-        """Add a unique row id field to an entity"""
-        raise NotImplementedError()
+    def drop_record_index(self, entity: EntityType) -> EntityType:
+        """Drop a unique row id field to an entity"""
+        raise NotImplementedError(f"drop_record_index not implemented in {self.__class__}")
 
     @classmethod
     def _raise_notimplemented_error(

src/dve/core_engine/backends/base/rules.py

@@ -12,6 +12,10 @@
 from dve.core_engine.type_hints import ExpressionArray, MultiExpression
 from dve.parser.type_hints import URI
 
+import polars as pl
+
+from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME
+
 BRACKETS = {"(": ")", "{": "}", "[": "]", "<": ">"}
 """A mapping of opening brackets to their closing counterpart."""
 STRING_START_CHARS = {'"', "'"}

src/dve/core_engine/backends/base/utilities.py

@@ -16,6 +16,7 @@
 from pydantic import BaseModel
 from pydantic.fields import ModelField
 
+from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME
 import dve.parser.file_handling as fh
 from dve.common.error_utils import (
     BackgroundMessageWriter,
@@ -29,6 +30,7 @@
 )
 from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import (
     duckdb_read_parquet,
+    duckdb_record_index,
     duckdb_write_parquet,
     get_duckdb_type_from_annotation,
     relation_is_empty,
@@ -53,7 +55,7 @@ def __call__(self, row: pd.Series):
         self.errors.extend(self.row_validator(row.to_dict())[1])  # type: ignore
         return row  # no op
 
-
+@duckdb_record_index
 @duckdb_write_parquet
 @duckdb_read_parquet
 class DuckDBDataContract(BaseDataContract[DuckDBPyRelation]):
@@ -144,10 +146,12 @@ def apply_data_contract(
                     fld.name: get_duckdb_type_from_annotation(fld.annotation)
                     for fld in entity_fields.values()
                 }
+                ddb_schema[RECORD_INDEX_COLUMN_NAME] = get_duckdb_type_from_annotation(int)
                 polars_schema: dict[str, PolarsType] = {
                     fld.name: get_polars_type_from_annotation(fld.annotation)
                     for fld in entity_fields.values()
                 }
+                polars_schema[RECORD_INDEX_COLUMN_NAME] = get_polars_type_from_annotation(int)
                 if relation_is_empty(relation):
                     self.logger.warning(f"+ Empty relation for {entity_name}")
                     empty_df = pl.DataFrame([], schema=polars_schema)  # type: ignore # pylint: disable=W0612
@@ -169,6 +173,9 @@ def apply_data_contract(
                         msg_count += len(msgs)
 
                 self.logger.info(f"Data contract found {msg_count} issues in {entity_name}")
+                
+                if not RECORD_INDEX_COLUMN_NAME in relation.columns:
+                    relation = self.add_record_index(relation)
 
                 casting_statements = [
                     (

src/dve/core_engine/backends/implementations/duckdb/contract.py

@@ -12,13 +12,14 @@
 
 import duckdb.typing as ddbtyp
 import numpy as np
-from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from duckdb import DuckDBPyConnection, DuckDBPyRelation, StarExpression
 from duckdb.typing import DuckDBPyType
 from pandas import DataFrame
 from pydantic import BaseModel
 from typing_extensions import Annotated, get_args, get_origin, get_type_hints
 
 from dve.core_engine.backends.base.utilities import _get_non_heterogenous_type
+from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME
 from dve.core_engine.type_hints import URI
 from dve.parser.file_handling.service import LocalFilesystemImplementation, _get_implementation
 
@@ -286,3 +287,19 @@ def duckdb_rel_to_dictionaries(
     cols: tuple[str] = tuple(entity.columns)  # type: ignore
     while rows := entity.fetchmany(batch_size):
         yield from (dict(zip(cols, rw)) for rw in rows)
+
+def _add_duckdb_record_index(self, entity: DuckDBPyRelation) -> DuckDBPyRelation:
+    if RECORD_INDEX_COLUMN_NAME in entity.columns:
+        return entity
+
+    return entity.select(f"*, row_number() OVER () as {RECORD_INDEX_COLUMN_NAME}")
+
+def _drop_duckdb_record_index(self, entity: DuckDBPyRelation) -> DuckDBPyRelation:
+    if RECORD_INDEX_COLUMN_NAME not in entity.columns:
+        return entity
+    return entity.select(StarExpression(exclude=[RECORD_INDEX_COLUMN_NAME]))
+
+def duckdb_record_index(cls):
+    setattr(cls, "add_record_index", _add_duckdb_record_index)
+    setattr(cls, "drop_record_index", _drop_duckdb_record_index)
+    return cls

src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py

@@ -6,23 +6,25 @@
 
 import duckdb as ddb
 import polars as pl
-from duckdb import DuckDBPyConnection, DuckDBPyRelation, default_connection, read_csv
+from duckdb import DuckDBPyConnection, DuckDBPyRelation, StarExpression, default_connection, read_csv
 from pydantic import BaseModel
 
 from dve.core_engine.backends.base.reader import BaseFileReader, read_function
 from dve.core_engine.backends.exceptions import EmptyFileError, MessageBearingError
 from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import (
+    duckdb_record_index,
     duckdb_write_parquet,
     get_duckdb_type_from_annotation,
 )
 from dve.core_engine.backends.implementations.duckdb.types import SQLType
 from dve.core_engine.backends.readers.utilities import check_csv_header_expected
-from dve.core_engine.backends.utilities import get_polars_type_from_annotation
+from dve.core_engine.backends.utilities import get_polars_type_from_annotation, polars_record_index
+from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME
 from dve.core_engine.message import FeedbackMessage
 from dve.core_engine.type_hints import URI, EntityName
 from dve.parser.file_handling import get_content_length
 
-
+@duckdb_record_index
 @duckdb_write_parquet
 class DuckDBCSVReader(BaseFileReader):
     """A reader for CSV files including the ability to compare the passed model
@@ -109,9 +111,9 @@ def read_to_relation(  # pylint: disable=unused-argument
         }
 
         reader_options["columns"] = ddb_schema
-        return read_csv(resource, **reader_options)
-
+        return self.add_record_index(read_csv(resource, **reader_options, parallel=False))
 
+@polars_record_index
 class PolarsToDuckDBCSVReader(DuckDBCSVReader):
     """
     Utilises the polars lazy csv reader which is then converted into a DuckDBPyRelation object.
@@ -142,10 +144,11 @@ def read_to_relation(  # pylint: disable=unused-argument
             for fld in schema.__fields__.values()
         }
         reader_options["dtypes"] = polars_types
+        
 
         # there is a raise_if_empty arg for 0.18+. Future reference when upgrading. Makes L85
         # redundant
-        df = pl.scan_csv(resource, **reader_options).select(list(polars_types.keys()))  # type: ignore  # pylint: disable=W0612
+        df = self.add_record_index(pl.scan_csv(resource, **reader_options).select(list(polars_types.keys())))  # type: ignore  # pylint: disable=W0612
 
         return ddb.sql("SELECT * FROM df")
 
@@ -189,8 +192,8 @@ def __init__(
     def read_to_relation(  # pylint: disable=unused-argument
         self, resource: URI, entity_name: EntityName, schema: type[BaseModel]
     ) -> DuckDBPyRelation:
-        entity = super().read_to_relation(resource=resource, entity_name=entity_name, schema=schema)
-        entity = entity.distinct()
+        entity: DuckDBPyRelation = super().read_to_relation(resource=resource, entity_name=entity_name, schema=schema)
+        entity = entity.select(StarExpression(exclude=[RECORD_INDEX_COLUMN_NAME])).distinct()
         no_records = entity.shape[0]
 
         if no_records != 1:
@@ -219,4 +222,4 @@ def read_to_relation(  # pylint: disable=unused-argument
                 ],
             )
 
-        return entity
+        return entity.select(f"*, 1 as {RECORD_INDEX_COLUMN_NAME}")