diff --git a/docs/advanced_guidance/json_schemas/contract/components/readable_entity.schema.json b/docs/advanced_guidance/json_schemas/contract/components/readable_entity.schema.json index f79e27b..91f0060 100644 --- a/docs/advanced_guidance/json_schemas/contract/components/readable_entity.schema.json +++ b/docs/advanced_guidance/json_schemas/contract/components/readable_entity.schema.json @@ -28,6 +28,11 @@ "$ref": "#/$defs/xmlReaderConfig", "description": "Configuration for loading an XML file.", "required": ["reader"] + }, + ".parquet": { + "$ref": "#/$defs/parquetReaderConfig", + "description": "Configuration for loading a Parquet file.", + "required": ["reader"] } }, "minProperties": 1 @@ -150,6 +155,32 @@ } } ] + }, + "parquetReaderConfig": { + "oneOf": [ + { + "type": "object", + "properties": { + "reader": { + "const": "DuckDBParquetReader" + }, + "kwargs": { + "$ref": "reader_constraints/ddb_parquet_reader.schema.json" + } + } + }, + { + "type": "object", + "properties": { + "reader": { + "const": "SparkParquetReader" + }, + "kwargs": { + "$ref": "reader_constraints/spark_parquet_reader.schema.json" + } + } + } + ] } }, "required": ["fields", "reader_config"] diff --git a/docs/advanced_guidance/json_schemas/contract/components/reader_constraints/ddb_parquet_reader.schema.json b/docs/advanced_guidance/json_schemas/contract/components/reader_constraints/ddb_parquet_reader.schema.json new file mode 100644 index 0000000..c4176dd --- /dev/null +++ b/docs/advanced_guidance/json_schemas/contract/components/reader_constraints/ddb_parquet_reader.schema.json @@ -0,0 +1,12 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "data-ingest:contract/components/reader_constraints/ddb_parquet_reader.schema.json", + "title": "Keyword Arguments a DuckDBParquetReader", + "description": "Arguments to contol how the DuckDB Parquet Reader interacts with a parquet files", + "type": "object", + "anyOf": [ + { + "$ref": "global_parquet_reader_args.schema.json" + } + ] +} \ No newline at end of file diff --git a/docs/advanced_guidance/json_schemas/contract/components/reader_constraints/global_parquet_reader_args.schema.json b/docs/advanced_guidance/json_schemas/contract/components/reader_constraints/global_parquet_reader_args.schema.json new file mode 100644 index 0000000..45ef3b9 --- /dev/null +++ b/docs/advanced_guidance/json_schemas/contract/components/reader_constraints/global_parquet_reader_args.schema.json @@ -0,0 +1,21 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "data-ingest:contract/components/reader_constraints/global_parquet_reader_args.schema.json", + "title": "Keyword Arguments used across all Parquet readers", + "description": "Arguments present in all Parquet readers available.", + "type": "object", + "properties": { + "hive_partitioning": { + "type": "boolean", + "description": "Infer statistics and schema from Hive partitioned URL and use them to prune reads." + }, + "field_check_error_code": { + "type": "string", + "description": "Error code to raise when fields are missing or unexpected." + }, + "field_check_error_message": { + "type": "string", + "description": "Error message to raise when fields are missing or unexpected." + } + } +} \ No newline at end of file diff --git a/docs/advanced_guidance/json_schemas/contract/components/reader_constraints/spark_parquet_reader.schema.json b/docs/advanced_guidance/json_schemas/contract/components/reader_constraints/spark_parquet_reader.schema.json new file mode 100644 index 0000000..d60ce1e --- /dev/null +++ b/docs/advanced_guidance/json_schemas/contract/components/reader_constraints/spark_parquet_reader.schema.json @@ -0,0 +1,32 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "data-ingest:contract/components/reader_constraints/spark_parquet_reader.schema.json", + "title": "Keyword Arguments a SparkParquetReader", + "description": "Arguments to contol how the PySpark Parquet Reader interacts with a parquet files", + "type": "object", + "anyOf": [ + { + "$ref": "global_parquet_reader_args.schema.json" + } + ], + "properties": { + "datetime_rebase_mode": { + "type": "string", + "description": "The datetimeRebaseMode option allows to specify the rebasing mode for the values of the DATE, TIMESTAMP_MILLIS, TIMESTAMP_MICROS logical types from the Julian to Proleptic Gregorian calendar. Default is `\"EXCEPTION\"`", + "enum": [ + "EXCEPTION", + "CORRECTED", + "LEGACY" + ] + }, + "int96_rebase_mode": { + "type": "string", + "description": "The int96RebaseMode option allows to specify the rebasing mode for INT96 timestamps from the Julian to Proleptic Gregorian calendar. Default is `\"EXCEPTION\"`.", + "enum": [ + "EXCEPTION", + "CORRECTED", + "LEGACY" + ] + } + } +} \ No newline at end of file diff --git a/docs/advanced_guidance/package_documentation/readers.md b/docs/advanced_guidance/package_documentation/readers.md index 4f19571..1900911 100644 --- a/docs/advanced_guidance/package_documentation/readers.md +++ b/docs/advanced_guidance/package_documentation/readers.md @@ -54,6 +54,25 @@ members: - __init__ +## Parquet + +=== "DuckDB" + + ::: dve.core_engine.backends.implementations.duckdb.readers.parquet.DuckDBParquetReader + options: + heading_level: 3 + members: + - __init__ + +=== "Spark" + + ::: dve.core_engine.backends.implementations.spark.readers.parquet.SparkParquetReader + options: + heading_level: 3 + members: + - __init__ + + ## XML === "Base" diff --git a/docs/user_guidance/file_transformation.md b/docs/user_guidance/file_transformation.md index 41dc7de..f792066 100644 --- a/docs/user_guidance/file_transformation.md +++ b/docs/user_guidance/file_transformation.md @@ -159,8 +159,9 @@ The secondary use of the File Transformation stage is the ability to normalise y ## Supported Formats -| Format | DuckDB | Spark | Version Available | -| ------- | ------------------ | ------------------ | ----------------- | -| `.csv` | :white_check_mark: | :white_check_mark: | >= 0.1.0 | -| `.json` | :white_check_mark: | :white_check_mark: | >= 0.1.0 | -| `.xml` | :white_check_mark: | :white_check_mark: | >= 0.1.0 | +| Format | DuckDB | Spark | Version Available | +| ---------- | ------------------ | ------------------ | ----------------- | +| `.csv` | :white_check_mark: | :white_check_mark: | >= 0.1.0 | +| `.json` | :white_check_mark: | :white_check_mark: | >= 0.1.0 | +| `.xml` | :white_check_mark: | :white_check_mark: | >= 0.1.0 | +| `.parquet` | :white_check_mark: | :white_check_mark: | >= 0.10.0 | diff --git a/pyproject.toml b/pyproject.toml index 0ad62d3..9944889 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -125,6 +125,10 @@ module = "polars.*" follow_imports = "skip" # ^language server knows what's going on, but mypy can't find attributes on Self? type +[[tool.mypy.overrides]] +module = "pyarrow.*" +ignore_missing_imports = true + [tool.black] line-length = 100 diff --git a/scripts/testdata_gen/gen_flights_parquet_testdata.py b/scripts/testdata_gen/gen_flights_parquet_testdata.py new file mode 100644 index 0000000..db2b47a --- /dev/null +++ b/scripts/testdata_gen/gen_flights_parquet_testdata.py @@ -0,0 +1,39 @@ +import polars as pl +from pathlib import Path + + +def main(): + data = [ + { + "flight_id": 1, + "plane_id": 1, + "passengers": [ + {"flight_id": 1, "passenger_id": 1, "passenger_name": "Alice"}, + {"flight_id": 1, "passenger_id": 2, "passenger_name": "Bob"} + ] + }, + { + "flight_id": 2, + "plane_id": 1, + "passengers": [ + {"flight_id": 2, "passenger_id": 3, "passenger_name": "Charlie"}, + {"flight_id": 2, "passenger_id": 4, "passenger_name": "Diana"} + ] + }, + { + "flight_id": 3, + "plane_id": 2, + "passengers": [ + {"flight_id": 3, "passenger_id": 5, "passenger_name": "Eve"}, + {"flight_id": 3, "passenger_id": 6, "passenger_name": "Frank"} + ] + } + ] + + df = pl.DataFrame(data) + df.write_parquet( + Path(Path(__file__).parent.parent.parent, "tests", "testdata", "flights", "flights.parquet") + ) + +if __name__ == "__main__": + main() diff --git a/src/dve/core_engine/backends/base/reader.py b/src/dve/core_engine/backends/base/reader.py index ae0e99f..5c15138 100644 --- a/src/dve/core_engine/backends/base/reader.py +++ b/src/dve/core_engine/backends/base/reader.py @@ -5,6 +5,7 @@ from inspect import ismethod from typing import Any, ClassVar, Optional, TypeVar +import pyarrow.parquet as pq from pydantic import BaseModel from typing_extensions import Protocol @@ -172,8 +173,27 @@ def _check_likely_text_file(resource: URI) -> bool: return True def raise_if_not_sensible_file(self, resource: URI, entity_name: str): - """Sense check that the file is a text file. Raise error if doesn't + """Sense check that the file is a text file or a valid parquet file. Raise error if doesn't appear to be the case.""" + if resource.endswith(".parquet"): + try: + pq.ParquetFile(resource) + except Exception as exc: + raise MessageBearingError( + "The submitted file doesn't appear to be a valid parquet format", + messages=[ + FeedbackMessage( + entity=entity_name, + record=None, + failure_type="submission", + error_location="Whole File", + error_code="MalformedFile", + error_message="The resource doesn't seem to be a valid parquet file." + ) + ] + ) from exc + return + if not self._check_likely_text_file(resource): raise MessageBearingError( "The submitted file doesn't appear to be text", diff --git a/src/dve/core_engine/backends/exceptions.py b/src/dve/core_engine/backends/exceptions.py index bb58516..e9a0c80 100644 --- a/src/dve/core_engine/backends/exceptions.py +++ b/src/dve/core_engine/backends/exceptions.py @@ -57,6 +57,30 @@ def __init__( ) +class UnableToParseParquetError(MessageBearingError): + """An error raised when unable to parse a CSV file""" + + def __init__( + self, entity_name: str, field_check_error_message: str, field_check_error_code: str + ): + super().__init__( + messages=[ + FeedbackMessage( + entity="parquet_structure", + record={ + entity_name: "Unable to parse file. Please check the structure of the file." + }, + failure_type="submission", + is_informational=False, + error_type="parquet read", + error_location=entity_name, + error_message=field_check_error_message, + error_code=field_check_error_code, + ) + ] + ) + + class BackendErrorMixin(ABC, BackendError): """A mixin used to create backend error type.""" diff --git a/src/dve/core_engine/backends/implementations/duckdb/__init__.py b/src/dve/core_engine/backends/implementations/duckdb/__init__.py index 996ec80..60dcf34 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/__init__.py +++ b/src/dve/core_engine/backends/implementations/duckdb/__init__.py @@ -7,6 +7,7 @@ from .readers import ( DuckDBCSVReader, DuckDBCSVRepeatingHeaderReader, + DuckDBParquetReader, DuckDBXMLStreamReader, PolarsToDuckDBCSVReader, ) @@ -16,6 +17,7 @@ register_reader(DuckDBCSVReader) register_reader(DuckDBCSVRepeatingHeaderReader) register_reader(DuckDBJSONReader) +register_reader(DuckDBParquetReader) register_reader(DuckDBXMLStreamReader) register_reader(PolarsToDuckDBCSVReader) diff --git a/src/dve/core_engine/backends/implementations/duckdb/contract.py b/src/dve/core_engine/backends/implementations/duckdb/contract.py index d9bb9bc..5dc3679 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/contract.py +++ b/src/dve/core_engine/backends/implementations/duckdb/contract.py @@ -30,7 +30,7 @@ duckdb_read_parquet, duckdb_record_index, duckdb_write_parquet, - get_duckdb_cast_statement_from_annotation, + generate_duckdb_casting_statements_from_model, get_duckdb_type_from_annotation, relation_is_empty, ) @@ -38,7 +38,7 @@ from dve.core_engine.backends.metadata.contract import DataContractMetadata from dve.core_engine.backends.types import StageSuccessful from dve.core_engine.backends.utilities import get_polars_type_from_annotation, stringify_model -from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME +from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME, SKIP_CONTRACT_CASTING from dve.core_engine.message import FeedbackMessage from dve.core_engine.type_hints import URI, EntityLocations from dve.core_engine.validation import RowValidator, apply_row_validator_helper @@ -165,29 +165,26 @@ def apply_data_contract( if RECORD_INDEX_COLUMN_NAME not in relation.columns: relation = self.add_record_index(relation) - casting_statements = [ - ( - get_duckdb_cast_statement_from_annotation(column, mdl_fld.annotation) - + f""" AS "{column}" """ - if column in relation.columns - else f"CAST(NULL AS {ddb_schema[column]}) AS {column}" + if list( + contract_metadata.reader_metadata[entity_name].keys() + )[0] not in SKIP_CONTRACT_CASTING: + casting_statements = generate_duckdb_casting_statements_from_model( + model_fields=entity_fields, + rel=relation, + ddb_schema=ddb_schema, + row_index_present=True, ) - for column, mdl_fld in entity_fields.items() - ] - casting_statements.append( - f"CAST({RECORD_INDEX_COLUMN_NAME} AS {get_duckdb_type_from_annotation(int)}) AS {RECORD_INDEX_COLUMN_NAME}" # pylint: disable=C0301 - ) - try: - relation = relation.project(", ".join(casting_statements)) - except Exception as err: # pylint: disable=broad-except - successful = False - self.logger.error(f"Error in casting relation: {err}") - dump_processing_errors( - working_dir, - "data_contract", - [generate_error_casting_entity_message(entity_name)], - ) - continue + try: + relation = relation.project(", ".join(casting_statements)) + except Exception as err: # pylint: disable=broad-except + successful = False + self.logger.error(f"Error in casting relation: {err}") + dump_processing_errors( + working_dir, + "data_contract", + [generate_error_casting_entity_message(entity_name)], + ) + continue if self.debug: # count will force evaluation - only done in debug diff --git a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py index 588cd7e..9cae060 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py +++ b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py @@ -17,6 +17,7 @@ from duckdb.typing import DuckDBPyType from pandas import DataFrame from pydantic import BaseModel +from pydantic.fields import FieldInfo from typing_extensions import Annotated, get_args, get_origin, get_type_hints from dve.common.error_utils import get_feedback_errors_uri @@ -484,3 +485,27 @@ def get_duckdb_cast_statement_from_annotation( stmt = f"TRIM({quoted_name})" return _cast_as_ddb_type(stmt, type_) if parent_element else stmt raise ValueError(f"No equivalent DuckDB type for {type_annotation!r}") + + +def generate_duckdb_casting_statements_from_model( + model_fields: dict[str, FieldInfo], + rel: DuckDBPyRelation, + ddb_schema: dict[str, Any], + row_index_present: bool = False, +) -> list[str]: + """Generate duckdb casting statement from pydantic model fields""" + casting_statements = [ + ( + get_duckdb_cast_statement_from_annotation(column, mdl_fld.annotation) + + f""" AS "{column}" """ + if column in rel.columns + else f"CAST(NULL AS {ddb_schema[column]}) AS {column}" + ) + for column, mdl_fld in model_fields.items() + ] + if row_index_present: + casting_statements.append( + f"CAST({RECORD_INDEX_COLUMN_NAME} AS {get_duckdb_type_from_annotation(int)}) AS {RECORD_INDEX_COLUMN_NAME}" # pylint: disable=C0301 + ) + + return casting_statements diff --git a/src/dve/core_engine/backends/implementations/duckdb/readers/__init__.py b/src/dve/core_engine/backends/implementations/duckdb/readers/__init__.py index b41cc67..20685e3 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/readers/__init__.py +++ b/src/dve/core_engine/backends/implementations/duckdb/readers/__init__.py @@ -2,12 +2,14 @@ from .csv import DuckDBCSVReader, DuckDBCSVRepeatingHeaderReader, PolarsToDuckDBCSVReader from .json import DuckDBJSONReader +from .parquet import DuckDBParquetReader from .xml import DuckDBXMLStreamReader __all__ = [ "DuckDBCSVReader", "DuckDBCSVRepeatingHeaderReader", "DuckDBJSONReader", + "DuckDBParquetReader", "DuckDBXMLStreamReader", "PolarsToDuckDBCSVReader", ] diff --git a/src/dve/core_engine/backends/implementations/duckdb/readers/parquet.py b/src/dve/core_engine/backends/implementations/duckdb/readers/parquet.py new file mode 100644 index 0000000..6f36915 --- /dev/null +++ b/src/dve/core_engine/backends/implementations/duckdb/readers/parquet.py @@ -0,0 +1,96 @@ +"""Parquet readers for DuckDB""" + +from collections.abc import Iterator +from typing import Any, Optional + +import duckdb as ddb +from duckdb import BinderException, DuckDBPyConnection, DuckDBPyRelation +from pydantic import BaseModel + +from dve.core_engine.backends.base.reader import read_function +from dve.core_engine.backends.exceptions import EmptyFileError, UnableToParseParquetError +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.parquet import BaseParquetReader +from dve.core_engine.backends.readers.utilities import get_parquet_metadata_row_count +from dve.core_engine.type_hints import URI, EntityName + + +@duckdb_record_index +@duckdb_write_parquet +class DuckDBParquetReader(BaseParquetReader): + """A DuckDB parquet reader.""" + + def __init__( + self, + *, + hive_partitioning: bool = False, + binary_as_string: bool = False, + connection: Optional[DuckDBPyConnection] = None, + field_check_error_code: str = "ParquetFieldMismatch", + field_check_error_message: str = "The submitted fields are invalid", + **_, + ): + self.binary_as_string = binary_as_string + """ + Parquet files generated by legacy writers do not correctly + set the UTF8 flag for strings, causing string columns to be + loaded as BLOB instead. Set this to true to load binary columns + as strings. + """ + self.connection = connection if connection else ddb.connect(":memory:") + """DuckDB connection. Creates an in memory ddb connection if not supplied.""" + super().__init__( + hive_partioning=hive_partitioning, + field_check_error_code=field_check_error_code, + field_check_error_message=field_check_error_message, + ) + + def read_to_py_iterator( + self, + resource: URI, + entity_name: EntityName, + schema: type[BaseModel], + all_model_fields: Optional[set[str]] = None, # pylint: disable=W0613 + ) -> Iterator[dict[str, Any]]: + """Creates an iterable object of rows as dictionaries""" + # TODO - these are actually loading entire relations into memory (csv & json imps as well) + # TODO - no out of the box offering for iterating records of a ddb relation + return self.read_to_relation(resource, entity_name, schema).pl().iter_rows(named=True) + + @read_function(DuckDBPyRelation) + def read_to_relation( + self, + resource: URI, + entity_name: EntityName, + schema: type[BaseModel], + all_model_fields: Optional[set[str]] = None, # pylint: disable=W0613 + ): + """Read a Parquet file to a Polars DataFrame""" + if get_parquet_metadata_row_count(resource) == 0: + raise EmptyFileError(f"File at {resource} is empty.") + + reader_options: dict[str, Any] = { + "hive_partitioning": self.hive_partioning, + "binary_as_string": self.binary_as_string, + } + + ddb_schema: dict[str, SQLType] = { + name: str(get_duckdb_type_from_annotation(fld.annotation)) # type: ignore + for name, fld in schema.model_fields.items() + } + + try: + relation = self.add_record_index( + self.connection.read_parquet(resource, **reader_options).select(*ddb_schema.keys()) + ) + except BinderException as exc: + raise UnableToParseParquetError( + entity_name, self.field_check_error_message, self.field_check_error_code + ) from exc + + return relation diff --git a/src/dve/core_engine/backends/implementations/spark/__init__.py b/src/dve/core_engine/backends/implementations/spark/__init__.py index 80cf4e4..e322269 100644 --- a/src/dve/core_engine/backends/implementations/spark/__init__.py +++ b/src/dve/core_engine/backends/implementations/spark/__init__.py @@ -4,12 +4,19 @@ from .backend import SparkBackend from .contract import SparkDataContract -from .readers import SparkCSVReader, SparkJSONReader, SparkXMLReader, SparkXMLStreamReader +from .readers import ( + SparkCSVReader, + SparkJSONReader, + SparkParquetReader, + SparkXMLReader, + SparkXMLStreamReader, +) from .reference_data import SparkRefDataLoader from .rules import SparkStepImplementations register_reader(SparkCSVReader) register_reader(SparkJSONReader) +register_reader(SparkParquetReader) register_reader(SparkXMLReader) register_reader(SparkXMLStreamReader) diff --git a/src/dve/core_engine/backends/implementations/spark/contract.py b/src/dve/core_engine/backends/implementations/spark/contract.py index 432a731..51d1971 100644 --- a/src/dve/core_engine/backends/implementations/spark/contract.py +++ b/src/dve/core_engine/backends/implementations/spark/contract.py @@ -37,7 +37,7 @@ from dve.core_engine.backends.metadata.contract import DataContractMetadata from dve.core_engine.backends.readers import CSVFileReader from dve.core_engine.backends.types import StageSuccessful -from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME +from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME, SKIP_CONTRACT_CASTING from dve.core_engine.type_hints import URI, EntityLocations, EntityName COMPLEX_TYPES: set[type[DataType]] = {StructType, ArrayType, MapType} @@ -148,31 +148,37 @@ def apply_data_contract( msg_count += len(batch) self.logger.info(f"Data contract found {msg_count} issues in {entity_name}") - try: - record_df = record_df.select( - *[ - ( - get_spark_cast_statement_from_annotation( - fld, fld_info.annotation - ).alias(fld) - if fld in record_df.columns - else lit(None) - .cast(get_type_from_annotation(fld_info.annotation)) - .alias(fld) - ) - for fld, fld_info in entity_fields.items() - ], - col(RECORD_INDEX_COLUMN_NAME).cast(LongType()).alias(RECORD_INDEX_COLUMN_NAME), - ) - except Exception as err: # pylint: disable=broad-except - successful = False - self.logger.error(f"Error in converting to dataframe: {err}") - dump_processing_errors( - working_dir, - "data_contract", - [generate_error_casting_entity_message(entity_name)], - ) - continue + if ( + list(contract_metadata.reader_metadata[entity_name].keys())[0] + not in SKIP_CONTRACT_CASTING + ): + try: + record_df = record_df.select( + *[ + ( + get_spark_cast_statement_from_annotation( + fld, fld_info.annotation + ).alias(fld) + if fld in record_df.columns + else lit(None) + .cast(get_type_from_annotation(fld_info.annotation)) + .alias(fld) + ) + for fld, fld_info in entity_fields.items() + ], + col(RECORD_INDEX_COLUMN_NAME) + .cast(LongType()) + .alias(RECORD_INDEX_COLUMN_NAME), + ) + except Exception as err: # pylint: disable=broad-except + successful = False + self.logger.error(f"Error in converting to dataframe: {err}") + dump_processing_errors( + working_dir, + "data_contract", + [generate_error_casting_entity_message(entity_name)], + ) + continue if self.debug: # Note, the count will realise the dataframe, so only do this diff --git a/src/dve/core_engine/backends/implementations/spark/readers/__init__.py b/src/dve/core_engine/backends/implementations/spark/readers/__init__.py index e89dc91..756ea3b 100644 --- a/src/dve/core_engine/backends/implementations/spark/readers/__init__.py +++ b/src/dve/core_engine/backends/implementations/spark/readers/__init__.py @@ -2,6 +2,7 @@ from dve.core_engine.backends.implementations.spark.readers.csv import SparkCSVReader from dve.core_engine.backends.implementations.spark.readers.json import SparkJSONReader +from dve.core_engine.backends.implementations.spark.readers.parquet import SparkParquetReader from dve.core_engine.backends.implementations.spark.readers.xml import ( SparkXMLReader, SparkXMLStreamReader, @@ -10,6 +11,7 @@ __all__ = [ "SparkCSVReader", "SparkJSONReader", + "SparkParquetReader", "SparkXMLReader", "SparkXMLStreamReader", ] diff --git a/src/dve/core_engine/backends/implementations/spark/readers/parquet.py b/src/dve/core_engine/backends/implementations/spark/readers/parquet.py new file mode 100644 index 0000000..a5f10e0 --- /dev/null +++ b/src/dve/core_engine/backends/implementations/spark/readers/parquet.py @@ -0,0 +1,96 @@ +"""Spark Parquet Readers""" + +from collections.abc import Iterator +from typing import Any, Literal, Optional + +from pydantic import BaseModel +from pyspark.errors import PySparkException +from pyspark.sql import DataFrame, SparkSession + +from dve.core_engine.backends.base.reader import read_function +from dve.core_engine.backends.exceptions import EmptyFileError, UnableToParseParquetError +from dve.core_engine.backends.implementations.spark.spark_helpers import ( + spark_record_index, + spark_write_parquet, +) +from dve.core_engine.backends.readers.parquet import BaseParquetReader +from dve.core_engine.backends.readers.utilities import get_parquet_metadata_row_count +from dve.core_engine.type_hints import URI, EntityName + + +@spark_record_index +@spark_write_parquet +class SparkParquetReader(BaseParquetReader): + """ + A Spark parquet reader. + + Specific reader args listed for PySpark 3.5.2: + https://archive.apache.org/dist/spark/docs/3.5.2/sql-data-sources-parquet.html#data-source-option. + """ + + def __init__( + self, + *, + datetime_rebase_mode: Literal["EXCEPTION", "CORRECTED", "LEGACY"] = "EXCEPTION", + int96_rebase_mode: Literal["EXCEPTION", "CORRECTED", "LEGACY"] = "EXCEPTION", + spark_session: Optional[SparkSession] = None, + field_check_error_code: str = "ParquetFieldMismatch", + field_check_error_message: str = "The submitted fields are invalid", + **_, + ): + self.datetime_rebase_mode = datetime_rebase_mode + """ + The datetimeRebaseMode option allows to specify the rebasing mode for the values of the + DATE, TIMESTAMP_MILLIS, TIMESTAMP_MICROS logical types from the Julian to Proleptic + Gregorian calendar. Default is `"EXCEPTION"`. + """ + self.int96_rebase_mode = int96_rebase_mode + """ + The int96RebaseMode option allows to specify the rebasing mode for INT96 timestamps from + the Julian to Proleptic Gregorian calendar. Default is `"EXCEPTION"`. + """ + self.spark_session = spark_session if spark_session else SparkSession.builder.getOrCreate() # type: ignore # pylint: disable=C0301 + """Spark Sesssion to utilise.""" + + super().__init__( + hive_partitioning=True, + field_check_error_code=field_check_error_code, + field_check_error_message=field_check_error_message, + ) + + def read_to_py_iterator( + self, + resource: URI, + entity_name: EntityName, + schema: type[BaseModel], + all_model_fields: Optional[set[str]] = None, # pylint: disable=W0613 + ) -> Iterator[dict[URI, Any]]: + df = self.read_to_dataframe(resource, entity_name, schema) + yield from (record.asDict(True) for record in df.toLocalIterator()) + + @read_function(DataFrame) + def read_to_dataframe( + self, + resource: URI, + entity_name: EntityName, + schema: type[BaseModel], # pylint: disable=W0613 + all_model_fields: Optional[set[str]] = None, # pylint: disable=W0613 + ) -> DataFrame: + """Read a parquet file into a PySpark DataFrame object.""" + if get_parquet_metadata_row_count(resource) == 0: + raise EmptyFileError(f"File at {resource} is empty.") + + read_options: dict[str, Any] = { + "datetimeRebaseMode": self.datetime_rebase_mode, + "int96RebaseMode": self.int96_rebase_mode, + } + + try: + reader = self.spark_session.read.format("parquet").options(**read_options) + df = self.add_record_index(reader.load(resource)) + except PySparkException as exc: + raise UnableToParseParquetError( + entity_name, self.field_check_error_message, self.field_check_error_code + ) from exc + + return df diff --git a/src/dve/core_engine/backends/readers/parquet.py b/src/dve/core_engine/backends/readers/parquet.py new file mode 100644 index 0000000..40f2b73 --- /dev/null +++ b/src/dve/core_engine/backends/readers/parquet.py @@ -0,0 +1,23 @@ +"""Core Parquet Reader""" + +from dve.core_engine.backends.base.reader import BaseFileReader + + +class BaseParquetReader(BaseFileReader): + """A base reader for parquet files""" + + def __init__( + self, + *, + hive_partitioning: bool = False, + field_check_error_code: str = "ParquetFieldMismatch", + field_check_error_message: str = "The submitted header is invalid", + **_, + ): + """Init function for the base parquet reader""" + self.hive_partioning = hive_partitioning + """Infer statistics and schema from Hive partitioned URL and use them to prune reads.""" + self.field_check_error_code = field_check_error_code + """Error code to raise when fields are missing or unexpected""" + self.field_check_error_message = field_check_error_message + """Error message to raise when fields are missing or unexpected""" diff --git a/src/dve/core_engine/backends/readers/utilities.py b/src/dve/core_engine/backends/readers/utilities.py index 3948d70..136955c 100644 --- a/src/dve/core_engine/backends/readers/utilities.py +++ b/src/dve/core_engine/backends/readers/utilities.py @@ -1,8 +1,11 @@ """General utilities for file readers""" +import glob from collections.abc import Iterable +from pathlib import Path from typing import Optional +import pyarrow.parquet as pq from pydantic import BaseModel from dve.core_engine.backends.exceptions import MessageBearingError @@ -74,3 +77,17 @@ def raise_message_bearing_error_on_header_differences( def get_all_model_fields(models: Iterable[type[BaseModel]]) -> set[str]: """Return all field names from all available models""" return {field for model in models for field in model.model_fields.keys()} + + +def get_parquet_metadata_row_count(resource: URI) -> int: + """Return metadata for a parquet file.""" + resource_path = Path(resource) + if resource_path.is_dir(): + content = glob.glob(f"{resource}/*.parquet") + if len(content) > 0: + resource = content[0] + else: + return 0 + + # todo - may need a attr check here in case num_rows not populated + return pq.ParquetFile(resource).metadata.num_rows diff --git a/src/dve/core_engine/constants.py b/src/dve/core_engine/constants.py index a2a4a65..ed31573 100644 --- a/src/dve/core_engine/constants.py +++ b/src/dve/core_engine/constants.py @@ -6,3 +6,6 @@ CONTRACT_ERROR_VALUE_FIELD_NAME: str = "__error_value" """The name of the field that can be used to extract the field value that caused a pydantic validation error""" + +SKIP_CONTRACT_CASTING: tuple[str] = (".parquet",) +"""File formats that can skip Data Contract casting as not required""" diff --git a/tests/features/flights.feature b/tests/features/flights.feature new file mode 100644 index 0000000..ae15281 --- /dev/null +++ b/tests/features/flights.feature @@ -0,0 +1,76 @@ +Feature: Pipeline tests using the books dataset + Tests for the processing framework which use the flights dataset. + + These tests are primarily around parquet submissions and ensuring the DVE handles + them correctly. + + Scenario: Validate a perfect flights parquet submission (duckdb) + Given I submit the flights file flights.parquet for processing + And A duckdb pipeline is configured with schema file 'flights.dischema.json' + And I add initial audit entries for the submission + Then the latest audit record for the submission is marked with processing status file_transformation + When I run the file transformation phase + Then the flights entity is stored as a parquet after the file_transformation phase + And the schema for flights entity matches the following + | column | dtype | + | flight_id | Int64 | + | plane_id | Int64 | + And the passengers entity is stored as a parquet after the file_transformation phase + # TODO - fix this as it's currently not working as expected + # And the schema for passengers entity matches the following + # | column | dtype | + # | flight_id | Int64 | + # | passenger_id | Int64 | + # | passenger_name | String | + And the latest audit record for the submission is marked with processing status data_contract + When I run the data contract phase + Then there are no record rejections from the data_contract phase + And there are no file rejections from the data_contract phase + When I run the business rules phase + Then the flights entity is stored as a parquet after the business_rules phase + And the passengers entity is stored as a parquet after the business_rules phase + And the latest audit record for the submission is marked with processing status error_report + When I run the error report phase + Then An error report is produced + And The statistics entry for the submission shows the following information + | parameter | value | + | record_count | 3 | + | number_submission_rejections | 0 | + | number_record_rejections | 0 | + | number_warnings | 0 | + + + Scenario: Validate a perfect flights parquet submission (spark) + Given I submit the flights file flights.parquet for processing + And A spark pipeline is configured with schema file 'flights_spark.dischema.json' + And I add initial audit entries for the submission + Then the latest audit record for the submission is marked with processing status file_transformation + When I run the file transformation phase + Then the flights entity is stored as a parquet after the file_transformation phase + And the schema for flights entity matches the following + | column | dtype | + | flight_id | Int64 | + | plane_id | Int64 | + And the passengers entity is stored as a parquet after the file_transformation phase + # TODO - fix this as it's currently not working as expected + # And the schema for passengers entity matches the following + # | column | dtype | + # | flight_id | Int64 | + # | passenger_id | Int64 | + # | passenger_name | String | + And the latest audit record for the submission is marked with processing status data_contract + When I run the data contract phase + Then there are no record rejections from the data_contract phase + And there are no file rejections from the data_contract phase + When I run the business rules phase + Then the flights entity is stored as a parquet after the business_rules phase + And the passengers entity is stored as a parquet after the business_rules phase + And the latest audit record for the submission is marked with processing status error_report + When I run the error report phase + Then An error report is produced + And The statistics entry for the submission shows the following information + | parameter | value | + | record_count | 3 | + | number_submission_rejections | 0 | + | number_record_rejections | 0 | + | number_warnings | 0 | diff --git a/tests/features/steps/steps_post_pipeline.py b/tests/features/steps/steps_post_pipeline.py index 906445e..ec86915 100644 --- a/tests/features/steps/steps_post_pipeline.py +++ b/tests/features/steps/steps_post_pipeline.py @@ -117,3 +117,26 @@ def check_error_aggregates_persisted(context): processing_location = get_processing_location(context) agg_file = Path(processing_location, "audit", "error_aggregates.parquet") assert agg_file.exists() and agg_file.is_file() + + +@then("the schema for {entity} entity matches the following") +def check_schema_matches(context, entity: str): + tbl: Table = context.table + + entity_path = Path( + get_processing_location(context), + "transform", + f"{entity}" + ) + try: + current_schema = pl.read_parquet_schema(entity_path) + except IsADirectoryError: + current_schema = pl.read_parquet(entity_path / "*.parquet").schema + + if tbl is None: + raise ValueError("No table supplied in step") + + row: Row + for row in tbl: + record: dict[str, str] = row.as_dict() + assert record["dtype"] == str(current_schema[record["column"]]) diff --git a/tests/test_core_engine/test_backends/test_readers/reader_utils.py b/tests/test_core_engine/test_backends/test_readers/reader_utils.py new file mode 100644 index 0000000..c4d69f8 --- /dev/null +++ b/tests/test_core_engine/test_backends/test_readers/reader_utils.py @@ -0,0 +1,24 @@ +"""Reader utils""" + +import tempfile +from pathlib import Path + +import polars as pl +import pytest +from pydantic import BaseModel + + +class TestParquetModel(BaseModel): + col1: str + col2: int + +@pytest.fixture(scope="function") +def test_parquet(): + with tempfile.TemporaryDirectory() as tdir: + df = pl.DataFrame( + [{"col1": "abc", "col2": 123}] + ) + write_path = Path(tdir, "test_parquet.parquet") + df.write_parquet(write_path) + + yield write_path, df.schema diff --git a/tests/test_core_engine/test_backends/test_readers/test_duckdb/__init__.py b/tests/test_core_engine/test_backends/test_readers/test_duckdb/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_core_engine/test_backends/test_readers/test_duckdb/test_ddb_parquet.py b/tests/test_core_engine/test_backends/test_readers/test_duckdb/test_ddb_parquet.py new file mode 100644 index 0000000..0dfc59c --- /dev/null +++ b/tests/test_core_engine/test_backends/test_readers/test_duckdb/test_ddb_parquet.py @@ -0,0 +1,23 @@ +from dve.core_engine.backends.implementations.duckdb.readers.parquet import DuckDBParquetReader +from dve.core_engine.backends.utilities import stringify_model + +import duckdb as ddb +import polars as pl + +from ..reader_utils import TestParquetModel, test_parquet + + +def test_ddb_parquet_read_all_defaults(test_parquet): + test_file_path, _expected_schema = test_parquet + + reader = DuckDBParquetReader() + rel: ddb.DuckDBPyRelation = reader.read_to_relation( + test_file_path.as_posix(), "test", TestParquetModel + ) + ignore_fields = ["__record_index__"] + + materialised_df = rel.pl() + assert materialised_df.shape[0] == 1 + assert all( + v == _expected_schema.get(k) or k in ignore_fields for k, v in materialised_df.schema.items() + ) diff --git a/tests/test_core_engine/test_backends/test_readers/test_spark/__init__.py b/tests/test_core_engine/test_backends/test_readers/test_spark/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_core_engine/test_backends/test_readers/test_spark/test_spark_parquet.py b/tests/test_core_engine/test_backends/test_readers/test_spark/test_spark_parquet.py new file mode 100644 index 0000000..82e2181 --- /dev/null +++ b/tests/test_core_engine/test_backends/test_readers/test_spark/test_spark_parquet.py @@ -0,0 +1,48 @@ +import tempfile +from pathlib import Path + +from dve.core_engine.backends.implementations.spark.readers import SparkParquetReader +from dve.core_engine.backends.utilities import stringify_model + +import pyspark.sql.functions as psf +import pyspark.sql.types as pst +import pytest +from pyspark.sql import DataFrame, Row, SparkSession + +from ..reader_utils import TestParquetModel + + +@pytest.fixture(scope="module") +def test_spark_parquet(spark: SparkSession): + with tempfile.TemporaryDirectory() as tdir: + schema = pst.StructType( + [ + pst.StructField("col1", pst.StringType()), + pst.StructField("col2", pst.IntegerType()), + ] + ) + df = spark.createDataFrame( + [ + Row( + col1="abc", + col2=123, + ), + ], + schema, + ) + write_path = Path(tdir, "test_parquet.parquet") + df.coalesce(1).write.parquet(write_path.as_posix()) + + yield write_path, schema + + +def test_spark_parquet_read_all_defaults(test_spark_parquet): + test_file_path, _expected_schema = test_spark_parquet + + reader = SparkParquetReader() + df: DataFrame = reader.read_to_dataframe(test_file_path.as_posix(), "test", TestParquetModel) + ignore_fields = ["__record_index__"] + df_fields = [f for f in df.schema.fields if f.name not in ignore_fields] + + assert df.count() == 1 + assert df_fields == _expected_schema.fields diff --git a/tests/testdata/flights/flights.dischema.json b/tests/testdata/flights/flights.dischema.json new file mode 100644 index 0000000..a9d9f1a --- /dev/null +++ b/tests/testdata/flights/flights.dischema.json @@ -0,0 +1,40 @@ +{ + "contract": { + "schemas": { + "passengers": { + "fields": { + "flight_id": "int", + "passenger_id": "int", + "passenger_name": "str" + } + } + }, + "datasets": { + "flights": { + "fields": { + "flight_id": "int", + "plane_id": "int" + }, + "reader_config": { + ".parquet": { + "reader": "DuckDBParquetReader" + } + } + }, + "passengers": { + "fields": { + "passengers": { + "model": "passengers", + "is_array": true + } + }, + "reader_config": { + ".parquet": { + "reader": "DuckDBParquetReader" + } + } + } + } + }, + "transformations": {} +} \ No newline at end of file diff --git a/tests/testdata/flights/flights.parquet b/tests/testdata/flights/flights.parquet new file mode 100644 index 0000000..d036964 Binary files /dev/null and b/tests/testdata/flights/flights.parquet differ diff --git a/tests/testdata/flights/flights_spark.dischema.json b/tests/testdata/flights/flights_spark.dischema.json new file mode 100644 index 0000000..82fbcfd --- /dev/null +++ b/tests/testdata/flights/flights_spark.dischema.json @@ -0,0 +1,40 @@ +{ + "contract": { + "schemas": { + "passengers": { + "fields": { + "flight_id": "int", + "passenger_id": "int", + "passenger_name": "str" + } + } + }, + "datasets": { + "flights": { + "fields": { + "flight_id": "int", + "plane_id": "int" + }, + "reader_config": { + ".parquet": { + "reader": "SparkParquetReader" + } + } + }, + "passengers": { + "fields": { + "passengers": { + "model": "passengers", + "is_array": true + } + }, + "reader_config": { + ".parquet": { + "reader": "SparkParquetReader" + } + } + } + } + }, + "transformations": {} +} \ No newline at end of file