Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion dist/unshimmed-common-from-single-shim.txt
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ com/nvidia/spark/rapids/iceberg/spark/GpuSparkUtil$.class
com/nvidia/spark/rapids/iceberg/spark/RapidsSparkCatalog.class
com/nvidia/spark/rapids/iceberg/spark/RapidsSparkSessionCatalog.class
com/nvidia/spark/rapids/iceberg/spark/source/RapidsSparkTable.class
org/apache/iceberg/aws/s3/IcebergS3InputFileAccess.class
org/apache/iceberg/data/GpuFileHelpers.class
org/apache/iceberg/io/GpuClusteredWriterBridge.class
org/apache/iceberg/io/GpuFanoutWriterBridge.class
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,28 +32,30 @@
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.URI;
import java.util.List;
import java.util.OptionalLong;

/**
* S3-backed {@link RapidsInputFile} that delegates byte-range reads to
* {@link IcebergS3RangeCopier}. The supplied {@link FileIO} is only used for
* its property map and any per-prefix storage-credential overlays.
*
* <p>The package-private S3 file access is isolated in {@link IcebergS3InputFileAccess}.
*/
public final class IcebergS3InputFile implements RapidsInputFile {
private static final Logger LOG = LoggerFactory.getLogger(IcebergS3InputFile.class);

private final IcebergInputFile delegate;
private final URI s3Uri;
private final String s3Bucket;
private final String s3Key;
private final IcebergS3Client icebergS3Client;

private IcebergS3InputFile(
IcebergInputFile delegate, URI s3Uri, IcebergS3Client icebergS3Client) {
IcebergInputFile delegate,
String s3Bucket,
String s3Key,
IcebergS3Client icebergS3Client) {
this.delegate = delegate;
this.s3Uri = s3Uri;
this.s3Bucket = s3Bucket;
this.s3Key = s3Key;
this.icebergS3Client = icebergS3Client;
}

Expand All @@ -64,25 +66,27 @@ public static RapidsInputFile maybeCreate(InputFile inputFile, FileIO fileIO) {
if (!RapidsInputFiles.isS3PerfEnabled()) {
return delegate;
}
URI s3Uri = IcebergS3InputFileAccess.s3Uri(inputFile);
if (s3Uri == null) {
if (!(inputFile instanceof BaseS3File)) {
return delegate;
}
S3URI s3Uri = ((BaseS3File) inputFile).uri();
String s3Bucket = s3Uri.bucket();
String s3Key = s3Uri.key();
// Iceberg < 1.7 does not have SupportsStorageCredentials; ShimUtils returns
// the per-prefix credential overlays (or an empty map on 1.6).
IcebergS3Client icebergS3Client = IcebergS3RangeCopier.resolveClient(
s3Uri.toString(),
inputFile.location(),
fileIO.properties(),
ShimUtils.storageCredentialOverlays(fileIO));
if (icebergS3Client == null) {
if (TaskContext.get() != null) {
GpuTaskMetrics$.MODULE$.get().recordPerfioS3IcebergFallback();
}
LOG.debug("IcebergS3RangeCopier path disabled for {}", s3Uri);
LOG.debug("IcebergS3RangeCopier path disabled for {}", inputFile.location());
return delegate;
}
LOG.debug("IcebergS3RangeCopier path active for {}", s3Uri);
return new IcebergS3InputFile(delegate, s3Uri, icebergS3Client);
LOG.debug("IcebergS3RangeCopier path active for {}", inputFile.location());
return new IcebergS3InputFile(delegate, s3Bucket, s3Key, icebergS3Client);
}

@Override
Expand Down Expand Up @@ -117,7 +121,8 @@ public InputFile getDelegate() {
@Override
public void readVectored(HostMemoryBuffer output, List<CopyRange> copyRanges)
throws IOException {
IcebergS3RangeCopier.copyToHMB(icebergS3Client, output, s3Uri, copyRanges);
IcebergS3RangeCopier.copyToHMB(
icebergS3Client, output, s3Bucket, s3Key, copyRanges);
}

/**
Expand All @@ -133,6 +138,7 @@ public void readTail(long length, HostMemoryBuffer output) throws IOException {
if (length < 0) {
throw new IllegalArgumentException("length must be non-negative");
}
IcebergS3RangeCopier.copyTailToHMB(icebergS3Client, output, s3Uri, length, /*dstOffset*/ 0L);
IcebergS3RangeCopier.copyTailToHMB(
icebergS3Client, output, s3Bucket, s3Key, length, /*dstOffset*/ 0L);
}
}

This file was deleted.

32 changes: 31 additions & 1 deletion integration_tests/src/main/python/iceberg/iceberg_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
import pytest

from asserts import assert_equal_with_local_sort, assert_gpu_and_cpu_are_equal_collect, assert_gpu_and_cpu_row_counts_equal, assert_gpu_fallback_collect, assert_spark_exception
from conftest import is_iceberg_remote_catalog
from conftest import is_iceberg_remote_catalog, is_iceberg_rest_catalog
from data_gen import *
from iceberg import get_full_table_name, iceberg_unsupported_mark, _build_tblprops, \
_BASE_TBLPROPS_SQL, create_iceberg_table
Expand Down Expand Up @@ -629,6 +629,36 @@ def setup_iceberg_table(spark):
lambda spark: spark.sql("SELECT * FROM {}".format(table)),
conf={'spark.rapids.sql.format.parquet.reader.type': reader_type})

@iceberg
@ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering
@pytest.mark.parametrize('reader_type', rapids_reader_types)
@pytest.mark.skipif(not is_iceberg_rest_catalog(),
reason="S3 path handling is exercised only with the REST catalog")
def test_iceberg_parquet_read_from_uri_invalid_s3_path(spark_tmp_table_factory, reader_type):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Assert effective PerfIO S3 enablement and preferably verify that no Iceberg fallback was recorded. Otherwise this test does not really exercise the previously-broken call path.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add a config in test to ensure perf io has been enabled.

table = get_full_table_name(spark_tmp_table_factory)
tmp_view = spark_tmp_table_factory.get()

def setup_iceberg_table(spark):
# A raw space is valid in an S3 object key but invalid in a URI. The RAPIDS reader must
# retain Iceberg's original key for the S3 request instead of using a URI-encoded path.
warehouse = spark.conf.get('spark.sql.catalog.spark_catalog.warehouse').rstrip('/')
data_path = f'{warehouse}/{spark_tmp_table_factory.get()} uri invalid path/data'
df = two_col_df(spark, long_gen, string_gen).sortWithinPartitions('b')
df.createOrReplaceTempView(tmp_view)
props = _build_tblprops({'write.data.path': data_path})
props_sql = ", ".join(f"'{k}' = '{v}'" for k, v in props.items())
spark.sql(f"CREATE TABLE {table} USING ICEBERG TBLPROPERTIES ({props_sql}) "
f"AS SELECT * FROM {tmp_view}")

with_cpu_session(setup_iceberg_table)
assert with_gpu_session(
lambda spark:
spark._jvm.com.nvidia.spark.rapids.fileio.RapidsInputFiles.isS3PerfEnabled()), \
"PerfIO S3 must be enabled at Spark startup for REST catalog tests"
assert_gpu_and_cpu_are_equal_collect(
lambda spark: spark.sql(f"SELECT * FROM {table}"),
conf={'spark.rapids.sql.format.parquet.reader.type': reader_type})

@iceberg
@ignore_order(local=True) # Iceberg plans with a thread pool and is not deterministic in file ordering
@pytest.mark.parametrize('reader_type', rapids_reader_types)
Expand Down
6 changes: 3 additions & 3 deletions jenkins/spark-tests.sh
Original file line number Diff line number Diff line change
Expand Up @@ -347,9 +347,8 @@ run_iceberg_tests() {
echo "!!! Running iceberg tests with rest catalog"
ICEBERG_REST_JARS="org.apache.iceberg:iceberg-spark-runtime-${ICEBERG_SPARK_VER}_${SCALA_BINARY_VER}:${ICEBERG_VERSION},\
org.apache.iceberg:iceberg-aws-bundle:${ICEBERG_VERSION}"
# filecache.enabled is a startup-only config, so it must be set here via
# PYSP_TEST_ env var rather than as a session-level Spark config, because
# FileCacheManager is initialized at executor startup time.
# filecache.enabled and perfio.s3.enabled are startup-only configs, so they must
# be set here via PYSP_TEST_ env vars rather than as session-level Spark configs.
env \
HOST_NAME=$PROJECT_REPO_HOST \
EXPECTED_ICEBERG_VERSION=${ICEBERG_VERSION} \
Expand All @@ -358,6 +357,7 @@ org.apache.iceberg:iceberg-aws-bundle:${ICEBERG_VERSION}"
PYSP_TEST_spark_driver_memory=1G \
PYSP_TEST_spark_executor_memory=2G \
PYSP_TEST_spark_rapids_filecache_enabled=true \
PYSP_TEST_spark_rapids_perfio_s3_enabled=true \
PYSP_TEST_spark_jars_packages="${ICEBERG_REST_JARS}" \
PYSP_TEST_spark_jars_ivySettings="${WORKSPACE}/jenkins/ivysettings.xml" \
PYSP_TEST_spark_sql_extensions="org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions" \
Expand Down
Loading