Writers

Base

geoembed.io.writers.base

Abstract base class for embedding writers.

EmbeddingWriter

Bases: ABC

Abstract interface for persisting embedding results.

Implementations handle different storage backends (Parquet, Delta Lake, etc.) while presenting a uniform interface to the pipeline.

write_batch(chip_ids, embeddings) abstractmethod

Write a batch of embeddings to storage.

Parameters:

Name	Type	Description	Default
chip_ids	list[str]	List of chip identifiers.	required
embeddings	ndarray	Array of shape (N, D) where N = len(chip_ids) and D = embedding dim.	required

Source code in src/geoembed/io/writers/base.py

@abstractmethod
def write_batch(self, chip_ids: list[str], embeddings: np.ndarray) -> None:
    """
    Write a batch of embeddings to storage.

    Args:
        chip_ids: List of chip identifiers.
        embeddings: Array of shape (N, D) where N = len(chip_ids) and D = embedding dim.
    """
    ...

finalize() abstractmethod

Perform any cleanup or optimisation after all batches have been written.

E.g., Z-ordering a Delta table, closing file handles, flushing buffers.

Source code in src/geoembed/io/writers/base.py

@abstractmethod
def finalize(self) -> None:
    """
    Perform any cleanup or optimisation after all batches have been written.

    E.g., Z-ordering a Delta table, closing file handles, flushing buffers.
    """
    ...

Parquet Writer

geoembed.io.writers.parquet_writer

Parquet-based embedding writer for local and cloud storage.

ParquetWriter(path, mode='overwrite')

Bases: EmbeddingWriter

Writes embeddings to a local Parquet file.

Accumulates batches in memory and writes once on finalize(), or appends to an existing file if mode="append".

Suitable for local development and small-to-medium datasets.

Initialise the Parquet writer.

Parameters:

Name	Type	Description	Default
path	str | Path	Output file path.	required
mode	str	"overwrite" replaces the file, "append" adds to existing.	'overwrite'

Source code in src/geoembed/io/writers/parquet_writer.py

def __init__(self, path: str | Path, mode: str = "overwrite"):
    """
    Initialise the Parquet writer.

    Args:
        path: Output file path.
        mode: "overwrite" replaces the file, "append" adds to existing.
    """
    self.path = Path(path)
    self.mode = mode
    self._batches: list[pd.DataFrame] = []

write_batch(chip_ids, embeddings)

Buffer a batch of embeddings.

Source code in src/geoembed/io/writers/parquet_writer.py

def write_batch(self, chip_ids: list[str], embeddings: np.ndarray) -> None:
    """Buffer a batch of embeddings."""
    df = pd.DataFrame({
        "chip_id": chip_ids,
        "embedding": list(embeddings),
    })
    self._batches.append(df)

finalize()

Write all buffered batches to the Parquet file.

Source code in src/geoembed/io/writers/parquet_writer.py

def finalize(self) -> None:
    """Write all buffered batches to the Parquet file."""
    if not self._batches:
        return

    combined = pd.concat(self._batches, ignore_index=True)

    if self.mode == "append" and self.path.exists():
        existing = pd.read_parquet(self.path)
        combined = pd.concat([existing, combined], ignore_index=True)

    self.path.parent.mkdir(parents=True, exist_ok=True)
    combined.to_parquet(self.path, index=False)
    print(f"Wrote {len(combined)} embeddings to {self.path}")
    self._batches.clear()

Delta Writer

geoembed.io.writers.delta_writer

Delta Lake embedding writer for Databricks / Unity Catalog.

DeltaWriter(table_name, mode='append', optimize_write=True, cluster_columns=None)

Bases: EmbeddingWriter

Writes embeddings to a Delta Lake table via Spark.

Designed for production use on Databricks with Unity Catalog governance. Supports append mode and optional liquid clustering.

Requires: geoembed[spark] extras.

Initialise the Delta writer.

Parameters:

Name	Type	Description	Default
table_name	str	Full Unity Catalog table name (e.g., "catalog.schema.table").	required
mode	str	Spark write mode ("append" or "overwrite").	'append'
optimize_write	bool	Enable Delta's optimizeWrite for auto-compaction.	True
cluster_columns	list[str] | None	Columns for liquid clustering (e.g., ["chip_id"]). Liquid clustering is incremental and automatic — no manual OPTIMIZE runs needed after initial setup.	None

Source code in src/geoembed/io/writers/delta_writer.py

def __init__(
    self,
    table_name: str,
    mode: str = "append",
    optimize_write: bool = True,
    cluster_columns: list[str] | None = None,
):
    """
    Initialise the Delta writer.

    Args:
        table_name: Full Unity Catalog table name (e.g., "catalog.schema.table").
        mode: Spark write mode ("append" or "overwrite").
        optimize_write: Enable Delta's optimizeWrite for auto-compaction.
        cluster_columns: Columns for liquid clustering (e.g., ["chip_id"]).
            Liquid clustering is incremental and automatic — no manual
            OPTIMIZE runs needed after initial setup.
    """
    self.table_name = table_name
    self.mode = mode
    self.optimize_write = optimize_write
    self.cluster_columns = cluster_columns
    self._spark = None

write_batch(chip_ids, embeddings)

Write a batch of embeddings directly to Delta.

Source code in src/geoembed/io/writers/delta_writer.py

def write_batch(self, chip_ids: list[str], embeddings: np.ndarray) -> None:
    """Write a batch of embeddings directly to Delta."""
    from geoembed.core.naming import quote_table_name

    spark = self._get_spark()

    df = pd.DataFrame({
        "chip_id": chip_ids,
        "embedding": list(embeddings.astype(float)),
    })

    sdf = spark.createDataFrame(df)

    writer = sdf.write.format("delta").mode(self.mode)
    if self.optimize_write:
        writer = writer.option("optimizeWrite", "true")
    writer.saveAsTable(quote_table_name(self.table_name))

finalize()

Apply liquid clustering if configured.

Source code in src/geoembed/io/writers/delta_writer.py

def finalize(self) -> None:
    """Apply liquid clustering if configured."""
    if self.cluster_columns:
        from geoembed.core.naming import quote_table_name

        spark = self._get_spark()
        quoted = quote_table_name(self.table_name)
        cols = ", ".join(self.cluster_columns)
        spark.sql(f"ALTER TABLE {quoted} CLUSTER BY ({cols})")
        print(f"Liquid clustering enabled on {self.table_name} by ({cols})")