TechPrufer
Get started
Google
Google logoGoogle · Feature trail

High-Performance Data Pipelines

2 posts · Aug 2024 – Oct 2025

What this trail covers

This thread tracks the evolution of building high-performance data pipelines for machine learning, focusing on efficient data loading, storage formats, and embedding generation. It introduces Grain, a data loading library for JAX, and ArrayRecord, a highly efficient file format. The thread also covers methods for converting existing datasets (like TFRecord) to ArrayRecord using tools like TensorFlow Datasets CLI and Apache Beam. The Gemini API now offers an experimental text embedding model, which has been generally available since July 2025, supporting over 100 languages with a 2048 maximum input token length and utilizing Matryoshka Representation Learning (MRL) for scalable output dimensions. Pricing is $0.15 per 1M input tokens, with free and paid tiers available. Support for the Batch API is also planned.

All Google trails
  1. Building High-Performance Data Pipelines with Grain and ArrayRecord- Google Developers BlogThis post introduces Grain, a data loading library for JAX, and ArrayRecord, a file format based on Riegeli, designed for high-performance data pipelines in ML. It details how Grain uses multiprocessing for parallel data loading and offers stateful, checkpointable iterators for reproducibility. ArrayRecord is highlighted for its efficient random access via a metadata index and chunked structure, enabling true global shuffling when paired with Grain. The post also provides instructions for converting TFRecord datasets to ArrayRecord using `tfds build` or Apache Beam.
  2. From Fine-Tuning to Production: A Scalable Embedding Pipeline with Dataflow- Google Developers BlogThis post introduces EmbeddingGemma, a 308M parameter open embedding model, and demonstrates its integration into a scalable, real-time knowledge ingestion pipeline using Google Cloud Dataflow and AlloyDB. It highlights the benefits of data locality, a unified system for autoscaling and monitoring, and simplicity through Apache Beam's MLTransform for embedding generation and VectorDatabaseWriteTransform for writing to vector databases. The post provides code examples for configuring SentenceTransformerEmbeddings with Gemma and AlloyDBVectorWriterConfig.
  3. Gemini Embedding now generally available in the Gemini API- Google Developers BlogAnnounces the general availability of the `gemini-embedding-001` text embedding model via the Gemini API and Vertex AI. Highlights its top performance on the MTEB Multilingual leaderboard, its versatility across domains, and its support for over 100 languages with a 2048 token limit. Details the use of Matryoshka Representation Learning (MRL) for scalable output dimensions (3072, 1536, 768). Provides pricing information ($0.15 per 1M input tokens) and links to documentation and a quickstart notebook. Notes the deprecation of older embedding models and the upcoming support for the Batch API.
  4. Simplified Dataflow Connectors with Managed I/O- Google Developers BlogThis post introduces Google Cloud Dataflow's Managed I/O feature, which simplifies the use of Apache Beam I/O connectors. Managed I/O automates SDK upgrades for connectors, provides standardized Java and Python APIs for various connectors (e.g., BigQuery, Kafka, Iceberg), and automatically optimizes connectors for the Dataflow runner. Benchmarks demonstrate linear scaling and efficient resource utilization for Iceberg and Kafka sinks, showcasing improved throughput, low latency, and backlog.
  5. State-of-the-art text embedding via the Gemini API- Google Developers BlogIntroduces a new experimental Gemini Embedding text model (gemini-embedding-exp-03-07) available via the Gemini API. This model is trained on the Gemini model, achieves top ranks on the MTEB Multilingual leaderboard, and offers improved features including an 8K token input limit, 3K output dimensions with Matryoshka Representation Learning (MRL) for scalable storage, and expanded language support (over 100 languages). It is presented as a unified model surpassing previous task-specific models and is compatible with the existing `embed_content` endpoint.
  7. See the Similarity: Personalizing Visual Search with Multimodal Embeddings- Google Developers BlogThis post introduces multimodal embeddings, which represent text, images, and video in a shared vector space. It details two practical applications: searching internal Google Slides presentations (over 775,000 slides) using Vertex AI's Vector Search with the ScaNN algorithm for enterprise-scale data, and an artist's personal illustration collection (approx. 250 images) using Cloud Firestore Vector Search with a KNN algorithm. The post provides code examples for generating embeddings, storing them in Firestore and Cloud Storage, and performing similarity searches. It also discusses the trade-offs between KNN and ScaNN for different scales and mentions UMAP for visualizing embeddings.
  8. Inference with Gemma using Dataflow and vLLM- Google Developers BlogIntroduces the integration of vLLM with Apache Beam's Dataflow for efficient LLM inference. Details vLLM's continuous batching technique for optimizing LLM serving by dynamically recomputing batches as tokens are generated, contrasting it with traditional fixed-size batching which leads to wasted compute. Explains how Dataflow's model manager simplifies deploying large models by allowing control over model copies, enabling a single vLLM instance to be shared across worker processes. Provides a Python code snippet demonstrating the integration and reports a >23x performance improvement in vCPU hours for processing 10,000 prompts using vLLM compared to a naive batching strategy.
  9. Gemma for Streaming ML with Dataflow- Google Developers BlogThis post details the integration of Gemma 2 models within an Apache Beam pipeline running on Google Dataflow for real-time processing of customer support chat data. It demonstrates how to use Gemma for sentiment analysis (classifying as positive, neutral, or negative), conversation summarization, and generating contextually appropriate responses for negative sentiment chats. The post highlights the benefits of embedding LLMs directly into streaming pipelines for data locality, unified autoscaling, and simplified monitoring. It also discusses challenges with non-deterministic LLM outputs and the need for response validation, drawing parallels to processing unstructured data. The implementation uses KerasNLP and requires a custom model handler for Gemma, with a recommendation to use JAX as the Keras backend and GPUs on Dataflow ML for production performance.