fix(ingestion): materialize TABLESAMPLE into temp tables for Snowflake SQLAlchemy profiler#18253
fix(ingestion): materialize TABLESAMPLE into temp tables for Snowflake SQLAlchemy profiler#18253sgomezvillamor wants to merge 21 commits into
Conversation
… profiler Root cause: SQLAlchemy profiler queries massive production tables directly, causing memory exhaustion. A 102B row table triggered 40-CTE cross-join queries that hung the process and spiked memory to 10GB+. Fix: Implement temp table sampling matching GE profiler's behavior: - Create temporary tables with TABLESAMPLE for large tables (>sample_size rows) - Use conditional sampling strategy matching GE exactly: * BLOCK + BERNOULLI for very large tables (>50M rows) * BERNOULLI only for smaller tables - Cache row counts to avoid redundant COUNT(*) queries - Clean up temp tables after profiling Config-driven via existing use_sampling (default: True) and sample_size (default: 10,000 rows) settings. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
github-actions
Bot
added
the
ingestion
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Connector Tests ResultsAll connector tests passed for commit To skip connector tests, add the Autogenerated by the connector-tests CI pipeline. |
- Replace SELECT COUNT(*) with INFORMATION_SCHEMA.TABLES.ROW_COUNT - Avoids full table scans on billion-row tables (instant vs minutes/timeout) - Default to sampling when row count unavailable (conservative approach) - Add warning when metadata query fails - No additional permissions required (same as existing connector usage) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
- Log execution time for combined queries (INFO level) - Log execution time for fallback queries (DEBUG level) - Shows number of queries combined in each batch - Helps diagnose slow profiling operations Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
…ution timing Changes to snowflake.py: - Use INFORMATION_SCHEMA.TABLES.ROW_COUNT instead of SELECT COUNT(*) for instant metadata lookup (no full table scan) - Default to sampling when row count is unavailable (conservative approach) - Add warning report when row count query fails Changes to sqlalchemy_query_combiner.py: - Add PerfTimer for query execution timing - Add 5-char random query IDs to correlate log traces - Log query execution before/after with timing and query ID - Use INFO level for successful queries, WARNING for failures Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
Filter by CURRENT_DATABASE() in INFORMATION_SCHEMA.TABLES query to avoid ambiguity when multiple databases have the same schema.table combination. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
Snowflake GE profiler was setting custom_sql for sampling which prevented SQLAlchemy profiler from using its adapter-based temp table sampling. Now custom_sql is only set for method=ge, and the Snowflake adapter properly creates session-scoped temp tables with quoted names for case preservation. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
Switches Snowflake SQLAlchemy profiler from temp tables to inline TABLESAMPLE queries, matching the proven GE profiler approach. Previous temp table implementation had schema prefix issues and deviated from GE profiler's battle-tested sampling strategy. GE uses inline TABLESAMPLE queries for Snowflake (temp tables are only for BigQuery). Changes: - Removed temp table creation logic from SnowflakeAdapter - Added _create_table_from_custom_sql() to wrap TABLESAMPLE as subquery - Removed method == "ge" check in snowflake_profiler.py to enable custom_sql for SQLAlchemy profiler - Added architectural comments explaining design decisions This approach: - Uses snowflake_profiler.py to decide on sampling (consistent with GE) - Passes TABLESAMPLE query via context.custom_sql - Avoids schema prefix and connection lifecycle issues with temp tables - Maintains backwards compatibility with GE profiler Future consideration: Evaluate temp tables for custom_sql like BigQuery and Athena adapters, but would require handling Snowflake's session-scoped temp table model and schema prefix complexity. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
…pdate type signatures - Add SEED parameter to TABLESAMPLE queries for deterministic sampling - Generates seed from hash(table_identifier) to ensure consistency - All profiling queries now operate on the same sample data - Prevents non-coherent metrics (e.g., min from one sample > max from another) - Update type signatures from sa.Table to sa.sql.FromClause - Allows Snowflake adapter to use inline Subquery for TABLESAMPLE - Updated ProfilingContext.sql_table to accept both Table and Subquery - Updated all method signatures in sqlalchemy_profiler, adapters, and query_combiner_runner - Use getattr() for .schema/.name access where FromClause doesn't guarantee these attributes Testing: - All unit tests passing (4/4 Snowflake adapter tests) - Integration tests show TABLESAMPLE with SEED working correctly - All linting and mypy checks passing Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
…e SQLAlchemy profiler Replace the inline subquery approach with temp table materialization. The previous approach re-evaluated TABLESAMPLE on every CTE in every combined query, causing 70-800x slowdown vs GE. Now the sample is materialized once into a session-scoped temp table and all profiling queries run against that small table. Key changes: - Gate custom_sql generation to GE-only in snowflake_profiler.py - Snowflake adapter creates temp table via CREATE OR REPLACE TEMPORARY TABLE - Two-tier BLOCK+BERNOULLI for very large tables, BERNOULLI-only fallback for views - No SEED needed (materialized once, no cross-query determinism needed) - Revert SEED from GE path (wasn't on master, added only for consistency) - Conservative sampling when INFORMATION_SCHEMA row count unavailable - Revert sa.sql.FromClause back to sa.Table across all adapter signatures - Revert defensive getattr patterns back to direct attribute access Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
![aikido-pr-checks[bot]](https://avatars.githubusercontent.com/in/898896?s=60&v=4){kind=small}
sgomezvillamor
changed the title
…p table sampling Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
maggiehays
added
the
needs-review
| executed_sql = str(mock_conn.execute.call_args[0][0]) | ||
| assert "CREATE OR REPLACE TEMPORARY TABLE" in executed_sql | ||
| assert "dh_sample_" in executed_sql | ||
| assert "TABLESAMPLE BERNOULLI" in executed_sql |
There was a problem hiding this comment.
Should fix: this test runs the real sampling arithmetic but never checks the computed percentages — only "TABLESAMPLE BERNOULLI" in sql. Drop the * 100, change a constant, or emit a >100% fraction (which Snowflake rejects) and this stays green. Since the percentage math is the whole point of the PR, assert the exact substring for row_count=100_000 / sample_size=10_000 → TABLESAMPLE BERNOULLI (10.00000000), and in test_sampled_temp_table_block_bernoulli_sql assert BLOCK (10.00000000) + BERNOULLI (0.10000000). Ideally also add a guard that block_pc <= 100.
There was a problem hiding this comment.
Good call — addressed in a4c6143:
- Tests now assert exact percentage substrings (
TABLESAMPLE BERNOULLI (10.00000000),BLOCK (10.00000000)+BERNOULLI (0.10000000)) so any math regression will be caught. - Added
min(block_pc, 100)as a defensive cap. Theuse_block_presampleguard (row_count > sample_size * overgeneration_factor) already guaranteesblock_pc < 100, but the explicit cap makes the invariant self-documenting.
| return context | ||
|
|
||
| # Determine if sampling is needed | ||
| row_count = self._get_row_count_from_metadata(context, conn) |
There was a problem hiding this comment.
Should fix (perf, on-theme): this issues a fresh INFORMATION_SCHEMA round-trip per table, but the value is already known upstream — the GE path gates on table.rows_count (fetched during schema crawl), and ProfilingContext.row_count exists as a field but is never populated. BigQueryAdapter._setup_sampling already shows the intended pattern: row_count = context.row_count or self._get_quick_row_count(...). Thread table.rows_count into the context and fall back to _get_row_count_from_metadata only when it's None — removes one metadata query per profiled table.
There was a problem hiding this comment.
Good catch — addressed in 726b76d:
snowflake_profiler.py: now passes"row_count": table.rows_countinget_batch_kwargssqlalchemy_profiler.py: threadsrow_countthrough_generate_single_profile→ProfilingContextsnowflake.pyadapter:row_count = context.row_count or self._get_row_count_from_metadata(context, conn)— skips the INFORMATION_SCHEMA query when the count is already known from the schema crawl- Added
test_setup_profiling_uses_context_row_countthat asserts_get_row_count_from_metadatais not called whencontext.row_countis pre-populated
| # custom_sql is or may still be used for some use cases (not non-sampling). | ||
| # Materialize it into a temp table so profiling queries run against a small | ||
| # concrete table, same as what GX does internally. | ||
| if context.custom_sql: |
There was a problem hiding this comment.
Suggestion: since custom_sql is now gated GE-only in snowflake_profiler.py, nothing populates context.custom_sql on the SQLAlchemy Snowflake path — this branch is only reachable from its unit test today. Either remove it until a real producer exists, or add a comment naming the future producer (e.g. anticipated partition support), so it doesn't read as a live capability. Related: _create_temp_table_from_custom_sql sets is_sampled=True for a filtered query that isn't a statistical sample, which would mislabel the profile as SAMPLE — worth fixing if/when this path goes live.
There was a problem hiding this comment.
Confirmed — custom_sql is indeed GE-only (gated by method == "ge" in snowflake_profiler.py). Replaced the dead handling branch with an assert in 3e90221:
if context.custom_sql:→assert not context.custom_sqlwith a message explaining it's GE-only- Removed the now-unreachable
_create_temp_table_from_custom_sqlmethod - Updated the test to verify the assert fires when
custom_sqlis passed
This also sidesteps the is_sampled=True mislabeling issue you flagged — since the method no longer exists.
| block_profiling_min_rows = 100 * estimated_block_row_count | ||
| overgeneration_factor = 1000 | ||
|
|
||
| tablename = f'"{context.schema}"."{context.table}"' |
There was a problem hiding this comment.
Minor (pre-existing pattern): "{schema}"."{table}" isn't escaped for embedded double-quotes — a legitimately quoted Snowflake identifier containing a " would break out of the quoting. Not introduced by this PR (the same pattern exists elsewhere in the profiler) and not what aikido flagged, but worth a .replace('"', '""') on each part while you're here.
There was a problem hiding this comment.
Fixed for the SQLAlchemy path in 1023e67:
- Made
db_nameoptional inget_quoted_identifier_for_table(returns"schema"."table"whenNone) - SQLAlchemy adapter now uses
SnowflakeIdentifierBuilder.get_quoted_identifier_for_table(db_name=None, ...)instead of raw f-string quoting - Added two tests for the
db_name=Nonevariant inTestSnowflakeIdentifierQuoting
Not fixing the GE path (snowflake_profiler.py:120) since GE profiler is being removed soon.
| self.report.combined_queries_issued += 1 | ||
| sa_res = _sa_execute_underlying_method(queue_item.conn, combined_query) | ||
| logger.info( | ||
| f"[{query_id}] Executing combined query ({len(pending_queue)} queries combined): {str(combined_query)}" |
There was a problem hiding this comment.
Suggestion (log noise): the [query_id] correlation + PerfTimer are genuinely useful, but logging the full str(combined_query) at INFO on every combined query (and the fallback equivalent) is a big verbosity jump for a utility used across all SQL profilers — a large run emits one long SQL line per batch at normal verbosity. Consider keeping the SQL body at DEBUG and logging only the [query_id] + timing at INFO.
There was a problem hiding this comment.
Addressed in a1a7d8b — SQL body moved to DEBUG, query_id + count/timing stays at INFO for both the combined and fallback paths.
maggiehays
added
pending-submitter-merge
and removed
needs-review
- Assert exact percentage substrings (e.g. BERNOULLI (10.00000000)) instead of just checking the keyword is present, so math regressions are caught. - Add min(block_pc, 100) cap as a defensive guard against >100% fractions. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…ke adapter custom_sql is GE-only (gated by `method == "ge"` in snowflake_profiler.py), so the SQLAlchemy adapter should never receive it. Replace the handling branch with an assert that documents this invariant, and remove the now-unreachable _create_temp_table_from_custom_sql method. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…p INFORMATION_SCHEMA query Thread table.rows_count (from the schema crawl) through batch_kwargs → _generate_single_profile → ProfilingContext.row_count, so the Snowflake adapter can skip the redundant INFORMATION_SCHEMA round-trip when the count is already known. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…wn and conservative sampling kicks in Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…E identifiers Use SnowflakeIdentifierBuilder.get_quoted_identifier_for_table in the SQLAlchemy adapter instead of raw f-string quoting. Made db_name optional in get_quoted_identifier_for_table since the adapter path only has schema + table (no database). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
… sampling fallback Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…wflake adapter Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Keep query_id + timing at INFO, move full SQL text to DEBUG to reduce log noise on large profiling runs. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…tifier_for_table call Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…in CTE batching tests Replace fragile caplog-based SQL assertions with before_cursor_execute event listeners that capture actual SQL sent to the database. Applied to both test_queries_are_batched_with_multiple_ctes and test_strategic_batching_with_multiple_flush_points. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
3 participants
Problem
The Snowflake SQLAlchemy profiler is 70-800x slower than the GE profiler on large tables because it re-evaluates
TABLESAMPLE BERNOULLI(...)on every CTE in every combined query (the query combiner batches multiple scalar queries into a single cross-joined CTE query, and each CTE re-scans the source table). GE avoids this by materializing the sample into a temp table once, then querying the small temp table.Solution
Materialize the TABLESAMPLE sample into a session-scoped Snowflake temporary table once, then run all profiling queries against that small materialized table. This matches what GE/GX does internally when given a
custom_sqlbatch kwarg.Key Changes
1. Gate
custom_sqlto GE-only (snowflake_profiler.py)custom_sqlTABLESAMPLE query is now only generated whenprofiling.method == "ge"custom_sql, so all its profiling queries run against that materialized tablecustom_sqlisNone— the adapter handles sampling directly2. Temp table materialization (
adapters/snowflake.py)setup_profiling()queriesINFORMATION_SCHEMA.TABLESfor instant row count (parameterized query)CREATE OR REPLACE TEMPORARY TABLE dh_sample_<uuid> AS SELECT * FROM ... TABLESAMPLE BERNOULLI(...)TABLESAMPLE BERNOULLIto work on viewsINFORMATION_SCHEMArow count is unavailable, assumes table is large and samples anywaycustom_sqlpath materializes into a temp table too (same as what GX does internally)3. Partition spec fix (
sqlalchemy_profiler.py)elif custom_sql:toelif custom_sql or context.is_sampled:so the partition spec is set correctly when the adapter performs sampling via temp table4. Query combiner observability (
sqlalchemy_query_combiner.py)PerfTimerinstrumentation and short query IDs to combined and fallback query execution5. Others
custom_sqlfield inProfilingContextadapters/__init__.py)Testing
./gradlew :metadata-ingestion:lintFix— passes./gradlew :metadata-ingestion:testQuick— passesconnector-tests):test_profiling,test_profiling_ge,test_profiling_table_level,test_profiling_table_level_ge— all passNotes
use_sampling=True(default),sample_size=10000(default)