change Parquet for DuckDB in EpiAutoGP (#1061)

* change Paquet for DuckDB

* add method for DUCKDB

* duckdb test julia -> julia

* rm R date conversion step

* add local test for interop date types

* pin to julia 1.11 + move interop test to end

* add project instantiation for interop test

Author: SamuelBrand1

.github/workflows/test.yaml

@@ -29,6 +29,7 @@ jobs:
           XLA_FLAGS: --xla_force_host_platform_device_count=2
         run: |
           uv run pytest \
+            --ignore=pipelines/tests/test_epiautogp_parquet_interop.py \
             --cov=pipelines --cov-report=term --cov-report=xml:pipelines-coverage.xml .
 
       - name: Upload results to Codecov
@@ -95,6 +96,8 @@ jobs:
 
       - name: Set up Julia
         uses: julia-actions/setup-julia@v3
+        with:
+          version: "1.11"
 
       - name: Run tests
         uses: julia-actions/julia-runtest@v1
@@ -116,3 +119,17 @@ jobs:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           verbose: true
+
+      - name: Set up R with hewr
+        uses: ./.github/actions/setup-hewr
+
+      - name: Set up Pyrenew-HEW via UV
+        uses: ./.github/actions/setup-pyrenew-hew
+
+      - name: Instantiate EpiAutoGP Julia project
+        run: |
+          julia --project=EpiAutoGP -e 'using Pkg; Pkg.instantiate()'
+
+      - name: Run EpiAutoGP parquet interop test
+        run: |
+          uv run pytest pipelines/tests/test_epiautogp_parquet_interop.py -q

EpiAutoGP/Project.toml

@@ -1,34 +1,36 @@
 name = "EpiAutoGP"
 uuid = "c2940010-6b35-4be1-8bbf-9fa0d9979e50"
-version = "0.1.0"
 authors = ["Sam Brand (USI1) <usi1@cdc.gov>"]
+version = "0.1.0"
 
 [deps]
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+DBInterface = "a10d1c49-ce27-4219-8d33-6db1a4562965"
 DataFramesMeta = "1313f7d8-7da2-5740-9ea0-a2ca25f37964"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+DuckDB = "d2f5444f-75bc-4fdf-ac35-56f514c445e1"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 NowcastAutoGP = "7e9f7f4b-f590-4c14-8324-de4fcbed18f7"
-Parquet = "626c502c-15b0-58ad-a749-f091afb673ae"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StructTypes = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
-NowcastAutoGP = {rev = "v0.3.0", url = "https://github.com/CDCgov/NowcastAutoGP"}
+NowcastAutoGP = {rev = "v0.3.0", url = "https://github.com/CDCgov/NowcastAutoGP.git"}
 
 [compat]
 ArgParse = "1.2.0"
 CSV = "0.10.15"
+DBInterface = "2.6.1"
 DataFramesMeta = "0.15.4"
 Dates = "1.11.0"
+DuckDB = "1.5.2"
 JSON3 = "1.14.3"
 Logging = "1.11.0"
 NowcastAutoGP = "0.3.0"
-Parquet = "0.8.6"
 Random = "1.11.0"
 Statistics = "1.11.1"
 StructTypes = "1.11.0"

EpiAutoGP/src/EpiAutoGP.jl

@@ -1,6 +1,6 @@
 module EpiAutoGP
 using NowcastAutoGP # Core modeling package
-using CSV, DataFramesMeta, Dates, JSON3, StructTypes, Parquet # Data handling packages
+using CSV, DataFramesMeta, Dates, DBInterface, DuckDB, JSON3, StructTypes # Data handling packages
 using ArgParse # Command-line argument parsing
 using Statistics # For modeling functions
 

EpiAutoGP/src/output.jl

@@ -117,6 +117,24 @@ function create_forecast_df(results::NamedTuple, output_type::PipelineOutput)
     return forecast_df
 end
 
+function _quote_duckdb_string(value::AbstractString)
+    return "'" * replace(value, "'" => "''") * "'"
+end
+
+function _write_parquet_with_duckdb(path::AbstractString, table)
+    con = DBInterface.connect(DuckDB.DB, ":memory:")
+    return try
+        DuckDB.register_data_frame(con, table, "forecast_samples")
+        DBInterface.execute(
+            con,
+            "COPY forecast_samples TO $(_quote_duckdb_string(path)) (FORMAT parquet)"
+        )
+    finally
+        DBInterface.close(con)
+    end
+end
+
+
 """
     create_forecast_output(input, results, output_dir, output_type; kwargs...) -> DataFrame
 
@@ -219,7 +237,7 @@ function create_forecast_output(
         Dict(
             "observed" => "observed_ed_visits",
             "other" => "other_ed_visits",
-            "pct" => "prop_disease_ed_visits",
+            "pct" => "prop_disease_ed_visits"
         )[input.ed_visit_type]
     end
 
@@ -236,21 +254,11 @@ function create_forecast_output(
     forecast_df[!, :geo_value] .= input.location
     forecast_df[!, :disease] .= input.pathogen
 
-    # Add metadata columns for hubverse compatibility
-    forecast_df[!, :geo_value] .= input.location
-    forecast_df[!, :disease] .= input.pathogen
-
-    # Convert date column to string for parquet compatibility
-    # Julia parquet does not support writing Dates
-    forecast_df[!, :date] = string.(forecast_df[!, :date])
-
     # Save as parquet if requested
     if save_output
-        # Use model-specific naming with frequency prefix
-        # This matches the convention: {frequency}_{model}_samples_{target_letter}.parquet
-        parquet_path = joinpath(output_dir, "samples_raw.parquet")
+        parquet_path = joinpath(output_dir, "samples.parquet")
         mkpath(dirname(parquet_path))
-        Parquet.write_parquet(parquet_path, forecast_df)
+        _write_parquet_with_duckdb(parquet_path, forecast_df)
 
         @info "Saved pipeline forecast samples to $parquet_path"
     end

EpiAutoGP/test/runtests.jl

@@ -2,6 +2,8 @@ using Test, EpiAutoGP
 using JSON3
 using ArgParse
 using Dates
+using DBInterface
+using DuckDB
 using CSV
 using DataFramesMeta
 using Random

EpiAutoGP/test/test_output.jl

@@ -78,4 +78,67 @@ end
             rm(tmpdir, recursive = true)
         end
     end
+
+    @testset "PipelineOutput writes samples parquet with Date column" begin
+        input = EpiAutoGPInput(
+            [Date("2024-01-01")],
+            [100.0],
+            "COVID-19",
+            "CA",
+            "nhsn",
+            "epiweekly",
+            "observed",
+            Date("2024-01-01"),
+            Date[],
+            Vector{Real}[]
+        )
+
+        forecast_dates = [Date("2024-01-08"), Date("2024-01-15")]
+        forecasts = reshape([10.0, 20.0, 30.0, 40.0], 2, 2)
+        results = (forecast_dates = forecast_dates, forecasts = forecasts)
+
+        tmpdir = mktempdir()
+        try
+            result_df = create_forecast_output(
+                input, results, tmpdir, PipelineOutput();
+                save_output = true
+            )
+
+            parquet_path = joinpath(tmpdir, "samples.parquet")
+            @test isfile(parquet_path)
+            @test eltype(result_df.date) == Date
+            @test propertynames(result_df) == [
+                :date,
+                Symbol(".value"),
+                Symbol(".draw"),
+                Symbol(".variable"),
+                :resolution,
+                :geo_value,
+                :disease
+            ]
+
+            con = DBInterface.connect(DuckDB.DB, ":memory:")
+            try
+                read_df = DataFrame(DBInterface.execute(
+                    con,
+                    "SELECT * FROM read_parquet($(EpiAutoGP._quote_duckdb_string(parquet_path)))"
+                ))
+
+                @test eltype(read_df.date) == Date
+                @test propertynames(read_df) == propertynames(result_df)
+                @test read_df.date == forecast_dates[[1, 2, 1, 2]]
+                @test read_df[!, Symbol(".draw")] == [1, 1, 2, 2]
+                @test read_df[!, Symbol(".value")] == [10.0, 20.0, 30.0, 40.0]
+                @test all(read_df[!, Symbol(".variable")] .==
+                          "observed_hospital_admissions")
+                @test all(read_df.resolution .== "epiweekly")
+                @test all(read_df.geo_value .== "CA")
+                @test all(read_df.disease .== "COVID-19")
+            finally
+                DBInterface.close(con)
+            end
+        finally
+            rm(tmpdir, recursive = true)
+        end
+    end
 end

pipelines/epiautogp/create_samples_from_epiautogp_fit_dir.R

@@ -24,20 +24,9 @@
     load_credentials,
     make_figures_from_model_fit_dir,
     model_fit_dir_to_hub_tbl,
-    run_r_script,
 )
 
 
-def create_samples_from_epiautogp_fit_dir(model_fit_dir: Path) -> None:
-    """Create samples.parquet from an EpiAutoGP model fit directory using R."""
-    run_r_script(
-        "pipelines/epiautogp/create_samples_from_epiautogp_fit_dir.R",
-        [str(model_fit_dir)],
-        function_name="create_samples_from_epiautogp_fit_dir",
-    )
-    return None
-
-
 @dataclass
 class ModelPaths:
     """
@@ -150,8 +139,6 @@ def post_process_forecast(self) -> None:
         self.logger.info("Processing forecast and generating plots...")
         model_fit_dir = Path(self.model_run_dir, self.model_name)
 
-        create_samples_from_epiautogp_fit_dir(model_fit_dir=model_fit_dir)
-
         make_figures_from_model_fit_dir(
             model_fit_dir=model_fit_dir,
             save_figs=True,