def test_map_element_at_ansi_fail(data_gen):
message = "org.apache.spark.SparkNoSuchElementException" if (not is_before_spark_330() or is_databricks104_or_later()) else "java.util.NoSuchElementException"
# For 3.3.0+ strictIndexOperator should not affect element_at
test_conf=copy_and_update(ansi_enabled_conf, {'spark.sql.ansi.strictIndexOperator': 'false'})
assert_gpu_and_cpu_error(
lambda spark: unary_op_df(spark, data_gen).selectExpr(
'element_at(a, "NOT_FOUND")').collect(),
conf=test_conf,
error_message=message)
def test_array_element_at_ansi_fail_invalid_index(index):
message = "ArrayIndexOutOfBoundsException" if is_before_spark_330() else "SparkArrayIndexOutOfBoundsException"
if isinstance(index, int):
test_func = lambda spark: unary_op_df(spark, ArrayGen(int_gen)).select(
element_at(col('a'), index)).collect()
else:
test_func = lambda spark: two_col_df(spark, ArrayGen(int_gen), index).selectExpr(
'element_at(a, b)').collect()
# For 3.3.0+ strictIndexOperator should not affect element_at
test_conf=copy_and_update(ansi_enabled_conf, {'spark.sql.ansi.strictIndexOperator': 'false'})
assert_gpu_and_cpu_error(
test_func,
conf=test_conf,
error_message=message)
@pytest.mark.parametrize('data_gen', vals, ids=idfn)
def test_timeadd(data_gen):
days, seconds = data_gen
assert_gpu_and_cpu_are_equal_collect(
# We are starting at year 0005 to make sure we don't go before year 0001
# and beyond year 10000 while doing TimeAdd
lambda spark: unary_op_df(
spark,
TimestampGen(start=datetime(5, 1, 1, tzinfo=timezone.utc),
end=datetime(15, 1, 1, tzinfo=timezone.utc)),
seed=1).selectExpr("a + (interval {} days {} seconds)".format(
days, seconds)))
@pytest.mark.skipif(
is_before_spark_330(),
reason='DayTimeInterval is not supported before Pyspark 3.3.0')
def test_timeadd_daytime_column():
gen_list = [
# timestamp column max year is 1000
('t', TimestampGen(end=datetime(1000, 1, 1, tzinfo=timezone.utc))),
# max days is 8000 year, so added result will not be out of range
('d',
DayTimeIntervalGen(min_value=timedelta(days=0),
max_value=timedelta(days=8000 * 365)))
]
assert_gpu_and_cpu_are_equal_collect(
lambda spark: gen_df(spark, gen_list).selectExpr(
"t + d", "t + INTERVAL '1 02:03:04' DAY TO SECOND"))
all_confs = reader_confs.copy()
all_confs.update({'spark.sql.sources.useV1SourceList': v1_enabled_list,
'spark.sql.caseSensitive': case_sensitive})
# This is a hack to get the type in a slightly less verbose way
extra_struct_gen = StructGen([('nested_col', orc_gen), ("nested_non_existing", orc_gen)])
extra_gen_list = [("top_pri", orc_gen),
("top_non_existing_mid", orc_gen),
("TOP_AR", ArrayGen(extra_struct_gen, max_length=10)),
("top_ST", extra_struct_gen),
("top_non_existing_end", orc_gen)]
rs = StructGen(extra_gen_list, nullable=False).data_type
assert_gpu_and_cpu_are_equal_collect(
lambda spark : spark.read.schema(rs).orc(data_path),
conf=all_confs)
@pytest.mark.skipif(is_before_spark_330(), reason='Hidden file metadata columns are a new feature of Spark 330')
@allow_non_gpu(any = True)
@pytest.mark.parametrize('metadata_column', ["file_path", "file_name", "file_size", "file_modification_time"])
def test_orc_scan_with_hidden_metadata_fallback(spark_tmp_path, metadata_column):
data_path = spark_tmp_path + "/hidden_metadata.orc"
with_cpu_session(lambda spark : spark.range(10) \
.selectExpr("id", "id % 3 as p") \
.write \
.partitionBy("p") \
.mode("overwrite") \
.orc(data_path))
def do_orc_scan(spark):
df = spark.read.orc(data_path).selectExpr("id", "_metadata.{}".format(metadata_column))
return df
# Merged "test_nested_array_item" with this one since arrays as literals is supported
@pytest.mark.parametrize('data_gen', array_item_test_gens, ids=idfn)
def test_array_item(data_gen):
assert_gpu_and_cpu_are_equal_collect(
lambda spark: two_col_df(spark, data_gen, array_index_gen).selectExpr(
'a[b]',
'a[0]',
'a[1]',
'a[null]',
'a[3]',
'a[50]',
'a[-1]'))
# No need to test this for multiple data types for array. Only one is enough
@pytest.mark.skipif(is_before_spark_330(), reason="'strictIndexOperator' is introduced from Spark 3.3.0")
@pytest.mark.parametrize('strict_index_enabled', [True, False])
@pytest.mark.parametrize('index', [-2, 100, array_neg_index_gen, array_out_index_gen], ids=idfn)
def test_array_item_with_strict_index(strict_index_enabled, index):
message = "SparkArrayIndexOutOfBoundsException"
if isinstance(index, int):
test_df = lambda spark: unary_op_df(spark, ArrayGen(int_gen)).select(col('a')[index])
else:
test_df = lambda spark: two_col_df(spark, ArrayGen(int_gen), index).selectExpr('a[b]')
test_conf=copy_and_update(
ansi_enabled_conf, {'spark.sql.ansi.strictIndexOperator': strict_index_enabled})
if strict_index_enabled:
assert_gpu_and_cpu_error(
lambda spark: test_df(spark).collect(),
@pytest.mark.parametrize('parquet_gens', parquet_write_gens_list, ids=idfn)
def test_write_empty_parquet_round_trip(spark_tmp_path, parquet_gens):
def create_empty_df(spark, path):
gen_list = [('_c' + str(i), gen) for i, gen in enumerate(parquet_gens)]
return gen_df(spark, gen_list, length=0).write.parquet(path)
data_path = spark_tmp_path + '/PARQUET_DATA'
assert_gpu_and_cpu_writes_are_equal_collect(
create_empty_df,
lambda spark, path: spark.read.parquet(path),
data_path,
conf=writer_confs)
# should fallback when trying to write field ID metadata
@pytest.mark.skipif(is_before_spark_330(),
reason='Field ID is not supported before Spark 330')
@allow_non_gpu('DataWritingCommandExec')
def test_parquet_write_field_id(spark_tmp_path):
data_path = spark_tmp_path + '/PARQUET_DATA'
schema = StructType([
StructField("c1", IntegerType(), metadata={'parquet.field.id': 1}),
])
data = [
(1, ),
(2, ),
(3, ),
]
assert_gpu_fallback_write(
lambda spark, path: spark.createDataFrame(data, schema).coalesce(
1).write.mode("overwrite").parquet(path),
"spark.rapids.sql.castFloatToString.enabled" : "true",
"spark.sql.legacy.castComplexTypesToString.enabled": legacy}
)
# The bug SPARK-37451 only affects the following versions
def is_neg_dec_scale_bug_version():
return ("3.1.1" <= spark_version() < "3.1.3") or ("3.2.0" <= spark_version() < "3.2.1")
@pytest.mark.skipif(is_neg_dec_scale_bug_version(), reason="RAPIDS doesn't support casting string to decimal for negative scale decimal in this version of Spark because of SPARK-37451")
def test_cast_string_to_negative_scale_decimal():
assert_gpu_and_cpu_are_equal_collect(
lambda spark: unary_op_df(spark, StringGen("[0-9]{9}")).select(
f.col('a').cast(DecimalType(8, -3))))
@pytest.mark.skipif(is_before_spark_330(), reason="ansi cast throws exception only in 3.3.0+")
@pytest.mark.parametrize('type', [DoubleType(), FloatType()])
def test_cast_double_to_timestamp(type):
def fun(spark):
data=[float("inf"),float("-inf"),float("nan")]
df = spark.createDataFrame(data, DoubleType())
return df.select(f.col('value').cast(TimestampType())).collect()
assert_gpu_and_cpu_error(fun, {"spark.sql.ansi.enabled": True}, "java.time.DateTimeException")
@pytest.mark.skipif(is_before_spark_330(), reason='DayTimeInterval is not supported before Pyspark 3.3.0')
def test_cast_day_time_interval_to_string():
_assert_cast_to_string_equal(DayTimeIntervalGen(start_field='day', end_field='day', special_cases=[MIN_DAY_TIME_INTERVAL, MAX_DAY_TIME_INTERVAL, timedelta(seconds=0)]), {})
_assert_cast_to_string_equal(DayTimeIntervalGen(start_field='day', end_field='hour', special_cases=[MIN_DAY_TIME_INTERVAL, MAX_DAY_TIME_INTERVAL, timedelta(seconds=0)]), {})
_assert_cast_to_string_equal(DayTimeIntervalGen(start_field='day', end_field='minute', special_cases=[MIN_DAY_TIME_INTERVAL, MAX_DAY_TIME_INTERVAL, timedelta(seconds=0)]), {})
_assert_cast_to_string_equal(DayTimeIntervalGen(start_field='day', end_field='second', special_cases=[MIN_DAY_TIME_INTERVAL, MAX_DAY_TIME_INTERVAL, timedelta(seconds=0)]), {})
_assert_cast_to_string_equal(DayTimeIntervalGen(start_field='hour', end_field='hour', special_cases=[MIN_DAY_TIME_INTERVAL, MAX_DAY_TIME_INTERVAL, timedelta(seconds=0)]), {})