tensorframe 0.1.0

Arrow-backed DataFrames with native tensor column support for PyTorch/JAX

TensorFrame

Arrow-backed DataFrames with native tensor column support for PyTorch/JAX.

TensorFrame bridges the gap between tabular data processing (Pandas/Polars) and deep learning tensor execution (PyTorch/JAX) — without falling back to slow Python object arrays.

Install

pip install tensorframe

Quick Start

import tensorframe as tf
import numpy as np
import pandas as pd

# Create a DataFrame with a 1024-D embedding column
embeddings = np.random.randn(1000, 1024).astype("float32")
df = pd.DataFrame({
    "patient_id": range(1000),
    "age":        np.random.randint(20, 80, 1000).astype(float),
    "features":   pd.Series(tf.TensorArray(embeddings)),
})

# L2-normalise directly on the column
df["norm"] = df["features"].tensor.l2_normalize()

# Filter rows by cosine similarity to a query vector
target = np.random.randn(1024).astype("float32")
scores  = df["norm"].tensor.dot(target)
filtered = df[scores > 0.5]

# Zero-friction PyTorch DataLoader
loader = tf.to_pytorch_dataloader(
    filtered,
    features=["norm", "age"],
    targets="patient_id",
    batch_size=64,
    shuffle=True,
)
for batch in loader:
    x = batch["norm"]      # torch.Tensor (64, 1024)
    y = batch["target"]    # torch.Tensor (64,)

Read Parquet (auto-detects tensor columns)

df = tf.read_parquet("multimodal_biomarkers.parquet")
# Any FixedSizeList columns are automatically wrapped as TensorArray

Key Features

Feature	Details
Arrow-backed storage	FixedSizeList — one contiguous C memory block per column
Zero-copy to numpy	np.frombuffer wraps the Arrow buffer — no allocation
Efficient PyTorch bridge	Single memcpy from Arrow → writeable numpy → torch.from_numpy
.tensor accessor	l2_normalize, dot, cosine_similarity, mean, norm, reshape
Parquet round-trip	Tensor shape metadata preserved through write/read cycles
NIfTI support	tf.read_nifti_batch(paths) for volumetric medical imaging
Out-of-core iteration	tf.iter_chunks, ChunkedDataLoader for large datasets

Requirements

• Python ≥ 3.9
• numpy, pandas, pyarrow
• Optional: torch (for GPU ops and DataLoader), nibabel (NIfTI I/O)

Author

Md. Shahinur Hasan — [e-mail]

GitHub: https://github.com/Shahinurhasans/TensorFrame

License

Apache 2.0