From b3d7a9d4404402d528bc68471969a1ec3197846a Mon Sep 17 00:00:00 2001 From: michael-corey Date: Mon, 20 Apr 2026 10:18:09 -0500 Subject: [PATCH] adding index field --- PLAN.md | 84 -- generic_loader/PARTITION_DESIGN.md | 938 ----------------------- generic_loader/load_folder.py | 125 ++- generic_loader/load_sas.py | 140 ++++ generic_loader/sample_config.yaml | 9 + generic_loader/sample_folder_config.yaml | 20 + 6 files changed, 289 insertions(+), 1027 deletions(-) delete mode 100644 PLAN.md delete mode 100644 generic_loader/PARTITION_DESIGN.md diff --git a/PLAN.md b/PLAN.md deleted file mode 100644 index 582ae3d..0000000 --- a/PLAN.md +++ /dev/null @@ -1,84 +0,0 @@ -# Plan: LIST Partition Support for generic_loader - -## Objective - -Produce an implementation-ready design for adding PostgreSQL LIST partitioning to the generic loader flows in `generic_loader/load_sas.py` and `generic_loader/load_folder.py`, driven by YAML configuration and compatible with the current streaming `COPY` load path. - -## Current context - -- `generic_loader/load_sas.py` - - `LoaderConfig` currently parses `filename`, `schemaname`, `tablename`, `if_exists`, `include`, and `exclude`. - - Schema inference is based on `read_sas_preview()` plus `infer_schema()`. - - `render_create_table()` emits one non-partitioned `CREATE TABLE` statement. - - `create_table()` handles `if_exists=fail|replace|append` for a single table. - - `copy_dataframes()` streams rows into the target table with `COPY ... FROM STDIN`. -- `generic_loader/load_folder.py` - - `FolderConfig` carries folder defaults. - - `ClusterSpec` carries resolved per-cluster load settings. - - `_ExplicitPattern` stores optional per-cluster overrides for explicit matches. - - `discover_clusters()` resolves cluster inheritance for `if_exists`, `include`, and `exclude`. - - `load_cluster()` infers schema from the first file, creates one table, then streams all files into it. -- Existing warnings are emitted to stderr as `[warn] ...`; the codebase does not currently use the `logging` module. -- This task is design-only in architect mode; no production Python changes are being made here. - -## Assumptions and constraints - -- PostgreSQL native LIST partitioning only; no range/hash partitioning and no automatic creation of missing partitions during append. -- Data continues to be copied to the parent table so PostgreSQL performs routing; `copy_dataframes()` behavior should remain unchanged. -- Accurate partition creation requires a complete discovery pass across all incoming rows unless the preview is known to already contain the full dataset. -- Folder-level partition settings should resolve into concrete per-cluster settings using the same inheritance style as current folder defaults. -- `if_exists=append` must validate compatibility and skip partition creation. -- Documentation must be detailed enough for an implementer to modify code, a QA lead to derive test scenarios, and a docs lead to update user-facing instructions without guessing. - -## Files and systems likely affected - -- `generic_loader/load_sas.py` -- `generic_loader/load_folder.py` -- `generic_loader/sample_config.yaml` -- `generic_loader/sample_folder_config.yaml` -- `generic_loader/PARTITION_DESIGN.md` -- Potentially module/CLI docstrings in `generic_loader/load_sas.py` and `generic_loader/load_folder.py` - -## Implementation approach - -1. Extend YAML and dataclass config surfaces with `partition_by` and `max_partitions`. -2. Add partition-planning helpers that: - - validate partition columns, - - normalize partition values consistently with the existing `COPY` preparation rules, - - discover cascading unique values across one file or an entire cluster, - - count resulting child partition tables and emit a threshold warning. -3. Extend DDL rendering so the parent table can be declared with `PARTITION BY LIST (...)` and child tables can be emitted recursively with `CREATE TABLE ... PARTITION OF ... FOR VALUES IN (...)`. -4. Extend table creation rules: - - `replace` drops the parent with `CASCADE` when partitioning is enabled and recreates the full tree, - - `fail` errors if the parent exists, - - `append` validates schema plus partition-key compatibility and does not create partitions. -5. Extend dry-run output so partitioned loads print the full ordered DDL set and perform the required partition discovery pass. -6. Extend folder orchestration so per-cluster partition settings inherit or override folder defaults in the same style as current config resolution. - -## Risks and edge cases - -- High-cardinality partition columns can generate very large partition trees, long DDL output, and slow Postgres planning. -- Empty strings in text columns currently become `NULL` on load because of `COPY ... NULL ''`; partition discovery must mirror that behavior or routing will be wrong. -- Different raw values can collide after sanitization or truncation; deterministic disambiguation is required. -- `NULL` partition values need explicit support in both DDL generation and child-table naming. -- Partitioned dry-runs become more expensive because they require scanning full source data rather than using only the schema preview. -- Multi-file clusters can still fail later on schema differences outside partition columns unless compatibility checks are broadened deliberately. - -## Acceptance criteria - -- The design document specifies exact YAML changes, dataclass changes, new helper functions, modified functions, algorithms, error handling, dry-run behavior, and `if_exists` semantics. -- Single-file and folder flows are both covered, including per-cluster inheritance/override behavior. -- Child-table naming, literal rendering, warning semantics, and append-mode validation are precise enough to implement directly. -- The design explicitly identifies what remains unchanged, especially the `COPY` routing path. - -## Validation strategy - -- Cross-check the plan against the current call graph and responsibilities in `load_sas.py` and `load_folder.py`. -- Prefer minimal-regression changes that preserve existing non-partitioned behavior. -- Include pseudocode and concrete examples for recursive partition DDL generation, cascading value discovery, null handling, and dry-run output. - -## Documentation updates required - -- Create `generic_loader/PARTITION_DESIGN.md` as the primary implementation-ready design artifact. -- Include exact sample YAML snippets for single-file and folder loaders. -- Document the dry-run cost change for partitioned loads and the `append` limitation that partitions are not auto-created. diff --git a/generic_loader/PARTITION_DESIGN.md b/generic_loader/PARTITION_DESIGN.md deleted file mode 100644 index b4af9d1..0000000 --- a/generic_loader/PARTITION_DESIGN.md +++ /dev/null @@ -1,938 +0,0 @@ -# Partition Feature Design for generic_loader - -## 1. Objective - -Add PostgreSQL LIST partitioning support to [`load_sas.py`](generic_loader/load_sas.py) and [`load_folder.py`](generic_loader/load_folder.py) without changing the existing streaming `COPY` data path in [`copy_dataframes()`](generic_loader/load_sas.py:1028). The feature must be YAML-driven, must support cascading partition levels, and must keep non-partitioned behavior unchanged. - -## 2. Current baseline - -### Single-file loader - -The single-file path is centered in [`generic_loader/load_sas.py`](generic_loader/load_sas.py): - -- [`LoaderConfig`](generic_loader/load_sas.py:273) stores file path, target schema/table, `if_exists`, and column filters. -- [`load_config()`](generic_loader/load_sas.py:350) parses YAML. -- [`read_sas_preview()`](generic_loader/load_sas.py:430) reads a bounded preview for schema inference. -- [`infer_schema()`](generic_loader/load_sas.py:637) infers Postgres column types. -- [`render_create_table()`](generic_loader/load_sas.py:756) renders one non-partitioned `CREATE TABLE` statement. -- [`create_table()`](generic_loader/load_sas.py:890) executes table creation or append/replace checks. -- [`copy_dataframes()`](generic_loader/load_sas.py:1028) streams chunks into the target table via `COPY ... FROM STDIN`. - -### Folder loader - -The folder path is centered in [`generic_loader/load_folder.py`](generic_loader/load_folder.py): - -- [`ClusterSpec`](generic_loader/load_folder.py:137) stores resolved per-cluster load settings. -- [`_ExplicitPattern`](generic_loader/load_folder.py:148) stores raw per-cluster YAML overrides. -- [`FolderConfig`](generic_loader/load_folder.py:160) stores folder defaults. -- [`load_folder_config()`](generic_loader/load_folder.py:200) parses folder YAML. -- [`discover_clusters()`](generic_loader/load_folder.py:295) resolves inheritance and groups files. -- [`load_cluster()`](generic_loader/load_folder.py:385) creates a table from the first file and streams every file in the cluster into it. - -### Important current behaviors to preserve - -- [`copy_dataframes()`](generic_loader/load_sas.py:1028) copies into exactly one qualified table name and should remain unchanged. -- [`create_table()`](generic_loader/load_sas.py:890) owns `if_exists` semantics and should remain the single gate for fail/replace/append behavior. -- Warnings are currently emitted to stderr as `[warn] ...`, for example in [`_assert_schema_compatible()`](generic_loader/load_sas.py:826), and the feature should follow that pattern instead of introducing a repository-wide logging refactor. - -## 3. Scope and non-goals - -### In scope - -- Optional YAML `partition_by` support. -- Configurable `max_partitions` threshold with default `10000`. -- Single-level and multi-level cascading LIST partitions. -- Partition value discovery from the incoming dataset at runtime. -- Recursive DDL generation for parent and child partitions. -- Folder-level defaults plus per-cluster overrides. -- Dry-run output for the full DDL tree. - -### Explicitly out of scope for this implementation - -- RANGE or HASH partitioning. -- Expression-based partition keys. -- Changing row-routing behavior in [`copy_dataframes()`](generic_loader/load_sas.py:1028). -- Automatically creating missing partitions in `append` mode. -- Reworking manifest validation in [`validate_against_manifest()`](generic_loader/load_sas.py:1102). - -## 4. YAML schema changes - -## 4.1 Single-file config - -Update the sample shape documented by [`generic_loader/sample_config.yaml`](generic_loader/sample_config.yaml) to include `partition_by` and `max_partitions`. - -### Proposed exact example - -```yaml -filename: samples/sample_kitchensink.xpt -schemaname: public -tablename: kitchensink - -# Optional. If set, only these columns are loaded. Mutually exclusive with exclude. -# include: -# - ID -# - INTCOL -# - DATECOL - -# Optional. Columns to drop. -# exclude: -# - ALLNULL - -# Optional. Create cascading LIST partitions in this order. -# Omit or set [] for no partitioning. -partition_by: - - state - - zip - -# Optional. Warn if the load would create more than this many partition tables. -# The load continues. Default: 10000. -max_partitions: 10000 - -# What to do if the target table already exists: fail | replace | append -# Defaults to fail. -if_exists: append -``` - -### Parsing and validation rules - -1. `partition_by` is optional. -2. Omitted, `null`, or `[]` means "not partitioned". -3. When present and non-empty, it must be a YAML sequence of non-empty strings. -4. Order matters. `['state', 'zip']` means `state` is level 1 and `zip` is level 2. -5. Duplicate names are invalid. -6. If `include` is present, every `partition_by` column must be included. -7. If `exclude` is present, no `partition_by` column may be excluded. -8. `max_partitions` is optional and defaults to `10000`. -9. `max_partitions` must be an integer greater than `0`. - -## 4.2 Folder config - -Update the sample shape documented by [`generic_loader/sample_folder_config.yaml`](generic_loader/sample_folder_config.yaml) to include folder defaults and per-cluster overrides. - -### Proposed exact example - -```yaml -folder: samples/folder_test -schemaname: public - -# Applied when creating the first file of each cluster. -# One of: fail | replace | append. Default: fail. -if_exists: replace - -# When true (default), any file not matched by an explicit pattern below is -# auto-grouped with its peers. -auto_detect: true - -# Optional folder-level column filter. -# include: -# - ID -# - INTCOL -# exclude: -# - ALLNULL - -# Optional folder default for LIST partitioning. -partition_by: - - state - - zip - -# Optional folder default threshold. Default: 10000. -max_partitions: 10000 - -clusters: - - pattern: '^group_a\d+\.xpt$' - tablename: group_a - # Inherits folder-level partition_by and max_partitions. - - - pattern: '^group_b\d+\.xpt$' - tablename: group_b - partition_by: - - state - max_partitions: 2000 - - - pattern: '^standalone\.xpt$' - tablename: standalone - partition_by: [] # Explicit opt-out of the folder default. -``` - -### Folder override rules - -1. Folder-level `partition_by` and `max_partitions` behave as defaults. -2. In an explicit cluster entry: - - if `partition_by` is omitted, inherit the folder-level value; - - if `partition_by` is a non-empty list, replace the folder-level value; - - if `partition_by: []`, explicitly disable partitioning for that cluster. -3. Cluster-level `max_partitions` overrides the folder-level threshold when present. -4. The resolved per-cluster rules should follow the same pattern already used by [`discover_clusters()`](generic_loader/load_folder.py:295) for `if_exists`, `include`, and `exclude`. - -## 5. Dataclass changes - -## 5.1 Existing public config dataclasses - -### [`LoaderConfig`](generic_loader/load_sas.py:273) - -Add: - -- `partition_by: Optional[List[str]] = None` -- `max_partitions: int = 10000` - -### [`ClusterSpec`](generic_loader/load_folder.py:137) - -Add resolved fields: - -- `partition_by: Optional[List[str]]` -- `max_partitions: int` - -### [`_ExplicitPattern`](generic_loader/load_folder.py:148) - -Add raw optional override fields: - -- `partition_by: Optional[List[str]] = None` -- `max_partitions: Optional[int] = None` - -Notes: - -- Preserve `partition_by=[]` when it appears in the YAML so [`discover_clusters()`](generic_loader/load_folder.py:295) can distinguish explicit disable from inheritance. -- `max_partitions` remains `None` when omitted so folder inheritance can resolve it later. - -### [`FolderConfig`](generic_loader/load_folder.py:160) - -Add: - -- `partition_by: Optional[List[str]] = None` -- `max_partitions: int = 10000` - -## 5.2 Recommended new internal helper dataclasses - -These are not required to be public, but they make the implementation substantially safer and clearer. - -### Recommended [`PartitionNode`](generic_loader/PARTITION_DESIGN.md) - -Suggested fields: - -- `field_name: str` -- `value: Any` - The normalized value Postgres will see during `COPY`; use `None` for SQL `NULL`. -- `table_name: str` -- `children: List[PartitionNode] = field(default_factory=list)` - -### Recommended [`PartitionPlan`](generic_loader/PARTITION_DESIGN.md) - -Suggested fields: - -- `fields: List[str]` -- `roots: List[PartitionNode]` -- `total_partition_tables: int` - -The implementation can use nested dicts instead, but an explicit plan object reduces naming, recursion, and dry-run bugs. - -## 6. New functions needed - -The exact names may vary, but the design should introduce helpers with the responsibilities below. - -## 6.1 Config parsing helpers - -### Recommended [`_parse_partition_by()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Parse `partition_by` from YAML. -- Enforce list-of-strings validation. -- Normalize omitted/empty top-level values to `None`. -- Preserve cluster-level empty list `[]` long enough for override resolution. - -### Recommended [`_parse_max_partitions()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Parse and validate `max_partitions`. -- Enforce positive integer semantics. - -## 6.2 Partition validation helpers - -### Recommended [`_validate_partition_columns()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Ensure every requested partition column exists after [`apply_column_filter()`](generic_loader/load_sas.py:468). -- Fail early if a partition column was removed by `include` or `exclude`. -- Produce context-rich errors that name the config, file, or cluster. - -### Recommended [`_assert_partition_compatible()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- In `append` mode, verify that the existing parent table is LIST-partitioned on the same ordered keys. -- Reuse [`SchemaCompatibilityError`](generic_loader/load_sas.py:307) for incompatibility. - -Expected catalog check: - -- Query `pg_partitioned_table` for `partstrat`. -- Query `pg_attribute` using `partattrs` order to get the parent key columns. -- Require `partstrat = 'l'`. -- Require the ordered key list to exactly equal the resolved `partition_by` list. - -## 6.3 Partition discovery helpers - -### Recommended [`discover_partition_values()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Scan an iterable of filtered DataFrames. -- Normalize the partition columns the same way [`_prepare_for_copy()`](generic_loader/load_sas.py:943) will normalize them for `COPY`. -- Build a cascading partition tree scoped by parent value. -- Count the child partition tables that will be created. - -Suggested input shape: - -- `dfs: Iterable[pd.DataFrame]` -- `columns: Dict[str, ColumnSpec]` -- `partition_by: List[str]` -- `root_table_name: str` - -Suggested output shape: - -- `PartitionPlan` - -### Recommended [`_warn_if_partition_count_exceeds()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Emit `[warn] ...` to stderr if `plan.total_partition_tables > max_partitions`. -- Never abort the load. - -## 6.4 Naming and literal helpers - -### Recommended [`_sanitize_partition_token()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Convert a normalized partition value into a safe, deterministic table-name suffix. - -### Recommended [`_build_partition_table_name()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Combine parent name and sanitized token. -- Enforce Postgres identifier-length limits. -- Resolve collisions deterministically. - -### Recommended [`_render_partition_literal()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Render one value for `FOR VALUES IN (...)`. -- Preserve the exact routed value Postgres will see during `COPY`. - -## 6.5 DDL rendering helpers - -### Recommended [`render_partition_ddl()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Render child `CREATE TABLE ... PARTITION OF ...` statements recursively. - -### Recommended [`render_create_table_statements()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Return the full ordered statement list for dry-run and actual execution. -- Keep the root statement first. -- Append recursive child statements afterward. - -## 6.6 Optional shared warning helper - -### Recommended [`_warn()`](generic_loader/PARTITION_DESIGN.md) - -Purpose: - -- Centralize the existing `[warn] ...` stderr behavior. -- Let both existing schema warnings and new partition warnings share one implementation. - -## 7. Modified functions - -## 7.1 [`load_config()`](generic_loader/load_sas.py:350) - -Modify to: - -1. Parse `partition_by`. -2. Parse `max_partitions`. -3. Validate include/exclude conflicts with `partition_by`. -4. Return the new fields in [`LoaderConfig`](generic_loader/load_sas.py:273). - -## 7.2 [`render_create_table()`](generic_loader/load_sas.py:756) - -Modify signature to accept optional partition metadata: - -- `partition_by: Optional[List[str]] = None` - -Behavior: - -- If `partition_by` is falsy, keep current output unchanged. -- If `partition_by` is present, append `PARTITION BY LIST ()` to the parent statement. -- This function should still render only the parent statement; child statements belong in [`render_partition_ddl()`](generic_loader/PARTITION_DESIGN.md). - -Example parent output: - -```sql -CREATE TABLE "public"."customers" ( - "state" TEXT, - "zip" TEXT, - "name" TEXT -) PARTITION BY LIST ("state"); -``` - -## 7.3 [`_drop_table()`](generic_loader/load_sas.py:771) - -Add an optional `cascade: bool = False` parameter. - -Behavior: - -- Non-partitioned replace keeps current plain `DROP TABLE` behavior. -- Partitioned replace uses `DROP TABLE CASCADE` so the parent drop removes all partitions. - -## 7.4 [`create_table()`](generic_loader/load_sas.py:890) - -Extend signature to accept: - -- `partition_by: Optional[List[str]] = None` -- `partition_plan: Optional[PartitionPlan] = None` - -Behavior: - -1. Preserve current `if_exists` validation. -2. For non-partitioned loads, preserve current behavior. -3. For partitioned loads: - - `fail`: if the parent table exists, raise [`TableExistsError`](generic_loader/load_sas.py:303). - - `replace`: if the parent exists, drop it with `CASCADE`, then recreate the full tree. - - `append`: run [`_assert_schema_compatible()`](generic_loader/load_sas.py:826) plus [`_assert_partition_compatible()`](generic_loader/PARTITION_DESIGN.md), then return without creating any partitions. -4. When creation is needed, execute the full statement list returned by [`render_create_table_statements()`](generic_loader/PARTITION_DESIGN.md). -5. Reject `partition_by` without a computed `partition_plan` when creation or dry-run rendering needs it. - -## 7.5 [`_prepare_for_copy()`](generic_loader/load_sas.py:943) - -Recommended refactor: - -- Extract or share the per-column normalization logic so partition discovery can use the same conversion rules. -- Do not change external behavior of the returned DataFrame. - -Reason: - -- The partition discovery pass must reason about the same values Postgres will actually receive. -- The most important special case is text columns, where empty strings currently become SQL `NULL` because [`copy_dataframes()`](generic_loader/load_sas.py:1028) uses `NULL ''`. - -## 7.6 [`main()`](generic_loader/load_sas.py:1192) - -Modify the single-file flow as follows: - -1. Load config. -2. Read preview and infer schema exactly as today. -3. Validate that partition columns exist after filtering. -4. If `partition_by` is set and the operation needs creation or dry-run rendering, run a full discovery pass over the file to build a `PartitionPlan`. -5. In dry-run mode, print the full DDL statement list rather than only the parent statement. -6. In live mode, pass `partition_by` and `partition_plan` into [`create_table()`](generic_loader/load_sas.py:890). -7. Keep [`copy_dataframes()`](generic_loader/load_sas.py:1028) unchanged so data is copied to the parent table and Postgres routes rows automatically. - -## 7.7 [`load_folder_config()`](generic_loader/load_folder.py:200) - -Modify to: - -1. Parse folder-level `partition_by` and `max_partitions`. -2. Parse per-cluster `partition_by` and `max_partitions`. -3. Validate include/exclude conflicts against the applicable partition list where possible. -4. Preserve explicit `partition_by: []` so cluster discovery can treat it as "disable inheritance". - -## 7.8 [`discover_clusters()`](generic_loader/load_folder.py:295) - -Modify to resolve per-cluster partition settings. - -For each resolved [`ClusterSpec`](generic_loader/load_folder.py:137): - -- `partition_by = patt.partition_by if patt.partition_by is not None else cfg.partition_by` -- `max_partitions = patt.max_partitions if patt.max_partitions is not None else cfg.max_partitions` -- normalize resolved empty list to `None` before storing on the final [`ClusterSpec`](generic_loader/load_folder.py:137) - -## 7.9 [`load_cluster()`](generic_loader/load_folder.py:385) - -Modify the cluster load order to: - -1. Infer schema from the first file exactly as today. -2. Validate partition columns against that schema. -3. If the cluster is partitioned and the operation is not append-only verification, scan all files in the cluster to build one shared `PartitionPlan`. -4. Call [`create_table()`](generic_loader/load_sas.py:890) with resolved `partition_by` and `partition_plan`. -5. Stream all files into the parent table exactly as today. - -## 7.10 [`main()`](generic_loader/load_folder.py:496) - -Modify dry-run behavior: - -- keep cluster discovery output; -- for each loadable cluster, print full DDL, not only one `CREATE TABLE` statement; -- when a cluster is partitioned, perform partition discovery across every file in that cluster, not only the first file. - -Also update the `--dry-run` help text because the current wording says the schema is inferred from only the first file of the cluster. - -## 8. Partition value discovery algorithm - -## 8.1 High-level rules - -1. Discovery operates on filtered data, meaning after the same column filter logic used by [`apply_column_filter()`](generic_loader/load_sas.py:468). -2. Discovery must use the same semantic values that Postgres will see during `COPY`, not raw pandas object identity. -3. The scan should be streaming and chunk-based to avoid materializing the full file or cluster in memory. -4. The resulting tree must scope each level under its parent so deeper values are not treated as globally unique. - -## 8.2 Normalization rules for partition keys - -Partition discovery should normalize each partition column using the same type-aware logic already embodied in [`_prepare_for_copy()`](generic_loader/load_sas.py:943), with the following behavior: - -- Integer-like columns (`INTEGER`, `BIGINT`, `SMALLINT`): coerce object values through numeric conversion, treat blank strings and NaN as `NULL`. -- Floating/numeric columns (`DOUBLE PRECISION`, `REAL`, `NUMERIC`): numeric conversion, NaN becomes `NULL`. -- Date columns: normalize to `datetime.date` or `NULL`. -- Timestamp columns: normalize to `datetime.datetime` or `NULL`. -- Time columns: normalize through the existing time conversion path or `NULL`. -- Text-like columns: `None`, pandas nulls, and `''` all become semantic `NULL`, because [`copy_dataframes()`](generic_loader/load_sas.py:1028) sends empty strings with `NULL ''`. -- Boolean columns: normalize to `True`, `False`, or `NULL`. - -This means partition discovery deduplicates on the routed value, not the raw source representation. For example, `'00123'` and `123` in an integer partition column should produce one partition value `123`, not two separate partitions. - -## 8.3 Discovery pseudocode - -```python -def discover_partition_values(dfs, columns, partition_by, root_table_name): - validate_partition_columns(columns, partition_by) - - root = PartitionPlan(fields=partition_by, roots=[], total_partition_tables=0) - root_index = {} # normalized value -> PartitionNode for depth 0 - - for df in dfs: - if df.empty: - continue - - part_df = df[partition_by].copy() - part_df = normalize_partition_frame(part_df, columns) - unique_paths = part_df.drop_duplicates() - - for path in unique_paths.itertuples(index=False, name=None): - parent_table = root_table_name - parent_children = root.roots - parent_index = root_index - - for depth, value in enumerate(path): - field_name = partition_by[depth] - if value not in parent_index: - child_table = build_partition_table_name(parent_table, value) - node = PartitionNode( - field_name=field_name, - value=value, - table_name=child_table, - ) - parent_index[value] = node - parent_children.append(node) - root.total_partition_tables += 1 - - node = parent_index[value] - parent_table = node.table_name - parent_children = node.children - parent_index = getattr(node, "_index", {}) - - sort_every_node_deterministically(root) - return root -``` - -## 8.4 Efficient implementation notes - -- The scan should retain only the partition columns for the current chunk after filtering. -- The in-memory structure should grow only with the number of unique partition nodes, not the number of rows. -- Reading partition values from the preview frame is only valid when that frame is known to contain the entire dataset. In the current CLI flow, the preview is normally not exhaustive, so partitioned loads should perform a full chunked scan. -- A future optimization may add optional reader-level column pruning to [`iter_sas_chunks()`](generic_loader/load_sas.py:447) and [`read_sas_preview()`](generic_loader/load_sas.py:430), but that is not required for correctness. - -## 9. DDL generation algorithm - -## 9.1 Root table - -If `partition_by` is set, the parent statement produced by [`render_create_table()`](generic_loader/load_sas.py:756) must end with: - -```sql -PARTITION BY LIST ("") -``` - -The parent still contains the full column list. - -## 9.2 Child tables - -For each discovered node: - -- if it is not the last partition level, create a child partition that is itself subpartitioned by the next field; -- if it is the last partition level, create a leaf partition with no further `PARTITION BY` clause. - -Examples for `partition_by: [state, zip]`: - -```sql -CREATE TABLE "public"."customers_ca" -PARTITION OF "public"."customers" -FOR VALUES IN ('CA') -PARTITION BY LIST ("zip"); - -CREATE TABLE "public"."customers_ca_60601" -PARTITION OF "public"."customers_ca" -FOR VALUES IN ('60601'); -``` - -## 9.3 DDL rendering pseudocode - -```python -def render_create_table_statements(schema, table, columns, partition_by, plan): - statements = [render_create_table(schema, table, columns, partition_by=partition_by)] - if partition_by: - statements.extend(render_partition_ddl(schema, table, columns, partition_by, plan.roots, depth=0)) - return statements - - -def render_partition_ddl(schema, parent_table, columns, partition_by, nodes, depth): - field_name = partition_by[depth] - next_field = partition_by[depth + 1] if depth + 1 < len(partition_by) else None - field_spec = columns[field_name] - statements = [] - - for node in nodes: - literal = render_partition_literal(node.value, field_spec) - if next_field is None: - statements.append( - f'CREATE TABLE {qualified(schema, node.table_name)} ' - f'PARTITION OF {qualified(schema, parent_table)} ' - f'FOR VALUES IN ({literal});' - ) - else: - statements.append( - f'CREATE TABLE {qualified(schema, node.table_name)} ' - f'PARTITION OF {qualified(schema, parent_table)} ' - f'FOR VALUES IN ({literal}) ' - f'PARTITION BY LIST ({quote_ident(next_field)});' - ) - statements.extend( - render_partition_ddl( - schema, - node.table_name, - columns, - partition_by, - node.children, - depth + 1, - ) - ) - - return statements -``` - -## 9.4 Statement order - -Emit statements in this order: - -1. parent table; -2. each level-1 child; -3. that child’s descendants before moving to the next sibling. - -This depth-first order guarantees that every parent exists before its children are created. - -## 10. Table-name sanitization rules - -The child-table name rule must be deterministic and explicit. - -## 10.1 Base token generation - -For each normalized partition value: - -1. Convert to a display token: - - `None` -> `null` - - `datetime.date`, `datetime.time`, `datetime.datetime` -> `isoformat()` string - - everything else -> `str(value)` -2. Lowercase the token. -3. Replace every run of one or more non-alphanumeric characters with `_`. -4. Trim leading and trailing `_`. -5. If the result is empty, use `value`. - -Examples: - -- `CA` -> `ca` -- `New York` -> `new_york` -- `60601-1234` -> `60601_1234` -- `NULL` -> `null` -- `***` -> `value` - -## 10.2 Final child name - -Child names are: - -```text -{parent_table}_{sanitized_token} -``` - -Examples: - -- `customers` + `CA` -> `customers_ca` -- `customers_ca` + `60601` -> `customers_ca_60601` - -## 10.3 Length limit - -Postgres identifiers are limited to 63 bytes. The implementation should treat 63 characters as the working limit because the loader currently emits ASCII-only sanitized suffixes. - -Rules: - -1. If `len(parent_table) >= 62`, fail fast with a clear error because there is no room for `_x`. -2. Otherwise, reserve `len(parent_table) + 1` characters for the prefix and underscore. -3. Truncate only the sanitized token, not the parent prefix. -4. If truncation makes two child names collide, append a deterministic short hash. - -## 10.4 Collision handling - -Different raw values can sanitize to the same token, for example: - -- `A-B` -> `a_b` -- `A B` -> `a_b` - -Recommended collision rule: - -1. First candidate: `parent_a_b` -2. On collision, append `_` derived from the exact normalized value for that node. -3. Re-truncate the base token as needed so the final name still fits the 63-character limit. - -Example: - -- `parent_a_b` -- `parent_a_b_f15c2d19` - -This keeps names deterministic across runs and avoids dependence on discovery order. - -## 11. Partition literal rendering rules - -The `FOR VALUES IN (...)` clause must use the exact routed value Postgres will receive after loader normalization. - -Recommended rendering rules: - -- `NULL` -> `NULL` -- text -> single-quoted with internal quotes escaped -- integers / numerics -> unquoted numeric literal -- boolean -> `TRUE` or `FALSE` -- date -> `DATE 'YYYY-MM-DD'` -- timestamp -> `TIMESTAMP 'YYYY-MM-DD HH:MM:SS'` -- time -> `TIME 'HH:MM:SS'` - -Important special case: - -- text `''` must not render as `''`; it must render as `NULL` because [`copy_dataframes()`](generic_loader/load_sas.py:1028) uses `NULL ''`. - -## 12. `if_exists` interaction - -## 12.1 `fail` - -- If the parent table exists, behavior is unchanged: raise [`TableExistsError`](generic_loader/load_sas.py:303). -- No partition compatibility inspection is needed because the operation stops immediately. - -## 12.2 `replace` - -- If the parent table exists and the config is partitioned, execute `DROP TABLE CASCADE`. -- Recreate the parent plus every partition statement in one transaction. -- If any statement fails, let the outer transaction rollback preserve atomicity. - -## 12.3 `append` - -Required behavior: - -1. Run [`_assert_schema_compatible()`](generic_loader/load_sas.py:826) on the parent table exactly as today. -2. If `partition_by` is configured, also verify that the parent is LIST-partitioned on the same ordered keys. -3. Do not create any partitions. -4. Copy rows to the parent table and let Postgres route them. - -Accepted limitation for v1: - -- If the existing partition tree does not contain a leaf partition for some incoming value, Postgres will fail during `COPY` with a native partition-routing error. -- This design does not require preflight catalog validation of every leaf partition because that adds significant scope and catalog-parsing complexity. - -## 13. Dry-run behavior - -## 13.1 Single-file loader - -Current dry-run behavior in [`main()`](generic_loader/load_sas.py:1192) prints only one statement from [`render_create_table()`](generic_loader/load_sas.py:756). For partitioned configs it should change to: - -1. infer schema from the preview as today; -2. run full partition discovery over the file; -3. warn on stderr if `total_partition_tables > max_partitions`; -4. print the full ordered DDL statement list to stdout; -5. open no database connection. - -Output format recommendation: - -- print statements separated by one blank line for readability; -- do not print extra prose on stdout, so the output remains easy to paste into SQL tooling. - -## 13.2 Folder loader - -Current dry-run behavior in [`main()`](generic_loader/load_folder.py:496) prints one `CREATE TABLE` per cluster based on the first file only. For partitioned clusters it should change to: - -1. keep printing the discovered cluster summary; -2. for each loadable cluster, print a header such as `--- DDL for cluster 'group_a' ---`; -3. infer schema from the first file as today; -4. if the cluster is partitioned, scan all files in that cluster to build one shared `PartitionPlan`; -5. print the full ordered DDL statement list. - -Important documentation note: - -- Partitioned dry-runs are now full-data scans over the partition columns and can take much longer than non-partitioned dry-runs. - -## 14. Error handling - -The implementation should handle failures at the earliest safe point with clear messages. - -## 14.1 Config-time errors - -Raise `ValueError` from [`load_config()`](generic_loader/load_sas.py:350) or [`load_folder_config()`](generic_loader/load_folder.py:200) for: - -- `partition_by` not being a list -- empty or non-string items inside `partition_by` -- duplicate partition column names -- `max_partitions <= 0` -- `include` omitting a partition column -- `exclude` removing a partition column -- cluster config specifying an invalid override shape - -## 14.2 Runtime validation errors before DDL - -Raise `ValueError` with file/cluster context for: - -- partition column not present after filtering -- partition column absent from the inferred schema -- parent table name too long to create child suffixes safely -- a partition value that cannot be normalized or rendered into SQL - -## 14.3 Append-time compatibility errors - -Raise [`SchemaCompatibilityError`](generic_loader/load_sas.py:307) for: - -- parent column mismatch detected by [`_assert_schema_compatible()`](generic_loader/load_sas.py:826) -- existing parent not being partitioned when `partition_by` is configured -- existing parent using a partition strategy other than LIST -- existing parent using a different ordered key list - -## 14.4 Warning-only conditions - -Emit `[warn] ...` to stderr, but continue, for: - -- `total_partition_tables > max_partitions` -- existing warnings already emitted by [`_assert_schema_compatible()`](generic_loader/load_sas.py:826) - -Recommended warning message: - -```text -[warn] partition plan for public.customers will create 12,431 partition tables, exceeding max_partitions=10,000 -``` - -## 14.5 Postgres runtime errors left to bubble - -Do not swallow driver/database exceptions for: - -- DDL execution failures -- `COPY` failures caused by missing append-mode partitions -- any transaction failure during live loading - -The outer transaction handling in [`main()`](generic_loader/load_sas.py:1192) and [`main()`](generic_loader/load_folder.py:496) should remain responsible for rollback. - -## 15. Detailed single-file flow after the change - -```text -load_config --> read_sas_preview --> apply_column_filter --> infer_schema --> validate partition columns --> if validate flag: run manifest validation --> if partitioned and (dry-run or create needed): discover partition values from full file --> if dry-run: print full DDL and exit --> connect --> create_table (with partition metadata) --> copy_dataframes to parent table --> commit / rollback exactly as today -``` - -Notes: - -- A partitioned live load usually requires one preview read, one full discovery pass, and one full load pass. -- This is a deliberate tradeoff to ensure the full partition tree exists before any row is copied. - -## 16. Detailed folder flow after the change - -For each cluster in [`load_cluster()`](generic_loader/load_folder.py:385): - -```text -infer schema from first file preview --> validate partition columns --> if partitioned and creation is needed: discover partition values across all files in the cluster --> create_table (with partition metadata) --> stream every file to the parent table --> for later files, keep the existing append-mode schema compatibility check -``` - -Notes: - -- The partition plan is cluster-wide, not file-by-file. -- All files in the cluster must route into one shared partition tree under the same parent table. - -## 17. What remains unchanged - -- [`infer_schema()`](generic_loader/load_sas.py:637) keeps its current type-inference behavior. -- [`copy_dataframes()`](generic_loader/load_sas.py:1028) remains unchanged and still copies to the parent table. -- [`assert_schema_compatible()`](generic_loader/load_sas.py:874) remains the public wrapper for append compatibility. -- Non-partitioned configs should continue to produce exactly one `CREATE TABLE` statement and the same load behavior as today. - -## 18. Implementation sequencing - -Recommended implementation order: - -1. Extend config dataclasses and parsers. -2. Add partition parsing/validation helpers. -3. Add internal partition plan data structure. -4. Add partition discovery and literal/name helpers. -5. Extend DDL rendering. -6. Extend [`create_table()`](generic_loader/load_sas.py:890) and [`_drop_table()`](generic_loader/load_sas.py:771). -7. Wire the single-file flow. -8. Wire the folder flow and inheritance rules. -9. Update dry-run/help text and sample YAML files. - -## 19. QA and validation matrix - -The implementation should be validated against at least these scenarios: - -1. Non-partitioned single-file load still behaves exactly as before. -2. Single-level text partitioning creates one child per unique value. -3. Multi-level cascading partitioning scopes child values to their parent. -4. `NULL` partition values create `FOR VALUES IN (NULL)` partitions. -5. Text empty strings route to the `NULL` partition, not `''`. -6. Sanitization collision (`A-B` vs `A B`) resolves deterministically. -7. Very long child names truncate correctly and still remain unique. -8. `max_partitions` warning appears but the load continues. -9. `replace` drops the parent with `CASCADE` and recreates the full tree. -10. `append` rejects a parent with the wrong partition strategy or key order. -11. Folder-level `partition_by` is inherited by auto-detected clusters. -12. Explicit cluster `partition_by` overrides folder defaults. -13. Explicit cluster `partition_by: []` disables a folder default. -14. Dry-run prints the full DDL tree and opens no connection. -15. Partitioned folder dry-run scans all files in the cluster, not just the first one. - -## 20. Documentation updates required - -In addition to implementing the code, update: - -- [`generic_loader/sample_config.yaml`](generic_loader/sample_config.yaml) with `partition_by` and `max_partitions` comments and examples. -- [`generic_loader/sample_folder_config.yaml`](generic_loader/sample_folder_config.yaml) with folder defaults, cluster overrides, and explicit opt-out examples. -- The module-level usage text in [`load_sas.py`](generic_loader/load_sas.py) so dry-run docs mention full DDL for partitioned tables. -- The module-level usage text in [`load_folder.py`](generic_loader/load_folder.py) so dry-run docs mention cluster-wide partition discovery. - -## 21. Final design summary - -The safest low-regression approach is: - -1. keep the current schema inference path unchanged; -2. add a separate full-data partition discovery pass for partitioned loads; -3. render one parent `CREATE TABLE` plus recursive `PARTITION OF` child statements; -4. create or replace the full tree before copying any data; -5. leave [`copy_dataframes()`](generic_loader/load_sas.py:1028) unchanged so PostgreSQL handles routing; -6. keep `append` mode strict about parent compatibility and intentionally do not auto-create missing partitions. - -That approach satisfies the feature requirements while containing code churn to config parsing, DDL rendering, runtime planning, and folder integration. diff --git a/generic_loader/load_folder.py b/generic_loader/load_folder.py index bf2e619..5136fc1 100644 --- a/generic_loader/load_folder.py +++ b/generic_loader/load_folder.py @@ -130,11 +130,13 @@ from load_sas import ( assert_schema_compatible, connect, copy_dataframes, + create_indexes, create_table, discover_partition_values_chunked, infer_schema, iter_sas_chunks, read_sas_preview, + render_create_indexes, render_create_table, render_partition_ddl, ) @@ -152,8 +154,9 @@ SAS_EXTENSIONS = (".sas7bdat", ".xpt", ".xport") class ClusterSpec: """Resolved per-cluster load settings. - ``partition_by`` and ``max_partitions`` are resolved from the folder - defaults and any per-cluster overrides during :func:`discover_clusters`. + ``partition_by``, ``max_partitions``, and ``indexes`` are resolved from + the folder defaults and any per-cluster overrides during + :func:`discover_clusters`. """ tablename: str @@ -165,6 +168,7 @@ class ClusterSpec: pattern: Optional[str] = None partition_by: List[str] = field(default_factory=list) max_partitions: int = 10_000 + indexes: List[str] = field(default_factory=list) @dataclass @@ -174,6 +178,7 @@ class _ExplicitPattern: ``partition_by`` defaults to ``None`` meaning "inherit from folder level". An explicit empty list ``[]`` means "disable partitioning for this cluster". ``max_partitions`` defaults to ``None`` meaning "inherit from folder level". + ``indexes`` defaults to ``None`` meaning "inherit from folder level". """ pattern: re.Pattern @@ -184,14 +189,15 @@ class _ExplicitPattern: exclude: Optional[List[str]] = None partition_by: Optional[List[str]] = None max_partitions: Optional[int] = None + indexes: Optional[List[str]] = None @dataclass class FolderConfig: """Folder-level configuration parsed from YAML. - ``partition_by`` and ``max_partitions`` serve as defaults for every - cluster unless overridden at the cluster level. + ``partition_by``, ``max_partitions``, and ``indexes`` serve as defaults + for every cluster unless overridden at the cluster level. """ folder: Path @@ -203,6 +209,7 @@ class FolderConfig: explicit: List[_ExplicitPattern] = field(default_factory=list) partition_by: List[str] = field(default_factory=list) max_partitions: int = 10_000 + indexes: List[str] = field(default_factory=list) # --------------------------------------------------------------------------- @@ -312,6 +319,57 @@ def _validate_partition_vs_columns( ) +def _parse_indexes( + raw_value: Any, where: str, *, allow_none: bool = False +) -> Optional[List[str]]: + """Parse an ``indexes`` value from YAML. + + Returns a list of non-empty, unique column name strings. When + ``allow_none`` is True (used for per-cluster entries), an omitted key + returns ``None`` to signal "inherit from folder level". An explicit + empty list ``[]`` always returns ``[]``. + """ + if raw_value is None: + return None if allow_none else [] + if isinstance(raw_value, str): + if not raw_value.strip(): + raise ValueError(f"{where}: 'indexes' string must be non-empty.") + return [raw_value.strip()] + if isinstance(raw_value, list): + if len(raw_value) == 0: + return [] + result: List[str] = [] + for i, item in enumerate(raw_value): + if not isinstance(item, str) or not item.strip(): + raise ValueError( + f"{where}: 'indexes[{i}]' must be a non-empty string." + ) + result.append(str(item).strip()) + if len(result) != len(set(result)): + raise ValueError( + f"{where}: 'indexes' contains duplicate column names." + ) + return result + raise ValueError( + f"{where}: 'indexes' must be a string or list of strings." + ) + + +def _validate_indexes_vs_columns( + indexes: List[str], + exclude: Optional[List[str]], + where: str, +) -> None: + """Raise if any ``indexes`` column is in the ``exclude`` list.""" + if not indexes or exclude is None: + return + excluded_idx = [c for c in indexes if c in exclude] + if excluded_idx: + raise ValueError( + f"{where}: 'exclude' removes index columns: {excluded_idx}" + ) + + def load_folder_config(path: Path) -> FolderConfig: """Parse and validate the folder-level YAML config at ``path``. @@ -350,6 +408,10 @@ def load_folder_config(path: Path) -> FolderConfig: ) _validate_partition_vs_columns(partition_by, exclude, f"Config {path}") + # -- folder-level index settings ---------------------------------------- + indexes = _parse_indexes(raw.get("indexes"), f"Config {path}") + _validate_indexes_vs_columns(indexes, exclude, f"Config {path}") + explicit: List[_ExplicitPattern] = [] clusters_raw = raw.get("clusters") or [] if not isinstance(clusters_raw, list): @@ -384,6 +446,13 @@ def load_folder_config(path: Path) -> FolderConfig: effective_pb = c_partition_by if c_partition_by is not None else partition_by _validate_partition_vs_columns(effective_pb, effective_exclude, where) + # -- per-cluster index settings ------------------------------------- + c_indexes = _parse_indexes( + entry.get("indexes"), where, allow_none=True + ) + effective_idx = c_indexes if c_indexes is not None else indexes + _validate_indexes_vs_columns(effective_idx, effective_exclude, where) + explicit.append( _ExplicitPattern( pattern=compiled, @@ -394,6 +463,7 @@ def load_folder_config(path: Path) -> FolderConfig: exclude=c_exclude, partition_by=c_partition_by, max_partitions=c_max_partitions, + indexes=c_indexes, ) ) @@ -407,6 +477,7 @@ def load_folder_config(path: Path) -> FolderConfig: explicit=explicit, partition_by=partition_by, max_partitions=max_partitions, + indexes=indexes, ) @@ -482,6 +553,11 @@ def discover_clusters(cfg: FolderConfig) -> List[ClusterSpec]: patt.max_partitions if patt.max_partitions is not None else cfg.max_partitions ) + # Resolve indexes: None = inherit folder, [] = disable, list = override + resolved_idx = ( + patt.indexes if patt.indexes is not None + else cfg.indexes + ) matched = [f for f in remaining if patt.pattern.search(f.name)] if not matched: @@ -498,6 +574,7 @@ def discover_clusters(cfg: FolderConfig) -> List[ClusterSpec]: pattern=patt.raw_pattern, partition_by=resolved_pb, max_partitions=resolved_mp, + indexes=resolved_idx, ) ) continue @@ -513,6 +590,7 @@ def discover_clusters(cfg: FolderConfig) -> List[ClusterSpec]: pattern=patt.raw_pattern, partition_by=resolved_pb, max_partitions=resolved_mp, + indexes=resolved_idx, ) ) @@ -532,6 +610,7 @@ def discover_clusters(cfg: FolderConfig) -> List[ClusterSpec]: source="auto", partition_by=cfg.partition_by, max_partitions=cfg.max_partitions, + indexes=cfg.indexes, ) ) @@ -596,6 +675,17 @@ def load_cluster(conn, cluster: ClusterSpec, schemaname: str) -> int: first, *rest = cluster.files first_columns = _infer_cluster_schema(first, cluster.include, cluster.exclude) + # -- Validate index columns early --------------------------------------- + if cluster.indexes: + missing_icols = [ + c for c in cluster.indexes if c not in first_columns + ] + if missing_icols: + raise ValueError( + f"cluster {cluster.tablename!r}: indexes references " + f"columns not present in the inferred schema: {missing_icols}" + ) + # -- Partition support -------------------------------------------------- partition_values: Optional[dict] = None if cluster.partition_by: @@ -657,6 +747,10 @@ def load_cluster(conn, cluster: ClusterSpec, schemaname: str) -> int: cluster.include, cluster.exclude, ) + # -- Index support ------------------------------------------------------ + if cluster.indexes: + create_indexes(conn, schemaname, cluster.tablename, cluster.indexes) + return total @@ -733,9 +827,12 @@ def _describe_cluster(cluster: ClusterSpec) -> str: parts = "" if cluster.partition_by: parts = f"\n partition_by: {cluster.partition_by}" + idx = "" + if cluster.indexes: + idx = f"\n indexes: {cluster.indexes}" return ( f"cluster {cluster.tablename!r} [{src}] if_exists={cluster.if_exists}\n" - f" files: {files}{parts}" + f" files: {files}{parts}{idx}" ) @@ -812,6 +909,24 @@ def main(argv: Optional[List[str]] = None) -> int: for stmt in child_stmts: print() print(stmt) + # Print CREATE INDEX DDL when the cluster has indexes. + if c.indexes: + missing_icols = [ + col for col in c.indexes if col not in columns + ] + if missing_icols: + print( + f" [error] indexes references columns not in " + f"schema: {missing_icols}", + file=sys.stderr, + ) + else: + idx_stmts = render_create_indexes( + cfg.schemaname, c.tablename, c.indexes, + ) + for stmt in idx_stmts: + print() + print(stmt) print() return 0 diff --git a/generic_loader/load_sas.py b/generic_loader/load_sas.py index 1e0c963..74f5ff8 100644 --- a/generic_loader/load_sas.py +++ b/generic_loader/load_sas.py @@ -289,6 +289,7 @@ class LoaderConfig: exclude: Optional[List[str]] = None partition_by: List[str] = field(default_factory=list) max_partitions: int = 10_000 + indexes: List[str] = field(default_factory=list) @dataclass @@ -455,6 +456,50 @@ def load_config(path: Path) -> LoaderConfig: f"got {max_partitions}" ) + # -- indexes ------------------------------------------------------------ + raw_idx = raw.get("indexes") + if raw_idx is None or (isinstance(raw_idx, list) and len(raw_idx) == 0): + indexes: List[str] = [] + elif isinstance(raw_idx, str): + if not raw_idx.strip(): + raise ValueError(f"Config {path}: 'indexes' string must be non-empty.") + indexes = [raw_idx.strip()] + elif isinstance(raw_idx, list): + indexes = [] + for i, item in enumerate(raw_idx): + if not isinstance(item, str) or not item.strip(): + raise ValueError( + f"Config {path}: 'indexes[{i}]' must be a non-empty string." + ) + indexes.append(str(item).strip()) + if len(indexes) != len(set(indexes)): + raise ValueError( + f"Config {path}: 'indexes' contains duplicate column names." + ) + else: + raise ValueError( + f"Config {path}: 'indexes' must be a string or list of strings." + ) + + # Validate indexes vs include/exclude + if indexes: + inc_list = [str(c) for c in include] if include is not None else None + exc_list = [str(c) for c in exclude] if exclude is not None else None + if exc_list is not None: + excluded_idx = [c for c in indexes if c in exc_list] + if excluded_idx: + raise ValueError( + f"Config {path}: 'exclude' removes index columns: " + f"{excluded_idx}" + ) + if inc_list is not None: + missing_in_include = [c for c in indexes if c not in inc_list] + if missing_in_include: + raise ValueError( + f"Config {path}: 'include' omits index columns: " + f"{missing_in_include}" + ) + return LoaderConfig( filename=filename, schemaname=schemaname, @@ -464,6 +509,7 @@ def load_config(path: Path) -> LoaderConfig: exclude=[str(c) for c in exclude] if exclude is not None else None, partition_by=partition_by, max_partitions=max_partitions, + indexes=indexes, ) @@ -1515,6 +1561,81 @@ def _render_partition_ddl_recursive( statements.append(stmt) +# --------------------------------------------------------------------------- +# Index support +# --------------------------------------------------------------------------- + + +def render_create_indexes( + schema: str, + tablename: str, + indexes: List[str], +) -> List[str]: + """Generate ``CREATE INDEX IF NOT EXISTS`` DDL for each column in *indexes*. + + Each index is a simple B-tree index on a single column. The index name + follows the pattern ``ix_{tablename}_{column}`` (raw, unsanitized names + wrapped with :func:`_quote_ident`). The table reference is fully + qualified as ``schema.tablename``. + + If the generated index name exceeds PostgreSQL's 63-character identifier + limit, it is truncated and a short hash suffix is appended to preserve + uniqueness (similar to partition name truncation). + + Returns a list of SQL strings, one per index. + """ + stmts: List[str] = [] + for col in indexes: + idx_name = f"ix_{tablename}_{col}" + if len(idx_name) > _PG_IDENT_MAX_LEN: + # Truncate and append an 8-char hash for uniqueness. + name_hash = hashlib.sha256(idx_name.encode()).hexdigest()[:8] + # 9 = 1 underscore + 8 hash chars + truncated = idx_name[: _PG_IDENT_MAX_LEN - 9].rstrip("_") + idx_name = f"{truncated}_{name_hash}" + stmt = ( + f"CREATE INDEX IF NOT EXISTS {_quote_ident(idx_name)} " + f"ON {_qualified(schema, tablename)} ({_quote_ident(col)});" + ) + stmts.append(stmt) + return stmts + + +def create_indexes( + conn, + schema: str, + tablename: str, + indexes: List[str], +) -> None: + """Execute ``CREATE INDEX IF NOT EXISTS`` for each column in *indexes*. + + Calls :func:`render_create_indexes` to obtain the DDL, executes each + statement, commits immediately after each successful creation, and logs + progress to stderr. If an individual index creation fails (e.g. a name + collision unrelated to ``IF NOT EXISTS``), the transaction is rolled back + (affecting only the failed statement) and the remaining indexes are still + attempted. + """ + stmts = render_create_indexes(schema, tablename, indexes) + with conn.cursor() as cur: + for stmt, col in zip(stmts, indexes): + try: + cur.execute(stmt) + conn.commit() + print( + f"[info] created index ix_{tablename}_{col} " + f"on {schema}.{tablename}({col})", + file=sys.stderr, + ) + except Exception as exc: + conn.rollback() + print( + f"[warn] failed to create index ix_{tablename}_{col} " + f"on {schema}.{tablename}({col}): {exc}", + file=sys.stderr, + ) + + # --------------------------------------------------------------------------- # COPY loading # --------------------------------------------------------------------------- @@ -1817,6 +1938,15 @@ def main(argv: Optional[List[str]] = None) -> int: f"(filtered) schema: {missing_pcols}" ) + # Validate index columns exist in the schema after filtering. + if cfg.indexes: + missing_icols = [c for c in cfg.indexes if c not in columns] + if missing_icols: + raise ValueError( + f"indexes references columns not present in the " + f"(filtered) schema: {missing_icols}" + ) + if args.validate: manifest_path = cfg.filename.with_suffix("").with_suffix(".expected.json") # The above strips .xpt then appends .expected.json, e.g. @@ -1874,6 +2004,14 @@ def main(argv: Optional[List[str]] = None) -> int: for stmt in child_stmts: print() print(stmt) + # Print CREATE INDEX DDL if indexes are configured. + if cfg.indexes: + idx_stmts = render_create_indexes( + cfg.schemaname, cfg.tablename, cfg.indexes, + ) + for stmt in idx_stmts: + print() + print(stmt) return 0 # Release the preview frame before opening the stream - lets the GC reclaim @@ -1906,6 +2044,8 @@ def main(argv: Optional[List[str]] = None) -> int: conn, cfg.schemaname, cfg.tablename, _filtered_chunks(), columns ) conn.commit() + if cfg.indexes: + create_indexes(conn, cfg.schemaname, cfg.tablename, cfg.indexes) except Exception: conn.rollback() raise diff --git a/generic_loader/sample_config.yaml b/generic_loader/sample_config.yaml index 0f33604..c487769 100644 --- a/generic_loader/sample_config.yaml +++ b/generic_loader/sample_config.yaml @@ -29,3 +29,12 @@ if_exists: append # max_partitions: Warning threshold for total partition count (default: 10000). # If the number of partitions exceeds this, a warning is logged but loading continues. # max_partitions: 10000 + +# indexes: Create B-tree indexes on these columns after data loading. +# Indexes are created with IF NOT EXISTS for safe use with append mode. +# Single column: +# indexes: state +# Multiple columns (one index per column): +# indexes: +# - state +# - zip diff --git a/generic_loader/sample_folder_config.yaml b/generic_loader/sample_folder_config.yaml index fb03eb8..066d840 100644 --- a/generic_loader/sample_folder_config.yaml +++ b/generic_loader/sample_folder_config.yaml @@ -45,6 +45,16 @@ auto_detect: true # (default: 10000). Inherited by all clusters unless overridden per-cluster. # max_partitions: 10000 +# Folder-level indexes: Create B-tree indexes on these columns after data +# loading. Inherited by all clusters unless overridden per-cluster. +# Indexes are created with IF NOT EXISTS for safe use with append mode. +# Single column: +# indexes: state +# Multiple columns (one index per column): +# indexes: +# - state +# - zip + # Explicit cluster patterns. Each pattern is matched against the file # *basename*. Files matched by a pattern are pulled out of the auto-detect # pool, so explicit and auto clusters compose cleanly. @@ -72,6 +82,16 @@ clusters: # - year # max_partitions: 500 + # Per-cluster indexes override. Takes precedence over the folder-level + # indexes default above. An explicit empty list disables indexing for + # this cluster even when the folder default has indexes. + # + # - pattern: '^group_d\d+\.xpt$' + # tablename: group_d + # indexes: + # - region + # - year + # With only the gq pattern explicit, auto_detect: true will still bucket # group_b1.xpt + group_b2.xpt into a "group_b" cluster and the lone # standalone.xpt into a "standalone" cluster. See generate_sample_folder.py