Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 25 additions & 24 deletions bin/build_dataset_package.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,9 @@

logger = logging.getLogger(__name__)

# Known columns present in each Hive-partitioned parquet table.
# Known columns present in each Delta table.
# fact, fact_resource and issue do not have start_date or end_date.
PARQUET_COLUMNS = {
TABLE_COLUMNS = {
"entity": [
"dataset", "end_date", "entity", "entry_date", "geometry",
"json", "name", "organisation_entity", "point", "prefix", "quality",
Expand All @@ -29,18 +29,22 @@
"entity", "entry_date", "entry_number", "field", "issue_type",
"line_number", "dataset", "resource", "value", "message",
],
"column_field": [
"dataset", "resource", "column", "field", "entry_date",
],
"dataset_resource": [
"dataset", "resource", "entry_date", "entity_count",
"entry_count", "line_count", "mime_type", "internal_path", "internal_mime_type",
],
}

# CSV-sourced tables — schema still read dynamically from SQLite after create_database()
CSV_SQLITE_TABLES = ["dataset_resource", "column_field", "old_entity"]
CSV_SQLITE_TABLES = ["old_entity"]


def _csv_sources(collection_data_path, collection, dataset, entity_min, entity_max):
def _csv_sources(collection_data_path, collection, entity_min, entity_max):
"""Return (sqlite_table, path, where_clause) tuples for tables loaded from CSV."""
base = f"{collection_data_path}/{collection}-collection"
return [
("dataset_resource", f"{base}/var/dataset-resource/{dataset}/*.csv", None),
("column_field", f"{base}/var/column-field/{dataset}/*.csv", None),
(
"old_entity",
f"{collection_data_path}/config/pipeline/{collection}/old-entity.csv",
Expand Down Expand Up @@ -132,41 +136,38 @@ def build_dataset_package(
logger.info(f"Loading tables from {base_path} into {output_path}")
conn = duckdb.connect()
conn.execute("INSTALL httpfs; LOAD httpfs;")
conn.execute("INSTALL delta; LOAD delta;")
conn.execute("INSTALL sqlite; LOAD sqlite;")
if isinstance(base_path, S3Path) or isinstance(AnyPath(collection_data_path), S3Path):
conn.execute("CREATE SECRET (TYPE S3, PROVIDER CREDENTIAL_CHAIN);")
conn.execute(f"ATTACH DATABASE '{output_path}' AS sqlite_db (TYPE SQLITE);")

# Load parquet tables — scan only the target dataset's partition directory
# to avoid hive partition schema mismatches across different dataset partitions
for table_name, cols in PARQUET_COLUMNS.items():
dataset_partition_path = base_path / table_name / f"dataset={dataset}"
# Load delta tables — each table is a Delta Lake table at {base_path}/{table_name}/
# partitioned by dataset. Filter to the target dataset via WHERE clause.
for table_name, cols in TABLE_COLUMNS.items():
delta_table_path = base_path / table_name
delta_log_path = delta_table_path / "_delta_log"

if not dataset_partition_path.exists():
logger.debug(f"No directory at {dataset_partition_path}, skipping '{table_name}'")
if not delta_log_path.exists():
logger.debug(f"No Delta table at {delta_table_path}, skipping '{table_name}'")
continue

scan_path = f"{dataset_partition_path}/**/*.parquet"
# dataset is a virtual hive column — supply it as a literal in the same
# position to keep col_list and select_list aligned
scan_path = str(delta_table_path)
col_list = ", ".join(f'"{c}"' for c in cols)
select_list = ", ".join(
f"'{dataset}' AS \"dataset\"" if c == "dataset" else f'"{c}"'
for c in cols
)

logger.info(f"Loading parquet files into '{table_name}'")
logger.info(f"Loading delta table '{table_name}' for dataset '{dataset}'")
conn.execute(f"""
INSERT INTO sqlite_db."{table_name}" ({col_list})
SELECT {select_list}
FROM parquet_scan('{scan_path}')
SELECT {col_list}
FROM delta_scan('{scan_path}')
WHERE dataset = '{dataset}'
""")

# Load CSV tables from the collection data path
if collection_data_path and collection:
logger.info(f"Loading CSV tables from {collection_data_path}")
for sqlite_table, path, where_clause in _csv_sources(
collection_data_path, collection, dataset, entity_min, entity_max
collection_data_path, collection, entity_min, entity_max
):
_load_csv_table(conn, sqlite_table, path, csv_table_columns, where_clause)
else:
Expand Down
12 changes: 5 additions & 7 deletions bin/package.sh
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
#!/bin/bash
# Load script - builds a dataset package directly from pre-built parquet tables in S3
# Load script - builds a dataset package directly from Delta tables in S3
# Requires S3 access via PARQUET_DATASETS_BUCKET

set -e
Expand Down Expand Up @@ -39,8 +39,6 @@ DATASET_DIR=${DATASET_DIR:-"dataset/"}
FLATTENED_DIR=${FLATTENED_DIR:-"flattened/"}
OUTPUT_LOG_DIR=${OUTPUT_LOG_DIR:-"log/"}

# Optional: bucket to save outputs to
COLLECTION_DATASET_BUCKET_NAME=${COLLECTION_DATASET_BUCKET_NAME:-""}

SQLITE_FILE="${DATASET_DIR}${DATASET_NAME}.sqlite3"
CSV_FILE="${DATASET_DIR}${DATASET_NAME}.csv"
Expand Down Expand Up @@ -106,15 +104,15 @@ digital-land expectations-dataset-checkpoint \

echo "Dataset $DATASET_NAME built successfully"

# Step 8: Save outputs to S3 if bucket is configured
if [ -n "$COLLECTION_DATASET_BUCKET_NAME" ]; then
echo "Step 8: Saving outputs to S3 bucket: $COLLECTION_DATASET_BUCKET_NAME"
# Step 8: Save outputs to S3
if [ -n "$COLLECTION_DATA_BUCKET" ]; then
echo "Step 8: Saving outputs to S3 bucket: $COLLECTION_DATA_BUCKET"
make save-dataset
make save-expectations
make save-performance
echo "All outputs saved to S3 successfully"
else
echo "Step 8: Skipping S3 upload (no COLLECTION_DATASET_BUCKET_NAME configured)"
echo "Step 8: Skipping S3 upload (no COLLECTION_DATA_BUCKET configured)"
fi

echo "Load complete!"
1 change: 1 addition & 0 deletions requirements-dev.txt
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
pytest>=8.0.0
pytest-mock>=3.12.0
pyarrow>=14.0.0
deltalake

# Code formatting and linting
black>=24.0.0
Expand Down
80 changes: 49 additions & 31 deletions tests/acceptance/test_build_dataset_package.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,9 @@
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq
import pytest
from click.testing import CliRunner
from deltalake import write_deltalake

from build_dataset_package import run_command
from digital_land.specification import Specification
Expand All @@ -28,9 +28,9 @@
# Helpers
# ---------------------------------------------------------------------------

def _write_parquet(path: Path, table: pa.Table):
path.parent.mkdir(parents=True, exist_ok=True)
pq.write_table(table, path)
def _write_delta(path: Path, table: pa.Table):
path.mkdir(parents=True, exist_ok=True)
write_deltalake(str(path), table, mode="append")


# ---------------------------------------------------------------------------
Expand All @@ -47,23 +47,22 @@ def specification_dir(tmp_path_factory):

@pytest.fixture()
def parquet_dir(tmp_path):
"""Build a minimal Hive-partitioned parquet tree under tmp_path.
"""Build a minimal Delta table tree under tmp_path.

Two datasets are written so tests can verify the dataset filter works.
The dataset column is derived from the directory name by DuckDB's
hive_partitioning — it is not a data column inside the parquet files.
dataset is a real column in each Delta table (not a hive partition directory).
"""
null_str2 = pa.array([None, None], type=pa.string())
null_str1 = pa.array([None], type=pa.string())

# entity table — two rows for DATASET, one decoy row for OTHER_DATASET
_write_parquet(
tmp_path / "entity" / f"dataset={DATASET}" / "data.parquet",
_write_delta(
tmp_path / "entity",
pa.table({
"dataset": pa.array([DATASET, DATASET], type=pa.string()),
"end_date": null_str2,
"entity": pa.array([1, 2], type=pa.int64()),
"entry_date": null_str2,
"geojson": null_str2,
"geometry": null_str2,
"json": null_str2,
"name": pa.array(["Entity One", "Entity Two"], type=pa.string()),
Expand All @@ -76,13 +75,13 @@ def parquet_dir(tmp_path):
"typology": null_str2,
}),
)
_write_parquet(
tmp_path / "entity" / f"dataset={OTHER_DATASET}" / "data.parquet",
_write_delta(
tmp_path / "entity",
pa.table({
"dataset": pa.array([OTHER_DATASET], type=pa.string()),
"end_date": null_str1,
"entity": pa.array([99], type=pa.int64()),
"entry_date": null_str1,
"geojson": null_str1,
"geometry": null_str1,
"json": null_str1,
"name": pa.array(["Should Not Appear"], type=pa.string()),
Expand All @@ -97,9 +96,10 @@ def parquet_dir(tmp_path):
)

# fact table — two rows for DATASET only
_write_parquet(
tmp_path / "fact" / f"dataset={DATASET}" / "data.parquet",
_write_delta(
tmp_path / "fact",
pa.table({
"dataset": pa.array([DATASET, DATASET], type=pa.string()),
"entity": pa.array([1, 2], type=pa.int64()),
"entry_date": null_str2,
"fact": pa.array(["fact-1", "fact-2"], type=pa.string()),
Expand All @@ -110,6 +110,38 @@ def parquet_dir(tmp_path):
}),
)

# dataset_resource table
_write_delta(
tmp_path / "dataset_resource",
pa.table({
"dataset": pa.array([DATASET], type=pa.string()),
"resource": pa.array(["abc123"], type=pa.string()),
"start_date": null_str1,
"end_date": null_str1,
"entry_date": null_str1,
"entity_count": pa.array([None], type=pa.int64()),
"entry_count": pa.array([None], type=pa.int64()),
"line_count": pa.array([None], type=pa.int64()),
"mime_type": null_str1,
"internal_path": null_str1,
"internal_mime_type": null_str1,
}),
)

# column_field table
_write_delta(
tmp_path / "column_field",
pa.table({
"dataset": pa.array([DATASET], type=pa.string()),
"resource": pa.array(["abc123"], type=pa.string()),
"column": pa.array(["Name"], type=pa.string()),
"field": pa.array(["name"], type=pa.string()),
"start_date": null_str1,
"end_date": null_str1,
"entry_date": null_str1,
}),
)

return tmp_path


Expand All @@ -124,20 +156,6 @@ def collection_dir(tmp_path, specification_dir):
entity_min = int(spec.get_dataset_entity_min(DATASET))
entity_max = int(spec.get_dataset_entity_max(DATASET))

base = tmp_path / "collection" / f"{COLLECTION_NAME}-collection"

# dataset-resource
dr_dir = base / "var" / "dataset-resource" / DATASET
dr_dir.mkdir(parents=True)
(dr_dir / "resource1.csv").write_text("resource,dataset\nabc123,brownfield-land\n")

# column-field
cf_dir = base / "var" / "column-field" / DATASET
cf_dir.mkdir(parents=True)
(cf_dir / "resource1.csv").write_text(
"dataset,resource,column,field\nbrownfield-land,abc123,Name,name\n"
)

# old-entity — one row in range, one outside
oe_dir = tmp_path / "collection" / "config" / "pipeline" / COLLECTION_NAME
oe_dir.mkdir(parents=True)
Expand Down Expand Up @@ -264,7 +282,7 @@ def test_tables_with_no_parquet_data_remain_empty(parquet_dir, collection_dir, s


def test_dataset_resource_csv_is_loaded(parquet_dir, collection_dir, specification_dir, tmp_path):
"""dataset-resource rows from the collection CSV should be loaded."""
"""dataset-resource rows from the Delta table should be loaded."""
output_path = tmp_path / "output" / f"{DATASET}.sqlite3"

result = CliRunner().invoke(run_command, [
Expand All @@ -285,7 +303,7 @@ def test_dataset_resource_csv_is_loaded(parquet_dir, collection_dir, specificati


def test_column_field_csv_is_loaded(parquet_dir, collection_dir, specification_dir, tmp_path):
"""column-field rows from the collection CSV should be loaded."""
"""column-field rows from the Delta table should be loaded."""
output_path = tmp_path / "output" / f"{DATASET}.sqlite3"

result = CliRunner().invoke(run_command, [
Expand Down
Loading