[BUG] limit function is producing inconsistent result when type is Byte, Long, Boolean and Timestamp

[razajafri]

Describe the bug
Using the limit inside a function that reads a parquet file results in inconsistent results.

Steps/Code to reproduce bug

@pytest.mark.parametrize('data_gen', [ByteGen(), LongGen(), BooleanGen(), TimestampGen()], ids=idfn)
@pytest.mark.parametrize('ts_write', ['INT96', 'TIMESTAMP_MICROS', 'TIMESTAMP_MILLIS'])
@pytest.mark.parametrize('enableVectorized', ['true', 'false'], ids=idfn)
@allow_non_gpu('CollectLimitExec', 'DataWritingCommandExec')
@ignore_order
def test_cache_columnar(spark_tmp_path, data_gen, enableVectorized, ts_write):
    data_path_gpu = spark_tmp_path + '/PARQUET_DATA'
    def read_parquet_cached(data_path):
        def write_read_parquet_cached(spark):
            df = unary_op_df(spark, data_gen).select(f.col("a"))
            df.write.mode('overwrite').parquet(data_path)
            cached = spark.read.parquet(data_path)
            cached.count()
            return debug_df(cached.limit(50))
        return write_read_parquet_cached
    # rapids-spark doesn't support LEGACY read for parquet
    conf={'spark.sql.legacy.parquet.datetimeRebaseModeInWrite': 'CORRECTED',
          'spark.sql.legacy.parquet.datetimeRebaseModeInRead' : 'CORRECTED',
          'spark.sql.inMemoryColumnarStorage.enableVectorizedReader' : enableVectorized,
          'spark.sql.parquet.outputTimestampType': ts_write}
    assert_gpu_and_cpu_are_equal_collect(read_parquet_cached(data_path_gpu), conf)

Expected behavior
The above test should pass against Spark 3.0.0. and Spark 3.1.0

[tgravescs]

the ones I see failing are:
FAILED src/main/python/parquet_test.py::test_cache_columnar[true-INT96-Timestamp][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED src/main/python/parquet_test.py::test_cache_columnar[false-INT96-Timestamp][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]

Which are both INT96 types, so is this a dup of #1007 or were you seeing more fail?

[razajafri]

@tgravescs run against Spark 3.1.0. I will update the description. I see a lot more.

FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-INT96-Byte][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-INT96-Long][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-INT96-Boolean][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-INT96-Timestamp][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-TIMESTAMP_MICROS-Byte][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-TIMESTAMP_MICROS-Long][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-TIMESTAMP_MICROS-Boolean][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-TIMESTAMP_MILLIS-Byte][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-TIMESTAMP_MILLIS-Long][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[true-TIMESTAMP_MILLIS-Boolean][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-INT96-Byte][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-INT96-Long][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-INT96-Boolean][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-INT96-Timestamp][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-TIMESTAMP_MICROS-Byte][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-TIMESTAMP_MICROS-Long][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-TIMESTAMP_MICROS-Boolean][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-TIMESTAMP_MILLIS-Byte][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-TIMESTAMP_MILLIS-Long][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]
FAILED ../../src/main/python/cache_test.py::test_cache_columnar_123[false-TIMESTAMP_MILLIS-Boolean][IGNORE_ORDER, ALLOW_NON_GPU(CollectLimitExec,DataWritingCommandExec)]

[tgravescs]

ok looks like it requires more then 1 executor or by default --master local[*]

[tgravescs]

this is because you are writing different data on each iteration (cpu and gpu):

def read_parquet_cached(data_path):
        def write_read_parquet_cached(spark):
            df = unary_op_df(spark, data_gen).select(f.col("a"))
            df.write.mode('overwrite').parquet(data_path)
            cached = spark.read.parquet(data_path)
            cached.count()
            return debug_df(cached.limit(50))
        return write_read_parquet_cached

So the data is going to be different because wrote different data each time.
You need to write a single file in advance and read on both cpu and gpu