Read parquet files from AWS S3

.github/workflows/ci.yml

@@ -55,6 +55,11 @@ jobs:
       - name: Compile once again but without optional deps
         run: mix compile --force --warnings-as-errors --no-optional-deps
 
+      - name: Run cloud integration tests
+        run: |
+          mix localstack.setup
+          mix test --only cloud_integration
+
   format:
     runs-on: ubuntu-latest
     name: mix format

lib/explorer/data_frame.ex

@@ -107,11 +107,18 @@ defmodule Explorer.DataFrame do
       # path to a file in disk
       Explorer.DataFrame.from_parquet("/path/to/file.parquet")
 
-      # path to a URL schema (with optional configuration) (coming soon)
-      Explorer.DataFrame.from_parquet("s3://bucket/file.parquet", config: FSS.S3.config(credentials))
+      # path to a URL schema (with optional configuration)
+      Explorer.DataFrame.from_parquet("s3://bucket/file.parquet", config: FSS.S3.Entry.config_from_system_env())
 
-      # path to a filesystem specification entry (coming soon)
-      Explorer.DataFrame.from_parquet(FSS.S3.entry("s3://bucket/file.parquet", credentials))
+      # it's possible to configure using keyword lists
+      Explorer.DataFrame.from_parquet("s3://bucket/file.parquet", config: [access_key_id: "my-key", secret_access_key: "my-secret"])
+
+      # a FSS entry with its config inside
+      Explorer.DataFrame.from_parquet(%FSS.S3.Entry{config: %FSS.S3.Config{}})
+
+  The `:config` option of `from_*` functions is only required if the filename is a path
+  to a remote resource. In case it's a FSS entry, the requirement is that the config is passed
+  inside the entry struct.
 
   ## Selecting columns and access
 
@@ -719,8 +726,7 @@ defmodule Explorer.DataFrame do
   defp normalise_entry(%S3.Entry{config: %S3.Config{}} = entry, nil), do: {:ok, entry}
 
   defp normalise_entry("s3://" <> _rest = entry, config) do
-    config = s3_config(config)
-    {:ok, %S3.Entry{url: entry, config: config}}
+    S3.Entry.parse(entry, config: config)
   end
 
   defp normalise_entry("file://" <> path, _config), do: {:ok, %Local.Entry{path: path}}
@@ -729,12 +735,6 @@ defmodule Explorer.DataFrame do
     {:ok, %Local.Entry{path: filepath}}
   end
 
-  defp s3_config(%S3.Config{} = config), do: config
-
-  defp s3_config(other) do
-    raise ArgumentError, "expected a valid FSS.S3.config/1 in :config, got: #{inspect(other)}"
-  end
-
   @doc """
   Similar to `from_parquet/2` but raises if there is a problem reading the Parquet file.
   """

lib/explorer/polars_backend/data_frame.ex

@@ -221,8 +221,12 @@ defmodule Explorer.PolarsBackend.DataFrame do
   end
 
   @impl true
-  def from_parquet(%S3.Entry{}, _max_rows, _columns) do
-    raise "S3 is not supported yet"
+  def from_parquet(%S3.Entry{} = entry, max_rows, columns) do
+    # We first read using a lazy dataframe, then we collect.
+    with {:ok, ldf} <- Native.lf_from_parquet_cloud(entry, max_rows, columns),
+         {:ok, df} <- Native.lf_collect(ldf) do
+      {:ok, Shared.create_dataframe(df)}
+    end
   end
 
   @impl true

lib/explorer/polars_backend/lazy_frame.ex

@@ -146,19 +146,16 @@ defmodule Explorer.PolarsBackend.LazyFrame do
   defp char_byte(<<char::utf8>>), do: char
 
   @impl true
-  def from_parquet(%S3.Entry{}, _max_rows, _columns) do
-    raise "S3 is not supported yet"
+  def from_parquet(%S3.Entry{} = entry, max_rows, columns) do
+    case Native.lf_from_parquet_cloud(entry, max_rows, columns) do
+      {:ok, df} -> {:ok, Shared.create_dataframe(df)}
+      {:error, error} -> {:error, error}
+    end
   end
 
   @impl true
   def from_parquet(%Local.Entry{} = entry, max_rows, columns) do
-    if columns do
-      raise ArgumentError,
-            "`columns` is not supported by Polars' lazy backend. " <>
-              "Consider using `select/2` after reading the parquet file"
-    end
-
-    case Native.lf_from_parquet(entry.path, max_rows) do
+    case Native.lf_from_parquet(entry.path, max_rows, columns) do
       {:ok, df} -> {:ok, Shared.create_dataframe(df)}
       {:error, error} -> {:error, error}
     end

lib/explorer/polars_backend/native.ex

@@ -175,7 +175,8 @@ defmodule Explorer.PolarsBackend.Native do
   def lf_slice(_df, _offset, _length), do: err()
   def lf_from_ipc(_filename), do: err()
   def lf_from_ndjson(_filename, _infer_schema_length, _batch_size), do: err()
-  def lf_from_parquet(_filename, _stop_after_n_rows), do: err()
+  def lf_from_parquet(_filename, _stop_after_n_rows, _maybe_columns), do: err()
+  def lf_from_parquet_cloud(_ex_s3_entry, _stop_after_n_rows, _maybe_columns), do: err()
 
   def lf_from_csv(
         _filename,

lib/fss.ex

@@ -13,19 +13,91 @@ defmodule FSS do
 
   defmodule S3 do
     defmodule Config do
-      defstruct [:region, :access_key_id, :secret_access_key]
+      defstruct [
+        :access_key_id,
+        :region,
+        :secret_access_key,
+        :endpoint,
+        :token
+      ]
 
       @type t :: %__MODULE__{
-              region: String.t(),
               access_key_id: String.t(),
-              secret_access_key: String.t()
+              region: String.t(),
+              secret_access_key: String.t(),
+              endpoint: String.t() | nil,
+              token: String.t() | nil
             }
     end
 
     defmodule Entry do
-      defstruct [:url, :config]
+      defstruct [:bucket, :key, :port, :config]
+
+      @type t :: %__MODULE__{
+              bucket: String.t(),
+              key: String.t(),
+              port: pos_integer(),
+              config: Config.t()
+            }
+
+      def config_from_system_env() do
+        %FSS.S3.Config{
+          access_key_id: System.get_env("AWS_ACCESS_KEY_ID"),
+          secret_access_key: System.get_env("AWS_SECRET_ACCESS_KEY"),
+          region: System.get_env("AWS_REGION", System.get_env("AWS_DEFAULT_REGION")),
+          token: System.get_env("AWS_SESSION_TOKEN")
+        }
+      end
+
+      def parse(url, opts \\ []) do
+        opts = Keyword.validate!(opts, config: nil)
+
+        uri = URI.parse(url)
+
+        case uri do
+          %{scheme: "s3", host: bucket, path: "/" <> key} when is_binary(bucket) ->
+            config =
+              opts
+              |> Keyword.fetch!(:config)
+              |> case do
+                nil ->
+                  config_from_system_env()
+
+                %Config{} = config ->
+                  config
+
+                config when is_list(config) or is_map(config) ->
+                  struct!(config_from_system_env(), config)
+
+                other ->
+                  raise ArgumentError,
+                        "expect configuration to be a %FSS.S3.Config{} struct, a keyword list or a map. Instead got #{inspect(other)}"
+              end
+              |> validate_config!()
+
+            {:ok, %__MODULE__{bucket: bucket, key: key, port: uri.port, config: config}}
+
+          _ ->
+            {:error,
+             ArgumentError.exception(
+               "expected s3://<bucket>/<key> URL, got: " <>
+                 URI.to_string(uri)
+             )}
+        end
+      end
+
+      defp validate_config!(%Config{} = config) do
+        access = config.access_key_id
+        secret = config.secret_access_key
+        region = config.region
+
+        if is_nil(access) or is_nil(secret) or is_nil(region) or
+             access == "" or secret == "" or region == "" do
+          raise "missing configuration keys or region to access the S3 API"
+        end
 
-      @type t :: %__MODULE__{url: String.t(), config: Config.t()}
+        config
+      end
     end
   end
 end

mix.exs

@@ -17,10 +17,11 @@ defmodule Explorer.MixProject do
       package: package(),
       deps: deps(),
       docs: docs(),
-      preferred_cli_env: [ci: :test],
+      preferred_cli_env: [ci: :test, "localstack.setup": :test],
       aliases: [
         "rust.lint": ["cmd cargo clippy --manifest-path=native/explorer/Cargo.toml -- -Dwarnings"],
         "rust.fmt": ["cmd cargo fmt --manifest-path=native/explorer/Cargo.toml --all"],
+        "localstack.setup": ["cmd ./test/support/setup-localstack.sh"],
         ci: ["format", "rust.fmt", "rust.lint", "test"]
       ]
     ]