Browse Source
S3: Add tests for PyArrow with native S3 filesystem (#7508)
S3: Add tests for PyArrow with native S3 filesystem (#7508)
* PyArrow native S3 filesystem * add sse-s3 tests * update * minor * ENABLE_SSE_S3 * Update test_pyarrow_native_s3.py * clean up * refactoring * Update test_pyarrow_native_s3.pypull/7510/head
Chris Lu
6 days ago
committed by
GitHub
GitHub
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 1008 additions and 5 deletions
-
22.github/workflows/s3-parquet-tests.yml
-
92test/s3/parquet/Makefile
-
87test/s3/parquet/README.md
-
134test/s3/parquet/example_pyarrow_native.py
-
41test/s3/parquet/parquet_test_utils.py
-
383test/s3/parquet/test_pyarrow_native_s3.py
-
254test/s3/parquet/test_sse_s3_compatibility.py
@ -0,0 +1,134 @@ |
|||||
|
#!/usr/bin/env python3 |
||||
|
# /// script |
||||
|
# dependencies = [ |
||||
|
# "pyarrow>=22", |
||||
|
# "boto3>=1.28.0", |
||||
|
# ] |
||||
|
# /// |
||||
|
|
||||
|
""" |
||||
|
Simple example of using PyArrow's native S3 filesystem with SeaweedFS. |
||||
|
|
||||
|
This is a minimal example demonstrating how to write and read Parquet files |
||||
|
using PyArrow's built-in S3FileSystem without any additional dependencies |
||||
|
like s3fs. |
||||
|
|
||||
|
Usage: |
||||
|
# Set environment variables |
||||
|
export S3_ENDPOINT_URL=localhost:8333 |
||||
|
export S3_ACCESS_KEY=some_access_key1 |
||||
|
export S3_SECRET_KEY=some_secret_key1 |
||||
|
export BUCKET_NAME=test-parquet-bucket |
||||
|
|
||||
|
# Run the script |
||||
|
python3 example_pyarrow_native.py |
||||
|
|
||||
|
# Or run with uv (if available) |
||||
|
uv run example_pyarrow_native.py |
||||
|
""" |
||||
|
|
||||
|
import os |
||||
|
import secrets |
||||
|
|
||||
|
import pyarrow as pa |
||||
|
import pyarrow.dataset as pads |
||||
|
import pyarrow.fs as pafs |
||||
|
import pyarrow.parquet as pq |
||||
|
|
||||
|
from parquet_test_utils import create_sample_table |
||||
|
|
||||
|
# Configuration |
||||
|
BUCKET_NAME = os.getenv("BUCKET_NAME", "test-parquet-bucket") |
||||
|
S3_ENDPOINT_URL = os.getenv("S3_ENDPOINT_URL", "localhost:8333") |
||||
|
S3_ACCESS_KEY = os.getenv("S3_ACCESS_KEY", "some_access_key1") |
||||
|
S3_SECRET_KEY = os.getenv("S3_SECRET_KEY", "some_secret_key1") |
||||
|
|
||||
|
# Determine scheme from endpoint |
||||
|
if S3_ENDPOINT_URL.startswith("http://"): |
||||
|
scheme = "http" |
||||
|
endpoint = S3_ENDPOINT_URL[7:] |
||||
|
elif S3_ENDPOINT_URL.startswith("https://"): |
||||
|
scheme = "https" |
||||
|
endpoint = S3_ENDPOINT_URL[8:] |
||||
|
else: |
||||
|
scheme = "http" # Default to http for localhost |
||||
|
endpoint = S3_ENDPOINT_URL |
||||
|
|
||||
|
print(f"Connecting to S3 endpoint: {scheme}://{endpoint}") |
||||
|
|
||||
|
# Initialize PyArrow's NATIVE S3 filesystem |
||||
|
s3 = pafs.S3FileSystem( |
||||
|
access_key=S3_ACCESS_KEY, |
||||
|
secret_key=S3_SECRET_KEY, |
||||
|
endpoint_override=endpoint, |
||||
|
scheme=scheme, |
||||
|
allow_bucket_creation=True, |
||||
|
allow_bucket_deletion=True, |
||||
|
) |
||||
|
|
||||
|
print("✓ Connected to S3 endpoint") |
||||
|
|
||||
|
|
||||
|
# Create bucket if needed (using boto3) |
||||
|
try: |
||||
|
import boto3 |
||||
|
from botocore.exceptions import ClientError |
||||
|
|
||||
|
s3_client = boto3.client( |
||||
|
's3', |
||||
|
endpoint_url=f"{scheme}://{endpoint}", |
||||
|
aws_access_key_id=S3_ACCESS_KEY, |
||||
|
aws_secret_access_key=S3_SECRET_KEY, |
||||
|
region_name='us-east-1', |
||||
|
) |
||||
|
|
||||
|
try: |
||||
|
s3_client.head_bucket(Bucket=BUCKET_NAME) |
||||
|
print(f"✓ Bucket exists: {BUCKET_NAME}") |
||||
|
except ClientError as e: |
||||
|
if e.response['Error']['Code'] == '404': |
||||
|
print(f"Creating bucket: {BUCKET_NAME}") |
||||
|
s3_client.create_bucket(Bucket=BUCKET_NAME) |
||||
|
print(f"✓ Bucket created: {BUCKET_NAME}") |
||||
|
else: |
||||
|
raise |
||||
|
except ImportError: |
||||
|
print("Warning: boto3 not available, assuming bucket exists") |
||||
|
|
||||
|
# Generate a unique filename |
||||
|
filename = f"{BUCKET_NAME}/dataset-{secrets.token_hex(8)}/test.parquet" |
||||
|
|
||||
|
print(f"\nWriting Parquet dataset to: {filename}") |
||||
|
|
||||
|
# Write dataset |
||||
|
table = create_sample_table(200_000) |
||||
|
pads.write_dataset( |
||||
|
table, |
||||
|
filename, |
||||
|
filesystem=s3, |
||||
|
format="parquet", |
||||
|
) |
||||
|
|
||||
|
print(f"✓ Wrote {table.num_rows:,} rows") |
||||
|
|
||||
|
# Read with pq.read_table |
||||
|
print("\nReading with pq.read_table...") |
||||
|
table_read = pq.read_table(filename, filesystem=s3) |
||||
|
print(f"✓ Read {table_read.num_rows:,} rows") |
||||
|
|
||||
|
# Read with pq.ParquetDataset |
||||
|
print("\nReading with pq.ParquetDataset...") |
||||
|
dataset = pq.ParquetDataset(filename, filesystem=s3) |
||||
|
table_dataset = dataset.read() |
||||
|
print(f"✓ Read {table_dataset.num_rows:,} rows") |
||||
|
|
||||
|
# Read with pads.dataset |
||||
|
print("\nReading with pads.dataset...") |
||||
|
dataset_pads = pads.dataset(filename, filesystem=s3) |
||||
|
table_pads = dataset_pads.to_table() |
||||
|
print(f"✓ Read {table_pads.num_rows:,} rows") |
||||
|
|
||||
|
print("\n✅ All operations completed successfully!") |
||||
|
print(f"\nFile written to: {filename}") |
||||
|
print("You can verify the file using the SeaweedFS S3 API or weed shell") |
||||
|
|
||||
@ -0,0 +1,41 @@ |
|||||
|
""" |
||||
|
Shared utility functions for PyArrow Parquet tests. |
||||
|
|
||||
|
This module provides common test utilities used across multiple test scripts |
||||
|
to avoid code duplication and ensure consistency. |
||||
|
""" |
||||
|
|
||||
|
import pyarrow as pa |
||||
|
|
||||
|
|
||||
|
def create_sample_table(num_rows: int = 5) -> pa.Table: |
||||
|
"""Create a sample PyArrow table for testing. |
||||
|
|
||||
|
Args: |
||||
|
num_rows: Number of rows to generate (default: 5) |
||||
|
|
||||
|
Returns: |
||||
|
PyArrow Table with test data containing: |
||||
|
- id: int64 sequential IDs (0 to num_rows-1) |
||||
|
- name: string user names (user_0, user_1, ...) |
||||
|
- value: float64 values (id * 1.5) |
||||
|
- flag: bool alternating True/False based on even/odd id |
||||
|
|
||||
|
Example: |
||||
|
>>> table = create_sample_table(3) |
||||
|
>>> print(table) |
||||
|
pyarrow.Table |
||||
|
id: int64 |
||||
|
name: string |
||||
|
value: double |
||||
|
flag: bool |
||||
|
""" |
||||
|
return pa.table( |
||||
|
{ |
||||
|
"id": pa.array(range(num_rows), type=pa.int64()), |
||||
|
"name": pa.array([f"user_{i}" for i in range(num_rows)], type=pa.string()), |
||||
|
"value": pa.array([float(i) * 1.5 for i in range(num_rows)], type=pa.float64()), |
||||
|
"flag": pa.array([i % 2 == 0 for i in range(num_rows)], type=pa.bool_()), |
||||
|
} |
||||
|
) |
||||
|
|
||||
@ -0,0 +1,383 @@ |
|||||
|
#!/usr/bin/env python3 |
||||
|
""" |
||||
|
Test script for PyArrow's NATIVE S3 filesystem with SeaweedFS. |
||||
|
|
||||
|
This test uses PyArrow's built-in S3FileSystem (pyarrow.fs.S3FileSystem) |
||||
|
instead of s3fs, providing a pure PyArrow solution for reading and writing |
||||
|
Parquet files to S3-compatible storage. |
||||
|
|
||||
|
Requirements: |
||||
|
- pyarrow>=10.0.0 |
||||
|
|
||||
|
Environment Variables: |
||||
|
S3_ENDPOINT_URL: S3 endpoint (default: localhost:8333) |
||||
|
S3_ACCESS_KEY: S3 access key (default: some_access_key1) |
||||
|
S3_SECRET_KEY: S3 secret key (default: some_secret_key1) |
||||
|
BUCKET_NAME: S3 bucket name (default: test-parquet-bucket) |
||||
|
TEST_QUICK: Run only small/quick tests (default: 0, set to 1 for quick mode) |
||||
|
|
||||
|
Usage: |
||||
|
# Run with default environment variables |
||||
|
python3 test_pyarrow_native_s3.py |
||||
|
|
||||
|
# Run with custom environment variables |
||||
|
S3_ENDPOINT_URL=localhost:8333 \ |
||||
|
S3_ACCESS_KEY=mykey \ |
||||
|
S3_SECRET_KEY=mysecret \ |
||||
|
BUCKET_NAME=mybucket \ |
||||
|
python3 test_pyarrow_native_s3.py |
||||
|
""" |
||||
|
|
||||
|
import os |
||||
|
import secrets |
||||
|
import sys |
||||
|
import logging |
||||
|
from typing import Optional |
||||
|
|
||||
|
import pyarrow as pa |
||||
|
import pyarrow.dataset as pads |
||||
|
import pyarrow.fs as pafs |
||||
|
import pyarrow.parquet as pq |
||||
|
|
||||
|
try: |
||||
|
import boto3 |
||||
|
from botocore.exceptions import ClientError |
||||
|
HAS_BOTO3 = True |
||||
|
except ImportError: |
||||
|
HAS_BOTO3 = False |
||||
|
|
||||
|
from parquet_test_utils import create_sample_table |
||||
|
|
||||
|
logging.basicConfig(level=logging.INFO, format="%(message)s") |
||||
|
|
||||
|
# Configuration from environment variables with defaults |
||||
|
S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL", "localhost:8333") |
||||
|
S3_ACCESS_KEY = os.environ.get("S3_ACCESS_KEY", "some_access_key1") |
||||
|
S3_SECRET_KEY = os.environ.get("S3_SECRET_KEY", "some_secret_key1") |
||||
|
BUCKET_NAME = os.getenv("BUCKET_NAME", "test-parquet-bucket") |
||||
|
TEST_QUICK = os.getenv("TEST_QUICK", "0") == "1" |
||||
|
|
||||
|
# Create randomized test directory |
||||
|
TEST_RUN_ID = secrets.token_hex(8) |
||||
|
TEST_DIR = f"parquet-native-tests/{TEST_RUN_ID}" |
||||
|
|
||||
|
# Test file sizes |
||||
|
TEST_SIZES = { |
||||
|
"small": 5, |
||||
|
"large": 200_000, # This will create multiple row groups |
||||
|
} |
||||
|
|
||||
|
# Filter to only small tests if quick mode is enabled |
||||
|
if TEST_QUICK: |
||||
|
TEST_SIZES = {"small": TEST_SIZES["small"]} |
||||
|
logging.info("Quick test mode enabled - running only small tests") |
||||
|
|
||||
|
|
||||
|
def init_s3_filesystem() -> tuple[Optional[pafs.S3FileSystem], str, str]: |
||||
|
"""Initialize PyArrow's native S3 filesystem. |
||||
|
|
||||
|
Returns: |
||||
|
tuple: (S3FileSystem instance, scheme, endpoint) |
||||
|
""" |
||||
|
try: |
||||
|
logging.info("Initializing PyArrow S3FileSystem...") |
||||
|
logging.info(f" Endpoint: {S3_ENDPOINT_URL}") |
||||
|
logging.info(f" Bucket: {BUCKET_NAME}") |
||||
|
|
||||
|
# Determine scheme from endpoint |
||||
|
if S3_ENDPOINT_URL.startswith("http://"): |
||||
|
scheme = "http" |
||||
|
endpoint = S3_ENDPOINT_URL[7:] # Remove http:// |
||||
|
elif S3_ENDPOINT_URL.startswith("https://"): |
||||
|
scheme = "https" |
||||
|
endpoint = S3_ENDPOINT_URL[8:] # Remove https:// |
||||
|
else: |
||||
|
# Default to http for localhost |
||||
|
scheme = "http" |
||||
|
endpoint = S3_ENDPOINT_URL |
||||
|
|
||||
|
# Enable bucket creation and deletion for testing |
||||
|
s3 = pafs.S3FileSystem( |
||||
|
access_key=S3_ACCESS_KEY, |
||||
|
secret_key=S3_SECRET_KEY, |
||||
|
endpoint_override=endpoint, |
||||
|
scheme=scheme, |
||||
|
allow_bucket_creation=True, |
||||
|
allow_bucket_deletion=True, |
||||
|
) |
||||
|
|
||||
|
logging.info("✓ PyArrow S3FileSystem initialized successfully\n") |
||||
|
return s3, scheme, endpoint |
||||
|
except Exception: |
||||
|
logging.exception("✗ Failed to initialize PyArrow S3FileSystem") |
||||
|
return None, "", "" |
||||
|
|
||||
|
|
||||
|
def ensure_bucket_exists_boto3(scheme: str, endpoint: str) -> bool: |
||||
|
"""Ensure the test bucket exists using boto3.""" |
||||
|
if not HAS_BOTO3: |
||||
|
logging.error("boto3 is required for bucket creation") |
||||
|
return False |
||||
|
|
||||
|
try: |
||||
|
# Create boto3 client |
||||
|
endpoint_url = f"{scheme}://{endpoint}" |
||||
|
s3_client = boto3.client( |
||||
|
's3', |
||||
|
endpoint_url=endpoint_url, |
||||
|
aws_access_key_id=S3_ACCESS_KEY, |
||||
|
aws_secret_access_key=S3_SECRET_KEY, |
||||
|
region_name='us-east-1', |
||||
|
) |
||||
|
|
||||
|
# Check if bucket exists |
||||
|
try: |
||||
|
s3_client.head_bucket(Bucket=BUCKET_NAME) |
||||
|
logging.info(f"✓ Bucket exists: {BUCKET_NAME}") |
||||
|
return True |
||||
|
except ClientError as e: |
||||
|
error_code = e.response['Error']['Code'] |
||||
|
if error_code == '404': |
||||
|
# Bucket doesn't exist, create it |
||||
|
logging.info(f"Creating bucket: {BUCKET_NAME}") |
||||
|
s3_client.create_bucket(Bucket=BUCKET_NAME) |
||||
|
logging.info(f"✓ Bucket created: {BUCKET_NAME}") |
||||
|
return True |
||||
|
else: |
||||
|
raise |
||||
|
except Exception: |
||||
|
logging.exception("✗ Failed to create/check bucket") |
||||
|
return False |
||||
|
|
||||
|
|
||||
|
def ensure_bucket_exists(s3: pafs.S3FileSystem) -> bool: |
||||
|
"""Ensure the test bucket exists using PyArrow's native S3FileSystem.""" |
||||
|
try: |
||||
|
# Check if bucket exists by trying to list it |
||||
|
try: |
||||
|
file_info = s3.get_file_info(BUCKET_NAME) |
||||
|
if file_info.type == pafs.FileType.Directory: |
||||
|
logging.info(f"✓ Bucket exists: {BUCKET_NAME}") |
||||
|
return True |
||||
|
except OSError as e: |
||||
|
# OSError typically means bucket not found or network/permission issues |
||||
|
error_msg = str(e).lower() |
||||
|
if "not found" in error_msg or "does not exist" in error_msg or "nosuchbucket" in error_msg: |
||||
|
logging.debug(f"Bucket '{BUCKET_NAME}' not found, will attempt creation: {e}") |
||||
|
else: |
||||
|
# Log other OSErrors (network, auth, etc.) for debugging |
||||
|
logging.debug(f"Error checking bucket '{BUCKET_NAME}', will attempt creation anyway: {type(e).__name__}: {e}") |
||||
|
except Exception as e: |
||||
|
# Catch any other unexpected exceptions and log them |
||||
|
logging.debug(f"Unexpected error checking bucket '{BUCKET_NAME}', will attempt creation: {type(e).__name__}: {e}") |
||||
|
|
||||
|
# Try to create the bucket |
||||
|
logging.info(f"Creating bucket: {BUCKET_NAME}") |
||||
|
s3.create_dir(BUCKET_NAME) |
||||
|
logging.info(f"✓ Bucket created: {BUCKET_NAME}") |
||||
|
return True |
||||
|
except Exception: |
||||
|
logging.exception(f"✗ Failed to create/check bucket '{BUCKET_NAME}' with PyArrow") |
||||
|
return False |
||||
|
|
||||
|
|
||||
|
def test_write_and_read(s3: pafs.S3FileSystem, test_name: str, num_rows: int) -> tuple[bool, str]: |
||||
|
"""Test writing and reading a Parquet dataset using PyArrow's native S3 filesystem.""" |
||||
|
try: |
||||
|
table = create_sample_table(num_rows) |
||||
|
|
||||
|
# Write using pads.write_dataset |
||||
|
filename = f"{BUCKET_NAME}/{TEST_DIR}/{test_name}/data.parquet" |
||||
|
logging.info(f" Writing {num_rows:,} rows to {filename}...") |
||||
|
|
||||
|
pads.write_dataset( |
||||
|
table, |
||||
|
filename, |
||||
|
filesystem=s3, |
||||
|
format="parquet", |
||||
|
) |
||||
|
logging.info(" ✓ Write completed") |
||||
|
|
||||
|
# Test Method 1: Read with pq.read_table |
||||
|
logging.info(" Reading with pq.read_table...") |
||||
|
table_read = pq.read_table(filename, filesystem=s3) |
||||
|
if table_read.num_rows != num_rows: |
||||
|
return False, f"pq.read_table: Row count mismatch (expected {num_rows}, got {table_read.num_rows})" |
||||
|
|
||||
|
# Check schema first |
||||
|
if not table_read.schema.equals(table.schema): |
||||
|
return False, f"pq.read_table: Schema mismatch (expected {table.schema}, got {table_read.schema})" |
||||
|
|
||||
|
# Sort both tables by 'id' column before comparison to handle potential row order differences |
||||
|
table_sorted = table.sort_by([('id', 'ascending')]) |
||||
|
table_read_sorted = table_read.sort_by([('id', 'ascending')]) |
||||
|
|
||||
|
if not table_read_sorted.equals(table_sorted): |
||||
|
# Provide more detailed error information |
||||
|
error_details = [] |
||||
|
for col_name in table.column_names: |
||||
|
col_original = table_sorted.column(col_name) |
||||
|
col_read = table_read_sorted.column(col_name) |
||||
|
if not col_original.equals(col_read): |
||||
|
error_details.append(f"column '{col_name}' differs") |
||||
|
return False, f"pq.read_table: Table contents mismatch ({', '.join(error_details)})" |
||||
|
logging.info(f" ✓ pq.read_table: {table_read.num_rows:,} rows") |
||||
|
|
||||
|
# Test Method 2: Read with pq.ParquetDataset |
||||
|
logging.info(" Reading with pq.ParquetDataset...") |
||||
|
dataset = pq.ParquetDataset(filename, filesystem=s3) |
||||
|
table_dataset = dataset.read() |
||||
|
if table_dataset.num_rows != num_rows: |
||||
|
return False, f"pq.ParquetDataset: Row count mismatch (expected {num_rows}, got {table_dataset.num_rows})" |
||||
|
|
||||
|
# Sort before comparison |
||||
|
table_dataset_sorted = table_dataset.sort_by([('id', 'ascending')]) |
||||
|
if not table_dataset_sorted.equals(table_sorted): |
||||
|
error_details = [] |
||||
|
for col_name in table.column_names: |
||||
|
col_original = table_sorted.column(col_name) |
||||
|
col_read = table_dataset_sorted.column(col_name) |
||||
|
if not col_original.equals(col_read): |
||||
|
error_details.append(f"column '{col_name}' differs") |
||||
|
return False, f"pq.ParquetDataset: Table contents mismatch ({', '.join(error_details)})" |
||||
|
logging.info(f" ✓ pq.ParquetDataset: {table_dataset.num_rows:,} rows") |
||||
|
|
||||
|
# Test Method 3: Read with pads.dataset |
||||
|
logging.info(" Reading with pads.dataset...") |
||||
|
dataset_pads = pads.dataset(filename, filesystem=s3) |
||||
|
table_pads = dataset_pads.to_table() |
||||
|
if table_pads.num_rows != num_rows: |
||||
|
return False, f"pads.dataset: Row count mismatch (expected {num_rows}, got {table_pads.num_rows})" |
||||
|
|
||||
|
# Sort before comparison |
||||
|
table_pads_sorted = table_pads.sort_by([('id', 'ascending')]) |
||||
|
if not table_pads_sorted.equals(table_sorted): |
||||
|
error_details = [] |
||||
|
for col_name in table.column_names: |
||||
|
col_original = table_sorted.column(col_name) |
||||
|
col_read = table_pads_sorted.column(col_name) |
||||
|
if not col_original.equals(col_read): |
||||
|
error_details.append(f"column '{col_name}' differs") |
||||
|
return False, f"pads.dataset: Table contents mismatch ({', '.join(error_details)})" |
||||
|
logging.info(f" ✓ pads.dataset: {table_pads.num_rows:,} rows") |
||||
|
|
||||
|
return True, "All read methods passed" |
||||
|
|
||||
|
except Exception as exc: |
||||
|
logging.exception(" ✗ Test failed") |
||||
|
return False, f"{type(exc).__name__}: {exc}" |
||||
|
|
||||
|
|
||||
|
def cleanup_test_files(s3: pafs.S3FileSystem) -> None: |
||||
|
"""Clean up test files from S3. |
||||
|
|
||||
|
Note: We cannot use s3.delete_dir() directly because SeaweedFS uses implicit |
||||
|
directories (path prefixes without physical directory objects). PyArrow's |
||||
|
delete_dir() attempts to delete the directory marker itself, which fails with |
||||
|
"INTERNAL_FAILURE" on SeaweedFS. Instead, we list and delete files individually, |
||||
|
letting implicit directories disappear automatically. |
||||
|
""" |
||||
|
try: |
||||
|
test_path = f"{BUCKET_NAME}/{TEST_DIR}" |
||||
|
logging.info(f"Cleaning up test directory: {test_path}") |
||||
|
|
||||
|
# List and delete files individually to handle implicit directories |
||||
|
try: |
||||
|
file_selector = pafs.FileSelector(test_path, recursive=True) |
||||
|
files = s3.get_file_info(file_selector) |
||||
|
|
||||
|
# Delete files first (not directories) |
||||
|
for file_info in files: |
||||
|
if file_info.type == pafs.FileType.File: |
||||
|
s3.delete_file(file_info.path) |
||||
|
logging.debug(f" Deleted file: {file_info.path}") |
||||
|
|
||||
|
logging.info("✓ Test directory cleaned up") |
||||
|
except OSError as e: |
||||
|
# Handle the case where the path doesn't exist or is inaccessible |
||||
|
if "does not exist" in str(e).lower() or "not found" in str(e).lower(): |
||||
|
logging.info("✓ Test directory already clean or doesn't exist") |
||||
|
else: |
||||
|
raise |
||||
|
except Exception: |
||||
|
logging.exception("Failed to cleanup test directory") |
||||
|
|
||||
|
|
||||
|
def main(): |
||||
|
"""Run all tests with PyArrow's native S3 filesystem.""" |
||||
|
print("=" * 80) |
||||
|
print("PyArrow Native S3 Filesystem Tests for SeaweedFS") |
||||
|
print("Testing Parquet Files with Multiple Row Groups") |
||||
|
if TEST_QUICK: |
||||
|
print("*** QUICK TEST MODE - Small files only ***") |
||||
|
print("=" * 80 + "\n") |
||||
|
|
||||
|
print("Configuration:") |
||||
|
print(f" S3 Endpoint: {S3_ENDPOINT_URL}") |
||||
|
print(f" Access Key: {S3_ACCESS_KEY}") |
||||
|
print(f" Bucket: {BUCKET_NAME}") |
||||
|
print(f" Test Directory: {TEST_DIR}") |
||||
|
print(f" Quick Mode: {'Yes (small files only)' if TEST_QUICK else 'No (all file sizes)'}") |
||||
|
print(f" PyArrow Version: {pa.__version__}") |
||||
|
print() |
||||
|
|
||||
|
# Initialize S3 filesystem |
||||
|
s3, scheme, endpoint = init_s3_filesystem() |
||||
|
if s3 is None: |
||||
|
print("Cannot proceed without S3 connection") |
||||
|
return 1 |
||||
|
|
||||
|
# Ensure bucket exists - try PyArrow first, fall back to boto3 |
||||
|
bucket_created = ensure_bucket_exists(s3) |
||||
|
if not bucket_created: |
||||
|
logging.info("Trying to create bucket with boto3...") |
||||
|
bucket_created = ensure_bucket_exists_boto3(scheme, endpoint) |
||||
|
|
||||
|
if not bucket_created: |
||||
|
print("Cannot proceed without bucket") |
||||
|
return 1 |
||||
|
|
||||
|
results = [] |
||||
|
|
||||
|
# Test all file sizes |
||||
|
for size_name, num_rows in TEST_SIZES.items(): |
||||
|
print(f"\n{'='*80}") |
||||
|
print(f"Testing with {size_name} files ({num_rows:,} rows)") |
||||
|
print(f"{'='*80}\n") |
||||
|
|
||||
|
test_name = f"{size_name}_test" |
||||
|
success, message = test_write_and_read(s3, test_name, num_rows) |
||||
|
results.append((test_name, success, message)) |
||||
|
|
||||
|
status = "✓ PASS" if success else "✗ FAIL" |
||||
|
print(f"\n{status}: {message}\n") |
||||
|
|
||||
|
# Summary |
||||
|
print("\n" + "=" * 80) |
||||
|
print("SUMMARY") |
||||
|
print("=" * 80) |
||||
|
passed = sum(1 for _, success, _ in results if success) |
||||
|
total = len(results) |
||||
|
print(f"\nTotal: {passed}/{total} passed\n") |
||||
|
|
||||
|
for test_name, success, message in results: |
||||
|
status = "✓" if success else "✗" |
||||
|
print(f" {status} {test_name}: {message}") |
||||
|
|
||||
|
print("\n" + "=" * 80) |
||||
|
if passed == total: |
||||
|
print("✓ ALL TESTS PASSED!") |
||||
|
else: |
||||
|
print(f"✗ {total - passed} test(s) failed") |
||||
|
|
||||
|
print("=" * 80 + "\n") |
||||
|
|
||||
|
# Cleanup |
||||
|
cleanup_test_files(s3) |
||||
|
|
||||
|
return 0 if passed == total else 1 |
||||
|
|
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
sys.exit(main()) |
||||
|
|
||||
@ -0,0 +1,254 @@ |
|||||
|
#!/usr/bin/env python3 |
||||
|
""" |
||||
|
Test script for SSE-S3 compatibility with PyArrow native S3 filesystem. |
||||
|
|
||||
|
This test specifically targets the SSE-S3 multipart upload bug where |
||||
|
SeaweedFS panics with "bad IV length" when reading multipart uploads |
||||
|
that were encrypted with bucket-default SSE-S3. |
||||
|
|
||||
|
Requirements: |
||||
|
- pyarrow>=10.0.0 |
||||
|
- boto3>=1.28.0 |
||||
|
|
||||
|
Environment Variables: |
||||
|
S3_ENDPOINT_URL: S3 endpoint (default: localhost:8333) |
||||
|
S3_ACCESS_KEY: S3 access key (default: some_access_key1) |
||||
|
S3_SECRET_KEY: S3 secret key (default: some_secret_key1) |
||||
|
BUCKET_NAME: S3 bucket name (default: test-parquet-bucket) |
||||
|
|
||||
|
Usage: |
||||
|
# Start SeaweedFS with SSE-S3 enabled |
||||
|
make start-seaweedfs-ci ENABLE_SSE_S3=true |
||||
|
|
||||
|
# Run the test |
||||
|
python3 test_sse_s3_compatibility.py |
||||
|
""" |
||||
|
|
||||
|
import os |
||||
|
import secrets |
||||
|
import sys |
||||
|
import logging |
||||
|
from typing import Optional |
||||
|
|
||||
|
import pyarrow as pa |
||||
|
import pyarrow.dataset as pads |
||||
|
import pyarrow.fs as pafs |
||||
|
import pyarrow.parquet as pq |
||||
|
|
||||
|
try: |
||||
|
import boto3 |
||||
|
from botocore.exceptions import ClientError |
||||
|
HAS_BOTO3 = True |
||||
|
except ImportError: |
||||
|
HAS_BOTO3 = False |
||||
|
logging.exception("boto3 is required for this test") |
||||
|
sys.exit(1) |
||||
|
|
||||
|
from parquet_test_utils import create_sample_table |
||||
|
|
||||
|
logging.basicConfig(level=logging.INFO, format="%(message)s") |
||||
|
|
||||
|
# Configuration |
||||
|
S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL", "localhost:8333") |
||||
|
S3_ACCESS_KEY = os.environ.get("S3_ACCESS_KEY", "some_access_key1") |
||||
|
S3_SECRET_KEY = os.environ.get("S3_SECRET_KEY", "some_secret_key1") |
||||
|
BUCKET_NAME = os.getenv("BUCKET_NAME", "test-parquet-bucket") |
||||
|
|
||||
|
TEST_RUN_ID = secrets.token_hex(8) |
||||
|
TEST_DIR = f"sse-s3-tests/{TEST_RUN_ID}" |
||||
|
|
||||
|
# Test sizes designed to trigger multipart uploads |
||||
|
# PyArrow typically uses 5MB chunks, so these sizes should trigger multipart |
||||
|
TEST_SIZES = { |
||||
|
"tiny": 10, # Single part |
||||
|
"small": 1_000, # Single part |
||||
|
"medium": 50_000, # Single part (~1.5MB) |
||||
|
"large": 200_000, # Multiple parts (~6MB) |
||||
|
"very_large": 500_000, # Multiple parts (~15MB) |
||||
|
} |
||||
|
|
||||
|
|
||||
|
def init_s3_filesystem() -> tuple[Optional[pafs.S3FileSystem], str, str]: |
||||
|
"""Initialize PyArrow's native S3 filesystem.""" |
||||
|
try: |
||||
|
logging.info("Initializing PyArrow S3FileSystem...") |
||||
|
|
||||
|
# Determine scheme from endpoint |
||||
|
if S3_ENDPOINT_URL.startswith("http://"): |
||||
|
scheme = "http" |
||||
|
endpoint = S3_ENDPOINT_URL[7:] |
||||
|
elif S3_ENDPOINT_URL.startswith("https://"): |
||||
|
scheme = "https" |
||||
|
endpoint = S3_ENDPOINT_URL[8:] |
||||
|
else: |
||||
|
scheme = "http" |
||||
|
endpoint = S3_ENDPOINT_URL |
||||
|
|
||||
|
s3 = pafs.S3FileSystem( |
||||
|
access_key=S3_ACCESS_KEY, |
||||
|
secret_key=S3_SECRET_KEY, |
||||
|
endpoint_override=endpoint, |
||||
|
scheme=scheme, |
||||
|
allow_bucket_creation=True, |
||||
|
allow_bucket_deletion=True, |
||||
|
) |
||||
|
|
||||
|
logging.info("✓ PyArrow S3FileSystem initialized\n") |
||||
|
return s3, scheme, endpoint |
||||
|
except Exception: |
||||
|
logging.exception("✗ Failed to initialize PyArrow S3FileSystem") |
||||
|
return None, "", "" |
||||
|
|
||||
|
|
||||
|
def ensure_bucket_exists(scheme: str, endpoint: str) -> bool: |
||||
|
"""Ensure the test bucket exists using boto3.""" |
||||
|
try: |
||||
|
endpoint_url = f"{scheme}://{endpoint}" |
||||
|
s3_client = boto3.client( |
||||
|
's3', |
||||
|
endpoint_url=endpoint_url, |
||||
|
aws_access_key_id=S3_ACCESS_KEY, |
||||
|
aws_secret_access_key=S3_SECRET_KEY, |
||||
|
region_name='us-east-1', |
||||
|
) |
||||
|
|
||||
|
try: |
||||
|
s3_client.head_bucket(Bucket=BUCKET_NAME) |
||||
|
logging.info(f"✓ Bucket exists: {BUCKET_NAME}") |
||||
|
except ClientError as e: |
||||
|
error_code = e.response['Error']['Code'] |
||||
|
if error_code == '404': |
||||
|
logging.info(f"Creating bucket: {BUCKET_NAME}") |
||||
|
s3_client.create_bucket(Bucket=BUCKET_NAME) |
||||
|
logging.info(f"✓ Bucket created: {BUCKET_NAME}") |
||||
|
else: |
||||
|
logging.exception("✗ Failed to access bucket") |
||||
|
return False |
||||
|
|
||||
|
# Note: SeaweedFS doesn't support GetBucketEncryption API |
||||
|
# so we can't verify if SSE-S3 is enabled via API |
||||
|
# We assume it's configured correctly in the s3.json config file |
||||
|
logging.info("✓ Assuming SSE-S3 is configured in s3.json") |
||||
|
return True |
||||
|
|
||||
|
except Exception: |
||||
|
logging.exception("✗ Failed to check bucket") |
||||
|
return False |
||||
|
|
||||
|
|
||||
|
def test_write_read_with_sse( |
||||
|
s3: pafs.S3FileSystem, |
||||
|
test_name: str, |
||||
|
num_rows: int |
||||
|
) -> tuple[bool, str, int]: |
||||
|
"""Test writing and reading with SSE-S3 encryption.""" |
||||
|
try: |
||||
|
table = create_sample_table(num_rows) |
||||
|
filename = f"{BUCKET_NAME}/{TEST_DIR}/{test_name}/data.parquet" |
||||
|
|
||||
|
logging.info(f" Writing {num_rows:,} rows...") |
||||
|
pads.write_dataset( |
||||
|
table, |
||||
|
filename, |
||||
|
filesystem=s3, |
||||
|
format="parquet", |
||||
|
) |
||||
|
|
||||
|
logging.info(" Reading back...") |
||||
|
table_read = pq.read_table(filename, filesystem=s3) |
||||
|
|
||||
|
if table_read.num_rows != num_rows: |
||||
|
return False, f"Row count mismatch: {table_read.num_rows} != {num_rows}", 0 |
||||
|
|
||||
|
return True, "Success", table_read.num_rows |
||||
|
|
||||
|
except Exception as e: |
||||
|
error_msg = f"{type(e).__name__}: {e!s}" |
||||
|
logging.exception(" ✗ Failed") |
||||
|
return False, error_msg, 0 |
||||
|
|
||||
|
|
||||
|
def main(): |
||||
|
"""Run SSE-S3 compatibility tests.""" |
||||
|
print("=" * 80) |
||||
|
print("SSE-S3 Compatibility Tests for PyArrow Native S3") |
||||
|
print("Testing Multipart Upload Encryption") |
||||
|
print("=" * 80 + "\n") |
||||
|
|
||||
|
print("Configuration:") |
||||
|
print(f" S3 Endpoint: {S3_ENDPOINT_URL}") |
||||
|
print(f" Bucket: {BUCKET_NAME}") |
||||
|
print(f" Test Directory: {TEST_DIR}") |
||||
|
print(f" PyArrow Version: {pa.__version__}") |
||||
|
print() |
||||
|
|
||||
|
# Initialize |
||||
|
s3, scheme, endpoint = init_s3_filesystem() |
||||
|
if s3 is None: |
||||
|
print("Cannot proceed without S3 connection") |
||||
|
return 1 |
||||
|
|
||||
|
# Check bucket and SSE-S3 |
||||
|
if not ensure_bucket_exists(scheme, endpoint): |
||||
|
print("\n⚠ WARNING: Failed to access or create the test bucket!") |
||||
|
print("This test requires a reachable bucket with SSE-S3 enabled.") |
||||
|
print("Please ensure SeaweedFS is running with: make start-seaweedfs-ci ENABLE_SSE_S3=true") |
||||
|
return 1 |
||||
|
|
||||
|
print() |
||||
|
results = [] |
||||
|
|
||||
|
# Test all sizes |
||||
|
for size_name, num_rows in TEST_SIZES.items(): |
||||
|
print(f"\n{'='*80}") |
||||
|
print(f"Testing {size_name} dataset ({num_rows:,} rows)") |
||||
|
print(f"{'='*80}") |
||||
|
|
||||
|
success, message, rows_read = test_write_read_with_sse( |
||||
|
s3, size_name, num_rows |
||||
|
) |
||||
|
results.append((size_name, num_rows, success, message, rows_read)) |
||||
|
|
||||
|
if success: |
||||
|
print(f" ✓ SUCCESS: Read {rows_read:,} rows") |
||||
|
else: |
||||
|
print(f" ✗ FAILED: {message}") |
||||
|
|
||||
|
# Summary |
||||
|
print("\n" + "=" * 80) |
||||
|
print("SUMMARY") |
||||
|
print("=" * 80) |
||||
|
|
||||
|
passed = sum(1 for _, _, success, _, _ in results if success) |
||||
|
total = len(results) |
||||
|
print(f"\nTotal: {passed}/{total} tests passed\n") |
||||
|
|
||||
|
print(f"{'Size':<15} {'Rows':>10} {'Status':<10} {'Rows Read':>10} {'Message':<40}") |
||||
|
print("-" * 90) |
||||
|
for size_name, num_rows, success, message, rows_read in results: |
||||
|
status = "✓ PASS" if success else "✗ FAIL" |
||||
|
rows_str = f"{rows_read:,}" if success else "N/A" |
||||
|
print(f"{size_name:<15} {num_rows:>10,} {status:<10} {rows_str:>10} {message[:40]}") |
||||
|
|
||||
|
print("\n" + "=" * 80) |
||||
|
if passed == total: |
||||
|
print("✓ ALL TESTS PASSED WITH SSE-S3!") |
||||
|
print("\nThis means:") |
||||
|
print(" - SSE-S3 encryption is working correctly") |
||||
|
print(" - PyArrow native S3 filesystem is compatible") |
||||
|
print(" - Multipart uploads are handled properly") |
||||
|
else: |
||||
|
print(f"✗ {total - passed} test(s) failed") |
||||
|
print("\nPossible issues:") |
||||
|
print(" - SSE-S3 multipart upload bug with empty IV") |
||||
|
print(" - Encryption/decryption mismatch") |
||||
|
print(" - File corruption during upload") |
||||
|
|
||||
|
print("=" * 80 + "\n") |
||||
|
|
||||
|
return 0 if passed == total else 1 |
||||
|
|
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
sys.exit(main()) |
||||
|
|
||||
Write
Preview
Loading…
Cancel
Save
Reference in new issue