diff --git a/src/dataset.jl b/src/dataset.jl index cf9d4ea..c4a4d8d 100644 --- a/src/dataset.jl +++ b/src/dataset.jl @@ -154,6 +154,8 @@ struct DatasetPartitions end end +Base.IteratorSize(::Type{DatasetPartitions}) = Base.SizeUnknown() + function iterated_partition(partitions::DatasetPartitions, cursor) partition = nothing walker, root, files, fileidx, step = cursor diff --git a/test/runtests.jl b/test/runtests.jl index 10a3455..a03605f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,4 +12,5 @@ const julia_parcompat = joinpath(artifact"julia_parcompat", readdir(artifact"jul include("test_cursors.jl") include("test_writer.jl") include("test_dataset_types.jl") + include("test_dataset_partitions.jl") end diff --git a/test/test_dataset_partitions.jl b/test/test_dataset_partitions.jl new file mode 100644 index 0000000..db19705 --- /dev/null +++ b/test/test_dataset_partitions.jl @@ -0,0 +1,36 @@ +using Parquet, Tables, Test + +@testset "Collect dataset partitions" begin + fixture = joinpath(@__DIR__, "booltest", "alltypes_plain.snappy.parquet") + mktempdir() do path + for name in ("a.parquet", "b.parquet", "_common_metadata", "_metadata") + cp(fixture, joinpath(path, name)) + end + for factory in (Parquet.Dataset, Parquet.read_parquet), (filter, expected) in ( + (p -> true, 2), + (p -> endswith(p, "a.parquet"), 1), + (p -> false, 0), + ) + dataset = factory(path; filter=filter) + try + nmanual = 0 + for partition in Tables.partitions(dataset) + nmanual += 1 + close(partition) + end + @test nmanual == expected + partitions = collect(Tables.partitions(dataset)) + try + @test partitions isa AbstractVector + @test length(partitions) == expected + @test all(p -> p isa Parquet.Table, partitions) + @test all(p -> Tables.getcolumn(p, :id) == Int32[6, 7], partitions) + finally + foreach(close, partitions) + end + finally + close(dataset) + end + end + end +end