thds.tabularasa 0.13.0__py3-none-any.whl

thds/tabularasa/schema/extract_from_parquet.py

@@ -0,0 +1,132 @@
+"""Should be able to 'extract' a Schema object from an existing parquet file.
+
+You might want to use this to convert a Parquet file directly to
+SQLite without additional ceremony - use `tabularasa.to_sqlite` in
+conjunction with this.
+"""
+
+import typing as ty
+from functools import partial
+from pathlib import Path
+
+import pyarrow as pa
+import pyarrow.lib
+import pyarrow.parquet as pq
+
+from . import metaschema as ms
+
+_LEAF_MAPPINGS = {
+    # DType doesn't currently support everything:
+    # https://arrow.apache.org/docs/python/api/datatypes.html
+    pa.string(): ms.DType.STR,
+    pa.bool_(): ms.DType.BOOL,
+    # ints
+    pa.int8(): ms.DType.INT8,
+    pa.int16(): ms.DType.INT16,
+    pa.int32(): ms.DType.INT32,
+    pa.int64(): ms.DType.INT64,
+    pa.uint8(): ms.DType.UINT8,
+    pa.uint16(): ms.DType.UINT16,
+    pa.uint32(): ms.DType.UINT32,
+    pa.uint64(): ms.DType.UINT64,
+    # reals
+    pa.float16(): ms.DType.FLOAT32,  # should we support float16 in DType?
+    pa.float32(): ms.DType.FLOAT32,
+    pa.float64(): ms.DType.FLOAT64,
+    # dates/times
+    pa.date32(): ms.DType.DATE,
+    pa.date64(): ms.DType.DATETIME,
+    pa.timestamp("s"): ms.DType.DATETIME,
+    pa.timestamp("ms"): ms.DType.DATETIME,
+    pa.timestamp("us"): ms.DType.DATETIME,
+    pa.timestamp("ns"): ms.DType.DATETIME,
+}
+ColumnType = ty.Union[ms.DType, ms.AnonCustomType, ms.CustomType, ms.ArrayType, ms.MappingType]
+
+
+def pyarrow_type_to_dtype(pyarrow_type: pyarrow.lib.DataType) -> ColumnType:
+    if pyarrow_type in _LEAF_MAPPINGS:
+        return _LEAF_MAPPINGS[pyarrow_type]
+    if pa.types.is_map(pyarrow_type):
+        key_type = pyarrow_type_to_dtype(pyarrow_type.key_type)
+        assert not isinstance(key_type, (ms._RawArrayType, ms._RawMappingType))
+        return ms.MappingType(
+            keys=key_type,
+            values=pyarrow_type_to_dtype(pyarrow_type.item_type),
+        )
+    if pa.types.is_list(pyarrow_type):
+        return ms.ArrayType(
+            values=pyarrow_type_to_dtype(pyarrow_type.value_type),
+        )
+    if pa.types.is_struct(pyarrow_type):
+        # TODO support these as though they were mappings, possibly?
+        raise ValueError("Structs are not yet supported by tabularasa.")
+    raise ValueError(f"Unsupported pyarrow type: {pyarrow_type}")
+
+
+def _decide_field_nullability(pyarrow_field: pyarrow.lib.Field, pq_file: pq.ParquetFile) -> bool:
+    if not pyarrow_field.nullable:
+        # if the incoming schema is certain about this, then
+        # maintain their declaration without inspecting the actual file.
+        return False
+    # otherwise, infer it from the data
+    for batch in pq_file.iter_batches(columns=[pyarrow_field.name]):
+        if batch[pyarrow_field.name].null_count:
+            return True
+    # the impact of saying this is False if no nulls were found but it
+    # was theoretically possible is low - we're trying to create a
+    # schema based on the data we _already have_, rather than the
+    # entire world of possible data.
+    return False
+
+
+def pyarrow_field_to_column(pq_file: pq.ParquetFile, pyarrow_field: pyarrow.lib.Field) -> ms.Column:
+    """Convert a pyarrow field to a Column object."""
+    return ms.Column(
+        name=pyarrow_field.name,
+        type=pyarrow_type_to_dtype(pyarrow_field.type),
+        nullable=_decide_field_nullability(pyarrow_field, pq_file),
+        doc=pyarrow_field.metadata and pyarrow_field.metadata.get("doc", "") or "autoextracted",
+    )
+
+
+def define_table_from_parquet(
+    pq_file: Path,
+    name: str,
+    *,
+    primary_key: ty.Optional[ms.IdTuple] = None,
+    indexes: ty.Collection[ms.IdTuple] = tuple(),
+) -> ms.Table:
+    """Extract a table from parquet into a Schema object.
+
+    The filename will be embedded in the doc field.
+    """
+    pq_schema = pq.read_schema(pq_file)
+
+    columns = list(map(partial(pyarrow_field_to_column, pq.ParquetFile(pq_file)), pq_schema))
+    valid_colnames = {column.name for column in columns}
+
+    # validate that primary_key and indexes match
+    def _validate(id_tuple: ms.IdTuple, descrip: str):
+        for identifier in id_tuple:
+            if identifier not in valid_colnames:
+                raise ValueError(
+                    f"Cannot specify name {identifier} as part of {descrip}"
+                    " since it is not a valid column name."
+                    f" Options are: {valid_colnames}"
+                )
+
+    for keys in indexes:
+        _validate(keys, f"index {keys}")
+    if primary_key:
+        _validate(primary_key, f"primary key {primary_key}")
+
+    return ms.Table(
+        name=name,
+        columns=columns,
+        doc=str(pq_file),
+        dependencies=None,
+        transient=True,
+        indexes=list(indexes),
+        primary_key=primary_key,
+    )

thds/tabularasa/schema/files.py

@@ -0,0 +1,215 @@
+import csv
+import datetime
+import os
+from enum import Enum
+from pathlib import Path
+from typing import Dict, List, Literal, Optional, Set, Tuple, Union
+
+import pkg_resources
+from pydantic import AnyUrl, BaseModel, Extra, Field
+
+from .util import DashedIdentifier, DocumentedMixin, DottedIdentifier, HexStr, PathStr
+
+
+class CSVQuotingConvention(Enum):
+    QUOTE_ALL = "quote_all"
+    QUOTE_NONE = "quote_none"
+    QUOTE_NONNUMERIC = "quote_nonnumeric"
+    QUOTE_MINIMAL = "quote_minimal"
+
+
+UpdateFrequency = Literal["Yearly", "Quarterly", "Monthly", "Biannual"]
+
+
+def quarter(date: datetime.date) -> int:
+    return (date.month - 1) // 3 + 1
+
+
+def half(date: datetime.date) -> int:
+    return (date.month - 1) // 6 + 1
+
+
+def _get_tail(freq: UpdateFrequency, date: datetime.date) -> Tuple[int, ...]:
+    if freq == "Yearly":
+        return ()
+    if freq == "Quarterly":
+        return (quarter(date),)
+    if freq == "Monthly":
+        return (date.month,)
+    return (half(date),)
+
+
+def _date_tuple(date: datetime.date, freq: UpdateFrequency) -> Tuple[int, ...]:
+    return (date.year, *_get_tail(freq, date))
+
+
+current_date = datetime.date.today()
+
+
+class FileSourceMixin(BaseModel, extra=Extra.forbid):
+    authority: Optional[str] = None
+    url: Optional[AnyUrl] = None
+    landing_page: Optional[AnyUrl] = None
+    last_checked: Optional[datetime.date] = None
+    last_updated: Optional[datetime.date] = None
+    update_frequency: Optional[UpdateFrequency] = None
+    is_open_access: Optional[bool] = None
+    doc: Optional[str] = None
+
+    def needs_update(self, current_date: datetime.date) -> bool:
+        if self.update_frequency is not None:
+            if self.last_updated is None:
+                return True
+            return _date_tuple(current_date, self.update_frequency) > _date_tuple(
+                self.last_updated, self.update_frequency
+            )
+        return False
+
+
+class LocalFileSourceMixin(FileSourceMixin):
+    filename: PathStr
+    package: Optional[DottedIdentifier] = None
+    encoding: Optional[str] = None
+
+    @property
+    def full_path(self) -> Path:
+        return Path(
+            self.filename
+            if self.package is None
+            else pkg_resources.resource_filename(self.package, self.filename)
+        )
+
+    @property
+    def file_handle(self):
+        if self.package is None:
+            return open(self.filename, "rb")
+        else:
+            if pkg_resources.resource_isdir(self.package, self.filename):
+                raise IsADirectoryError(pkg_resources.resource_filename(self.package, self.filename))
+            return pkg_resources.resource_stream(self.package, self.filename)
+
+    @property
+    def is_dir(self) -> bool:
+        if self.package is None:
+            return os.path.isdir(self.filename)
+        else:
+            return pkg_resources.resource_isdir(self.package, self.filename)
+
+
+class TabularFileSource(DocumentedMixin, LocalFileSourceMixin):
+    delimiter: Optional[str] = csv.excel.delimiter
+    quotechar: Optional[str] = csv.excel.quotechar
+    escapechar: Optional[str] = csv.excel.escapechar
+    doublequote: Optional[bool] = csv.excel.doublequote
+    skipinitialspace: Optional[bool] = csv.excel.skipinitialspace
+    lineterminator: Optional[str] = csv.excel.lineterminator
+    skiprows: Optional[int] = None
+    quoting: Optional[CSVQuotingConvention] = CSVQuotingConvention.QUOTE_MINIMAL
+    package: Optional[DottedIdentifier] = None
+    # Fairly conservative choice - only empty string is treated as explicitly null,
+    # and only on nullable columns
+    na_values: Optional[Set[str]] = Field(default_factory=lambda: {""})
+
+    @property
+    def csv_dialect(self) -> csv.Dialect:
+        # This is ugly but required for flexibility when using pandas.read_csv.
+        # The reason is that while read_csv allows passing all the attributes of a csv.Dialect,
+        # it does _not_ allow passing multi-char lineterminators (as are present in csv.excel as '\r\n'
+        # and this is often required in practice). The simple-seeming thing then is to pass the excel
+        # dialect and let the keyword args override it, but pandas overrides in the _opposite_ direction:
+        # the dialect overrides the keyword args. So in order to inject e.g. different delimiters or
+        # quoting conventions while keeping the '\r\n' lineterminator of the excel dialect, we have to
+        # build a new dialect object here.
+        kwargs = self.csv_reader_kwargs
+
+        if all(v == getattr(csv.excel, k) for k, v in kwargs.items()):
+            return csv.excel()
+        else:
+            dialect_name = "csv_dialect%d" % abs(hash(tuple(sorted(kwargs.items()))))
+            try:
+                dialect = csv.get_dialect(dialect_name)
+            except Exception:
+                dialect_cls = type(dialect_name, (csv.excel,), kwargs)
+                csv.register_dialect(dialect_name, dialect_cls)
+                dialect = dialect_cls()
+
+            return dialect
+
+    @property
+    def csv_reader_kwargs(self) -> Dict[str, Union[str, int, bool]]:
+        kw: Dict[str, Union[str, int, bool]] = {}
+        if self.quoting is not None:
+            kw.update(quoting=getattr(csv, self.quoting.name))
+
+        for name in [
+            "delimiter",
+            "quotechar",
+            "escapechar",
+            "doublequote",
+            "skipinitialspace",
+            "lineterminator",
+            "skiprows",
+        ]:
+            value = self.__dict__.get(name)
+            if value is not None:
+                kw[name] = value
+
+        return kw
+
+
+class VersionControlledPath(BaseModel, extra=Extra.forbid):
+    name: PathStr
+    md5: Optional[HexStr] = None
+
+
+class ADLSDataSpec(FileSourceMixin):
+    adls_account: DashedIdentifier
+    adls_filesystem: DashedIdentifier
+    paths: List[VersionControlledPath] = Field(min_items=1)
+    ordered: bool = False
+
+
+class LocalDataSpec(LocalFileSourceMixin):
+    order: List[PathStr] = Field(default_factory=list, unique_items=True)
+    package: Optional[DottedIdentifier] = None
+
+    def list_dir(self):
+        if self.package is None:
+            return os.listdir(self.filename)
+        else:
+            return pkg_resources.resource_listdir(self.package, self.filename)
+
+    @property
+    def all_data_specs(self) -> List["LocalDataSpec"]:
+        if self.is_dir:
+            return [
+                LocalDataSpec(package=self.package, filename=os.path.join(self.filename, filename))
+                for filename in self.list_dir()
+            ]
+        else:
+            return [self]
+
+    @property
+    def ordered_data_specs(self) -> List["LocalDataSpec"]:
+        if not self.order:
+            raise AttributeError("No `order` is set - `ordered_data_specs` are not available")
+        spec_order = {os.path.basename(spec.filename): spec for spec in self.all_data_specs}
+        if set(spec_order.keys()) != set(self.order):
+            raise ValueError(
+                f"`order` does not match the file names in the spec: {set(self.order)} != {set(spec_order.keys())}"
+            )
+        return [spec_order[name] for name in self.order]
+
+
+class RemoteBlobStoreSpec(BaseModel, extra=Extra.forbid):
+    adls_account: DashedIdentifier
+    adls_filesystem: DashedIdentifier
+    path: PathStr
+
+    def data_spec(self, md5: str, extension: Optional[str] = None) -> ADLSDataSpec:
+        ext = extension or ""
+        return ADLSDataSpec(
+            adls_account=self.adls_account,
+            adls_filesystem=self.adls_filesystem,
+            paths=[VersionControlledPath(name=f"{self.path.rstrip('/')}/{md5}{ext}", md5=md5)],
+        )