chrislu
2 weeks ago
3 changed files with 506 additions and 6 deletions
-
35test/s3/parquet/Makefile
-
31test/s3/parquet/README.md
-
446test/s3/parquet/test_cross_filesystem_compatibility.py
@ -0,0 +1,446 @@ |
|||||
|
#!/usr/bin/env python3 |
||||
|
""" |
||||
|
Cross-filesystem compatibility tests for PyArrow Parquet files. |
||||
|
|
||||
|
This test verifies that Parquet files written using one filesystem implementation |
||||
|
(s3fs or PyArrow native S3) can be correctly read using the other implementation. |
||||
|
|
||||
|
Test Matrix: |
||||
|
- Write with s3fs → Read with PyArrow native S3 |
||||
|
- Write with PyArrow native S3 → Read with s3fs |
||||
|
|
||||
|
Requirements: |
||||
|
- pyarrow>=22.0.0 |
||||
|
- s3fs>=2024.12.0 |
||||
|
- boto3>=1.40.0 |
||||
|
|
||||
|
Environment Variables: |
||||
|
S3_ENDPOINT_URL: S3 endpoint (default: http://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_cross_filesystem_compatibility.py |
||||
|
|
||||
|
# Run with custom environment variables |
||||
|
S3_ENDPOINT_URL=http://localhost:8333 \ |
||||
|
S3_ACCESS_KEY=mykey \ |
||||
|
S3_SECRET_KEY=mysecret \ |
||||
|
BUCKET_NAME=mybucket \ |
||||
|
python3 test_cross_filesystem_compatibility.py |
||||
|
""" |
||||
|
|
||||
|
import os |
||||
|
import secrets |
||||
|
import sys |
||||
|
import logging |
||||
|
from typing import Optional, Tuple |
||||
|
|
||||
|
import pyarrow as pa |
||||
|
import pyarrow.dataset as pads |
||||
|
import pyarrow.fs as pafs |
||||
|
import pyarrow.parquet as pq |
||||
|
import s3fs |
||||
|
|
||||
|
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", "http://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-cross-fs-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_s3fs() -> Optional[s3fs.S3FileSystem]: |
||||
|
"""Initialize s3fs filesystem.""" |
||||
|
try: |
||||
|
logging.info("Initializing s3fs...") |
||||
|
fs = s3fs.S3FileSystem( |
||||
|
client_kwargs={"endpoint_url": S3_ENDPOINT_URL}, |
||||
|
key=S3_ACCESS_KEY, |
||||
|
secret=S3_SECRET_KEY, |
||||
|
use_listings_cache=False, |
||||
|
) |
||||
|
logging.info("✓ s3fs initialized successfully") |
||||
|
return fs |
||||
|
except Exception: |
||||
|
logging.exception("✗ Failed to initialize s3fs") |
||||
|
return None |
||||
|
|
||||
|
|
||||
|
def init_pyarrow_s3() -> Tuple[Optional[pafs.S3FileSystem], str, str]: |
||||
|
"""Initialize PyArrow's native S3 filesystem. |
||||
|
|
||||
|
Returns: |
||||
|
tuple: (S3FileSystem instance, scheme, endpoint) |
||||
|
""" |
||||
|
try: |
||||
|
logging.info("Initializing PyArrow S3FileSystem...") |
||||
|
|
||||
|
# 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") |
||||
|
return s3, scheme, endpoint |
||||
|
except Exception: |
||||
|
logging.exception("✗ Failed to initialize PyArrow S3FileSystem") |
||||
|
return None, "", "" |
||||
|
|
||||
|
|
||||
|
def ensure_bucket_exists(s3fs_fs: s3fs.S3FileSystem, pyarrow_s3: pafs.S3FileSystem) -> bool: |
||||
|
"""Ensure the test bucket exists using s3fs.""" |
||||
|
try: |
||||
|
if not s3fs_fs.exists(BUCKET_NAME): |
||||
|
logging.info(f"Creating bucket: {BUCKET_NAME}") |
||||
|
try: |
||||
|
s3fs_fs.mkdir(BUCKET_NAME) |
||||
|
logging.info(f"✓ Bucket created: {BUCKET_NAME}") |
||||
|
except FileExistsError: |
||||
|
# Bucket was created between the check and mkdir call |
||||
|
logging.info(f"✓ Bucket exists: {BUCKET_NAME}") |
||||
|
else: |
||||
|
logging.info(f"✓ Bucket exists: {BUCKET_NAME}") |
||||
|
return True |
||||
|
except Exception: |
||||
|
logging.exception("✗ Failed to create/check bucket") |
||||
|
return False |
||||
|
|
||||
|
|
||||
|
def write_with_s3fs(table: pa.Table, path: str, s3fs_fs: s3fs.S3FileSystem) -> bool: |
||||
|
"""Write Parquet file using s3fs filesystem.""" |
||||
|
try: |
||||
|
pads.write_dataset(table, path, format="parquet", filesystem=s3fs_fs) |
||||
|
return True |
||||
|
except Exception: |
||||
|
logging.exception("✗ Failed to write with s3fs") |
||||
|
return False |
||||
|
|
||||
|
|
||||
|
def write_with_pyarrow_s3(table: pa.Table, path: str, pyarrow_s3: pafs.S3FileSystem) -> bool: |
||||
|
"""Write Parquet file using PyArrow native S3 filesystem.""" |
||||
|
try: |
||||
|
pads.write_dataset(table, path, format="parquet", filesystem=pyarrow_s3) |
||||
|
return True |
||||
|
except Exception: |
||||
|
logging.exception("✗ Failed to write with PyArrow S3") |
||||
|
return False |
||||
|
|
||||
|
|
||||
|
def read_with_s3fs(path: str, s3fs_fs: s3fs.S3FileSystem) -> Tuple[bool, Optional[pa.Table], str]: |
||||
|
"""Read Parquet file using s3fs filesystem with multiple methods.""" |
||||
|
errors = [] |
||||
|
|
||||
|
# Try pq.read_table |
||||
|
try: |
||||
|
table = pq.read_table(path, filesystem=s3fs_fs) |
||||
|
return True, table, "pq.read_table" |
||||
|
except Exception as e: |
||||
|
errors.append(f"pq.read_table: {type(e).__name__}: {e}") |
||||
|
|
||||
|
# Try pq.ParquetDataset |
||||
|
try: |
||||
|
dataset = pq.ParquetDataset(path, filesystem=s3fs_fs) |
||||
|
table = dataset.read() |
||||
|
return True, table, "pq.ParquetDataset" |
||||
|
except Exception as e: |
||||
|
errors.append(f"pq.ParquetDataset: {type(e).__name__}: {e}") |
||||
|
|
||||
|
# Try pads.dataset |
||||
|
try: |
||||
|
dataset = pads.dataset(path, format="parquet", filesystem=s3fs_fs) |
||||
|
table = dataset.to_table() |
||||
|
return True, table, "pads.dataset" |
||||
|
except Exception as e: |
||||
|
errors.append(f"pads.dataset: {type(e).__name__}: {e}") |
||||
|
|
||||
|
return False, None, " | ".join(errors) |
||||
|
|
||||
|
|
||||
|
def read_with_pyarrow_s3(path: str, pyarrow_s3: pafs.S3FileSystem) -> Tuple[bool, Optional[pa.Table], str]: |
||||
|
"""Read Parquet file using PyArrow native S3 filesystem with multiple methods.""" |
||||
|
errors = [] |
||||
|
|
||||
|
# Try pq.read_table |
||||
|
try: |
||||
|
table = pq.read_table(path, filesystem=pyarrow_s3) |
||||
|
return True, table, "pq.read_table" |
||||
|
except Exception as e: |
||||
|
errors.append(f"pq.read_table: {type(e).__name__}: {e}") |
||||
|
|
||||
|
# Try pq.ParquetDataset |
||||
|
try: |
||||
|
dataset = pq.ParquetDataset(path, filesystem=pyarrow_s3) |
||||
|
table = dataset.read() |
||||
|
return True, table, "pq.ParquetDataset" |
||||
|
except Exception as e: |
||||
|
errors.append(f"pq.ParquetDataset: {type(e).__name__}: {e}") |
||||
|
|
||||
|
# Try pads.dataset |
||||
|
try: |
||||
|
dataset = pads.dataset(path, filesystem=pyarrow_s3) |
||||
|
table = dataset.to_table() |
||||
|
return True, table, "pads.dataset" |
||||
|
except Exception as e: |
||||
|
errors.append(f"pads.dataset: {type(e).__name__}: {e}") |
||||
|
|
||||
|
return False, None, " | ".join(errors) |
||||
|
|
||||
|
|
||||
|
def verify_table_integrity(original: pa.Table, read: pa.Table) -> Tuple[bool, str]: |
||||
|
"""Verify that read table matches the original table.""" |
||||
|
# Check row count |
||||
|
if read.num_rows != original.num_rows: |
||||
|
return False, f"Row count mismatch: expected {original.num_rows}, got {read.num_rows}" |
||||
|
|
||||
|
# Check schema |
||||
|
if not read.schema.equals(original.schema): |
||||
|
return False, f"Schema mismatch: expected {original.schema}, got {read.schema}" |
||||
|
|
||||
|
# Sort both tables by 'id' column before comparison to handle potential row order differences |
||||
|
original_sorted = original.sort_by([('id', 'ascending')]) |
||||
|
read_sorted = read.sort_by([('id', 'ascending')]) |
||||
|
|
||||
|
# Check data equality |
||||
|
if not read_sorted.equals(original_sorted): |
||||
|
# Provide detailed error information |
||||
|
error_details = [] |
||||
|
for col_name in original.column_names: |
||||
|
col_original = original_sorted.column(col_name) |
||||
|
col_read = read_sorted.column(col_name) |
||||
|
if not col_original.equals(col_read): |
||||
|
error_details.append(f"column '{col_name}' differs") |
||||
|
return False, f"Data mismatch: {', '.join(error_details)}" |
||||
|
|
||||
|
return True, "Data verified successfully" |
||||
|
|
||||
|
|
||||
|
def test_write_s3fs_read_pyarrow( |
||||
|
test_name: str, |
||||
|
num_rows: int, |
||||
|
s3fs_fs: s3fs.S3FileSystem, |
||||
|
pyarrow_s3: pafs.S3FileSystem |
||||
|
) -> Tuple[bool, str]: |
||||
|
"""Test: Write with s3fs, read with PyArrow native S3.""" |
||||
|
try: |
||||
|
table = create_sample_table(num_rows) |
||||
|
path = f"{BUCKET_NAME}/{TEST_DIR}/{test_name}/data.parquet" |
||||
|
|
||||
|
# Write with s3fs |
||||
|
logging.info(f" Writing {num_rows:,} rows with s3fs to {path}...") |
||||
|
if not write_with_s3fs(table, path, s3fs_fs): |
||||
|
return False, "Write with s3fs failed" |
||||
|
logging.info(" ✓ Write completed") |
||||
|
|
||||
|
# Read with PyArrow native S3 |
||||
|
logging.info(" Reading with PyArrow native S3...") |
||||
|
success, read_table, method = read_with_pyarrow_s3(path, pyarrow_s3) |
||||
|
if not success: |
||||
|
return False, f"Read with PyArrow S3 failed: {method}" |
||||
|
logging.info(f" ✓ Read {read_table.num_rows:,} rows using {method}") |
||||
|
|
||||
|
# Verify data integrity |
||||
|
verify_success, verify_msg = verify_table_integrity(table, read_table) |
||||
|
if not verify_success: |
||||
|
return False, f"Verification failed: {verify_msg}" |
||||
|
logging.info(f" ✓ {verify_msg}") |
||||
|
|
||||
|
return True, f"s3fs→PyArrow: {method}" |
||||
|
|
||||
|
except Exception as e: |
||||
|
logging.exception(" ✗ Test failed") |
||||
|
return False, f"{type(e).__name__}: {e}" |
||||
|
|
||||
|
|
||||
|
def test_write_pyarrow_read_s3fs( |
||||
|
test_name: str, |
||||
|
num_rows: int, |
||||
|
s3fs_fs: s3fs.S3FileSystem, |
||||
|
pyarrow_s3: pafs.S3FileSystem |
||||
|
) -> Tuple[bool, str]: |
||||
|
"""Test: Write with PyArrow native S3, read with s3fs.""" |
||||
|
try: |
||||
|
table = create_sample_table(num_rows) |
||||
|
path = f"{BUCKET_NAME}/{TEST_DIR}/{test_name}/data.parquet" |
||||
|
|
||||
|
# Write with PyArrow native S3 |
||||
|
logging.info(f" Writing {num_rows:,} rows with PyArrow native S3 to {path}...") |
||||
|
if not write_with_pyarrow_s3(table, path, pyarrow_s3): |
||||
|
return False, "Write with PyArrow S3 failed" |
||||
|
logging.info(" ✓ Write completed") |
||||
|
|
||||
|
# Read with s3fs |
||||
|
logging.info(" Reading with s3fs...") |
||||
|
success, read_table, method = read_with_s3fs(path, s3fs_fs) |
||||
|
if not success: |
||||
|
return False, f"Read with s3fs failed: {method}" |
||||
|
logging.info(f" ✓ Read {read_table.num_rows:,} rows using {method}") |
||||
|
|
||||
|
# Verify data integrity |
||||
|
verify_success, verify_msg = verify_table_integrity(table, read_table) |
||||
|
if not verify_success: |
||||
|
return False, f"Verification failed: {verify_msg}" |
||||
|
logging.info(f" ✓ {verify_msg}") |
||||
|
|
||||
|
return True, f"PyArrow→s3fs: {method}" |
||||
|
|
||||
|
except Exception as e: |
||||
|
logging.exception(" ✗ Test failed") |
||||
|
return False, f"{type(e).__name__}: {e}" |
||||
|
|
||||
|
|
||||
|
def cleanup_test_files(s3fs_fs: s3fs.S3FileSystem) -> None: |
||||
|
"""Clean up test files from S3.""" |
||||
|
try: |
||||
|
test_path = f"{BUCKET_NAME}/{TEST_DIR}" |
||||
|
if s3fs_fs.exists(test_path): |
||||
|
logging.info(f"Cleaning up test directory: {test_path}") |
||||
|
s3fs_fs.rm(test_path, recursive=True) |
||||
|
logging.info("✓ Test directory cleaned up") |
||||
|
except Exception: |
||||
|
logging.exception("Failed to cleanup test directory") |
||||
|
|
||||
|
|
||||
|
def main(): |
||||
|
"""Run cross-filesystem compatibility tests.""" |
||||
|
print("=" * 80) |
||||
|
print("Cross-Filesystem Compatibility Tests for PyArrow Parquet") |
||||
|
print("Testing: s3fs ↔ PyArrow Native S3 Filesystem") |
||||
|
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 both filesystems |
||||
|
s3fs_fs = init_s3fs() |
||||
|
if s3fs_fs is None: |
||||
|
print("Cannot proceed without s3fs connection") |
||||
|
return 1 |
||||
|
|
||||
|
pyarrow_s3, scheme, endpoint = init_pyarrow_s3() |
||||
|
if pyarrow_s3 is None: |
||||
|
print("Cannot proceed without PyArrow S3 connection") |
||||
|
return 1 |
||||
|
|
||||
|
print() |
||||
|
|
||||
|
# Ensure bucket exists |
||||
|
if not ensure_bucket_exists(s3fs_fs, pyarrow_s3): |
||||
|
print("Cannot proceed without bucket") |
||||
|
return 1 |
||||
|
|
||||
|
print() |
||||
|
|
||||
|
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 1: Write with s3fs, read with PyArrow native S3 |
||||
|
test_name = f"{size_name}_s3fs_to_pyarrow" |
||||
|
print(f"Test: Write with s3fs → Read with PyArrow native S3") |
||||
|
success, message = test_write_s3fs_read_pyarrow( |
||||
|
test_name, num_rows, s3fs_fs, pyarrow_s3 |
||||
|
) |
||||
|
results.append((test_name, success, message)) |
||||
|
status = "✓ PASS" if success else "✗ FAIL" |
||||
|
print(f"{status}: {message}\n") |
||||
|
|
||||
|
# Test 2: Write with PyArrow native S3, read with s3fs |
||||
|
test_name = f"{size_name}_pyarrow_to_s3fs" |
||||
|
print(f"Test: Write with PyArrow native S3 → Read with s3fs") |
||||
|
success, message = test_write_pyarrow_read_s3fs( |
||||
|
test_name, num_rows, s3fs_fs, pyarrow_s3 |
||||
|
) |
||||
|
results.append((test_name, success, message)) |
||||
|
status = "✓ PASS" if success else "✗ FAIL" |
||||
|
print(f"{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 CROSS-FILESYSTEM TESTS PASSED!") |
||||
|
print("\nConclusion: Files written with s3fs and PyArrow native S3 are") |
||||
|
print("fully compatible and can be read by either filesystem implementation.") |
||||
|
else: |
||||
|
print(f"✗ {total - passed} test(s) failed") |
||||
|
|
||||
|
print("=" * 80 + "\n") |
||||
|
|
||||
|
# Cleanup |
||||
|
cleanup_test_files(s3fs_fs) |
||||
|
|
||||
|
return 0 if passed == total else 1 |
||||
|
|
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
sys.exit(main()) |
||||
|
|
||||
Write
Preview
Loading…
Cancel
Save
Reference in new issue