[python-package] Add support for PyArrow-backed Pandas DataFrames by maxzw · Pull Request #7263 · lightgbm-org/LightGBM

maxzw · 2026-05-07T11:39:46Z

NOTE: full support for all metadata fields depends on #7260 (position)

Problem

LightGBM currently converts all pandas DataFrames to NumPy arrays during Dataset construction, even when the DataFrame uses PyArrow-backed columns (ArrowDtype). This forces an unnecessary copy of data that's already in Arrow's columnar format, missing an opportunity for zero-copy data transfer to the C++ backend.

With pandas 2.0+ increasingly adopting PyArrow as the default backend for better memory efficiency and interoperability, users working with Arrow-backed DataFrames face a performance penalty. This PR enables LightGBM to detect and leverage PyArrow-backed columns, maintaining data in the Arrow format through the entire pipeline.

Approach

The key design principle was to keep pandas-specific processing (categorical handling, type validation) in the pandas-specific code path while cleanly passing Arrow tables downstream.

Integration Point: Modified _data_from_pandas() to detect PyArrow-backed columns and convert the entire DataFrame to an Arrow Table instead of a NumPy array. This keeps all existing categorical feature logic intact while letting the downstream C++ interface naturally handle Arrow tables via the existing _create_from_arrow_table() path.

Zero-Copy Strategy: The implementation handles DataFrames with mixed dtypes (Arrow + NumPy + nullable extension types) by extracting underlying buffers directly where possible, avoiding intermediate copies. Each column type is handled optimally:

Arrow-backed columns: Extract the underlying Arrow array with .__arrow_array__()
Nullable extension types: Pass through PyArrow's null-aware conversion
NumPy-backed columns: Use .values for a zero-copy view

Only dtypes already supported by LightGBM's Arrow interface (int, float, bool) are allowed, ensuring consistency with the existing Arrow code path.

Implementation note: An alternative approach was considered that would use pa.Table.from_pandas() for all-Arrow DataFrames:

if all(_is_pd_arrow_dtype(dtype) for dtype in data.dtypes):
    table = pa_Table.from_pandas(data, preserve_index=False)

However, benchmarking showed this provided no meaningful speedup over the unified column-by-column implementation, so the latter approach was chosen.

Changes

Core Logic (basic.py)

_pandas_to_arrow_table() (new): Converts pandas DataFrames to Arrow Tables with zero-copy optimizations. Uses a unified implementation that handles mixed dtypes efficiently:
- Arrow-backed columns: Extract underlying array directly (.__arrow_array__())
- Nullable extension types (Int64Dtype, Float64Dtype): Pass through PyArrow's null-aware conversion
- NumPy-backed columns: Use .values for zero-copy view
- Builds the Arrow table with pa.Table.from_arrays() for consistent performance across all DataFrame types
_data_from_pandas(): Updated return type from np.ndarray to Union[np.ndarray, pa.Table]. Early-exits with Arrow table when any PyArrow-backed columns detected, preserving all existing categorical logic.
_check_for_bad_arrow_dtypes() (new): Centralized validation for Arrow type compatibility, replacing duplicated checks in Dataset and _InnerPredictor.
_is_pd_arrow_dtype(), _is_arrow_backed_pd_series(), _extract_arrow_array() (new): Helper functions for detecting and extracting Arrow arrays from pandas Series.
Metadata field support: Extended set_label(), set_weight(), set_group(), set_position() to accept PyArrow-backed Series, maintaining API consistency.

Compatibility Layer (compat.py)

Added imports for pd.ArrowDtype and pa.DataType with fallbacks

Tests

Basic Dataset functionality (test_basic.py):

All PyArrow integer types (int8/16/32/64, uint8/32/64)
All PyArrow float types (float32/64)
Boolean PyArrow columns
Mixed DataFrames (Arrow + NumPy + nullable extension types)
Categorical feature handling with Arrow backend
Edge cases: empty DataFrames, all-null columns, numeric column names
Metadata fields (label, weight, group, init_score) as Arrow-backed Series
Invalid dtypes (string, timestamp, date, time, duration) properly rejected

Sklearn integration (test_sklearn.py):

LGBMClassifier fit/predict with Arrow DataFrames
LGBMRegressor fit/predict with Arrow DataFrames
Cross-validation with Arrow-backed data

Benchmarks

All benchmarks ran 50 iterations per configuration on datasets ranging from 10K to 1M rows. Four scenarios tested:

Numpy baseline: Standard pandas DataFrame backed by NumPy arrays
Full arrow: All columns as PyArrow-backed
Mixed arrow/numpy: 50% PyArrow + 50% NumPy columns
Mixed arrow/nullable: 50% PyArrow + 50% nullable extension types

Dataset Construction Time

PyArrow-backed DataFrames provide speedups up to 43% in dataset construction for large datasets (50K+ rows). Minimal impact for small datasets (10K rows: ~6%).

Benchmark code and logs

The benchmark below was ran on a 16GB MacBook Pro M2 Pro.

import time
import numpy as np
import pandas as pd
import lightgbm as lgb
import matplotlib.pyplot as plt

# Dataset sizes to test
SIZES = [
    (10_000, 10),  # 10K rows, 10 cols
    (50_000, 20),  # 50K rows, 20 cols
    (100_000, 20),  # 100K rows, 20 cols
    (500_000, 50),  # 500K rows, 50 cols
    (1_000_000, 50),  # 1M rows, 50 cols
]

N_ITER = 50  # Iterations per configuration

print("=" * 80)
print("CONSTRUCTION TIME BENCHMARK")
print("=" * 80)

# Collect all results for plotting
results = {"sizes": [], "numpy": [], "arrow": [], "mixed": [], "nullable": []}

for n_rows, n_cols in SIZES:
    data = np.random.RandomState(42).random((n_rows, n_cols))
    y = np.random.RandomState(42).randint(0, 2, n_rows)

    print(f"\nDataset: {n_rows:,} rows x {n_cols} cols ({N_ITER} iterations)")
    print("-" * 80)

    # Baseline: Numpy-backed DataFrame
    numpy_times = []
    for _ in range(N_ITER):
        df = pd.DataFrame(data)
        start = time.perf_counter()
        ds = lgb.Dataset(df, label=y).construct()
        numpy_times.append((time.perf_counter() - start) * 1000)
    numpy_mean = np.mean(numpy_times)

    print(f"  Numpy baseline:     {numpy_mean:>6.2f} ms")

    # Full arrow: All columns arrow-backed (fast path)
    arrow_times = []
    for _ in range(N_ITER):
        df = pd.DataFrame(data, dtype="float64[pyarrow]")
        start = time.perf_counter()
        ds = lgb.Dataset(df, label=y).construct()
        arrow_times.append((time.perf_counter() - start) * 1000)
    arrow_mean = np.mean(arrow_times)
    arrow_pct = ((arrow_mean - numpy_mean) / numpy_mean) * 100

    print(f"  Full arrow:         {arrow_mean:>6.2f} ms ({arrow_pct:+.1f}%)")

    # Mixed arrow/numpy: 50% arrow + 50% numpy
    mixed_times = []
    for _ in range(N_ITER):
        df = pd.DataFrame(data)
        half = n_cols // 2
        for i in range(half):
            df[i] = df[i].astype("float64[pyarrow]")
        start = time.perf_counter()
        ds = lgb.Dataset(df, label=y).construct()
        mixed_times.append((time.perf_counter() - start) * 1000)
    mixed_mean = np.mean(mixed_times)
    mixed_pct = ((mixed_mean - numpy_mean) / numpy_mean) * 100

    print(f"  Mixed arrow/numpy:  {mixed_mean:>6.2f} ms ({mixed_pct:+.1f}%)")

    # Mixed arrow/nullable: 50% arrow + 50% nullable
    nullable_times = []
    for _ in range(N_ITER):
        df = pd.DataFrame(data)
        half = n_cols // 2
        for i in range(half):
            df[i] = df[i].astype("float64[pyarrow]")
        for i in range(half, n_cols):
            df[i] = df[i].astype(pd.Float64Dtype())
        start = time.perf_counter()
        ds = lgb.Dataset(df, label=y).construct()
        nullable_times.append((time.perf_counter() - start) * 1000)
    nullable_mean = np.mean(nullable_times)
    nullable_pct = ((nullable_mean - numpy_mean) / numpy_mean) * 100

    print(f"  Mixed arrow/null:   {nullable_mean:>6.2f} ms ({nullable_pct:+.1f}%)")

    # Store results
    results["sizes"].append(f"{n_rows // 1000}K")
    results["numpy"].append(numpy_times)
    results["arrow"].append(arrow_times)
    results["mixed"].append(mixed_times)
    results["nullable"].append(nullable_times)

# Create subplots - one per dataset size for independent y axes
print("\nGenerating plot...")
n_sizes = len(results["sizes"])
fig, axes = plt.subplots(1, n_sizes, figsize=(4 * n_sizes, 6), sharey=False)

if n_sizes == 1:
    axes = [axes]

for i, (ax, size_label) in enumerate(zip(axes, results["sizes"])):
    data_to_plot = [results["numpy"][i], results["arrow"][i], results["mixed"][i], results["nullable"][i]]
    bp = ax.boxplot(data_to_plot, patch_artist=True, tick_labels=["Numpy", "Arrow", "50% Arrow\n50% Numpy", "50% Arrow\n50% Null"])

    # Color the boxes
    colors = ["lightblue", "lightgreen", "lightyellow", "lightcoral"]
    for patch, color in zip(bp["boxes"], colors):
        patch.set_facecolor(color)

    ax.set_title(f"{size_label} rows")
    ax.set_ylabel("Construction Time (ms)")
    ax.grid(True, alpha=0.3)

fig.suptitle("Dataset Construction Time: Numpy vs Arrow Backends", fontsize=14)
plt.tight_layout()
plt.savefig("construction_time.png", dpi=150)

Logs:

================================================================================
CONSTRUCTION TIME BENCHMARK
================================================================================

Dataset: 10,000 rows x 10 cols (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:       2.11 ms
  Full arrow:           1.98 ms (-6.1%)
  Mixed arrow/numpy:    2.00 ms (-5.0%)
  Mixed arrow/null:     2.05 ms (-2.7%)

Dataset: 50,000 rows x 20 cols (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:      22.64 ms
  Full arrow:          14.29 ms (-36.9%)
  Mixed arrow/numpy:   15.37 ms (-32.1%)
  Mixed arrow/null:    15.87 ms (-29.9%)

Dataset: 100,000 rows x 20 cols (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:      45.23 ms
  Full arrow:          29.17 ms (-35.5%)
  Mixed arrow/numpy:   32.19 ms (-28.8%)
  Mixed arrow/null:    30.94 ms (-31.6%)

Dataset: 500,000 rows x 50 cols (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:     292.34 ms
  Full arrow:         166.31 ms (-43.1%)
  Mixed arrow/numpy:  192.65 ms (-34.1%)
  Mixed arrow/null:   186.97 ms (-36.0%)

Dataset: 1,000,000 rows x 50 cols (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:     365.08 ms
  Full arrow:         225.18 ms (-38.3%)
  Mixed arrow/numpy:  283.62 ms (-22.3%)
  Mixed arrow/null:   265.41 ms (-27.3%)

Generating plot...

End-to-End Training Time

PyArrow provides speedups up to 10% for end-to-end training on large datasets (500K-1M rows). Minimal impact for small to medium datasets.

Benchmark code and logs

The benchmark below was ran on a 16GB MacBook Pro M2 Pro.

import time
import numpy as np
import pandas as pd
import lightgbm as lgb
import matplotlib.pyplot as plt

# Dataset sizes to test
SIZES = [
    (10_000, 10, 50),  # 10K rows, 10 cols, 50 rounds
    (50_000, 20, 100),  # 50K rows, 20 cols, 100 rounds
    (100_000, 20, 100),  # 100K rows, 20 cols, 100 rounds
    (500_000, 50, 100),  # 500K rows, 50 cols, 100 rounds
    (1_000_000, 50, 100),  # 1M rows, 50 cols, 100 rounds
]

N_ITER = 50  # Iterations per configuration
params = {"objective": "binary", "verbose": -1, "seed": 42}

print("=" * 80)
print("END-TO-END TRAINING BENCHMARK")
print("=" * 80)

# Collect all results for plotting
results = {"sizes": [], "numpy": [], "arrow": [], "mixed": [], "nullable": []}

for n_rows, n_cols, n_rounds in SIZES:
    data = np.random.RandomState(42).random((n_rows, n_cols))
    y = np.random.RandomState(42).randint(0, 2, n_rows)

    print(f"\nDataset: {n_rows:,} rows x {n_cols} cols, {n_rounds} rounds ({N_ITER} iterations)")
    print("-" * 80)

    # Baseline: Numpy-backed DataFrame
    numpy_times = []
    for _ in range(N_ITER):
        df = pd.DataFrame(data)
        start = time.perf_counter()
        ds = lgb.Dataset(df, label=y).construct()
        bst = lgb.train(params, ds, num_boost_round=n_rounds)
        numpy_times.append((time.perf_counter() - start) * 1000)
    numpy_mean = np.mean(numpy_times)

    print(f"  Numpy baseline:     {numpy_mean:>7.1f} ms")

    # Full arrow: All columns arrow-backed
    arrow_times = []
    for _ in range(N_ITER):
        df = pd.DataFrame(data, dtype="float32[pyarrow]")
        start = time.perf_counter()
        ds = lgb.Dataset(df, label=y).construct()
        bst = lgb.train(params, ds, num_boost_round=n_rounds)
        arrow_times.append((time.perf_counter() - start) * 1000)
    arrow_mean = np.mean(arrow_times)
    arrow_pct = ((arrow_mean - numpy_mean) / numpy_mean) * 100

    print(f"  Full arrow:         {arrow_mean:>7.1f} ms ({arrow_pct:+.1f}%)")

    # Mixed arrow/numpy: 50% arrow + 50% numpy
    mixed_times = []
    for _ in range(N_ITER):
        df = pd.DataFrame(data)
        half = n_cols // 2
        for i in range(half):
            df[i] = df[i].astype("float32[pyarrow]")
        start = time.perf_counter()
        ds = lgb.Dataset(df, label=y).construct()
        bst = lgb.train(params, ds, num_boost_round=n_rounds)
        mixed_times.append((time.perf_counter() - start) * 1000)
    mixed_mean = np.mean(mixed_times)
    mixed_pct = ((mixed_mean - numpy_mean) / numpy_mean) * 100

    print(f"  Mixed arrow/numpy:  {mixed_mean:>7.1f} ms ({mixed_pct:+.1f}%)")

    # Mixed arrow/nullable: 50% arrow + 50% nullable
    nullable_times = []
    for _ in range(N_ITER):
        df = pd.DataFrame(data)
        half = n_cols // 2
        for i in range(half):
            df[i] = df[i].astype("float32[pyarrow]")
        for i in range(half, n_cols):
            df[i] = df[i].astype(pd.Float64Dtype())
        start = time.perf_counter()
        ds = lgb.Dataset(df, label=y).construct()
        bst = lgb.train(params, ds, num_boost_round=n_rounds)
        nullable_times.append((time.perf_counter() - start) * 1000)
    nullable_mean = np.mean(nullable_times)
    nullable_pct = ((nullable_mean - numpy_mean) / numpy_mean) * 100

    print(f"  Mixed arrow/null:   {nullable_mean:>7.1f} ms ({nullable_pct:+.1f}%)")

    # Store results
    results["sizes"].append(f"{n_rows // 1000}K")
    results["numpy"].append(numpy_times)
    results["arrow"].append(arrow_times)
    results["mixed"].append(mixed_times)
    results["nullable"].append(nullable_times)

# Create subplots - one per dataset size for independent y axes
print("\nGenerating plot...")
n_sizes = len(results["sizes"])
fig, axes = plt.subplots(1, n_sizes, figsize=(4 * n_sizes, 6), sharey=False)

if n_sizes == 1:
    axes = [axes]

for i, (ax, size_label) in enumerate(zip(axes, results["sizes"])):
    data_to_plot = [results["numpy"][i], results["arrow"][i], results["mixed"][i], results["nullable"][i]]
    bp = ax.boxplot(data_to_plot, patch_artist=True, tick_labels=["Numpy", "Arrow", "50% Arrow\n50% Numpy", "50% Arrow\n50% Null"])

    # Color the boxes
    colors = ["lightblue", "lightgreen", "lightyellow", "lightcoral"]
    for patch, color in zip(bp["boxes"], colors):
        patch.set_facecolor(color)

    ax.set_title(f"{size_label} rows")
    ax.set_ylabel("Training Time (ms)")
    ax.grid(True, alpha=0.3)

fig.suptitle("End-to-End Training: Numpy vs Arrow Backends", fontsize=14)
plt.tight_layout()
plt.savefig("e2e_training.png", dpi=150)

Logs:

================================================================================
END-TO-END TRAINING BENCHMARK
================================================================================

Dataset: 10,000 rows x 10 cols, 50 rounds (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:       267.2 ms
  Full arrow:           267.4 ms (+0.1%)
  Mixed arrow/numpy:    267.2 ms (-0.0%)
  Mixed arrow/null:     266.9 ms (-0.1%)

Dataset: 50,000 rows x 20 cols, 100 rounds (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:       595.5 ms
  Full arrow:           588.1 ms (-1.3%)
  Mixed arrow/numpy:    591.9 ms (-0.6%)
  Mixed arrow/null:     590.4 ms (-0.8%)

Dataset: 100,000 rows x 20 cols, 100 rounds (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:       655.6 ms
  Full arrow:           678.4 ms (+3.5%)
  Mixed arrow/numpy:    743.3 ms (+13.4%)
  Mixed arrow/null:     781.0 ms (+19.1%)

Dataset: 500,000 rows x 50 cols, 100 rounds (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:      1352.6 ms
  Full arrow:          1241.3 ms (-8.2%)
  Mixed arrow/numpy:   1266.3 ms (-6.4%)
  Mixed arrow/null:    1259.1 ms (-6.9%)

Dataset: 1,000,000 rows x 50 cols, 100 rounds (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:      2077.7 ms
  Full arrow:          1882.1 ms (-9.4%)
  Mixed arrow/numpy:   1879.3 ms (-9.6%)
  Mixed arrow/null:    1890.9 ms (-9.0%)

Generating plot...

Sklearn Integration

PyArrow has minimal impact on sklearn integration performance, with speedups up to 11% for large datasets (500K rows). Performance is within ±2% for most dataset sizes.

Benchmark code and logs

The benchmark below was ran on a 16GB MacBook Pro M2 Pro.

import time
import numpy as np
import pandas as pd
import lightgbm as lgb
import matplotlib.pyplot as plt

# Dataset sizes to test
SIZES = [
    (10_000, 10, 50),  # 10K rows, 10 cols, 50 estimators
    (50_000, 20, 50),  # 50K rows, 20 cols, 50 estimators
    (100_000, 20, 100),  # 100K rows, 20 cols, 100 estimators
    (500_000, 50, 100),  # 500K rows, 50 cols, 100 estimators
    (1_000_000, 50, 100),  # 1M rows, 50 cols, 100 estimators
]

N_ITER = 50  # Iterations per configuration

print("=" * 80)
print("SKLEARN INTEGRATION BENCHMARK")
print("=" * 80)

# Collect all results for plotting
results = {"sizes": [], "numpy": [], "arrow": [], "mixed": [], "nullable": []}

for n_rows, n_cols, n_estimators in SIZES:
    data = np.random.RandomState(42).random((n_rows, n_cols))
    y = np.random.RandomState(42).randint(0, 2, n_rows)

    print(f"\nDataset: {n_rows:,} rows x {n_cols} cols, {n_estimators} estimators ({N_ITER} iterations)")
    print("-" * 80)

    # Baseline: Numpy-backed DataFrame
    numpy_times = []
    for _ in range(N_ITER):
        df = pd.DataFrame(data)
        start = time.perf_counter()
        clf = lgb.LGBMClassifier(n_estimators=n_estimators, verbose=-1, random_state=42)
        clf.fit(df, y)
        pred = clf.predict(df)
        numpy_times.append((time.perf_counter() - start) * 1000)
    numpy_mean = np.mean(numpy_times)

    print(f"  Numpy baseline:     {numpy_mean:>7.1f} ms")

    # Full arrow: All columns arrow-backed
    arrow_times = []
    for _ in range(N_ITER):
        df_arrow = pd.DataFrame(data, dtype="float32[pyarrow]")
        y_arrow = pd.Series(y, dtype="int32[pyarrow]")
        start = time.perf_counter()
        clf = lgb.LGBMClassifier(n_estimators=n_estimators, verbose=-1, random_state=42)
        clf.fit(df_arrow, y_arrow)
        pred = clf.predict(df_arrow)
        arrow_times.append((time.perf_counter() - start) * 1000)
    arrow_mean = np.mean(arrow_times)
    arrow_pct = ((arrow_mean - numpy_mean) / numpy_mean) * 100

    print(f"  Full arrow:         {arrow_mean:>7.1f} ms ({arrow_pct:+.1f}%)")

    # Mixed arrow/numpy: 50% arrow + 50% numpy
    mixed_times = []
    for _ in range(N_ITER):
        df_mixed = pd.DataFrame(data)
        half = n_cols // 2
        for i in range(half):
            df_mixed[i] = df_mixed[i].astype("float32[pyarrow]")
        start = time.perf_counter()
        clf = lgb.LGBMClassifier(n_estimators=n_estimators, verbose=-1, random_state=42)
        clf.fit(df_mixed, y)
        pred = clf.predict(df_mixed)
        mixed_times.append((time.perf_counter() - start) * 1000)
    mixed_mean = np.mean(mixed_times)
    mixed_pct = ((mixed_mean - numpy_mean) / numpy_mean) * 100

    print(f"  Mixed arrow/numpy:  {mixed_mean:>7.1f} ms ({mixed_pct:+.1f}%)")

    # Mixed arrow/nullable: 50% arrow + 50% nullable
    nullable_times = []
    for _ in range(N_ITER):
        df_nullable = pd.DataFrame(data)
        half = n_cols // 2
        for i in range(half):
            df_nullable[i] = df_nullable[i].astype("float32[pyarrow]")
        for i in range(half, n_cols):
            df_nullable[i] = df_nullable[i].astype(pd.Float64Dtype())
        start = time.perf_counter()
        clf = lgb.LGBMClassifier(n_estimators=n_estimators, verbose=-1, random_state=42)
        clf.fit(df_nullable, y)
        pred = clf.predict(df_nullable)
        nullable_times.append((time.perf_counter() - start) * 1000)
    nullable_mean = np.mean(nullable_times)
    nullable_pct = ((nullable_mean - numpy_mean) / numpy_mean) * 100

    print(f"  Mixed arrow/null:   {nullable_mean:>7.1f} ms ({nullable_pct:+.1f}%)")

    # Store results
    results["sizes"].append(f"{n_rows // 1000}K")
    results["numpy"].append(numpy_times)
    results["arrow"].append(arrow_times)
    results["mixed"].append(mixed_times)
    results["nullable"].append(nullable_times)

# Create subplots - one per dataset size for independent y axes
print("\nGenerating plot...")
n_sizes = len(results["sizes"])
fig, axes = plt.subplots(1, n_sizes, figsize=(4 * n_sizes, 6), sharey=False)

if n_sizes == 1:
    axes = [axes]

for i, (ax, size_label) in enumerate(zip(axes, results["sizes"])):
    data_to_plot = [results["numpy"][i], results["arrow"][i], results["mixed"][i], results["nullable"][i]]
    bp = ax.boxplot(data_to_plot, patch_artist=True, tick_labels=["Numpy", "Arrow", "50% Arrow\n50% Numpy", "50% Arrow\n50% Null"])

    # Color the boxes
    colors = ["lightblue", "lightgreen", "lightyellow", "lightcoral"]
    for patch, color in zip(bp["boxes"], colors):
        patch.set_facecolor(color)

    ax.set_title(f"{size_label} rows")
    ax.set_ylabel("Fit + Predict Time (ms)")
    ax.grid(True, alpha=0.3)

fig.suptitle("Sklearn Integration: Numpy vs Arrow Backends", fontsize=14)
plt.tight_layout()
plt.savefig("sklearn_integration.png", dpi=150)

Logs:

================================================================================
SKLEARN INTEGRATION BENCHMARK
================================================================================

Dataset: 10,000 rows x 10 cols, 50 estimators (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:       272.2 ms
  Full arrow:           275.9 ms (+1.4%)
  Mixed arrow/numpy:    272.2 ms (-0.0%)
  Mixed arrow/null:     274.5 ms (+0.8%)

Dataset: 50,000 rows x 20 cols, 50 estimators (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:       327.1 ms
  Full arrow:           321.0 ms (-1.9%)
  Mixed arrow/numpy:    322.8 ms (-1.3%)
  Mixed arrow/null:     325.2 ms (-0.6%)

Dataset: 100,000 rows x 20 cols, 100 estimators (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:       692.9 ms
  Full arrow:           680.7 ms (-1.8%)
  Mixed arrow/numpy:    690.1 ms (-0.4%)
  Mixed arrow/null:     725.2 ms (+4.7%)

Dataset: 500,000 rows x 50 cols, 100 estimators (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:      1554.4 ms
  Full arrow:          1385.9 ms (-10.8%)
  Mixed arrow/numpy:   1435.5 ms (-7.7%)
  Mixed arrow/null:    1431.4 ms (-7.9%)

Dataset: 1,000,000 rows x 50 cols, 100 estimators (50 iterations)
--------------------------------------------------------------------------------
  Numpy baseline:      2242.0 ms
  Full arrow:          2228.1 ms (-0.6%)
  Mixed arrow/numpy:   2322.6 ms (+3.6%)
  Mixed arrow/null:    2226.8 ms (-0.7%)

Generating plot...

Compatibility

Requires pyarrow and cffi installed (same as existing Arrow support).

Open Questions

Should PyArrow support be mentioned in the docstrings that already mention pandas DataFrame/Series?
Should PyArrow support be documented in user-facing guides?
Should the benchmark scripts be committed to the repo?
Feedback on benchmark methodology?
Are there missing edge cases / tests that should still be added?
Should the minimum pandas/pyarrow versions where this works be documented?

maxzw · 2026-05-07T11:47:56Z

@jameslamb I'd appreciate your feedback on this draft PR!

Ic3fr0g · 2026-05-09T13:32:01Z

I raised a PR(#7264) that actually addresses categorical handling for polars dataframes and while you raise a good point about the numpy copy it might be more performant if pandas-2.0+ is converted to polars df before training. Here's a way to do it with zero-copy (with full dtype support provided the previous PR goes live) -

import polars as pl
import pandas as pd
import pyarrow as pa

# 1. Ensure Pandas 2.0+ and pyarrow are installed.
# 2. Create Pandas DataFrame backed by Arrow
df_pd = pd.DataFrame({'a': [1, 2], 'b': [3, 4]}).convert_dtypes(dtype_backend="pyarrow")

# 3. Use pl.from_pandas()
# This uses Arrow memory directly if possible (zero-copy).
pl_df = pl.from_pandas(df_pd)

Additionally, pandas-3.0+ natively implements _arrow_c_stream so it'll work with a minor PR to flip the if...elif branching for pandas / pyarrow dataframes in basic.py

maxzw · 2026-05-10T16:29:35Z

@Ic3fr0g that does work, but I don't think polars should be a dependency for users that just use pandas 2.0.

jameslamb · 2026-05-12T03:26:25Z

Just letting you know, looks like scikit-learn adopted narwhals in its recent nightlies and that caused a regression in support for pyarrow inputs. You may see it show up in tests here.

See https://github.com/lightgbm-org/LightGBM/pull/7268/changes#r3223494907 for a workaround and follow #7269 for a resolution.

borchero · 2026-05-16T19:38:37Z

Hey @maxzw! I have recently opened #7275 which will improve the Arrow integration of LightGBM. Once that is merged, I would like to repurpose this PR (feel free to open a new one entirely) as follows:

Via [c++][python-package] Leverage Arrow PyCapsule Interface for Arrow interoperability (including polars support) #7275, support for PyArrow-based pandas data frames has essentially been added (by relying on narwhals and the Arrow PyCapsule interface). However, that PR still sidesteps the new integration for pandas data frames by explicitly checking for them.
In this PR, we should allow pandas data frames to go through the "narwhals-route" unless they are exclusively backed by NumPy arrays (otherwise, we'd have a severe performance/memory usage regression).

Overall, this should make the PR much more straightforward.

maxzw · 2026-06-04T13:24:24Z

Hey @maxzw! I have recently opened #7275 which will improve the Arrow integration of LightGBM. Once that is merged, I would like to repurpose this PR (feel free to open a new one entirely) as follows:

Via [c++][python-package] Leverage Arrow PyCapsule Interface for Arrow interoperability #7275, support for PyArrow-based pandas data frames has essentially been added (by relying on narwhals and the Arrow PyCapsule interface). However, that PR still sidesteps the new integration for pandas data frames by explicitly checking for them.

In this PR, we should allow pandas data frames to go through the "narwhals-route" unless they are exclusively backed by NumPy arrays (otherwise, we'd have a severe performance/memory usage regression).

Overall, this should make the PR much more straightforward.

Hi @borchero thanks for the heads-up. I'll close this one and see - depending on the referenced PR's - if I can be of help filling the remaining gaps.

jameslamb added in progress feature labels May 11, 2026

maxzw added 2 commits May 13, 2026 09:20

Add support for PyArrow backed Pandas columns

2cc4190

Skip tests for old pandas versions

035f680

maxzw force-pushed the feat/pandas-pyarrow-backed-columns branch from e42f4f0 to 035f680 Compare May 13, 2026 07:20

borchero mentioned this pull request May 14, 2026

[c++][python-package] Leverage Arrow PyCapsule Interface for Arrow interoperability (including polars support) #7275

Merged

maxzw closed this Jun 4, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[python-package] Add support for PyArrow-backed Pandas DataFrames#7263

[python-package] Add support for PyArrow-backed Pandas DataFrames#7263
maxzw wants to merge 2 commits into
lightgbm-org:masterfrom
maxzw:feat/pandas-pyarrow-backed-columns

maxzw commented May 7, 2026 •

edited

Loading

Uh oh!

maxzw commented May 7, 2026

Uh oh!

Ic3fr0g commented May 9, 2026

Uh oh!

maxzw commented May 10, 2026

Uh oh!

jameslamb commented May 12, 2026

Uh oh!

borchero commented May 16, 2026

Uh oh!

maxzw commented Jun 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

maxzw commented May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Approach

Changes

Core Logic (basic.py)

Compatibility Layer (compat.py)

Tests

Benchmarks

Dataset Construction Time

End-to-End Training Time

Sklearn Integration

Compatibility

Open Questions

Uh oh!

maxzw commented May 7, 2026

Uh oh!

Ic3fr0g commented May 9, 2026

Uh oh!

maxzw commented May 10, 2026

Uh oh!

jameslamb commented May 12, 2026

Uh oh!

borchero commented May 16, 2026

Uh oh!

maxzw commented Jun 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

maxzw commented May 7, 2026 •

edited

Loading