diff --git a/.gitignore b/.gitignore
index 4e890ab..6cc3760 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,3 +5,6 @@
 /docs/build/
 Manifest.toml
 /.vscode/*
+*.arrow
+*.m
+*.csv
diff --git a/Project.toml b/Project.toml
index 1329a90..699437b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -9,12 +9,17 @@ JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 ParametricOptInterface = "0ce4ce61-57bf-432b-a095-efac525d185e"
 
 [compat]
-julia = "1"
+JuMP = "1"
+Arrow = "2"
+ParametricOptInterface = "0.5"
+julia = "1.6"
 
 [extras]
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
+PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "DelimitedFiles", "HiGHS"]
+test = ["Test", "DelimitedFiles", "Downloads", "HiGHS", "PowerModels"]
diff --git a/docs/make.jl b/docs/make.jl
index d8e279a..6cbb92d 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -14,12 +14,7 @@ makedocs(;
         edit_link="main",
         assets=String[],
     ),
-    pages=[
-        "Home" => "index.md",
-    ],
+    pages=["Home" => "index.md"],
 )
 
-deploydocs(;
-    repo="github.com/andrewrosemberg/L2O.jl",
-    devbranch="main",
-)
+deploydocs(; repo="github.com/andrewrosemberg/L2O.jl", devbranch="main")
diff --git a/examples/powermodels/pg_lib.jl b/examples/powermodels/pg_lib.jl
new file mode 100644
index 0000000..a097570
--- /dev/null
+++ b/examples/powermodels/pg_lib.jl
@@ -0,0 +1,79 @@
+using Downloads
+using PowerModels
+using JuMP, HiGHS
+import ParametricOptInterface as POI
+
+"""
+    createvarrefs!(sp::JuMP.Model, pm::AbstractPowerModel)
+
+create ref for anonimous variables on model
+"""
+function createvarrefs!(sp::JuMP.Model, pm::AbstractPowerModel)
+    for listvarref in values(PowerModels.var(pm))
+        for variableref in values(listvarref)
+            if typeof(variableref) == JuMP.VariableRef
+                sp[Symbol(name(variableref))] = variableref
+            end
+        end
+    end
+end
+
+"""
+    generate_dataset_pglib(data_dir::AbstractString, case_name::AbstractString; download_files::Bool=true, filetype::Type{RecorderFile},
+    num_p::Int=10
+)
+
+Generate dataset for pglib case_name with num_p problems and save it in data_dir
+"""
+function generate_dataset_pglib(
+    data_dir,
+    case_name;
+    filetype=CSVFile,
+    download_files=true,
+    num_p=10,
+)
+    case_file_path = joinpath(data_dir, case_name)
+    if download_files && !isfile(case_file_path)
+        Downloads.download(
+            "https://raw.githubusercontent.com/power-grid-lib/pglib-opf/01681386d084d8bd03b429abcd1ee6966f68b9a3/" *
+            case_name,
+            case_file_path,
+        )
+    end
+
+    # Read data
+    network_data = PowerModels.parse_file(case_file_path)
+
+    # The problem to iterate over
+    model = Model(() -> POI.Optimizer(HiGHS.Optimizer()))
+
+    # Save original load value and Link POI
+    original_load = network_data["load"]["1"]["pd"]
+    network_data["load"]["1"]["pd"] = p = @variable(model, _p in POI.Parameter(1.0))
+
+    # Instantiate the model
+    pm = instantiate_model(
+        network_data,
+        DCPPowerModel,
+        PowerModels.build_opf;
+        setting=Dict("output" => Dict("duals" => true)),
+        jump_model=model,
+    )
+
+    # The problem iterator
+    problem_iterator = ProblemIterator(
+        collect(1:num_p), Dict(p => collect(1.0:num_p) .* original_load)
+    )
+
+    # Create ref for anonimous variables on model
+    createvarrefs!(model, pm)
+
+    # Solve the problem and return the number of successfull solves
+    file = joinpath(data_dir, "test.$(string(filetype))")
+    number_generators = length(network_data["gen"])
+    recorder = Recorder{filetype}(
+        file; primal_variables=[Symbol("0_pg[$i]") for i in 1:number_generators]
+    )
+    return solve_batch(model, problem_iterator, recorder), number_generators
+end
+
diff --git a/src/L2O.jl b/src/L2O.jl
index 498fe41..5e8cf52 100644
--- a/src/L2O.jl
+++ b/src/L2O.jl
@@ -7,7 +7,6 @@ import Base: string
 
 export ArrowFile, CSVFile, ProblemIterator, Recorder, solve_batch
 
-
 include("datasetgen.jl")
 include("csvrecorder.jl")
 include("arrowrecorder.jl")
diff --git a/src/arrowrecorder.jl b/src/arrowrecorder.jl
index 71a1a44..e690c16 100644
--- a/src/arrowrecorder.jl
+++ b/src/arrowrecorder.jl
@@ -8,23 +8,24 @@ Base.string(::Type{ArrowFile}) = "arrow"
 Record optimization problem solution to an Arrow file.
 """
 function record(recorder::Recorder{ArrowFile}, model::JuMP.Model, id::T) where {T<:Integer}
-    if !isfile(recorder.filename)
-        Arrow.append(
-            recorder.filename, (;
-                id = T[], 
-                zip(recorder.primal_variables, fill(Float64[], length(recorder.primal_variables)))..., 
-                zip(Symbol.("dual_" .* string.(recorder.dual_variables)), fill(Float64[], length(recorder.dual_variables)))..., 
-            )
-        )
-    end
-
-    Arrow.append(
-        recorder.filename, (;
-            id = [id], 
-            zip(recorder.primal_variables, [[MOI.get(model, MOI.VariablePrimal(), model[p])] for p in recorder.primal_variables])..., 
-            zip(Symbol.("dual_" .* string.(recorder.dual_variables)), [[MOI.get(model, MOI.ConstraintDual(), model[p])] for p in recorder.dual_variables])..., 
-        )
+    return Arrow.append(
+        recorder.filename,
+        (;
+            id=[id],
+            zip(
+                recorder.primal_variables,
+                [
+                    [MOI.get(model, MOI.VariablePrimal(), model[p])] for
+                    p in recorder.primal_variables
+                ],
+            )...,
+            zip(
+                Symbol.("dual_" .* string.(recorder.dual_variables)),
+                [
+                    [MOI.get(model, MOI.ConstraintDual(), model[p])] for
+                    p in recorder.dual_variables
+                ],
+            )...,
+        ),
     )
-
-
-end
\ No newline at end of file
+end
diff --git a/src/csvrecorder.jl b/src/csvrecorder.jl
index 75a3c0a..788f0d0 100644
--- a/src/csvrecorder.jl
+++ b/src/csvrecorder.jl
@@ -32,4 +32,4 @@ function record(recorder::Recorder{CSVFile}, model::JuMP.Model, id::Int64)
         end
         write(f, "\n")
     end
-end
\ No newline at end of file
+end
diff --git a/src/datasetgen.jl b/src/datasetgen.jl
index 988061e..f10d871 100644
--- a/src/datasetgen.jl
+++ b/src/datasetgen.jl
@@ -11,7 +11,12 @@ mutable struct Recorder{T<:RecorderFile}
     dual_variables::AbstractArray{Symbol}
     filterfn::Function
 
-    function Recorder{T}(filename::String; primal_variables=[], dual_variables=[], filterfn=(model)-> termination_status(model) == MOI.OPTIMAL) where T<:RecorderFile
+    function Recorder{T}(
+        filename::String;
+        primal_variables=[],
+        dual_variables=[],
+        filterfn=(model) -> termination_status(model) == MOI.OPTIMAL,
+    ) where {T<:RecorderFile}
         return new{T}(filename, primal_variables, dual_variables, filterfn)
     end
 end
@@ -21,14 +26,16 @@ end
 
 Iterator for optimization problem instances.
 """
-struct ProblemIterator{T<:Real, Z<:Integer}
+struct ProblemIterator{T<:Real,Z<:Integer}
     ids::Vector{Z}
-    pairs::Dict{VariableRef, Vector{T}}
-    function ProblemIterator(ids::Vector{Z}, pairs::Dict{VariableRef, Vector{T}}) where {T<:Real, Z<:Integer}
+    pairs::Dict{VariableRef,Vector{T}}
+    function ProblemIterator(
+        ids::Vector{Z}, pairs::Dict{VariableRef,Vector{T}}
+    ) where {T<:Real,Z<:Integer}
         for (p, val) in pairs
             @assert length(ids) == length(val)
         end
-        return new{T, Z}(ids, pairs)
+        return new{T,Z}(ids, pairs)
     end
 end
 
@@ -37,17 +44,8 @@ end
 
 Update the value of a parameter in a JuMP model.
 """
-function update_model!(model::JuMP.Model, p::VariableRef, val::T) where {T<:Real}
-    MOI.set(model, POI.ParameterValue(), p, val)
-end
-
-"""
-    update_model!(model::JuMP.Model, p::VariableRef, val::AbstractArray{Real})
-
-Update the value of a parameter in a JuMP model.
-"""
-function update_model!(model::JuMP.Model, p::VariableRef, val::AbstractArray{T}) where {T<:Real}
-    MOI.set(model, POI.ParameterValue(), p, val)
+function update_model!(model::JuMP.Model, p::VariableRef, val)
+    return MOI.set(model, POI.ParameterValue(), p, val)
 end
 
 """
@@ -66,13 +64,16 @@ end
 
 Solve an optimization problem and record the solution.
 """
-function solve_and_record(model::JuMP.Model, problem_iterator::ProblemIterator, recorder::Recorder, idx::Integer)
+function solve_and_record(
+    model::JuMP.Model, problem_iterator::ProblemIterator, recorder::Recorder, idx::Integer
+)
     update_model!(model, problem_iterator.pairs, idx)
     optimize!(model)
     if recorder.filterfn(model)
         record(recorder, model, problem_iterator.ids[idx])
+        return 1
     end
-    return nothing
+    return 0
 end
 
 """
@@ -80,9 +81,11 @@ end
 
 Solve a batch of optimization problems and record the solutions.
 """
-function solve_batch(model::JuMP.Model, problem_iterator::ProblemIterator, recorder::Recorder)
-    for idx in 1:length(problem_iterator.ids)
-        solve_and_record(model, problem_iterator, recorder, idx)
-    end
-    return nothing
+function solve_batch(
+    model::JuMP.Model, problem_iterator::ProblemIterator, recorder::Recorder
+)
+    return sum(
+        solve_and_record(model, problem_iterator, recorder, idx) for
+        idx in 1:length(problem_iterator.ids)
+    ) / length(problem_iterator.ids)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 0c75066..1758bb9 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,23 +5,46 @@ using DelimitedFiles
 using JuMP, HiGHS
 import ParametricOptInterface as POI
 
-function testdataset_gen(path)
-    @testset "Dataset generation: $filetype" for filetype in [CSVFile, ArrowFile]
+include(
+    joinpath(
+        dirname(dirname(@__FILE__)), "examples", "powermodels", "pg_lib.jl"
+    ),
+)
+
+"""
+    testdataset_gen(path::String)
+
+Test dataset generation for different filetypes
+"""
+function testdataset_gen(path::String)
+    @testset "Type: $filetype" for filetype in [CSVFile, ArrowFile]
+        # The problem to iterate over
         model = Model(() -> POI.Optimizer(HiGHS.Optimizer()))
         @variable(model, x)
         p = @variable(model, _p in POI.Parameter(1.0))
         @constraint(model, cons, x + _p >= 3)
         @objective(model, Min, 2x)
 
+        # The problem iterator
         num_p = 10
+        @test_throws AssertionError ProblemIterator(collect(1:num_p), Dict(p => collect(1.0:3.0)))
+        @test_throws MethodError ProblemIterator(collect(1.0:3.0), Dict(p => collect(1.0:3.0)))
         problem_iterator = ProblemIterator(collect(1:num_p), Dict(p => collect(1.0:num_p)))
-        file = joinpath(path, "test.$(string(filetype))")
-        recorder = Recorder{filetype}(file, primal_variables=[:x], dual_variables=[:cons])
+
+        # The recorder
+        file = joinpath(path, "test.$(string(filetype))") # file path
+        @test Recorder{filetype}(file; primal_variables=[:x]) isa Recorder{filetype}
+        @test Recorder{filetype}(file; dual_variables=[:cons]) isa Recorder{filetype}
+        recorder = Recorder{filetype}(file; primal_variables=[:x], dual_variables=[:cons])
+
+        # Solve all problems and record solutions
         solve_batch(model, problem_iterator, recorder)
+
+        # Check if file exists and has the correct number of rows and columns
         if filetype == CSVFile
             file1 = joinpath(path, "test.csv")
             @test isfile(file1)
-            @test length(readdlm(file1, ',')[:, 1]) == num_p+1
+            @test length(readdlm(file1, ',')[:, 1]) == num_p + 1
             @test length(readdlm(file1, ',')[1, :]) == 3
             rm(file1)
         else
@@ -35,7 +58,30 @@ function testdataset_gen(path)
 end
 
 @testset "L2O.jl" begin
-    mktempdir() do path
-        testdataset_gen(path)
+    @testset "Dataset Generation" begin
+        mktempdir() do path
+            # Different filetypes
+            testdataset_gen(path)
+            # Pglib
+            @testset "pg_lib case" begin
+                # Define test case from pglib
+                case_name = "pglib_opf_case5_pjm.m"
+
+                # Define number of problems
+                num_p = 10
+
+                # Generate dataset
+                success_solves, number_generators = generate_dataset_pglib(
+                    path, case_name; num_p=num_p
+                )
+
+                # Check if the number of successfull solves is equal to the number of problems saved
+                file = joinpath(path, "test.csv")
+                @test isfile(file)
+                @test length(readdlm(file, ',')[:, 1]) == num_p * success_solves + 1
+                @test length(readdlm(file, ',')[1, :]) == number_generators + 1
+                rm(file)
+            end
+        end
     end
 end