# Copyright (C) 2026, University of Maryland. All Rights Reserved.
# Authors: Tiago de Conto, Amelia Grace Holcomb
# For commercial licensing inquiries, contact UM Ventures at umdtechtransfer@umd.edu
import os, re, glob, h5py, h3, json
import shutil
import functools
import numpy as np
import pandas as pd
import geopandas as gpd
import h3pandas
import dask
import dask.dataframe
import dask_geopandas
import dask.bag as dbg
import pyarrow.parquet as pq
from typing import Union, List, Dict, Optional, Tuple, Any, Callable, Iterable
from earthaccess.store import EarthAccessFile
from dask.distributed import progress
from .config import GEDI_BEAMS, GH3_DEFAULT_DOWNLOAD_DIR, GH3_DEFAULT_TMP_DIR, GH3_DEFAULT_SOC_DIR, GH3_DEFAULT_H3_DIR, GEDI_PRODUCTS, GEDI_START_DATE, BUILD_LOG_FILENAME, PARTITION_META_FILENAME, _get_versioned, _GEDI_L2A_ESSENTIALS, _PRODUCT_QUALITY_FLAGS
from .utils import now, json_read, json_write, to_geojson, parquet_append_columns, parquet_merge_files, parquet_join_columns, parquet_fill_columns, read_parquet_schema, h5_is_valid, get_dask_client, generate_manifest, check_nan_only_columns, h3_partition_bbox, parse_h3_partition_dirname, AtomicFileWriter
from .h3utils import intersect_h3_geometries, h3_index_df, fix_h3_geometry
from .gedidriver import GEDIFile, add_special_columns, soc_file_tree, dask_h5_merged, gedi_vars_expand, gedi_vars_from_h5, gedi_vars_static, gedi_subset, validate_soc_files, load_h5, load_h5_merged, expand_var_wildcards, gedi_release_rank, gedi_release_tag
from .daac import gedi_download
from .logging_config import get_logger
from .validation import validate_h3_params, validate_product_vars, validate_directory_exists
from .exceptions import (H3ValidationError, GediValidationError, GediFileError, GediMergeError, GediError,
GediDatabaseError, GediProcessingError)
logger = get_logger(__name__)
def _init_earthaccess_worker():
"""Pre-authenticate earthaccess in Dask worker processes."""
import earthaccess
if not earthaccess.__auth__.authenticated:
earthaccess.login(strategy='netrc', persist=False)
# Install HTTP timeouts on the cloned-per-thread sessions used by
# earthaccess._download_file. Without this, a dead CloudFront edge
# connection blocks the worker indefinitely (CLOSE-WAIT recv hang).
from .daac import _install_request_timeouts
_install_request_timeouts()
def _complete_essentials_and_flags(product_vars: Dict, version=None, ensure_l2a: bool = True) -> Dict:
"""Union the release's L2A essentials and per-product quality flags into
every explicit variable list, purging names that belong only to another
release.
Shared by the download paths (:func:`_ensure_download_essentials`) and the
build path (:func:`_expand_product_vars`) so the two cannot drift. A
``minimal``/``default`` keyword resolved under a different release (a
provisional fallback, or a ``product_vars`` persisted by an older run)
bakes that release's essential/flag names into the list; a plain union
would only add the correct names alongside, leaving a stale name that
does not exist in the real archive and fails every beam that requests it.
Only names from the known essentials / quality-flag tables are purged —
names shared across releases (``degrade_flag``) and explicit non-preset
requests are never touched. ``None`` ("all variables") lists are left
alone. ``ensure_l2a=False`` never adds L2A but still completes the lists
that are present.
"""
essentials = _get_versioned(_GEDI_L2A_ESSENTIALS, version)
if 'L2A' in product_vars:
if product_vars['L2A'] is not None:
stale = {v for vs in _GEDI_L2A_ESSENTIALS.values() for v in vs} - set(essentials)
kept = [v for v in product_vars['L2A'] if v not in stale]
product_vars['L2A'] = list(dict.fromkeys(kept + list(essentials)))
elif ensure_l2a:
product_vars['L2A'] = list(essentials)
for prod, val in product_vars.items():
flag_map = _PRODUCT_QUALITY_FLAGS.get(prod)
if flag_map and val is not None:
flags = _get_versioned(flag_map, version)
flag_names = {f for f, _ in flags}
stale = {f for fl in flag_map.values() for f, _ in fl} - flag_names
val[:] = [v for v in val if v not in stale]
for flag_name, _condition in flags:
if flag_name not in val:
val.append(flag_name)
return product_vars
def _ensure_download_essentials(product_vars: Dict, version=None, ensure_l2a: bool = True) -> Dict:
"""Complete an expanded ``product_vars`` so the files it fetches can be built.
Applies the build-time completion (:func:`_complete_essentials_and_flags`:
L2A essentials, per-product quality flags, stale-release purge) at download
time, so a subset fetched with an *explicit* L2A list (or via S3 ETL, which
stores only what was asked for) is not missing the columns every build
needs, and appends ``shot_number`` to every explicit list.
"""
_complete_essentials_and_flags(product_vars, version=version, ensure_l2a=ensure_l2a)
for val in product_vars.values():
if val is not None and 'shot_number' not in val:
val.append('shot_number')
return product_vars
def _withhold_flags_absent_from_db(product_vars: Dict, h3_dir: str, version=None) -> Dict[str, List[str]]:
"""Drop quality flags an existing database does not store from an update's lists, in place.
:func:`_complete_essentials_and_flags` appends every quality flag of the
release to each list. A database built before a flag joined that table
stores no column for it, and an update must not start one: the merge
keeps the existing files' column set, so only new cells would hold the
column, and every read naming it (``--quality`` among them) would fail
on the older files. The flag is left to a variable update, which adds it
to every file at once.
A-priori: reads only the build log's ``h3_columns``. A database without
one (a fresh build) and a product it does not hold yet (no column with
its suffix: all its columns are new) are left alone. A product counts as
held by any column with its suffix, not by ``root_file_<p>``, which a
variable update of an older gedih3 did not write. A withheld flag is
logged with the variable update that adds it.
Parameters
----------
product_vars : dict
Expanded ``{product: [variable, ...]}`` lists, edited in place.
h3_dir : str
Database root.
version : int, optional
GEDI release of the database.
Returns
-------
dict
``{product: [flag, ...]}`` withheld, empty when nothing was.
"""
log_path = os.path.join(h3_dir, BUILD_LOG_FILENAME)
from .utils import json_read_cached
columns = set((json_read_cached(log_path) or {}).get('h3_columns') or []) if os.path.isfile(log_path) else set()
withheld = {}
for prod, val in product_vars.items():
flag_map = _PRODUCT_QUALITY_FLAGS.get(prod)
suffix = f"_{prod.lower()}"
if not columns or not flag_map or val is None or not any(c.endswith(suffix) for c in columns):
continue
absent = [f for f, _ in _get_versioned(flag_map, version) if f in val and f"{f}{suffix}" not in columns]
if absent:
val[:] = [v for v in val if v not in absent]
withheld[prod] = absent
if withheld:
desc = ' '.join(f"-{p.lower()} {' '.join(fs)}" for p, fs in sorted(withheld.items()))
logger.warning(
f"The database predates quality flag(s) {withheld}: new rows are written without them, so "
f"every file keeps one column set, and --quality skips them until a variable update adds them "
f"to the whole database (a run that also adds them as new variables does so after this step): "
f"gh3_build -i <soc> -o {h3_dir} {desc}"
)
return withheld
[docs]
def download_soc(product_vars: Dict, spatial=None, temporal=None, direct_access=False, update=False, version=None, odir=GH3_DEFAULT_SOC_DIR, n_jobs=5, on_granule_complete=None, ensure_l2a=True):
"""
Download GEDI HDF5 files in SOC (Science Operation Center) format.
Expands variable specifications, ensures L2A essentials and shot_number
are included, then delegates to ``gedi_download`` for retrieval.
Parameters
----------
product_vars : dict
Mapping of GEDI product codes (e.g., 'L2A', 'L4A') to variable
lists. Accepts 'default', 'minimal', or explicit variable names.
spatial : GeoDataFrame, list, or str, optional
Spatial filter (vector file, bbox as [W,S,E,N], or ISO3 code).
temporal : tuple of str, optional
Temporal range as (start_date, end_date) in 'YYYY-MM-DD' format.
direct_access : bool, default False
If True, use S3 streaming instead of downloading to disk.
update : bool, default False
If True, resume a previous download (skip already-downloaded files).
version : int or str, optional
GEDI data version (e.g., 3). If None, uses the package default (``GEDI_DEFAULT_VERSION``).
odir : str
Output directory for downloaded HDF5 files.
n_jobs : int, default 5
Number of parallel download workers.
on_granule_complete : callable, optional
Callback ``(granule_info_dict, status_str) -> None`` for per-granule
progress tracking. Passed through to ``gedi_download()``; see its
docstring for the ``granule_info_dict`` contract (includes a
``path`` key).
ensure_l2a : bool, default True
If True, automatically add L2A essentials when L2A is not in
product_vars. Set to False for variable-only updates where L2A
data already exists in the target database.
Returns
-------
list
List of downloaded SOC file paths or EarthAccessFile objects.
"""
product_vars = gedi_vars_expand(product_vars, version=version)
product_vars = _ensure_download_essentials(product_vars, version=version, ensure_l2a=ensure_l2a)
soc_files = gedi_download(
product_vars=product_vars,
odir=None if direct_access else odir,
spatial=spatial, temporal=temporal,
version=version,
resume=update, n_jobs=n_jobs,
to_list=direct_access,
on_granule_complete=on_granule_complete,
)
return soc_files
def _subset_s3_file(granule, prod, product_vars, odir, file_idx, n_files):
"""Subset a single S3 file to a compact local HDF5. Process-safe via granule serialization.
Accepts a serializable DataGranule object (not an EarthAccessFile handle).
Each worker re-authenticates and opens its own S3 handle, enabling true
multiprocessing via Dask workers instead of GIL-constrained threads.
"""
import earthaccess
from .logging_config import get_logger
logger = get_logger(__name__)
# Parse filename from granule data link
gf = GEDIFile(granule.data_links()[0])
local_dir = os.path.join(odir, str(gf.date.year), gf.date.strftime('%j'))
os.makedirs(local_dir, exist_ok=True)
local_path = os.path.join(local_dir, gf.full_name)
# Track whether we computed a union of vars from an existing file
union_vars = None
vars_for_prod = product_vars.get(prod)
if os.path.exists(local_path):
try:
existing_vars = set(gedi_vars_from_h5(local_path))
needed_vars = set(vars_for_prod) if vars_for_prod else set()
if needed_vars.issubset(existing_vars):
logger.debug(f"[{file_idx}/{n_files}] Skipping {gf.full_name} (already exists with required variables)")
return local_path
# Re-download with union of existing + requested
union_vars = sorted(existing_vars | needed_vars)
logger.debug(f"[{file_idx}/{n_files}] Re-downloading {gf.full_name} (missing {len(needed_vars - existing_vars)} variables)")
os.unlink(local_path)
except Exception:
logger.warning(f"[{file_idx}/{n_files}] Could not read existing {gf.full_name}, re-downloading")
os.unlink(local_path)
# Re-authenticate in worker process (same pattern as daac.py:379)
if not earthaccess.__auth__.authenticated:
earthaccess.login(strategy='netrc', persist=False)
# Open S3 handle with retry (transient failures common during worker startup)
from .exceptions import RETRY_DEFAULTS
max_attempts = RETRY_DEFAULTS['max_attempts']
initial_wait = RETRY_DEFAULTS['initial_wait']
max_wait = RETRY_DEFAULTS['max_wait']
s3_files = None
for attempt in range(1, max_attempts + 1):
try:
s3_files = earthaccess.open([granule], show_progress=False,
open_kwargs={'block_size': 16 * 1024 * 1024})
if s3_files:
break
except Exception as e:
if attempt == max_attempts:
logger.warning(f"[{file_idx}/{n_files}] Failed to open S3 handle for "
f"{gf.full_name} after {max_attempts} attempts: {e}")
return None
wait_time = min(initial_wait * (2 ** (attempt - 1)), max_wait)
logger.debug(f"[{file_idx}/{n_files}] S3 open attempt {attempt} failed for "
f"{gf.full_name}, retrying in {wait_time:.1f}s: {e}")
import time
time.sleep(wait_time)
if not s3_files:
logger.warning(f"[{file_idx}/{n_files}] No S3 handle returned for "
f"{gf.full_name} after {max_attempts} attempts")
return None
s3_file = s3_files[0]
# Use union vars if we merged with existing file, otherwise resolve from product_vars
if union_vars is not None:
vars_for_prod = union_vars
else:
vars_for_prod = product_vars.get(prod)
if vars_for_prod is None:
# NASA release file on S3 → static manifest is canonical and free.
# Falls back to remote H5 enumeration if no manifest ships for
# this (product, version).
vars_for_prod = gedi_vars_static(prod, version=gf.version)
if vars_for_prod is None:
vars_for_prod = gedi_vars_from_h5(s3_file)
logger.debug(f"[{file_idx}/{n_files}] Subsetting {gf.full_name} ({len(vars_for_prod)} vars)")
try:
result = gedi_subset(s3_file, local_path, vars_for_prod)
if result is None:
logger.warning(f"Subsetting produced no output for {gf.full_name}")
return None
return local_path
except Exception as e:
logger.warning(f"Failed to subset {gf.full_name}: {e}")
if os.path.exists(local_path):
os.unlink(local_path)
return None
[docs]
def s3_etl_subset(product_vars, spatial=None, temporal=None, version=None, odir=None, ensure_l2a=True):
"""
ETL-style S3 build: search for GEDI granules, then dispatch Dask workers
to open remote HDF5 files, extract only selected variables via range
requests, and write compact local HDF5 files.
This transfers and stores significantly less data than full downloads
(10-50x smaller depending on variable selection). Workers run as true
separate processes via Dask (GIL-free), matching the parallelism of
download mode.
Parameters
----------
product_vars : dict
Mapping of GEDI product codes (e.g., 'L2A', 'L4A') to variable
lists. Accepts 'default', 'minimal', or explicit variable names.
spatial : GeoDataFrame, list, or str, optional
Spatial filter (vector file, bbox as [W,S,E,N], or ISO3 code).
temporal : tuple of str, optional
Temporal range as (start_date, end_date) in 'YYYY-MM-DD' format.
version : int or str, optional
GEDI data version. If None, uses the package default (``GEDI_DEFAULT_VERSION``).
odir : str
Output directory for compact local HDF5 files in SOC structure.
ensure_l2a : bool, default True
If True, automatically add L2A essentials when L2A is not in
product_vars. Set to False for variable-only updates where L2A
data already exists in the target database.
Returns
-------
str
Path to the output directory containing compact HDF5 files.
Raises
------
GediFileError
If no GEDI files are found on S3 for the given parameters.
"""
from .daac import GEDIAccessor
# Expand variable specifications and complete them with the L2A
# essentials, shot_number and per-product quality flags — the compact
# files written here hold *only* what is listed, so anything the build
# needs must be requested now.
product_vars = gedi_vars_expand(product_vars, version=version)
product_vars = _ensure_download_essentials(product_vars, version=version, ensure_l2a=ensure_l2a)
# Search for granules per product (no S3 handles opened yet)
logger.info("Searching NASA DAAC for GEDI granules")
gass = GEDIAccessor(authenticate=True, spatial=spatial, temporal=temporal)
for prod in product_vars:
prod_version = version if version is not None else GEDI_PRODUCTS.get(prod.upper(), {}).get('version')
gass.search_data(product=prod, version=prod_version)
# Resolve None variables (dump all) from the static per-product manifest;
# NASA release files on S3 share the canonical schema, so no remote
# HDF5 metadata round-trip is needed. Falls back to opening one S3
# handle if no manifest ships for this (product, version).
for prod, prod_vars in product_vars.items():
if prod_vars is not None:
continue
granules = gass.product_files.get(prod, [])
if not granules:
continue
static_vars = gedi_vars_static(prod, version=version)
if static_vars is not None:
product_vars[prod] = static_vars
logger.info(f"Discovered {len(product_vars[prod])} variables for {prod} (static manifest)")
continue
s3_files = gass.link_s3(product=prod)
if s3_files:
product_vars[prod] = gedi_vars_from_h5(s3_files[0])
logger.info(f"Discovered {len(product_vars[prod])} variables for {prod} (HDF5 introspection)")
# Build flat task list of (granule, product) from search results
os.makedirs(odir, exist_ok=True)
tasks = []
for prod, granules in gass.product_files.items():
for granule in granules:
tasks.append((granule, prod))
n_files = len(tasks)
if n_files == 0:
raise GediFileError("No GEDI files found on S3 for the given parameters")
# Always-parallel dispatch to Dask workers (true multiprocessing,
# GIL-free). gedih3 CLIs (gh3_build, gh3_doctor, …) all create a
# Client at startup, so this assertion fires only for direct
# library/notebook callers — which can wrap in ``with Client(...)``.
client = get_dask_client()
if client is None:
raise GediError(
"s3_etl_subset requires a registered dask.distributed Client. "
"CLI tools create one at startup; library callers must wrap "
"in `with dask.distributed.Client(...) as client: ...`."
)
completed_count = 0
failed = 0
n_workers = sum(client.nthreads().values())
logger.info(f"Subsetting {n_files} files from S3 with Dask ({n_workers} workers)")
client.run(_init_earthaccess_worker)
futures = [
client.submit(
_subset_s3_file, granule, prod, product_vars,
odir, idx + 1, n_files, pure=False,
)
for idx, (granule, prod) in enumerate(tasks)
]
from distributed import as_completed as dask_as_completed
from tqdm import tqdm as tqdm_bar
pbar = tqdm_bar(total=n_files, desc="S3 ETL subset", unit="file")
try:
for future, result in dask_as_completed(futures, with_results=True):
completed_count += 1
if result is None:
failed += 1
pbar.set_postfix(failed=failed, refresh=False)
pbar.update(1)
finally:
pbar.close()
if failed > 0:
logger.warning(f"S3 ETL completed with {failed}/{n_files} failures")
# R2 producer-driven refresh: s3_etl_subset writes HDF5s to the SOC
# tree, so it must persist the manifest before returning. Without
# this, the next consumer (gh3_build, gh3_doctor) re-pays the
# multi-million-file recursive glob on cold GPFS.
from .gedidriver import write_soc_manifest
n_manifest = write_soc_manifest(odir)
if n_manifest:
logger.info(f"SOC manifest refreshed ({n_manifest} files)")
return odir
[docs]
def h3_part_files(df, dir_path, res=12, part=3, lat_col='lat_lowestmode', lon_col='lon_lowestmode', roi_tiles=[]):
"""
Write a DataFrame to H3-partitioned parquet files.
Indexes the DataFrame by H3 cell, groups rows by partition cell, and
writes each group as a separate parquet file. If a file already exists
for a partition, new rows are appended.
Parameters
----------
df : pandas.DataFrame
Input DataFrame with coordinate columns and ``root_beam``,
``root_file`` metadata columns.
dir_path : str
Base directory for output partition subdirectories.
res : int, default 12
H3 resolution for shot-level indexing.
part : int, default 3
H3 resolution for file partitioning.
lat_col : str, default 'lat_lowestmode'
Column name for latitude values.
lon_col : str, default 'lon_lowestmode'
Column name for longitude values.
roi_tiles : list of str, optional
If non-empty, only write partitions whose H3 cell ID is in this list.
Returns
-------
list of str or None
List of written parquet file paths, or None if the input is empty.
"""
if df.empty:
return
df = h3_index_df(df, res=res, part=part, lat_col=lat_col, lon_col=lon_col)
df = df.reset_index().set_index(f'h3_{part:02d}')
files = []
for i in df.index.unique():
if len(roi_tiles) > 0 and i not in roi_tiles:
continue
hex_path = os.path.join(dir_path,i)
hex_df = df.loc[[i]]
gedi_name = re.sub('\\.h5$',f'.{hex_df.root_beam.iloc[0]}.parquet', hex_df.root_file.iloc[0])
f = os.path.join(hex_path, gedi_name)
if f.endswith('.parquet'):
os.makedirs(hex_path, exist_ok=True)
if os.path.exists(f):
parquet_append_columns(hex_df, f)
else:
hex_df.to_parquet(f, engine='pyarrow', index=True, compression='zstd')
files.append(f)
del hex_df
del df
return files
[docs]
def h3_skip_part(h3_dir, h3_part, gedi_file, cols=None):
"""
Check whether an H3 partition already contains data from a GEDI granule.
Reads the partition's merged metadata to determine if the granule has
already been indexed and, optionally, if the requested columns are
already present.
Parameters
----------
h3_dir : str
Root directory of the H3 database.
h3_part : str
H3 cell ID of the partition to check.
gedi_file : str
Path or filename of the GEDI HDF5 file to test against.
cols : list of str, optional
If provided, also verify that these columns already exist in the
partition. The partition is only skipped when both the granule
is present and all requested columns are available.
Returns
-------
bool
True if the partition should be skipped (granule already indexed
and columns present), False otherwise.
"""
res = h3.get_resolution(h3_part)
meta_file = os.path.join(h3_dir, f"h3_{res:02d}={h3_part}", f"{h3_part}{PARTITION_META_FILENAME}")
if not os.path.exists(meta_file):
return False
metadata = json_read(meta_file)
skip_cols = True
if cols:
existing_cols = set(metadata['columns'])
cols = set(cols) - {'year', f'h3_{res:02d}'}
skip_cols = cols.issubset(existing_cols)
gf = GEDIFile(gedi_file)
granule_id = {'orbit':gf.orbit, 'granule':gf.orbit_granule, 'track':gf.track}
return skip_cols and granule_id in metadata['granules']
[docs]
def h3_add_skip_column(df, h3_dir):
"""
Add a ``_skip`` boolean column indicating partitions to skip.
Designed for use with ``dask.dataframe.map_partitions``. Checks each
unique H3 partition cell individually, since a single beam can span
multiple H3 cells. Delegates to ``h3_skip_part`` per cell.
Parameters
----------
df : pandas.DataFrame
A single Dask partition containing H3 partition column(s)
(``h3_XX``) and ``root_file_l2a``.
h3_dir : str
Root directory of the H3 database.
Returns
-------
pandas.DataFrame
Input DataFrame with an added ``_skip`` column (True if the
cell's granule data already exists in the database).
"""
if df.empty:
df['_skip'] = True
return df
h3_col = sorted([c for c in df.columns if re.match(r'h3_\d{2}', c)])[0]
gedi_file = df['root_file_l2a'].iloc[0]
cols = df.columns.tolist()
# Check each unique H3 partition cell individually
skip_map = {}
for h3_part in df[h3_col].unique():
skip_map[h3_part] = h3_skip_part(h3_dir=h3_dir, h3_part=h3_part, gedi_file=gedi_file, cols=cols)
df = df.assign(_skip=df[h3_col].map(skip_map))
return df
@dask.delayed
def dh3_merge_metadata(h3_subdir):
return h3_merge_metadata(h3_subdir)
def _dest_holds_fragments(out_file: str, files: List[str]) -> bool:
"""True when ``out_file`` provably already contains every fragment in ``files``.
Two gates, cheapest first. (1) Granules: each Stage 1 fragment's basename
encodes its granule (``_FRAGMENT_BASENAME_RE``) and the destination's
per-year metadata JSON lists the granules merged into it — the invariant
``h3_skip_part`` uses to skip extraction; one small JSON read, which
rejects the common case (a destination rewritten after extraction, whose
new granules are not listed) without opening a parquet. (2) Shots: a
granule is per-granule but a fragment is per-beam, so a destination that
holds some of a granule's beams would pass (1) for a missing beam; only
``shot_number`` membership proves the rows are there. That reads one
column, and only for a partition that already passed (1) — in practice
just the interrupted-cleanup case the skip exists for.
Anything unprovable answers False (no fragments, a missing or unreadable
metadata file, a legacy ``part.N`` fragment name, an unreadable column),
and the caller then merges: shot dedup keeps the rows already in the
destination, so a refused skip costs a rewrite, never rows.
Parameters
----------
out_file : str
Existing destination parquet of one ``(cell, year)`` partition.
files : list of str
Fragment paths about to be merged into it.
Returns
-------
bool
"""
if not files:
return False
meta_path = out_file.replace('.parquet', PARTITION_META_FILENAME)
try:
granules = json_read(meta_path).get('granules') or []
except Exception:
return False
have = set()
for g in granules:
try:
have.add((int(g['orbit']), int(g['granule']), int(g['track'])))
except (KeyError, TypeError, ValueError):
continue
for f in files:
m = _FRAGMENT_BASENAME_RE.match(os.path.basename(f))
if m is None or (int(m.group(1)), int(m.group(2)), int(m.group(3))) not in have:
return False
try:
import pyarrow.compute as pc
dest_shots = pq.read_table(out_file, columns=['shot_number'])['shot_number']
for f in files:
frag_shots = pq.read_table(f, columns=['shot_number'])['shot_number']
if not pc.all(pc.is_in(frag_shots, value_set=dest_shots.cast(frag_shots.type))).as_py():
return False
except Exception:
return False
return True
def _promote_orphan_merge_result(h3_file):
"""Promote ``<h3_file>.tmp`` when it is the only copy of the partition-year.
Before #46, a merge into an existing destination wrote the complete
merged result to ``<h3_file>.tmp`` (by ``parquet_merge_files``' final
atomic rename), deleted the destination and every fragment, and only
then renamed the result into place. A kill in between left this file as
the only copy. Merges no longer write ``<h3_file>.tmp``, and the
``AtomicFileWriter`` rewrites that do keep the destination in place
while their ``.tmp`` exists, so a ``<h3_file>.tmp`` with no destination
beside it is such an orphan. Promoting it (and rewriting its metadata,
which that kill skipped) recovers every row; deleting it, as the
stale-tmp sweep and ``gh3_doctor --check orphans`` would, loses them.
(A destination deleted on purpose while an ``AtomicFileWriter`` rewrite
of it was interrupted would also look like this; nothing does that.)
Parameters
----------
h3_file : str
Final path of the partition-year parquet.
Returns
-------
bool
True if an orphan was promoted.
"""
orphan = h3_file + '.tmp'
if os.path.exists(h3_file) or not os.path.exists(orphan):
return False
try:
pq.ParquetFile(orphan).metadata # readability check
except Exception as e:
logger.warning(f"Orphaned merge result {orphan} is unreadable ({e!r}); not promoting it")
return False
os.replace(orphan, h3_file)
h3_write_metadata(h3_file)
logger.warning(f"Recovered {h3_file} from the orphaned result of an interrupted merge")
return True
[docs]
def h3_merge_files(in_dir, out_dir, rm_src=True, replace=False, replace_granules=None):
"""
Merge multiple parquet files for an H3 partition into a single file.
Reads all parquet files in ``in_dir``, merges them (deduplicating shots
when appending to an existing file), writes the result to ``out_dir``,
and generates a sidecar metadata JSON file.
Parameters
----------
in_dir : str
Input directory containing parquet fragment files for one
H3 cell/year combination.
out_dir : str
Output directory where the merged file will be written,
preserving the ``<h3_cell>/<year>/`` structure.
rm_src : bool, default True
If True, remove the source directory after a successful merge.
replace : bool, default False
If True, overwrite existing output files. If False, merge new
data into any existing output file.
replace_granules : dict, optional
``{(orbit, granule, track): {product: basename}}``: re-index these
granules (#47). The existing file's rows from them are dropped
unless they already come from the given files, and the fragments'
rows take their place (``parquet_merge_files``). The existing file
is rewritten even when ``in_dir`` holds no fragment (the granules'
shots left this cell-year), the "already merged" skip never applies,
and a file left with no row is removed with its metadata (and its
year directory, when empty). Re-running gives the same result. An
unreadable existing file is handled as on any merge: the fragments
alone are written (without fragments, it is left as it is). Not
with ``replace=True`` (:class:`GediValidationError`).
Returns
-------
str or None
Path to the merged parquet file, or None if ``in_dir`` contained
no parquet files and there was no orphaned merge result to promote
(see ``_promote_orphan_merge_result``), or the file was left empty
by ``replace_granules`` and removed.
"""
# The output path is derived from in_dir's parents: with no trailing
# separator they would be off by one level, and the merge a silent no-op.
in_dir = os.path.join(in_dir, '')
files = glob.glob(os.path.join(in_dir,'*.parquet'))
year_dir = os.path.dirname(in_dir)
year = os.path.basename(year_dir.rstrip('/'))
h3_dir = os.path.dirname(year_dir)
h3part = os.path.basename(h3_dir.rstrip('/'))
odir = os.path.join(out_dir, h3part, year)
oname = f'{h3part.split('=')[-1]}.{year.split('=')[-1]}.0.parquet'
out_file = os.path.join(odir, oname)
h3_file = out_file
# Before the no-input return: the merge that left an orphan has already
# consumed its fragments, so the resume that should promote it is called
# on an empty in_dir.
promoted = not replace and _promote_orphan_merge_result(h3_file)
if replace_granules and replace:
raise GediValidationError("h3_merge_files: replace_granules needs replace=False (it edits the existing file)")
replacing = bool(replace_granules) and os.path.exists(out_file)
if len(files) == 0 and not replacing:
return h3_file if promoted else None
# Preventative: drop 0-byte source fragments before they reach
# parquet_merge_files (which fails the whole merge on the first bad
# file). 0-byte parquets are the dominant SIGKILL-leftover class —
# ``AtomicFileWriter`` writes to ``.tmp`` then ``os.replace``s, but a
# worker killed between the empty-file creation and any actual data
# write leaves a final-named 0-byte parquet. One ``stat`` per source
# fragment, effectively free since the file open hits the same
# metadata anyway. The worker logs each unlinked fragment so the
# operator can correlate with which (granule × beam) is being lost.
_filtered = []
for f in files:
try:
if os.path.getsize(f) == 0:
logger.warning(
f"h3_merge_files: dropping 0-byte source fragment {f}"
)
try:
os.unlink(f)
except OSError:
pass
continue
except OSError:
# File vanished between glob and stat — race with another
# worker's rm_src cleanup. Drop it silently; merge proceeds
# with the rest.
continue
_filtered.append(f)
files = _filtered
if len(files) == 0 and not replacing:
return h3_file if promoted else None
os.makedirs(odir, exist_ok=True)
# Per-partition stale tmp cleanup: a prior crash between .merge.tmp write
# and os.replace leaves orphaned <out_file>.tmp / .merge.tmp here. Scoping
# this to odir keeps the cost O(1) per merge (one listdir on the partition's
# own dir) rather than O(N_partitions) globbed on the driver. A
# <out_file>.tmp that was the only copy has already been promoted above,
# so what is left here is redundant with the destination and fragments.
# (Merges no longer write <out_file>.tmp; older ones did, see #46.)
try:
for _name in os.listdir(odir):
if _name.endswith('.merge.tmp') or _name.endswith('.parquet.tmp'):
try:
os.unlink(os.path.join(odir, _name))
except OSError:
pass
except OSError:
pass
# Disk-canonical skip: this merge already completed in a prior run and
# only the source-cleanup step was interrupted — clean up and return, no
# need to re-merge identical data through the dedup path. Being newer than
# every fragment is NOT proof on its own: any tool that rewrites database
# files between extraction and merge (dtype_drift --fix, gh3_update,
# backfill, product fills) makes every destination newer, and trusting
# mtime then deletes unmerged fragments. The proof is the destination's
# own metadata listing every fragment's granule (_dest_holds_fragments);
# the mtime test only short-circuits that JSON read. A corrupt newer-than-
# source dest must NOT short-circuit either (we'd return a broken partition).
# Never while replacing: the destination's metadata lists the granules
# being replaced (their old rows), which is no proof of the new ones.
if not replace and not replacing and os.path.exists(out_file):
try:
out_mtime = os.path.getmtime(out_file)
if (all(os.path.getmtime(f) <= out_mtime for f in files)
and _dest_holds_fragments(out_file, files)):
try:
pq.ParquetFile(out_file).metadata # readability check
if rm_src:
shutil.rmtree(in_dir, ignore_errors=True)
return h3_file
except Exception:
# Corrupt dest — fall through to the merge path, which
# detects the same corruption below and overwrites it.
pass
except OSError:
pass
if is_temp := (os.path.exists(out_file) and not replace):
# Validate the existing dest is readable before adding it to flist.
# A prior crash mid-write can leave a corrupt parquet here; trying to
# merge through it would raise and abort the whole merge phase.
# Treat corruption as "no usable dest" — merge only the fragments and
# overwrite the bad dest atomically via .merge.tmp + os.replace.
try:
pq.ParquetFile(out_file).metadata # cheap header check
# The existing partition goes FIRST and the order is kept:
# parquet_merge_files takes its target schema from files[0]
# (the database's schema must win), and check_shots dedup
# keeps the first occurrence (existing rows must win).
files = [out_file] + [f for f in files if f != out_file]
except Exception as e:
logger.warning(
f"Existing partition {h3_file} is unreadable ({e!r}); "
f"discarding it and merging tmp fragments fresh"
)
is_temp = False # treat as no dest, full overwrite
if not files:
return None # replacing an unreadable destination, with no fragment to merge
# Derive bbox from the H3 partition geometry directly (no data scan).
# Buffered enough to safely contain all level-N children at any depth
# via the empirical icosahedral-distortion factor in h3_partition_bbox.
cell_id, parent_res = parse_h3_partition_dirname(h3part)
bbox = h3_partition_bbox(cell_id, parent_res) if cell_id is not None else None
# Capture streaming stats from the merge so h3_write_metadata doesn't
# have to re-read the just-written file (saves ~1.5-2 GB peak memory
# per dense partition + the GPFS read-back I/O).
# parquet_merge_files writes <dest>.merge.tmp and renames it over the
# destination atomically. rm_src=False keeps the fragments until the
# metadata is written (the rmtree below); deleting inputs before the
# merged rows were in place lost the partition-year on a kill (#46).
stats = parquet_merge_files(
out_file, files, check_shots=is_temp, rm_src=False, bbox=bbox,
replace_granules=replace_granules, replace_in=out_file if is_temp else None,
)
if replacing and not stats['shot_count']:
# Every row was a replaced granule's: no file, no metadata, no year.
# Metadata first: a kill in between leaves a 0-row file, which the
# next run removes, rather than metadata listing gone granules.
for path in (h3_file.replace('.parquet', PARTITION_META_FILENAME), out_file):
try:
os.unlink(path)
except FileNotFoundError:
pass
try:
os.rmdir(odir)
except OSError:
pass
if rm_src:
shutil.rmtree(in_dir, ignore_errors=True)
return None
# Metadata before the fragment cleanup: a kill between the two leaves a
# destination whose metadata lists every fragment's granule, so the next
# run takes the _dest_holds_fragments skip instead of re-merging.
meta_file = h3_write_metadata(h3_file, stats=stats)
if rm_src:
shutil.rmtree(in_dir, ignore_errors=True)
return h3_file
@dask.delayed
def dh3_merge_files(in_dir, out_dir, rm_src=True, replace=False):
return h3_merge_files(in_dir=in_dir, out_dir=out_dir, rm_src=rm_src, replace=replace)
def _expand_product_vars(
product_vars: Dict[str, List[str]],
soc_files: List[Dict[str, str]],
version: Optional[int] = None
) -> Dict[str, List[str]]:
"""
Expand product variable specifications and ensure L2A essentials are included.
Parameters
----------
product_vars : dict
Raw product variable specifications (may contain 'default', 'minimal', etc.)
soc_files : list of dict
List of SOC file dictionaries to sample for 'all' variable expansion
version : int or None
GEDI data version. If None, auto-detected from the first SOC file.
Returns
-------
dict
Expanded product variables with L2A essentials included
"""
# Auto-detect version from first SOC file if not specified
if version is None and soc_files:
for soc_dict in soc_files:
for prod, path in soc_dict.items():
try:
version = GEDIFile(path).version
break
except Exception:
continue
if version is not None:
break
product_vars = gedi_vars_expand(product_vars, version=version)
def _first_valid_sample(product):
# Pick the first file for this product that opens cleanly and has at
# least one BEAM group, so wildcard / "all variables" introspection
# never trips on a corrupt head-of-list granule.
for sf in soc_files:
f = sf.get(product)
if f and h5_is_valid(f):
return f
return None
# Expand wildcard patterns (e.g. 'rh_*', 'geolocation/sensitivity_a?')
# against available HDF5 variables before further processing.
for k, val in product_vars.items():
if val is not None and any(any(c in v for c in ('*', '?', '[', ']')) for v in val):
file = _first_valid_sample(k)
if file:
available = gedi_vars_from_h5(file)
product_vars[k] = expand_var_wildcards(val, available)
# L2A essentials + per-product quality flags, stale-release names purged.
# (None means "all variables" — essentials and flags already included.)
_complete_essentials_and_flags(product_vars, version=version, ensure_l2a=True)
for k, val in product_vars.items():
if val is None:
file = _first_valid_sample(k)
if file is None:
raise GediFileError(
f"No valid SOC file found for product {k} to introspect variables"
)
product_vars[k] = gedi_vars_from_h5(file)
return product_vars
def _granule_id_from_l2a_path(path: str) -> Optional[Tuple[int, int, int]]:
"""Parse (orbit, granule, track) from a GEDI L2A HDF5 filename.
Cheap path-only parse that mirrors GEDIFile but skips the os.path.getsize
check, which is expensive on shared filesystems. Returns None if the
basename doesn't match the GEDI naming convention.
"""
try:
fl = os.path.basename(str(path)).split('_')
return (int(fl[3][1:]), int(fl[4]), int(fl[5][1:]))
except (IndexError, ValueError, AttributeError):
return None
def _granule_ids_in_fragment(parquet_file: str) -> set:
"""Read root_file_l2a from a single tmp parquet fragment, return granule IDs.
Tries column statistics first (no data read needed when min==max, which
holds for the common case of one granule per fragment under by_beam=True).
Falls back to reading the column. Returns an empty set on read failure.
"""
try:
pf = pq.ParquetFile(parquet_file)
if 'root_file_l2a' in pf.schema_arrow.names:
col_idx = pf.schema_arrow.get_field_index('root_file_l2a')
try:
stats = pf.metadata.row_group(0).column(col_idx).statistics
if stats is not None and stats.has_min_max and stats.min == stats.max:
gid = _granule_id_from_l2a_path(stats.min)
return {gid} if gid is not None else set()
except Exception:
pass
tbl = pq.read_table(parquet_file, columns=['root_file_l2a'])
out = set()
for path in tbl['root_file_l2a'].to_pylist():
if path is None:
continue
gid = _granule_id_from_l2a_path(path)
if gid is not None:
out.add(gid)
return out
except Exception:
pass
return set()
# Match fragment basenames written by stage 1 (v0.8.0+):
# ``O{orbit:05d}_G{granule:02d}_T{track:05d}.{beam}.parquet``
# (see _create_h3_dataframe). When the filename matches, BOTH the granule
# ID and the beam are recoverable without any parquet I/O — just string
# parsing on the worker, microseconds per file. The beam capture is
# load-bearing for the reconcile's partial-granule detection: without it
# a granule with a single beam fragment on disk would be marked INDEXED
# and the remaining 7 beams' shots would be silently dropped on resume.
_FRAGMENT_BASENAME_RE = re.compile(r'^O(\d+)_G(\d+)_T(\d+)\.([A-Za-z0-9_]+)\.parquet$')
# Sentinel used in the per-granule beam set returned by
# ``_process_h3_partition`` for legacy fragments whose filename does not
# encode the beam (pre-v0.8.0 ``part.NNN.parquet``). The reconcile treats
# the sentinel as "trust this granule is complete" — legacy builds rarely
# resume and reaching into parquet contents to recover the beam is more
# I/O than the rare case warrants.
_LEGACY_BEAM_SENTINEL = '*'
# ---------------------------------------------------------------------------
# Streaming partition-write shared helpers
# ---------------------------------------------------------------------------
#
# The streaming writer (_write_partitioned_streaming) emits one of these
# completion sentinels per successful (granule × beam) task, after every
# leaf parquet for that task has been atomically committed via
# AtomicFileWriter. The reconcile then trusts the sentinel as proof that
# the (granule × beam) is fully on disk — eliminating the legacy
# "any-beam-fragment-equals-complete-granule" data-loss path (Agent 3
# adversarial review #E.1).
#
# Two sentinel kinds, both zero-byte, both meaning "read successfully and
# every leaf is committed": ``<frag>.done`` = at least one leaf committed;
# ``<frag>.empty`` = zero rows (outside the ROI, empty beam, already-covered
# cells). Resume skips both; reconcile Pass C only counts a granule as
# on-disk when at least one of its beams is ``.done``, because an all-empty
# granule proves nothing about rows under a different scope.
_COMPLETE_SENTINEL_DIRNAME = '_complete'
def _frag_name_granule_key(frag_name: Optional[str]) -> Optional[Tuple[int, int, int]]:
"""``(orbit, granule, track)`` of a Stage 1 fragment name, else ``None``.
Inverse of :func:`_granule_beam_frag_name` as far as the granule goes;
same key the build log uses.
"""
m = _FRAGMENT_BASENAME_RE.match(f'{frag_name}.parquet') if frag_name else None
return (int(m.group(1)), int(m.group(2)), int(m.group(3))) if m else None
def _granule_beam_frag_name(soc_dict: Dict[str, str], beam: str) -> Optional[str]:
"""Stable basename (without ``.parquet``) for one (granule, beam) tuple.
Matches ``_FRAGMENT_BASENAME_RE`` exactly so the reconcile, merge, and
legacy ``to_parquet(name_function=...)`` paths all see identical
fragment paths. Returns ``None`` when the source HDF5 filename cannot
be parsed (matches the legacy fallback at the prior inline builder
site — caller falls back to dask-default naming, which in the
streaming path means we just skip the (granule × beam) since opaque
names would not round-trip through the reconcile cleanly).
"""
try:
path = next(iter(soc_dict.values()))
fl = os.path.basename(str(path)).split('_')
orbit = int(fl[3][1:])
granule = int(fl[4])
track = int(fl[5][1:])
return f"O{orbit:05d}_G{granule:02d}_T{track:05d}.{beam}"
except Exception:
return None
def _complete_sentinel_path(tmp_dir: str, frag_name: str, empty: bool = False) -> str:
"""Path of the per-(granule × beam) completion sentinel (``.empty`` for a zero-row task).
Lives under ``tmp_dir/_complete/`` (one directory, all sentinels) so
the reconcile can enumerate completions via a single ``os.scandir``
on the sentinel dir instead of a recursive walk of the partition
tree. ``frag_name`` matches ``_FRAGMENT_BASENAME_RE`` (no ``.parquet``
suffix) so the sentinel basename uniquely identifies the task.
"""
return os.path.join(tmp_dir, _COMPLETE_SENTINEL_DIRNAME,
f'{frag_name}.empty' if empty else f'{frag_name}.done')
def _emit_complete_sentinel(tmp_dir: str, frag_name: str, empty: bool = False) -> None:
"""Touch the completion sentinel for one (granule × beam). Idempotent.
``empty=True`` writes the ``.empty`` kind (task read OK, zero rows);
otherwise the ``.done`` kind (at least one leaf committed).
Atomic via ``open(... 'x')`` semantics — concurrent emitters on shared
GPFS race-create the same file; only one wins, the others observe
``FileExistsError`` and treat it as "already done". No need for
AtomicFileWriter here: the file is zero-byte (its existence is the
signal); a partial write cannot leave a half-emitted sentinel.
"""
path = _complete_sentinel_path(tmp_dir, frag_name, empty=empty)
parent = os.path.dirname(path)
os.makedirs(parent, exist_ok=True)
try:
with open(path, 'x'):
pass
except FileExistsError:
pass
# Per-failed-merge sentinel dir, parallel to ``_COMPLETE_SENTINEL_DIRNAME``.
# Each failed (h3_partition × year) merge writes one file naming the
# partition and the error so L1 resume + L2 doctor can recover without
# scanning the tmp tree. Atomic per-failure (no append-to-shared-file
# torn-line risk that plain ``open('a')`` would have on SIGKILL).
_MERGE_FAILURES_DIRNAME = '_merge_failures'
def _merge_failure_sentinel_name(tmp_dir: str, partition_dir: str) -> str:
"""Stable filename derived from the partition's path relative to tmp_dir.
``tmp/partitions/h3_03=8366c1fffffffff/year=2019`` →
``h3_03=8366c1fffffffff__year=2019.fail`` — readable and unique under
the canonical tree layout (no slashes left to escape).
"""
rel = os.path.relpath(partition_dir.rstrip('/'), tmp_dir)
return rel.replace('/', '__') + '.fail'
def _merge_failure_sentinel_path(tmp_dir: str, partition_dir: str) -> str:
return os.path.join(
tmp_dir, _MERGE_FAILURES_DIRNAME,
_merge_failure_sentinel_name(tmp_dir, partition_dir),
)
def _emit_merge_failure_sentinel(tmp_dir: str, partition_dir: str, error: BaseException) -> None:
"""Persist a one-shot record of a failed merge.
Contents: two lines — the absolute ``partition_dir`` and the formatted
exception. ``AtomicFileWriter`` is overkill here since the file is
small + single-write + per-failure-unique; ``open(path, 'w')`` is
sufficient and the worst case (write interrupted by SIGKILL) just
leaves a truncated record that L1 resume re-derives by re-attempting
the merge.
"""
path = _merge_failure_sentinel_path(tmp_dir, partition_dir)
parent = os.path.dirname(path)
os.makedirs(parent, exist_ok=True)
body = f"{partition_dir}\n{type(error).__name__}: {error}\n"
try:
with open(path, 'w') as f:
f.write(body)
except OSError:
# The failure sentinel is a safety net, not a correctness contract;
# losing one is non-fatal (next resume re-discovers via re-merge).
pass
_GRANULE_FAILURES_FILENAME = '_granule_failures.jsonl'
# Per-build sidecar enumerating granules whose status must be flipped back
# from INDEXED → MERGE_FAILED because their merged partition broke on a
# recoverable, known-bad fragment class. Written by ``_merge_and_finalize``
# during the failure handler; folded into the build-log JSON by the CLI
# (``apply_merge_failures_to_logger``) after merge returns. Decouples the
# merge driver from the H3BuildLogger so ``_merge_and_finalize`` keeps a
# narrow signature ``(tmp_dir, h3_dir)``.
_MERGE_FAILED_GRANULES_FILENAME = '_merge_failed_granules.jsonl'
# Marker tokens we use to recognize the fragment-corruption classes the
# resume layer can recover from (delete fragment + re-extract granule).
# Other exception kinds (disk full on dest, missing schema field, etc.)
# leave the granule status alone — those are infrastructure issues, not
# data issues, and flipping them would force pointless re-extraction.
_RECOVERABLE_FRAGMENT_ERROR_MARKERS = (
'Parquet file size is 0 bytes',
'Parquet magic bytes not found in footer',
'Couldn\'t deserialize thrift', # truncated footer mid-Thrift
)
def _is_recoverable_fragment_error(exc: BaseException) -> bool:
"""True when the exception text matches a known partial-write artifact
class. See ``_RECOVERABLE_FRAGMENT_ERROR_MARKERS`` for the curated list.
"""
msg = str(exc)
return any(marker in msg for marker in _RECOVERABLE_FRAGMENT_ERROR_MARKERS)
def _granules_in_partition_dir(partition_dir: str) -> List[Tuple[int, int, int]]:
"""Parse (orbit, granule, track) tuples from fragment basenames in one
h3_*/year=* partition directory. Returns a deduplicated list.
O(N_fragments_in_dir). Never iterates other partitions.
"""
out: set = set()
try:
entries = list(os.scandir(partition_dir))
except OSError:
return []
for e in entries:
if not e.is_file() or not e.name.endswith('.parquet'):
continue
m = _FRAGMENT_BASENAME_RE.match(e.name)
if m:
out.add((int(m.group(1)), int(m.group(2)), int(m.group(3))))
return sorted(out)
def _emit_merge_failed_granules(tmp_dir: str, partition_dir: str,
granules: List[Tuple[int, int, int]],
error: BaseException) -> None:
"""Append granule-flip-back records to the merge-failed-granules JSONL.
One record per granule (not per partition) so the CLI fold can drive
status updates directly. The whole list for a single failed partition
is written in one ``open + write`` to keep the per-failure cost O(1)
on the hot path.
"""
if not granules:
return
path = os.path.join(tmp_dir, _MERGE_FAILED_GRANULES_FILENAME)
err_str = f"{type(error).__name__}: {error}"
lines = [
json.dumps({
'orbit': g[0], 'granule': g[1], 'track': g[2],
'partition_dir': partition_dir,
'error': err_str,
})
for g in granules
]
try:
with open(path, 'a') as f:
f.write('\n'.join(lines) + '\n')
except OSError:
# Safety net, not a correctness contract — the next resume's merge
# will still re-discover via re-failing the same merge.
pass
def _read_merge_failed_granules(tmp_dir: str) -> List[Dict[str, Any]]:
"""Read every recorded granule-flip-back record. O(N_failures)."""
path = os.path.join(tmp_dir, _MERGE_FAILED_GRANULES_FILENAME)
out: List[Dict[str, Any]] = []
if not os.path.isfile(path):
return out
try:
with open(path, 'r') as f:
for line in f:
line = line.strip()
if not line:
continue
try:
out.append(json.loads(line))
except json.JSONDecodeError:
continue
except OSError:
pass
return out
[docs]
def preclean_merge_failures(tmp_dir: str) -> Dict[str, int]:
"""L1 resume pre-clean: act on every recorded ``_merge_failures`` sentinel.
For each partition the prior run's merge marked failed:
1. Stat each ``*.parquet`` in the partition — unlink if 0-byte (the
canonical class B artifact). Truncated-but-nonzero parquets (class C)
are unlinked only when ``pq.ParquetFile`` cannot open them — checked
cheaply by opening the footer once. The cost is bounded by the
failure list, never the full healthy tree. Each unlinked fragment's
``_complete/`` sentinel goes too, so the next Stage 1 re-runs that
(granule × beam) instead of skipping it as done.
2. Unlink any co-located ``*.tmp`` / ``*.merge.tmp`` siblings. These
survive SIGKILL when ``AtomicFileWriter.__exit__`` never runs and
would otherwise pollute future merges.
3. Delete the failure sentinel itself so re-running the pre-clean is
idempotent (next run only re-acts on freshly-failed merges).
Step 1 probes each fragment in O(1) (``_parquet_tail_ok``: the magic bytes
at both ends and a plausible footer length), which catches 0-byte,
truncated and never-finished files. The full footer parse is reserved for
partitions whose recorded error is a footer-level artifact (thrift) or
is missing:
parsing a 1.4k-column footer costs seconds per fragment on GPFS, and a
serial parse of every fragment behind a 14k-sentinel backlog (e.g. a
batch of schema failures, whose fragments are healthy) was ~58 h on the
driver. Partitions fan out over the registered dask Client.
Returns ``{'partitions_cleaned': N, 'parquets_removed': N, 'tmps_removed': N}``.
Companion to ``apply_merge_failures_to_logger`` — calling this without
the granule flip-back would unlink fragments and leave their granules
marked INDEXED, permanently dropping rows. ``build_h3db`` runs it before
Stage 1 (so the lost tasks are re-extracted in the same run) and
``_merge_and_finalize`` again at merge entry.
"""
out = {'partitions_cleaned': 0, 'parquets_removed': 0, 'tmps_removed': 0}
failures = _scan_merge_failure_sentinels(tmp_dir)
if not failures:
return out
# Only a footer-level error needs the full parse; the O(1) probe covers
# every other partial-write artifact. An empty record (a sentinel cut
# short) proves nothing, so it gets the full parse too.
items = [(d, not err or 'thrift' in err.lower()) for d, err in failures.items()]
client = get_dask_client()
if client is not None and len(items) > 1:
from .parallel import parallel_map
results = (r for _, r in parallel_map(items, _preclean_partition, desc='merge pre-clean',
unit='part', batch_size=256, tmp_dir=tmp_dir))
else:
results = (_preclean_partition(it, tmp_dir=tmp_dir) for it in items)
for r in results:
if isinstance(r, Exception):
logger.warning(f"Merge pre-clean failed for a partition: {type(r).__name__}: {r}")
continue
for k, v in r.items():
out[k] += v
return out
def _parquet_tail_ok(path: str, size: int) -> bool:
"""O(1) integrity probe: ``PAR1`` at both ends and a footer length that fits.
Parquet writes its footer last, so a write cut short (SIGKILL, full disk)
leaves a file without the trailing magic. Reads 12 bytes; parses nothing.
"""
if size < 12:
return False
try:
with open(path, 'rb') as fh:
head = fh.read(4)
fh.seek(-8, os.SEEK_END)
tail = fh.read(8)
except OSError:
return False
footer_len = int.from_bytes(tail[:4], 'little')
return head == b'PAR1' and tail[4:] == b'PAR1' and 0 < footer_len <= size - 12
def _preclean_partition(item, *, tmp_dir: str) -> Dict[str, int]:
"""Worker: steps 1–3 of :func:`preclean_merge_failures` for one partition.
``item`` is ``(partition_dir, full_probe)``; ``full_probe`` adds the footer
parse to the O(1) probe, for partitions whose failure was footer-level.
"""
partition_dir, full_probe = item
out = {'partitions_cleaned': 0, 'parquets_removed': 0, 'tmps_removed': 0}
if not os.path.isdir(partition_dir):
# Partition was deleted between runs (cleanup, manual rm); just
# drop the sentinel so it doesn't fire again.
try:
os.unlink(_merge_failure_sentinel_path(tmp_dir, partition_dir))
except OSError:
pass
return out
try:
entries = list(os.scandir(partition_dir))
except OSError:
return out
for e in entries:
if not e.is_file():
continue
name = e.name
try:
size = e.stat().st_size
except OSError:
continue
if name.endswith('.tmp') or name.endswith('.merge.tmp'):
# AtomicFileWriter orphan from SIGKILL — always safe to remove.
try:
os.unlink(e.path)
out['tmps_removed'] += 1
except OSError:
pass
continue
if not name.endswith('.parquet'):
continue
should_remove = not _parquet_tail_ok(e.path, size)
if not should_remove and full_probe:
try:
pq.ParquetFile(e.path).metadata
except Exception:
should_remove = True
if should_remove:
try:
os.unlink(e.path)
out['parquets_removed'] += 1
except OSError:
pass
try:
os.unlink(_complete_sentinel_path(tmp_dir, name[:-len('.parquet')]))
except OSError:
pass # (a fragment's task is a data task: never an ``.empty`` sentinel)
# Drop the sentinel — the cleanup acted; next merge will re-emit if
# it fails again. Keeping it would loop the pre-clean forever.
try:
os.unlink(_merge_failure_sentinel_path(tmp_dir, partition_dir))
except OSError:
pass
out['partitions_cleaned'] += 1
return out
[docs]
def apply_merge_failures_to_logger(h3_logger, tmp_dir: str) -> int:
"""Fold the merge-failed-granules sidecar into the build-log's granule
status: every recorded granule becomes ``MERGE_FAILED``. Returns the
number of entries changed.
Any prior status flips, not only ``INDEXED``: a granule extracted by the
failing run is still ``PENDING`` here, and ``set_post_build_info`` would
otherwise mark it ``INDEXED`` from its rows in other partitions. A
recorded granule the log does not track yet (S3 mode registers none) is
added. ``set_post_build_info``, the reconcile and ``log_state`` all keep
``MERGE_FAILED``; only :func:`_release_merge_failed`, after a Stage 1 run
that re-extracted the granule, clears it.
Idempotent. Removes the sidecar after the fold. The CLI calls it at
startup (a run killed before its fold leaves the sidecar behind) and
right after every merge. Does NOT call ``h3_logger.save_log`` — the
caller controls save cadence.
"""
records = _read_merge_failed_granules(tmp_dir)
if not records:
# Nothing parseable (empty, or one torn line from a kill mid-append):
# drop it, or it keeps ``_cleanup_merged_tmp`` from ever clearing
# tmp/partitions.
try:
os.unlink(os.path.join(tmp_dir, _MERGE_FAILED_GRANULES_FILENAME))
except OSError:
pass
return 0
if getattr(h3_logger, 'granule_info', None) is None:
h3_logger.granule_info = []
by_key = {(g.get('orbit'), g.get('granule'), g.get('track')): g for g in h3_logger.granule_info}
flipped = 0
for rec in records:
try:
key = (int(rec['orbit']), int(rec['granule']), int(rec['track']))
except (KeyError, ValueError, TypeError):
continue
g = by_key.get(key)
if g is None:
g = by_key[key] = {'orbit': key[0], 'granule': key[1], 'track': key[2]}
h3_logger.granule_info.append(g)
if g.get('status') != 'MERGE_FAILED':
g['status'] = 'MERGE_FAILED'
flipped += 1
# Truncate the sidecar — next resume re-derives only if a new merge
# fails. Leaving stale records would cause an old MERGE_FAILED flip
# to keep firing every resume forever.
try:
os.unlink(os.path.join(tmp_dir, _MERGE_FAILED_GRANULES_FILENAME))
except OSError:
pass
return flipped
def _release_merge_failed(h3_logger, listed=None) -> int:
"""``MERGE_FAILED`` → ``PENDING`` for the granules a Stage 1 run re-extracted.
Stage 1 never skips a ``MERGE_FAILED`` granule, and the pre-clean at its
start reopened the (granule × beam) tasks whose fragments were lost, so
once it returns their rows went through this run's merge. ``PENDING``
lets ``set_post_build_info`` mark them ``INDEXED``. Call it before the
post-merge fold, which re-flags any whose merge failed again.
Parameters
----------
h3_logger : H3BuildLogger
Build log whose ``granule_info`` is updated in place (not saved).
listed : set of tuple, optional
``(orbit, granule, track)`` of every granule Stage 1 listed. One it
could not see (source gone, excluded) was not re-extracted and keeps
its flag. ``None`` (no local listing: S3 mode) releases all.
Returns
-------
int
Number of granules released.
"""
n = 0
for g in getattr(h3_logger, 'granule_info', None) or []:
if g.get('status') == 'MERGE_FAILED' and (
listed is None or (g.get('orbit'), g.get('granule'), g.get('track')) in listed):
g['status'] = 'PENDING'
n += 1
return n
# Pattern: HDF5 "object 'X' doesn't exist" / "Unable to synchronously open
# object (object 'X' doesn't exist)" — used by ``_classify_load_h5_failure``
# to recognize the missing-variable case so the end-of-build advisory
# can offer a precise recipe.
# ``str(KeyError(msg))`` is the repr of ``msg``, which backslash-escapes the
# quotes h5py puts around the object name (``object \'x\' doesn\'t exist``).
_MISSING_VAR_RE = re.compile(r"object\s+\\?['\"]([^'\"\\]+)\\?['\"]\s+doesn\\?'?t\s+exist", re.IGNORECASE)
def _classify_load_h5_failure(exc: BaseException, soc_dict: Dict[str, str]) -> Dict[str, Any]:
"""Structured classification of a Stage1 ``load_h5_merged`` failure.
Today we only specialize the ``missing_var`` case (NASA-side L2A/L2B/L4A
schema variance across orbit clusters — orbits O20752–O20767 of L2A,
which lack ``l2a_quality_flag_rel3_a10`` despite the shipped manifest
claiming they have it). Other failures get a generic ``other`` kind
with the raw exception text — still queryable, just not auto-recoverable.
Returns a JSON-serializable dict so the driver can append it directly
to the per-build ``_granule_failures.jsonl`` sidecar.
"""
msg = str(exc)
kind: str = 'other'
var: Optional[str] = None
product: Optional[str] = None
if isinstance(exc, KeyError) or 'KeyError' in type(exc).__name__:
m = _MISSING_VAR_RE.search(msg)
if m:
kind = 'missing_var'
var = m.group(1)
# Best-effort product inference: which product file does the source path
# belong to? soc_dict keys are the product codes (L2A/L2B/L4A); pick the
# first one whose path appears in the message, else None.
if msg:
for prod_code, h5_path in (soc_dict or {}).items():
try:
if h5_path and os.path.basename(str(h5_path)) in msg:
product = prod_code
break
except Exception:
continue
return {
'kind': kind,
'var': var,
'product': product,
'error_type': type(exc).__name__,
'error_message': msg,
}
def _append_granule_failure(tmp_dir: str, frag_name: str, failure: Dict[str, Any]) -> None:
"""Append one failure record to ``tmp_dir/_granule_failures.jsonl``.
Single-writer (driver thread) so no concurrency guard needed. Append-only
and O(1) per failure: a SIGKILL loses at most the in-flight line, and the
hot path never rewrites the file. The same task failing on several runs
appends several lines; :func:`_read_granule_failures` collapses them to
one record per ``frag_name`` and :func:`_compact_granule_failures` rewrites
the file once at the end of Stage 1. The file is a forensics sidecar read
by the end-of-build advisory and ``gh3_doctor``'s ``tmp_partitions_health``;
nothing folds it into the build log.
"""
path = os.path.join(tmp_dir, _GRANULE_FAILURES_FILENAME)
record = {'frag_name': frag_name, **failure}
try:
with open(path, 'a') as f:
f.write(json.dumps(record) + '\n')
except OSError:
# Same logic as merge-failure sentinel: this is a safety net, not a
# correctness contract — the in-memory ``n_fail`` counter is still
# accurate, and the WARN log line still surfaces the error.
pass
def _scan_granule_failures(tmp_dir: str) -> Tuple[List[Dict[str, Any]], int, int]:
"""Collapse the failure sidecar to ``(records, n_lines, n_invalid)``.
One record per ``frag_name`` in first-seen order, carrying the LAST
line's content and ``runs`` = total occurrences (a compacted record's own
``runs`` counts, a plain line counts 1). ``n_lines`` is the number of
valid lines read and ``n_invalid`` the torn / non-object lines skipped
(blank lines are ignored), so callers can tell whether a rewrite would
change the file. A record without ``frag_name`` is never merged with
another: it is keyed by its line number and stays visible.
"""
path = os.path.join(tmp_dir, _GRANULE_FAILURES_FILENAME)
by_key: Dict[Any, Dict[str, Any]] = {}
n_lines = n_invalid = 0
if not os.path.isfile(path):
return [], 0, 0
try:
with open(path, 'r') as f:
for line in f:
line = line.strip()
if not line:
continue
try:
rec = json.loads(line)
except json.JSONDecodeError:
# Tolerate the rare torn last line from a SIGKILL during
# the writer's flush — same principle as parquet_merge
# tolerating corrupt fragments.
n_invalid += 1
continue
if not isinstance(rec, dict):
n_invalid += 1
continue
n_lines += 1
key = rec.get('frag_name')
if key is None:
key = ('__line__', n_lines)
try:
runs = max(int(rec.get('runs', 1)), 1)
except (TypeError, ValueError):
runs = 1
prev = by_key.get(key)
if prev is not None:
runs += prev['runs']
rec['runs'] = runs
by_key[key] = rec # dict keeps the first-seen position
except OSError:
pass
return list(by_key.values()), n_lines, n_invalid
def _read_granule_failures(tmp_dir: str) -> List[Dict[str, Any]]:
"""Read the recorded failures, one record per distinct task.
O(N_failures); never iterates partitions. Records are keyed by
``frag_name`` (a Stage 1 ``(granule x beam)`` fragment name, or a granule
id from the variable-add fan). The latest record per key wins and its
``runs`` field is the total number of times that key was recorded, so
``len(result)`` is the number of distinct failed tasks.
"""
return _scan_granule_failures(tmp_dir)[0]
def _compact_granule_failures(tmp_dir: str, recovered: Iterable[Any] = ()) -> None:
"""Rewrite the failure sidecar once: drop recovered keys, merge duplicates.
``recovered`` holds keys whose task succeeded this run; only those are
dropped (a key of any other namespace, e.g. a variable-fan granule id,
never matches and is kept). One compacted record per remaining key
carries the latest error and its ``runs`` count. The file is deleted when
nothing remains, so ``_cleanup_merged_tmp`` can remove the tmp tree. No-op
when the rewrite would change nothing. A crash before this runs leaves the
append-only file, which :func:`_read_granule_failures` still reads
correctly.
"""
path = os.path.join(tmp_dir, _GRANULE_FAILURES_FILENAME)
recovered = set(recovered)
try:
records, n_lines, n_invalid = _scan_granule_failures(tmp_dir)
if not os.path.isfile(path):
return
kept = [r for r in records if r.get('frag_name') not in recovered]
if len(kept) == n_lines and len(kept) == len(records) and not n_invalid:
return # one clean line per key, nothing recovered
if not kept:
os.unlink(path)
return
with AtomicFileWriter(path) as tmp_path:
with open(tmp_path, 'w') as f:
for rec in kept:
f.write(json.dumps(rec) + '\n')
except OSError as e:
logger.warning(f"Could not compact {path}: {e}")
def _scan_merge_failure_sentinels(tmp_dir: str) -> Dict[str, str]:
"""Return ``{partition_dir: error_string}`` for every recorded failure.
O(N_failures) — never iterates healthy partitions. Used by L1 resume
and L2 doctor (``tmp_partitions_health``).
"""
out: Dict[str, str] = {}
sentinel_dir = os.path.join(tmp_dir, _MERGE_FAILURES_DIRNAME)
if not os.path.isdir(sentinel_dir):
return out
try:
entries = list(os.scandir(sentinel_dir))
except OSError:
return out
for entry in entries:
if not entry.is_file() or not entry.name.endswith('.fail'):
continue
try:
with open(entry.path, 'r') as f:
lines = f.read().splitlines()
except OSError:
continue
if not lines:
continue
partition_dir = lines[0].strip()
err = lines[1].strip() if len(lines) >= 2 else ''
if partition_dir:
out[partition_dir] = err
return out
def _scan_complete_sentinels(tmp_dir: str, with_empty: bool = False):
"""Return the set of frag_names with an emitted completion sentinel.
One ``os.scandir`` over ``tmp_dir/_complete/``; O(n_completed_tasks)
rather than O(n_fragments). Counts both kinds (``.done`` and ``.empty``).
Empty set if the sentinel dir doesn't exist yet (fresh build or
pre-migration legacy tmp tree).
With ``with_empty=True`` returns ``(all_complete, empty_only)`` where
``empty_only`` is the subset whose only sentinel is ``.empty`` (zero rows).
"""
sentinel_dir = os.path.join(tmp_dir, _COMPLETE_SENTINEL_DIRNAME)
done: set = set()
empty: set = set()
try:
with os.scandir(sentinel_dir) as it:
for e in it:
if not e.is_file(follow_symlinks=False):
continue
if e.name.endswith('.done'):
done.add(e.name[:-len('.done')])
elif e.name.endswith('.empty'):
empty.add(e.name[:-len('.empty')])
except OSError:
pass
if with_empty:
return done | empty, empty - done
return done | empty
def _canonical_write_schema(meta_df, part: int) -> Any:
"""Build the canonical pyarrow schema for streaming per-leaf writes.
Mirrors dask's ``_meta_nonempty → pyarrow_schema_dispatch`` chain used
by ``dd.to_parquet``. Building schema once on the driver and passing
it to every per-leaf write enforces column-order and nullable-dtype
parity across all fragments — without this guard, per-leaf
schema-inference from data would produce subtle divergences (e.g.
``Int64`` nullable vs ``int64``, datetime tz) that break
parquet_merge_files's schema union at the merge phase.
Parameters
----------
meta_df : pandas / geopandas DataFrame
Empty frame with the canonical column set, dtypes, and index name
that every per-leaf write should match. Should already include
``geometry`` + ``datetime`` + ``year`` (post-add_special_columns,
post year-assign) and the partition columns (we drop them here).
part : int
Partition resolution; drives the ``h3_{part:02d}`` partition
column name to drop from the body schema.
Returns
-------
pyarrow.Schema
Schema with partition columns removed. For GeoDataFrame input the
schema carries the GeoParquet ``geo`` metadata so per-leaf cast
preserves the geometry encoding contract.
"""
drop_cols = [f'h3_{part:02d}', 'year']
body = meta_df.drop(columns=drop_cols, errors='ignore').head(0)
if isinstance(body, gpd.GeoDataFrame):
from geopandas.io.arrow import _geopandas_to_arrow
schema = _geopandas_to_arrow(body, index=True).schema
else:
import pyarrow as pa
schema = pa.Schema.from_pandas(body, preserve_index=True)
return _schema_resolve_null_fields(schema)
def _schema_resolve_null_fields(schema):
"""Retype any ``null`` field in *schema* as ``string``.
The canonical schema is inferred from a row-less frame, and on pandas < 3
an empty object column carries no type information, so pyarrow infers
``null`` where the same column with data infers ``string``. Casting real
strings into a null-typed field then fails every per-leaf write with
``ArrowNotImplementedError: Unsupported cast from string to null``.
pandas 3 keeps its dedicated string dtype through ``head(0)`` and so never
produces the null fields this repairs.
Parameters
----------
schema : pyarrow.Schema
Schema inferred from the empty meta frame.
Returns
-------
pyarrow.Schema
Same schema with null-typed fields retyped to ``string``, metadata
(including the GeoParquet ``geo`` key) preserved.
"""
import pyarrow as pa
null_positions = [i for i, field in enumerate(schema) if pa.types.is_null(field.type)]
for i in null_positions:
schema = schema.set(i, schema.field(i).with_type(pa.string()))
return schema
def _parquet_storage_type(arrow_type):
"""``arrow_type`` as it reads back from Parquet.
Parquet has no seconds timestamp unit, so pyarrow stores ``timestamp[s]``
(what the row-less Stage 1 meta infers for ``datetime``) as
``timestamp[ms]``. Comparisons between a derived schema and files on disk
must use the stored form, or every database reports a phantom drift.
"""
import pyarrow as pa
if pa.types.is_timestamp(arrow_type) and arrow_type.unit == 's':
return pa.timestamp('ms', tz=arrow_type.tz)
return arrow_type
# Arrow offset width is a representation, not a value type: pandas 3 infers
# ``large_string`` where pandas 2 inferred ``string`` (both are supported), and
# Parquet stores both as the same UTF8 byte array.
_OFFSET_WIDTH_TWINS = {'large_string': 'string', 'large_binary': 'binary'}
def _same_value_type(a: str, b: str) -> bool:
"""Whether two stored dtype strings hold the same values, offset width aside."""
return _OFFSET_WIDTH_TWINS.get(a, a) == _OFFSET_WIDTH_TWINS.get(b, b)
def _stage1_write_schema(ddf, lat_col: str, lon_col: str, dat_col: str, part: int):
"""The pyarrow schema Stage 1 writes every leaf in, from a source ddf's meta.
Applies the same transforms the streaming worker applies to data
(special columns, ``year``) to the row-less meta, then
:func:`_canonical_write_schema`. Single source of truth for "what a build
writes", shared by the build and :func:`_source_write_schema`.
"""
meta = ddf._meta.copy()
meta = add_special_columns(meta, lon_col=lon_col, lat_col=lat_col, dat_col=dat_col)
if 'datetime' in meta.columns:
meta = meta.assign(year=meta['datetime'].dt.year.astype('int32'))
return _canonical_write_schema(meta, part=part)
def _source_write_schema(soc_dict: Dict[str, str], product_vars: Dict[str, List[str]],
res: int, part: int, version: int):
"""Stage 1 write schema for one source granule, exactly as a build derives it.
Used to answer "which dtypes would ``gh3_build`` write for this release"
without building anything: one granule goes through the build's own
variable expansion, meta inference (one HDF5 read on the driver) and
:func:`_stage1_write_schema`.
Parameters
----------
soc_dict : dict
``{product: hdf5_path}`` for one granule, all of release ``version``.
product_vars : dict
Products and variables as recorded in the build log.
res, part : int
H3 index and partition levels of the database.
version : int
GEDI release of the database (and of ``soc_dict``).
Returns
-------
pyarrow.Schema
"""
import copy
pv = _expand_product_vars(copy.deepcopy(product_vars), [soc_dict], version=version)
soc = {p: f for p, f in soc_dict.items() if p in pv}
ddf, lat_col, lon_col, dat_col, _ = _create_h3_dataframe([soc], pv, res, part)
return _stage1_write_schema(ddf, lat_col, lon_col, dat_col, part)
def _align_schema_to_db(schema, h3_dir: str):
"""Retype *schema* fields to the dtypes the existing database records.
The Stage 1 schema is derived from the source HDF5 of the build's
release. When an update meets a database whose recorded dtype for a
column differs (e.g. a database built by an older gedih3 whose schema
came from another release), new data must still land: fragments are
written in the database's types so every later merge is homogeneous.
Casting is left to the per-leaf ``Table.cast`` (``safe=True``), so a value
that does not fit fails its task instead of being truncated.
A-priori: reads only the build log's ``h3_columns_dtypes``, no parquet I/O.
Parameters
----------
schema : pyarrow.Schema
Source-derived canonical write schema.
h3_dir : str
Database root. No build log, or no recorded dtypes, is a no-op.
Returns
-------
tuple of (pyarrow.Schema, dict)
The aligned schema (metadata preserved) and ``{column: (source, db)}``
for every retyped column whose values differ in type. An offset-width
difference alone (``large_string`` vs ``string``) is aligned but not
reported. A recorded ``null`` is never a target: real values cannot
cast to it, so every leaf write would fail.
"""
import pyarrow as pa
log_path = os.path.join(h3_dir, BUILD_LOG_FILENAME)
if not os.path.isfile(log_path):
return schema, {}
db_dtypes = json_read(log_path).get('h3_columns_dtypes') or {}
drift = {}
for i, field in enumerate(schema):
db_type = db_dtypes.get(field.name)
src_type = str(_parquet_storage_type(field.type))
if db_type in (None, 'null', src_type):
continue
try:
target = pa.type_for_alias(db_type)
except (ValueError, KeyError):
continue
schema = schema.set(i, field.with_type(target))
if not _same_value_type(src_type, db_type):
drift[field.name] = (str(field.type), db_type)
return schema, drift
def _sample_partition_schema(h3_dir: str, log: dict):
"""Footer schema of one existing partition file, or ``None``.
Located a-priori from the build log's ``h3_partition_ids`` and partition
level: one cell dir, one year dir, one footer read — never a tree walk.
"""
part = log.get('h3_partition_level')
for cell in (log.get('h3_partition_ids') or [])[:5]:
try:
cell_dir = os.path.join(h3_dir, f"h3_{part:02d}={cell}")
year_dirs = sorted(e.path for e in os.scandir(cell_dir) if e.is_dir() and e.name.startswith('year='))
for ydir in year_dirs:
pfs = sorted(e.path for e in os.scandir(ydir) if e.name.endswith('.parquet'))
if pfs:
return pq.read_schema(pfs[0])
except (OSError, TypeError, ValueError):
continue
return None
def _complete_schema_from_db(schema, h3_dir: str):
"""Complete the Stage 1 write schema with the database's columns, in its column order.
A build that admits granules whose later products (L2B, L4A, L4C…) are not
published yet must still write fragments carrying the database's full
column set: the merge requires every destination column in every fragment,
and a new cell must not start with a narrower schema than its neighbours.
Its files must also keep the database's column order: a source sample
lacking some products orders columns differently, and readers (``gh3_load``
over files of both orders) reject a column order that differs from their
metadata. The order comes from one partition footer, located a-priori;
types come from the build log's ``h3_columns_dtypes``, or from that footer
for a dtype string that is not a plain pyarrow alias.
Parameters
----------
schema : pyarrow.Schema
Stage 1 write schema (already aligned by :func:`_align_schema_to_db`).
h3_dir : str
Database root. No build log is a no-op.
Returns
-------
tuple of (pyarrow.Schema, list of str)
The completed schema (metadata preserved) and the appended column names.
Raises
------
GediValidationError
When a database column's type can be resolved from neither the build
log nor a partition footer (the fragments would fail every merge).
"""
import pyarrow as pa
log_path = os.path.join(h3_dir, BUILD_LOG_FILENAME)
if not os.path.isfile(log_path):
return schema, []
log = json_read(log_path)
db_dtypes = log.get('h3_columns_dtypes') or {}
footer = _sample_partition_schema(h3_dir, log) or pa.schema([])
# Legacy logs record columns without dtypes: types then come from the footer.
missing = [c for c in sorted(set(db_dtypes) | set(log.get('h3_columns') or [])) if c not in schema.names]
added = []
for name in missing:
try:
typ = pa.type_for_alias(db_dtypes[name])
except (ValueError, KeyError, TypeError):
if name not in footer.names:
raise GediValidationError(
f"Cannot resolve the database type of column {name!r} (build log records "
f"{db_dtypes.get(name)!r}); refusing to write fragments that would fail every merge"
)
typ = footer.field(name).type
schema = schema.append(pa.field(name, typ, nullable=True))
added.append(name)
# The database's column order; columns it does not hold yet go last.
order = [n for n in footer.names if n in schema.names]
order += [n for n in schema.names if n not in set(order)]
if order != schema.names:
schema = pa.schema([schema.field(n) for n in order], metadata=schema.metadata)
return schema, added
def _conform_table_to_schema(table, schema):
"""Reorder ``table`` to ``schema``, adding all-null columns it lacks, then cast.
Used only by builds that admit granules missing some products: their
leaves carry no columns for those products. A column the schema does not
have is an error, never silently dropped. The cast is the same safe cast
the plain path applies.
"""
import pyarrow as pa
present = {name: i for i, name in enumerate(table.column_names)} # one pass: leaves have ~1.5k columns
extra = set(present) - set(schema.names)
if extra:
raise GediValidationError(f"Leaf has columns outside the write schema: {sorted(extra)}")
n = table.num_rows
arrays = [table.column(present[f.name]) if f.name in present else pa.nulls(n, f.type) for f in schema]
return pa.Table.from_arrays(arrays, names=schema.names).cast(schema)
def _derive_merged_output_paths(merge_progress_file: str, h3_dir: str) -> List[str]:
"""Derive absolute output parquet paths from the merge-progress file.
Pure in-memory transform of the merge-progress lines (one tmp partition
dir per line) into the deterministic final paths via the
``h3_merge_files`` naming contract:
``<tmp>/h3_<p>=X/year=Y`` → ``<h3_dir>/h3_<p>=X/year=Y/X.Y.0.parquet``.
Used by ``_merge_and_finalize`` to (a) write the final database
manifest without walking the tree at end-of-merge, and (b) refresh
the manifest incrementally during long merge phases so consumers
reading mid-build see partial-but-fresh state. Zero GPFS metadata ops.
"""
out: List[str] = []
if not os.path.exists(merge_progress_file):
return out
try:
with open(merge_progress_file, 'r') as f:
for line in f:
tmp_p = line.strip()
if not tmp_p:
continue
year_bn = os.path.basename(tmp_p.rstrip('/'))
ydir = os.path.dirname(tmp_p.rstrip('/'))
h3part = os.path.basename(ydir.rstrip('/'))
if not h3part.startswith('h3_') or '=' not in h3part:
continue
h3val = h3part.split('=', 1)[-1]
yval = year_bn.split('=', 1)[-1] if '=' in year_bn else year_bn
out.append(os.path.join(
h3_dir, h3part, year_bn, f'{h3val}.{yval}.0.parquet'
))
except OSError:
return []
return out
def _scan_partition_meta_granules(
partition_dir: str,
*,
meta_filename: str = PARTITION_META_FILENAME,
) -> set:
"""Worker: parse granule IDs from every PARTITION_META JSON under one
h3_* partition. Module-level so it pickles for dask.
Looks in two locations:
1. ``partition_dir/*<meta>`` — partition-level meta (no year subdir)
2. ``partition_dir/*/*<meta>`` — year-level meta (h3_*/year=*/...)
Returns a set of ``(orbit, granule, track)`` tuples; empty on any read
failure (the caller treats per-partition errors as non-fatal, matching
the existing reconcile semantics).
"""
out: set = set()
for pat in (f'*{meta_filename}', os.path.join('*', f'*{meta_filename}')):
for mf in glob.glob(os.path.join(partition_dir, pat)):
try:
for g in (json_read(mf) or {}).get('granules', []):
try:
out.add((g['orbit'], g['granule'], g['track']))
except (KeyError, TypeError):
continue
except Exception:
continue
return out
def _list_year_subdirs(h3_dir: str) -> List[str]:
"""Return ``<h3_dir>/year=*/`` paths (with trailing separator).
Worker-side body of the merge-phase tmp partition listing. Called once
per ``h3_*`` dir via ``client.map`` so the cumulative readdir latency
on shared GPFS is parallelized across all workers instead of being
paid serially on the driver. ``os.scandir`` is preferred over
``glob``/``listdir`` because its DirEntry caches the directory bit
from the parent readdir — no extra ``stat()`` round-trip per child.
"""
out: List[str] = []
try:
with os.scandir(h3_dir) as it:
for e in it:
if e.name.startswith('year=') and e.is_dir(follow_symlinks=False):
out.append(e.path + os.sep)
except OSError:
pass
return out
def _has_unmerged_fragments(tmp_dir: str) -> bool:
"""True if any ``<tmp_dir>/h3_*/year=*/*.parquet`` fragment awaits a merge.
Stops at the first fragment. Looks for files, not just ``h3_*`` dirs: a
clean merge that kept tmp for forensics leaves the emptied dirs behind,
and merging those would only drop the bbox index for nothing.
"""
try:
with os.scandir(tmp_dir) as top:
h3_dirs = [e.path for e in top if e.is_dir(follow_symlinks=False) and e.name.startswith('h3_')]
except OSError:
return False
for h3d in h3_dirs:
for ydir in _list_year_subdirs(h3d):
try:
with os.scandir(ydir) as it:
if any(e.name.endswith('.parquet') for e in it):
return True
except OSError:
continue
return False
def _process_h3_partition(h3_dir: str) -> Dict[Tuple[int, int, int], set]:
"""Return ``{(orbit, granule, track): set(beam_str, ...)}`` for granules
represented under ``h3_dir/year=*/*.parquet``.
The beam set is what makes this partial-resume safe. A granule with
only some of its beams on disk (e.g. one (granule × beam) task finished
writing, the others were still pending when the build was killed)
would otherwise be indistinguishable from a fully-extracted granule.
The reconcile aggregates beam sets across all h3_* partitions and
only flips INDEXED when every expected beam is present — without that,
the missing beams' shots are silently dropped on the next stage-1 run.
Worker-side body of the resume reconcile, run as one Dask task per
``h3_*`` tmp partition. All cluster parallelism comes from
``client.map`` over h3 partitions — there is intentionally no
thread/process pool inside the task itself: every worker spawning N
threads multiplies cluster-wide concurrency by N×nworkers, and on a
shared GPFS that overwhelms the metadata server long before per-task
throughput catches up.
Fast path: if the fragment basename matches the v0.8.0+ naming
convention ``O{orbit}_G{granule}_T{track}.{beam}.parquet``, both
granule ID and beam are parsed from the filename — no parquet I/O,
microseconds per file.
Fallback: legacy ``part.NNN.parquet`` names lack beam info, so we
read parquet column statistics (``_granule_ids_in_fragment``) — ~90
ms cold per file on GPFS — and tag the granule with
``_LEGACY_BEAM_SENTINEL`` so the reconcile treats it as complete.
"""
out: Dict[Tuple[int, int, int], set] = {}
fallback_paths: list = []
try:
for ye in os.scandir(h3_dir):
if not (ye.is_dir(follow_symlinks=False) and ye.name.startswith('year=')):
continue
try:
for fe in os.scandir(ye.path):
if not (fe.is_file(follow_symlinks=False) and fe.name.endswith('.parquet')):
continue
m = _FRAGMENT_BASENAME_RE.match(fe.name)
if m is not None:
key = (int(m.group(1)), int(m.group(2)), int(m.group(3)))
out.setdefault(key, set()).add(m.group(4))
else:
fallback_paths.append(fe.path)
except OSError:
continue
except OSError:
return out
for p in fallback_paths:
for gid in _granule_ids_in_fragment(p):
out.setdefault(gid, set()).add(_LEGACY_BEAM_SENTINEL)
return out
def _reconcile_granules_from_disk(h3_dir: str, h3_logger, tmp_dir: Optional[str] = None,
expected_scope: Optional[str] = None) -> int:
"""Mark granules INDEXED based on what's already on disk.
Scans the finalized partition metadata under ``h3_dir`` AND the tmp fragment
tree under ``tmp_dir`` to identify granules whose data is already present.
Flips matching ``H3BuildLogger.granule_info`` entries to ``status='INDEXED'``
so the next stage-1 invocation skips them via ``get_finished_granules()``.
Idempotent. Safe to call before every build. Caller is responsible for
persisting the change with ``h3_logger.save_log(...)``.
Parameters
----------
h3_dir : str
Final H3 database directory (contains ``h3_*/<year>/*.parquet`` and
per-partition ``*.metadata.json`` sidecars).
h3_logger : H3BuildLogger
Live build-log object whose ``granule_info`` will be mutated in place.
tmp_dir : str, optional
Temporary partitions directory (typically ``<build_tmp>/partitions``).
Skipped if None or non-existent.
expected_scope : str, optional
Stage 1 scope fingerprint (:func:`_stage1_scope_fingerprint`) the next
Stage 1 will run under. Completion sentinels under ``tmp_dir`` are
trusted for flips only when ``_complete_scope.json`` records exactly
this value: sentinels from another scope (a narrower ROI, other
variables) prove nothing about this one, and the driver discards them
later than this runs. ``None`` (default) trusts no sentinel; the
granules stay non-INDEXED and Stage 1 skips or re-reads their tasks
by sentinel/scope. Finalized-metadata (Pass A) and the legacy flips
for trees with neither ``_complete/``, a scope record nor a
``_complete.stale.*`` sibling are unaffected. A resume of an
interrupted spatial- or temporal-only expansion re-reads its tasks
(the log then holds the union scope, so the fingerprint differs).
Returns
-------
int
Number of granule entries flipped from non-INDEXED to INDEXED.
"""
from .logger import GRANULE_STATUS_NO_DATA
if not hasattr(h3_logger, 'granule_info') or not h3_logger.granule_info:
return 0
# Short-circuit: nothing to flip → no need to touch disk.
# NO_DATA counts as finished: read in full, nothing on disk to find.
if not any(g.get('status') not in ('INDEXED', GRANULE_STATUS_NO_DATA) for g in h3_logger.granule_info):
logger.info(
"Resume reconciliation: build log shows no pending granules, "
"skipping disk scan"
)
return 0
# ``indexed_ids`` holds granules confirmed fully on disk and safe to skip
# on the next stage 1 run.
indexed_ids: set = set()
# ``granule_beams`` aggregates beam sets across every h3_* tmp partition
# so we can distinguish fully-extracted granules from partial ones that
# were killed mid-write. Without this completeness check a granule with
# only one beam fragment would be flipped INDEXED and the remaining 7
# beams' shots would be silently dropped on resume.
granule_beams: Dict[Tuple[int, int, int], set] = {}
# Pass A — finalized partition metadata JSONs.
# Granules named here are inside an h3 partition that has already been
# merged and finalized → all their beams' rows are already consolidated
# into the final parquet, so they're complete by construction.
#
# Sourcing the partition list:
# * Prefer the manifest sentinel (one cached read) — at continental
# scale this is the only way to keep Pass A O(N_partitions) instead
# of O(N_finalized_files) (two recursive globs over the entire DB
# tree, which was the dominant resume cost on million-partition
# databases).
# * Fall back to a single ``os.scandir`` on h3_dir for legacy DBs
# without a manifest sentinel — still avoids the recursive globs.
#
# Dispatch the per-partition meta read across workers when a dask
# client is registered; serial loop otherwise (small DB or library /
# notebook context with no cluster).
partition_dirs_set: set = set()
try:
from .utils import _read_manifest as _read_db_manifest
_mp = _read_db_manifest(h3_dir)
except Exception:
_mp = None
if _mp:
for _rel in _mp:
_head = _rel.split('/', 1)[0]
if _head.startswith('h3_'):
partition_dirs_set.add(os.path.join(h3_dir, _head))
if not partition_dirs_set:
try:
for e in os.scandir(h3_dir):
if e.is_dir(follow_symlinks=False) and e.name.startswith('h3_'):
partition_dirs_set.add(e.path)
except OSError:
pass
partition_dirs_list = sorted(partition_dirs_set)
if partition_dirs_list:
client = None
try:
client = get_dask_client()
except Exception:
pass
if client is not None and len(partition_dirs_list) > 100:
from .parallel import parallel_map
for _, _gids in parallel_map(
partition_dirs_list,
_scan_partition_meta_granules,
desc='Reconcile Pass A',
unit='part',
meta_filename=PARTITION_META_FILENAME,
):
if isinstance(_gids, Exception):
continue
indexed_ids.update(_gids)
else:
for _pd in partition_dirs_list:
indexed_ids.update(
_scan_partition_meta_granules(
_pd, meta_filename=PARTITION_META_FILENAME,
)
)
# Evidence a NO_DATA granule may be flipped back on: finalized metadata
# only. Sentinels (Pass C) also exist for the zero-row task that earned
# NO_DATA, so they say nothing about rows.
metadata_ids = set(indexed_ids)
# Pass B — tmp fragments. One Dask task per h3_* partition, with the
# parquet metadata reads inside each task parallelized via a thread
# pool (see _process_h3_partition).
if tmp_dir and os.path.isdir(tmp_dir):
try:
h3_dirs = sorted(
e.path for e in os.scandir(tmp_dir)
if e.is_dir(follow_symlinks=False) and e.name.startswith('h3_')
)
except OSError:
h3_dirs = []
if h3_dirs:
logger.info(
f"Reconciling granule status from {len(h3_dirs)} h3_* tmp partitions"
)
client = None
try:
client = get_dask_client()
except Exception:
pass
def _merge_partition_data(data: Dict[Tuple[int, int, int], set]) -> None:
for gid, beams in data.items():
granule_beams.setdefault(gid, set()).update(beams)
from tqdm import tqdm as tqdm_bar
if client is not None:
from dask.distributed import as_completed as dask_as_completed
futures = client.map(_process_h3_partition, h3_dirs, pure=False)
pbar = tqdm_bar(
total=len(futures),
desc="Reconcile partitions", unit="dir",
)
try:
for fut in dask_as_completed(futures):
try:
_merge_partition_data(fut.result())
except Exception:
pass
finally:
fut.release()
pbar.update(1)
finally:
pbar.close()
else:
# In-process fallback for tiny scenarios (no client).
for d in tqdm_bar(h3_dirs, desc="Reconcile partitions", unit="dir"):
try:
_merge_partition_data(_process_h3_partition(d))
except Exception:
continue
# Pass C — scan streaming completion sentinels under tmp_dir/_complete/.
# Each sentinel proves one (granule × beam) task ran to completion under
# the streaming writer (all leaves committed atomically + sentinel
# emitted as the final step). Sentinels are the authoritative
# completeness signal going forward.
sentinel_beams: Dict[Tuple[int, int, int], set] = {}
data_gids: set = set() # granules with at least one ``.done`` (data) sentinel
sentinel_dir_exists = False
if tmp_dir and os.path.isdir(tmp_dir):
sentinel_dir_exists = os.path.isdir(
os.path.join(tmp_dir, _COMPLETE_SENTINEL_DIRNAME)
)
if sentinel_dir_exists:
all_frags, empty_frags = _scan_complete_sentinels(tmp_dir, with_empty=True)
for frag_name in all_frags:
m = _FRAGMENT_BASENAME_RE.match(f'{frag_name}.parquet')
if m is None:
continue
gid = (int(m.group(1)), int(m.group(2)), int(m.group(3)))
sentinel_beams.setdefault(gid, set()).add(m.group(4))
if frag_name not in empty_frags:
data_gids.add(gid)
# Decide reconcile mode.
#
# sentinel_mode=True → the streaming writer has run at least once on
# this tmp tree, so sentinels are the authoritative completeness signal.
# Fragment-presence alone (Pass B) is NO LONGER sufficient because the
# streaming writer may have written some leaves of a (granule × beam)
# before being killed mid-task — only the sentinel proves all leaves
# were committed.
#
# sentinel_mode=False → tmp tree predates the streaming writer (built
# by the legacy ddf.to_parquet path). Fall back to the pre-existing
# fragment-presence heuristic and EMIT sentinels for any granule we
# flip INDEXED — that bridges the legacy tree into the sentinel model
# so subsequent resumes are sentinel-aware.
#
# A scope record or a ``_complete.stale.*`` sibling also marks a streaming
# tree: ``_check_scope_fingerprint`` renames ``_complete/`` away on a scope
# change, and a run killed before its first task finishes leaves old-scope
# fragments with no ``_complete/``. That must never read as legacy.
streaming_tree = sentinel_dir_exists
if not streaming_tree and tmp_dir and os.path.isdir(tmp_dir):
streaming_tree = os.path.exists(os.path.join(tmp_dir, _SCOPE_FILENAME))
if not streaming_tree:
try:
streaming_tree = any(
e.name.startswith(_COMPLETE_SENTINEL_DIRNAME + '.stale.') for e in os.scandir(tmp_dir)
)
except OSError:
pass
sentinel_mode = streaming_tree
expected_beams = set(GEDI_BEAMS)
n_partial = 0
n_migrated_sentinels = 0
migration_emit_pairs: List[Tuple[str, str]] = [] # (frag_name, beam) — for clarity in logs
if sentinel_mode and (expected_scope is None or _read_scope_fingerprint(tmp_dir) != expected_scope):
logger.info(
"Resume reconciliation: completion sentinels were not produced under the current "
"Stage 1 scope (or the scope is unknown); not trusting them for granule status"
)
elif sentinel_mode:
# AUTHORITATIVE PATH: granules complete only when every expected
# beam has its sentinel emitted (and at least one is a data sentinel). Fragment-presence in granule_beams
# is ignored for completeness; we still report partials based on it
# as a diagnostic.
for gid, beams in sentinel_beams.items():
if gid in indexed_ids:
continue
# All beams complete AND at least one carried data: an all-empty
# granule only proves emptiness under the scope that ran, so it
# stays as-is (Stage 1 skips its tasks via the sentinel scan when
# the scope matches, and re-reads them when it changed).
if expected_beams.issubset(beams) and gid in data_gids:
indexed_ids.add(gid)
elif not expected_beams.issubset(beams):
n_partial += 1
# Also surface fragment-on-disk-but-no-sentinel granules as partial
# in the diagnostic count (they will be re-extracted on next run).
for gid in granule_beams:
if gid in indexed_ids:
continue
if gid not in sentinel_beams:
n_partial += 1
else:
# LEGACY MIGRATION PATH: use the pre-streaming fragment-presence
# heuristic. For every granule flipped INDEXED here, emit sentinels
# for each of its beams so subsequent reconciles are sentinel-aware.
for gid, beams in granule_beams.items():
if gid in indexed_ids:
continue
if _LEGACY_BEAM_SENTINEL in beams or expected_beams.issubset(beams):
indexed_ids.add(gid)
# Migration: synthesise a frag_name per beam in `beams` and
# emit its sentinel. For the legacy-sentinel case we don't
# know the actual beam, so emit for the full GEDI_BEAMS set
# (consistent with the legacy "trust this granule" semantics).
emit_beams = set(GEDI_BEAMS) if _LEGACY_BEAM_SENTINEL in beams else beams
orbit, granule, track = gid
for beam in emit_beams:
frag_name = f"O{orbit:05d}_G{granule:02d}_T{track:05d}.{beam}"
if tmp_dir:
_emit_complete_sentinel(tmp_dir, frag_name)
migration_emit_pairs.append((frag_name, beam))
n_migrated_sentinels += 1
else:
n_partial += 1
if n_migrated_sentinels:
logger.info(
f"Resume reconciliation (migration): emitted {n_migrated_sentinels} "
f"completion sentinels for legacy tmp fragments; subsequent reconciles "
f"will use sentinel-authoritative mode"
)
if n_partial:
logger.info(
f"Resume reconciliation: {n_partial} granule(s) found with partial "
f"beam coverage on disk; leaving non-INDEXED for re-extraction"
)
if not indexed_ids:
return 0
n_flipped = 0
for g in h3_logger.granule_info:
key = (g['orbit'], g['granule'], g['track'])
status = g.get('status')
# MERGE_FAILED is found on disk by construction (its other partitions
# merged); flipping it would skip the re-extraction that recovers it.
if key in indexed_ids and status not in ('INDEXED', 'MERGE_FAILED', GRANULE_STATUS_NO_DATA):
g['status'] = 'INDEXED'
n_flipped += 1
elif status == GRANULE_STATUS_NO_DATA and key in metadata_ids:
g['status'] = 'INDEXED' # rows now exist (a wider scope found some)
g.pop('products', None) # rebuilt by set_post_build_info / lazy upgrade
g.pop('fill_attempts', None)
n_flipped += 1
if n_flipped:
logger.info(
f"Resume reconciliation: {n_flipped} granules flipped to INDEXED "
f"({len(indexed_ids)} unique granules found on disk)"
)
else:
logger.info(
f"Resume reconciliation: no new granules to flip "
f"({len(indexed_ids)} unique granules already accounted for)"
)
return n_flipped
def _filter_soc_files_by_temporal(all_soc_files, temporal):
"""Drop ``soc_file_tree`` entries whose granule date falls outside ``temporal``.
``temporal`` (``-t0``/``-t1``) is only ever forwarded to the CMR search
inside ``download_soc``/``s3_etl_subset`` -- the local-directory and
pre-acquired-list branches in ``build_h3db`` list every granule with the
requested products regardless of date, and neither ``_filter_granules``
nor ``_apply_spatial_filter`` downstream check dates either. Without this,
building from an already-downloaded SOC tree silently processes the
entire archive instead of the requested range.
Uses the same ``datetime >= start`` / ``datetime <= end`` bound
``gh3_extract``'s query filter applies (see ``cliutils.build_query_string``)
so a local build respects ``-t0``/``-t1`` the same way building from S3 or
after a fresh download does.
Parameters
----------
all_soc_files : list of dict
``{product: path}`` mappings from :func:`gedidriver.soc_file_tree`.
temporal : tuple of str, optional
``(start_date, end_date)`` as accepted by :func:`utils.parse_temporal`.
Either bound may be ``None`` for an open-ended range. ``None`` itself
is a no-op.
Returns
-------
list of dict
The subset of ``all_soc_files`` whose granule date is in range.
Entries whose date cannot be parsed are kept (fail open) rather than
silently dropped.
"""
if not temporal or not all_soc_files:
return all_soc_files
from datetime import datetime, timedelta
from .utils import parse_temporal
start, end = parse_temporal(temporal)
start_dt = datetime.strptime(start, '%Y-%m-%d') if start else None
# end is a date-only bound but granule dates carry a time-of-day component;
# advance to the start of the next day so the entire end date is included.
end_dt = datetime.strptime(end, '%Y-%m-%d') + timedelta(days=1) if end else None
if start_dt is None and end_dt is None:
return all_soc_files
def _in_range(entry):
path = next(iter(entry.values()))
try:
p = path.path if isinstance(path, EarthAccessFile) else path
d = GEDIFile(p).date
except Exception:
return True
if start_dt is not None and d < start_dt:
return False
if end_dt is not None and d >= end_dt:
return False
return True
filtered = [e for e in all_soc_files if _in_range(e)]
dropped = len(all_soc_files) - len(filtered)
if dropped:
logger.info(
f"Temporal filter [{start or '-inf'}, {end or '+inf'}]: "
f"excluded {dropped} granule(s) outside the requested range"
)
return filtered
# Products published with or before L2A. A phased build (allow_missing_products)
# requires these; only the later products (L2B, L4A, L4C…) may be missing.
_PRODUCTS_BEFORE_L2B = frozenset({'L1B', 'L2A'})
def _drop_partition_fragments(h3_part_dir: str, *, keys: set, tmp_dir: str) -> int:
"""Worker: delete the fragments (and completion sentinels) of ``keys`` under one tmp ``h3_*`` dir."""
n = 0
for ydir in _list_year_subdirs(h3_part_dir):
try:
entries = list(os.scandir(ydir))
except OSError:
continue
for e in entries:
m = _FRAGMENT_BASENAME_RE.match(e.name)
if m and (int(m[1]), int(m[2]), int(m[3])) in keys:
for path in (e.path, _complete_sentinel_path(tmp_dir, e.name[:-len('.parquet')])):
try:
os.unlink(path)
except OSError:
pass
n += 1
return n
def _drop_deferred_fragments(tmp_dir: str, keys: set) -> int:
"""Delete tmp fragments of granules a spatial-only expansion leaves for the next build.
A run killed during Stage 1 can leave some of a granule's (granule ×
beam) fragments in ``tmp_dir``. If this run defers that granule
(``build_h3db(only_granules=)``), merging those fragments would list it
in the partition metadata, and ``set_post_build_info`` (or the next
reconcile) would mark it ``INDEXED`` with its other beams never read.
Its fragments go instead: the next build reads it over the whole region.
Parameters
----------
tmp_dir : str
Stage 1 fragment root (``<tmp>/partitions``).
keys : set of tuple
``(orbit, granule, track)`` of the deferred granules.
Returns
-------
int
Number of fragments removed.
"""
if not keys or not _has_unmerged_fragments(tmp_dir):
return 0
try:
with os.scandir(tmp_dir) as top:
h3_dirs = [e.path for e in top if e.is_dir(follow_symlinks=False) and e.name.startswith('h3_')]
except OSError:
return 0
client = get_dask_client()
if client is not None and len(h3_dirs) > 1:
from .parallel import parallel_map
results = (r for _, r in parallel_map(h3_dirs, _drop_partition_fragments, desc='dropping deferred fragments',
unit='part', batch_size=256, keys=keys, tmp_dir=tmp_dir))
else:
results = (_drop_partition_fragments(d, keys=keys, tmp_dir=tmp_dir) for d in h3_dirs)
n = sum(r for r in results if isinstance(r, int))
if n:
logger.info(f"Dropped {n} tmp fragment(s) of granules left for the next build")
return n
def _filter_granules(
prod_soc_files: List[Dict[str, str]],
product_vars: Dict[str, List[str]],
skip_granules: Optional[List[Dict]] = None,
required_products: Optional[set] = None,
only_granules: Optional[List[Dict]] = None,
) -> List[Dict[str, str]]:
"""
Filter SOC files to exclude incomplete, corrupted, or already-processed granules.
The product and skip checks are a-priori — set lookups on the driver, the
granule key parsed from the filename (no stat, no HDF5 open) — so only the
surviving granules reach the parallel ``h5_is_valid`` check. On an update
that is the new granules alone, not the whole SOC listing.
Parameters
----------
prod_soc_files : list of dict
List of product file dictionaries
product_vars : dict
Requested products and their variables
skip_granules : list of dict, optional
Granule identifiers to skip
required_products : set of str, optional
Products a granule must have. ``None`` requires every product in
``product_vars``; a subset (``{'L2A'}``) admits granules whose later
products are not published yet.
only_granules : list of dict, optional
If given, every other granule is left out (and counted as deferred):
see ``H3BuildLogger.stage1_only_granules``.
Returns
-------
list of dict
Filtered list of valid SOC file dictionaries
"""
required = set(product_vars) if required_products is None else set(required_products)
skip = {(g['orbit'], g['granule'], g['track']) for g in (skip_granules or [])}
only = None if only_granules is None else {(g['orbit'], g['granule'], g['track']) for g in only_granules}
def _granule_key(prod):
path = next(iter(prod.values()))
key = _granule_id_from_l2a_path(getattr(path, 'path', path))
if key is None:
gf = GEDIFile(path)
key = (gf.orbit, gf.orbit_granule, gf.track)
return key
candidates = [p for p in prod_soc_files
if p and required.issubset(p) and (not skip or _granule_key(p) not in skip)]
if only is not None:
n_listed = len(candidates)
candidates = [p for p in candidates if _granule_key(p) in only]
if len(candidates) < n_listed:
logger.warning(
f"{n_listed - len(candidates)} granule(s) not finished under the stored region are left for "
f"the next build: this update reads only the region it adds, and they need the whole region. "
f"Run gh3_build again once it completes to index them.")
def _filter_soc_file(prod):
if any(isinstance(f, EarthAccessFile) for f in prod.values()):
return prod
for f in prod.values():
if not h5_is_valid(f):
return None
return prod
logger.info("Checking for incomplete, corrupted, or existing granules to skip")
if not candidates:
n_skipped = len(prod_soc_files)
if n_skipped:
logger.info(f"Skipped {n_skipped}/{n_skipped} granules (already indexed or incomplete)")
return []
bag_task = (
dbg.from_sequence(candidates, partition_size=100)
.map(_filter_soc_file)
.filter(lambda x: x is not None)
.persist()
)
# Per-batch tqdm bar over the bag's futures — renders reliably in TTY,
# SSH, and log-redirected contexts where distributed.progress() can be
# invisible. Mirrors the partition-merge pattern below.
from dask.distributed import futures_of, as_completed as dask_as_completed
from tqdm import tqdm as tqdm_bar
bag_futures = futures_of(bag_task)
if bag_futures:
pbar = tqdm_bar(dask_as_completed(bag_futures), total=len(bag_futures),
desc="Checking SOC files", unit="batch")
for _ in pbar:
pass
pbar.close()
soc_files = list(bag_task.compute())
del bag_task
n_skipped = len(prod_soc_files) - len(soc_files)
if n_skipped > 0:
logger.info(f"Skipped {n_skipped}/{len(prod_soc_files)} granules (already indexed, incomplete, or corrupted)")
return soc_files
def _create_h3_dataframe(
soc_files: List[Dict[str, str]],
product_vars: Dict[str, List[str]],
res: int,
part: int
) -> Tuple[dask_geopandas.GeoDataFrame, str, str, str, List[str]]:
"""
Create a Dask GeoDataFrame with H3 indexing from SOC files.
Parameters
----------
soc_files : list of dict
List of valid SOC file dictionaries
product_vars : dict
Products and variables to extract
res : int
H3 resolution for indexing
part : int
H3 resolution for partitioning
Returns
-------
tuple
(dask_geopandas.GeoDataFrame, lat_col, lon_col, dat_col, frag_names)
``frag_names[i]`` is a stable, content-derived basename for Dask
partition ``i``, used as ``name_function`` in ``to_parquet`` so re-runs
write to the same path and overwrite in place (no shot duplication).
"""
logger.info(f"Found {len(soc_files)} new GEDI granules with requested products")
ddf = dask_h5_merged(soc_files, product_vars, shots=None, dropna=True, by_beam=True, suffix_all=True)
# Build per-partition stable names matching dask_h5_merged's by_beam=True
# enumeration: itertools.product(soc_files, GEDI_BEAMS). Each Dask partition
# i corresponds to one (granule, beam) tuple. Delegated to the shared
# _granule_beam_frag_name helper so the streaming writer
# (_write_partitioned_streaming) and the legacy to_parquet path agree on
# the same filename convention by construction.
import itertools as _it
frag_names: List[str] = []
for _soc, _beam in _it.product(soc_files, GEDI_BEAMS):
name = _granule_beam_frag_name(_soc, _beam)
if name is None:
# Fallback: opaque but deterministic per-partition name. Only
# fires when the source HDF5 path can't be parsed (never happens
# for NASA-formatted granules). Legacy behaviour preserved.
name = f"part.{len(frag_names):08d}"
frag_names.append(name)
lat_col = 'lat_lowestmode'
lon_col = 'lon_lowestmode'
dat_col = 'delta_time'
if 'lat_lowestmode_l2a' in ddf.columns:
lat_col += '_l2a'
if 'lon_lowestmode_l2a' in ddf.columns:
lon_col += '_l2a'
if 'delta_time_l2a' in ddf.columns:
dat_col += '_l2a'
logger.info(f"Indexing H3 at resolution {res}, partitioning at {part}")
ddf = ddf.map_partitions(h3_index_df, res=res, part=part, lat_col=lat_col, lon_col=lon_col)
return ddf, lat_col, lon_col, dat_col, frag_names
def _apply_spatial_filter(
ddf: dask.dataframe.DataFrame,
spatial,
part: int,
h3_dir: str
) -> dask.dataframe.DataFrame:
"""
Apply spatial and existing-data filters to the Dask DataFrame.
Parameters
----------
ddf : dask.dataframe.DataFrame
Input Dask DataFrame with H3 index
spatial : GeoDataFrame, list, or str
Spatial filter geometry
part : int
H3 partition resolution
h3_dir : str
Path to existing H3 database (for skip detection)
Returns
-------
dask.dataframe.DataFrame
Filtered Dask DataFrame
"""
h3_tiles = []
if spatial is not None:
h3_tiles = intersect_h3_geometries(spatial, res=part)
if len(h3_tiles) > 0:
logger.info("Removing H3 partitions outside spatial filter")
ddf = ddf[ddf[f'h3_{part:02d}'].isin(h3_tiles)]
build_log = os.path.join(h3_dir, BUILD_LOG_FILENAME)
if os.path.exists(build_log):
logger.info("Checking for existing indexed GEDI data to skip")
_meta = ddf._meta.copy()
_meta['_skip'] = False
ddf = ddf.map_partitions(h3_add_skip_column, h3_dir=h3_dir, meta=_meta)
ddf = ddf[~ddf['_skip']]
ddf = ddf.drop(columns=['_skip'])
return ddf
# Top-level sidecar (not inside ``_complete/``): renaming ``_complete/`` away
# must not lose it, and writing it must not create ``_complete/`` (whose mere
# existence flips reconcile into sentinel mode).
_SCOPE_FILENAME = '_complete_scope.json'
def _discard_complete_sentinels(tmp_dir: str, desc: str) -> None:
"""Atomically retire ``tmp_dir/_complete/``; raises ``OSError`` if it cannot.
``_remove_tree_fanout`` is best-effort and not atomic, so a partial or
swallowed failure could leave stale sentinels behind. Renaming first
makes the discard all-or-nothing: either ``_complete/`` is gone (nothing
reads the renamed ``_complete.stale.*`` sibling) or the rename raised.
A missing directory is a no-op. The renamed tree is then removed
best-effort; a leftover is harmless and goes with the whole-tree cleanup.
"""
import uuid
src = os.path.join(tmp_dir, _COMPLETE_SENTINEL_DIRNAME)
dst = f"{src}.stale.{uuid.uuid4().hex}"
try:
os.rename(src, dst)
except FileNotFoundError:
return
_remove_tree_fanout(dst, desc=desc)
def _discard_unscoped_sentinels(tmp_dir: str) -> bool:
"""Retire ``tmp_dir/_complete/`` when no scope record vouches for it.
:func:`_check_scope_fingerprint` adopts such a legacy set as in-flight
resume state. The CLI calls this first when the build log says the last
build COMPLETED: the sentinels are then leftovers, and adopting them would
make a wider ROI skip the tasks that were empty under the old one.
Returns True if sentinels were discarded; raises ``GediFileError`` if
they could not be.
"""
if os.path.exists(os.path.join(tmp_dir, _SCOPE_FILENAME)):
return False
if not os.path.isdir(os.path.join(tmp_dir, _COMPLETE_SENTINEL_DIRNAME)):
return False
try:
_discard_complete_sentinels(tmp_dir, desc="Clearing unscoped completion sentinels")
except OSError as e:
raise GediFileError(
f"Cannot discard stale completion sentinels in {tmp_dir}: {type(e).__name__}: {e}. "
f"Remove {os.path.join(tmp_dir, _COMPLETE_SENTINEL_DIRNAME)} manually and re-run."
) from e
return True
def _scope_fingerprint(spatial_h3_tiles, res: int, part: int, product_vars: Dict[str, List[str]]) -> str:
"""Hash of the inputs that decide which rows a (granule x beam) task yields."""
import hashlib
h = hashlib.sha256()
payload = {
'tiles': None if spatial_h3_tiles is None else sorted(spatial_h3_tiles),
'res': res, 'part': part,
'vars': {k: sorted(list(v or [])) for k, v in sorted(product_vars.items())},
}
h.update(json.dumps(payload, sort_keys=True).encode())
return h.hexdigest()
def _spatial_tiles(spatial, part: int) -> Optional[List[str]]:
"""Stage 1 tile set for an ROI at partition level ``part`` (None = no filter)."""
if spatial is None:
return None
tiles = intersect_h3_geometries(spatial, res=part)
return list(tiles) if len(tiles) > 0 else None
def _stage1_scope_fingerprint(spatial, res: int, part: int, product_vars: Dict[str, List[str]]) -> str:
"""Scope fingerprint of a Stage 1 run, from its spatial filter, H3 levels and variables.
``build_h3db`` computes it over the *expanded* variables and hands it to
``_write_partitioned_streaming``. The expansion edits the build log's own
``product_vars`` in place, so a resumed CLI reads the expanded form back
and its pre-reconcile value agrees with the recorded one. (On a fresh
build the CLI still holds raw specs, but there is no tmp tree to trust.)
"""
return _scope_fingerprint(_spatial_tiles(spatial, part), res, part, product_vars)
def _read_scope_fingerprint(tmp_dir: str) -> Optional[str]:
"""Fingerprint recorded in ``tmp_dir/_complete_scope.json`` (None if absent/unreadable)."""
try:
with open(os.path.join(tmp_dir, _SCOPE_FILENAME)) as fh:
return json.load(fh).get('fingerprint')
except (OSError, ValueError, AttributeError):
return None
def _check_scope_fingerprint(tmp_dir: str, fingerprint: str) -> None:
"""Invalidate stale completion sentinels when the Stage 1 scope changed.
A sentinel only proves a task finished under the scope it ran with (a task
outside the old ROI is "complete" with zero rows). The scope fingerprint
lives in ``tmp_dir/_complete_scope.json``. On mismatch the sentinels are
retired atomically (:func:`_discard_complete_sentinels`) *before* the new
scope is written, so every task re-runs; leaf writes are idempotent (``<frag_name>.parquet`` is replaced
atomically). A missing file next to existing sentinels is a tmp tree from
an older version: adopt the current scope and keep the sentinels.
"""
path = os.path.join(tmp_dir, _SCOPE_FILENAME)
stored = _read_scope_fingerprint(tmp_dir)
if stored == fingerprint:
return
if stored is not None or os.path.isfile(path):
# Changed scope, or a scope file too damaged to prove anything: re-run.
logger.warning(
"Streaming resume: the spatial scope, H3 levels or variables changed since the "
"previous run (or its record is unreadable); discarding its completion sentinels "
"so every task is re-run."
)
try:
_discard_complete_sentinels(tmp_dir, desc="Clearing stale completion sentinels")
except OSError as e:
# Never stamp the new scope over sentinels we could not retire.
raise GediFileError(
f"Cannot discard stale completion sentinels in {tmp_dir}: {type(e).__name__}: {e}. "
f"Remove {os.path.join(tmp_dir, _COMPLETE_SENTINEL_DIRNAME)} manually and re-run."
) from e
# else: no record at all — a fresh tmp tree, or one from an older version
# whose sentinels we keep (legacy behaviour) by adopting the current scope.
os.makedirs(tmp_dir, exist_ok=True)
with AtomicFileWriter(path) as tmp_path:
with open(tmp_path, 'w') as fh:
json.dump({'fingerprint': fingerprint}, fh)
def _write_one_granule_beam(
task: Tuple[Dict[str, str], str, str],
*,
product_vars: Dict[str, List[str]],
res: int,
part: int,
tmp_dir: str,
h3_dir: str,
lat_col: str,
lon_col: str,
dat_col: str,
spatial_h3_tiles: Optional[List[str]] = None,
skip_check_enabled: bool = False,
schema: Any = None,
fill_missing_columns: bool = False,
) -> Dict[str, Any]:
"""Worker-side body of the streaming partition write.
Loads ONE (granule × beam) HDF5, applies the same transformation chain
the legacy dask graph used (h3_index_df → optional spatial filter →
h3_add_skip_column → add_special_columns → year synthesis → groupby on
[h3_{part:02d}, year]), writes one parquet leaf per (h3 cell × year)
group via AtomicFileWriter + GeoDataFrame.to_parquet, then emits a
per-(granule × beam) completion sentinel only AFTER every leaf is
committed. The sentinel means "this task was read successfully and all
its leaves (possibly zero) are committed" (``.done`` = at least one leaf,
``.empty`` = zero rows); the reconcile trusts it as
proof that the (granule × beam) is fully on disk — eliminating the
legacy "any-beam-fragment-equals-complete-granule" data-loss path on
kill-mid-write resume. A task that was read but yielded no rows
(empty beam, outside the ROI, already-covered cells) still gets a
sentinel, so resume does not re-read it.
Parameters
----------
task : (soc_dict, beam, frag_name)
``soc_dict`` maps product code → HDF5 path for one granule.
``beam`` is one of ``GEDI_BEAMS``. ``frag_name`` matches
``_FRAGMENT_BASENAME_RE`` and is the per-task basename used both
for the leaf parquet files and the completion sentinel.
product_vars, res, part, lat_col, lon_col, dat_col
Identical to the legacy chain's kwargs.
tmp_dir, h3_dir
Output tmp tree root and existing-h3-db root.
spatial_h3_tiles
Driver-broadcast list of H3 cell IDs (at resolution ``part``) to
keep. ``None`` disables spatial filtering. Replaces the legacy
``ddf[ddf[h3_part_col].isin(h3_tiles)]`` at _apply_spatial_filter.
skip_check_enabled
When True, runs ``h3_add_skip_column`` to drop rows whose target
h3 partition already has finalized data in ``h3_dir``.
schema
Driver-built pyarrow Schema used for every per-leaf write to
force column-order + dtype parity across fragments. Without
this, per-leaf schema inference would drift between fragments
and break the merge phase's schema union. See
``_canonical_write_schema``.
fill_missing_columns
When True (a build admitting granules with missing products), a
leaf lacking some schema columns gets them as all-null
(:func:`_conform_table_to_schema`) instead of failing the cast.
Returns
-------
dict
``{'frag_name': str, 'leaves': int, 'rows': int, 'skipped': bool,
'error': Optional[str], 'complete': bool}``. ``complete`` is True
exactly when the completion sentinel was emitted (positive evidence
the task was read in full); the driver uses it to retire an earlier
failure record. ``skipped=True`` covers
empty-after-load, empty-after-index, empty-after-spatial-filter,
empty-after-skip-check, and no-leaves-written; the completion
sentinel IS emitted in each of these (the task was read and
genuinely produced no data) as the ``.empty`` kind, with
``leaves == 0`` and ``empty=True``; a task with >= 1 leaf writes
``.done`` and ``empty=False``. No sentinel is
emitted on a load error (``error`` set) or when no product file
exists for the task (``load_h5_merged`` returned ``None``).
"""
soc_dict, beam, frag_name = task
stats = {'frag_name': frag_name, 'leaves': 0, 'rows': 0, 'skipped': False,
'error': None, 'failure': None, 'complete': False, 'empty': False}
# 1) Load HDF5 for one (granule, beam) — identical contract to
# dask_h5_merged(by_beam=True)'s inner load_by_beam closure.
try:
df = load_h5_merged(
soc_dict, product_vars=product_vars,
which_beams=[beam], shots=None,
dropna=True, suffix_all=True,
)
except Exception as e:
# Mirrors the legacy graph's _load_h5_merged exception swallow —
# corrupt h5 returns an empty meta upstream rather than failing the
# whole job. Streaming surfaces the error in stats for visibility,
# plus a structured ``failure`` record so the driver can persist it
# for the end-of-build advisory and gh3_doctor without needing to grep
# the WARN log lines later.
stats['skipped'] = True
stats['error'] = f"load_h5_merged: {type(e).__name__}: {e}"
stats['failure'] = _classify_load_h5_failure(e, soc_dict)
return stats
if df is None:
# No product file present for this task at all ("no source"), as
# opposed to a file that was read and held zero rows. Leave it
# sentinel-less so a later resume retries once the source exists.
stats['skipped'] = True
return stats
if df.empty:
# Read successfully, zero rows: record completion (``.empty`` kind)
# so resume never re-reads this task.
_emit_complete_sentinel(tmp_dir, frag_name, empty=True)
stats['complete'] = True
stats['skipped'] = True
stats['empty'] = True
return stats
# 2) H3 index — same call as legacy ddf.map_partitions(h3_index_df, ...).
df = h3_index_df(df, res=res, part=part, lat_col=lat_col, lon_col=lon_col)
if df.empty:
_emit_complete_sentinel(tmp_dir, frag_name, empty=True)
stats['complete'] = True
stats['skipped'] = True
stats['empty'] = True
return stats
h3_part_col = f'h3_{part:02d}'
# 3) Spatial filter — replaces _apply_spatial_filter's isin branch.
# Tile set is precomputed driver-side and scattered (constant per build).
if spatial_h3_tiles is not None:
df = df[df[h3_part_col].isin(spatial_h3_tiles)]
if df.empty:
_emit_complete_sentinel(tmp_dir, frag_name, empty=True)
stats['complete'] = True
stats['skipped'] = True
stats['empty'] = True
return stats
# 4) Skip-existing-data filter — replaces _apply_spatial_filter's
# h3_add_skip_column branch. h3_add_skip_column reads from h3_dir
# to detect cells whose finalized data already covers this granule.
if skip_check_enabled and 'root_file_l2a' in df.columns:
df = h3_add_skip_column(df, h3_dir=h3_dir)
df = df[~df['_skip']].drop(columns=['_skip'])
if df.empty:
_emit_complete_sentinel(tmp_dir, frag_name, empty=True)
stats['complete'] = True
stats['skipped'] = True
stats['empty'] = True
return stats
# 5) Special columns + year — same calls as legacy.
df = add_special_columns(df, lon_col=lon_col, lat_col=lat_col, dat_col=dat_col)
df = df.assign(year=df['datetime'].dt.year)
# 6) Groupby + per-leaf atomic write. observed=True + sort=False mirror
# the legacy partition_on=[h3_part, year] semantics — partition
# columns are stored in directory names, dropped from file body.
leaves_written = 0
rows_written = 0
# Lazy import — only the streaming write path uses this private hook
# into geopandas's arrow conversion. Kept identical to what dask's
# GeoArrowEngine._pandas_to_arrow_table calls under the hood, so the
# streaming output matches the legacy chain byte-for-byte (modulo
# row order within identical input).
from geopandas.io.arrow import _geopandas_to_arrow
for (h3_cell, year), leaf_df in df.groupby([h3_part_col, 'year'], sort=False, observed=True):
if leaf_df.empty:
continue
leaf_dir = os.path.join(tmp_dir,
f'{h3_part_col}={h3_cell}',
f'year={int(year)}')
out_path = os.path.join(leaf_dir, f'{frag_name}.parquet')
body = leaf_df.drop(columns=[h3_part_col, 'year'])
# Convert to pyarrow Table via the geopandas hook (carries the
# GeoParquet ``geo`` schema metadata + WKB geometry encoding),
# then cast to the canonical driver-built schema to lock down
# column order and nullable-dtype tagging across all fragments.
table = _geopandas_to_arrow(body, index=True)
if schema is not None:
if fill_missing_columns and table.schema.names != schema.names:
table = _conform_table_to_schema(table, schema)
else:
table = table.cast(schema)
with AtomicFileWriter(out_path) as tmp_path:
pq.write_table(table, tmp_path, compression='zstd')
leaves_written += 1
rows_written += len(body)
# 7) Completion sentinel — the load-bearing data-loss guard. Emitted
# only AFTER every leaf is committed (AtomicFileWriter.__exit__
# succeeded). If the worker dies between leaves, no sentinel is
# emitted → reconcile leaves the granule non-INDEXED → next resume
# re-extracts the (granule × beam) idempotently. Emitted even when
# zero leaves were written: "read OK, all leaves (possibly zero)
# committed". Errors never reach here (exception path returns early).
# ``.done`` when a leaf was committed, ``.empty`` for zero leaves.
_emit_complete_sentinel(tmp_dir, frag_name, empty=leaves_written == 0)
stats['complete'] = True
if leaves_written == 0:
stats['skipped'] = True
stats['empty'] = True
stats['leaves'] = leaves_written
stats['rows'] = rows_written
return stats
def _streaming_enabled() -> bool:
"""Whether the partition-write phase should use the streaming path.
Streaming is the default since the v0.9.5 cutover. Set
``GH3_WRITE_STREAMING={0,false,off,no}`` to opt back into the legacy
``ddf.to_parquet`` path (kept for one release cycle as a diagnostic
fallback; will be removed in v0.10.0).
"""
val = os.environ.get('GH3_WRITE_STREAMING', '').strip().lower()
if val in ('0', 'false', 'off', 'no'):
return False
return True
def _streaming_batch_size(n_workers: Optional[int] = None) -> int:
"""Rolling-window inflight target for the streaming driver.
Defaults to ``max(n_workers * 2, 100)`` when ``n_workers`` is known —
just enough buffering above the 1-task-per-worker minimum to absorb
task-time variance without hiding stragglers. A 2000-inflight default
(an earlier design) was 8× overkill on this cluster: most queued tasks
would sit invisible behind the rolling-window's release-as-completed
loop, defeating the visibility goal.
Override via ``GH3_WRITE_STREAMING_BATCH`` env. ``n_workers=None``
falls back to a static ``500`` (~2× the gsapp cluster's typical
capacity) when the driver hasn't connected to the scheduler yet.
"""
val = os.environ.get('GH3_WRITE_STREAMING_BATCH')
if val is not None:
try:
return max(1, int(val))
except ValueError:
pass
if n_workers and n_workers > 0:
return max(n_workers * 2, 100)
return 500
def _write_partitioned_streaming(
ddf: dask.dataframe.DataFrame,
soc_files: List[Dict[str, str]],
product_vars: Dict[str, List[str]],
res: int,
part: int,
tmp_dir: str,
h3_dir: str,
spatial,
lat_col: str,
lon_col: str,
dat_col: str,
inflight_target: Optional[int] = None,
allow_missing_products: bool = False,
stage1_outcome_callback: Optional[Callable[[set, set], None]] = None,
scope_fingerprint: Optional[str] = None,
skip_check: Optional[bool] = None,
) -> bool:
"""Streaming replacement for the legacy ``ddf.to_parquet().persist()``.
Emits one per-(granule × beam) ``client.submit`` task and drains via
``as_completed`` with a rolling inflight window of
``inflight_target`` futures. Each completed future is released
immediately, so worker memory plateaus at the per-task working set
instead of accumulating across the whole graph the way the to-parquet-
barrier-bound legacy path does (see dask/dask#4894, #5159, #8377).
Correctness contract vs. the legacy path:
* Fragment basenames identical (``_granule_beam_frag_name``).
* Output schema identical (driver-built canonical pyarrow schema,
cast per leaf — matches dask's GeoArrowEngine pipeline).
* Hive layout identical (``h3_{part:02d}={cell}/year={year}/``).
* Resume-safe: only granules whose worker emitted a ``.done``
completion sentinel are recognized as finished by the reconcile.
Parameters
----------
ddf
Lazy ddf returned by ``_create_h3_dataframe``. Used here only to
derive the canonical write schema from its ``_meta`` (no
``.persist()``; the lazy graph is discarded after the schema
probe).
soc_files, product_vars, res, part, lat_col, lon_col, dat_col
Same inputs as the legacy chain.
tmp_dir, h3_dir
Output tree root + existing-h3-db root.
spatial
Same spatial filter input as ``_apply_spatial_filter`` —
intersected to a tile list ONCE driver-side and scattered.
inflight_target
Max in-flight futures at any moment. ``None`` → env-driven
default (``GH3_WRITE_STREAMING_BATCH``).
allow_missing_products
Some granules lack products the build requests. The write schema
is completed with the database's columns
(:func:`_complete_schema_from_db`) and leaves missing product
columns get them as all-null, so every fragment still merges.
stage1_outcome_callback
Called once, only after every task was drained, with
``(incomplete_keys, empty_keys)`` of ``(orbit, granule, track)``:
granules with any task not proven complete (error, no product file,
a lost future) and granules whose every beam task is proven complete
AND empty (zero leaves now, or an ``.empty`` sentinel on disk). Never called when the drain is interrupted,
so a caller holding no outcome must treat that as "nothing proven".
The CLI's ``NO_DATA`` marking rests on it.
skip_check
``False`` writes every row, even of a granule the database's
metadata already lists (a granule replacement, #47). ``None``
enables the skip check whenever the database has partitions.
"""
import itertools as _it
import time as _time
from dask.distributed import as_completed as dask_as_completed
from tqdm import tqdm as tqdm_bar
client = get_dask_client()
if client is None:
raise GediError(
"_write_partitioned_streaming requires a registered dask "
"Client (wrap your call in `with Client(...): ...`)."
)
# Inflight-target sizing: tied to the live worker count so the rolling
# window stays just-large-enough to keep workers busy without hiding
# stragglers in a deep queue. Caller can override (e.g. for the
# LocalCluster-based regression test which runs with 2 workers).
if inflight_target is None:
try:
n_workers_live = len(client.scheduler_info().get('workers', {})) or None
except Exception:
n_workers_live = None
inflight_target = _streaming_batch_size(n_workers=n_workers_live)
logger.info(f"Writing partitioned H3 data (streaming) to: {tmp_dir}")
os.makedirs(tmp_dir, exist_ok=True)
# ── Driver-side preflight: build everything that should be broadcast ──
# The merge phase's _merge_and_finalize uses the same pattern at
# gh3builder.py:_merge_and_finalize for its h3-dir scatter list.
# 1) Canonical schema — apply the streaming-side transforms to the
# ddf._meta and let _canonical_write_schema produce the pyarrow
# schema we'll cast each leaf to. Locks column order + dtypes.
canonical_schema = _stage1_write_schema(ddf, lat_col, lon_col, dat_col, part)
canonical_schema, dtype_drift = _align_schema_to_db(canonical_schema, h3_dir)
if dtype_drift:
cols = ', '.join(f"{c} (source {a}, database {b})" for c, (a, b) in sorted(dtype_drift.items()))
logger.warning(
f"Source schema differs from the existing database for {len(dtype_drift)} column(s): {cols}. "
f"New data is written in the database's types (lossless cast; a value that does not fit "
f"fails its task instead of being truncated). To move the database to the source types: "
f"gh3_doctor -i {h3_dir} --check dtype_drift --fix"
)
if allow_missing_products:
canonical_schema, added = _complete_schema_from_db(canonical_schema, h3_dir)
if added:
logger.info(
f"Write schema completed with {len(added)} database column(s) the source sample lacks; "
f"granules without those products store them as null until backfilled"
)
absent = [p for p in product_vars
if not any(n.lower().endswith(f"_{p.lower()}") for n in canonical_schema.names)]
if absent:
# Neither a granule of this batch nor the database defines these
# products' columns: files written now would lack them, and a later
# merge would drop them from complete granules. Refuse up front.
raise GediValidationError(
f"No granule in this batch has {absent} and the database stores no such columns yet; "
f"build at least one granule with every product first (without --allow-missing-products)"
)
# 2) Spatial tile set — replaces _apply_spatial_filter's isin branch.
spatial_h3_tiles: Optional[List[str]] = None
if spatial is not None:
tiles = intersect_h3_geometries(spatial, res=part)
if len(tiles) > 0:
logger.info("Pre-computing spatial H3 tile list (driver-side)")
spatial_h3_tiles = list(tiles)
# 3) Skip-check enablement gate — only fire h3_add_skip_column when
# the destination h3_dir already has an indexed build log AND a
# persisted partition list. Without the persisted list, every
# h3_skip_part call would just stat-storm GPFS for no benefit
# (Agent 3 adversarial review #J.2).
build_log_path = os.path.join(h3_dir, BUILD_LOG_FILENAME)
skip_check_enabled = False
if skip_check is False:
pass
elif os.path.exists(build_log_path):
try:
log = json_read(build_log_path) or {}
skip_check_enabled = bool(log.get('h3_partition_ids'))
except Exception:
skip_check_enabled = False
if skip_check_enabled:
logger.info("Skip-check enabled (existing finalized partitions detected)")
# 4) Skip already-completed granule×beam tasks via sentinel scan. A
# resume picks up exactly where the previous run stopped — completed
# tasks are not re-submitted, partial tasks (no sentinel) are.
_check_scope_fingerprint(
tmp_dir, scope_fingerprint or _scope_fingerprint(spatial_h3_tiles, res, part, product_vars))
completed_frags, empty_frags = _scan_complete_sentinels(tmp_dir, with_empty=True)
if completed_frags:
logger.info(
f"Streaming resume: {len(completed_frags)} (granule × beam) sentinel(s) "
f"found on disk; matching tasks will be skipped at submission time."
)
# ── Bake broadcast values into a functools.partial ──────────────────
# Iteration history on this driver's broadcast-kwarg handling:
# 1. scatter(value, broadcast=True) — scattered iterables element-wise,
# created thousands of stray futures.
# 2. scatter([value], broadcast=True)[0] — fixed cardinality but
# broadcast=True hung waiting for every SSH-tunneled worker ACK.
# 3. client.map(fn, tasks, **kwargs) inlining — dask treats certain
# kwarg names as scheduler TaskState keys (a kwarg named
# "spatial_h3_tiles" became <TaskState 'spatial_h3_tiles'
# processing> stuck-in-processing); workers then crashed on
# AssertionError when trying to resolve the dependency.
#
# Final approach: bake all broadcast kwargs into a ``functools.partial``
# closure around the worker function. ``client.map(partial_fn, tasks)``
# then sees a single callable + the per-task iterable — dask cannot
# introspect or split the partial's captured kwargs, so there is no
# opportunity to misinterpret them as scheduler tasks or scatter them
# element-wise.
#
# Memory accounting:
# * The partial is pickled ONCE into the Blockwise layer (one ~210 KB
# blob covering spatial_h3_tiles + product_vars + the pyarrow
# schema), and per-task entries are tiny refs into that layer.
# * Total scheduler-side graph footprint for a 532k-task build:
# ~530 MB — negligible on a 1 TB hub. Drains as fut.release()
# runs per completion below.
import functools as _functools
worker_fn = _functools.partial(
_write_one_granule_beam,
product_vars=product_vars,
res=res, part=part,
tmp_dir=tmp_dir, h3_dir=h3_dir,
lat_col=lat_col, lon_col=lon_col, dat_col=dat_col,
spatial_h3_tiles=spatial_h3_tiles,
skip_check_enabled=skip_check_enabled,
schema=canonical_schema,
fill_missing_columns=allow_missing_products,
)
logger.info(
f"Driver: kwargs baked into partial (no scatter, no separate TaskStates). "
f"spatial_h3_tiles={len(spatial_h3_tiles) if spatial_h3_tiles else 0} cells, "
f"product_vars={len(product_vars)} products, "
f"schema cols={len(canonical_schema.names) if canonical_schema is not None else 0}."
)
# ── Build task stream ─────────────────────────────────────────────
# Generator over (soc_dict, beam, frag_name). Skips tasks whose
# sentinel is already on disk (resume fast-path).
logger.info("Driver: building task list...")
# Keys that failed on an earlier run (small: O(N_failures)). A task that
# later succeeds is "recovered" and its stale record is dropped by one
# compaction at the end of this phase. Only Stage 1 frag_names can ever
# land in ``recovered``; records keyed by anything else are left alone.
previously_failed = {r.get('frag_name') for r in _read_granule_failures(tmp_dir)}
recovered: set = set()
# Per-granule proof for the outcome callback. A beam is proven EMPTY when
# its task completed with zero leaves this run, or was skipped on an
# ``.empty`` sentinel; a beam that wrote (or has a ``.done`` sentinel) rules
# the granule out. O(N_granules), not O(N_tasks).
empty_beams: Dict[Tuple[int, int, int], set] = {}
has_data: set = set()
incomplete_keys: set = set()
lost_unattributed = False
def _note_complete(frag_name, empty):
gk = _frag_name_granule_key(frag_name)
if gk is None:
return
if empty:
empty_beams.setdefault(gk, set()).add(frag_name.rsplit('.', 1)[1])
else:
has_data.add(gk)
def _outcome():
full = set(GEDI_BEAMS)
return incomplete_keys, {k for k, b in empty_beams.items()
if full.issubset(b) and k not in has_data and k not in incomplete_keys}
def _task_stream():
for soc, beam in _it.product(soc_files, GEDI_BEAMS):
frag_name = _granule_beam_frag_name(soc, beam)
if frag_name is None:
continue # opaque soc filename — never happens for NASA granules
if frag_name in completed_frags:
_note_complete(frag_name, frag_name in empty_frags)
if frag_name in previously_failed:
recovered.add(frag_name)
continue
yield (soc, beam, frag_name)
tasks = list(_task_stream()) # materialize so we know the total
total = len(tasks)
if total == 0:
_compact_granule_failures(tmp_dir, recovered)
if stage1_outcome_callback is not None:
stage1_outcome_callback(*_outcome())
logger.info("Streaming write: no remaining tasks (all granules already complete)")
return any(
entry.is_dir() and entry.name.startswith('h3_')
for entry in os.scandir(tmp_dir)
) if os.path.isdir(tmp_dir) else False
logger.info(
f"Streaming write: {len(soc_files)} granules × {len(GEDI_BEAMS)} beams = "
f"{total} tasks (skipped_by_resume={len(completed_frags)})"
)
# ── Submit the entire batch via client.map ─────────────────────────
# client.map(fn, iterable, **kwargs) builds a Blockwise graph layer
# where the function + kwargs live ONCE in the scheduler's graph
# (one ~210 KB blob covering spatial_h3_tiles + product_vars + the
# canonical pyarrow schema), and per-task entries are tiny refs into
# that layer. Total scheduler-side footprint for the full 534k-task
# build: ~535 MB — negligible on a 1 TB hub.
#
# Why this beats the prior rolling-window submit-then-as_completed-add
# pattern:
# * Dashboard sees the full 534k-task scope from the start →
# real progress visualization instead of a fake 100-task window.
# * Worker saturation is dask-managed (via the scheduler's queue),
# so every worker is busy as long as tasks remain — no manual
# inflight knob to mis-tune.
# * The "scheduler-side serialization explosion" concern from
# Agent 3 #D.2 was specific to ``for t in tasks: client.submit(
# fn, t, **kwargs)`` which serializes kwargs once per submit
# call. ``client.map`` shares kwargs across the batch (single
# entry in the graph) so the explosion does not apply.
#
# Memory drains as ``fut.release()`` runs per completion below.
logger.info(f"Driver: submitting {total} tasks via client.map (partial-wrapped fn)...")
all_futures = client.map(worker_fn, tasks, pure=False)
logger.info(
f"Driver: submission complete — {len(all_futures)} futures in graph. "
f"Scheduler will distribute across all live workers."
)
ac = dask_as_completed(all_futures)
pbar = tqdm_bar(total=total, desc="Stage1 partition writes", unit="task")
n_ok = n_fail = n_leaves = 0
# Periodic-progress INFO is opt-in only — tqdm.set_postfix already shows
# ok/fail/leaves on the terminal, so emitting the same data to the log
# every 60s was operator clutter (~840 lines per 14h continental build).
# Set ``GH3_LOG_PROGRESS=1`` to re-enable for detached / tail-followed
# log workflows. WARN/ERROR lines (per-failure + end-of-phase summary)
# remain unconditional — those are actionable, not progress noise.
log_progress = os.environ.get('GH3_LOG_PROGRESS', '').strip().lower() in ('1', 'true', 'yes', 'on')
log_every_seconds = 60
next_log_t = _time.monotonic() + log_every_seconds if log_progress else float('inf')
try:
for fut in ac:
try:
result = fut.result()
if result.get('error') or not result.get('complete'):
_gk = _frag_name_granule_key(result.get('frag_name'))
if _gk is None:
lost_unattributed = True
else:
incomplete_keys.add(_gk)
else:
_note_complete(result.get('frag_name'), bool(result.get('empty')))
if result.get('error'):
n_fail += 1
logger.warning(
f"Stage1 task {result.get('frag_name')}: {result['error']}"
)
# Persist a structured failure record (driver-side,
# single-writer, append-only) so the build log's
# per-granule status can be enriched with a precise
# cause at finalize without log-grep. See
# _classify_load_h5_failure + _append_granule_failure.
failure = result.get('failure')
if failure:
_append_granule_failure(tmp_dir, result.get('frag_name'), failure)
else:
n_ok += 1
n_leaves += result.get('leaves', 0)
if result.get('complete') and result.get('frag_name') in previously_failed:
recovered.add(result.get('frag_name'))
except Exception as e:
n_fail += 1
logger.warning(f"Stage1 task raised: {type(e).__name__}: {e}")
# The future carries no task: find it by key (rare path, so
# the O(N) scan is paid only on failure). Unknown => the
# outcome cannot be attributed and is withheld.
_gk = None
try:
for _f, _t in zip(all_futures, tasks):
if _f.key == fut.key:
_gk = _frag_name_granule_key(_t[2])
break
except Exception:
_gk = None
if _gk is None:
lost_unattributed = True
else:
incomplete_keys.add(_gk)
finally:
fut.release()
pbar.update(1)
pbar.set_postfix(ok=n_ok, fail=n_fail, leaves=n_leaves)
if log_progress:
now = _time.monotonic()
if now >= next_log_t:
done = n_ok + n_fail
logger.info(
f"Streaming write: {done}/{total} done "
f"({100*done/total:.1f}%) — ok={n_ok} fail={n_fail} "
f"leaves={n_leaves}"
)
next_log_t = now + log_every_seconds
finally:
pbar.close()
# One rewrite: drop records of tasks that succeeded this run, merge
# duplicate lines. Must precede the merge's ``_cleanup_merged_tmp``.
_compact_granule_failures(tmp_dir, recovered)
# Reached only when the loop ran to its end (an interrupt or driver error
# propagates past here). Drained fully = every submitted task reported.
if stage1_outcome_callback is not None and not lost_unattributed and (n_ok + n_fail) == total:
stage1_outcome_callback(*_outcome())
if n_fail:
logger.error(
f"Streaming write: {n_fail}/{total} tasks failed. Their granules "
f"remain non-INDEXED and will be retried on the next resume."
)
# O(1) emptiness check — identical to the legacy _write_partitioned tail.
try:
return any(
entry.is_dir() and entry.name.startswith('h3_')
for entry in os.scandir(tmp_dir)
)
except FileNotFoundError:
return False
def _write_partitioned(
ddf: dask.dataframe.DataFrame,
tmp_dir: str,
part: int,
lat_col: str,
lon_col: str,
dat_col: str,
frag_names: Optional[List[str]] = None,
) -> bool:
"""
Write H3-indexed data to partitioned parquet files.
Parameters
----------
ddf : dask.dataframe.DataFrame
Input Dask DataFrame with H3 index
tmp_dir : str
Temporary directory for output
part : int
H3 partition resolution
lat_col, lon_col, dat_col : str
Column names for coordinates and datetime
frag_names : list of str, optional
Stable per-Dask-partition basenames (without ``.parquet`` extension)
used as ``name_function`` in ``to_parquet``. When provided, re-runs of
the same granule overwrite the same file in place (idempotent resume,
no shot duplication). Falls back to dask defaults when None.
Returns
-------
bool
True if any partition file was written, False if the tmp tree is
empty (e.g. all granules filtered out). The caller only needs an
emptiness check; building a full list of written paths via a recursive
glob scales as O(n_files) and is unbounded for global builds (millions
of fragments × shared filesystem stat latency).
"""
logger.info("Adding date and geometry columns to H3 database")
ddf = ddf.map_partitions(add_special_columns, lon_col=lon_col, lat_col=lat_col, dat_col=dat_col)
ddf = ddf.assign(year=ddf.datetime.dt.year)
ddf = dask_geopandas.from_dask_dataframe(ddf)
logger.info(f"Writing partitioned H3 data to temporary directory: {tmp_dir}")
# overwrite=False preserves any tmp fragments from prior killed runs;
# name_function with granule-derived names ensures re-extracted granules
# overwrite their own files in place rather than appending duplicates.
_to_parquet_kwargs = dict(
write_index=True,
overwrite=False,
compression='zstd',
partition_on=[f'h3_{part:02d}', 'year'],
# write_metadata_file=False: gh3 maintains its own gedih3_build_log.json
# sidecar. Forcing False prevents the global _metadata aggregation task,
# which holds every per-partition write output in scheduler memory until
# all millions of writes complete (the global-build memory plateau).
write_metadata_file=False,
compute=False,
)
if frag_names is not None:
# Closure captures frag_names; dask passes the partition index i.
_to_parquet_kwargs['name_function'] = lambda i: f"{frag_names[i]}.parquet"
write_task = ddf.to_parquet(tmp_dir, **_to_parquet_kwargs)
# optimize_graph=True (default) fuses read_h5 → h3_index → add_skip →
# add_special_columns → from_dask_dataframe → to_parquet into one task per
# input partition. Without fusion the scheduler eagerly produces reads
# ahead of writes and tens of thousands of intermediate dataframes pile up
# on workers. Fusion does NOT shuffle or reduce partition cardinality; with
# partition_on=[...] every input partition still emits its own file in each
# leaf hive directory (see dask/dask#8445 + #8487).
write_task = write_task.persist()
try:
progress(write_task)
# Check for worker errors before cancelling
from distributed import futures_of
futures = futures_of(write_task)
errors = [f for f in futures if f.status == 'error']
if errors:
errors[0].result() # Re-raises the original worker exception
finally:
logger.debug("Clearing dask workers")
client = get_dask_client()
client.cancel(write_task, force=True)
del write_task, ddf
# O(1) emptiness check: at least one h3_* leaf dir means stage 1 wrote
# something. Avoid recursive glob over the full tmp tree (would walk
# millions of fragments on a global build).
try:
return any(
entry.is_dir() and entry.name.startswith('h3_')
for entry in os.scandir(tmp_dir)
)
except FileNotFoundError:
return False
# Below this many entries a single driver-side ``shutil.rmtree`` is cheaper
# than dispatching; above it (e.g. one ``_complete/`` sentinel per
# granule × beam at continental scale) the removals fan out across workers.
_REMOVE_FANOUT_MIN_ENTRIES = 10000
def _remove_path(path):
"""Worker: best-effort removal of one file or directory tree. Picklable
top-level fn for parallel cleanup of large sentinel/fragment sets via
``parallel_map``. Idempotent — a missing path is success (already
cleaned). Returns the error string on failure, ``None`` otherwise."""
try:
if os.path.isdir(path) and not os.path.islink(path):
shutil.rmtree(path)
else:
os.unlink(path)
except FileNotFoundError:
pass
except OSError as e:
return f"{type(e).__name__}: {e}"
return None
def _remove_tree_fanout(root: str, desc: str) -> None:
"""Remove ``root`` and everything under it without a serial driver sweep.
One ``os.scandir`` of ``root`` (a single readdir), then — when it holds at
least ``_REMOVE_FANOUT_MIN_ENTRIES`` entries and a dask Client is
registered — each entry is removed on a worker via :func:`_remove_path`
in batches, and the emptied ``root`` last. Narrow trees, or no Client,
take a plain ``shutil.rmtree``. Best-effort: failures are logged as a
WARNING and never raised — cleanup must not fail a finished build.
Parameters
----------
root : str
Directory to remove. A missing directory is a no-op.
desc : str
Label for the ``parallel_map`` dispatch log line.
"""
if not os.path.isdir(root):
return
try:
entries = [e.path for e in os.scandir(root)]
if len(entries) >= _REMOVE_FANOUT_MIN_ENTRIES and get_dask_client() is not None:
from .parallel import parallel_map
for path, err in parallel_map(entries, _remove_path, desc=desc,
unit='entry', batch_size=2000):
if err:
logger.debug(f"Could not remove {path}: {err}")
shutil.rmtree(root)
except Exception as e:
logger.warning(f"Could not clean up {root}: {type(e).__name__}: {e}")
def _cleanup_merged_tmp(tmp_dir: str, merge_failed: bool) -> None:
"""Drop the Stage 1 scaffolding under ``tmp_dir`` once a merge succeeded.
Every successful ``h3_merge_files`` already deleted its fragment dir
(``rm_src=True``), so what remains is resume/forensics scaffolding:
``_merge_progress.txt``, one ``_complete/`` sentinel per granule × beam,
the emptied ``h3_*`` dirs, and the failure sidecars. Once every partition
is merged the finalized partition metadata is the resume record
(reconcile Pass A), so none of it is needed — and a stale
``_merge_progress.txt`` is actively harmful: it is the L2 merge-resume
signal, and would turn the next update's ``gh3_build`` into a silent
merge-only no-op.
Tiers, each gated on O(1) signals the build already persisted:
* any merge failed in this run → keep everything (the next resume needs
the progress file, the unmerged fragments and the failure records);
* merge clean → drop ``_merge_progress.txt`` first (a crash mid-cleanup
must not leave the one hazardous file behind). It goes even when a
failure sidecar remains: every partition it lists is merged, and a
kept copy would also make the next merge skip those partitions' new
fragments. ``_complete/`` goes right after, for the same reason: the
sentinels are an in-flight Stage 1 resume record, and a kept set would
make a later build with a wider ROI skip tasks that were empty under
the old one;
* then keep the rest while ``_merge_failed_granules.jsonl`` awaits the
CLI fold into the build log, or ``_granule_failures.jsonl`` holds
Stage 1 forensics (read by the end-of-build advisory and
``gh3_doctor``'s ``tmp_partitions_health``); otherwise remove
``tmp_dir`` entirely.
Parameters
----------
tmp_dir : str
The Stage 1 partitions directory (``<build tmp>/partitions``). Only
this directory is touched — never its parent, which may be a
user-supplied ``--tmpdir`` holding unrelated files.
merge_failed : bool
Whether any partition merge failed in this run.
"""
if merge_failed:
logger.info(f"Keeping {tmp_dir} for resume: merge failures are pending recovery")
return
try:
os.unlink(os.path.join(tmp_dir, '_merge_progress.txt'))
except FileNotFoundError:
pass
except OSError as e:
logger.warning(f"Could not remove stale merge progress file in {tmp_dir}: {e}")
return
try:
_discard_complete_sentinels(tmp_dir, desc="Cleaning completion sentinels")
except OSError as e:
logger.warning(f"Could not discard completion sentinels in {tmp_dir}: {e}")
return
if os.path.exists(os.path.join(tmp_dir, _MERGE_FAILED_GRANULES_FILENAME)):
logger.info(f"Keeping {tmp_dir}: merge-failed granules await the build-log fold")
return
if os.path.exists(os.path.join(tmp_dir, _GRANULE_FAILURES_FILENAME)):
logger.info(
f"Keeping {tmp_dir} for Stage 1 failure forensics "
f"(gh3_doctor -i <database> -t {tmp_dir} --check tmp_partitions_health); "
f"delete it once reviewed"
)
return
logger.info(f"Cleaning up build scaffolding in {tmp_dir}")
_remove_tree_fanout(tmp_dir, desc="Cleaning tmp partitions")
def _merge_and_finalize(
tmp_dir: str,
h3_dir: str,
replacement_tmp: Optional[str] = None,
extra_dirs: Optional[List[str]] = None,
replacement_digest: Optional[str] = None,
) -> List[str]:
"""
Merge temporary partitioned files into the final H3 database.
Tracks progress via an append-only ``_merge_progress.txt`` file in
``tmp_dir``. On resume, already-merged partitions are skipped. When every
partition merged, ``tmp_dir`` is cleaned up before returning — see
:func:`_cleanup_merged_tmp` for what is kept, and when.
Parameters
----------
tmp_dir : str
Temporary directory containing partitioned files
h3_dir : str
Final output directory for H3 database
replacement_tmp : str, optional
The build tmp dir holding a granule replacement's journal (#47), as
passed to :func:`read_replacement_journal`: every merge replaces the
journaled granules, read once per worker process and checked against
``replacement_digest`` (:func:`_h3_merge_replacing`), so no task
carries the whole target map.
extra_dirs : list of str, optional
More ``<tmp_dir>/h3_<p>=<cell>/year=<year>`` targets, with or without
fragments: the (cell, year) files holding a replaced granule's old
rows, which are rewritten even when no new row lands there.
replacement_digest : str, optional
:func:`_targets_digest` of the targets the caller indexed, so every
merge uses exactly those; when omitted, the journal's own digest.
Returns
-------
list of str
List of merged H3 parquet file paths (without those a replacement
left empty and removed)
"""
logger.info(f"Merging H3 partitions into final database path: {h3_dir}")
# Pre-clean: act on any _merge_failures sentinels left by a prior run.
# Removes the 0-byte / truncated parquets that would otherwise re-fail
# this merge for the same partitions. Idempotent + bounded by the
# failure-list size (never scans the healthy tree). Paired with the
# CLI-side ``apply_merge_failures_to_logger`` fold so the granules in
# those partitions are PENDING for Stage 1 to re-extract.
_pc = preclean_merge_failures(tmp_dir)
if any(_pc.values()):
logger.info(
f"Merge pre-clean: removed {_pc['parquets_removed']} bad "
f"parquets + {_pc['tmps_removed']} .tmp orphans across "
f"{_pc['partitions_cleaned']} partition(s)."
)
# Drop any pre-existing bbox index BEFORE the first worker rewrites a
# partition file: the merge changes data envelopes, a stale index
# silently under-selects (the one unacceptable failure), and a crash
# mid-merge must not leave it behind. Deleting at entry (not at the
# tail) also covers concurrent readers during the hours-long merge
# window and every resume path. Rebuild with gh3_bbox_index after the
# build; absence only costs speed.
from .gh3driver import invalidate_bbox_index
if invalidate_bbox_index(h3_dir):
logger.info("Removed stale _bbox_index.parquet (rebuild with gh3_bbox_index)")
# List candidate tmp partitions. Two-level walk:
# 1. Driver-side os.scandir of <tmp_dir> for h3_* — one readdir, fast.
# 2. Per-h3 year=*/ listing fanned out across all workers via
# client.map(_list_year_subdirs); cumulative readdir latency on
# shared GPFS is parallelized instead of paid serially on the driver.
# Empty year dirs are NOT filtered here — h3_merge_files short-circuits
# on no-input via `if len(files) == 0: return`.
client = get_dask_client()
try:
h3_part_dirs = sorted(
e.path for e in os.scandir(tmp_dir)
if e.is_dir(follow_symlinks=False) and e.name.startswith('h3_')
)
except OSError:
h3_part_dirs = []
if h3_part_dirs and client is not None:
year_lists = client.gather(client.map(_list_year_subdirs, h3_part_dirs, pure=False))
tmp_h3_dirs = [p for sub in year_lists for p in sub]
else:
tmp_h3_dirs = [p for d in h3_part_dirs for p in _list_year_subdirs(d)]
if extra_dirs:
listed = {d.rstrip('/') for d in tmp_h3_dirs}
tmp_h3_dirs += [d + '/' for d in sorted({d.rstrip('/') for d in extra_dirs} - listed)]
os.makedirs(h3_dir, exist_ok=True)
# Stale .merge.tmp cleanup is delegated to h3_merge_files (worker-side,
# scoped to the single partition's odir). The previous global sweep
# `glob(h3_dir/h3_*/*/*.merge.tmp)` walked the entire final database from
# the driver — another serial GPFS pass that scaled with DB size for no
# benefit (parquet_merge_files already overwrites its own stale tmp).
# Resume support: skip already-merged partitions
merge_progress_file = os.path.join(tmp_dir, '_merge_progress.txt')
merged_parts = set()
if os.path.exists(merge_progress_file):
with open(merge_progress_file, 'r') as f:
merged_parts = {line.strip() for line in f if line.strip()}
if merged_parts:
logger.info(f"Resuming merge: {len(merged_parts)} partitions already merged")
remaining_dirs = [d for d in tmp_h3_dirs if d.rstrip('/') not in merged_parts]
failed_count = 0
if not remaining_dirs:
logger.info("All partitions already merged (resume)")
else:
from dask.distributed import as_completed as dask_as_completed
from tqdm import tqdm as tqdm_bar
# Submit every remaining merge to the scheduler at once and let dask
# distribute work across all available workers. No driver-side
# throttling: scaling decisions belong to the scheduler.
# pure=False: merges read fragments, write outputs and delete
# sources — dask must never key-cache their results/exceptions
# (this phase keeps every future alive for its whole duration, so
# any same-process resubmission of a dir would otherwise get the
# cached outcome without re-executing).
# A replacement's targets come from its journal, not inside each task:
# a 10k-granule map added ~7 ms of serialization per task.
merge_fn = (functools.partial(_h3_merge_replacing, replacement_tmp=replacement_tmp,
digest=replacement_digest or _replacement_digest(replacement_tmp))
if replacement_tmp else h3_merge_files)
futures_list = client.map(
merge_fn,
remaining_dirs,
out_dir=h3_dir,
rm_src=True,
replace=False,
pure=False,
)
futures: Dict[Any, str] = dict(zip(futures_list, remaining_dirs))
merged_count = 0
pbar = tqdm_bar(total=len(remaining_dirs), desc="Merging partitions", unit="part")
# Incremental manifest refresh: at continental scale the merge
# phase runs for hours, during which any consumer reading the
# database (gh3_load, gh3_extract) would see a stale manifest
# sentinel from the last full refresh — including the previous
# build, if any. Periodic in-merge refreshes give those consumers
# a fresh view of the partially-built DB. Trigger: every
# ``manifest_refresh_every`` successful merges. The refresh itself
# is O(N_merged_so_far) — pure in-memory derivation from
# ``_merge_progress.txt`` + one atomic file write; no tree walk.
# Skip in tests / library callers that disable via env var.
manifest_refresh_every = max(
1, int(os.environ.get('GH3_MANIFEST_REFRESH_EVERY', '1000'))
)
_next_refresh_at = manifest_refresh_every
# Do NOT call future.release() per completion — keeping the futures
# alive (via futures_list) makes the scheduler retain finished task
# records, so the dashboard progress bar fills monotonically (X done
# / 47k total) instead of draining as tasks are released. Memory cost
# is ~tens of MB on the scheduler for the entire merge phase.
for future in dask_as_completed(futures_list):
d = futures.pop(future)
try:
future.result()
with open(merge_progress_file, 'a') as f:
f.write(d.rstrip('/') + '\n')
merged_count += 1
# Incremental manifest refresh — see comment block above
# the loop. Bounded by ``manifest_refresh_every`` so the
# total cost across a 50k-merge run is ~50 refreshes,
# each <1s — negligible against the multi-minute merge.
if merged_count >= _next_refresh_at:
try:
_paths = _derive_merged_output_paths(merge_progress_file, h3_dir)
if _paths:
generate_manifest(h3_dir, files=_paths, tree_shape='h3db')
except Exception as _re:
# Manifest refresh is best-effort — never let a
# write failure abort an in-flight merge.
logger.debug(f"Incremental manifest refresh skipped: {_re}")
_next_refresh_at = merged_count + manifest_refresh_every
except Exception as e:
failed_count += 1
# Preserve the h3_cell + year in the relative path so the
# operator can locate the failing partition without scanning
# the tmp tree. Combined with the [file=<path>] suffix that
# parquet_merge_files now attaches inside ``e``, a single
# grep on the WARN line gives both the partition and the
# exact bad fragment.
logger.warning(
f"Merge failed for {os.path.relpath(d, tmp_dir)}: "
f"{type(e).__name__}: {e}"
)
# Persist an atomic per-failure sentinel so L1 resume +
# L2 doctor can recover without re-running the full merge
# to rediscover what broke. See _emit_merge_failure_sentinel.
_emit_merge_failure_sentinel(tmp_dir, d, e)
# If the failure is a known-bad-fragment class (0-byte /
# truncated parquet), parse the granules whose fragments
# live in this partition and record them for the CLI to
# flip back from INDEXED → MERGE_FAILED. Without this,
# L1 resume pre-clean (Phase 2.2) would unlink the bad
# fragment but leave the granule INDEXED, permanently
# dropping its rows. Only fires for the curated marker
# set; infrastructure errors (disk full, etc.) skip this
# and the granule stays INDEXED.
if _is_recoverable_fragment_error(e):
grans = _granules_in_partition_dir(d)
_emit_merge_failed_granules(tmp_dir, d, grans, e)
pbar.update(1)
pbar.set_postfix(ok=merged_count, fail=failed_count)
pbar.close()
del futures, futures_list
if failed_count > 0:
logger.error(f"{failed_count} partition merges failed. Re-run to retry.")
# Derive h3_files + h3_subdirs from the authoritative _merge_progress.txt
# instead of two driver-side GPFS scans (``glob('h3_*/*/*.parquet')`` +
# ``glob('h3_*/')``). Zero GPFS metadata ops — pure in-memory derivation
# via the deterministic ``h3_merge_files`` naming contract. See
# ``_derive_merged_output_paths``.
h3_files: List[str] = sorted(
_derive_merged_output_paths(merge_progress_file, h3_dir)
)
h3_subdirs_set: set = {
os.path.dirname(os.path.dirname(p)) + '/' for p in h3_files
}
if replacement_tmp:
# A replacement may have emptied some; their cells still get their
# summary recompiled (or removed) below.
h3_files = [f for f in h3_files if os.path.exists(f)]
# Check a sample of merged files for NaN-only columns
if h3_files:
import geopandas as _gpd
for f in h3_files[:5]:
try:
_sample = _gpd.read_parquet(f)
check_nan_only_columns(_sample, context=f"Build partition {os.path.basename(f)}: ", logger=logger)
except Exception:
pass
logger.info("Compiling H3 metadata files")
h3_subdirs = sorted(h3_subdirs_set)
meta_tasks = [dh3_merge_metadata(i) for i in h3_subdirs]
meta_tasks = dask.persist(*meta_tasks, optimize_graph=False)
progress(meta_tasks)
meta_files = list(dask.compute(*meta_tasks))
del meta_tasks
# Generate manifest for accelerated file listing. (The bbox index was
# already invalidated at merge ENTRY — before the first partition
# write — so nothing here can bump the root dir mtime after this.)
logger.info("Generating file manifest")
generate_manifest(h3_dir, tree_shape='h3db')
# Last: h3_files was derived from _merge_progress.txt above, and the
# cleanup deletes it.
_cleanup_merged_tmp(tmp_dir, merge_failed=failed_count > 0)
return h3_files
def _scan_year_file_for_update(year_pf, new_cols):
"""Worker: cheap footer + sidecar scan for one ``(cell, year)`` file.
Returns the granule list when the file still needs the new columns,
``None`` when it's already done / has no sidecar / hit a recoverable
read error (corrupt parquet footer, malformed sidecar JSON). The
driver maps this across every year file, then uses the returned
granule lists to inline the per-task h5 paths via the driver-side
``all_soc`` dict — no ``client.scatter`` needed (the fresh-build
streaming driver established that scatter-broadcast hangs on
SSH-tunneled clusters and the inlined-args pattern is the
established fix).
Worker-side ``try/except`` swallows the recoverable errors so the
driver-side ``as_completed(..., with_results=True)`` collection
loop doesn't crash on a single corrupt fragment — matches the
pattern in ``s3_etl_subset`` (line 342) which collects ``None``
sentinels from failed workers and aggregates a single warning.
"""
try:
new_cols = set(new_cols)
if new_cols:
existing_cols = set(read_parquet_schema(year_pf)['column'].tolist())
if new_cols.issubset(existing_cols):
return None
year_meta_path = year_pf.replace('.parquet', PARTITION_META_FILENAME)
if not os.path.exists(year_meta_path):
return None
meta = json_read(year_meta_path)
return meta.get('granules', []) or None
except Exception as e:
logger.warning(f"Scan failed for {os.path.basename(year_pf)}: {type(e).__name__}: {e}")
return None
# ── Variable-add inverted fan-out: shared helpers ────────────────────────
# The variable-add path reads each unique granule h5 EXACTLY ONCE (the
# fresh build's per-(cell,year) approach re-read each granule ~40× — a
# DB-wide 39.75× redundancy measured on the production tree, 2.64M reads
# for 66.5k unique granules). Stage 1 (``_var_fan_granule``) reads a
# granule once and fans its shots to every owning (cell, year) as a tiny
# intermediate parquet under ``tmp_dir/_var_frags/``; Stage 2
# (``_var_merge_cell_year``) concats a cell-year's fragments and joins
# them into the BASE parquet (single self-contained file per partition —
# no sidecars). Shots are routed to cells by matching ``shot_number``
# against the existing base parquet (never by recomputing H3 from the
# new product's coords — coordinate drift between products would
# misroute boundary shots and silently NaN the left-join).
_VAR_FRAG_DIRNAME = '_var_frags'
_VAR_FAN_COMPLETE_DIRNAME = '_var_fan_complete'
_VAR_MERGE_PROGRESS_FILENAME = '_var_merge_progress.txt'
def _var_frag_token(year_pf: str) -> str:
"""Reversible ``h3_<part>=<cell>__year=<YYYY>`` token for a year_pf.
Used as the per-(cell, year) fragment subdir name so Stage 2 can find
a cell-year's fragments from the same base path Stage 1 fanned into.
"""
ydir = os.path.dirname(year_pf)
return f"{os.path.basename(os.path.dirname(ydir))}__{os.path.basename(ydir)}"
def _var_fragment_dir(tmp_dir: str, year_pf: str) -> str:
return os.path.join(tmp_dir, _VAR_FRAG_DIRNAME, _var_frag_token(year_pf))
def _var_fan_sentinel_path(tmp_dir: str, granule_key: str) -> str:
return os.path.join(tmp_dir, _VAR_FAN_COMPLETE_DIRNAME, f'{granule_key}.done')
def _emit_var_fan_sentinel(tmp_dir: str, granule_key: str) -> None:
"""Touch the per-granule fan-complete sentinel. Idempotent + atomic.
Mirrors ``_emit_complete_sentinel``: zero-byte file, existence IS the
signal, ``open(..., 'x')`` is atomic on shared GPFS so concurrent
emitters can't leave a half-written sentinel.
"""
path = _var_fan_sentinel_path(tmp_dir, granule_key)
os.makedirs(os.path.dirname(path), exist_ok=True)
try:
with open(path, 'x'):
pass
except FileExistsError:
pass
def _cached_base_shots(year_pf: str):
"""Per-worker cached sorted ``shot_number`` array for one base parquet.
A granule fans to ~40 cells, and each cell-year base is needed by every
granule crossing it — so the same base shot array is read many times by
one worker. The LRU cache (per worker process) collapses those to one
read per (cell, year) the worker touches. Bounded at 256 arrays (~few
hundred MB worst case). Keyed by the file's mtime and size too, so a
worker that outlives one update (a persistent cluster, a notebook) never
routes against a file rewritten since. A read error raises — the granule
is then failed and retried, never silently given zero rows.
"""
st = os.stat(year_pf)
return _base_shots(year_pf, st.st_mtime_ns, st.st_size)
@functools.lru_cache(maxsize=256)
def _base_shots(year_pf: str, _mtime_ns: int, _size: int):
return np.sort(pd.read_parquet(year_pf, columns=['shot_number'])['shot_number'].to_numpy())
def _var_fan_granule(task, *, new_product_vars, tmp_dir, include_source=False, split_products=False):
"""Stage 1 worker: read ONE granule's h5(s) once, fan shots to cells.
``task`` = ``(granule_key, h5_by_prod, year_pfs)`` where:
* ``granule_key`` — ``O{orbit:05d}_{granule:02d}_T{track:05d}`` (the
soc_file_tree pivot key; also the fragment/sentinel basename).
* ``h5_by_prod`` — ``{product: h5_path}`` for THIS granule, already
filtered driver-side to products in ``new_product_vars``.
* ``year_pfs`` — base parquet paths this granule contributes to
(its owning ``(cell, year)`` units, pre-resolved driver-side).
Reads each product h5 ONCE (``shots=None`` — the ``[:]`` full-column
read is GPFS-optimal; a ``shots=`` filter saves no disk I/O and only
adds a per-beam isin, see benchmark). Joins products on shot_number,
suffix-renames, then for each owning cell-year writes
``tmp_dir/_var_frags/<token>/<granule_key>.parquet`` containing only
the rows whose shot_number is in that base's shot set — atomically,
only when non-empty.
Resume: a ``.done`` sentinel short-circuits an already-fanned granule.
Failures are classified (``_classify_load_h5_failure``) and returned
so the driver can record them without aborting the job; NO sentinel is
written on failure, so the next resume retries the granule.
``include_source=True`` also emits each product's ``root_file_<prod>``
and ``root_beam_<prod>`` columns, as a build writes them; a collection
of products emits them for those products only. The product fill needs
them, and so does a variable add of a product the database had no
source columns for (#68): they are what a granule replacement matches.
``split_products=True`` (product fill) fans each product on its own —
fragments ``<granule_key>.<prod>.parquet``, each product's frame freed
before the next is read — instead of outer-joining them first. The join
upcasts integer and flag columns (and doubles them in memory); a fill
merges the per-product fragments as separate patches instead.
"""
granule_key, h5_by_prod, year_pfs = task
if os.path.exists(_var_fan_sentinel_path(tmp_dir, granule_key)):
return {'granule': granule_key, 'skipped': True, 'fragments': 0}
def _fan(frame, token):
"""Write ``frame``'s rows to each owning cell-year; returns fragments written."""
frame = frame.reset_index() # shot_number back as a column
sn = frame['shot_number'].to_numpy()
n = 0
for year_pf in year_pfs:
base_shots = _cached_base_shots(year_pf)
if base_shots.size == 0:
continue
mask = np.isin(sn, base_shots)
if not mask.any():
continue
frag_dir = _var_fragment_dir(tmp_dir, year_pf)
os.makedirs(frag_dir, exist_ok=True)
with AtomicFileWriter(os.path.join(frag_dir, f'{token}.parquet')) as tmp_path:
frame.loc[mask].to_parquet(tmp_path, engine='pyarrow', index=False, compression='zstd')
n += 1
return n
# Read each product h5 ONCE; outer-join on the shot_number index (or,
# split_products, fan each product as soon as it is read).
new_df = None
n_frag = 0
prod_frags = {}
for prod, h5_path in h5_by_prod.items():
var_list = [v for v in (new_product_vars.get(prod) or []) if v != 'shot_number']
if not var_list:
continue
try:
df = load_h5(h5_path, columns=['shot_number'] + var_list, shots=None,
include_source=(include_source if isinstance(include_source, bool)
else prod in include_source))
except Exception as exc:
return {
'granule': granule_key,
'error': f"{type(exc).__name__}: {exc}",
'failure': _classify_load_h5_failure(exc, h5_by_prod),
}
if split_products:
prod_frags[prod] = 0
if df is None or df.empty:
continue
suffix = f"_{prod.lower()}"
df = df.rename(columns=lambda x: x if x.endswith(suffix) else f"{x}{suffix}")
if split_products:
prod_frags[prod] = _fan(df, f'{granule_key}.{prod}')
n_frag += prod_frags[prod]
del df
continue
new_df = df if new_df is None else new_df.join(df, how='outer')
if new_df is not None and not new_df.empty:
n_frag += _fan(new_df, granule_key)
# A structurally-empty granule still gets its sentinel so a variable-add
# resume does not keep retrying it. A fill product that routed no row
# wrote nothing: no sentinel, so a retry re-reads it instead of reporting
# it done.
if not split_products or (prod_frags and all(prod_frags.values())):
_emit_var_fan_sentinel(tmp_dir, granule_key)
out = {'granule': granule_key, 'fragments': n_frag}
if split_products:
out['product_fragments'] = prod_frags
return out
def _var_merge_cell_year(year_pf, *, tmp_dir, fill=False, new_schema=None):
"""Stage 2 worker: merge a cell-year's fanned fragments into its base.
Concats every ``<granule_key>.parquet`` Stage 1 wrote for this
``(cell, year)`` (each already cell-specific — read exactly once),
dedups on shot_number, and column-joins the result into the BASE
parquet via ``parquet_join_columns`` (rowgroup-wise, bounded memory,
atomic ``.join.tmp`` + ``os.replace``). ``parquet_join_columns``
filters columns already present in the base, so a re-run after a
partial merge never duplicates columns. Refreshes the per-year
metadata JSON afterward.
On a successful merge the worker deletes THIS cell-year's fragment
directory immediately — the same ``rm_src=True`` discipline the
fresh-build merge (`h3_merge_files`) uses. This drops the consumed
fragments distributed across all workers as the merge proceeds,
instead of leaving every fragment for one giant serial
``shutil.rmtree`` on the driver at end-of-run (which on a continental
fan-out is millions of files and pins the driver in disk-sleep for
hours while blocking the COMPLETED build-log write). Resume-safe:
once the base has the new columns, a re-run finds no fragments and
returns ``None`` (the skip-check also catches it via the base
schema), so deleting fragments after the atomic join never loses
recoverable work.
Returns the base path on a real merge, ``None`` when the cell-year had
no fragments (granule(s) produced no matching shots here).
``fill=True`` (product backfill) writes the fragments into columns the
base already has, only where the base is null
(``parquet_fill_columns``: rowgroup-wise, atomic, existing values never
overwritten) instead of joining new columns.
``new_schema`` (variable update): the new columns' schema, one for every
file of the run (:func:`_var_new_columns_schema`). The fragments are
conformed to it, and a file with no fragment still gets the columns, as
null: every file must end with the same columns in the same types, or
readers reject the mix, and a later product backfill (which never adds
columns) could not fill them.
"""
frag_dir = _var_fragment_dir(tmp_dir, year_pf)
nv_path = os.path.join(frag_dir, '_newvars.parquet')
try:
# Exclude the consolidated _newvars.parquet: a resume after a
# crash-between-write-and-progress would otherwise fold it back
# into itself (harmless but wasteful).
frags = [os.path.join(frag_dir, f) for f in os.listdir(frag_dir)
if f.endswith('.parquet') and f != '_newvars.parquet']
except FileNotFoundError:
frags = []
if not frags and (fill or new_schema is None
or set(new_schema.names) <= set(read_parquet_schema(year_pf)['column'])):
return None
if fill:
# Per-product fragments (granules own disjoint shots): one compact
# Arrow table per product, no pandas round trip, no consolidated file.
import pyarrow as pa
by_prod: Dict[str, List[str]] = {}
for f in frags:
by_prod.setdefault(os.path.basename(f).rsplit('.', 2)[-2], []).append(f)
patches = [pa.concat_tables([pq.read_table(f) for f in fs], promote_options='permissive')
for fs in by_prod.values()]
parquet_fill_columns(year_pf, patches, key_col='shot_number', append_new=False)
del patches
else:
if frags:
cat = pd.concat([pd.read_parquet(f) for f in frags], ignore_index=True)
cat = cat.drop_duplicates(subset='shot_number', keep='first')
if new_schema is None:
with AtomicFileWriter(nv_path) as tmp_path:
cat.to_parquet(tmp_path, engine='pyarrow', index=False, compression='zstd')
else:
import pyarrow as pa
table = (_conform_table_to_schema(pa.Table.from_pandas(cat, preserve_index=False), new_schema)
if frags else new_schema.empty_table())
os.makedirs(frag_dir, exist_ok=True)
with AtomicFileWriter(nv_path) as tmp_path:
pq.write_table(table, tmp_path, compression='zstd')
del table
if frags:
del cat
parquet_join_columns([year_pf, nv_path], year_pf, key_col='shot_number')
# Refresh ONLY the column list + dtypes in the per-year metadata. A
# variable-add changes columns, never shots/dates/granules — so patch
# the existing JSON from a footer-only schema read instead of calling
# h3_write_metadata (which re-reads root_file_l2a + datetime to
# recompute stats that have not changed).
_refresh_year_columns_meta(year_pf)
# rm_src: drop this cell-year's consumed fragments now (distributed,
# one rmtree per worker) rather than accumulating all of them for a
# serial driver-side sweep at the end. The base already carries the
# columns atomically, so this is safe post-join.
shutil.rmtree(frag_dir, ignore_errors=True)
return year_pf
def _refresh_year_columns_meta(year_pf: str) -> None:
"""Patch ``columns`` + ``column_dtypes`` in a per-year metadata JSON.
Footer-only schema read; preserves all other fields (shot_count,
shot_range, date_range, granules — unchanged by a variable-add). If the
metadata JSON is absent (legacy partition), falls back to the full
``h3_write_metadata`` so the cache is still populated.
"""
meta_path = year_pf.replace('.parquet', PARTITION_META_FILENAME)
if not os.path.exists(meta_path):
h3_write_metadata(year_pf)
return
try:
meta = json_read(meta_path) or {}
schema = read_parquet_schema(year_pf)
cols = schema['column'].tolist()
meta['columns'] = cols
meta['column_dtypes'] = dict(zip(cols, schema['dtype'].tolist()))
meta['last_modified'] = now()
json_write(meta, meta_path, rewrite=True)
except Exception as e:
logger.warning(
f"Column-meta refresh failed for {os.path.basename(year_pf)}: "
f"{type(e).__name__}: {e}; falling back to full metadata write"
)
h3_write_metadata(year_pf)
# The per-product provenance columns a build writes (``load_h5(include_source=True)``,
# product-suffixed): the source file's basename and the beam.
_SOURCE_COLUMNS = ('root_file', 'root_beam')
def _var_source_schema(h5_path, *, prod, variables, include_source=False):
"""Worker: the schema a variable update's fan writes for one product.
Reads ``variables`` of one granule's ``prod`` file exactly as
:func:`_var_fan_granule` does (``load_h5``, product-suffixed names), so
the types are the ones its fragments carry, before any outer join
between products can upcast them.
"""
import pyarrow as pa
df = load_h5(h5_path, columns=['shot_number'] + list(variables), shots=None, include_source=include_source)
suffix = f"_{prod.lower()}"
df = df.rename(columns=lambda x: x if x.endswith(suffix) else f"{x}{suffix}")
return pa.Schema.from_pandas(df.reset_index(), preserve_index=False).remove_metadata()
def _var_new_columns_schema(client, gran_h5, product_vars, source_products=()):
"""One schema for every file a variable update writes, from one granule per product.
Deterministic (the first granule key holding the product), so a resumed
run derives the same schema. A product no granule has a file for is left
out, with a warning: no file gets its columns, so the files stay alike.
``source_products`` also get their source columns (``_SOURCE_COLUMNS``).
"""
import pyarrow as pa
futures = {}
for prod, variables in product_vars.items():
variables = [v for v in (variables or []) if v != 'shot_number']
if not variables:
continue
src = next((gran_h5[k][prod] for k in sorted(gran_h5) if prod in gran_h5[k]), None)
if src is None:
logger.warning(f"Variable update: no granule has a local {prod} file; {prod} variables are not added. "
f"Download {prod} and re-run to add them")
continue
futures[prod] = client.submit(_var_source_schema, src, prod=prod, variables=variables,
include_source=prod in source_products, pure=False)
fields = {}
for prod, fut in futures.items():
for field in fut.result():
fields.setdefault(field.name, field)
if 'shot_number' not in fields:
return None
return pa.schema([fields.pop('shot_number')] + list(fields.values()))
def _fan_merge_products(h3_dir, gran_h5, gran_year_pfs, product_vars, tmp_dir, *, fill=False, year_pfs=None,
source_products=()):
"""Stages 1-3 of the inverted granule fan-out, shared by variable-add and product-fill.
Reads each granule's product h5(s) exactly once and fans its shots to the
(cell, year) base files that own them (``_var_fan_granule``), then merges
each touched cell-year's fragments into its base (``_var_merge_cell_year``:
rowgroup-wise, atomic, fragments dropped per merge), then refreshes the
per-cell metadata and the manifest. Resume state (per-granule ``.done``
sentinels, ``_var_merge_progress.txt``) lives under ``tmp_dir``; it is
removed only after a fully clean run.
Parameters
----------
h3_dir : str
Database root.
gran_h5 : dict
``{granule_key: {product: h5_path}}`` — the products to read per granule.
gran_year_pfs : dict
``{granule_key: [base year_pf, ...]}`` — the files each granule owns rows in.
product_vars : dict
``{product: [variable, ...]}`` to read (unsuffixed names).
tmp_dir : str
Scratch root for fragments and resume state (one per mode).
fill : bool, default False
False joins new columns (variable-add). True writes into existing
null cells only (product backfill) and reads ``root_file_<prod>`` too.
year_pfs : list of str, optional
Variable-add: every file still lacking the new columns. Files no
granule routes rows to are merged too, so they get the columns as
null (:func:`_var_merge_cell_year`).
source_products : collection of str, optional
Variable-add: products whose ``root_file_<prod>`` / ``root_beam_<prod>``
columns are written too (a fill always writes them).
Returns
-------
dict
``{'updated_files': [...], 'failed_granules': set, 'failed_products':
{granule: set}, 'fragments': int}``. A granule is failed when its read
failed or any of its files failed to merge; everything else it owns was
written. ``failed_products`` (fill only): products of an otherwise
successful granule that routed no row.
"""
from dask.distributed import as_completed as dask_as_completed
from tqdm import tqdm as tqdm_bar
client = get_dask_client()
if client is None:
raise GediError("_fan_merge_products requires a registered dask Client")
os.makedirs(tmp_dir, exist_ok=True)
# ── Stage 1: fan — read each granule ONCE, write per-cell fragments ──
fan_fn = functools.partial(
_var_fan_granule, new_product_vars=product_vars, tmp_dir=tmp_dir,
include_source=True if fill else frozenset(source_products), split_products=fill,
)
fan_tasks = [(ot, gran_h5[ot], gran_year_pfs[ot]) for ot in gran_year_pfs]
fan_futures = client.map(fan_fn, fan_tasks, pure=False)
fan_prev_failed = {r.get('frag_name') for r in _read_granule_failures(tmp_dir)}
fan_recovered: set = set()
fut_to_gran = {f: t[0] for f, t in zip(fan_futures, fan_tasks)}
failed_granules = set()
failed_products: Dict[str, set] = {}
n_frag = n_fan_fail = 0
pbar = tqdm_bar(total=len(fan_tasks), desc="Stage1 granule fan", unit="granule")
try:
for fut in dask_as_completed(fan_futures):
try:
res = fut.result()
if res.get('error'):
n_fan_fail += 1
failed_granules.add(res.get('granule'))
logger.warning(f"Stage1 granule {res.get('granule')}: {res['error']}")
failure = res.get('failure')
if failure:
_append_granule_failure(tmp_dir, res.get('granule'), failure)
else:
n_frag += res.get('fragments', 0)
# A fill product listed by these files but routing no row
# to them filled nothing: leave it pending, never "filled".
empty = sorted(p for p, n in (res.get('product_fragments') or {}).items() if not n)
# failed_granules / failed_products track later-stage and
# pending-product failures separately; recover only a
# granule whose read left nothing pending.
if res.get('granule') in fan_prev_failed and not (fill and empty):
fan_recovered.add(res.get('granule'))
if fill and empty:
failed_products.setdefault(res.get('granule'), set()).update(empty)
logger.warning(f"Product backfill: granule {res.get('granule')} {', '.join(empty)} "
f"matched no rows in the files that list it; left pending")
except Exception as e:
n_fan_fail += 1
failed_granules.add(fut_to_gran.get(fut))
logger.warning(f"Stage1 granule task raised: {type(e).__name__}: {e}")
finally:
fut.release()
pbar.update(1)
pbar.set_postfix(fragments=n_frag, failed=n_fan_fail)
finally:
pbar.close()
_compact_granule_failures(tmp_dir, fan_recovered)
if n_fan_fail:
logger.error(
f"Stage1: {n_fan_fail} granule(s) failed to read. Their fragments "
f"are absent; affected cells keep their prior columns. Re-run to retry."
)
# ── Stage 2: merge — concat fragments into each base parquet ────────
touched_year_pfs = sorted({yp for yps in gran_year_pfs.values() for yp in yps} | set(year_pfs or ()))
merge_progress_file = os.path.join(tmp_dir, _VAR_MERGE_PROGRESS_FILENAME)
merged_set = set()
# A fill never skips a file as "already merged": re-filling is idempotent
# (null cells only), and a retry must merge a granule whose fragment only
# arrived now into a file an earlier attempt already filled for others.
if not fill and os.path.exists(merge_progress_file):
with open(merge_progress_file) as f:
merged_set = {ln.strip() for ln in f if ln.strip()}
if merged_set:
logger.info(f"Resuming merge: {len(merged_set)} (cell, year) already merged")
remaining = [yp for yp in touched_year_pfs if yp not in merged_set]
new_schema = None if fill else _var_new_columns_schema(client, gran_h5, product_vars,
source_products=source_products)
merge_fn = functools.partial(_var_merge_cell_year, tmp_dir=tmp_dir, fill=fill, new_schema=new_schema)
updated_files: List[str] = []
touched_cells = set()
n_merge_fail = 0
failed_year_pfs = set()
if remaining:
merge_futures = client.map(merge_fn, remaining, pure=False)
fut_to_pf = dict(zip(merge_futures, remaining))
pbar = tqdm_bar(total=len(remaining), desc="Stage2 cell-year merge", unit="file")
progress_fh = open(merge_progress_file, 'a')
try:
for fut in dask_as_completed(merge_futures):
pf = fut_to_pf[fut]
try:
res = fut.result()
if res is not None:
updated_files.append(res)
touched_cells.add(os.path.dirname(os.path.dirname(res)))
# Record progress whether or not rows landed: an empty
# (no-fragment) cell-year is "done" for resume purposes.
progress_fh.write(pf + '\n')
progress_fh.flush()
except Exception as e:
n_merge_fail += 1
failed_year_pfs.add(pf)
logger.warning(
f"Stage2 merge failed for {os.path.relpath(pf, h3_dir)}: "
f"{type(e).__name__}: {e}"
)
finally:
fut.release()
pbar.update(1)
pbar.set_postfix(updated=len(updated_files), failed=n_merge_fail)
finally:
progress_fh.close()
pbar.close()
if n_merge_fail:
logger.error(f"Stage2: {n_merge_fail} cell-year merge(s) failed. Re-run to retry.")
logger.info(
f"Merged {len(updated_files)} (cell, year) files across "
f"{len(touched_cells)} cells (fragments written={n_frag})"
)
# ── Stage 3: per-cell metadata aggregation + manifest ───────────────
if touched_cells:
logger.info(f"Merging per-cell metadata for {len(touched_cells)} updated cells")
meta_futures = client.map(h3_merge_metadata, sorted(touched_cells), pure=False)
meta_pbar = tqdm_bar(total=len(touched_cells), desc="Merging cell metadata", unit="cell")
try:
for f in dask_as_completed(meta_futures):
try:
f.result()
except Exception as e:
logger.warning(f"Cell metadata merge failed: {e}")
meta_pbar.update(1)
finally:
meta_pbar.close()
generate_manifest(h3_dir, tree_shape='h3db')
# ── Cleanup: drop the leftover sentinel/progress scaffolding ────────
# The per-cell-year fragment dirs were already deleted INSIDE
# _var_merge_cell_year as each merge succeeded (rm_src discipline,
# distributed across workers) — so _var_frags is empty here and its
# rmtree is a cheap no-op rather than a multi-hour serial sweep of
# millions of files on the driver. What remains is the
# _var_fan_complete sentinel dir + the merge-progress file, only
# meaningful for resuming an interrupted run. Drop them on a fully
# clean run; leave them if anything failed so the next resume reuses
# completed work. Fan-out cleanup is parallelized to avoid the same
# serial-rmtree trap (66k+ sentinel files at continental scale).
if updated_files and n_fan_fail == 0 and n_merge_fail == 0:
frag_root = os.path.join(tmp_dir, _VAR_FRAG_DIRNAME)
if os.path.isdir(frag_root):
try:
shutil.rmtree(frag_root) # empty after per-merge rm_src — cheap
except OSError as e:
logger.warning(f"Could not clean up {frag_root}: {e}")
_remove_tree_fanout(os.path.join(tmp_dir, _VAR_FAN_COMPLETE_DIRNAME),
desc="Cleaning fan sentinels")
p = os.path.join(tmp_dir, _VAR_MERGE_PROGRESS_FILENAME)
try:
if os.path.exists(p):
os.unlink(p)
except OSError:
pass
failed_granules |= {g for g, yps in gran_year_pfs.items() if failed_year_pfs.intersection(yps)}
failed_granules.discard(None)
return {'updated_files': updated_files, 'failed_granules': failed_granules,
'failed_products': failed_products, 'fragments': n_frag}
def _build_add_variables(h3_dir, new_product_vars, soc_source=None, version=None,
tmp_dir=None, exclude=None, releases_callback=None):
"""Add new variables to an existing H3 database via inverted fan-out.
**Architecture: read each granule ONCE, fan its shots to all owning
cells, then merge per-(cell, year) into the BASE parquet.** The
legacy per-(cell, year) path re-read every granule h5 once per cell
that listed it — a DB-wide 39.75× redundancy measured on the
production tree (2.64M reads for 66.5k unique granules), and h5 reads
are ~93% of the runtime (cold GPFS chunk I/O). Inverting the
granule→cell relationship (free from the metadata granule lists)
collapses that to one read per granule.
Shots are routed to cells by matching ``shot_number`` against the
existing base parquets — never by recomputing H3 from the new
product's coordinates (coordinate drift between products would
misroute boundary shots and silently NaN the left-join). One
self-contained parquet per partition is preserved (no sidecars).
Stages (all scatter-free — per-task args inlined, constants baked
into a ``functools.partial``, per the fresh-build streaming-driver
contract in ``_write_partitioned_streaming``):
0. Driver enumerates ``(cell, year)`` base parquets (parallel
year-subdir listing), skip-checks each via
``_scan_year_file_for_update`` (drops units whose base already
has all new columns → re-run is a fast no-op), builds the SOC
tree once, and inverts to ``granule → [base year_pf...]``.
1. ``client.map(_var_fan_granule, ...)`` — one task per unique
granule: read its h5(s) once, fan shots to owning cell-years as
tiny fragments under ``tmp_dir/_var_frags/``. Resume via
per-granule ``.done`` sentinels.
2. ``client.map(_var_merge_cell_year, ...)`` — one task per
touched ``(cell, year)``: concat its fragments and join into the
base parquet (rowgroup-wise, atomic). Resume via an append-only
``_var_merge_progress.txt``.
3. Per-cell ``h3_merge_metadata`` + manifest refresh. Fragments
cleaned at the end.
Parameters
----------
h3_dir : str
Root directory of the H3 database.
new_product_vars : dict
Product code → list of new variable names to add.
soc_source : str, list, or None
Source for HDF5 files. ``None`` (S3 ETL mode) is not supported
on this path.
version : int or None
GEDI data version for the SOC tree's filename filter.
tmp_dir : str or None
Scratch root for intermediate fragments + sentinels. Defaults to
``<h3_dir>/.tmp_var_update`` when not supplied by the caller.
exclude : list of str, optional
fnmatch patterns of source files to ignore (as ``gh3_build --exclude``),
so the new variables come from the files the database was built from.
releases_callback : callable, optional
Called once, before the first write, with ``{key: {product: tag}}``
(``gedi_release_tag``) of the files read for each product new to the
database, as :func:`_build_fill_products` does.
Notes
-----
Each granule's latest listed release is read. A product new to the
database (no column of it in the build log's ``h3_columns``) gets
``root_file_<prod>`` and ``root_beam_<prod>`` too, as a build writes them
(#68): a granule replacement matches its rows by them, and the file each
row came from is known. A product an earlier variable add added without
them does not get them (its columns may come from an older release than
the one read now); :func:`_replace_granules_build` refuses to replace it.
A product new to the database also gets its release's quality flags
(``_PRODUCT_QUALITY_FLAGS``), as a build stores them, in every file.
Returns
-------
list of str or None
Base year-file paths updated this run, or None when nothing changed.
"""
logger.info("Variable-only expansion detected: inverted granule fan-out")
from dask.distributed import as_completed as dask_as_completed
from tqdm import tqdm as tqdm_bar
from .parallel import parallel_map
client = get_dask_client()
if client is None:
raise GediError(
"_build_add_variables requires a registered dask Client "
"(wrap your call in `with Client(...): ...`)."
)
if tmp_dir is None:
tmp_dir = os.path.join(h3_dir, '.tmp_var_update')
os.makedirs(tmp_dir, exist_ok=True)
# ── Stage 0a: enumerate (cell, year) base parquets ──────────────────
# Driver scandirs cell dirs; per-cell year-subdir listing is fanned
# across workers (same shape as the fresh-build merge enumeration).
try:
h3_part_dirs = sorted(
e.path for e in os.scandir(h3_dir)
if e.is_dir(follow_symlinks=False) and e.name.startswith('h3_')
)
except OSError:
h3_part_dirs = []
if not h3_part_dirs:
logger.info("No existing partitions to update")
return None
year_lists = client.gather(client.map(_list_year_subdirs, h3_part_dirs, pure=False))
year_files = [
pf
for ydirs in year_lists
for ydir in ydirs
for pf in glob.glob(os.path.join(ydir, '*.parquet'))
]
if not year_files:
logger.info("No (cell, year) parquet files found to update")
return None
logger.info(
f"Found {len(year_files)} (cell, year) parquet files across "
f"{len(h3_part_dirs)} cells"
)
# ── Stage 0b: skip-check + granule lists (parallel) ─────────────────
# Target column set (with product suffix) for the base skip-check.
new_cols = set()
for prod, var_list in new_product_vars.items():
if var_list:
suffix = f"_{prod.lower()}"
new_cols.update(
v if v.endswith(suffix) else f"{v}{suffix}"
for v in var_list if v != 'shot_number'
)
# A product new to the database gets its source columns too (#68), from
# the file its values are read from. Not one it already holds columns of
# without them (added by an earlier variable add): its other columns may
# come from an older release, which a source column naming this one would
# hide from a granule replacement. A file lacking them is not done.
from .utils import json_read_cached
h3_columns = (json_read_cached(os.path.join(h3_dir, BUILD_LOG_FILENAME)) or {}).get('h3_columns') or []
if not h3_columns: # a log too old to list them: one footer answers
h3_columns = list(read_parquet_schema(year_files[0])['column'])
source_products = sorted(p for p, v in new_product_vars.items()
if v and h3_columns and not any(c.endswith(f"_{p.lower()}") for c in h3_columns))
for prod in source_products:
new_cols.update(f"{c}_{prod.lower()}" for c in _SOURCE_COLUMNS)
# ...and its quality flags, as a build stores them (_complete_essentials_and_flags):
# every file gets them in this run. A product the database holds is left as it is:
# updates withhold the flags it lacks (_withhold_flags_absent_from_db).
new_product_vars = {p: (list(v) if v is not None else None) for p, v in new_product_vars.items()}
_complete_essentials_and_flags({p: new_product_vars[p] for p in source_products}, version=version,
ensure_l2a=False)
for prod in source_products:
suffix = f"_{prod.lower()}"
new_cols.update(f"{v}{suffix}" for v in new_product_vars[prod] if v != 'shot_number')
if source_products:
logger.info(f"Variable update: {', '.join(source_products)} rows also get their source file and beam "
f"columns ({', '.join(f'root_file_{p.lower()}' for p in source_products)})")
logger.info("Scanning year files for resume state and granule lists")
scan_futures = client.map(
_scan_year_file_for_update, year_files, new_cols=list(new_cols), pure=False
)
fut_idx = {f: i for i, f in enumerate(scan_futures)}
scan_results: List[Any] = [None] * len(year_files)
scan_pbar = tqdm_bar(total=len(year_files), desc="Scanning year files", unit="file")
try:
for fut, result in dask_as_completed(scan_futures, with_results=True):
scan_results[fut_idx[fut]] = result
scan_pbar.update(1)
finally:
scan_pbar.close()
del scan_futures, fut_idx
# ── Stage 0c: build SOC tree once, invert granule → [base year_pf] ──
# Partial granules too: in a phased database a granule may lack another,
# later product and must still get the new variables of the ones it has.
if isinstance(soc_source, str):
glob_kwargs = {'version': version} if version is not None else None
all_soc = soc_file_tree(soc_source, to_list=False, glob_kwargs=glob_kwargs, require_all=False,
exclude=exclude)
elif isinstance(soc_source, list):
all_soc = soc_file_tree(soc_source, to_list=False, require_all=False, exclude=exclude)
elif soc_source is None:
logger.warning("Variable-only update requires a local SOC source (no S3 ETL support on this path)")
return None
else:
return None
if not all_soc:
logger.warning("SOC tree is empty — nothing to join")
return None
logger.info(f"SOC tree built ({len(all_soc)} orb_tracks)")
gran_year_pfs: Dict[str, List[str]] = {} # granule_key → owning base year_pfs
gran_h5: Dict[str, Dict[str, str]] = {} # granule_key → {prod: h5_path}
skipped = 0
no_soc = 0
for year_pf, granules in zip(year_files, scan_results):
if granules is None:
skipped += 1
continue
for g in granules:
ot = f"O{g['orbit']:05d}_{g['granule']:02d}_T{g['track']:05d}"
soc_files = all_soc.get(ot)
if not soc_files:
no_soc += 1
continue
h5_by_prod = {
p: soc_files[p] for p in new_product_vars
if p in soc_files and new_product_vars[p]
}
if not h5_by_prod:
continue
gran_h5[ot] = h5_by_prod
gran_year_pfs.setdefault(ot, []).append(year_pf)
if not gran_year_pfs:
logger.info(
f"Nothing to update (already-done={skipped}, no-SOC-match={no_soc}). "
f"Variable columns already present in all base files."
)
return None
n_naive = sum(len(v) for v in gran_year_pfs.values())
logger.info(
f"Inverted map: {len(gran_year_pfs)} unique granules → "
f"{n_naive} (cell, year) edges "
f"({n_naive / len(gran_year_pfs):.1f}× read redundancy avoided; "
f"already-done={skipped})"
)
if releases_callback is not None and source_products:
# Before the first write, as a fill records them: rows of these files
# may be committed by a run killed part-way.
tags = {}
for ot, h5 in gran_h5.items():
t = {p: tag for p in source_products
if p in h5 and (tag := gedi_release_tag(str(getattr(h5[p], 'path', h5[p]))))}
if t:
orbit, granule, track = ot.split('_')
tags[(int(orbit[1:]), int(granule), int(track[1:]))] = t
releases_callback(tags)
pending = [pf for pf, granules in zip(year_files, scan_results) if granules is not None]
return _fan_merge_products(h3_dir, gran_h5, gran_year_pfs, new_product_vars, tmp_dir,
year_pfs=pending, source_products=source_products)['updated_files'] or None
def _granule_key_str(key) -> str:
"""``(orbit, granule, track)`` → the ``soc_file_tree`` pivot key."""
return f"O{key[0]:05d}_{key[1]:02d}_T{key[2]:05d}"
def _year_file_target_granules(year_pf, *, targets):
"""Worker: the target granule keys a (cell, year) file's metadata lists.
Reads only the per-year metadata JSON. ``targets`` is a frozenset of
``soc_file_tree`` keys, baked into a ``functools.partial`` by the caller.
A missing file or sidecar (the cell has no data that year) is ``None``.
"""
meta_path = year_pf.replace('.parquet', PARTITION_META_FILENAME)
try:
meta = json_read(meta_path)
except FileNotFoundError:
return None # any other error raises: an unreadable sidecar is not "no data"
hits = [k for k in (_granule_key_str((g['orbit'], g['granule'], g['track']))
for g in (meta or {}).get('granules', [])) if k in targets]
return hits or None
def _product_fill_vars(h3_dir: str, products, soc_files=None, version=None) -> Dict[str, List[str]]:
"""Variables to read per product so a fill writes exactly what a build wrote.
The build log's per-product variable lists, expanded the way Stage 1
expands them (:func:`_expand_product_vars`: presets, wildcards, L2A
essentials, quality flags) with ``soc_files`` as the sample for wildcard
or all-variable lists; without samples, the sample-free part of that
expansion. The loader adds ``root_file`` / ``root_beam`` itself, and
2-D datasets expand to their per-index columns there, as in a build.
"""
import copy
log = json_read(os.path.join(h3_dir, BUILD_LOG_FILENAME)) or {}
recorded = {p: copy.deepcopy(((log.get('products') or {}).get(p) or {}).get('variables')) for p in products}
if version is None:
version = log.get('gedi_version')
if soc_files:
expanded = _expand_product_vars(recorded, soc_files, version=version)
else:
expanded = gedi_vars_expand(recorded, version=version)
_complete_essentials_and_flags(expanded, version=version, ensure_l2a=False)
return {p: [v for v in (expanded.get(p) or []) if v != 'shot_number'] for p in products}
_PRODUCT_FILL_DIRNAME = '_product_fill'
def _build_fill_products(h3_dir, targets, soc_source, version=None, tmp_dir=None, exclude=None,
releases_callback=None):
"""Backfill products published after their granule's rows were indexed.
A granule indexed while some of its products were not yet published (a
``gh3_build --allow-missing-products`` run) stores those products'
columns as null. Once the product files exist, this writes their values
into those rows — only into cells that are null, never overwriting a
value — at a fraction of a rebuild's cost:
* **Targets are a-priori.** The caller passes them from the build log's
per-product ``MISSING_SOURCE`` status; nothing scans the database.
* **Candidate files are a-priori.** One ``{cell}.{year}.0.parquet`` per
partition id in the build log and per year a target granule covers
(from its file name). Workers read only those files' metadata JSON to
find which ones list a target granule.
* **Each product file is read once.** The shared inverted fan-out
(:func:`_fan_merge_products`, as the variable-add path uses) fans a
granule's shots to the files that own them, then fills each touched
file rowgroup-wise and atomically (:func:`parquet_fill_columns`).
Files that hold no target granule are never opened or rewritten.
Parameters
----------
h3_dir : str
Database root (with a build log).
targets : dict
``{(orbit, granule, track): iterable of products}`` to fill.
soc_source : str or list
Local SOC directory, or a listing already made
(``soc_file_tree(..., to_list=True)`` entries) to reuse without a walk.
version : int, optional
GEDI release of the database, pinning the SOC listing.
tmp_dir : str, optional
Scratch root (the build's ``--tmpdir``; default ``<h3_dir>/.tmp``).
Fragments and resume state live in its own ``_product_fill``
subdirectory, apart from a variable update's. Resume state is keyed to
the exact target set (granule → products): a run with different
targets clears that subdirectory first, so a leftover per-granule
sentinel can never stand in for a product it did not fill.
exclude : list of str, optional
fnmatch patterns of source files to ignore (as ``gh3_build --exclude``).
releases_callback : callable, optional
Called once, before the first write, with ``{key: {product: tag}}``
(``gedi_release_tag``) of the files the fill reads, so the caller can
persist them (:meth:`H3BuildLogger.record_releases`) ahead of any row
they produce.
Returns
-------
dict
``{'filled': {key: [products]}, 'unavailable': {key: [products]},
'unlocated': {key: [products]}, 'failed': {key: [products]},
'updated_files': [paths]}`` —
``unavailable``: no product file yet; ``unlocated``: the file exists
but no database file lists the granule (it cannot be filled);
``failed``: a read or merge error, or no row matched (retry). See
:meth:`H3BuildLogger.record_fill_results`.
"""
client = get_dask_client()
if client is None:
raise GediError("_build_fill_products requires a registered dask Client")
targets = {tuple(k): set(v) for k, v in targets.items() if v}
out = {'filled': {}, 'unavailable': {}, 'unlocated': {}, 'failed': {}, 'updated_files': []}
if not targets:
return out
tmp_dir = os.path.join(tmp_dir or os.path.join(h3_dir, '.tmp'), _PRODUCT_FILL_DIRNAME)
# ── Stage 0a: product files now on disk, one listing ──
if isinstance(soc_source, list):
soc_tree = {}
for entry in soc_source:
path = next(iter(entry.values()), None)
key = _granule_id_from_l2a_path(getattr(path, 'path', path)) if path else None
if key is not None:
soc_tree[_granule_key_str(key)] = entry
else:
glob_kwargs = {'version': version} if version is not None else None
# Partial granules too: a target may still lack another product.
soc_tree = soc_file_tree(soc_source, to_list=False, glob_kwargs=glob_kwargs, exclude=exclude,
require_all=False)
gran_h5: Dict[str, Dict[str, str]] = {}
years = set()
for key, prods in targets.items():
entry = soc_tree.get(_granule_key_str(key)) or {}
h5 = {p: entry[p] for p in prods if p in entry}
if len(h5) < len(prods):
out['unavailable'][key] = sorted(prods.difference(h5))
if not h5:
continue
gran_h5[_granule_key_str(key)] = h5
try:
first = next(iter(h5.values()))
stamp = os.path.basename(str(getattr(first, 'path', first))).split('_')[2]
year, doy = int(stamp[:4]), int(stamp[4:7])
except (IndexError, ValueError):
year, doy = None, 0
if year is not None:
years.add(year)
if doy >= 365: # a granule starting on the last day can run past midnight
years.add(year + 1)
if not gran_h5:
logger.info(f"Product backfill: none of {len(targets)} pending granule(s) has its product files yet")
return out
# ── Stage 0b: which (cell, year) files list each granule ──
log = json_read(os.path.join(h3_dir, BUILD_LOG_FILENAME)) or {}
part = log.get('h3_partition_level')
cells = log.get('h3_partition_ids') or []
candidates = [
os.path.join(h3_dir, f"h3_{part:02d}={c}", f"year={y}", f"{c}.{y}.0.parquet")
for c in cells for y in sorted(years)
]
logger.info(
f"Product backfill: {len(gran_h5)} granule(s) have new product files; locating them in "
f"{len(candidates)} candidate (cell, year) file(s), metadata only"
)
scan_fn = functools.partial(_year_file_target_granules, targets=frozenset(gran_h5))
gran_year_pfs: Dict[str, List[str]] = {}
from .parallel import parallel_map
unreadable = 0
for year_pf, hits in parallel_map(candidates, scan_fn, desc="Locating granules", unit="file",
batch_size=500 if len(candidates) > 10000 else 0):
if isinstance(hits, Exception):
unreadable += 1
logger.warning(f"Product backfill: could not read metadata of {year_pf}: {hits}")
continue
if not hits:
continue
for k in hits:
gran_year_pfs.setdefault(k, []).append(year_pf)
if unreadable:
# A granule's rows could sit in an unreadable file: filling the rest
# and reporting it filled would hide those rows. Retry the whole set.
raise GediError(f"Product backfill: {unreadable} metadata file(s) unreadable; nothing filled, retry later")
by_str = {_granule_key_str(k): k for k in targets}
for ot in list(gran_h5):
if ot not in gran_year_pfs:
out['unlocated'][by_str[ot]] = sorted(gran_h5.pop(ot))
if not gran_h5:
logger.warning("Product backfill: no database file lists the pending granule(s); nothing to fill")
return out
# ── Resume state belongs to one target set ──
signature = json.dumps(sorted((ot, sorted(h5)) for ot, h5 in gran_h5.items()))
sig_path = os.path.join(tmp_dir, '_fill_targets.json')
try:
with open(sig_path) as f:
previous = f.read()
except OSError:
previous = None
if previous != signature and os.path.isdir(tmp_dir):
logger.info(f"Product backfill: discarding resume state of a different target set in {tmp_dir}")
_remove_tree_fanout(tmp_dir, desc="Cleaning stale fill state")
if previous != signature:
os.makedirs(tmp_dir, exist_ok=True)
with AtomicFileWriter(sig_path) as tmp_path:
with open(tmp_path, 'w') as f:
f.write(signature)
# ── Stages 1-3: read each product file once, fill the files that own the rows ──
if releases_callback is not None:
# Before the first write: a fill killed or failed part-way leaves
# committed rows from these files (#47).
tags = {by_str[ot]: {p: t for p, f in h5.items() if (t := gedi_release_tag(str(getattr(f, 'path', f))))}
for ot, h5 in gran_h5.items()}
releases_callback({k: v for k, v in tags.items() if v})
products = sorted({p for h5 in gran_h5.values() for p in h5})
fill_vars = _product_fill_vars(h3_dir, products, soc_files=list(gran_h5.values()), version=version)
res = _fan_merge_products(h3_dir, gran_h5, gran_year_pfs, fill_vars, tmp_dir, fill=True)
for ot, h5 in gran_h5.items():
failed = set(h5) if ot in res['failed_granules'] else res['failed_products'].get(ot, set()) & set(h5)
for bucket, prods in (('failed', failed), ('filled', set(h5) - failed)):
if prods:
out[bucket][by_str[ot]] = sorted(prods)
out['updated_files'] = res['updated_files']
logger.info(
f"Product backfill: filled {len(out['filled'])} granule(s) across {len(res['updated_files'])} file(s); "
f"{len(out['failed'])} failed (retried next run), {len(out['unavailable'])} still awaiting products"
+ (f", {len(out['unlocated'])} not found in any database file" if out['unlocated'] else "")
)
return out
def _resolve_build_version(h3_dir: str, soc_source=None) -> int:
"""GEDI release for a build called with ``version=None``.
Resolution order, all a-priori (no HDF5 open): the existing database's
build log ``gedi_version`` (an update must stay on its release) >
:func:`gedih3.logger.resolve_soc_version` for a local SOC tree, or the
``_V00N`` shared by every file of a pre-acquired list >
``GEDI_DEFAULT_VERSION``. Never a random file sample: a tree that holds
two releases side by side would make that a coin flip.
Parameters
----------
h3_dir : str
Target database root; its build log, when present, is authoritative.
soc_source : str, list, or None
Local SOC directory, pre-acquired file list, or ``None`` (S3).
Returns
-------
int
The resolved release number.
"""
from .config import GEDI_DEFAULT_VERSION
from .logger import resolve_soc_version
log_path = os.path.join(h3_dir, BUILD_LOG_FILENAME)
if os.path.isfile(log_path):
v = json_read(log_path).get('gedi_version')
if v is not None:
return int(v)
if isinstance(soc_source, str):
v = resolve_soc_version(soc_source)
if v is not None:
return int(v)
elif isinstance(soc_source, list):
# Filenames only (no stat): a single-release list pins itself; a mixed
# one falls through, and the version filter reports what it drops.
releases = {m.group(1) for p in soc_source
if (m := re.search(r'_V(\d{3})[^/]*\.h5$', str(getattr(p, 'path', p))))}
if len(releases) == 1:
return int(releases.pop())
return GEDI_DEFAULT_VERSION
# Scratch of a granule replacement (#47) under the build's tmp dir, and its
# journal: the target set, written before the first destructive step. The
# marker is written before the first rewrite of a database file.
_REPLACE_DIRNAME = '_replace'
_REPLACE_JOURNAL = '_replace_targets.json'
_REPLACE_MERGING = '_merging'
def _year_file_replace_targets(year_pf, *, targets, products=()):
"""Worker: :func:`_year_file_target_granules`, for a replacement.
A year file whose metadata is missing cannot be ruled out as holding a
target's rows (a lost sidecar, or a file a killed replacement emptied
between its two unlinks), so it is rewritten too rather than read as "no
data that year": the merge decides from the rows' own source columns.
Returns ``(hits, unsourced)``: ``unsourced``, the ``products`` with no
``root_file_<prod>`` column in a file that holds a target (its footer,
the file's own truth: files of one database can differ, #68).
"""
if os.path.exists(year_pf) and not os.path.exists(year_pf.replace('.parquet', PARTITION_META_FILENAME)):
hits = ['(no metadata)']
else:
hits = _year_file_target_granules(year_pf, targets=targets)
if not hits or not products:
return hits, []
names = set(read_parquet_schema(year_pf)['column'])
return hits, sorted(p for p in products if f"root_file_{p.lower()}" not in names)
_ROOT_FILE_KEY_RE = re.compile(r'_O(\d+)_(\d{2})_T(\d+)_')
def _year_file_source_releases(year_pf, *, products, targets, version=None):
"""Worker: the source releases a (cell, year) file's rows name, per target granule.
Skips a file whose metadata lists none of ``targets`` (``soc_file_tree``
keys); one without metadata is read, and ``targets=None`` reads every file
for every granule. Otherwise reads only the ``root_file_<prod>`` columns,
dictionary-encoded and in batches (a few hundred distinct names in the
largest production file: under a second, ~150 MB), and parses each name:
the granule key and the release tag (``gedi_release_tag``) are in it.
Only names of GEDI ``version`` count. Returns ``{granule key: {product:
[tags]}}`` for the targets; a product with more than one tag for a granule
(a replacement abandoned half way) is left to the caller.
"""
import pyarrow.compute as pc
if (targets is not None and os.path.exists(year_pf.replace('.parquet', PARTITION_META_FILENAME))
and not _year_file_target_granules(year_pf, targets=targets)): # no sidecar: read it anyway
return {}
pf = pq.ParquetFile(year_pf)
cols = {f"root_file_{p.lower()}": p for p in products}
present = [c for c in cols if c in pf.schema_arrow.names]
if not present:
return {}
pf = pq.ParquetFile(year_pf, read_dictionary=present)
names = {c: set() for c in present}
for batch in pf.iter_batches(batch_size=1_000_000, columns=present):
for c in present:
names[c].update(pc.unique(batch.column(c)).to_pylist())
suffix = f"_V{int(version):03d}" if version else None
out: Dict[str, Dict[str, set]] = {}
for c, values in names.items():
for name in values:
base = os.path.basename(str(name)) if name else ''
m = _ROOT_FILE_KEY_RE.search(base)
tag = gedi_release_tag(base) if m and (suffix is None or suffix in base) else None
key = f"O{int(m.group(1)):05d}_{int(m.group(2)):02d}_T{int(m.group(3)):05d}" if tag else None
if key is not None and (targets is None or key in targets):
out.setdefault(key, {}).setdefault(cols[c], set()).add(tag)
return {k: {p: sorted(t) for p, t in v.items()} for k, v in out.items()}
[docs]
def releases_from_rows(h3_dir: str, products, keys, version: Optional[int] = None):
"""The source release of granule products, read from the database's rows (#67).
A database built before gedih3 0.20.10 records no release per granule,
so re-released granules cannot be found from the build log. Every row
names its source file per product (``root_file_<prod>``), and the file
name holds the granule key and the release (PPDS, PGE, generation,
suffix): one worker task per (cell, year) file reads its metadata and,
if it lists one of ``keys``, those columns only
(:func:`_year_file_source_releases`). The files come from the
database's ``_manifest.txt`` (:func:`smart_glob`).
Parameters
----------
h3_dir : str
Database root.
products : iterable of str
Products to read, e.g. ``['L2A', 'L4A']``; one without a source
column is skipped.
keys : iterable of tuple
``(orbit, granule, track)`` of the granules to read.
version : int, optional
The database's GEDI version: names of another one are ignored.
Returns
-------
tuple
``(releases, mixed, unread, failed)``: ``{key: {product: tag}}`` for
products whose rows name one release; ``{key: {product: [tags]}}``
for those naming several (a replacement left half done: never
guessed); the keys a file that could not be read may hold, left out
of both (their rows were not all read: all of ``keys`` when that
cannot be told from the file's metadata); and the files that could
not be read (skipped and logged; ``None`` for a batch of them).
Requires a registered dask Client.
"""
from .parallel import parallel_map
from .utils import smart_glob
keys = set(keys)
targets = frozenset(_granule_key_str(k) for k in keys)
if not targets:
return {}, {}, set(), []
files = sorted(smart_glob(os.path.join(h3_dir, 'h3_*', 'year=*', '*.parquet')))
found: Dict[str, Dict[str, set]] = {}
failed = []
# A large target set would travel inside every task: read every granule
# then (a database recording none lists them all anyway), filter here.
shipped = targets if len(targets) <= _ROW_RELEASE_SHIPPED_TARGETS else None
fn = functools.partial(_year_file_source_releases, products=sorted(products), targets=shipped, version=version)
for year_pf, res in parallel_map(files, fn, desc="Reading source releases", unit="file",
batch_size=500 if len(files) > 10000 else 0):
if isinstance(res, Exception):
failed.append(year_pf)
logger.warning(f"Reading source releases: skipped {year_pf or 'a batch of files'} "
f"({type(res).__name__}: {res})")
continue
for key, prods in (res or {}).items():
if key in targets:
for p, tags in prods.items():
found.setdefault(key, {}).setdefault(p, set()).update(tags)
unread = _row_release_unread(failed, targets, keys)
releases, mixed = {}, {}
for key, prods in found.items():
k = (int(key[1:6]), int(key[7:9]), int(key[11:]))
if k in unread:
continue
for p, tags in prods.items():
if len(tags) == 1:
releases.setdefault(k, {})[p] = next(iter(tags))
else:
mixed.setdefault(k, {})[p] = sorted(tags, key=gedi_release_rank)
return releases, mixed, unread, failed
_ROW_RELEASE_SHIPPED_TARGETS = 2000
_ROW_RELEASE_MAX_FAILED_SIDECARS = 100
def _row_release_unread(failed, targets, keys):
"""The keys whose rows ``releases_from_rows`` could not all read.
The granules each unreadable file's metadata lists (read here, on the
driver: a few small sidecars); a file gone since the listing holds none.
Every key when a file has no readable
metadata, a whole batch failed (``None``), or too many files failed to
look each one up.
"""
if not failed:
return set()
if None in failed or len(failed) > _ROW_RELEASE_MAX_FAILED_SIDECARS:
return set(keys)
unread = set()
for year_pf in failed:
if not os.path.exists(year_pf):
continue # gone since the listing (a stale manifest): no rows to read
try:
hits = _year_file_target_granules(year_pf, targets=targets)
if hits is None and not os.path.exists(year_pf.replace('.parquet', PARTITION_META_FILENAME)):
return set(keys) # no metadata: the file may hold any of them
except Exception: # noqa: BLE001 - unreadable metadata: any of them
return set(keys)
unread.update((int(k[1:6]), int(k[7:9]), int(k[11:])) for k in hits or ())
return unread
[docs]
def replacement_abandon_hint(tmp_dir: str) -> str:
"""How to abandon an unfinished replacement, and what that leaves."""
rtmp = os.path.join(tmp_dir, _REPLACE_DIRNAME)
if os.path.exists(os.path.join(rtmp, _REPLACE_MERGING)):
return (f"removing {rtmp} abandons it, but some database files were already rewritten: those granules "
f"stay partly re-indexed until they are re-indexed again")
return f"removing {rtmp} abandons it safely: no database file was rewritten yet"
[docs]
def blind_replacement_targets(h3_dir: str, targets,
unsourced_by_key: Optional[Dict[Tuple[int, int, int], set]] = None
) -> Dict[Tuple[int, int, int], List[str]]:
"""Granules a replacement cannot tell old rows from new ones for (#68).
A replacement merge keeps a row when every ``root_file_<prod>`` column
names the given file. A product added by a variable update of an older
gedih3 has no such column, so when it is the only product of a granule
with a new release, the merge would keep the old rows while the build log
records the new release. A granule with another product re-released is
fine: its rows are replaced whole. One with no recorded release (indexed
before gedih3 recorded them) cannot be shown to be fine, so it counts.
Answered from the build log (``h3_columns``, the union of every file's
columns, and the recorded releases), and from ``unsourced_by_key`` for
files that differ.
Parameters
----------
h3_dir : str
Database root.
targets : dict
``{(orbit, granule, track): {product: file}}``, as passed to
:func:`_replace_granules_build`.
unsourced_by_key : dict, optional
``{(orbit, granule, track): products}`` lacking a source column in a
file holding that granule (a database expanded after a product was
added without one has it in some files only), from the files'
footers (:func:`_year_file_replace_targets`).
Returns
-------
dict
``{(orbit, granule, track): [product, ...]}``, the unsourced products
of each such granule.
"""
from .utils import json_read_cached
log = json_read_cached(os.path.join(h3_dir, BUILD_LOG_FILENAME)) or {}
h3_columns = set(log.get('h3_columns') or [])
by_key = {tuple(k): set(v) for k, v in (unsourced_by_key or {}).items() if v}
if h3_columns:
everywhere = {p for files in targets.values() for p in files if f"root_file_{p.lower()}" not in h3_columns}
if everywhere:
for key in targets:
by_key.setdefault(tuple(key), set()).update(everywhere)
if not by_key:
return {}
recorded = {(g['orbit'], g['granule'], g['track']): g.get('releases') or {} for g in log.get('granules') or []}
blind = {}
for key, files in targets.items():
unsourced = by_key.get(tuple(key), set()) & set(files)
if not unsourced:
continue
rel = recorded.get(tuple(key), {})
# Unknown (no recorded release) counts as unchanged: nothing proves the rows go.
if not any(rel.get(p) and rel[p] != gedi_release_tag(os.path.basename(str(getattr(files[p], 'path', files[p]))))
for p in files if p not in unsourced):
blind[tuple(key)] = sorted(unsourced)
return blind
[docs]
def blind_replacement_message(blind, tmp_dir: Optional[str] = None) -> str:
"""Why the granules of :func:`blind_replacement_targets` are not replaced."""
keys = sorted(_granule_key_str(k) for k in blind)
products = sorted({p for prods in blind.values() for p in prods})
msg = (f"{len(keys)} granule(s) ({', '.join(keys[:5])}{' ...' if len(keys) > 5 else ''}) cannot be replaced: "
f"their rows have no {', '.join(f'root_file_{p.lower()}' for p in products)} column ({', '.join(products)} "
f"added by a variable update of an older gedih3), and no other product of theirs is known to have a new "
f"release (none changed, or none recorded), so their old rows cannot be told from the new ones. Only a "
f"fresh build records the source of {', '.join(products)}")
if tmp_dir and os.path.exists(replacement_journal(tmp_dir)):
msg += f". An unfinished replacement is journaled: {replacement_abandon_hint(tmp_dir)}"
return msg
[docs]
def check_replacement_sources(h3_dir: str, targets, tmp_dir: Optional[str] = None,
unsourced_by_key: Optional[Dict[Tuple[int, int, int], set]] = None) -> None:
"""Refuse to replace a granule whose changed rows cannot be told from its old ones.
See :func:`blind_replacement_targets`; the parameters are the same, and
``tmp_dir`` (the build tmp dir) points the message at an unfinished
replacement's journal.
Raises
------
GediValidationError
If such a granule is among ``targets``. Nothing was changed.
"""
blind = blind_replacement_targets(h3_dir, targets, unsourced_by_key)
if blind:
raise GediValidationError("Cannot replace granules: " + blind_replacement_message(blind, tmp_dir)
+ ". This run changed nothing; leave them out")
[docs]
def replacement_journal(tmp_dir: str) -> str:
"""Path of the journal of a granule replacement under a build tmp dir.
Its presence means a replacement started there and has not finished:
``gh3_build`` resumes it before anything else.
"""
return os.path.join(tmp_dir, _REPLACE_DIRNAME, _REPLACE_JOURNAL)
[docs]
def read_replacement_journal(tmp_dir: str) -> Optional[Dict[Tuple[int, int, int], Dict[str, str]]]:
"""The target set of an unfinished replacement, or ``None`` when there is none.
``{(orbit, granule, track): {product: basename}}``, as passed to
:func:`_replace_granules_build`.
"""
try:
return _read_replacement_targets(tmp_dir)[1]
except FileNotFoundError:
return None
# Per worker process: the target map of the replacement read last, keyed by
# the journal's path and content digest (#47). One map at most is kept.
_REPLACE_MAP_CACHE: Dict[Tuple[str, str], Dict[Tuple[int, int, int], Dict[str, str]]] = {}
def _read_replacement_targets(replacement_tmp: str) -> Tuple[str, Dict[Tuple[int, int, int], Dict[str, str]]]:
"""``(digest, targets)`` of the replacement journaled under ``replacement_tmp``.
The digest is a SHA-256 of the journal's canonical JSON. Raises
``FileNotFoundError`` when there is no journal.
"""
with open(replacement_journal(replacement_tmp)) as f:
entries = json.load(f)
digest = _targets_digest(entries)
targets = {}
for ot, files in entries:
orbit, granule, track = ot.split('_')
targets[(int(orbit[1:]), int(granule), int(track[1:]))] = dict(files)
return digest, targets
def _targets_digest(entries) -> str:
"""SHA-256 of a replacement's journal entries, as written or as read back."""
import hashlib
return hashlib.sha256(json.dumps(entries).encode()).hexdigest()
def _replacement_digest(replacement_tmp: str) -> str:
"""Content digest of a replacement's journal (see :func:`_read_replacement_targets`)."""
try:
return _read_replacement_targets(replacement_tmp)[0]
except FileNotFoundError:
raise GediMergeError(f"Granule replacement: no journal under {replacement_tmp}; nothing was merged")
def _h3_merge_replacing(in_dir, out_dir, replacement_tmp, digest, **kwargs):
""":func:`h3_merge_files` replacing the granules journaled under ``replacement_tmp``.
The merge task of a granule replacement: the target map is read from the
journal once per worker process, rather than serialized into every task.
It is cached by the content ``digest`` the driver read, never by file
stamps, and a journal whose content differs (or is gone) fails the merge:
a task never merges with another replacement's map.
"""
key = (replacement_journal(replacement_tmp), digest)
rmap = _REPLACE_MAP_CACHE.get(key)
if rmap is None:
try:
found, rmap = _read_replacement_targets(replacement_tmp)
except FileNotFoundError:
found = None
if found != digest:
raise GediMergeError(f"Granule replacement: the journal under {replacement_tmp} changed or is gone "
f"during the merge; nothing was merged for {in_dir}")
_REPLACE_MAP_CACHE.clear()
_REPLACE_MAP_CACHE[key] = rmap
return h3_merge_files(in_dir, out_dir, replace_granules=rmap, **kwargs)
def _replace_granules_build(
h3_dir: str,
targets: Dict[Tuple[int, int, int], Dict[str, str]],
product_vars: Dict[str, List[str]],
res: int,
part: int,
spatial=None,
version: Optional[int] = None,
tmp_dir: Optional[str] = None,
skip_blind: bool = False,
) -> Dict[str, Any]:
"""Re-index granules from new source files, replacing their rows (#47).
For granules NASA re-released within the database's GEDI version:
1. **Old rows.** Every ``(cell, year)`` file whose metadata lists a
target is found from the metadata alone (every partition of the build
log, so it does not depend on a region, and needs no ring-1 rule).
A target whose rows cannot be told from its new ones is refused here
(:func:`check_replacement_sources`, with those files' footers).
2. **Journal.** The target set is written to
:func:`replacement_journal` before anything destructive; a run with a
different target set while one is unfinished is refused first.
3. **New rows.** Stage 1 reads the new files in full, under the
database's region, with the skip check off, into
``<tmp_dir>/_replace/partitions`` (its own sentinels, resumable).
4. **One rewrite per cell-year.** Only once Stage 1 has read every target
granule in full, each affected file is merged with
``h3_merge_files(replace_granules=)``: the old rows go, the new rows
take their place, a file left empty is removed. Each rewrite is
atomic and gives the same result when re-run, so a kill at any point
resumes from the journal.
The caller updates the build log (releases, statuses, partitions) and
removes the scratch directory once the log is saved.
Parameters
----------
h3_dir : str
Database root.
targets : dict
``{(orbit, granule, track): {product: path}}``: the new source file
of every product the database holds, for each granule to replace.
product_vars : dict
The database's products and variables (the build log's).
res, part : int
H3 resolution and partition level of the database.
spatial : optional
The database's region; ``None`` for the whole globe.
version : int, optional
The database's GEDI version.
tmp_dir : str, optional
The build's tmp dir (default ``<h3_dir>/.tmp``).
skip_blind : bool, default False
Leave out, rather than refuse, the targets whose old rows cannot be
told from their new ones (:func:`blind_replacement_targets`): for
targets chosen by ``auto``, not named ones.
Returns
-------
dict
``{'replaced': set of keys, 'empty': set of keys, 'files': list,
'blind': dict}``: ``empty``, the granules with no row in the region any
more (now ``NO_DATA``); ``files``, the database files written;
``blind``, ``{key: products}`` left out (``skip_blind``), or replaced
by a resumed run that had already started merging (an older gedih3
journaled it): their rows of those products keep their old values,
so the caller keeps those products' recorded releases.
Raises
------
GediValidationError
Another replacement is unfinished under ``tmp_dir``, or a target's
rows cannot be told from its new ones (:func:`check_replacement_sources`).
GediDatabaseError
A year file's metadata is unreadable.
GediProcessingError
Stage 1 did not read every target in full.
GediMergeError
A merge failed. In every case nothing is lost: re-run to resume.
"""
if get_dask_client() is None:
raise GediError("_replace_granules_build requires a registered dask Client")
rtmp = os.path.join(tmp_dir or os.path.join(h3_dir, '.tmp'), _REPLACE_DIRNAME)
parquet_dir = os.path.join(rtmp, 'partitions')
rmap = {tuple(k): {p: os.path.basename(str(getattr(f, 'path', f))) for p, f in files.items()}
for k, files in targets.items()}
entries = sorted([_granule_key_str(k), sorted(v.items())] for k, v in rmap.items())
journal = os.path.join(rtmp, _REPLACE_JOURNAL)
try:
with open(journal) as f:
previous = json.load(f)
except FileNotFoundError:
previous = None
if previous is not None and previous != json.loads(json.dumps(entries)):
raise GediValidationError(
f"Another granule replacement is unfinished in {rtmp}: gh3_build resumes it; finish it first, "
f"or remove that directory to abandon it")
# ── 1. The (cell, year) files holding their old rows: metadata only ──
log = json_read(os.path.join(h3_dir, BUILD_LOG_FILENAME)) or {}
# The log's partitions, and any partition directory on disk it misses: one
# readdir of the root, so a stale list cannot hide old rows.
cells = set(log.get('h3_partition_ids') or [])
prefix = f"h3_{part:02d}="
with os.scandir(h3_dir) as it:
cells.update(e.name[len(prefix):] for e in it if e.name.startswith(prefix) and e.is_dir())
years = set()
for files in targets.values():
try:
stamp = os.path.basename(str(next(iter(files.values())))).split('_')[2]
year, doy = int(stamp[:4]), int(stamp[4:7])
except (IndexError, ValueError, StopIteration):
raise GediValidationError(f"Cannot date the source files {files}")
years.add(year)
if doy >= 365: # a granule starting on the last day can run past midnight
years.add(year + 1)
candidates = [os.path.join(h3_dir, f"{prefix}{c}", f"year={y}", f"{c}.{y}.0.parquet")
for c in sorted(cells) for y in sorted(years)]
logger.info(f"Granule replacement: locating {len(rmap)} granule(s) in {len(candidates)} candidate "
f"(cell, year) file(s), metadata only")
from .parallel import parallel_map
scan_fn = functools.partial(_year_file_replace_targets, targets=frozenset(e[0] for e in entries),
products=sorted({p for files in rmap.values() for p in files}))
old_dirs = {} # tmp dir -> target keys its file holds (None: any)
by_str = {_granule_key_str(k): k for k in rmap}
unsourced_by_key: Dict[Tuple[int, int, int], set] = {}
for year_pf, found in parallel_map(candidates, scan_fn, desc="Locating granules", unit="file",
batch_size=500 if len(candidates) > 10000 else 0):
if isinstance(found, Exception):
raise GediDatabaseError(f"Granule replacement: cannot read the metadata of {year_pf} ({found}); "
f"the replacement did not go on")
hits, unsourced = found
for k in (rmap if hits == ['(no metadata)'] else [by_str[h] for h in hits or []]) if unsourced else ():
unsourced_by_key.setdefault(k, set()).update(unsourced)
if hits:
year_dir = os.path.dirname(year_pf)
old_dirs[os.path.join(parquet_dir, os.path.basename(os.path.dirname(year_dir)),
os.path.basename(year_dir))] = (
None if hits == ['(no metadata)'] else {by_str[h] for h in hits})
# A product whose rows record no source file, in any file of a granule (#68).
blind = blind_replacement_targets(h3_dir, rmap, unsourced_by_key)
if blind and previous is not None and os.path.exists(os.path.join(rtmp, _REPLACE_MERGING)):
# Journaled by an older gedih3 and already rewriting files: finish it,
# as that version would have; refusing would strand it half done.
logger.warning(f"Granule replacement: finishing the run started earlier, but "
f"{blind_replacement_message(blind)}; their rows of those products keep their old values")
elif blind and skip_blind and previous is None: # a journal already written holds them: refuse below
logger.warning(f"Granule replacement: leaving out {blind_replacement_message(blind)}")
for key in blind:
rmap.pop(key)
targets = {k: v for k, v in targets.items() if tuple(k) != key}
entries = sorted([_granule_key_str(k), sorted(v.items())] for k, v in rmap.items())
old_dirs = {d: keys for d, keys in old_dirs.items() if keys is None or keys & set(rmap)}
if not rmap:
return {'replaced': set(), 'empty': set(), 'files': [], 'blind': blind}
elif blind:
check_replacement_sources(h3_dir, rmap, tmp_dir=os.path.dirname(rtmp), unsourced_by_key=unsourced_by_key)
old_dirs = list(old_dirs)
# ── 2. Journal, before the first destructive step ──
if previous is None:
os.makedirs(rtmp, exist_ok=True)
with AtomicFileWriter(journal) as tmp_path:
with open(tmp_path, 'w') as f:
json.dump(entries, f)
from .gh3driver import invalidate_bbox_index
invalidate_bbox_index(h3_dir)
# ── 3. Stage 1: the new files, in full, without the skip check ──
product_vars = _expand_product_vars({p: list(v) for p, v in product_vars.items()}, list(targets.values()),
version=version)
_withhold_flags_absent_from_db(product_vars, h3_dir, version=version)
soc_files = [{p: f for p, f in files.items() if p in product_vars} for files in targets.values()]
_pc = preclean_merge_failures(parquet_dir)
if _pc['parquets_removed']:
logger.info(f"Granule replacement: reopened {_pc['parquets_removed']} corrupt fragment task(s)")
ddf, lat_col, lon_col, dat_col, _ = _create_h3_dataframe(soc_files, product_vars, res, part)
ddf = _apply_spatial_filter(ddf, spatial, part, h3_dir)
outcome = {}
_write_partitioned_streaming(
ddf, soc_files, product_vars, res, part, parquet_dir, h3_dir, spatial, lat_col, lon_col, dat_col,
stage1_outcome_callback=lambda incomplete, empty: outcome.update(incomplete=set(incomplete),
empty=set(empty)),
scope_fingerprint=_stage1_scope_fingerprint(spatial, res, part, product_vars),
skip_check=False,
)
if 'incomplete' not in outcome or outcome['incomplete']:
raise GediProcessingError(
f"Granule replacement: Stage 1 did not read every granule in full "
f"({len(outcome.get('incomplete') or [])} unfinished), so no file was rewritten in this run. Re-run "
f"gh3_build to resume it; if it keeps failing, {replacement_abandon_hint(os.path.dirname(rtmp))}")
# ── 4. One atomic rewrite per affected cell-year ──
with open(os.path.join(rtmp, _REPLACE_MERGING), 'w'):
pass
files = _merge_and_finalize(parquet_dir, h3_dir, replacement_tmp=os.path.dirname(rtmp), extra_dirs=old_dirs,
replacement_digest=_targets_digest(entries))
failed = _scan_merge_failure_sentinels(parquet_dir)
if failed:
raise GediMergeError(f"Granule replacement: {len(failed)} partition merge(s) failed; re-run gh3_build to "
f"finish the replacement")
return {'replaced': set(rmap), 'empty': outcome['empty'] & set(rmap), 'files': files, 'blind': blind}
[docs]
def build_h3db(
product_vars: Dict[str, List[str]],
res: int = 12,
part: int = 3,
spatial=None,
temporal=None,
soc_source: Union[str, List, None] = None,
version: Optional[int] = None,
tmp_dir: str = os.path.join(GH3_DEFAULT_TMP_DIR, 'gh3_build'),
h3_dir: str = GH3_DEFAULT_H3_DIR,
skip_granules: Optional[List[Dict]] = None,
status_callback: Optional[Callable[[str], None]] = None,
variable_only_update: bool = False,
exclude: Optional[List[str]] = None,
allow_missing_products: bool = False,
granule_products_callback: Optional[Callable[[Dict[Tuple[int, int, int], List[str]]], None]] = None,
stage1_outcome_callback: Optional[Callable[[set, set], None]] = None,
only_granules: Optional[List[Dict]] = None,
releases_callback: Optional[Callable[[Dict[Tuple[int, int, int], Dict[str, str]]], None]] = None,
) -> Optional[List[str]]:
"""
Build an H3-indexed GEDI database from local SOC files or S3 download.
Parameters
----------
product_vars : dict
Dictionary mapping GEDI product codes (e.g., 'L2A', 'L4A') to lists of
variable names to extract. Use 'default', 'minimal', or explicit lists.
res : int, default 12
H3 resolution level for indexing individual shots (0-15).
part : int, default 3
H3 resolution level for file partitioning (0-15, must be <= res).
spatial : GeoDataFrame, list, or str, optional
Spatial filter for the region of interest.
temporal : tuple, optional
Temporal filter as (start_date, end_date) in 'YYYY-MM-DD' format.
soc_source : str, list, or None
- ``None``: download from NASA S3 to a temp directory (cleaned up after build)
- ``str``: path to local directory containing GEDI SOC HDF5 files
- ``list``: pre-acquired list of file paths or EarthAccessFile objects
version : int or None
GEDI data release for every product. ``None`` resolves via
:func:`_resolve_build_version` (build log > SOC tree > package
default). Every SOC listing is filtered to this release, so the
Stage 1 schema is derived from files of exactly that release.
tmp_dir : str
Path to temporary directory for intermediate files.
h3_dir : str
Output path for the H3-indexed parquet database.
skip_granules : list of dict, optional
List of granule identifiers to skip (from previous builds).
status_callback : callable, optional
Called with status string at pipeline stage transitions.
variable_only_update : bool, default False
If True, only add new variable columns to existing partition files
via shot_number join (``_build_add_variables``). Skips full pipeline.
The caller (CLI) determines this from the build logger state.
exclude : list of str, optional
fnmatch patterns of source file basenames to ignore (e.g.
``['*_SGS.h5']``), in every listing. The patterns kept in an existing
database's build log (``gh3_build --exclude``) are always added.
allow_missing_products : bool, default False
Admit granules that have L2A but lack later products (L2B, L4A…
published after L2A). Their rows store those products' columns as
null, in the database's types, until :func:`_build_fill_products`
backfills them. Requires the streaming writer. Default: a granule
must have every requested product.
granule_products_callback : callable, optional
Called once, before any fragment is written, with
``{(orbit, granule, track): [missing products]}`` for every granule
this Stage 1 processes (an empty list when it has every product).
The CLI records ``MISSING_SOURCE`` from it, so the backfill knows its
targets without scanning the database, even after a crash mid-write.
stage1_outcome_callback : callable, optional
Called as ``(incomplete_keys, empty_keys)`` once Stage 1 drained
every task (see :func:`_write_partitioned_streaming`), before the
merge. Never called by the legacy writer (``GH3_WRITE_STREAMING``
off, deprecated), by an interrupted drain, or when no granule is left
to process.
only_granules : list of dict, optional
If given, Stage 1 reads only these granules and leaves every other one
for a later build (``H3BuildLogger.stage1_only_granules``: a
spatial-only expansion reads only the added region).
releases_callback : callable, optional
``variable_only_update`` only: called before the first write with the
source releases (``{(orbit, granule, track): {product: tag}}``) read
for each product new to the database, which gets source columns too
(see :func:`_build_add_variables`), so the caller can record them
(:meth:`H3BuildLogger.record_releases`).
Returns
-------
list of str or None
List of output parquet file paths, or None if no new data processed.
Raises
------
H3ValidationError
If H3 resolution or partition parameters are invalid
GediFileError
If local source directory doesn't exist or S3 returns no files
"""
# Validate parameters
logger.debug("Validating build parameters")
res, part = validate_h3_params(res, part)
if version is None:
version = _resolve_build_version(h3_dir, soc_source)
# The database's own exclusions (gh3_build --exclude, kept in its log) always apply.
_log_path = os.path.join(h3_dir, BUILD_LOG_FILENAME)
from .logger import stored_exclude
from .utils import json_read_cached
_stored = stored_exclude(json_read_cached(_log_path)) if os.path.isfile(_log_path) else []
exclude = list(dict.fromkeys([*_stored, *(exclude or [])])) or None
# Track temp directory for S3 download cleanup
_s3_tmp_dir = None
# Determine source mode and acquire SOC files
if soc_source is None:
# S3 ETL mode: open remote files via earthaccess, extract only selected
# variables via h5py range requests, write compact local HDF5 files.
# Transfers 10-50x less data than full download depending on variable selection.
_s3_tmp_dir = os.path.join(tmp_dir, '_s3_download')
os.makedirs(_s3_tmp_dir, exist_ok=True)
if variable_only_update:
# Download only new products (no L2A essentials), then join to existing partitions
logger.info(f"S3 variable-only update: adding {set(product_vars.keys())} to existing database")
try:
if status_callback:
status_callback('PROCESSING')
s3_etl_subset(
product_vars=product_vars.copy(),
spatial=spatial,
temporal=temporal,
version=version,
odir=_s3_tmp_dir,
ensure_l2a=False,
)
result = _build_add_variables(
h3_dir, product_vars,
soc_source=_s3_tmp_dir, version=version, tmp_dir=tmp_dir, exclude=exclude,
releases_callback=releases_callback,
)
if not result:
logger.warning(
"Variable update completed but no partition files were modified. "
"New columns may not have been added to the database."
)
return result
finally:
if os.path.exists(_s3_tmp_dir):
shutil.rmtree(_s3_tmp_dir, ignore_errors=True)
# Fresh build or mixed update: full pipeline with L2A
logger.info(f"ETL subset from S3 to temp directory: {_s3_tmp_dir}")
s3_etl_subset(
product_vars=product_vars.copy(),
spatial=spatial,
temporal=temporal,
version=version,
odir=_s3_tmp_dir,
)
soc_source = _s3_tmp_dir
logger.info("S3 ETL subset complete, building from local files")
if isinstance(soc_source, list):
# Pre-acquired file list (local or EarthAccessFile)
logger.info(f"Using {len(soc_source)} pre-acquired source files")
before = len(soc_source)
def _matches_version(p):
try:
pp = p.path if isinstance(p, EarthAccessFile) else p
return GEDIFile(pp).version == version
except Exception:
return False
soc_source = [p for p in soc_source if _matches_version(p)]
dropped = before - len(soc_source)
if dropped:
logger.warning(f"Dropped {dropped} pre-acquired file(s) not matching version V{version:03d}")
all_soc_files = soc_file_tree(soc_source, to_list=True, exclude=exclude,
require_all=not allow_missing_products)
elif isinstance(soc_source, str):
# Local directory mode (also used after S3 download to temp dir)
if not os.path.exists(soc_source):
raise GediFileError(f"SOC source directory not found: {soc_source}")
# Version is resolved at entry, so the glob filter always excludes
# other releases and soc_file_tree's pivot key (orb_track, no
# version) cannot collide same-orbit granules across releases.
logger.info(f"Listing source SOC files (GEDI release V{version:03d})")
all_soc_files = soc_file_tree(soc_source, to_list=True, glob_kwargs={'version': version}, exclude=exclude,
require_all=not allow_missing_products)
else:
raise GediValidationError(f"Invalid soc_source type: {type(soc_source)}")
# Neither branch above consults `temporal` -- only the S3/download paths
# forward it to the CMR search. Apply it here so a build from an
# already-downloaded SOC tree honors -t0/-t1 instead of silently
# processing every granule with the requested products, any date.
# (No-op for variable_only_update: that path re-discovers files from
# `soc_source` directly inside `_build_add_variables` and never reads
# `all_soc_files`, by design -- it backfills the product across every
# existing granule, not just a new date slice.)
all_soc_files = _filter_soc_files_by_temporal(all_soc_files, temporal)
try:
if status_callback:
status_callback('PROCESSING')
if variable_only_update:
logger.info(f"Variable-only update: adding {set(product_vars.keys())} to existing database")
result = _build_add_variables(h3_dir, product_vars, soc_source=soc_source, version=version,
tmp_dir=tmp_dir, exclude=exclude, releases_callback=releases_callback)
if not result:
logger.warning(
"Variable update completed but no partition files were modified. "
"New columns may not have been added to the database."
)
return result
# Expand variable specifications and ensure L2A essentials
product_vars = _expand_product_vars(product_vars, all_soc_files, version=version)
_withhold_flags_absent_from_db(product_vars, h3_dir, version=version)
# Scope fingerprint over the EXPANDED variables. The expansion edits
# the caller's dict in place, and that dict is the build log's: what
# a resumed CLI reads back (and fingerprints before reconcile) is this
# expanded form, so hashing the raw specs would never match it.
scope_fp = _stage1_scope_fingerprint(spatial, res, part, product_vars)
# Filter to only files with required products
prod_soc_files = [{k: val for k, val in i.items() if k in product_vars} for i in all_soc_files]
# Filter out incomplete, corrupted, or already-processed granules
# L2A and what is published no later (L1B): missing, those mean an
# incomplete download, not a product that is not out yet.
required = (_PRODUCTS_BEFORE_L2B & set(product_vars)) or None if allow_missing_products else None
soc_files = _filter_granules(prod_soc_files, product_vars, skip_granules,
required_products=required, only_granules=only_granules)
if only_granules is not None:
_admitted = {(g['orbit'], g['granule'], g['track']) for g in only_granules}
_listed = {_granule_id_from_l2a_path(getattr(f, 'path', f))
for f in (next(iter(p.values())) for p in prod_soc_files if p)}
_drop_deferred_fragments(os.path.join(tmp_dir, 'partitions'), _listed - _admitted - {None})
if len(soc_files) == 0:
parquet_dir = os.path.join(tmp_dir, 'partitions')
if _has_unmerged_fragments(parquet_dir):
# Resume reconcile flips a granule INDEXED once its sentinels
# prove Stage 1 finished it, merged or not. A run stopped after
# Stage 1 but before the merge recorded progress leaves those
# fragments only in tmp: merge them, or they are lost while the
# log says INDEXED.
logger.info("No new granules to process; merging fragments left by an earlier run")
if status_callback:
status_callback('MERGING')
return _merge_and_finalize(parquet_dir, h3_dir)
logger.info("No new granules to process")
return None
if allow_missing_products:
if not _streaming_enabled():
raise GediValidationError("allow_missing_products requires the streaming writer (GH3_WRITE_STREAMING)")
# Complete granules first: the schema sample (the first readable
# granule) then carries every product whenever one is available.
soc_files.sort(key=lambda s: len(s) < len(product_vars))
n_partial = sum(len(s) < len(product_vars) for s in soc_files)
if n_partial:
logger.info(
f"{n_partial}/{len(soc_files)} granule(s) lack some requested products; their "
f"columns stay null until the products are published and backfilled"
)
if granule_products_callback is not None:
processed = {}
for s_ in soc_files:
path = next(iter(s_.values()))
key = _granule_id_from_l2a_path(getattr(path, 'path', path))
if key is not None:
processed[key] = sorted(p for p in product_vars if p not in s_)
granule_products_callback(processed)
# Create H3-indexed Dask DataFrame
ddf, lat_col, lon_col, dat_col, frag_names = _create_h3_dataframe(soc_files, product_vars, res, part)
# Apply spatial and existing-data filters
os.makedirs(tmp_dir, exist_ok=True)
ddf = _apply_spatial_filter(ddf, spatial, part, h3_dir)
if status_callback:
status_callback('PARTITIONING')
# Write partitioned parquet files to a subdirectory of tmp_dir to avoid
# conflicting with Dask worker scratch space (dask-worker-space/dirlock)
# which causes PermissionError on Windows when overwrite=True deletes tmp_dir
parquet_dir = os.path.join(tmp_dir, 'partitions')
# A prior merge that failed on corrupt fragments: drop them with their
# completion sentinels before the task list is built, so this Stage 1
# re-extracts exactly those (granule × beam) tasks (their granules are
# MERGE_FAILED, never skipped) and the merge below gets them back.
_pc = preclean_merge_failures(parquet_dir)
if _pc['parquets_removed']:
logger.info(f"Reopened {_pc['parquets_removed']} corrupt fragment task(s) from a failed merge")
if _streaming_enabled():
# Streaming writer: client.map + as_completed with per-task atomic
# leaves and per-(granule × beam) completion sentinels. Bounded
# worker memory; resume-safe; identical fragment layout to the
# legacy path (see _write_partitioned_streaming docstring).
wrote_any = _write_partitioned_streaming(
ddf, soc_files, product_vars, res, part,
parquet_dir, h3_dir, spatial,
lat_col, lon_col, dat_col,
allow_missing_products=allow_missing_products,
stage1_outcome_callback=stage1_outcome_callback,
scope_fingerprint=scope_fp,
)
else:
wrote_any = _write_partitioned(
ddf, parquet_dir, part, lat_col, lon_col, dat_col,
frag_names=frag_names,
)
if not wrote_any:
logger.info("No new data to process")
return None
if status_callback:
status_callback('MERGING')
# Merge and finalize database
h3_files = _merge_and_finalize(parquet_dir, h3_dir)
return h3_files
finally:
# Clean up S3 temp download directory
if _s3_tmp_dir is not None and os.path.exists(_s3_tmp_dir):
logger.debug(f"Cleaning up S3 temp directory: {_s3_tmp_dir}")
shutil.rmtree(_s3_tmp_dir, ignore_errors=True)
[docs]
def merge_build_logs(log_file_1: str, log_file_2: str, output_log_file: str) -> dict:
"""
Merge two build log files from separate databases.
Combines granules, columns, partition IDs, and date ranges while validating
configuration consistency (gedi_version, h3_resolution_level, h3_partition_level).
Parameters
----------
log_file_1 : str
Path to first build log (base log)
log_file_2 : str
Path to second build log (log to merge in)
output_log_file : str
Path to write merged build log
Returns
-------
dict
Merged log dictionary
Raises
------
FileNotFoundError
If either log file does not exist
ValueError
If log configurations are incompatible (different gedi_version, h3_resolution_level, or h3_partition_level)
Examples
--------
>>> merged_log = merge_build_logs(
... f'/path/to/database_world/{BUILD_LOG_FILENAME}',
... f'/path/to/database_world_a10/{BUILD_LOG_FILENAME}',
... f'/path/to/database_world_merged/{BUILD_LOG_FILENAME}'
... )
"""
# Load both log files
if not os.path.exists(log_file_1):
raise GediFileError(f"Log file not found: {log_file_1}")
if not os.path.exists(log_file_2):
raise GediFileError(f"Log file not found: {log_file_2}")
log1 = json_read(log_file_1)
log2 = json_read(log_file_2)
# Validate configuration compatibility
for key in ['gedi_version', 'h3_resolution_level', 'h3_partition_level']:
if log1.get(key) != log2.get(key):
raise GediMergeError(
f"Incompatible {key}: log1={log1.get(key)}, log2={log2.get(key)}. "
f"Cannot merge logs with different configurations."
)
# Start with log1 as base
merged_log = log1.copy()
merged_log['last_modified'] = now()
# Source exclusions: a file either build left out stays out.
from .logger import stored_exclude
exclude = list(dict.fromkeys(stored_exclude(log1) + stored_exclude(log2)))
if exclude:
merged_log['exclude'] = exclude
# Merge products and variables
products_1 = log1.get('products', {})
products_2 = log2.get('products', {})
merged_products = products_1.copy()
for prod_key, prod_data in products_2.items():
if prod_key in merged_products:
# Product exists in both logs - merge variables
existing_vars = set(merged_products[prod_key].get('variables', []))
new_vars = set(prod_data.get('variables', []))
merged_vars = sorted(list(existing_vars | new_vars))
merged_products[prod_key]['variables'] = merged_vars
# Use COMPLETED status if both are complete
if merged_products[prod_key].get('status') == 'COMPLETED' and prod_data.get('status') == 'COMPLETED':
merged_products[prod_key]['status'] = 'COMPLETED'
# Update timestamp to most recent
ts1 = merged_products[prod_key].get('last_modified', '')
ts2 = prod_data.get('last_modified', '')
merged_products[prod_key]['last_modified'] = max(ts1, ts2) if ts1 and ts2 else (ts1 or ts2)
else:
# Product only in log2 - add it
merged_products[prod_key] = prod_data.copy()
merged_log['products'] = merged_products
# Merge granules (deduplicate)
granules_1 = log1.get('granules', [])
granules_2 = log2.get('granules', [])
merged_granules = granules_1.copy()
# Entries alike but for their recorded releases (one log predates them, or
# the databases read different files) become one entry. A product keeps a
# tag only if both carry one, the older: an unknown source may be older
# still, and a re-release must be reported, never hidden (#47).
def _sans_releases(e):
return {k: v for k, v in e.items() if k != 'releases'}
def _fold(a, b):
out = _sans_releases(a)
if 'releases' in a and 'releases' in b:
ra, rb = a['releases'], b['releases']
out['releases'] = {p: min(ra[p], rb[p], key=gedi_release_rank) for p in ra if p in rb}
return out
by_key = {}
for i, e in enumerate(merged_granules):
by_key.setdefault((e.get('orbit'), e.get('granule'), e.get('track')), []).append(('1', i))
unpaired = []
for g in granules_2:
key = (g.get('orbit'), g.get('granule'), g.get('track'))
pool = {'1': merged_granules, '2': unpaired}
same = next(((src, i) for src, i in by_key.get(key, ())
if _sans_releases(pool[src][i]) == _sans_releases(g)), None)
if same is None:
by_key.setdefault(key, []).append(('2', len(unpaired)))
unpaired.append(g)
else:
pool[same[0]][same[1]] = _fold(pool[same[0]][same[1]], g)
for g in unpaired:
if g in merged_granules:
continue
# NO_DATA ("no rows under that database's scope") loses to any other
# status for the same granule: the other log holds evidence of rows
# or of an unfinished read. Other collisions keep both entries.
key = (g.get('orbit'), g.get('granule'), g.get('track'))
twins = [i for i, e in enumerate(merged_granules)
if (e.get('orbit'), e.get('granule'), e.get('track')) == key
and 'NO_DATA' in (e.get('status'), g.get('status'))]
if not twins:
merged_granules.append(g)
elif g.get('status') != 'NO_DATA':
for i in reversed(twins[1:]):
del merged_granules[i]
merged_granules[twins[0]] = g
# NO_DATA was proven under one database's scope; the other never examined
# the granule, so unless both logs say so it is PENDING again.
def _no_data_keys(granules):
return {(e.get('orbit'), e.get('granule'), e.get('track'))
for e in granules if e.get('status') == 'NO_DATA'}
both = _no_data_keys(granules_1) & _no_data_keys(granules_2)
for i, e in enumerate(merged_granules):
if e.get('status') == 'NO_DATA' and (e.get('orbit'), e.get('granule'), e.get('track')) not in both:
merged_granules[i] = {**e, 'status': 'PENDING'}
if merged_granules:
merged_log['granules'] = merged_granules
# Merge h3_columns (deduplicate and sort)
cols_1 = set(log1.get('h3_columns', []))
cols_2 = set(log2.get('h3_columns', []))
merged_cols = sorted(list(cols_1 | cols_2))
if merged_cols:
merged_log['h3_columns'] = merged_cols
# Merge h3_columns_dtypes (post-merge invariant: identical schema
# across partitions, so the two logs should agree on overlapping
# columns; later log wins on conflict, matching merge_build_logs'
# general "log2 augments log1" semantics).
dtypes_1 = log1.get('h3_columns_dtypes') or {}
dtypes_2 = log2.get('h3_columns_dtypes') or {}
if dtypes_1 or dtypes_2:
merged_log['h3_columns_dtypes'] = {**dtypes_1, **dtypes_2}
# Merge h3_partition_ids (deduplicate and sort)
parts_1 = set(log1.get('h3_partition_ids', []))
parts_2 = set(log2.get('h3_partition_ids', []))
merged_parts = sorted(list(parts_1 | parts_2))
if merged_parts:
merged_log['h3_partition_ids'] = merged_parts
# Merge date_range (min start date, max end date)
dr1 = log1.get('date_range')
dr2 = log2.get('date_range')
if dr1 and dr2:
merged_log['date_range'] = [min(dr1[0], dr2[0]), max(dr1[1], dr2[1])]
elif dr2:
merged_log['date_range'] = dr2
# Merge spatial_filter (use union of geometries if both exist)
spatial_1 = log1.get('spatial_filter')
spatial_2 = log2.get('spatial_filter')
if spatial_1 and spatial_2:
# Combine feature collections
features_1 = spatial_1.get('features', [])
features_2 = spatial_2.get('features', [])
merged_features = features_1 + features_2
merged_log['spatial_filter'] = {
'type': 'FeatureCollection',
'features': merged_features
}
# Update bbox if available
if 'bbox' in spatial_1 and 'bbox' in spatial_2:
bbox1 = spatial_1['bbox']
bbox2 = spatial_2['bbox']
merged_log['spatial_filter']['bbox'] = [
min(bbox1[0], bbox2[0]), # min lon
min(bbox1[1], bbox2[1]), # min lat
max(bbox1[2], bbox2[2]), # max lon
max(bbox1[3], bbox2[3]) # max lat
]
elif spatial_2:
merged_log['spatial_filter'] = spatial_2
# Write merged log
os.makedirs(os.path.dirname(output_log_file), exist_ok=True)
json_write(merged_log, output_log_file, rewrite=True)
logger.info(f"Merged build logs: {log_file_1} + {log_file_2} -> {output_log_file}")
return merged_log
