apache_beam.ml.rag.ingestion.mysql_common module

apache_beam.ml.rag.ingestion.mysql_common.chunk_embedding_fn(chunk: Chunk) → str

Convert embedding to MySQL vector string format.

Formats dense embedding as a MySQL-compatible vector string. Example: [1.0, 2.0] -> ‘[1.0,2.0]’

Parameters:

chunk – Input Chunk object.

Returns:

MySQL vector string representation of the embedding.

Return type:

str

Raises:

ValueError – If chunk has no dense embedding.

class apache_beam.ml.rag.ingestion.mysql_common.ColumnSpec(column_name: str, python_type: Type, value_fn: Callable[[Chunk], Any], placeholder: str = '?')

Bases: object

Specification for mapping Chunk fields to MySQL columns for insertion.

Defines how to extract and format values from Chunks into MySQL database columns, handling the full pipeline from Python value to SQL insertion.

The insertion process works as follows: - value_fn extracts a value from the Chunk and formats it as needed - The value is stored in a NamedTuple field with the specified python_type - During SQL insertion, the value is bound to a ? placeholder

column_name

The column name in the database table.

Type:

str

python_type

Python type for the NamedTuple field that will hold the value. Must be compatible with RowCoder.

Type:

Type

value_fn

Function to extract and format the value from a Chunk. Takes a Chunk and returns a value of python_type.

Type:

Callable[[apache_beam.ml.rag.types.Chunk], Any]

placeholder

Optional placeholder to apply typecasts or functions to value ? placeholder e.g. “string_to_vector(?)” for vector columns.

Type:

str

Examples

Basic text column (uses standard JDBC type mapping):

>>> ColumnSpec.text(
...     column_name="content",
...     value_fn=lambda chunk: chunk.content.text
... )
... # Results in: INSERT INTO table (content) VALUES (?)

Timestamp from metadata:

>>> ColumnSpec(
...     column_name="created_at",
...     python_type=str,
...     value_fn=lambda chunk: chunk.metadata.get("timestamp")
... )
... # Results in: INSERT INTO table (created_at) VALUES (?)

Factory Methods:

text: Creates a text column specification. integer: Creates an integer column specification. float: Creates a float column specification. vector: Creates a vector column specification with string_to_vector(). json: Creates a JSON column specification.

column_name: str

python_type: Type

value_fn: Callable[[Chunk], Any]

placeholder: str = '?'

classmethod text(column_name: str, value_fn: Callable[[Chunk], Any]) → ColumnSpec

Create a text column specification.

classmethod integer(column_name: str, value_fn: Callable[[Chunk], Any]) → ColumnSpec

Create an integer column specification.

classmethod float(column_name: str, value_fn: Callable[[Chunk], Any]) → ColumnSpec

Create a float column specification.

classmethod vector(column_name: str, value_fn: ~typing.Callable[[~apache_beam.ml.rag.types.Chunk], ~typing.Any] = <function chunk_embedding_fn>) → ColumnSpec

Create a vector column specification with string_to_vector() function.

classmethod json(column_name: str, value_fn: Callable[[Chunk], Any]) → ColumnSpec

Create a JSON column specification.

apache_beam.ml.rag.ingestion.mysql_common.embedding_to_string(embedding: List[float]) → str

Convert embedding to MySQL vector string format.

class apache_beam.ml.rag.ingestion.mysql_common.ColumnSpecsBuilder

Bases: object

Builder for ColumnSpec’s with chainable methods.

static with_defaults() → ColumnSpecsBuilder

Add all default column specifications.

with_id_spec(column_name: str = 'id', python_type: ~typing.Type = <class 'str'>, convert_fn: ~typing.Callable[[str], ~typing.Any] | None = None) → ColumnSpecsBuilder

Add ID ColumnSpec with optional type and conversion.

Parameters:

• column_name – Name for the ID column (defaults to “id”)

• python_type – Python type for the column (defaults to str)

• convert_fn – Optional function to convert the chunk ID If None, uses ID as-is

Returns:

Self for method chaining

Example

>>> builder.with_id_spec(
...     column_name="doc_id",
...     python_type=int,
...     convert_fn=lambda id: int(id.split('_')[1])
... )

with_content_spec(column_name: str = 'content', python_type: ~typing.Type = <class 'str'>, convert_fn: ~typing.Callable[[str], ~typing.Any] | None = None) → ColumnSpecsBuilder

Add content ColumnSpec with optional type and conversion.

Parameters:

• column_name – Name for the content column (defaults to “content”)

• python_type – Python type for the column (defaults to str)

• convert_fn – Optional function to convert the content text If None, uses content text as-is

Returns:

Self for method chaining

Example

>>> builder.with_content_spec(
...     column_name="content_length",
...     python_type=int,
...     convert_fn=len  # Store content length instead of content
... )

with_metadata_spec(column_name: str = 'metadata', python_type: ~typing.Type = <class 'str'>, convert_fn: ~typing.Callable[[~typing.Dict[str, ~typing.Any]], ~typing.Any] | None = None) → ColumnSpecsBuilder

Add metadata ColumnSpec with optional type and conversion.

Parameters:

• column_name – Name for the metadata column (defaults to “metadata”)

• python_type – Python type for the column (defaults to str)

• convert_fn – Optional function to convert the metadata dictionary If None and python_type is str, converts to JSON string

Returns:

Self for method chaining

Example

>>> builder.with_metadata_spec(
...     column_name="meta_tags",
...     python_type=str,
...     convert_fn=lambda meta: ','.join(meta.keys())
... )

with_embedding_spec(column_name: str = 'embedding', convert_fn: ~typing.Callable[[~typing.List[float]], ~typing.Any] = <function embedding_to_string>) → ColumnSpecsBuilder

Add embedding ColumnSpec with optional conversion.

Parameters:

• column_name – Name for the embedding column (defaults to “embedding”)

• convert_fn – Optional function to convert the dense embedding values If None, uses default MySQL vector format

Returns:

Self for method chaining

Example

>>> builder.with_embedding_spec(
...     column_name="embedding_vector",
...     convert_fn=lambda values: '[' + ','.join(f"{x:.4f}"
...       for x in values) + ']'
... )

add_metadata_field(field: str, python_type: Type, column_name: str | None = None, convert_fn: Callable[[Any], Any] | None = None, default: Any | None = None) → ColumnSpecsBuilder

Add a ColumnSpec that extracts and converts a field from chunk metadata.

Parameters:

• field – Key to extract from chunk metadata

• python_type – Python type for the column (e.g. str, int, float)

• column_name – Name for the column (defaults to metadata field name)

• convert_fn – Optional function to convert the extracted value to desired type. If None, value is used as-is

• default – Default value if field is missing from metadata

Returns:

Self for chaining

Examples

Simple string field: >>> builder.add_metadata_field(“source”, str)

Integer with default: >>> builder.add_metadata_field( … field=”count”, … python_type=int, … column_name=”item_count”, … default=0 … )

Float with conversion and default: >>> builder.add_metadata_field( … field=”confidence”, … python_type=float, … convert_fn=lambda x: round(float(x), 2), … default=0.0 … )

Timestamp with conversion: >>> builder.add_metadata_field( … field=”created_at”, … python_type=str, … convert_fn=lambda ts: ts.replace(‘T’, ‘ ‘) … )

add_custom_column_spec(spec: ColumnSpec) → ColumnSpecsBuilder

Add a custom ColumnSpec to the builder.

Use this method when you need complete control over the ColumnSpec, including custom value extraction and type handling.

Parameters:

spec – A ColumnSpec instance defining the column name, type, value extraction, and optional MySQL function.

Returns:

Self for method chaining

Examples

Custom text column from chunk metadata: >>> builder.add_custom_column_spec( … ColumnSpec.text( … column_name=”source_and_id”, … value_fn=lambda chunk: … f”{chunk.metadata.get(‘source’)}_{chunk.id}” … ) … )

build() → List[ColumnSpec]

Build the final list of column specifications.

class apache_beam.ml.rag.ingestion.mysql_common.ConflictResolution(action: Literal['UPDATE', 'IGNORE'] = 'UPDATE', update_fields: List[str] | None = None, primary_key_field: str | None = None)

Bases: object

Specification for how to handle conflicts during insert.

Configures conflict handling behavior when inserting records that may violate unique constraints using MySQL’s ON DUPLICATE KEY UPDATE syntax.

MySQL automatically detects conflicts based on PRIMARY KEY or UNIQUE constraints defined on the table.

action

How to handle conflicts - either “UPDATE” or “IGNORE”. UPDATE: Updates existing record with new values. IGNORE: Skips conflicting records (uses no-op update).

Type:

Literal[‘UPDATE’, ‘IGNORE’]

update_fields

Optional list of fields to update on conflict. If None, all fields are updated (for UPDATE action only).

Type:

List[str] | None

primary_key_field

Required for IGNORE action. The primary key field name to use for the no-op update.

Type:

str | None

Examples

Update all fields on conflict: >>> ConflictResolution(action=”UPDATE”)

Update specific fields on conflict: >>> ConflictResolution( … action=”UPDATE”, … update_fields=[“embedding”, “content”] … )

Ignore conflicts with explicit primary key: >>> ConflictResolution( … action=”IGNORE”, … primary_key_field=”id” … )

Ignore conflicts with custom primary key: >>> ConflictResolution( … action=”IGNORE”, … primary_key_field=”custom_id” … )

action: Literal['UPDATE', 'IGNORE'] = 'UPDATE'

update_fields: List[str] | None = None

primary_key_field: str | None = None