AI-Powered Data Processing for Pandas & Spark

Welcome to openaivec - Transform your data analysis with OpenAI's language models! This library enables seamless integration of AI text processing, sentiment analysis, NLP tasks, and embeddings into your Pandas DataFrames and Apache Spark workflows for scalable data insights.

🚀 Quick Start Example

Transform your data with AI in just one line:

import pandas as pd
from openaivec import pandas_ext

# AI-powered data processing
fruits = pd.Series(["apple", "banana", "orange", "grape", "kiwi"])
fruits.ai.responses("Translate this fruit name into French.")
# Result: ['pomme', 'banane', 'orange', 'raisin', 'kiwi']

Perfect for data scientists, analysts, and ML engineers who want to leverage AI for text processing at scale.

📦 Installation

pipuvWith Spark Support

pip install openaivec

uv add openaivec

pip install "openaivec[spark]"
# or
uv add "openaivec[spark]"

🎯 Key Features

• 🚀 Vectorized Processing: Handle thousands of records in minutes, not hours
• 💰 Cost Efficient: Automatic deduplication reduces API costs by 50-90%
• 🔗 Seamless Integration: Works within existing pandas/Spark workflows
• 📈 Enterprise Scale: From 100s to millions of records
• 🤖 Advanced NLP: Pre-built tasks for sentiment analysis, translation, NER, and more

Links

• GitHub Repository
• PyPI Package
• Complete Documentation

📚 Examples & Tutorials

Get started with these comprehensive examples:

📓 Getting Started - Basic pandas integration and usage
📓 Customer Feedback Analysis - Sentiment analysis & prioritization
📓 Survey Data Transformation - Unstructured to structured data
📓 Spark Processing - Enterprise-scale distributed processing
📓 Async Workflows - High-performance async processing
📓 Prompt Engineering - Advanced prompting techniques
📓 FAQ Generation - Auto-generate FAQs from documents

📖 API Reference

Detailed documentation for all components:

🔗 pandas_ext - Pandas Series and DataFrame extensions
🔗 spark - Apache Spark UDF builders
🔗 responses - Batch response processing
🔗 embeddings - Batch embedding generation
🔗 prompt - Few-shot prompt building
🔗 util - Utility functions and helpers

Quick Start

Here is a simple example of how to use openaivec with pandas:

import pandas as pd
from openai import OpenAI
from openaivec import pandas_ext

from typing import List

# Set OpenAI Client (optional: this is default client if environment "OPENAI_API_KEY" is set)
pandas_ext.use(OpenAI())

# Set models for responses and embeddings(optional: these are default models)
pandas_ext.responses_model("gpt-4.1-nano")
pandas_ext.embeddings_model("text-embedding-3-small")


fruits: List[str] = ["apple", "banana", "orange", "grape", "kiwi", "mango", "peach", "pear", "pineapple", "strawberry"]
fruits_df = pd.DataFrame({"name": fruits})

frults_df is a pandas DataFrame with a single column name containing the names of fruits. We can mutate the Field name with the accessor ai to add a new column color with the color of each fruit.:

fruits_df.assign(
    color=lambda df: df["name"].ai.responses("What is the color of this fruit?")
)

The result is a new DataFrame with the same number of rows as fruits_df, but with an additional column color containing the color of each fruit. The ai accessor uses the OpenAI API to generate the responses for each fruit name in the name column.

name	color
apple	red
banana	yellow
orange	orange
grape	purple
kiwi	brown
mango	orange
peach	orange
pear	green
pineapple	brown
strawberry	red

Structured Output is also supported. For example, we will translate the name of each fruit into multiple languages. We can use the ai accessor to generate a new column translation with the translation of each fruit name into English, French, Japanese, Spanish, German, Italian, Portuguese and Russian:

from pydantic import BaseModel

class Translation(BaseModel):
    en: str  # English
    fr: str  # French
    ja: str  # Japanese
    es: str  # Spanish
    de: str  # German
    it: str  # Italian
    pt: str  # Portuguese
    ru: str  # Russian

fruits_df.assign(
    translation=lambda df: df["name"].ai.responses(
        instructions="Translate this fruit name into English, French, Japanese, Spanish, German, Italian, Portuguese and Russian.",
        response_format=Translation,
    )
)

name	translation
apple	en='Apple' fr='Pomme' ja='リンゴ' es='Manzana' de...
banana	en='Banana' fr='Banane' ja='バナナ' es='Banana' de...
orange	en='Orange' fr='Orange' ja='オレンジ' es='Naranja' de...
grape	en='Grape' fr='Raisin' ja='ブドウ' es='Uva' de='T...
kiwi	en='Kiwi' fr='Kiwi' ja='キウイ' es='Kiwi' de='Kiw...
mango	en='Mango' fr='Mangue' ja='マンゴー' es='Mango' de...
peach	en='Peach' fr='Pêche' ja='モモ' es='Durazno' de...
pear	en='Pear' fr='Poire' ja='梨' es='Pera' de='Birn...
pineapple	en='Pineapple' fr='Ananas' ja='パイナップル' es='Piñ...
strawberry	en='Strawberry' fr='Fraise' ja='イチゴ' es='Fresa...

Structured output can be extracted into separate columns using the extract method. For example, we can extract the translations into separate columns for each language:

fruits_df.assign(
    translation=lambda df: df["name"].ai.responses(
        instructions="Translate this fruit name into English, French, Japanese, Spanish, German, Italian, Portuguese and Russian.",
        response_format=Translation,
    )
).ai.extract("translation")

name	translation_en	translation_fr	translation_ja	translation_es	translation_de	translation_it	translation_pt	translation_ru
apple	Apple	Pomme	リンゴ	Manzana	Apfel	Mela	Maçã	Яблоко
banana	Banana	Banane	バナナ	Banana	Banane	Banana	Banana	Банан
orange	Orange	Orange	オレンジ	Naranja	Orange	Arancia	Laranja	Апельсин
grape	Grape	Raisin	ブドウ	Uva	Traube	Uva	Uva	Виноград
kiwi	Kiwi	Kiwi	キウイ	Kiwi	Kiwi	Kiwi	Kiwi	Киви
mango	Mango	Mangue	マンゴー	Mango	Mango	Mango	Manga	Манго
peach	Peach	Pêche	モモ	Durazno	Pfirsich	Pesca	Pêssego	Персик
pear	Pear	Poire	梨	Pera	Birne	Pera	Pêra	Груша
pineapple	Pineapple	Ananas	パイナップル	Piña	Ananas	Ananas	Abacaxi	Ананас
strawberry	Strawberry	Fraise	イチゴ	Fresa	Erdbeere	Fragola	Morango	Клубника