owid-catalog: Data Structures and Processing¶

def add(
    self,
    table: tables.Table,
    formats: list[FileFormat] = DEFAULT_FORMATS,
    repack: bool = True,
) -> None:
    """Add a table to this dataset.

    Saves the table to the dataset's directory in the specified format(s).
    By default, saves in multiple formats for compatibility.

    Args:
        table: The table to add to the dataset.
        formats: List of file formats to save (feather, parquet, csv).
            Defaults to DEFAULT_FORMATS (usually ["feather"]).
        repack: If True, optimize column dtypes to reduce file size
            (e.g. float64 -> float32). Set to False for very large dataframes
            if repacking fails or is too slow.

    Raises:
        PrimaryKeyMissing: If table has no primary key and OWID_STRICT is set.
        NonUniqueIndex: If table index has duplicates and OWID_STRICT is set.

    Example:
        ```python
        >>> ds.add(table)  # Save in default format
        >>> ds.add(table, formats=["csv"])  # Save only as CSV
        >>> ds.add(table, repack=False)  # Skip optimization
        ```
    """

    utils.validate_underscore(table.metadata.short_name, "Table's short_name")
    for col in list(table.columns) + list(table.index.names):
        utils.validate_underscore(col, "Variable's name")

    if not table.primary_key:
        if environ.get("OWID_STRICT"):
            raise PrimaryKeyMissing(
                f"Table `{table.metadata.short_name}` does not have a primary_key -- please use t.set_index([col, ...], verify_integrity=True) to indicate dimensions before saving"
            )
        else:
            warnings.warn(
                f"Table `{table.metadata.short_name}` does not have a primary_key -- please use t.set_index([col, ...], verify_integrity=True) to indicate dimensions before saving"
            )

    if not table.index.is_unique and environ.get("OWID_STRICT"):
        [(k, dups)] = table.index.value_counts().head(1).to_dict().items()
        raise NonUniqueIndex(
            f"Table `{table.metadata.short_name}` has duplicate values in the index -- could you have made a mistake?\n\n"
            f"e.g. key {k} is repeated {dups} times in the index"
        )

    # check Float64 and Int64 columns for np.nan
    # see: https://github.com/owid/etl/issues/1334
    for col, dtype in table.dtypes.items():
        if dtype in NULLABLE_DTYPES:
            # pandas nullable types like Float64 have their own pd.NA instead of np.nan
            # make sure we don't use wrong nan, otherwise dropna and other methods won't work
            assert np.isnan(table[col]).sum() == 0, (
                f"Column `{col}` is using np.nan, but it should be using pd.NA because it has type {table[col].dtype}"
            )

    # copy dataset metadata to the table
    table.metadata.dataset = self.metadata

    for format in formats:
        if format not in SUPPORTED_FORMATS:
            raise Exception(f"Format '{format}'' is not supported")

        table_filename = join(self.path, table.metadata.checked_name + f".{format}")
        table.to(table_filename, repack=repack)

def checksum(self) -> str:
    """Calculate MD5 checksum of all data and metadata in the dataset.

    Generates a checksum that includes the dataset's index file and all
    data files. Useful for detecting changes to the dataset.

    Returns:
        MD5 checksum as a hexadecimal string.

    Example:
        ```python
        >>> checksum = ds.checksum()
        >>> print(f"Dataset checksum: {checksum}")
        ```
    """
    _hash = hashlib.md5()
    _hash.update(checksum_file(self._index_file).digest())

    for data_file in self._data_files:
        _hash.update(checksum_file(data_file).digest())

        metadata_file = Path(data_file).with_suffix(".meta.json").as_posix()
        _hash.update(checksum_file(metadata_file).digest())

    return _hash.hexdigest()

def index(self, catalog_path: Path = Path("/")) -> pd.DataFrame:
    """Generate an index DataFrame describing all tables in this dataset.

    Creates a summary DataFrame with one row per table, including metadata
    like namespace, version, checksum, dimensions, and file paths.

    Args:
        catalog_path: Base path for calculating relative paths. Defaults to "/".

    Returns:
        DataFrame with columns: namespace, dataset, version, table, checksum, is_public,
        title, description, dimensions, path, channel, and formats.

    Example:
        ```python
        >>> index = ds.index()
        >>> print(index[["table", "dimensions", "checksum"]])
        ```
    """
    base = {
        "namespace": self.metadata.namespace,
        "dataset": self.metadata.short_name,
        "version": self.metadata.version,
        "checksum": self.checksum(),
        "is_public": self.metadata.is_public,
    }
    rows = []
    for metadata_file in self._metadata_files:
        with open(metadata_file) as istream:
            metadata = TableMeta.from_dict(json.load(istream))

        row = base.copy()

        assert metadata.short_name
        row["table"] = metadata.short_name

        # Content metadata (fallback to dataset-level if table-level not available)
        row["title"] = metadata.title or self.metadata.title
        row["description"] = metadata.description or self.metadata.description

        row["dimensions"] = json.dumps(metadata.primary_key)

        table_path = Path(self.path) / metadata.short_name
        relative_path = table_path.relative_to(catalog_path)
        row["path"] = relative_path.as_posix()
        row["channel"] = relative_path.parts[0]

        row["formats"] = [f for f in SUPPORTED_FORMATS if table_path.with_suffix(f".{f}").exists()]  # ty: ignore

        rows.append(row)

    return pd.DataFrame.from_records(rows)

def read(
    self,
    name: str | None = None,
    reset_index: bool = True,
    safe_types: bool = True,
    reset_metadata: Literal["keep", "keep_origins", "reset"] = "keep",
    load_data: bool = True,
) -> tables.Table:
    """Read a table from the dataset with performance options.

    This is an alternative to `ds[table_name]` with more control over
    loading behavior for performance optimization.

    Args:
        name: Name of the table to read. If None and dataset has only one
            table, reads that table automatically.
        reset_index: If True, don't set primary keys. This can make loading
            large multi-index datasets much faster. Default is True.
        safe_types: If True, convert numeric columns to nullable types
            (Float64, Int64) and categorical to string[pyarrow]. This increases
            memory usage but prevents type issues. Default is True.
        reset_metadata: Controls variable metadata reset behavior:
            - "keep": Leave metadata unchanged (default)
            - "keep_origins": Reset metadata but retain origins attribute
            - "reset": Reset all variable metadata
        load_data: If False, only load metadata without actual data. Useful
            when you only need to inspect metadata. Default is True.

    Returns:
        The loaded table with data and metadata.

    Raises:
        ValueError: If name is None but dataset contains multiple tables.
        KeyError: If the specified table name doesn't exist.

    Example:
        Read single table with safe defaults
        ```python
        table = ds.read()
        ```

        Keep index
        ```python
        >>> table = ds.read("population", reset_index=False)
        ```

        Faster, less memory
        ```python
        >>> table = ds.read("large_table", safe_types=False)
        ```

        Only metadata
        ```python
        >>> meta_only = ds.read(load_data=False)
        ```
    """
    if name is None:
        if len(self.table_names) == 1:
            name = self.table_names[0]
        else:
            raise ValueError("Multiple tables exist. Please specify the table name.")
    stem = self.path / Path(name)

    for format in SUPPORTED_FORMATS:
        path = stem.with_suffix(f".{format}")
        if path.exists():
            t = tables.Table.read(path, primary_key=[] if reset_index else None, load_data=load_data)
            t.metadata.dataset = self.metadata
            if safe_types and load_data:
                t = cast(tables.Table, to_safe_types(t))
            if reset_metadata in ["keep_origins", "reset"]:  # Handles "keep_origins" and "reset"
                t.metadata = TableMeta()
                for col in t.columns:
                    if reset_metadata == "keep_origins":  # Preserve 'origins' attribute
                        origins = t[col].metadata.origins if hasattr(t[col].metadata, "origins") else None
                        t[col].metadata = VariableMeta()
                        t[col].metadata.origins = origins  # Preserve 'origins' attribute
                    if reset_metadata == "reset":  # Reset all metadata
                        t[col].metadata = VariableMeta()
            return t

    raise KeyError(f"Table `{name}` not found, available tables: {', '.join(self.table_names[:10])}")

def update_metadata(
    self,
    metadata_path: Path,
    yaml_params: dict[str, Any] | None = None,
    if_origins_exist: SOURCE_EXISTS_OPTIONS = "replace",
    errors: Literal["ignore", "warn", "raise"] = "raise",
    extra_variables: Literal["raise", "ignore"] = "raise",
) -> None:
    """Update dataset and table metadata from a YAML file.

    Loads metadata from a .meta.yml file and updates the dataset's metadata
    and all referenced tables. This is the primary way to add rich metadata
    to datasets in the ETL workflow.

    Args:
        metadata_path: Path to the .meta.yml file with metadata definitions.
            See existing metadata files for examples of the expected structure.
        yaml_params: Additional parameters to pass to the YAML loader.
        if_origins_exist: How to handle existing origins:
            - "replace" (default): Replace existing origin with new one
            - "append": Append new origin to existing origins
            - "fail": Raise exception if origin already exists
        errors: How to handle errors during update:
            - "raise" (default): Raise exception on errors
            - "warn": Issue warning but continue processing
            - "ignore": Silently ignore errors
        extra_variables: How to handle variables in metadata not in dataset:
            - "raise" (default): Raise exception
            - "ignore": Skip extra variables

    Example:
        ```python
        >>> ds.update_metadata(Path("dataset.meta.yml"))
        >>> ds.update_metadata(
        ...     Path("dataset.meta.yml"),
        ...     if_origins_exist="append",
        ...     errors="warn"
        ... )
        ```
    """
    self.metadata.update_from_yaml(metadata_path)

    with open(metadata_path) as istream:
        metadata = yaml.safe_load(istream)
        for table_name in metadata.get("tables", {}).keys():
            try:
                table = self[table_name]
            except KeyError as e:
                if errors == "raise":
                    raise e
                else:
                    if errors == "warn":
                        warnings.warn(str(e))
                    continue
            table.update_metadata_from_yaml(
                metadata_path,
                table_name,
                if_origins_exist=if_origins_exist,
                yaml_params=yaml_params,
                extra_variables=extra_variables,
            )
            table._save_metadata(join(self.path, table.metadata.checked_name + ".meta.json"))

def astype(self, *args: Any, **kwargs: Any) -> Table:
    """Cast table columns to specified dtype(s).

    Convert one or more columns to a specified data type. Wrapper
    around pandas astype that returns a Table.

    Args:
        *args: Positional arguments passed to pandas.DataFrame.astype.
        **kwargs: Keyword arguments passed to pandas.DataFrame.astype.

    Returns:
        Table with columns cast to specified types.

    Example:
        Cast single column:
        ```python
        table = table.astype({"population": int})
        ```

        Cast multiple columns:
        ```python
        table = table.astype({"year": int, "gdp": float})
        ```

        Cast all columns:
        ```python
        table = table.astype(str)
        ```
    """
    return super().astype(*args, **kwargs)  # ty: ignore

def check_metadata(self, ignore_columns: list[str] | None = None) -> None:
    """Check that all variables in the table have origins."""
    if ignore_columns is None:
        if self.primary_key:
            ignore_columns = self.primary_key
        else:
            ignore_columns = ["year", "country"]

    for column in [column for column in self.columns if column not in ignore_columns]:
        if not self[column].metadata.origins:
            warnings.warn(f"Variable {column} has no origins.", warnings.NoOriginsWarning)

def copy(self, deep: bool = True) -> Table:
    """Create a copy of the table with all metadata.

    Args:
        deep: If True (default), make a deep copy of the data and metadata.
            If False, creates a shallow copy.

    Returns:
        A new Table with copied data and metadata.

    Example:
        ```python
        table_copy = table.copy()  # Deep copy
        table_copy = table.copy(deep=False)  # Shallow copy
        ```
    """
    # This could be causing this warning:
    #   Passing a BlockManager to Table is deprecated and will raise in a future version. Use public APIs instead.
    # but I'm not sure how to fix it
    tab = super().copy(deep=deep)
    return tab.copy_metadata(self)

def copy_metadata(self, from_table: Table, deep: bool = False) -> Table:
    """Copy metadata from another table to this table.

    Copies both table-level metadata and variable-level metadata for all
    matching columns. Useful for preserving metadata after transformations.

    Args:
        from_table: Source table to copy metadata from.
        deep: If True, make a deep copy of the metadata. Default is False.

    Returns:
        Self, for method chaining.

    Example:
        ```python
        new_table = Table(transformed_df)
        new_table.copy_metadata(original_table)
        ```
    """
    return copy_metadata(to_table=self, from_table=from_table, deep=deep)