feat: Hive partition pruning in distributed reads (#18)

Skip Parquet files whose Hive partition values (year=2024/month=01/)
don't match pipeline filter predicates. Extracts simple equality
predicates from filter ops and matches against partition key-value
pairs parsed from file paths.

Integrated into the Partitioner before size-balanced assignment —
fewer files to distribute means better balance and less IO.

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: cigrainger

lib/dux/remote/partition_pruner.ex

@@ -0,0 +1,124 @@
+defmodule Dux.Remote.PartitionPruner do
+  @moduledoc false
+
+  # Prunes Hive-partitioned Parquet files based on pipeline filter ops.
+  #
+  # Given a list of file paths like:
+  #   ["s3://bucket/year=2024/month=01/data.parquet",
+  #    "s3://bucket/year=2024/month=02/data.parquet",
+  #    "s3://bucket/year=2023/month=12/data.parquet"]
+  #
+  # And filter ops like:
+  #   [{:filter, "\"year\" = 2024"}]
+  #
+  # Returns only the files where year=2024.
+  #
+  # Only handles simple equality predicates on partition columns.
+  # Complex expressions are ignored (safe — we just read more data).
+
+  @doc """
+  Prune files whose Hive partition values don't match filter predicates.
+
+  Returns the filtered list of files. If no partition filters can be
+  extracted, returns the original list unchanged.
+  """
+  def prune(files, ops) do
+    filters = extract_partition_filters(ops)
+
+    if filters == [] do
+      files
+    else
+      Enum.filter(files, fn file ->
+        partitions = extract_hive_partitions(file)
+        matches_all?(partitions, filters)
+      end)
+    end
+  end
+
+  # ---------------------------------------------------------------------------
+  # Hive partition extraction from file paths
+  # ---------------------------------------------------------------------------
+
+  # Extract key=value pairs from a file path.
+  # "s3://bucket/year=2024/month=01/data.parquet" → %{"year" => "2024", "month" => "01"}
+  @doc false
+  def extract_hive_partitions(path) do
+    path
+    |> String.split("/")
+    |> Enum.flat_map(fn segment ->
+      case String.split(segment, "=", parts: 2) do
+        [key, value] when key != "" and value != "" -> [{key, value}]
+        _ -> []
+      end
+    end)
+    |> Map.new()
+  end
+
+  # ---------------------------------------------------------------------------
+  # Filter predicate extraction from ops
+  # ---------------------------------------------------------------------------
+
+  # Extract simple equality predicates from filter ops that match partition columns.
+  # Returns [{column_name, value}] for predicates like:
+  #   "year" = 2024        → {"year", "2024"}
+  #   "month" = '01'       → {"month", "01"}
+  #   year = 2024          → {"year", "2024"}
+  defp extract_partition_filters(ops) do
+    ops
+    |> Enum.flat_map(fn
+      {:filter, expr} when is_binary(expr) -> parse_equality_predicates(expr)
+      _ -> []
+    end)
+  end
+
+  # Parse simple equality predicates from a SQL WHERE clause string.
+  # Handles: col = val, "col" = val, col = 'val', and AND-connected predicates.
+  defp parse_equality_predicates(expr) do
+    # Split on AND (case-insensitive) to handle compound filters
+    expr
+    |> String.split(~r/\bAND\b/i)
+    |> Enum.flat_map(&parse_single_predicate/1)
+  end
+
+  defp parse_single_predicate(predicate) do
+    trimmed = String.trim(predicate)
+
+    # Match patterns like: "col" = val, col = val, "col" = 'val'
+    # Captures: optional quotes around column name, = operator, value
+    case Regex.run(
+           ~r/^"?([a-zA-Z_][a-zA-Z0-9_]*)"?\s*=\s*(.+)$/,
+           trimmed
+         ) do
+      [_, column, raw_value] ->
+        value = clean_value(String.trim(raw_value))
+        [{column, value}]
+
+      _ ->
+        []
+    end
+  end
+
+  # Strip quotes and whitespace from a SQL value literal
+  defp clean_value(val) do
+    val
+    |> String.trim("'")
+    |> String.trim("\"")
+    |> String.trim()
+  end
+
+  # ---------------------------------------------------------------------------
+  # Matching
+  # ---------------------------------------------------------------------------
+
+  # Check if a file's partition values match ALL filter predicates.
+  # If the file has no partition for a given filter column, it passes
+  # (the filter may apply to a non-partition column).
+  defp matches_all?(partitions, filters) do
+    Enum.all?(filters, fn {col, expected} ->
+      case Map.get(partitions, col) do
+        nil -> true
+        actual -> actual == expected
+      end
+    end)
+  end
+end

lib/dux/remote/partitioner.ex

@@ -1,6 +1,8 @@
 defmodule Dux.Remote.Partitioner do
   @moduledoc false
 
+  alias Dux.Remote.PartitionPruner
+
   # Assigns data partitions to workers.
   #
   # For Parquet glob sources and DuckLake file manifests, splits files
@@ -57,6 +59,18 @@ defmodule Dux.Remote.Partitioner do
   # ---------------------------------------------------------------------------
 
   defp distribute_parquet_files(files, workers, pipeline, source_opts, assign_opts) do
+    # Prune files whose Hive partition values don't match pipeline filters
+    files = PartitionPruner.prune(files, pipeline.ops)
+
+    if files == [] do
+      # All files pruned — return empty result from one worker
+      [{hd(workers), %{pipeline | source: {:sql, "SELECT 1 WHERE false"}}}]
+    else
+      distribute_pruned_files(files, workers, pipeline, source_opts, assign_opts)
+    end
+  end
+
+  defp distribute_pruned_files(files, workers, pipeline, source_opts, assign_opts) do
     files_with_sizes = fetch_file_sizes(files, assign_opts)
 
     assignments =

test/dux/coordinator_test.exs

@@ -307,6 +307,116 @@ defmodule Dux.CoordinatorTest do
     end
   end
 
+  describe "Partitioner partition pruning" do
+    test "prunes Hive-partitioned files based on filter ops" do
+      dir = tmp_path("hive_pruning")
+      File.mkdir_p!(dir)
+
+      try do
+        # Create Hive-partitioned structure
+        for year <- [2023, 2024], month <- ["01", "02"] do
+          part_dir = Path.join([dir, "year=#{year}", "month=#{month}"])
+          File.mkdir_p!(part_dir)
+
+          Dux.from_list([%{"x" => year * 100 + String.to_integer(month)}])
+          |> Dux.to_parquet(Path.join(part_dir, "data.parquet"))
+        end
+
+        pipeline =
+          Dux.from_parquet(Path.join(dir, "**/*.parquet"))
+          |> Dux.filter_with("year = 2024")
+
+        workers = [:w1, :w2]
+        assignments = Partitioner.assign(pipeline, workers)
+
+        # Should only have year=2024 files (2 files), not year=2023
+        all_files =
+          Enum.flat_map(assignments, fn {_w, p} ->
+            case p.source do
+              {:parquet, f, _} -> [f]
+              {:parquet_list, f, _} -> f
+            end
+          end)
+
+        assert length(all_files) == 2
+        assert Enum.all?(all_files, &String.contains?(&1, "year=2024"))
+      after
+        File.rm_rf!(dir)
+      end
+    end
+
+    test "prunes on multiple partition columns" do
+      dir = tmp_path("hive_multi_prune")
+      File.mkdir_p!(dir)
+
+      try do
+        for year <- [2023, 2024], month <- ["01", "02"] do
+          part_dir = Path.join([dir, "year=#{year}", "month=#{month}"])
+          File.mkdir_p!(part_dir)
+
+          Dux.from_list([%{"x" => 1}])
+          |> Dux.to_parquet(Path.join(part_dir, "data.parquet"))
+        end
+
+        pipeline =
+          Dux.from_parquet(Path.join(dir, "**/*.parquet"))
+          |> Dux.filter_with("year = 2024 AND month = '01'")
+
+        workers = [:w1, :w2]
+        assignments = Partitioner.assign(pipeline, workers)
+
+        all_files =
+          Enum.flat_map(assignments, fn {_w, p} ->
+            case p.source do
+              {:parquet, f, _} -> [f]
+              {:parquet_list, f, _} -> f
+            end
+          end)
+
+        assert length(all_files) == 1
+        assert hd(all_files) |> String.contains?("year=2024")
+        assert hd(all_files) |> String.contains?("month=01")
+      after
+        File.rm_rf!(dir)
+      end
+    end
+
+    test "no pruning when filter doesn't match partition columns" do
+      dir = tmp_path("hive_no_prune")
+      File.mkdir_p!(dir)
+
+      try do
+        for year <- [2023, 2024] do
+          part_dir = Path.join(dir, "year=#{year}")
+          File.mkdir_p!(part_dir)
+
+          Dux.from_list([%{"x" => year}])
+          |> Dux.to_parquet(Path.join(part_dir, "data.parquet"))
+        end
+
+        pipeline =
+          Dux.from_parquet(Path.join(dir, "**/*.parquet"))
+          |> Dux.filter_with("x > 100")
+
+        workers = [:w1, :w2]
+        assignments = Partitioner.assign(pipeline, workers)
+
+        all_files =
+          Enum.flat_map(assignments, fn {_w, p} ->
+            case p.source do
+              {:parquet, f, _} -> [f]
+              {:parquet_list, f, _} -> f
+            end
+          end)
+
+        # x is not a partition column, so all files pass through
+        assert length(all_files) == 2
+      after
+        File.rm_rf!(dir)
+      end
+    end
+  end
+
   # ---------------------------------------------------------------------------
   # Merger unit tests
   # ---------------------------------------------------------------------------