Data Jobs¶

import duckdb
import polars as pl

# Query Polars DataFrame with SQL
df = pl.read_parquet("data.parquet")
result = duckdb.sql("SELECT * FROM df WHERE value > 100").pl()

# Polars reads partitioned data efficiently
df = pl.scan_parquet("gs://bucket/data/**/*.parquet")
    .filter(pl.col("year") == 2024)
    .filter(pl.col("month") == 1)
    .collect()

# Pseudocode pattern with performance optimization
def process_with_checkpoint(partitions: list[str], checkpoint_path: str):
    """Process partitions with checkpoint tracking.

    Performance improvement:
    - Convert completed list to set for O(1) lookup instead of O(n)
    - This is critical when processing thousands of partitions
    - Example: 10,000 partitions goes from ~50M operations to ~10K operations
    """
    completed_list = load_checkpoint(checkpoint_path)
    # Convert to set for O(1) lookup performance
    completed = set(completed_list)

    for partition in partitions:
        if partition in completed:
            logger.info(f"Skipping {partition}, already completed")
            continue

        process_partition(partition)
        # Add to set and save
        completed.add(partition)
        save_checkpoint(checkpoint_path, partition)

from functools import lru_cache
from google.cloud import bigquery

@lru_cache(maxsize=1)
def get_bigquery_client(project: str) -> bigquery.Client:
    """Get cached BigQuery client instance.

    Performance: Client is created once and reused across calls.
    Saves ~100-500ms per call depending on service.
    """
    return bigquery.Client(project=project)

# Use in your code
def query_data(sql: str, project: str):
    client = get_bigquery_client(project)
    return client.query(sql).to_dataframe()

# Bad: O(n) lookup for each of m items = O(n*m)
completed_files = ["file1.csv", "file2.csv", ...]  # Could be 10,000+ items
for file in all_files:
    if file in completed_files:  # Linear search each time!
        continue

# Good: O(1) lookup for each of m items = O(n+m)
completed_files_set = set(completed_files)  # One-time O(n) conversion
for file in all_files:
    if file in completed_files_set:  # Constant time lookup!
        continue

# Performance gain: 10,000 files = ~50M operations → ~10K operations

from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
import logging

logger = logging.getLogger(__name__)

def process_files_efficiently(files: list[Path], max_workers: int = 4):
    """Process multiple files in parallel.

    Performance improvements:
    - Parallel processing using ThreadPoolExecutor
    - Batch size limits memory usage
    - Progress tracking without blocking
    """
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit all tasks
        futures = {executor.submit(process_file, f): f for f in files}

        # Process results as they complete
        for future in as_completed(futures):
            file = futures[future]
            try:
                result = future.result()
                logger.info(f"Processed {file}: {result}")
            except Exception as e:
                logger.error(f"Failed {file}: {e}")

# Polars handles most CSV edge cases
df = pl.read_csv(
    "file.csv",
    separator=";",
    encoding="utf-8",
    null_values=["", "NULL", "N/A"],
    try_parse_dates=True,
)

# Polars reads gzipped files directly
df = pl.read_csv("data.csv.gz")
df = pl.read_parquet("data.parquet.gz")

# For zip archives
import zipfile
with zipfile.ZipFile("archive.zip") as z:
    with z.open("data.csv") as f:
        df = pl.read_csv(f)

from pathlib import Path

# Build paths
data_dir = Path("/data")
input_file = data_dir / "input" / "file.csv"

# Glob patterns
parquet_files = data_dir.glob("**/*.parquet")

# File info
file.name       # "file.csv"
file.stem       # "file"
file.suffix     # ".csv"
file.parent     # Path("/data/input")

# Create directories
output_dir.mkdir(parents=True, exist_ok=True)

# Read/write
content = file.read_text(encoding="utf-8")
file.write_bytes(data)

import tempfile
from pathlib import Path

# Context manager cleans up automatically
with tempfile.TemporaryDirectory() as tmp:
    tmp_path = Path(tmp)
    output = tmp_path / "result.parquet"
    df.write_parquet(output)
    # Upload to GCS, etc.
# Directory deleted automatically

# Polars reads directly from GCS
df = pl.read_parquet("gs://bucket/path/file.parquet")
df = pl.scan_parquet("gs://bucket/path/**/*.parquet")

# For more control
from google.cloud import storage
client = storage.Client()
bucket = client.bucket("my-bucket")
blob = bucket.blob("path/file.csv")
blob.download_to_filename("/tmp/file.csv")

import paramiko
from pathlib import Path
from tenacity import retry, stop_after_attempt, wait_exponential

def download_from_sftp(
    host: str,
    username: str,
    key_path: Path,
    remote_path: str,
    local_path: Path,
    timeout: float = 30.0,
    verify_size: bool = True,
) -> None:
    """Download file from SFTP server with retry logic and proper timeouts.

    Performance improvements:
    - Configurable timeout prevents hanging connections
    - File size validation ensures complete download (can be disabled for text mode)
    - Proper nested try/finally ensures resource cleanup

    Note: tenacity is listed in the Resilience tools section above.

    Args:
        verify_size: If True, validates downloaded file size matches remote.
                     Set to False if transferring text files that may have
                     line ending conversions (CRLF <-> LF).
    """
    key = paramiko.RSAKey.from_private_key_file(str(key_path))

    transport = paramiko.Transport((host, 22))
    transport.connect(username=username, pkey=key, timeout=timeout)

    try:
        sftp = paramiko.SFTPClient.from_transport(transport)
        try:
            # Get remote file size for validation (optional)
            if verify_size:
                remote_size = sftp.stat(remote_path).st_size

            sftp.get(remote_path, str(local_path))

            # Verify download completed successfully
            if verify_size:
                local_size = local_path.stat().st_size
                if local_size != remote_size:
                    raise ValueError(
                        f"Downloaded size mismatch: {local_size} != {remote_size}. "
                        "This may occur with text mode transfers and line ending differences. "
                        "Set verify_size=False if this is expected."
                    )
        finally:
            sftp.close()
    finally:
        transport.close()


@retry(
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=2, max=10),
)
def download_from_sftp_with_retry(
    host: str,
    username: str,
    key_path: Path,
    remote_path: str,
    local_path: Path,
) -> None:
    """Wrapper with automatic retry for transient failures."""
    download_from_sftp(host, username, key_path, remote_path, local_path)

import httpx
from pathlib import Path

def download_file(
    url: str,
    output: Path,
    chunk_size: int = 1024 * 1024,  # 1MB chunks
    timeout: float = 30.0,
) -> None:
    """Download file with streaming and optimized performance.

    Performance improvements:
    - Configurable chunk size (default 1MB) for efficient memory usage
    - Explicit timeout prevents hanging on slow connections
    - Buffer writes for better I/O performance
    - Content-Length validation ensures complete download

    Args:
        url: URL to download from
        output: Local path to save file
        chunk_size: Size of chunks to download (default 1MB)
        timeout: Request timeout in seconds (default 30s)
    """
    with httpx.stream("GET", url, timeout=timeout, follow_redirects=True) as response:
        response.raise_for_status()

        # Get expected file size if available
        expected_size = int(response.headers.get("content-length", 0))
        downloaded = 0

        with output.open("wb") as f:
            for chunk in response.iter_bytes(chunk_size=chunk_size):
                f.write(chunk)
                downloaded += len(chunk)

        # Verify download if content-length was provided
        if expected_size > 0 and downloaded != expected_size:
            output.unlink()  # Delete incomplete file
            raise ValueError(
                f"Download incomplete: {downloaded} bytes received, "
                f"expected {expected_size} bytes"
            )

from pydantic_settings import BaseSettings
from pydantic import Field, SecretStr

class Settings(BaseSettings):
    # Required
    gcp_project: str
    input_bucket: str
    output_bucket: str

    # With defaults
    log_level: str = "INFO"
    dry_run: bool = False
    batch_size: int = 10000

    # Secrets (masked in logs)
    database_url: SecretStr
    api_key: SecretStr

    model_config = {
        "env_file": ".env",
        "env_file_encoding": "utf-8",
    }

settings = Settings()

from google.cloud import secretmanager
from functools import lru_cache

# Cache the client instance to avoid repeated initialization
@lru_cache(maxsize=1)
def _get_secret_manager_client() -> secretmanager.SecretManagerServiceClient:
    """Get a cached Secret Manager client.

    Performance improvement:
    - Client is created once and reused across calls
    - Reduces connection overhead and initialization time
    - Thread-safe with lru_cache
    """
    return secretmanager.SecretManagerServiceClient()


def get_secret(secret_id: str, project_id: str, version: str = "latest") -> str:
    """Fetch secret from GCP Secret Manager with client reuse.

    Performance improvements:
    - Reuses client instance instead of creating new one each call
    - Configurable version for flexibility
    - Proper error handling for missing secrets

    Args:
        secret_id: The secret identifier
        project_id: GCP project ID
        version: Secret version (default: "latest")

    Returns:
        The secret value as a string

    Raises:
        google.api_core.exceptions.NotFound: If secret doesn't exist
    """
    client = _get_secret_manager_client()
    name = f"projects/{project_id}/secrets/{secret_id}/versions/{version}"
    response = client.access_secret_version(request={"name": name})
    return response.payload.data.decode("utf-8")