snapshot of https://github.com/seaweedfs/seaweedfs.git (master, 2026-09-16; blobs > 131072B omitted)

test/s3/parquet/test_pyarrow_native_s3.py

@@ -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())
+