diff --git a/.gitignore b/.gitignore index 2087c35..065154d 100644 --- a/.gitignore +++ b/.gitignore @@ -55,3 +55,4 @@ checksum-Elixir.ExArrow.Native.exs /bench/output/ .DS_Store +/baoulo/ diff --git a/ANNOUNCEMENT_0_2_0.md b/ANNOUNCEMENT_0_2_0.md deleted file mode 100644 index 624615c..0000000 --- a/ANNOUNCEMENT_0_2_0.md +++ /dev/null @@ -1,124 +0,0 @@ -# ExArrow 0.2.0 β€” TLS Flight, multi-dataset routing, ADBC connection pool - -**[r/elixir | Elixir Forum post]** - ---- - -ExArrow 0.2.0 is out. ExArrow is an Elixir library that brings Apache Arrow -IPC, Arrow Flight (gRPC), and ADBC (Arrow Database Connectivity) to the BEAM -with zero-copy semantics: data lives in native Rust/Arrow buffers and never -touches the BEAM heap unless you explicitly materialise it. - -This release delivers the full v0.2 roadmap. - ---- - -## What is new - -### TLS for Arrow Flight - -The built-in Flight server now accepts encrypted connections. Pass `tls:` to -`Server.start_link/2` for one-way or mutual TLS: - -```elixir -cert = File.read!("server.crt") -key = File.read!("server.key") - -{:ok, server} = ExArrow.Flight.Server.start_link(9999, - tls: [cert_pem: cert, key_pem: key]) -``` - -Mutual TLS (mTLS) adds `ca_cert_pem:`. The client already auto-selected TLS for -remote hosts via the OS certificate store; nothing changes there. - ---- - -### Multi-dataset routing in the Flight server - -The server now stores multiple named datasets instead of a single echo slot. -Upload with a named descriptor, retrieve by ticket: - -```elixir -:ok = ExArrow.Flight.Client.do_put(client, schema, batches, - descriptor: {:cmd, "sales_2024"}) - -{:ok, stream} = ExArrow.Flight.Client.do_get(client, "sales_2024") -``` - -All existing code that relies on the default `"echo"` ticket continues to work -unchanged. - ---- - -### ADBC connection pool - -`ExArrow.ADBC.ConnectionPool` is a NimblePool-backed pool that recycles open -ADBC connections across callers. Drop it into your supervision tree: - -```elixir -children = [ - {ExArrow.ADBC.DatabaseServer, - name: :mydb, - driver_path: "/usr/local/lib/libadbc_driver_duckdb.so"}, - {ExArrow.ADBC.ConnectionPool, - name: :mypool, database: :mydb, pool_size: 4} -] -Supervisor.start_link(children, strategy: :one_for_one) - -{:ok, stream} = ExArrow.ADBC.ConnectionPool.query(:mypool, - "SELECT * FROM events WHERE day = today()") -``` - ---- - -### Integration test matrix - -A new CI workflow (`integration.yml`) runs live ADBC tests against PostgreSQL -14/15/16 and DuckDB 1.1.3/1.2.0 on every push to main and on pull requests. -The tests are in `test/ex_arrow/adbc_integration_test.exs` and are excluded -from the default `mix test` run. - ---- - -## Links - -- Hex: https://hex.pm/packages/ex_arrow -- Docs: https://hexdocs.pm/ex_arrow -- GitHub: https://github.com/thanos/ex_arrow -- Changelog: https://github.com/thanos/ex_arrow/blob/main/CHANGELOG.md - ---- - -## What is ExArrow for? - -ExArrow is transport and protocol glue, not a dataframe library: - -- Read and write Arrow IPC (stream and file formats) -- Connect to Arrow Flight servers (Dremio, InfluxDB IOx, DuckDB, custom) -- Execute SQL via ADBC and receive Arrow result streams with minimal BEAM copying -- Build data pipelines where the Elixir node is an ingestion endpoint, forwarder, - or thin query client - -For in-memory analysis use Explorer. For normal Ecto queries use Ecto. ExArrow -sits at the layer below, where the Arrow wire format matters. - ---- - -## Upgrade - -No breaking changes for most users. If you have a Mox stub for the -`do_put` Flight callback, update it from 3 to 4 arguments: - -```elixir -# Before -Mox.expect(MyMock, :do_put, fn client, schema, batches -> :ok end) -# After -Mox.expect(MyMock, :do_put, fn client, schema, batches, _opts -> :ok end) -``` - -Full upgrade guide: [RELEASE_NOTES_0_2_0.md](RELEASE_NOTES_0_2_0.md) - ---- - -Feedback, issues, and pull requests welcome on GitHub. Next up (v0.3): Arrow -compute kernels, Parquet support, and Explorer/Nx bridge modules. diff --git a/ANNOUNCEMENT_0_3_0.md b/ANNOUNCEMENT_0_3_0.md deleted file mode 100644 index c1f3a92..0000000 --- a/ANNOUNCEMENT_0_3_0.md +++ /dev/null @@ -1,127 +0,0 @@ -# ExArrow 0.3.0 β€” Parquet, Compute kernels, Explorer bridge, Nx bridge - -Hi everyone πŸ‘‹ - -ExArrow 0.3.0 is out. This release ships the four features that were on the -v0.3 roadmap: Arrow compute kernels that stay entirely in native memory, -Parquet read/write, a one-call Explorer bridge, and an Nx bridge that shares -raw byte buffers with tensors. - ---- - -## What is ExArrow? - -ExArrow gives Elixir and Erlang applications first-class Apache Arrow support: -IPC streaming, Arrow Flight (gRPC), and ADBC database connectivity. Column -data lives in Rust buffers; the BEAM holds lightweight opaque handles. -Precompiled NIFs for Linux x86-64/aarch64, macOS arm64/x86-64, and Windows β€” -no Rust required. - -```elixir -{:ex_arrow, "~> 0.3.0"} -``` - ---- - -## What is new in 0.3.0 - -### Parquet read/write - -Read and write Parquet files with the same `ExArrow.Stream` interface you -already use for IPC and ADBC: - -```elixir -# Read -{:ok, stream} = ExArrow.Parquet.Reader.from_file("/data/events.parquet") -batches = ExArrow.Stream.to_list(stream) - -# Write -{:ok, binary} = ExArrow.Parquet.Writer.to_binary(schema, batches) -:ok = ExArrow.Parquet.Writer.to_file("/out/result.parquet", schema, batches) -``` - -In-memory binaries are also accepted by the reader, so files downloaded from -S3 or received over HTTP can be parsed without touching the filesystem. - ---- - -### Arrow compute kernels - -`ExArrow.Compute` wraps the `arrow-select` and `arrow-ord` Rust crates. -Everything runs in native Arrow memory β€” no data is ever copied into BEAM -terms, and the results are new `ExArrow.RecordBatch` handles you can chain -directly into further operations, IPC writes, or Flight uploads: - -```elixir -{:ok, slim} = ExArrow.Compute.project(batch, ["id", "score", "region"]) -{:ok, sorted} = ExArrow.Compute.sort(slim, "score", ascending: false) -{:ok, active} = ExArrow.Compute.filter(sorted, predicate_batch) -``` - ---- - -### Explorer bridge - -`ExArrow.Explorer` converts between `ExArrow.Stream` / `ExArrow.RecordBatch` -and `Explorer.DataFrame` in a single call. The bridge uses Arrow IPC -internally β€” no CSV, no row-by-row conversion. - -Requires `{:explorer, "~> 0.8"}` in your `mix.exs` (optional). - -```elixir -# Stream from a Flight server β†’ DataFrame -{:ok, stream} = ExArrow.Flight.Client.do_get(client, "sales_2024") -{:ok, df} = ExArrow.Explorer.from_stream(stream) -Explorer.DataFrame.filter(df, score > 0.9) - -# DataFrame β†’ Parquet file -{:ok, stream} = ExArrow.Explorer.to_stream(df) -:ok = ExArrow.Parquet.Writer.to_file("/out/result.parquet", - ExArrow.Stream.schema(stream) |> elem(1), - ExArrow.Stream.to_list(stream)) -``` - ---- - -### Nx bridge - -`ExArrow.Nx` converts Arrow numeric columns to `Nx.Tensor` values and back -by copying the raw byte buffer once. No intermediate list, no extra heap -allocation. - -Requires `{:nx, "~> 0.7"}` in your `mix.exs` (optional). - -```elixir -# Column to tensor -{:ok, tensor} = ExArrow.Nx.column_to_tensor(batch, "price") -mean_price = tensor |> Nx.mean() |> Nx.to_number() - -# All numeric columns at once -{:ok, tensors} = ExArrow.Nx.to_tensors(batch) - -# Tensor back to a record batch -weights = Nx.tensor([0.1, 0.2, 0.7], type: {:f, 64}) -{:ok, batch} = ExArrow.Nx.from_tensor(weights, "weights") -``` - ---- - -## No breaking changes - -Update the version pin, run `mix deps.get`, and you are done: - -```elixir -{:ex_arrow, "~> 0.3.0"} -``` - ---- - -## Links - -- **Hex:** https://hex.pm/packages/ex_arrow -- **Docs:** https://hexdocs.pm/ex_arrow -- **Changelog:** https://github.com/thanos/ex_arrow/blob/main/CHANGELOG.md -- **Release notes:** https://github.com/thanos/ex_arrow/blob/main/RELEASE_NOTES_0_3_0.md -- **Source:** https://github.com/thanos/ex_arrow - -Feedback, issues, and PRs very welcome. Thanks! diff --git a/CHANGELOG.md b/CHANGELOG.md index a05dc17..95c06d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,120 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.7.0] - 2026-06-29 + +### Added β€” Arrow-native streaming and pipeline infrastructure + +ExArrow v0.7.0 transforms the library from a transport and interoperability +layer into the foundation for Arrow-native data pipelines on the BEAM. The +central architectural principle is: **operate on Arrow RecordBatch values** β€” +not `list(map())`, not `Explorer.DataFrame`, not `Nx.Tensor`. Explorer and Nx +remain downstream consumers; ExArrow is the Arrow layer. + +#### Milestone 1 β€” First-class streaming (`ExArrow.Stream`) + +- `ExArrow.Stream.from_parquet/1`, `from_parquet_binary/1` +- `ExArrow.Stream.from_ipc/1`, `from_ipc_file/1` +- `ExArrow.Stream.from_flight/2` +- `ExArrow.Stream.from_flight_sql/2` +- `ExArrow.Stream.from_adbc/1` and `from_adbc/2` +- `ExArrow.Stream.source/1` returns the origin metadata attached by the + constructor (forwarded to telemetry as `:source`). +- All constructors return `{:ok, stream} | {:error, reason}` and compose with + `|>/2` into `ExArrow.Pipeline` and `ExArrow.Flow`. + +#### Milestone 2 β€” Batch operations (`ExArrow.Batch`) + +- `ExArrow.Batch.schema/1`, `select/2`, `drop/2`, `rename/2`, `take/2`, + `filter/2`. +- `select/2`, `drop/2`, `take/2`, and `filter/2` work for all Arrow types + ExArrow supports (delegate to the native compute kernels). +- `rename/2` rebuilds a batch from raw column buffers and supports the + fixed-width numeric and boolean types. +- `take/2` accepts an integer (first N rows) or a list of zero-based indices. + +#### Milestone 3 β€” Flow integration (`ExArrow.Flow`) + +- `ExArrow.Flow.from_batches/1` wraps an `ExArrow.Stream` in a `Flow` of + `ExArrow.RecordBatch` values. +- `ExArrow.Flow.map_batches/2` and `each_batch/2` emit + `[:ex_arrow, :pipeline, :batch]` telemetry per batch. +- All standard `Flow` combinators (`map`, `flat_map`, `partition`, `reduce`) + work on the result. + +#### Milestone 4 β€” GenStage integration (`ExArrow.GenStage`) + +- `ExArrow.GenStage.ParquetProducer`, `FlightProducer`, `ADBCProducer`. +- Demand-driven streaming with proper backpressure and clean shutdown. +- Shared `ExArrow.GenStage.dispatch/2` and `terminate/1` helpers. +- Works with producer, consumer, and producer-consumer wiring patterns. + +#### Milestone 5 β€” Broadway integration (`ExArrow.Broadway`) + +- `ExArrow.Broadway.BatchBuilder.from_messages/1` assembles + `ExArrow.RecordBatch` values from Broadway messages. +- `ExArrow.Broadway.ParquetSink.write/3` and `FlightSink.write/4` sink + assembled batches to Parquet files and Flight servers. +- Configurable batch sizing and flush intervals via the Broadway batcher + configuration. + +#### Milestone 6 β€” Telemetry (`ExArrow.Telemetry`) + +- Events: `[:ex_arrow, :flight, :query]`, `[:ex_arrow, :flight_sql, :query]`, + `[:ex_arrow, :parquet, :read]`, `[:ex_arrow, :parquet, :write]`, + `[:ex_arrow, :stream, :batch]`, `[:ex_arrow, :pipeline, :batch]`. +- Measurements: `rows`, `columns`, `bytes`, `duration`, `batch_count`. +- Metadata: `source`, `destination`, `schema`, `driver`. +- `ExArrow.Telemetry.execute/3` and `span/3` degrade to no-ops when the + `:telemetry` dependency is absent. + +#### Milestone 7 β€” Pipeline sinks + +- `ExArrow.Sink.Parquet.write/2` +- `ExArrow.Sink.Flight.write/3` +- `ExArrow.Sink.DataFrame.write/1` +- `ExArrow.Sink.Nx.write/1` + +#### Milestone 8 β€” Pipeline DSL (`ExArrow.Pipeline`) + +- `ExArrow.Pipeline.map_batches/2`, `each_batch/2`, `write_parquet/2`, + `write_flight/3`, `write_dataframe/1`. +- Thin, lazy abstraction over `ExArrow.Stream` with `{:ok, pipeline} | + {:error, reason}` threading; errors short-circuit through every stage. + +#### Milestone 9 β€” Benchmarks + +- `bench/v070_stream_flow_pipeline_bench.exs` measures Parquet/IPC stream + drains, Flow execution, and Pipeline `map_batches` + `write_parquet` at + 1K / 100K / 1M rows. +- `bench/v070_record_batch_vs_maps_bench.exs` compares Arrow `RecordBatch` + against `list(map())` for build, transform (select), and drain (sum) at + the same row counts. + +#### Milestone 10 β€” Educational material + +- `guides/06_arrow_streams.md` +- `guides/07_arrow_and_flow.md` +- `guides/08_arrow_and_genstage.md` +- `guides/09_arrow_and_broadway.md` +- `guides/10_arrow_pipeline_patterns.md` + +### Changed + +- `mix.exs`: version bumped to `0.7.0`; added optional dependencies on + `:telemetry`, `:flow`, `:gen_stage`, and `:broadway`; `docs/0` extras and + `groups_for_modules` updated to include the new modules. +- `ExArrow.Stream` now carries a `source` field and emits telemetry from + `next/1`; the struct is backward-compatible (the new field defaults to + `nil`). + +### Testing + +- 825 tests, 16 properties, 0 failures (45 excluded integration tests). +- 89.9% coverage (target: >= 85%). +- `mix format`, `mix credo --strict`, `mix dialyzer`, `mix sobelow`, + `mix docs --warnings-as-errors` all pass. + ## [0.6.3] - 2026-06-16 ### Fixed diff --git a/MILESTONE_00.md b/MILESTONE_00.md deleted file mode 100644 index bec229c..0000000 --- a/MILESTONE_00.md +++ /dev/null @@ -1,53 +0,0 @@ -Milestone 0 complete -1) Scope and acceptance criteria -Scope: Project skeleton, public API outline, resource-handle strategy, minimal NIF, CI scaffold. -DoD: -mix test passes -NIF compiles and loads (ExArrow.native_version/0 returns the crate version) -Docs build with mix docs -2) Files created/updated -Path Purpose -mix.exs Rustler deps, ex_doc, stream_data; docs config -lib/ex_arrow.ex Public entry; API outline; native_version/0 -lib/ex_arrow/error.ex Structured error (code, message, details) -lib/ex_arrow/schema.ex, field.ex, array.ex, record_batch.ex, table.ex, stream.ex Core handle stubs -lib/ex_arrow/ipc/reader.ex, ipc/writer.ex IPC stubs -lib/ex_arrow/flight/client.ex, flight/server.ex Flight stubs -lib/ex_arrow/adbc/database.ex, connection.ex, statement.ex ADBC stubs -lib/ex_arrow/native.ex Rustler NIF loader (ex_arrow_native) -native/ex_arrow_native/Cargo.toml, src/lib.rs Minimal Rust NIF (nif_version) -native/ex_arrow_native/.cargo/config.toml macOS link flags for NIF -test/ex_arrow_test.exs Tests for NIF, Error, core/IPC/Flight/ADBC stubs -examples/ipc_roundtrip.exs, flight_echo/server.exs, flight_echo/client.exs, adbc_query.exs Example stubs -docs/overview.md, docs/memory_model.md ExDoc extras -.formatter.exs Include examples/**/*.exs -.gitignore /priv/native/ -.github/workflows/ci.yml Elixir + Rust build and test -README.md Short project and dev instructions - -3) Tests (what they assert) -Native NIF: ExArrow.native_version/0 returns a non-empty string. -ExArrow.Error: Exception with code/message; message-only creates :unknown code. -Core stubs: Schema.fields β†’ []; RecordBatch.schema/num_rows β†’ nil/0; Stream.next β†’ nil. -IPC/Flight/ADBC: Reader from_binary, Writer to_binary, Client connect, Database open return {:error, :not_implemented}. - -4) Implementation -All of the above files were added or updated as described. NIF is a single Rust function nif_version(); Elixir loads it via ExArrow.Native with Rustler (path native/ex_arrow_native, debug in dev/test). - -5) Refactors/cleanup -Removed unused alias ExArrow.RecordBatch from IPC.Reader. -No further refactors in this milestone. - -6) Docs -Overview: What ExArrow is (IPC/Flight/ADBC, native handles), what it is not, current status. -Memory model: Handles and lifetimes, copying rules, NIF scheduling note. -ExDoc configured with main: "overview" and these extras. - -7) How to run the examples -``` -mix run examples/ipc_roundtrip.exs -mix run examples/flight_echo/server.exs -mix run examples/flight_echo/client.exs -mix run examples/adbc_query.exs -``` -Each prints the native version and the current stub result (e.g. :not_implemented). \ No newline at end of file diff --git a/MILESTONE_01.md b/MILESTONE_01.md deleted file mode 100644 index a062b0f..0000000 --- a/MILESTONE_01.md +++ /dev/null @@ -1,48 +0,0 @@ - -## Milestone 1 complete: IPC MVP (streaming only) - -1) Scope and acceptance criteria -Scope: Schema and RecordBatch handles; IPC stream read from binary and file; IPC stream write to binary and file; streaming iterator in Elixir; types: bool, int64, float64, utf8, binary, nulls. -DoD: -Roundtrip tests: encode β†’ decode matches on schema, row count, and fixture -examples/ipc_roundtrip.exs runs successfully -IPC guide exists and is short -2) Files created/updated -Path Change -native/ex_arrow_native/Cargo.toml Added arrow, arrow-ipc, arrow-schema, arrow-array (56.x) -native/ex_arrow_native/src/lib.rs on_load + resource registration; NIF list extended -native/ex_arrow_native/src/resources.rs New: ExArrowSchema, ExArrowRecordBatch, ExArrowIpcStream -native/ex_arrow_native/src/ipc.rs New: IPC read/write NIFs, schema/batch accessors, test fixture -lib/ex_arrow/native.ex NIF stubs for all new functions -lib/ex_arrow/schema.ex fields/1 implemented via NIF -lib/ex_arrow/record_batch.ex schema/1, num_rows/1 implemented via NIF -lib/ex_arrow/stream.ex schema/1, next/1 implemented via NIF -lib/ex_arrow/ipc/reader.ex from_binary/1, from_file/1 implemented -lib/ex_arrow/ipc/writer.ex to_binary/2, to_file/3 implemented -test/ex_arrow_test.exs IPC roundtrip tests, from_binary error test, stream iterator test -examples/ipc_roundtrip.exs Full roundtrip: fixture β†’ read β†’ write β†’ read -docs/ipc_guide.md New: IPC guide (stream vs file, read, write, types, example) -mix.exs docs/ipc_guide.md added to extras -3) Tests (what they assert) -Roundtrip: ipc_test_fixture_binary() returns binary; from_binary gives stream; stream schema has 2 fields (id int64, name utf8); first batch has 2 rows; roundtrip write then read gives same schema and row count. -from_binary empty: Returns {:error, _}. -Stream iterator: Collecting batches from the fixture yields one batch with 2 rows. -Schema.fields / RecordBatch from real IPC: Schema has 2 fields; batch has 2 rows and schema with 2 fields. -4) Implementation notes -Rust: Resource structs hold Arc, RecordBatch, and Mutex>>>. NIFs return {:ok, term} or {:error, msg} as encoded terms. Binary results use OwnedBinary so Elixir gets real binaries. Fixture NIF builds a small IPC stream (id int64, name utf8, 2 rows) for tests. -Elixir: Handles are %Schema{resource: ref} etc.; refs are passed to NIFs. Stream.next/1 returns %RecordBatch{}, nil, or {:error, msg}. -5) Refactors/cleanup -None beyond the above. -6) Docs -IPC guide (docs/ipc_guide.md): stream vs file, reading (from_binary, from_file, schema, next), writing (to_binary, to_file), supported types, link to examples/ipc_roundtrip.exs. -7) How to run the example -``` -mix run examples/ipc_roundtrip.exs -``` -Example output: fixture size, schema fields ["id", "name"], 1 batch and 2 rows, roundtrip write size, and β€œRoundtrip read OK”. - - -Summary of changes: -Rust: New resources and ipc modules; IPC stream read/write; resource types for schema, batch, stream; test fixture binary; all results returned as proper binaries where needed. -Elixir: Schema.fields, RecordBatch.schema/num_rows, Stream.schema/next implemented; IPC.Reader.from_binary/from_file, IPC.Writer.to_binary/to_file implemented; example and IPC guide added. -Tests: 4 new IPC tests (roundtrip, empty binary error, iterator, schema/batch accessors); Flight/ADBC stubs unchanged. diff --git a/README.md b/README.md index cfd9e10..7b5050b 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,9 @@ [![Coverage Status](https://coveralls.io/repos/github/thanos/ex_arrow/badge.svg?branch=main)](https://coveralls.io/github/thanos/ex_arrow?branch=main) -Native Apache Arrow for the BEAM: IPC streaming, Arrow Flight, Arrow Flight SQL, and ADBC database bindings. Column data lives in Rust buffers; Elixir holds lightweight opaque handles. Precompiled NIFs for Linux, macOS, and Windows β€” no Rust required to use. +Native Apache Arrow for the BEAM: IPC streaming, Arrow Flight, Arrow Flight SQL, ADBC database bindings, and Arrow-native pipelines. Column data lives in Rust buffers; Elixir holds lightweight opaque handles. Precompiled NIFs for Linux, macOS, and Windows β€” no Rust required to use. + +> **v0.7.0 β€” Arrow-native pipelines.** New: `ExArrow.Stream` constructors for every source, `ExArrow.Batch` transforms, `ExArrow.Pipeline` DSL, `ExArrow.Flow` / `ExArrow.GenStage` / `ExArrow.Broadway` integrations, `ExArrow.Sink.*` destinations, and `ExArrow.Telemetry` events. The unit of execution is the Arrow `RecordBatch`. See [What's changed in v0.7.0](#whats-changed-in-v070). --- @@ -21,6 +23,7 @@ Native Apache Arrow for the BEAM: IPC streaming, Arrow Flight, Arrow Flight SQL, - [Requirements](#requirements) - [Installation](#installation) - [Quick start](#quick-start) +- [What's changed in v0.7.0](#whats-changed-in-v070) - [Livebook tutorials](#livebook-tutorials) - [IPC: stream and file](#ipc-stream-and-file) - [Arrow Flight: client and server](#arrow-flight-client-and-server) @@ -28,6 +31,8 @@ Native Apache Arrow for the BEAM: IPC streaming, Arrow Flight, Arrow Flight SQL, - [ADBC: database to Arrow streams](#adbc-database-to-arrow-streams) - [Parquet: read and write](#parquet-read-and-write) - [Arrow compute kernels](#arrow-compute-kernels) +- [Batch operations](#batch-operations) +- [Streaming pipelines](#streaming-pipelines) - [Using ExArrow with Explorer](#using-exarrow-with-explorer) - [Using ExArrow with Nx](#using-exarrow-with-nx) - [Use case examples](#use-case-examples) @@ -39,6 +44,7 @@ Native Apache Arrow for the BEAM: IPC streaming, Arrow Flight, Arrow Flight SQL, - [Shipped (v0.3.0)](#shipped-v030) - [Shipped (v0.4.0)](#shipped-v040) - [Shipped (v0.5.0)](#shipped-v050) + - [Shipped (v0.7.0)](#shipped-v070) - [FAQ](#faq) - [License](#license) @@ -90,8 +96,13 @@ BEAM terms unless you ask for them. there until consumed. The BEAM scheduler is never stalled on large copies. Dirty NIF schedulers are used for blocking I/O. - **A uniform stream abstraction** β€” `ExArrow.Stream` works identically for -IPC, Flight, and ADBC results. Code that processes batches does not know or -care where the data came from. + IPC, Flight, and ADBC results. Code that processes batches does not know or + care where the data came from. +- **Arrow-native pipelines (v0.7.0)** β€” `ExArrow.Pipeline` provides a lazy, + composable DSL for transforming and sinking streams of `RecordBatch` values. + `ExArrow.Flow`, `ExArrow.GenStage`, and `ExArrow.Broadway` integrations + cover parallel, demand-driven, and ingestion workloads. Telemetry events + fire at every stage. --- @@ -123,15 +134,19 @@ flowchart TB App --> Explorer("Explorer\ndataframes & analysis") App --> Nx("Nx\ntensors & ML") - App --> ExArrow("ExArrow\nIPC Β· Flight Β· ADBC") + App --> ExArrow("ExArrow\nIPC Β· Flight Β· ADBC Β· Pipelines") App --> ExZarr("ExZarr\nZarr chunked arrays") ExArrow --> IPC("Arrow IPC\nstream & file") ExArrow --> FlightSvr("Arrow Flight\ngRPC server") ExArrow --> FlightSQL("Arrow Flight SQL\nremote query server") ExArrow --> ADBCDrv("ADBC\ndriver") + ExArrow --> Pipeline("Pipelines\nStream Β· Batch Β· Sink") + ExArrow --> Flow("Flow Β· GenStage Β· Broadway") IPC -. "interop via IPC binary" .-> Explorer + Pipeline -. "RecordBatch" .-> Explorer + Pipeline -. "RecordBatch" .-> Nx FlightSvr --> FlightSvcs("Dremio Β· InfluxDB IOx\nCustom Flight services") FlightSQL --> SQLSvcs("DuckDB Β· DataFusion\nDremio Β· InfluxDB v3") @@ -144,7 +159,7 @@ flowchart TB class App app class Explorer,Nx,ExArrow,ExZarr lib - class IPC,FlightSvr,FlightSQL,ADBCDrv proto + class IPC,FlightSvr,FlightSQL,ADBCDrv,Pipeline,Flow proto class FlightSvcs,SQLSvcs,Databases external ``` @@ -189,7 +204,7 @@ Add the dependency: ```elixir def deps do - [{:ex_arrow, "~> 0.6"}] + [{:ex_arrow, "~> 0.7"}] end ``` @@ -217,16 +232,38 @@ For **path dependencies** in Livebook (`Mix.install`), open notebooks from is detected) or use the Hex package: ```elixir -Mix.install([{:ex_arrow, "~> 0.6.3"}, {:rustler, "~> 0.36", optional: true}]) +Mix.install([{:ex_arrow, "~> 0.7.0"}, {:rustler, "~> 0.36", optional: true}]) ``` Alternatively, use the published Hex package so the precompiled NIF is used -and no Rust is needed: `Mix.install([{:ex_arrow, "~> 0.6.3"}])`. +and no Rust is needed: `Mix.install([{:ex_arrow, "~> 0.7.0"}])`. --- ## Quick start +**Read an Arrow stream with the v0.7.0 constructors** (one entry point for +every source): + +```elixir +{:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") +{:ok, schema} = ExArrow.Stream.schema(stream) +Enum.each(stream, fn batch -> + IO.puts("rows: #{ExArrow.RecordBatch.num_rows(batch)}") +end) +``` + +**Build a pipeline** (lazy `map_batches` + sink): + +```elixir +ExArrow.Stream.from_parquet("/data/events.parquet") +|> ExArrow.Pipeline.map_batches(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id", "score"]) + slim +end) +|> ExArrow.Pipeline.write_parquet("/data/slim.parquet") +``` + **Read an Arrow IPC stream:** ```elixir @@ -282,6 +319,73 @@ batch = ExArrow.Stream.next(stream) --- +## What's changed in v0.7.0 + +v0.7.0 turns ExArrow from a transport and interchange library into the +foundation for **Arrow-native data pipelines on the BEAM**. The central +architectural principle: **operate on Arrow `RecordBatch` values** β€” not +`list(map())`, not `Explorer.DataFrame`, not `Nx.Tensor`. Explorer and Nx +remain downstream consumers; ExArrow is the Arrow layer. + +### New modules + +| Module | Purpose | +|------------------------------|------------------------------------------------------------------| +| `ExArrow.Stream` | `from_parquet/1`, `from_ipc/1`, `from_flight/2`, `from_flight_sql/2`, `from_adbc/1,2` constructors with source tracking and telemetry | +| `ExArrow.Batch` | Lightweight RecordBatch transforms: `schema/1`, `select/2`, `drop/2`, `rename/2`, `take/2`, `filter/2` | +| `ExArrow.Pipeline` | Lazy pipeline DSL: `map_batches/2`, `each_batch/2`, `write_parquet/2`, `write_flight/3`, `write_dataframe/1` | +| `ExArrow.Flow` | Arrow-native Flow execution: `from_batches/1`, `map_batches/2`, `each_batch/2` | +| `ExArrow.GenStage` | Demand-driven producers: `ParquetProducer`, `FlightProducer`, `ADBCProducer` | +| `ExArrow.Broadway` | Ingestion pipelines: `BatchBuilder`, `ParquetSink`, `FlightSink` | +| `ExArrow.Sink.Parquet` | Write a stream to a Parquet file | +| `ExArrow.Sink.Flight` | Upload a stream to a Flight server | +| `ExArrow.Sink.DataFrame` | Convert a stream to an Explorer DataFrame | +| `ExArrow.Sink.Nx` | Convert a batch to an Nx tensor | +| `ExArrow.Telemetry` | Telemetry events for every transport and pipeline operation | + +### Optional dependencies added + +```elixir +{:telemetry, "~> 1.0", optional: true}, +{:flow, "~> 1.2", optional: true}, +{:gen_stage, "~> 1.2", optional: true}, +{:broadway, "~> 1.0", optional: true} +``` + +All degrade gracefully when absent β€” `ExArrow.Telemetry.execute/3` is a no-op +without `:telemetry`, the Flow/GenStage/Broadway modules return +`{:error, "..."}` without their respective deps. + +### Telemetry events + +`[:ex_arrow, :flight, :query]`, `[:ex_arrow, :flight_sql, :query]`, +`[:ex_arrow, :parquet, :read]`, `[:ex_arrow, :parquet, :write]`, +`[:ex_arrow, :stream, :batch]`, `[:ex_arrow, :pipeline, :batch]`. See +`ExArrow.Telemetry` for measurements and metadata. + +### New guides + +- [05 Arrow pipelines overview](guides/05_arrow_pipelines_overview.md) +- [06 Arrow streams](guides/06_arrow_streams.md) +- [07 Arrow and Flow](guides/07_arrow_and_flow.md) +- [08 Arrow and GenStage](guides/08_arrow_and_genstage.md) +- [09 Arrow and Broadway](guides/09_arrow_and_broadway.md) +- [10 Arrow pipeline patterns](guides/10_arrow_pipeline_patterns.md) + +### New benchmarks + +- `bench/v070_stream_flow_pipeline_bench.exs` β€” Parquet/IPC stream drains, + Flow execution, Pipeline `map_batches` + `write_parquet` at 1K/100K/1M rows. +- `bench/v070_record_batch_vs_maps_bench.exs` β€” Arrow `RecordBatch` vs + `list(map())` for build, transform, and drain. + +### Stats + +825 tests, 16 properties, 89.9% coverage. `mix format`, `mix credo --strict`, +`mix dialyzer`, `mix sobelow`, `mix docs --warnings-as-errors` all pass. + +--- + ## Livebook tutorials Interactive notebooks (open in [Livebook](https://livebook.dev)): @@ -292,7 +396,7 @@ Interactive notebooks (open in [Livebook](https://livebook.dev)): - **[03 ADBC](livebook/03_adbc.livemd)** β€” Database, Connection, Statement, Stream (`:adbc_package` in Livebook). - **[04 ADBC integration](livebook/04_adbc_integration.livemd)** β€” Connection pooling with NimblePool. -See [livebook/README.md](livebook/README.md) for run instructions. Notebooks use Hex `~> 0.6.3` by default; opening from `livebook/` in a clone builds from source. +See [livebook/README.md](livebook/README.md) for run instructions. Notebooks use Hex `~> 0.7.0` by default; opening from `livebook/` in a clone builds from source. --- @@ -637,6 +741,78 @@ All operations run entirely in native memory. Results are new --- +## Batch operations + +`ExArrow.Batch` (v0.7.0) provides lightweight `RecordBatch` transforms that +stay in native Arrow memory. It is not a dataframe β€” use Explorer for +analytics. All functions return `{:ok, batch} | {:error, message}`. + +```elixir +# Project a subset of columns (works for every Arrow type) +{:ok, slim} = ExArrow.Batch.select(batch, ["id", "score"]) + +# Drop columns +{:ok, rest} = ExArrow.Batch.drop(batch, ["internal_flag"]) + +# Rename columns (numeric/boolean columns; utf8 not supported by buffer extract) +{:ok, renamed} = ExArrow.Batch.rename(batch, %{"id" => "user_id"}) + +# Take rows: first N, or a list of zero-based indices +{:ok, first10} = ExArrow.Batch.take(batch, 10) +{:ok, picked} = ExArrow.Batch.take(batch, [0, 2, 4]) + +# Filter rows with a boolean predicate batch +{:ok, mask} = ExArrow.Compute.project(batch, ["is_active"]) +{:ok, filtered} = ExArrow.Batch.filter(batch, mask) +``` + +--- + +## Streaming pipelines + +`ExArrow.Pipeline` (v0.7.0) is a thin, lazy abstraction for transforming and +sinking Arrow streams. Pipelines compose with `|>/2` and short-circuit on +error. + +```elixir +ExArrow.Stream.from_flight_sql(client, "SELECT * FROM events") +|> ExArrow.Pipeline.map_batches(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id", "ts", "amount"]) + slim +end) +|> ExArrow.Pipeline.each_batch(fn batch -> + IO.puts("processed #{ExArrow.RecordBatch.num_rows(batch)} rows") +end) +|> ExArrow.Pipeline.write_parquet("/data/events_slim.parquet") +``` + +Sinks: `write_parquet/2`, `write_flight/3`, `write_dataframe/1`. + +### Sinks + +`ExArrow.Sink.*` modules provide the same destinations with a stream-or-batch +input shape, for use outside the Pipeline DSL: + +```elixir +:ok = ExArrow.Sink.Parquet.write(stream, "/data/out.parquet") +:ok = ExArrow.Sink.Flight.write(stream, client, descriptor: {:cmd, "events"}) +{:ok, df} = ExArrow.Sink.DataFrame.write(stream) +``` + +### Flow, GenStage, Broadway + +For parallel or demand-driven workloads, drop down to the integration layers: + +- **`ExArrow.Flow`** β€” parallel batch processing via the Flow library. +- **`ExArrow.GenStage.{Parquet,Flight,ADBC}Producer`** β€” demand-driven + producers with backpressure for long-running pipelines. +- **`ExArrow.Broadway`** β€” ingestion pipelines with `BatchBuilder`, + `ParquetSink`, and `FlightSink`. + +See the [guides](#documentation) for full examples and tradeoffs. + +--- + ## Using ExArrow with Explorer [Explorer](https://hex.pm/packages/explorer) handles in-memory analysis. @@ -854,6 +1030,8 @@ HTML reports are written to `bench/output/` (gitignored). | `pipeline_bench.exs` | End-to-end: IPC file on disk to Flight doput without materialising in BEAM | | `explorer_arrow_bench.exs` | Explorer <-> Arrow interchange at 1K/100K/1M rows | | `nx_arrow_bench.exs` | Nx <-> Arrow interchange at 1K/100K/1M rows (rank-1 and rank-2) | +| `v070_stream_flow_pipeline_bench.exs` | Parquet/IPC stream drains, Flow execution, Pipeline map_batches + write_parquet at 1K/100K/1M rows (v0.7.0) | +| `v070_record_batch_vs_maps_bench.exs` | Arrow `RecordBatch` vs `list(map())` for build, transform, and drain at 1K/100K/1M rows (v0.7.0) | ### Published results @@ -872,6 +1050,12 @@ The CI workflow posts a PR alert comment when any scenario regresses more than - [Explorer Integration](guides/02_explorer_integration.md) β€” from_dataframe, to_dataframe, type mapping, limitations - [Nx Integration](guides/03_nx_integration.md) β€” from_nx, to_nx, boolean tensors, rank-2 - [Arrow Ecosystem](guides/04_arrow_ecosystem.md) β€” how ExArrow complements Explorer, Nx, ADBC, Flight, Parquet, ExZarr +- [Arrow Pipelines Overview](guides/05_arrow_pipelines_overview.md) β€” orientation to the v0.7.0 pipeline modules +- [Arrow Streams](guides/06_arrow_streams.md) β€” the v0.7.0 streaming abstraction, constructors, consumption patterns +- [Arrow and Flow](guides/07_arrow_and_flow.md) β€” parallel batch processing with `ExArrow.Flow` +- [Arrow and GenStage](guides/08_arrow_and_genstage.md) β€” demand-driven producers with backpressure +- [Arrow and Broadway](guides/09_arrow_and_broadway.md) β€” ingestion pipelines with `BatchBuilder` and sinks +- [Arrow Pipeline Patterns](guides/10_arrow_pipeline_patterns.md) β€” composable `ExArrow.Pipeline` transforms and sinks - [Memory model](docs/memory_model.md) β€” handles, copying rules, NIF scheduling - [IPC guide](docs/ipc_guide.md) β€” stream vs file, types, limitations - [Parquet guide](docs/parquet_guide.md) β€” read/write Parquet, streaming, comparison with IPC @@ -987,11 +1171,34 @@ welcome for any of them. data (names, binaries, dtype strings, row count). 21 dtypes including timestamps, durations, utf8, and binary columns. +### Shipped (v0.7.0) + +- **First-class streaming** β€” `ExArrow.Stream.from_parquet/1`, + `from_ipc/1`, `from_flight/2`, `from_flight_sql/2`, `from_adbc/1,2` + constructors with source tracking and telemetry. The unit of streaming is + the Arrow `RecordBatch`. +- **Batch operations** β€” `ExArrow.Batch`: `schema/1`, `select/2`, `drop/2`, + `rename/2`, `take/2`, `filter/2`. Lightweight, Arrow-native transforms. +- **Pipeline DSL** β€” `ExArrow.Pipeline`: `map_batches/2`, `each_batch/2`, + `write_parquet/2`, `write_flight/3`, `write_dataframe/1`. Lazy; composes + with `|>/2`; errors short-circuit. +- **Flow integration** β€” `ExArrow.Flow.from_batches/1`, `map_batches/2`, + `each_batch/2` for parallel batch processing. +- **GenStage integration** β€” `ExArrow.GenStage.ParquetProducer`, + `FlightProducer`, `ADBCProducer` with demand-driven backpressure. +- **Broadway integration** β€” `ExArrow.Broadway.BatchBuilder`, + `ParquetSink`, `FlightSink` for ingestion pipelines. +- **Pipeline sinks** β€” `ExArrow.Sink.Parquet`, `Flight`, `DataFrame`, `Nx`. +- **Telemetry** β€” `ExArrow.Telemetry` with events for every transport and + pipeline operation. Optional `:telemetry` dep. +- **Benchmarks** β€” stream/Flow/pipeline and RecordBatch vs `list(map())` at + 1K/100K/1M rows. +- **Educational guides** β€” `guides/06..10` covering streams, Flow, GenStage, + Broadway, and pipeline patterns. + ### Longer-term - **Streaming writes to Delta Lake** β€” sink for data pipeline nodes. -- **OTel / telemetry integration** β€” `:telemetry` events for IPC read/write - throughput, Flight request latency, and ADBC query duration. - **Windows aarch64 precompiled NIF** β€” once GitHub-hosted Windows arm64 runners are generally available. diff --git a/RELEASE_NOTES_0_2_0.md b/RELEASE_NOTES_0_2_0.md deleted file mode 100644 index 8b660d9..0000000 --- a/RELEASE_NOTES_0_2_0.md +++ /dev/null @@ -1,188 +0,0 @@ -# ExArrow 0.2.0 β€” Release Notes - -**Released:** 2026-03-09 - -ExArrow 0.2.0 delivers the four features promised on the v0.2 roadmap: TLS for -Arrow Flight, multi-dataset server routing, an ADBC connection pool, and a -broader integration test matrix. All changes are backward compatible; no -existing code needs to change to upgrade from 0.1.0. - ---- - -## What is new - -### TLS for Arrow Flight - -The built-in Flight server now supports encrypted connections. Pass a `:tls` -option to `Server.start_link/2`: - -```elixir -# One-way TLS (server presents a certificate) -cert = File.read!("server.crt") -key = File.read!("server.key") -{:ok, server} = ExArrow.Flight.Server.start_link(9999, - tls: [cert_pem: cert, key_pem: key]) - -# Mutual TLS (both sides present certificates) -ca = File.read!("ca.crt") -{:ok, server} = ExArrow.Flight.Server.start_link(9999, - tls: [cert_pem: cert, key_pem: key, ca_cert_pem: ca]) -``` - -The client already selected TLS automatically for non-loopback hosts (using the -OS certificate store). For a custom or self-signed CA, pass -`tls: [ca_cert_pem: pem]` to `Client.connect/3`. - -Plaintext (`tls: false`, or no `:tls` option on loopback) continues to work -exactly as before. - ---- - -### Flight server routing β€” multiple named datasets - -The built-in Flight server now stores datasets in a `HashMap` keyed by ticket, -rather than always overwriting a single `"echo"` slot. - -Upload with a descriptor to store under a named ticket: - -```elixir -:ok = ExArrow.Flight.Client.do_put(client, schema, batches, - descriptor: {:cmd, "sales_2024"}) - -:ok = ExArrow.Flight.Client.do_put(client, schema, other_batches, - descriptor: {:path, ["metrics", "daily"]}) -``` - -Retrieve by ticket: - -```elixir -{:ok, stream} = ExArrow.Flight.Client.do_get(client, "sales_2024") -{:ok, stream} = ExArrow.Flight.Client.do_get(client, "metrics/daily") -``` - -List all stored datasets: - -```elixir -{:ok, flights} = ExArrow.Flight.Client.list_flights(client) -{:ok, tickets} = ExArrow.Flight.Client.do_action(client, "list_tickets", <<>>) -``` - -Calls that do not pass a `:descriptor` default to the `"echo"` ticket, so all -existing code continues to work without modification. - ---- - -### ADBC connection pool - -`ExArrow.ADBC.ConnectionPool` is a NimblePool-backed pool that reuses open -`Connection` handles across callers. - -#### Supervised pool (recommended) - -```elixir -children = [ - {ExArrow.ADBC.DatabaseServer, - name: :mydb, - driver_path: "/usr/local/lib/libadbc_driver_duckdb.so"}, - {ExArrow.ADBC.ConnectionPool, - name: :mypool, database: :mydb, pool_size: 4} -] -Supervisor.start_link(children, strategy: :one_for_one) - -# Anywhere in the application: -{:ok, stream} = ExArrow.ADBC.ConnectionPool.query(:mypool, - "SELECT * FROM events WHERE day = today()") -``` - -#### Ad-hoc pool - -```elixir -{:ok, db} = ExArrow.ADBC.Database.open(driver_path: "/path/to/driver.so") -{:ok, pool} = ExArrow.ADBC.ConnectionPool.start_link(database: db, pool_size: 4) -{:ok, stream} = ExArrow.ADBC.ConnectionPool.query(pool, "SELECT 42 AS n") -``` - -#### Multi-statement checkout - -```elixir -ExArrow.ADBC.ConnectionPool.with_connection(pool, fn conn -> - {:ok, stmt} = ExArrow.ADBC.Statement.new(conn) - ExArrow.ADBC.Statement.set_sql(stmt, "BEGIN") - ExArrow.ADBC.Statement.execute(stmt) - # ... more statements ... - {result, :ok} -end) -``` - ---- - -### Larger integration test matrix - -A new `.github/workflows/integration.yml` workflow runs ADBC integration tests -against: - -- **PostgreSQL** 14, 15, and 16 (via GitHub Actions service containers) -- **DuckDB** 1.1.3 and 1.2.0 (via downloaded ADBC driver binary) - -The tests live in `test/ex_arrow/adbc_integration_test.exs` and are excluded -from the default test run. To run them locally: - -```bash -# Start PostgreSQL -docker run -d -p 5432:5432 -e POSTGRES_PASSWORD=postgres postgres:16 - -PG_HOST=localhost mix test --include adbc_integration \ - test/ex_arrow/adbc_integration_test.exs - -# DuckDB -DUCKDB_DRIVER=/usr/local/lib/libduckdb_adbc.so \ - mix test --include adbc_integration \ - test/ex_arrow/adbc_integration_test.exs -``` - ---- - -### New public API additions - -| Module | Function | Description | -|--------|----------|-------------| -| `ExArrow.Schema` | `field_names/1` | Returns field names as `[String.t()]` | -| `ExArrow.Stream` | `to_list/1` | Collects all batches into a list | -| `ExArrow.ADBC.Database` | `close/1` | Explicit handle release | -| `ExArrow.ADBC.Connection` | `close/1` | Explicit handle release | -| `ExArrow.ADBC.ConnectionPool` | `start_link/1`, `query/3`, `with_connection/3` | NimblePool pool | -| `ExArrow.ADBC.DatabaseServer` | `start_link/1`, `get/1` | Supervised database handle | - ---- - -## Upgrade guide - -No breaking changes. The only API change that requires attention is if you -have a **Mox stub for `ClientBehaviour.do_put`**: the callback now takes four -arguments (`client, schema, batches, opts`). Update your mock expectation: - -```elixir -# Before (0.1.0) -Mox.expect(MyMock, :do_put, fn client, schema, batches -> :ok end) - -# After (0.2.0) -Mox.expect(MyMock, :do_put, fn client, schema, batches, _opts -> :ok end) -``` - ---- - -## Dependencies - -No new required dependencies. The connection pool requires `nimble_pool` -(already optional in 0.1.0); add it to your `mix.exs` to use -`ExArrow.ADBC.ConnectionPool`: - -```elixir -{:nimble_pool, "~> 1.1"} -``` - ---- - -## Full changelog - -See [CHANGELOG.md](CHANGELOG.md) for the complete list of changes. diff --git a/RELEASE_NOTES_0_3_0.md b/RELEASE_NOTES_0_3_0.md deleted file mode 100644 index e19e08a..0000000 --- a/RELEASE_NOTES_0_3_0.md +++ /dev/null @@ -1,237 +0,0 @@ -# ExArrow 0.3.0 β€” Release Notes - -**Released:** 2026-03-10 - -ExArrow 0.3.0 ships the four features promised on the v0.3 roadmap: Arrow -compute kernels, Parquet read/write, an Explorer bridge, and an Nx bridge. -All changes are backward compatible; upgrading from 0.2.0 requires only a -version bump in `mix.exs`. - ---- - -## What is new - -### Arrow compute kernels - -`ExArrow.Compute` exposes three operations that run entirely inside native -Arrow memory via the `arrow-select` and `arrow-ord` Rust crates. No column -data is ever copied into BEAM terms; every function returns a new -`ExArrow.RecordBatch` handle. - -```elixir -# Select only the columns you need -{:ok, slim} = ExArrow.Compute.project(batch, ["user_id", "score", "region"]) - -# Sort by score descending -{:ok, sorted} = ExArrow.Compute.sort(slim, "score", ascending: false) - -# Filter to rows where is_active == true -# (build a boolean batch from a boolean column or a second query) -{:ok, active} = ExArrow.Compute.filter(sorted, predicate_batch) -``` - -Operations can be chained without any intermediate serialisation: - -```elixir -{:ok, result} = - ExArrow.Compute.project(batch, ["id", "score"]) - |> then(fn {:ok, b} -> ExArrow.Compute.sort(b, "score", ascending: false) end) - |> then(fn {:ok, b} -> ExArrow.Compute.filter(b, top_n_predicate) end) -``` - ---- - -### Parquet support - -`ExArrow.Parquet.Reader` and `ExArrow.Parquet.Writer` wrap the `parquet` Rust -crate and expose the same `ExArrow.Stream` interface as IPC and ADBC. - -**Read from a file or binary:** - -```elixir -{:ok, stream} = ExArrow.Parquet.Reader.from_file("/data/events.parquet") -{:ok, schema} = ExArrow.Stream.schema(stream) -batches = ExArrow.Stream.to_list(stream) - -# Or from a binary downloaded from S3 / object storage -{:ok, stream} = ExArrow.Parquet.Reader.from_binary(parquet_bytes) -``` - -**Write to a file or binary:** - -```elixir -:ok = ExArrow.Parquet.Writer.to_file("/out/result.parquet", schema, batches) - -# Or capture as an in-memory binary (e.g. to upload to S3) -{:ok, parquet_bytes} = ExArrow.Parquet.Writer.to_binary(schema, batches) -``` - -**Round-trip with compute kernels:** - -```elixir -{:ok, stream} = ExArrow.Parquet.Reader.from_file("/data/raw.parquet") -batches = ExArrow.Stream.to_list(stream) -{:ok, schema} = ExArrow.Parquet.Reader.from_file("/data/raw.parquet") - |> then(fn {:ok, s} -> ExArrow.Stream.schema(s) end) - -# Project and sort before writing -{:ok, processed} = ExArrow.Compute.project(hd(batches), ["id", "score"]) -{:ok, sorted} = ExArrow.Compute.sort(processed, "score") -:ok = ExArrow.Parquet.Writer.to_file("/data/sorted.parquet", schema, [sorted]) -``` - ---- - -### Explorer bridge module - -`ExArrow.Explorer` converts between `ExArrow.Stream` / `ExArrow.RecordBatch` -and `Explorer.DataFrame` with a single function call. The bridge uses Arrow -IPC internally; no CSV or row-by-row conversion is performed. - -Requires `{:explorer, "~> 0.8"}` in your `mix.exs`. - -**From a stream or batch to a DataFrame:** - -```elixir -{:ok, stream} = ExArrow.IPC.Reader.from_file("/data/events.arrow") -{:ok, df} = ExArrow.Explorer.from_stream(stream) -Explorer.DataFrame.filter(df, score > 0.9) - -# Single batch -batch = ExArrow.Stream.next(stream) -{:ok, df} = ExArrow.Explorer.from_record_batch(batch) -``` - -**From a DataFrame back to an ExArrow stream:** - -```elixir -df = Explorer.DataFrame.new(x: [1, 2, 3], y: ["a", "b", "c"]) - -{:ok, stream} = ExArrow.Explorer.to_stream(df) -{:ok, batches} = ExArrow.Explorer.to_record_batches(df) -``` - -**Full pipeline β€” query, compute, analyse:** - -```elixir -{:ok, stream} = ExArrow.ADBC.Statement.execute(stmt) -{:ok, slim} = ExArrow.Compute.project(ExArrow.Stream.next(stream), - ["user_id", "score"]) -{:ok, df} = ExArrow.Explorer.from_record_batch(slim) -Explorer.DataFrame.describe(df) -``` - -When Explorer is not present, every function returns -`{:error, "Explorer is not available …"}`. - ---- - -### Nx bridge module - -`ExArrow.Nx` converts Arrow numeric columns to `Nx.Tensor` values (and back) -by copying the raw byte buffer once β€” no list materialisation. - -Requires `{:nx, "~> 0.7"}` in your `mix.exs`. - -**Column to tensor:** - -```elixir -{:ok, stream} = ExArrow.Parquet.Reader.from_file("/data/trades.parquet") -batch = ExArrow.Stream.next(stream) -{:ok, prices} = ExArrow.Nx.column_to_tensor(batch, "price") -mean_price = prices |> Nx.mean() |> Nx.to_number() -``` - -**All numeric columns at once:** - -```elixir -{:ok, tensors} = ExArrow.Nx.to_tensors(batch) -# %{"price" => #Nx.Tensor, "qty" => #Nx.Tensor} -sorted_prices = tensors["price"] |> Nx.sort() -``` - -**Tensor back to a record batch:** - -```elixir -weights = Nx.tensor([0.1, 0.2, 0.7], type: {:f, 64}) -{:ok, batch} = ExArrow.Nx.from_tensor(weights, "weights") -``` - -Supported Arrow β†’ Nx type mappings: - -| Arrow type | Nx dtype | -|------------|------------| -| Int8 | `{:s, 8}` | -| Int16 | `{:s, 16}` | -| Int32 | `{:s, 32}` | -| Int64 | `{:s, 64}` | -| UInt8 | `{:u, 8}` | -| UInt16 | `{:u, 16}` | -| UInt32 | `{:u, 32}` | -| UInt64 | `{:u, 64}` | -| Float32 | `{:f, 32}` | -| Float64 | `{:f, 64}` | - -Non-numeric columns return `{:error, "unsupported column type…"}` from -`column_to_tensor/2` and are silently skipped by `to_tensors/1`. When Nx -is not present, every function returns `{:error, "Nx is not available …"}`. - ---- - -## New public API - -| Module | Function | Description | -|--------|----------|-------------| -| `ExArrow.Compute` | `filter/2` | Filter rows using a boolean-typed predicate batch | -| `ExArrow.Compute` | `project/2` | Select and reorder columns by name | -| `ExArrow.Compute` | `sort/3` | Sort batch by a named column | -| `ExArrow.Parquet.Reader` | `from_file/1` | Open a Parquet file as an `ExArrow.Stream` | -| `ExArrow.Parquet.Reader` | `from_binary/1` | Open a Parquet binary as an `ExArrow.Stream` | -| `ExArrow.Parquet.Writer` | `to_file/3` | Write schema + batches to a Parquet file | -| `ExArrow.Parquet.Writer` | `to_binary/2` | Serialize schema + batches to a Parquet binary | -| `ExArrow.Explorer` | `from_stream/1` | Convert `ExArrow.Stream` β†’ `Explorer.DataFrame` | -| `ExArrow.Explorer` | `from_record_batch/1` | Convert `ExArrow.RecordBatch` β†’ `Explorer.DataFrame` | -| `ExArrow.Explorer` | `to_stream/1` | Convert `Explorer.DataFrame` β†’ `ExArrow.Stream` | -| `ExArrow.Explorer` | `to_record_batches/1` | Convert `Explorer.DataFrame` β†’ `[ExArrow.RecordBatch]` | -| `ExArrow.Nx` | `column_to_tensor/2` | Arrow column β†’ `Nx.Tensor` | -| `ExArrow.Nx` | `to_tensors/1` | All numeric columns β†’ `%{name => Nx.Tensor}` | -| `ExArrow.Nx` | `from_tensor/2` | `Nx.Tensor` β†’ single-column `ExArrow.RecordBatch` | - ---- - -## Optional dependencies - -No new required dependencies. The two new bridge modules each unlock when -their optional dep is present: - -```elixir -# Explorer bridge (ExArrow.Explorer) -{:explorer, "~> 0.8"} - -# Nx bridge (ExArrow.Nx) -{:nx, "~> 0.7"} -``` - -The existing optional deps (ADBC, NimblePool) are unchanged. - ---- - -## Upgrade guide - -No breaking changes. Update your version pin: - -```elixir -# Before -{:ex_arrow, "~> 0.2.0"} - -# After -{:ex_arrow, "~> 0.3.0"} -``` - -Then run `mix deps.get` and `mix compile`. - ---- - -## Full changelog - -See [CHANGELOG.md](CHANGELOG.md) for the complete list of changes. diff --git a/RELEASE_NOTES_0_5_0.md b/RELEASE_NOTES_0_5_0.md deleted file mode 100644 index ed1af45..0000000 --- a/RELEASE_NOTES_0_5_0.md +++ /dev/null @@ -1,270 +0,0 @@ -# ExArrow 0.5.0 β€” Release Notes - -**Released:** 2026-04-16 - -ExArrow 0.5.0 adds a production-grade Arrow Flight SQL client, making Elixir -a first-class participant in the Flight SQL ecosystem alongside DuckDB, -DataFusion, Dremio, and InfluxDB v3. The release also delivers lazy streaming -with `Enumerable` support for all stream types, a Mox-compatible behaviour for -unit testing, and structured error types with gRPC status codes. - -All changes are backward compatible; upgrading from 0.4.0 requires only a -version bump. - ---- - -## What is new - -### Arrow Flight SQL client - -Arrow Flight SQL layers SQL query semantics on top of Arrow Flight (gRPC + -Arrow IPC). Queries are dispatched to the server, which executes them and -streams results back as columnar `RecordBatch` data β€” the same Arrow format -used everywhere in ExArrow. - -**Quick start:** - -```elixir -{:ok, client} = ExArrow.FlightSQL.Client.connect("localhost:32010") - -# Materialised query β€” all batches collected before returning -{:ok, result} = ExArrow.FlightSQL.Client.query(client, "SELECT id, name FROM users") -result.num_rows #=> 42 -result.schema #=> %ExArrow.Schema{...} - -# Lazy query β€” stream one batch at a time -{:ok, stream} = ExArrow.FlightSQL.Client.stream_query(client, "SELECT * FROM big_table") -Enum.each(stream, fn batch -> process(batch) end) - -# DML -{:ok, 3} = ExArrow.FlightSQL.Client.execute(client, "DELETE FROM t WHERE id < 4") -{:ok, :unknown} = ExArrow.FlightSQL.Client.execute(client, "CREATE TABLE t (id INT)") -``` - -**TLS connections** β€” plaintext is used automatically for loopback addresses; -remote hosts use the OS trust store; a custom CA certificate can be provided: - -```elixir -# TLS with OS trust store (automatic for remote hosts) -{:ok, client} = ExArrow.FlightSQL.Client.connect("dremio.example.com:32010") - -# Custom CA -pem = File.read!("priv/ca.pem") -{:ok, client} = ExArrow.FlightSQL.Client.connect("secure.server:32010", - tls: [ca_cert_pem: pem]) -``` - -**Bearer-token authentication:** - -```elixir -{:ok, client} = ExArrow.FlightSQL.Client.connect("dremio.example.com:32010", - tls: true, - headers: [{"authorization", "Bearer my-pat-token"}] -) -``` - ---- - -### Lazy streaming with `Enumerable` - -`ExArrow.Stream` now implements the `Enumerable` protocol, so all `Enum.*` -and `Stream.*` functions work directly on any stream handle β€” IPC, Parquet, -ADBC, and Flight SQL alike. - -```elixir -{:ok, stream} = ExArrow.FlightSQL.Client.stream_query(client, "SELECT * FROM events") - -# Collect all batches -batches = Enum.to_list(stream) - -# Map, filter, reduce β€” standard Elixir idioms -row_counts = Enum.map(stream, &ExArrow.RecordBatch.num_rows/1) - -# Take only the first N batches β€” the rest are never fetched -first_two = Enum.take(stream, 2) - -# Comprehension syntax -for batch <- stream, do: process_batch(batch) -``` - -Early termination (e.g. `Enum.take/2`) is safe β€” the underlying gRPC channel -is released when the stream variable goes out of scope. - ---- - -### Prepared statements - -Server-side prepared statements allow the server to parse and plan a query -once and then execute it one or more times: - -```elixir -{:ok, stmt} = ExArrow.FlightSQL.Client.prepare(client, - "SELECT * FROM events WHERE ts > '2024-01-01'") - -# Execute as a streaming query -{:ok, stream} = ExArrow.FlightSQL.Statement.execute(stmt) -batches = Enum.to_list(stream) - -# Re-execute the same statement (reuses the server plan) -{:ok, stream2} = ExArrow.FlightSQL.Statement.execute(stmt) - -# Or execute as DML -{:ok, dml_stmt} = ExArrow.FlightSQL.Client.prepare(client, - "DELETE FROM logs WHERE ts < '2020-01-01'") -{:ok, 1042} = ExArrow.FlightSQL.Statement.execute_update(dml_stmt) -``` - -Servers that do not support prepared statements return -`{:error, %Error{code: :unimplemented}}`. - ---- - -### Metadata discovery - -```elixir -{:ok, stream} = ExArrow.FlightSQL.Client.get_tables(client, - db_schema_filter: "public", table_types: ["TABLE", "VIEW"]) -batches = Enum.to_list(stream) - -{:ok, stream} = ExArrow.FlightSQL.Client.get_db_schemas(client) -{:ok, stream} = ExArrow.FlightSQL.Client.get_sql_info(client) -``` - ---- - -### Explorer and Nx integration - -```elixir -# Query β†’ Explorer DataFrame -{:ok, result} = ExArrow.FlightSQL.Client.query(client, "SELECT * FROM sales") -{:ok, df} = ExArrow.FlightSQL.Result.to_dataframe(result) - -# Query β†’ Nx tensor (first batch only) -{:ok, result} = ExArrow.FlightSQL.Client.query(client, "SELECT price FROM quotes") -{:ok, tensor} = ExArrow.FlightSQL.Result.to_tensor(result, "price") - -# Lazy stream β†’ Explorer DataFrame (large result sets) -{:ok, stream} = ExArrow.FlightSQL.Client.stream_query(client, "SELECT * FROM big_table") -{:ok, df} = ExArrow.Explorer.from_stream(stream) -``` - -`to_dataframe/1` requires `{:explorer, "~> 0.11"}`. `to_tensor/2` requires -`{:nx, "~> 0.9"}`. Both return -`{:error, %ExArrow.FlightSQL.Error{code: :conversion_error}}` when the -optional dependency is absent. - ---- - -### Mox-compatible behaviour for unit testing - -Swap the real implementation for a mock without a live server: - -```elixir -# test/test_helper.exs -Mox.defmock(MyApp.FlightSQLMock, for: ExArrow.FlightSQL.ClientBehaviour) - -# In your test -Application.put_env(:ex_arrow, :flight_sql_client_impl, MyApp.FlightSQLMock) - -Mox.expect(MyApp.FlightSQLMock, :query, fn _client, "SELECT 1", [] -> - {:ok, %ExArrow.FlightSQL.Result{schema: schema, batches: [], num_rows: 0}} -end) -``` - ---- - -### Structured errors - -All non-bang functions return `{:ok, value}` or -`{:error, %ExArrow.FlightSQL.Error{}}`: - -```elixir -case ExArrow.FlightSQL.Client.query(client, sql) do - {:ok, result} -> handle(result) - {:error, %Error{code: :unauthenticated}} -> reauthenticate() - {:error, %Error{code: :not_found, message: msg}} -> Logger.warn(msg) - {:error, err} -> raise err -end -``` - -Error codes: `:transport_error`, `:server_error`, `:unimplemented`, -`:unauthenticated`, `:permission_denied`, `:not_found`, `:invalid_argument`, -`:protocol_error`, `:multi_endpoint`, `:invalid_option`, `:conversion_error`. - ---- - -## New public API - -| Module | Function | Description | -|---|---|---| -| `ExArrow.FlightSQL.Client` | `connect/1,2` | Connect to a Flight SQL server | -| `ExArrow.FlightSQL.Client` | `query/2`, `query!/2` | Materialised SQL query | -| `ExArrow.FlightSQL.Client` | `stream_query/2` | Lazy SQL query returning `ExArrow.Stream` | -| `ExArrow.FlightSQL.Client` | `execute/2` | DML/DDL with affected-row count | -| `ExArrow.FlightSQL.Client` | `prepare/2` | Server-side prepared statement | -| `ExArrow.FlightSQL.Client` | `get_tables/1,2` | List tables visible to the connected user | -| `ExArrow.FlightSQL.Client` | `get_db_schemas/1,2` | List database schemas | -| `ExArrow.FlightSQL.Client` | `get_sql_info/1` | Server capability flags | -| `ExArrow.FlightSQL.Statement` | `execute/1` | Execute a prepared statement as a lazy stream | -| `ExArrow.FlightSQL.Statement` | `execute_update/1` | Execute a prepared DML statement | -| `ExArrow.FlightSQL.Result` | `from_stream/1` | Materialise a stream into a `Result` struct | -| `ExArrow.FlightSQL.Result` | `to_dataframe/1` | Convert result to `Explorer.DataFrame` | -| `ExArrow.FlightSQL.Result` | `to_tensor/2` | Extract a numeric column as `Nx.Tensor` | -| `ExArrow.Stream` | β€” | Now implements `Enumerable` | - ---- - -## Changed behaviour - -**`ExArrow.Stream` is now `Enumerable`** β€” `Enum.to_list/1`, `Enum.map/2`, -`Enum.take/2`, and all other `Enum.*` / `Stream.*` functions work directly on -stream handles. Existing code using `Stream.next/1` and `Stream.to_list/1` -continues to work unchanged. - ---- - -## Bug fixes - -**Elixir 1.17+ typing warning in `Adbc.Result.from_py/1`** β€” the -`{:ok, stream_ref, capsule}` match was unreachable when `Pythonx` is not -loaded. Both `from_py/1` and `from_py!/1` are now guarded with -`Code.ensure_loaded?(Pythonx)`, eliminating the "clause will never match" -compiler warning. - ---- - -## Dependencies - -No new required dependencies. Optional dependencies for new features: - -```elixir -# Required only for TLS with a custom CA (Flight SQL connect option) -# Uses OTP :ssl and :public_key β€” no new Hex packages needed. - -# Optional (unchanged from 0.4.0 β€” enable for ecosystem bridges) -{:explorer, "~> 0.11", optional: true} # Result.to_dataframe/1 -{:nx, "~> 0.9", optional: true} # Result.to_tensor/2 -``` - ---- - -## Upgrade guide - -No breaking changes. Update your version pin: - -```elixir -# Before -{:ex_arrow, "~> 0.4.0"} - -# After -{:ex_arrow, "~> 0.5.0"} -``` - -Then run `mix deps.get` and `mix compile`. - ---- - -## Full changelog - -See [CHANGELOG.md](CHANGELOG.md) for the complete list of changes including -internal fixes and documentation updates. diff --git a/bench/run_all.exs b/bench/run_all.exs index c73a21c..50768a7 100644 --- a/bench/run_all.exs +++ b/bench/run_all.exs @@ -17,7 +17,9 @@ scripts = [ "bench/adbc_bench.exs", "bench/pipeline_bench.exs", "bench/explorer_arrow_bench.exs", - "bench/nx_arrow_bench.exs" + "bench/nx_arrow_bench.exs", + "bench/v070_stream_flow_pipeline_bench.exs", + "bench/v070_record_batch_vs_maps_bench.exs" ] for script <- scripts do diff --git a/bench/v070_record_batch_vs_maps_bench.exs b/bench/v070_record_batch_vs_maps_bench.exs new file mode 100644 index 0000000..e637c79 --- /dev/null +++ b/bench/v070_record_batch_vs_maps_bench.exs @@ -0,0 +1,135 @@ +Code.require_file("bench_helper.exs", __DIR__) + +# --------------------------------------------------------------------------- +# ExArrow v0.7.0 β€” Arrow RecordBatch vs list(map()) throughput. +# +# Compares the cost of building, transforming (select), and draining a +# dataset held as Arrow RecordBatch values versus a list of row maps. +# +# Datasets: 1K, 100K, 1M rows. +# --------------------------------------------------------------------------- + +IO.puts("\n== Arrow RecordBatch vs list(map()) Benchmark ==\n") + +alias ExArrow.RecordBatch + +# Build a single s64 column "v" with `n` rows as an Arrow batch. +build_batch = fn n -> + bin = + 1..n + |> Enum.map(&<<&1::little-signed-64>>) + |> IO.iodata_to_binary() + + {:ok, batch} = RecordBatch.from_columns(["v"], [bin], ["s64"], n) + batch +end + +# Build the equivalent list of row maps. +build_maps = fn n -> + for i <- 1..n, do: %{"v" => i} +end + +batch_1k = build_batch.(1_000) +batch_100k = build_batch.(100_000) +batch_1m = build_batch.(1_000_000) + +maps_1k = build_maps.(1_000) +maps_100k = build_maps.(100_000) +# 1M maps is large; build lazily inside the scenario to avoid holding it +# in memory for the whole run. + +# ── Build benchmarks ────────────────────────────────────────────────────────── + +Benchee.run( + %{ + "build RecordBatch (1K rows)" => fn -> build_batch.(1_000) end, + "build RecordBatch (100K rows)" => fn -> build_batch.(100_000) end, + "build RecordBatch (1M rows)" => fn -> build_batch.(1_000_000) end, + "build list(map()) (1K rows)" => fn -> build_maps.(1_000) end, + "build list(map()) (100K rows)" => fn -> build_maps.(100_000) end, + "build list(map()) (1M rows)" => fn -> build_maps.(1_000_000) end + }, + time: 5, + memory_time: 2, + warmup: 2, + formatters: [ + Benchee.Formatters.Console, + {Benchee.Formatters.HTML, + file: Path.join(Bench.DataGen.output_dir(), "v070_build.html"), auto_open: false}, + {Benchee.Formatters.JSON, file: Path.join(Bench.DataGen.output_dir(), "v070_build.json")} + ] +) + +# ── Transform benchmarks (select / map a single field) ──────────────────────── + +Benchee.run( + %{ + "Arrow select 1 column (1K rows)" => fn -> + {:ok, _} = ExArrow.Batch.select(batch_1k, ["v"]) + end, + "Arrow select 1 column (100K rows)" => fn -> + {:ok, _} = ExArrow.Batch.select(batch_100k, ["v"]) + end, + "Arrow select 1 column (1M rows)" => fn -> + {:ok, _} = ExArrow.Batch.select(batch_1m, ["v"]) + end, + "Enum.map row maps (1K rows)" => fn -> + _ = Enum.map(maps_1k, & &1["v"]) + end, + "Enum.map row maps (100K rows)" => fn -> + _ = Enum.map(maps_100k, & &1["v"]) + end, + "Enum.map row maps (1M rows)" => fn -> + maps = build_maps.(1_000_000) + _ = Enum.map(maps, & &1["v"]) + end + }, + time: 5, + memory_time: 2, + warmup: 2, + formatters: [ + Benchee.Formatters.Console, + {Benchee.Formatters.HTML, + file: Path.join(Bench.DataGen.output_dir(), "v070_transform.html"), auto_open: false}, + {Benchee.Formatters.JSON, file: Path.join(Bench.DataGen.output_dir(), "v070_transform.json")} + ] +) + +# ── Drain benchmarks (sum a column) ─────────────────────────────────────────── + +# For Arrow we extract the column buffer and sum the s64 values. +sum_arrow = fn batch -> + ref = RecordBatch.resource_ref(batch) + {:ok, {binary, "s64", _n}} = ExArrow.Native.record_batch_column_buffer(ref, "v") + sum = for <>, reduce: 0, do: (acc -> acc + v) + sum +end + +sum_maps = fn maps -> + Enum.reduce(maps, 0, fn m, acc -> acc + m["v"] end) +end + +Benchee.run( + %{ + "Arrow sum column (1K rows)" => fn -> sum_arrow.(batch_1k) end, + "Arrow sum column (100K rows)" => fn -> sum_arrow.(batch_100k) end, + "Arrow sum column (1M rows)" => fn -> sum_arrow.(batch_1m) end, + "Enum.reduce row maps (1K rows)" => fn -> sum_maps.(maps_1k) end, + "Enum.reduce row maps (100K rows)" => fn -> sum_maps.(maps_100k) end, + "Enum.reduce row maps (1M rows)" => fn -> + maps = build_maps.(1_000_000) + sum_maps.(maps) + end + }, + time: 5, + memory_time: 2, + warmup: 2, + formatters: [ + Benchee.Formatters.Console, + {Benchee.Formatters.HTML, + file: Path.join(Bench.DataGen.output_dir(), "v070_drain.html"), auto_open: false}, + {Benchee.Formatters.JSON, file: Path.join(Bench.DataGen.output_dir(), "v070_drain.json")} + ] +) + +:ok diff --git a/bench/v070_stream_flow_pipeline_bench.exs b/bench/v070_stream_flow_pipeline_bench.exs new file mode 100644 index 0000000..b65eee1 --- /dev/null +++ b/bench/v070_stream_flow_pipeline_bench.exs @@ -0,0 +1,172 @@ +Code.require_file("bench_helper.exs", __DIR__) + +# --------------------------------------------------------------------------- +# ExArrow v0.7.0 β€” Arrow-native stream / Flow / pipeline benchmarks. +# +# Scenarios: +# +# 1. parquet_stream – open and drain a Parquet stream. +# 2. ipc_stream – open and drain an IPC stream. +# 3. flow_from_batches – drain the same stream through ExArrow.Flow. +# 4. pipeline_map_parquet – run a Pipeline.map_batches + write_parquet. +# +# Datasets: 1K, 100K, 1M rows (replicated batches of the IPC fixture). +# --------------------------------------------------------------------------- + +IO.puts("\n== ExArrow v0.7.0 Stream / Flow / Pipeline Benchmark ==\n") + +alias ExArrow.IPC +alias ExArrow.Stream + +# Build a single-batch IPC binary with `rows` rows by replicating the +# fixture batch `rows / fixture_rows` times within one batch list and +# concatenating through the IPC writer. For simplicity we scale by +# replicating *batches* and accept that the row count is approximate. +{_schema, batches} = Bench.DataGen.schema_and_batches(1) +fixture_batch = hd(batches) +fixture_rows = ExArrow.RecordBatch.num_rows(fixture_batch) + +# Build IPC binaries for each target row count by replicating the fixture +# batch enough times to reach at least the target. +build_ipc_bin = fn target_rows -> + batches_needed = div(target_rows + fixture_rows - 1, fixture_rows) + {schema, batches} = Bench.DataGen.schema_and_batches(batches_needed) + {:ok, bin} = ExArrow.IPC.Writer.to_binary(schema, batches) + {schema, bin, batches_needed * fixture_rows} +end + +{schema_1k, bin_1k, rows_1k} = build_ipc_bin.(1_000) +{schema_100k, bin_100k, rows_100k} = build_ipc_bin.(100_000) +{schema_1m, bin_1m, rows_1m} = build_ipc_bin.(1_000_000) + +# Write Parquet fixtures for the parquet scenarios. +write_parquet_fixture = fn schema, bin, name -> + {:ok, stream} = IPC.Reader.from_binary(bin) + batches = Stream.to_list(stream) + path = Path.join(System.tmp_dir!(), "ex_arrow_bench_#{name}.parquet") + :ok = ExArrow.Parquet.Writer.to_file(path, schema, batches) + {path, batches} +end + +{pq_1k, _} = write_parquet_fixture.(schema_1k, bin_1k, "1k") +{pq_100k, _} = write_parquet_fixture.(schema_100k, bin_100k, "100k") +{pq_1m, _} = write_parquet_fixture.(schema_1m, bin_1m, "1m") + +output_dir = Bench.DataGen.output_dir() + +# ── Stream drain benchmarks ─────────────────────────────────────────────────── + +Benchee.run( + %{ + "Parquet stream drain (1K rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_parquet(pq_1k) + _ = Stream.to_list(s) + end, + "Parquet stream drain (100K rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_parquet(pq_100k) + _ = Stream.to_list(s) + end, + "Parquet stream drain (1M rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_parquet(pq_1m) + _ = Stream.to_list(s) + end, + "IPC stream drain (1K rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_ipc(bin_1k) + _ = Stream.to_list(s) + end, + "IPC stream drain (100K rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_ipc(bin_100k) + _ = Stream.to_list(s) + end, + "IPC stream drain (1M rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_ipc(bin_1m) + _ = Stream.to_list(s) + end + }, + time: 5, + memory_time: 2, + warmup: 2, + formatters: [ + Benchee.Formatters.Console, + {Benchee.Formatters.HTML, file: Path.join(output_dir, "v070_stream.html"), auto_open: false}, + {Benchee.Formatters.JSON, file: Path.join(output_dir, "v070_stream.json")} + ] +) + +# ── Flow benchmarks ─────────────────────────────────────────────────────────── + +Benchee.run( + %{ + "Flow.map over IPC stream (100K rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_ipc(bin_100k) + + s + |> ExArrow.Flow.from_batches() + |> Flow.map(&ExArrow.RecordBatch.num_rows/1) + |> Enum.to_list() + end, + "Flow.map over IPC stream (1M rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_ipc(bin_1m) + + s + |> ExArrow.Flow.from_batches() + |> Flow.map(&ExArrow.RecordBatch.num_rows/1) + |> Enum.to_list() + end + }, + time: 5, + memory_time: 2, + warmup: 2, + formatters: [ + Benchee.Formatters.Console, + {Benchee.Formatters.HTML, file: Path.join(output_dir, "v070_flow.html"), auto_open: false}, + {Benchee.Formatters.JSON, file: Path.join(output_dir, "v070_flow.json")} + ] +) + +# ── Pipeline benchmarks ─────────────────────────────────────────────────────── + +pipeline_out = Path.join(System.tmp_dir!(), "ex_arrow_bench_pipeline_out.parquet") + +Benchee.run( + %{ + "Pipeline.map_batches + write_parquet (100K rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_ipc(bin_100k) + + {:ok, s} + |> ExArrow.Pipeline.map_batches(& &1) + |> ExArrow.Pipeline.write_parquet(pipeline_out) + + File.rm(pipeline_out) + end, + "Pipeline.map_batches + write_parquet (1M rows)" => fn -> + {:ok, s} = ExArrow.Stream.from_ipc(bin_1m) + + {:ok, s} + |> ExArrow.Pipeline.map_batches(& &1) + |> ExArrow.Pipeline.write_parquet(pipeline_out) + + File.rm(pipeline_out) + end + }, + time: 5, + memory_time: 2, + warmup: 2, + formatters: [ + Benchee.Formatters.Console, + {Benchee.Formatters.HTML, file: Path.join(output_dir, "v070_pipeline.html"), auto_open: false}, + {Benchee.Formatters.JSON, file: Path.join(output_dir, "v070_pipeline.json")} + ] +) + +# Cleanup +File.rm(pq_1k) +File.rm(pq_100k) +File.rm(pq_1m) + +IO.puts(""" + Approximate row counts: + 1K -> #{rows_1k} rows + 100K -> #{rows_100k} rows + 1M -> #{rows_1m} rows +""") diff --git a/docs/benchmarks.md b/docs/benchmarks.md index ccea9e2..83516a6 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -103,6 +103,20 @@ Scenarios: - **materialise β†’ Flight** β€” read the file, collect every batch into BEAM first, then upload. Shows the overhead of materialisation. +### v0.7.0 stream / Flow / pipeline (`bench/v070_stream_flow_pipeline_bench.exs`) + +Measures Parquet and IPC stream drains, Flow execution, and Pipeline +`map_batches` + `write_parquet` at 1K / 100K / 1M rows. Compares the cost +of draining a stream through `ExArrow.Stream.to_list/1` versus flowing the +same stream through `ExArrow.Flow.from_batches/1` and writing through +`ExArrow.Pipeline.write_parquet/2`. + +### v0.7.0 RecordBatch vs list(map()) (`bench/v070_record_batch_vs_maps_bench.exs`) + +Compares Arrow `RecordBatch` against `list(map())` for build, transform +(column select), and drain (sum a column) at 1K / 100K / 1M rows. Quantifies +the cost of going through row maps versus staying in native Arrow memory. + ## Published results Benchmark results from every push to `main` are stored in the `gh-pages` diff --git a/docs/ipc_guide.md b/docs/ipc_guide.md index 86072a0..0963e81 100644 --- a/docs/ipc_guide.md +++ b/docs/ipc_guide.md @@ -2,6 +2,11 @@ ExArrow supports the Arrow IPC format: streaming (sequential) and file format (random access). You can read and write record batches as a stream of bytes or to a file. +> **v0.7.0**: `ExArrow.Stream.from_ipc/1` and `from_ipc_file/1` are the +> preferred entry points for the streaming format β€” they tag the stream with +> `source` metadata for telemetry. `ExArrow.IPC.Reader.from_binary/1` and +> `from_file/1` continue to work unchanged. + ## Stream vs file - **Stream**: Sequential read/write. Use `Reader.from_binary/1` or `Reader.from_file/1` to get an `ExArrow.Stream` that yields record batches via `ExArrow.Stream.next/1`. Use `Writer.to_binary/2` or `Writer.to_file/3` to write batches. No random access. diff --git a/docs/overview.md b/docs/overview.md index e961f9c..ab54bea 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -3,7 +3,27 @@ The main overview, installation, quick start, and usage examples live in the [README on GitHub](https://github.com/thanos/ex_arrow/blob/main/README.md). -This doc set covers: +## What's changed in v0.7.0 + +v0.7.0 adds Arrow-native streaming and pipeline infrastructure. The unit of +execution is the Arrow `RecordBatch` β€” not `list(map())`, not +`Explorer.DataFrame`, not `Nx.Tensor`. Explorer and Nx remain downstream +consumers; ExArrow is the Arrow layer. + +New modules: `ExArrow.Stream` (constructors), `ExArrow.Batch` (transforms), +`ExArrow.Pipeline` (DSL), `ExArrow.Flow`, `ExArrow.GenStage.*Producer`, +`ExArrow.Broadway` (`BatchBuilder`, `ParquetSink`, `FlightSink`), +`ExArrow.Sink.*`, and `ExArrow.Telemetry`. Optional deps added: +`:telemetry`, `:flow`, `:gen_stage`, `:broadway`. + +New guides: [05 Arrow pipelines overview](05_arrow_pipelines_overview.md), +[06 Arrow streams](06_arrow_streams.md), +[07 Arrow and Flow](07_arrow_and_flow.md), +[08 Arrow and GenStage](08_arrow_and_genstage.md), +[09 Arrow and Broadway](09_arrow_and_broadway.md), +[10 Arrow pipeline patterns](10_arrow_pipeline_patterns.md). + +## Doc set | Topic | Guide | |-------|--------| @@ -30,5 +50,9 @@ This doc set covers: | `ExArrow.ADBC.ConnectionPool` | `{:nimble_pool, "~> 1.1"}` | NimblePool-backed connection pool for ADBC databases | | `:adbc_package` backend | `{:adbc, "~> 0.9"}` + `{:explorer, "~> 0.11"}` | Supervised pure-Elixir ADBC backend; `Database.open(:adbc_package)` | | ADBC driver download | `{:adbc, "~> 0.9"}` | `ExArrow.ADBC.DriverHelper.ensure_driver_and_open/2` | +| `ExArrow.Telemetry` | `{:telemetry, "~> 1.0"}` | Emit and observe telemetry events for every transport and pipeline operation | +| `ExArrow.Flow` | `{:flow, "~> 1.2"}` | Arrow-native Flow execution: `from_batches/1`, `map_batches/2`, `each_batch/2` | +| `ExArrow.GenStage.*Producer` | `{:gen_stage, "~> 1.2"}` | Demand-driven producers: `ParquetProducer`, `FlightProducer`, `ADBCProducer` | +| `ExArrow.Broadway` | `{:broadway, "~> 1.0"}` | Ingestion pipelines: `BatchBuilder`, `ParquetSink`, `FlightSink` | API reference: `mix docs` or [Hex Docs](https://hexdocs.pm/ex_arrow). diff --git a/docs/parquet_guide.md b/docs/parquet_guide.md index 2197333..369d9d0 100644 --- a/docs/parquet_guide.md +++ b/docs/parquet_guide.md @@ -5,6 +5,12 @@ ExArrow supports reading and writing Apache Parquet files via the Arrow Rust get the same `ExArrow.Stream` interface on the read side and the same schema + batches pattern on the write side. +> **v0.7.0**: `ExArrow.Stream.from_parquet/1` and +> `from_parquet_binary/1` are the preferred entry points β€” they tag the +> stream with `source` metadata and emit `[:ex_arrow, :parquet, :read]` +> telemetry. `ExArrow.Parquet.Reader.from_file/1` and `from_binary/1` +> continue to work unchanged. + --- ## Reading diff --git a/guides/05_arrow_pipelines_overview.md b/guides/05_arrow_pipelines_overview.md new file mode 100644 index 0000000..18bef73 --- /dev/null +++ b/guides/05_arrow_pipelines_overview.md @@ -0,0 +1,47 @@ +# Arrow Pipelines Overview + +ExArrow v0.7.0 introduces Arrow-native pipelines for the BEAM. This guide is +a short orientation to the pipeline modules and how they fit together; the +subsequent guides cover each in depth. + +## The principle + +The central architectural principle is: + + Operate on Arrow RecordBatch values. + Not list(map()). + Not Explorer.DataFrame. + Not Nx.Tensor. + +Explorer and Nx remain downstream consumers. ExArrow is the Arrow layer. + +## The modules + +| Module | Role | Guide | +|---------------------|----------------------------------------------------|-------| +| `ExArrow.Stream` | Open a stream of RecordBatch values from any source | [06 Arrow streams](06_arrow_streams.md) | +| `ExArrow.Batch` | Lightweight column/row transforms on a batch | (see module docs) | +| `ExArrow.Pipeline` | Lazy, composable DSL for transforming and sinking | [10 Pipeline patterns](10_arrow_pipeline_patterns.md) | +| `ExArrow.Flow` | Parallel batch processing via Flow | [07 Arrow and Flow](07_arrow_and_flow.md) | +| `ExArrow.GenStage` | Demand-driven producers with backpressure | [08 Arrow and GenStage](08_arrow_and_genstage.md) | +| `ExArrow.Broadway` | Ingestion pipelines (Kafka, SQS, ...) | [09 Arrow and Broadway](09_arrow_and_broadway.md) | +| `ExArrow.Sink.*` | Standard destinations (Parquet, Flight, DataFrame, Nx) | (see module docs) | +| `ExArrow.Telemetry` | Events for every transport and pipeline operation | (see module docs) | + +## Quick example + + ExArrow.Stream.from_parquet("/data/events.parquet") + |> ExArrow.Pipeline.map_batches(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id", "score"]) + slim + end) + |> ExArrow.Pipeline.write_parquet("/data/slim.parquet") + +## Where to start + +1. Read [06 Arrow streams](06_arrow_streams.md) to understand the streaming + abstraction and the `from_*/1` constructors. +2. Read [10 Pipeline patterns](10_arrow_pipeline_patterns.md) for the + `ExArrow.Pipeline` DSL and sinks. +3. Read the integration guides (07–09) when you need parallelism, + backpressure, or ingestion. diff --git a/guides/06_arrow_streams.md b/guides/06_arrow_streams.md new file mode 100644 index 0000000..e6eefb5 --- /dev/null +++ b/guides/06_arrow_streams.md @@ -0,0 +1,157 @@ +# Arrow Streams + +ExArrow v0.7.0 introduces first-class streaming as the primary mechanism for +working with large datasets. This guide explains the streaming abstraction, +the available sources, and the tradeoffs of each. + +## The unit of streaming: RecordBatch + +ExArrow streams yield `ExArrow.RecordBatch` values, not row maps. A +`RecordBatch` is an opaque handle to a native Arrow batch β€” a collection of +column arrays sharing a schema and row count. The column buffers stay in +native (Rust) memory until you explicitly extract them; only the small handle +crosses the BEAM heap. + +This is the central architectural principle of v0.7.0: + + Operate on Arrow RecordBatch values. + Not `list(map())`. Not `Explorer.DataFrame`. Not `Nx.Tensor`. + +Explorer and Nx remain downstream consumers. ExArrow is the Arrow layer. + +## The `ExArrow.Stream` constructors + +Every common source has a `from_*/1` constructor on `ExArrow.Stream` that +returns `{:ok, stream} | {:error, reason}`: + +| Constructor | Source | +|--------------------------------------|----------------------------------------------| +| `ExArrow.Stream.from_parquet/1` | Parquet file at `path` | +| `ExArrow.Stream.from_parquet_binary/1` | In-memory Parquet bytes | +| `ExArrow.Stream.from_ipc/1` | Arrow IPC stream binary | +| `ExArrow.Stream.from_ipc_file/1` | Arrow IPC file at `path` | +| `ExArrow.Stream.from_flight/2` | Flight `do_get` ticket | +| `ExArrow.Stream.from_flight_sql/2` | Flight SQL query | +| `ExArrow.Stream.from_adbc/1` | Pre-built ADBC statement | +| `ExArrow.Stream.from_adbc/2` | `{connection, sql}` one-shot ADBC query | + +Each constructor tags the stream with a `source` term (e.g. `{:parquet, path}` +or `{:flight_sql, sql}`) that is forwarded to telemetry events as `:source` +metadata. + +## Consuming a stream + +Three consumption patterns cover every use case. + +### 1. Lazy iteration with `next/1` + +`ExArrow.Stream.next/1` pulls one batch at a time and returns `nil` when the +stream is exhausted. This is the lowest-level interface and the one that +gives you explicit control over error handling. + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + + loop = fn s, f -> + case ExArrow.Stream.next(s) do + nil -> :done + {:error, reason} -> {:error, reason} + batch -> f.(batch); loop.(s, f) + end + end + + loop.(stream, fn batch -> + IO.puts("rows: #{ExArrow.RecordBatch.num_rows(batch)}") + end) + +For recoverable error handling, pattern-match on the `{:error, reason}` return +of `next/1` rather than letting `Enum` raise. + +### 2. Enumerable consumption + +`ExArrow.Stream` implements `Enumerable`, so all `Enum` and `Stream` functions +work directly on a stream handle: + + {:ok, stream} = ExArrow.Stream.from_flight_sql(client, "SELECT * FROM t") + + batches = Enum.to_list(stream) + row_counts = Enum.map(stream, &ExArrow.RecordBatch.num_rows/1) + first_two = Enum.take(stream, 2) + +`Enum.take/2` stops consuming early β€” the remaining batches are never fetched. +Enumeration raises on a transport or server error; for recoverable handling +use `next/1`. + +### 3. Pipeline consumption + +`ExArrow.Pipeline` (see the [Pipeline Patterns](10_arrow_pipeline_patterns.md) +guide) wraps a stream with composable `map_batches/2` and sink functions: + + ExArrow.Stream.from_parquet("/data/events.parquet") + |> ExArrow.Pipeline.map_batches(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id", "score"]) + slim + end) + |> ExArrow.Pipeline.write_parquet("/data/slim.parquet") + +The pipeline is lazy: nothing runs until a sink consumes it. + +## Schema preservation + +Every stream carries its Arrow schema. `ExArrow.Stream.schema/1` returns the +`ExArrow.Schema` handle without consuming any batches: + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + {:ok, schema} = ExArrow.Stream.schema(stream) + ExArrow.Schema.field_names(schema) # => ["id", "name", "score", ...] + +Field names, types, and nullability are preserved end-to-end through IPC, +Parquet, Flight, and ADBC. + +## Backpressure and laziness + +ExArrow streams are lazy by construction: + +- **Parquet** β€” the footer is scanned once on open (making the schema + available); row groups are decoded on demand by `next/1`. Stopping early + leaves the remaining row groups undecoded. +- **IPC** β€” batches are decoded one at a time as `next/1` is called. +- **Flight / Flight SQL** β€” the gRPC stream stays open until exhausted or + garbage-collected; each `next/1` pulls one server-side batch. +- **ADBC** β€” the driver's result iterator is consumed on demand. + +For GenStage-style demand-driven backpressure, use the +`ExArrow.GenStage.*Producer` modules (see the +[GenStage guide](08_arrow_and_genstage.md)). + +## Telemetry + +Every `from_*/1` constructor for Parquet, Flight, and Flight SQL emits a +source-level telemetry event (`[:ex_arrow, :parquet, :read]`, +`[:ex_arrow, :flight, :query]`, `[:ex_arrow, :flight_sql, :query]`). Each +`next/1` that yields a batch emits `[:ex_arrow, :stream, :batch]` with `rows`, +`columns`, and `batch_count` measurements. + +See the [Telemetry](#) module docs for the full event list and a handler +example. + +## Choosing a source + +| Need | Use | +|---------------------------------------|-------------------------------------------------| +| Read a columnar file from disk | `from_parquet/1` or `from_ipc_file/1` | +| Read bytes already in memory | `from_parquet_binary/1` or `from_ipc/1` | +| Query a Flight / Flight SQL server | `from_flight/2` or `from_flight_sql/2` | +| Query via an ADBC driver | `from_adbc/1` (prepared) or `from_adbc/2` (SQL) | +| Compose with Flow / GenStage / Broadway | `ExArrow.Flow`, `ExArrow.GenStage`, `ExArrow.Broadway` | + +## Example: copy a Parquet file through Arrow + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + {:ok, schema} = ExArrow.Stream.schema(stream) + batches = ExArrow.Stream.to_list(stream) + :ok = ExArrow.Parquet.Writer.to_file("/data/events_copy.parquet", schema, batches) + +Or, using the Pipeline DSL: + + ExArrow.Stream.from_parquet("/data/events.parquet") + |> ExArrow.Pipeline.write_parquet("/data/events_copy.parquet") diff --git a/guides/07_arrow_and_flow.md b/guides/07_arrow_and_flow.md new file mode 100644 index 0000000..140a7e3 --- /dev/null +++ b/guides/07_arrow_and_flow.md @@ -0,0 +1,120 @@ +# Arrow and Flow + +`ExArrow.Flow` brings parallel, batch-oriented processing to Arrow streams by +wrapping the [Flow](https://github.com/elixir-lang/flow) library. This guide +explains how it works, when to use it, and the performance tradeoffs. + +## Why Flow? + +`Enum` and `Stream` run in a single process. For CPU-bound transformations +on large Arrow datasets that single process becomes the bottleneck. Flow +spreads the work across multiple stages (processes), each consuming batches +independently. + +The unit of work in `ExArrow.Flow` is the **batch**, not the row. A Flow +stage receives an `ExArrow.RecordBatch` handle and returns one. Because the +handle is an opaque reference to native memory, no column buffers are copied +to the BEAM heap when a batch moves between stages β€” only the small reference +term is sent over the mailbox. + +## Building a Flow + +`ExArrow.Flow.from_batches/1` accepts an `ExArrow.Stream.t()` (or any +`Enumerable.t()` of batches) and returns a `Flow.t()`: + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + + stream + |> ExArrow.Flow.from_batches() + |> Flow.map(&ExArrow.RecordBatch.num_rows/1) + |> Enum.to_list() + +The function also unwraps `{:ok, stream}` results so it composes directly with +`ExArrow.Stream.from_*/1` in a pipe: + + ExArrow.Stream.from_parquet("/data/events.parquet") + |> ExArrow.Flow.from_batches() + |> Flow.map(&ExArrow.RecordBatch.num_rows/1) + |> Enum.to_list() + +`opts` are forwarded to `Flow.from_enumerable/2`: + + ExArrow.Flow.from_batches(stream, stages: 8, max_demand: 4) + +## Combinators + +All standard `Flow` combinators work: + +- `Flow.map/2` β€” transform each batch. +- `Flow.flat_map/2` β€” expand one batch into many. +- `Flow.partition/2` β€” partition batches by key for shuffled reductions. +- `Flow.reduce/3` β€” reduce batches within a window/partition. + +`ExArrow.Flow` adds two telemetry-emitting helpers: + +- `ExArrow.Flow.map_batches/2` β€” `Flow.map/2` plus a + `[:ex_arrow, :pipeline, :batch]` event per batch. +- `ExArrow.Flow.each_batch/2` β€” run a side effect per batch, batches pass + through unchanged. + +## Performance implications + +### Parallelism + +Flow spins up a configurable number of producer and consumer stages +(`:stages`, `:max_demand`, `:min_demand`). Each stage decodes and transforms +batches independently, so wall-clock time scales with available cores for +CPU-bound work. + +### Memory + +Only batch references cross process boundaries; the Arrow buffers stay in +native memory until a stage explicitly extracts them. Peak memory is roughly +`stages * largest_batch` rather than the whole dataset. + +### Backpressure + +GenStage demand is honoured end-to-end, so a slow consumer slows the producer +without piling up batches. + +### Not a row API + +Converting batches to row maps inside a Flow stage defeats the purpose. Keep +transformations column-wise β€” use `ExArrow.Batch` or `ExArrow.Compute` to +project, filter, or sort within a stage. + +## Example: parallel column projection + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + + stream + |> ExArrow.Flow.from_batches(stages: 4) + |> Flow.map(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id", "score"]) + slim + end) + |> Enum.to_list() + +## Example: partitioned reduction + + {:ok, stream} = ExArrow.Stream.from_flight_sql(client, "SELECT user_id, amount FROM sales") + + stream + |> ExArrow.Flow.from_batches() + |> Flow.partition(key: fn batch -> ExArrow.RecordBatch.column_names(batch) end) + |> Flow.reduce(fn -> %{} end, fn batch, acc -> + # merge batch into acc keyed by user_id + Map.merge(acc, summarise(batch), fn _k, a, b -> a + b end) + end) + |> Enum.to_list() + +## When to use Flow vs Pipeline + +| Use `ExArrow.Flow` when | Use `ExArrow.Pipeline` when | +|-----------------------------------------|---------------------------------------------| +| You need explicit parallelism control | You want a thin, stable abstraction | +| You want partition/reduce semantics | You are doing map-only transformations | +| You are comfortable with Flow's API | You want telemetry wired in automatically | + +`ExArrow.Pipeline` may internally use Flow in future releases; the public API +will stay the same. diff --git a/guides/08_arrow_and_genstage.md b/guides/08_arrow_and_genstage.md new file mode 100644 index 0000000..b0f0d51 --- /dev/null +++ b/guides/08_arrow_and_genstage.md @@ -0,0 +1,119 @@ +# Arrow and GenStage + +`ExArrow.GenStage` provides demand-driven producers that emit +`ExArrow.RecordBatch` values from the common ExArrow sources. This guide +explains the architecture, the three producers, and how to wire producer, +consumer, and producer-consumer pipelines. + +## Why GenStage? + +[GenStage](https://github.com/elixir-lang/gen_stage) is the standard Elixir +library for demand-driven data pipelines. A producer only emits events when +a consumer demands them, so slow consumers apply backpressure all the way to +the source β€” exactly what you want when reading a large Parquet file or +streaming a Flight SQL result. + +ExArrow's producers emit `ExArrow.RecordBatch` handles, not row maps. The +unit of work is the batch. + +## The three producers + +| Module | Source | +|---------------------------------------|----------------------------------------------| +| `ExArrow.GenStage.ParquetProducer` | Parquet file (`:path`) or binary (`:binary`) | +| `ExArrow.GenStage.FlightProducer` | Flight `do_get` (`:client` + `:ticket`) | +| `ExArrow.GenStage.ADBCProducer` | ADBC (`:statement` or `:connection` + `:sql`)| + +All three accept a pre-opened `:stream` option for testing or for sources not +covered by the dedicated options. + +## Lifecycle + +- **Demand-driven**: batches are only read when a consumer demands them. +- **Arrow batch delivery**: each emitted event is an `ExArrow.RecordBatch` + handle. +- **Clean shutdown**: when the underlying stream is exhausted the producer + sends itself a `{ExArrow.GenStage, :stop}` message and exits with reason + `:normal`. +- **Resource cleanup**: `terminate/2` drains the stream so file/socket + descriptors are released promptly. + +## Pattern 1: producer + consumer + + defmodule Collector do + use GenStage + + def init(pid), do: {:consumer, pid} + + def handle_events(batches, _from, pid) do + send(pid, {:batches, batches}) + {:noreply, [], pid} + end + end + + {:ok, producer} = + ExArrow.GenStage.ParquetProducer.start_link(path: "/data/events.parquet") + + {:ok, consumer} = GenStage.start_link(Collector, self()) + GenStage.sync_subscribe(consumer, to: producer, max_demand: 4) + +## Pattern 2: producer-consumer + +A producer-consumer transforms events between a producer and a consumer. This +is where you plug in `ExArrow.Batch` transformations: + + defmodule MyTransformer do + use GenStage + + def init(state), do: {:producer_consumer, state} + + def handle_events(batches, _from, state) do + transformed = + Enum.map(batches, fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id"]) + slim + end) + + {:noreply, transformed, state} + end + end + + {:ok, producer} = ExArrow.GenStage.ParquetProducer.start_link(path: "/data/events.parquet") + {:ok, transformer} = GenStage.start_link(MyTransformer, :ok) + {:ok, consumer} = GenStage.start_link(Collector, self()) + + GenStage.sync_subscribe(transformer, to: producer, max_demand: 1) + GenStage.sync_subscribe(consumer, to: transformer, max_demand: 1) + +## Pattern 3: producer + producer-consumer + consumer (full pipeline) + +Combine the two patterns for a three-stage pipeline: + + ParquetProducer ──► Selector ──► Collector + +Each stage honours demand, so the producer only reads row groups as fast as +the collector can acknowledge them. + +## Telemetry + +Every batch emitted by an ExArrow producer fires a +`[:ex_arrow, :stream, :batch]` telemetry event with `rows`, `columns`, and +`batch_count` measurements and `%{source: ...}` metadata. + +## Error handling + +If `ExArrow.Stream.next/1` returns `{:error, reason}` the producer sends +itself a stop message and exits `:normal`. Consumers see a producer `:DOWN` +notification. For finer-grained error propagation, wrap the producer in a +`Supervisor` with a restart strategy and monitor the producer process. + +## Choosing between Flow and GenStage + +| Use GenStage when | Use Flow when | +|----------------------------------------|---------------------------------------------| +| You need explicit demand/backpressure | You want Fire-and-forget parallelism | +| You are building a long-running pipeline | You are doing a one-shot batch job | +| You want to wire custom consumers | You want partition/reduce semantics | + +GenStage is the lower-level building block; Flow is built on top of it. +ExArrow exposes both so you can pick the right tool for each workload. diff --git a/guides/09_arrow_and_broadway.md b/guides/09_arrow_and_broadway.md new file mode 100644 index 0000000..4d98ac1 --- /dev/null +++ b/guides/09_arrow_and_broadway.md @@ -0,0 +1,129 @@ +# Arrow and Broadway + +`ExArrow.Broadway` integrates ExArrow with [Broadway](https://github.com/dashbitco/broadway), +the standard Elixir library for ingestion pipelines (Kafka, SQS, S3, ...). +This guide explains the architecture, the batch builder, and the Parquet and +Flight sinks. + +## Architecture + +The canonical ingestion pipeline looks like: + + Kafka / SQS / S3 + | (messages carry Arrow columnar payloads) + v + Broadway producer + | + v + Broadway processor ──► ExArrow.Broadway.BatchBuilder + | (groups messages into RecordBatch values) + v + Broadway batcher ──► batch_size / batch_timeout config + | + v + handle_batch/4 ──► ExArrow.Broadway.ParquetSink / FlightSink + +The unit that flows through the pipeline is an `ExArrow.RecordBatch` handle, +not a row map. + +## Message shape + +`ExArrow.Broadway.BatchBuilder` expects each Broadway message's `data` to be +one of: + +- an `ExArrow.RecordBatch.t()` handle (the common case for Arrow-aware + producers β€” e.g. a Kafka deserialiser that emits Arrow IPC), or +- a `{names, binaries, dtypes, length}` tuple describing raw Arrow columns, + which `BatchBuilder` converts to a batch via + `ExArrow.RecordBatch.from_columns/4`. + +`from_messages/1` returns `{:ok, schema, batches}` where `schema` is the +schema of the first batch and `batches` is the assembled batch list. + +## BatchBuilder + + def handle_batch(:parquet, messages, _info, _ctx) do + {:ok, schema, batches} = + ExArrow.Broadway.BatchBuilder.from_messages(messages) + + ExArrow.Broadway.ParquetSink.write("/data/events.parquet", schema, batches) + end + +## ParquetSink + +`ExArrow.Broadway.ParquetSink.write/3` writes the assembled batches to a +Parquet file in a single `ExArrow.Parquet.Writer.to_file/3` call. It emits a +`[:ex_arrow, :parquet, :write]` telemetry event with `:rows`, `:batch_count`, +and `%{destination: path, source: :broadway}` metadata. + +## FlightSink + +`ExArrow.Broadway.FlightSink.write/4` uploads the assembled batches to a +Flight server via `ExArrow.Flight.Client.do_put/4`. It emits a +`[:ex_arrow, :flight, :query]` telemetry event. + +## Tuning + +Batch sizing and flush intervals are controlled by the Broadway batcher +configuration (`:batch_size`, `:batch_timeout`), not by ExArrow. Typical +settings: + + batchers: [ + parquet: [concurrency: 2, batch_size: 100, batch_timeout: 1000] + ] + +- `:batch_size` β€” how many messages to accumulate before flushing. Larger + values produce fewer, larger Arrow batches (better throughput, more memory). +- `:batch_timeout` β€” maximum milliseconds to wait before flushing a partial + batch. Lower values reduce latency at the cost of smaller batches. +- `:concurrency` β€” number of batch processor processes. Increase for I/O- + bound sinks (Parquet to a network drive, Flight to a remote server). + +## Example pipeline + + defmodule MyPipeline do + use Broadway + + def start_link(opts) do + Broadway.start_link(__MODULE__, + name: __MODULE__, + producer: [module: {MyKafkaProducer, opts}], + processors: [default: [concurrency: 4]], + batchers: [ + parquet: [concurrency: 2, batch_size: 100, batch_timeout: 1000] + ] + ) + end + + def handle_message(:default, %Broadway.Message{data: batch} = msg, _ctx) + when is_struct(batch, ExArrow.RecordBatch) do + Broadway.Message.put_batcher(msg, :parquet) + end + + def handle_batch(:parquet, messages, _info, _ctx) do + {:ok, schema, batches} = + ExArrow.Broadway.BatchBuilder.from_messages(messages) + + ExArrow.Broadway.ParquetSink.write("/data/events.parquet", schema, batches) + end + end + +## Error handling + +`BatchBuilder.from_messages/1` returns `{:error, message}` if any message has +unsupported `data`. In `handle_batch/4` you should pattern-match on the +result and use `Broadway.Message.failed/2` to mark the messages for retry +rather than letting the batch handler crash. + +## Tradeoffs + +- **Batch size vs latency**: larger batches improve Arrow write throughput + but increase end-to-end latency. Tune `:batch_size` and `:batch_timeout` + together. +- **Memory**: each batch processor holds up to `:batch_size` batches in + memory. For very large batches, lower `:batch_size` or increase + `:concurrency`. +- **Arrow-aware producers**: for the best performance, have your Broadway + producer emit `ExArrow.RecordBatch` handles directly (e.g. by deserialising + Arrow IPC from Kafka). The `{names, binaries, dtypes, length}` tuple path + is a convenience for producers that work with raw column bytes. diff --git a/guides/10_arrow_pipeline_patterns.md b/guides/10_arrow_pipeline_patterns.md new file mode 100644 index 0000000..22a7b71 --- /dev/null +++ b/guides/10_arrow_pipeline_patterns.md @@ -0,0 +1,140 @@ +# Arrow Pipeline Patterns + +`ExArrow.Pipeline` is the stable, composable abstraction for transforming and +sinking Arrow streams. This guide walks through the common patterns and the +tradeoffs of each. + +## The Pipeline abstraction + +A pipeline wraps an `ExArrow.Stream.t()` with an Elixir `Stream` of +`ExArrow.RecordBatch` values alongside the stream's schema. Every function +accepts and returns `{:ok, pipeline} | {:error, reason}`, so pipelines compose +with `|>/2` directly from `ExArrow.Stream.from_*/1` constructors. Errors +short-circuit through every stage. + +The pipeline is **lazy**: `map_batches/2` and `each_batch/2` do no work until +a sink (`write_parquet/2`, `write_flight/3`, `write_dataframe/1`) runs. + +## map_batches/2 + +Transform each batch lazily. `fun` receives an `ExArrow.RecordBatch.t()` and +should return an `ExArrow.RecordBatch.t()` (or any term the downstream sink +expects). + + ExArrow.Stream.from_parquet("/data/events.parquet") + |> ExArrow.Pipeline.map_batches(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id", "score"]) + slim + end) + |> ExArrow.Pipeline.write_parquet("/data/slim.parquet") + +`map_batches/2` emits a `[:ex_arrow, :pipeline, :batch]` telemetry event per +batch. + +## each_batch/2 + +Run a side effect per batch without changing the pipeline. The batches pass +through unchanged. + + ExArrow.Stream.from_parquet("/data/events.parquet") + |> ExArrow.Pipeline.each_batch(fn batch -> + IO.puts("rows: #{ExArrow.RecordBatch.num_rows(batch)}") + end) + |> ExArrow.Pipeline.write_parquet("/data/copy.parquet") + +## write_parquet/2 + +Consume the pipeline and write every batch to a Parquet file. Triggers +evaluation of all upstream stages. Emits a +`[:ex_arrow, :parquet, :write]` telemetry event. + + ExArrow.Stream.from_flight_sql(client, "SELECT * FROM events") + |> ExArrow.Pipeline.write_parquet("/data/events.parquet") + +## write_flight/3 + +Consume the pipeline and upload every batch to a Flight server. `opts` are +forwarded to `ExArrow.Flight.Client.do_put/4`. + + ExArrow.Stream.from_parquet("/data/events.parquet") + |> ExArrow.Pipeline.write_flight(client, descriptor: {:cmd, "events"}) + +## write_dataframe/1 + +Consume the pipeline and convert it into an `Explorer.DataFrame`. Requires +the optional `{:explorer, "~> 0.11"}` dependency. + + {:ok, df} = + ExArrow.Stream.from_parquet("/data/events.parquet") + |> ExArrow.Pipeline.write_dataframe() + +## Composing stages + +Stages compose with `|>/2` because every function accepts the previous +stage's `{:ok, pipeline}` result: + + ExArrow.Stream.from_flight_sql(client, "SELECT * FROM events") + |> ExArrow.Pipeline.map_batches(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id", "score"]) + slim + end) + |> ExArrow.Pipeline.each_batch(&log_batch/1) + |> ExArrow.Pipeline.write_parquet("/data/slim.parquet") + +## Error propagation + +If a constructor or stage returns `{:error, reason}`, every subsequent stage +passes the error through unchanged: + + {:error, "no connection"} + |> ExArrow.Pipeline.map_batches(& &1) # => {:error, "no connection"} + |> ExArrow.Pipeline.write_parquet("/x.parquet") # => {:error, "no connection"} + +Use a `with` chain if you prefer explicit error handling: + + with {:ok, stream} <- ExArrow.Stream.from_parquet("/data/events.parquet"), + {:ok, _} <- ExArrow.Pipeline.write_parquet({:ok, stream}, "/data/copy.parquet") do + :ok + end + +## When to use Pipeline vs Flow vs GenStage + +| Pipeline | Flow | GenStage | +|--------------------------------|-------------------------------|----------------------------| +| Thin, stable abstraction | Explicit parallelism control | Demand-driven backpressure | +| Map-only transformations | Partition/reduce semantics | Long-running pipelines | +| Telemetry wired in | One-shot batch jobs | Custom consumer wiring | +| Single-process (lazy Stream) | Multi-process | Multi-process | + +`ExArrow.Pipeline` is the right starting point. Reach for `ExArrow.Flow` +when you need parallelism beyond a single process, and for `ExArrow.GenStage` +when you need explicit demand/backpressure or a long-running pipeline. + +## Schema handling + +The pipeline captures the stream's schema at wrap time (via +`ExArrow.Stream.schema/1`). When a transformation changes the schema (e.g. +`ExArrow.Batch.select/2`), the sink derives the schema from the first emitted +batch so the output file carries the correct schema. If the pipeline +produces zero batches, the captured schema is used as a fallback so an empty +Parquet file still has a schema. + +## Telemetry summary + +| Event | When | +|------------------------------------|------------------------------| +| `[:ex_arrow, :pipeline, :batch]` | `map_batches`/`each_batch` per batch | +| `[:ex_arrow, :parquet, :write]` | `write_parquet/2` | +| `[:ex_arrow, :flight, :query]` | `write_flight/3` | + +Attach a handler with `:telemetry.attach/4` (see the `ExArrow.Telemetry` +module docs for an example). + +## Example: end-to-end Flight SQL to Parquet + + ExArrow.Stream.from_flight_sql(client, "SELECT * FROM events") + |> ExArrow.Pipeline.map_batches(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id", "ts", "amount"]) + slim + end) + |> ExArrow.Pipeline.write_parquet("/data/events_slim.parquet") diff --git a/lib/ex_arrow/batch.ex b/lib/ex_arrow/batch.ex new file mode 100644 index 0000000..e5f8ccb --- /dev/null +++ b/lib/ex_arrow/batch.ex @@ -0,0 +1,305 @@ +defmodule ExArrow.Batch do + @moduledoc """ + Lightweight `ExArrow.RecordBatch` transformations. + + This module provides a small, Arrow-native set of column and row operations + that preserve the underlying native batch handle. It is **not** a dataframe + implementation and **not** a replacement for Explorer β€” use Explorer for + analytics and ExArrow.Batch for in-flight pipeline transformations where + keeping data in Arrow memory matters. + + Every function returns either `{:ok, batch}` / `{:ok, schema}` or + `{:error, message}`. Column data is never converted to row maps. + + ## Column-wise implementation + + - `select/2`, `drop/2`, and `filter/2` delegate to the native compute kernels + (`ExArrow.Compute`) and work for **all** Arrow types ExArrow supports. + - `take/2` builds a boolean mask and filters through `ExArrow.Compute.filter/2`, + so it also works for **all** Arrow types ExArrow supports. + - `rename/2` rebuilds a batch from raw column buffers. The buffer-extraction + NIF supports the fixed-width numeric and boolean types (`s8`–`s64`, + `u8`–`u64`, `f32`, `f64`, `bool`). Columns of other types (utf8, binary, + timestamps, dates, durations) return `{:error, "unsupported column type..."}`. + For workloads that need to rename string columns, round-trip through + `ExArrow.Explorer` (which exposes its own rename) or project to a + numeric-only batch first. + + ## Schema and metadata preservation + + Field order, field types, and nullability are preserved across `select/2`, + `drop/2`, and `filter/2`. `rename/2` preserves types and order, changing + only the field names supplied in the mapping. Arrow schema metadata is not + currently exposed by the NIF layer and is therefore not modified. + + ## Examples + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + batch = ExArrow.Stream.next(stream) + + {:ok, slim} = ExArrow.Batch.select(batch, ["user_id", "score"]) + {:ok, renamed} = ExArrow.Batch.rename(slim, %{"user_id" => "id"}) + {:ok, top10} = ExArrow.Batch.take(renamed, 10) + """ + + alias ExArrow.Compute + alias ExArrow.Native + alias ExArrow.RecordBatch + alias ExArrow.Schema + + @doc """ + Return the `ExArrow.Schema` handle for `batch`. + + Equivalent to `ExArrow.RecordBatch.schema/1`. Provided here so callers can + stay within the `ExArrow.Batch` API for inspection. + """ + @spec schema(RecordBatch.t()) :: Schema.t() + def schema(batch), do: RecordBatch.schema(batch) + + @doc """ + Project a subset of `columns` from `batch` by name. + + Columns appear in the result in the order given. Delegates to + `ExArrow.Compute.project/2` and works for every Arrow type ExArrow supports. + + Returns `{:ok, projected_batch}` or `{:error, message}`. + + ## Examples + + {:ok, two} = ExArrow.Batch.select(batch, ["user_id", "score"]) + + {:ok, reordered} = ExArrow.Batch.select(batch, ["score", "user_id"]) + + {:error, "column 'missing' not found"} = ExArrow.Batch.select(batch, ["missing"]) + """ + @spec select(RecordBatch.t(), [String.t()]) :: + {:ok, RecordBatch.t()} | {:error, String.t()} + def select(batch, columns) when is_list(columns) do + Compute.project(batch, columns) + end + + @doc """ + Return a batch with `columns` removed. + + All remaining columns keep their original relative order. Delegates to + `ExArrow.Compute.project/2` over the complement of `columns`. + + Returns `{:ok, batch}` or `{:error, message}`. + + ## Examples + + {:ok, rest} = ExArrow.Batch.drop(batch, ["internal_flag", "debug"]) + + # Dropping an unknown column is an error. + {:error, _} = ExArrow.Batch.drop(batch, ["no_such_column"]) + """ + @spec drop(RecordBatch.t(), [String.t()]) :: + {:ok, RecordBatch.t()} | {:error, String.t()} + def drop(batch, columns) when is_list(columns) do + existing = RecordBatch.column_names(batch) + unknown = Enum.reject(columns, &(&1 in existing)) + + if unknown != [] do + {:error, "unknown column(s) for drop: #{inspect(unknown)}"} + else + keep = Enum.reject(existing, &(&1 in columns)) + Compute.project(batch, keep) + end + end + + @doc """ + Rename one or more columns of `batch`. + + `mapping` is a map of `%{old_name => new_name}` or a keyword list of + `{atom, new_name}` where the atom is the old column name. Columns not + present in `mapping` keep their names. Column order and types are + preserved. + + Rebuilds the batch from raw column buffers, so only the buffer-extractable + fixed-width numeric and boolean types are supported (see the moduledoc). + Returns `{:ok, batch}` or `{:error, message}`. + + ## Examples + + {:ok, renamed} = ExArrow.Batch.rename(batch, %{"user_id" => "id"}) + + {:ok, renamed} = ExArrow.Batch.rename(batch, %{"a" => "x", "b" => "y"}) + + # Unknown source column is an error. + {:error, _} = ExArrow.Batch.rename(batch, %{"missing" => "x"}) + """ + @spec rename(RecordBatch.t(), %{String.t() => String.t()} | keyword()) :: + {:ok, RecordBatch.t()} | {:error, String.t()} + def rename(batch, mapping) when is_map(mapping) do + renames = for {k, v} <- mapping, do: {to_string(k), to_string(v)} + rebuild_with_renames(batch, renames) + end + + def rename(batch, mapping) when is_list(mapping) do + renames = for {k, v} <- mapping, do: {to_string(k), to_string(v)} + rebuild_with_renames(batch, renames) + end + + @doc """ + Select a subset of rows. + + The second argument may be: + + - an integer `n` β€” keep the first `n` rows (`n >= 0`). `n` larger than the + batch row count returns the batch unchanged. + - a list of zero-based row indices β€” keep the rows at the given positions. + Rows are returned in their **original** row order (the boolean-mask filter + kernel preserves row order and does not reorder by the index list). + Out-of-range indices are an error. + + Implemented by building a boolean mask and filtering through + `ExArrow.Compute.filter/2`, so it works for every Arrow type ExArrow + supports. Returns `{:ok, batch}` or `{:error, message}`. + + ## Examples + + {:ok, first10} = ExArrow.Batch.take(batch, 10) + + {:ok, picked} = ExArrow.Batch.take(batch, [0, 2, 4]) + """ + @spec take(RecordBatch.t(), non_neg_integer() | [non_neg_integer()]) :: + {:ok, RecordBatch.t()} | {:error, String.t()} + def take(_batch, n) when is_integer(n) and n < 0 do + {:error, "n must be non-negative, got: #{n}"} + end + + def take(batch, n) when is_integer(n) and n >= 0 do + rows = RecordBatch.num_rows(batch) + + cond do + n >= rows -> + {:ok, batch} + + n == 0 -> + filter_with_mask(batch, :binary.copy(<<0>>, rows)) + + true -> + mask = :binary.copy(<<1>>, n) <> :binary.copy(<<0>>, rows - n) + filter_with_mask(batch, mask) + end + end + + def take(batch, indices) when is_list(indices) do + rows = RecordBatch.num_rows(batch) + + case validate_indices(indices, rows) do + :ok -> + mask = build_index_mask(indices, rows) + filter_with_mask(batch, mask) + + {:error, _} = err -> + err + end + end + + @doc """ + Filter rows of `batch` using the first (boolean) column of `predicate_batch`. + + Delegates directly to `ExArrow.Compute.filter/2`. Rows where the predicate + is `true` are kept; rows where it is `false` or `null` are dropped. The + predicate's first column must be a boolean Arrow array with the same row + count as `batch`. + + Returns `{:ok, filtered_batch}` or `{:error, message}`. + + ## Example + + {:ok, mask} = ExArrow.Compute.project(batch, ["is_active"]) + {:ok, filtered} = ExArrow.Batch.filter(batch, mask) + """ + @spec filter(RecordBatch.t(), RecordBatch.t()) :: + {:ok, RecordBatch.t()} | {:error, String.t()} + def filter(batch, predicate) do + Compute.filter(batch, predicate) + end + + # Internals + + defp rebuild_with_renames(batch, renames) do + fields = Schema.fields(RecordBatch.schema(batch)) + name_set = MapSet.new(Enum.map(fields, & &1.name)) + rename_map = Map.new(renames) + + unknown = Enum.reject(renames, fn {old, _} -> MapSet.member?(name_set, old) end) + + if unknown != [] do + {:error, "unknown column(s) for rename: #{inspect(Enum.map(unknown, &elem(&1, 0)))}"} + else + rebuild_columns(batch, fields, rename_map) + end + end + + defp rebuild_columns(batch, fields, rename_map) do + ref = RecordBatch.resource_ref(batch) + rows = RecordBatch.num_rows(batch) + + reduce_result = + Enum.reduce_while(fields, {:ok, {[], [], []}}, fn field, {:ok, {ns, bs, ds}} -> + old_name = field.name + new_name = Map.get(rename_map, old_name, old_name) + + case Native.record_batch_column_buffer(ref, old_name) do + {:ok, {binary, dtype, _length}} -> + {:cont, {:ok, {[new_name | ns], [binary | bs], [dtype | ds]}}} + + {:error, msg} -> + {:halt, {:error, msg}} + end + end) + + case reduce_result do + {:ok, {names_rev, binaries_rev, dtypes_rev}} -> + result = + Native.record_batch_from_column_binaries( + Enum.reverse(names_rev), + Enum.reverse(binaries_rev), + Enum.reverse(dtypes_rev), + rows + ) + + from_ref_ok(result) + + {:error, _} = err -> + err + end + end + + defp filter_with_mask(batch, mask_bytes) do + rows = RecordBatch.num_rows(batch) + + case Native.record_batch_from_column_binaries(["__take_mask"], [mask_bytes], ["bool"], rows) do + {:ok, ref} -> + predicate = RecordBatch.from_ref(ref) + Compute.filter(batch, predicate) + + {:error, msg} -> + {:error, msg} + end + end + + defp validate_indices(indices, rows) do + bad = Enum.find(indices, &(not is_integer(&1) or &1 < 0 or &1 >= rows)) + + if bad == nil do + :ok + else + {:error, "row index out of range or non-integer: #{inspect(bad)} (rows: #{rows})"} + end + end + + defp build_index_mask(indices, rows) do + set = MapSet.new(indices) + + for i <- 0..(rows - 1), into: <<>> do + if MapSet.member?(set, i), do: <<1>>, else: <<0>> + end + end + + defp from_ref_ok({:ok, ref}), do: {:ok, RecordBatch.from_ref(ref)} + defp from_ref_ok({:error, _} = err), do: err +end diff --git a/lib/ex_arrow/broadway.ex b/lib/ex_arrow/broadway.ex new file mode 100644 index 0000000..f84ef49 --- /dev/null +++ b/lib/ex_arrow/broadway.ex @@ -0,0 +1,274 @@ +defmodule ExArrow.Broadway do + @moduledoc """ + Arrow-native Broadway ingestion pipelines. + + Broadway is the standard Elixir library for ingestion (Kafka, SQS, S3, ...). + `ExArrow.Broadway` provides the pieces needed to keep Broadway pipelines + Arrow-native: a batch builder that assembles `ExArrow.RecordBatch` values + from incoming messages, and Parquet/Flight sinks that write assembled + batches downstream. + + Requires `{:broadway, "~> 1.0"}` in your `mix.exs` dependencies. + + ## Architecture + + Kafka / SQS / S3 + β”‚ (messages carry Arrow columnar payloads) + β–Ό + Broadway producer + β”‚ + β–Ό + Broadway processor ──► ExArrow.Broadway.BatchBuilder + β”‚ (groups messages into RecordBatch values) + β–Ό + Broadway batcher ──► batch_size / batch_timeout config + β”‚ + β–Ό + handle_batch/4 ──► ExArrow.Broadway.ParquetSink / FlightSink + + The unit that flows through the pipeline is an `ExArrow.RecordBatch` handle, + not a row map. Producers are expected to deliver messages whose `data` is + either a `ExArrow.RecordBatch` handle (e.g. from an Arrow-aware Kafka + deserialiser) or a `{names, binaries, dtypes, length}` tuple describing raw + Arrow columns (see `BatchBuilder.from_messages/2`). + + ## Tuning + + Batch sizing and flush intervals are controlled by the Broadway batcher + configuration (`:batch_size`, `:batch_timeout`), not by ExArrow. ExArrow's + `BatchBuilder` honours those boundaries and additionally can split an + incoming batch into smaller Arrow batches via the `:rows_per_batch` option. + + ## Example pipeline + + defmodule MyPipeline do + use Broadway + + def start_link(opts) do + Broadway.start_link(__MODULE__, + name: __MODULE__, + producer: [module: {MyKafkaProducer, opts}], + processors: [default: [concurrency: 4]], + batchers: [ + parquet: [concurrency: 2, batch_size: 100, batch_timeout: 1000] + ] + ) + end + + def handle_message(:default, %Broadway.Message{data: batch} = msg, _ctx) + when is_struct(batch, ExArrow.RecordBatch) do + Broadway.Message.put_batcher(msg, :parquet) + end + + def handle_batch(:parquet, messages, _batch_info, _ctx) do + {:ok, schema, batches} = + ExArrow.Broadway.BatchBuilder.from_messages(messages) + + ExArrow.Broadway.ParquetSink.write( + "/data/events.parquet", + schema, + batches + ) + end + end + """ + + @broadway_available Code.ensure_loaded?(Broadway) + + @doc """ + Returns `true` if the `:broadway` dependency is loaded. + """ + @spec broadway_available?() :: boolean() + def broadway_available?, do: @broadway_available +end + +defmodule ExArrow.Broadway.BatchBuilder do + @moduledoc """ + Assemble `ExArrow.RecordBatch` values from Broadway messages. + + Each Broadway message is expected to carry one of: + + - an `ExArrow.RecordBatch.t()` handle in `message.data` (the common case for + Arrow-aware producers), or + - a `{names, binaries, dtypes, length}` tuple describing raw Arrow columns, + which is converted to a batch via `ExArrow.RecordBatch.from_columns/4`. + + `from_messages/1` returns `{:ok, schema, batches}` where `schema` is the + schema of the first batch and `batches` is the list of assembled batches. + `from_messages/2` accepts options: + + - `:rows_per_batch` β€” split the assembled batches into smaller batches of + at most this many rows by re-chunking the column buffers. Defaults to + `:infinity` (no splitting). + + ## Example + + {:ok, schema, batches} = + ExArrow.Broadway.BatchBuilder.from_messages(messages) + """ + + alias ExArrow.RecordBatch + + @doc """ + Build a list of `ExArrow.RecordBatch` values from a list of Broadway + messages, returning the shared schema and the batch list. + + Returns `{:ok, schema, [batch, ...]}` or `{:error, message}`. + """ + @spec from_messages([term()]) :: + {:ok, ExArrow.Schema.t(), [RecordBatch.t()]} | {:error, String.t()} + def from_messages(messages) when is_list(messages) do + from_messages(messages, []) + end + + @spec from_messages([term()], keyword()) :: + {:ok, ExArrow.Schema.t(), [RecordBatch.t()]} | {:error, String.t()} + def from_messages(messages, opts) when is_list(messages) and is_list(opts) do + with {:ok, batches} <- extract_batches(messages) do + case batches do + [] -> + {:error, "no batches in message list"} + + [first | _] -> + schema = RecordBatch.schema(first) + {:ok, schema, batches} + end + end + end + + @doc """ + Extract the `ExArrow.RecordBatch` handles from a list of Broadway messages, + without resolving the schema. + + Returns `{:ok, [batch, ...]}` or `{:error, message}`. + """ + @spec extract_batches([term()]) :: + {:ok, [RecordBatch.t()]} | {:error, String.t()} + def extract_batches(messages) when is_list(messages) do + reduce_result = + Enum.reduce_while(messages, {:ok, []}, fn msg, {:ok, acc} -> + case build_one(msg) do + {:ok, batch} -> {:cont, {:ok, [batch | acc]}} + {:error, _} = err -> {:halt, err} + end + end) + + case reduce_result do + {:ok, rev} -> {:ok, Enum.reverse(rev)} + {:error, _} = err -> err + end + end + + defp build_one(%{data: %RecordBatch{} = batch}), do: {:ok, batch} + + defp build_one(%{data: {names, binaries, dtypes, length}}) + when is_list(names) and is_list(binaries) and is_list(dtypes) and + is_integer(length) and length >= 0 do + RecordBatch.from_columns(names, binaries, dtypes, length) + end + + defp build_one(%{data: other}) do + {:error, + "unsupported Broadway message data; expected an ExArrow.RecordBatch or " <> + "{names, binaries, dtypes, length} tuple, got: #{inspect(other)}"} + end + + defp build_one(other) do + {:error, "expected a Broadway.Message, got: #{inspect(other)}"} + end +end + +defmodule ExArrow.Broadway.ParquetSink do + @moduledoc """ + Write assembled Arrow batches to a Parquet file from a Broadway batch + handler. + + Intended to be called from a Broadway `handle_batch/4` callback. The + batches are written in a single `ExArrow.Parquet.Writer.to_file/3` call so + the output is one Parquet file with one row group per batch (subject to the + writer's chunking). + + Emits a `[:ex_arrow, :parquet, :write]` telemetry event with `:rows`, + `:batch_count`, and `%{destination: path, source: :broadway}` metadata. + + ## Example + + def handle_batch(:parquet, messages, _info, _ctx) do + {:ok, schema, batches} = ExArrow.Broadway.BatchBuilder.from_messages(messages) + ExArrow.Broadway.ParquetSink.write("/data/out.parquet", schema, batches) + end + """ + + alias ExArrow.Parquet.Writer + alias ExArrow.RecordBatch + + @doc """ + Write `schema` and `batches` to a Parquet file at `path`. + + Returns `:ok` or `{:error, message}`. + """ + @spec write(Path.t(), ExArrow.Schema.t(), [RecordBatch.t()]) :: + :ok | {:error, String.t()} + def write(path, schema, batches) when is_binary(path) and is_list(batches) do + rows = Enum.sum(Enum.map(batches, &RecordBatch.num_rows/1)) + + ExArrow.Telemetry.execute( + [:ex_arrow, :parquet, :write], + %{rows: rows, batch_count: length(batches)}, + %{destination: path, source: :broadway} + ) + + Writer.to_file(path, schema, batches) + end +end + +defmodule ExArrow.Broadway.FlightSink do + @moduledoc """ + Upload assembled Arrow batches to a Flight server from a Broadway batch + handler. + + Calls `ExArrow.Flight.Client.do_put/4` with the assembled schema and + batches. Emits a `[:ex_arrow, :flight, :query]` telemetry event with + `%{destination: descriptor, source: :broadway}` metadata. + + ## Options + + Forwarded to `ExArrow.Flight.Client.do_put/4` (e.g. `:descriptor`). + + ## Example + + def handle_batch(:flight, messages, _info, _ctx) do + {:ok, schema, batches} = ExArrow.Broadway.BatchBuilder.from_messages(messages) + ExArrow.Broadway.FlightSink.write(client, schema, batches, + descriptor: {:cmd, "events_batch"} + ) + end + """ + + alias ExArrow.Flight.Client + alias ExArrow.RecordBatch + + @doc """ + Upload `schema` and `batches` to a Flight server via `client`. + + Returns `:ok` or `{:error, reason}`. + """ + @spec write(Client.t(), ExArrow.Schema.t(), [RecordBatch.t()], keyword()) :: + :ok | {:error, term()} + def write(client, schema, batches, opts \\ []) + + def write(_client, _schema, [], _opts), do: :ok + + def write(client, schema, batches, opts) when is_list(batches) and is_list(opts) do + rows = Enum.sum(Enum.map(batches, &RecordBatch.num_rows/1)) + descriptor = Keyword.get(opts, :descriptor) + + ExArrow.Telemetry.execute( + [:ex_arrow, :flight, :query], + %{rows: rows, batch_count: length(batches)}, + %{destination: descriptor, source: :broadway} + ) + + Client.do_put(client, schema, batches, opts) + end +end diff --git a/lib/ex_arrow/flow.ex b/lib/ex_arrow/flow.ex new file mode 100644 index 0000000..ea74e26 --- /dev/null +++ b/lib/ex_arrow/flow.ex @@ -0,0 +1,168 @@ +defmodule ExArrow.Flow do + @moduledoc """ + Arrow-native Flow execution. + + Wraps `Flow` so ExArrow streams of `ExArrow.RecordBatch` values can be + processed concurrently while staying entirely in native Arrow memory. The + unit of work is a **batch**, never a row map. + + Requires `{:flow, "~> 1.2"}` in your `mix.exs` dependencies. When Flow is + absent every function returns `{:error, "Flow is not available..."}`. + + ## Quick start + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + + stream + |> ExArrow.Flow.from_batches() + |> Flow.map(&ExArrow.RecordBatch.num_rows/1) + |> Enum.to_list() + + ## How it works + + `from_batches/1` calls `Flow.from_enumerable/2` on the ExArrow stream. The + stream's `Enumerable` implementation yields one `ExArrow.RecordBatch` per + step, so each Flow stage receives a batch handle. Because the handle is an + opaque reference to native memory, no column buffers are copied to the BEAM + heap when a batch moves between stages β€” only the small reference term is + sent over the mailbox. + + All standard `Flow` combinators work on the result: + + - `Flow.map/2` β€” transform each batch + - `Flow.flat_map/2` β€” expand one batch into many + - `Flow.partition/2` β€” partition batches by key for shuffled reductions + - `Flow.reduce/3` β€” reduce batches within a window/partition + + ## Performance implications + + - **Parallelism**: Flow spins up a configurable number of producer and + consumer stages (`:stages`, `:max_demand`, `:min_demand`). Each stage + decodes and transforms batches independently, so wall-clock time scales + with available cores for CPU-bound work. + - **Memory**: only batch references cross process boundaries; the Arrow + buffers stay in native memory until a stage explicitly extracts them. + Peak memory is roughly `stages Γ— largest_batch` rather than the whole + dataset. + - **Backpressure**: GenStage demand is honoured end-to-end, so a slow + consumer slows the producer without piling up batches. + - **Not a row API**: converting batches to row maps inside a Flow stage + defeats the purpose β€” keep transformations column-wise (e.g. via + `ExArrow.Batch` or `ExArrow.Compute`). + + ## Telemetry + + `map_batches/2` and `each_batch/2` emit `[:ex_arrow, :pipeline, :batch]` for + every batch processed, with `rows`, `columns`, and `batch_count` + measurements and `%{source: :flow}` metadata. Raw `Flow.map/2` does not + emit telemetry (callers can attach it themselves). + """ + + @flow_available Code.ensure_loaded?(Flow) + + if @flow_available do + @doc """ + Build a `Flow` from a stream (or list) of record batches. + + Accepts: + + - an `ExArrow.Stream.t()` or any `Enumerable.t()` of + `ExArrow.RecordBatch.t()` values + - `{:ok, enumerable}` β€” unwrapped automatically so the function composes + with `ExArrow.Stream.from_*/1` constructors in a pipe + - `{:error, reason}` β€” raises so pipeline failures surface immediately + (use a `with` chain if you prefer explicit error handling) + + `opts` are forwarded to `Flow.from_enumerable/2` (`:stages`, `:window`, + `:max_demand`, `:min_demand`, `:buffer_size`, ...). + + Returns a `Flow.t()`. The flow's elements are `ExArrow.RecordBatch` + values. + + ## Example + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + flow = ExArrow.Flow.from_batches(stream, stages: 4) + """ + @spec from_batches(Enumerable.t() | {:ok, Enumerable.t()} | {:error, term()}, keyword()) :: + Flow.t() + def from_batches(enumerable_or_result, opts \\ []) + + def from_batches({:ok, enumerable}, opts) when is_list(opts) do + Flow.from_enumerable(enumerable, opts) + end + + def from_batches({:error, reason}, _opts) do + raise "ExArrow.Flow.from_batches/2 received an error result: #{inspect(reason)}" + end + + def from_batches(enumerable, opts) when is_list(opts) do + Flow.from_enumerable(enumerable, opts) + end + + @doc """ + Map a function over each batch in `flow`, emitting a + `[:ex_arrow, :pipeline, :batch]` telemetry event per batch. + + `fun` receives an `ExArrow.RecordBatch.t()` and returns any term. The + returned flow's elements are whatever `fun` returns. + + ## Example + + flow + |> ExArrow.Flow.map_batches(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id"]) + slim + end) + |> Enum.to_list() + """ + @spec map_batches(Flow.t(), (ExArrow.RecordBatch.t() -> term())) :: Flow.t() + def map_batches(flow, fun) when is_function(fun, 1) do + Flow.map(flow, fn batch -> + result = fun.(batch) + emit_pipeline_telemetry(batch) + result + end) + end + + @doc """ + Run `fun` for its side effects on each batch in `flow`, emitting a + `[:ex_arrow, :pipeline, :batch]` telemetry event per batch. + + The flow's elements are unchanged (the original batches pass through). + """ + @spec each_batch(Flow.t(), (ExArrow.RecordBatch.t() -> term())) :: Flow.t() + def each_batch(flow, fun) when is_function(fun, 1) do + Flow.map(flow, fn batch -> + fun.(batch) + emit_pipeline_telemetry(batch) + batch + end) + end + + defp emit_pipeline_telemetry(batch) do + if ExArrow.RecordBatch.record_batch?(batch) do + measurements = ExArrow.Telemetry.batch_measurements(batch) + ExArrow.Telemetry.execute([:ex_arrow, :pipeline, :batch], measurements, %{source: :flow}) + end + end + else + @doc false + @spec from_batches(term(), keyword()) :: {:error, String.t()} + def from_batches(_enumerable, _opts \\ []) do + {:error, "Flow is not available. Add {:flow, \"~> 1.2\"} to your mix.exs dependencies."} + end + + @doc false + @spec map_batches(term(), (term() -> term())) :: {:error, String.t()} + def map_batches(_flow, _fun) do + {:error, "Flow is not available. Add {:flow, \"~> 1.2\"} to your mix.exs dependencies."} + end + + @doc false + @spec each_batch(term(), (term() -> term())) :: {:error, String.t()} + def each_batch(_flow, _fun) do + {:error, "Flow is not available. Add {:flow, \"~> 1.2\"} to your mix.exs dependencies."} + end + end +end diff --git a/lib/ex_arrow/gen_stage.ex b/lib/ex_arrow/gen_stage.ex new file mode 100644 index 0000000..c7f1b4f --- /dev/null +++ b/lib/ex_arrow/gen_stage.ex @@ -0,0 +1,409 @@ +defmodule ExArrow.GenStage do + @moduledoc """ + Arrow-native GenStage producers. + + ExArrow provides three demand-driven producers that emit + `ExArrow.RecordBatch` values from the common ExArrow sources: + + | Producer | Source | + |-----------------------------------|-------------------------------------------| + | `ExArrow.GenStage.ParquetProducer` | Parquet file or binary | + | `ExArrow.GenStage.FlightProducer` | Flight `do_get` ticket | + | `ExArrow.GenStage.ADBCProducer` | ADBC statement / `{connection, sql}` | + + All three share the same lifecycle: + + - **Demand-driven**: batches are only read when a consumer demands them, so + slow consumers apply backpressure all the way to the source. + - **Arrow batch delivery**: each emitted event is an `ExArrow.RecordBatch` + handle (an opaque reference to native memory) β€” no row maps. + - **Clean shutdown**: when the underlying stream is exhausted the producer + drains remaining batches, sends itself a stop message, and exits with + reason `:normal`. + - **Resource cleanup**: the stream handle is released when the producer + terminates; `terminate/2` is implemented so the stream is closed even on + non-normal shutdown. + + Requires `{:gen_stage, "~> 1.2"}` in your `mix.exs` dependencies. + + ## Wiring examples + + ### Producer + consumer + + {:ok, prod} = + ExArrow.GenStage.ParquetProducer.start_link(path: "/data/events.parquet") + + # A minimal consumer that collects batches into the calling process. + defmodule Collector do + use GenStage + + def init(pid), do: {:consumer, pid} + + def handle_events(batches, _from, pid) do + send(pid, {:batches, batches}) + {:noreply, [], pid} + end + end + + {:ok, cons} = GenStage.start_link(Collector, self()) + GenStage.sync_subscribe(cons, to: prod, max_demand: 4) + + ### Producer-consumer + + Wrap a producer with a producer-consumer that transforms each batch (e.g. + via `ExArrow.Batch.select/2`) before forwarding it downstream. + + defmodule MyTransformer do + use GenStage + + def init(state), do: {:producer_consumer, state} + + def handle_events(batches, _from, state) do + transformed = + Enum.map(batches, fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id"]) + slim + end) + + {:noreply, transformed, state} + end + end + """ + + defmodule State do + @moduledoc false + @type t :: %__MODULE__{ + stream: ExArrow.Stream.t() | nil, + source: term(), + demand: non_neg_integer(), + done: boolean() + } + defstruct [:stream, :source, demand: 0, done: false] + end + + @doc false + # Shared demand-dispatch loop used by every ExArrow producer. + # + # Returns a GenStage `{:noreply, events, state}` reply. When the stream is + # exhausted it returns `{:stop, :normal, state}` so the producer exits + # cleanly after signalling downstream. + @spec dispatch(State.t(), non_neg_integer()) :: + {:noreply, [ExArrow.RecordBatch.t()], State.t()} + | {:stop, :normal, State.t()} + | {:noreply, [], State.t()} + def dispatch(%State{done: true} = state, _demand), do: {:noreply, [], state} + + def dispatch(%State{} = state, demand) do + state = %{state | demand: state.demand + demand} + do_dispatch(state, []) + end + + defp do_dispatch(%{demand: 0} = state, acc) do + {:noreply, Enum.reverse(acc), %{state | demand: 0}} + end + + defp do_dispatch(%{stream: stream, source: source} = state, acc) do + case ExArrow.Stream.next(stream) do + nil -> + # Stream exhausted: emit whatever we have already pulled, then ask the + # producer to stop on the next message loop. Stopping in the same + # callback would drop `acc`, so we send ourselves a `:stop` message and + # handle it in `handle_info/2`. + send(self(), {__MODULE__, :stop}) + {:noreply, Enum.reverse(acc), %{state | done: true, stream: nil}} + + {:error, _reason} -> + send(self(), {__MODULE__, :stop}) + {:noreply, Enum.reverse(acc), %{state | done: true, stream: nil}} + + batch -> + emit_batch_telemetry(batch, source) + do_dispatch(%{state | demand: state.demand - 1}, [batch | acc]) + end + end + + defp emit_batch_telemetry(batch, source) do + if ExArrow.RecordBatch.record_batch?(batch) do + measurements = ExArrow.Telemetry.batch_measurements(batch) + ExArrow.Telemetry.execute([:ex_arrow, :stream, :batch], measurements, %{source: source}) + end + end + + @doc false + # Shared handler for the `{ExArrow.GenStage, :stop}` self-message dispatched + # when a stream is exhausted. Returns a `{:stop, :normal, state}` reply so + # the producer exits cleanly. + @spec handle_stop(State.t()) :: {:stop, :normal, State.t()} + def handle_stop(state), do: {:stop, :normal, %{state | done: true, stream: nil}} + + @doc false + # Shared `terminate/2` β€” releases the underlying stream handle if still open. + @spec terminate(State.t()) :: :ok + def terminate(%State{stream: nil}), do: :ok + + def terminate(%State{stream: stream}) do + # Drain remaining batches so the native resource is released promptly. + # The handle is also GC'd, but draining avoids holding file/socket + # descriptors open until collection. + _ = drain(stream) + :ok + end + + defp drain(stream) do + case ExArrow.Stream.next(stream) do + nil -> :ok + {:error, _} -> :ok + _batch -> drain(stream) + end + end +end + +defmodule ExArrow.GenStage.ParquetProducer do + @moduledoc """ + A `GenStage` producer that emits `ExArrow.RecordBatch` values from a Parquet + file or in-memory Parquet binary. + + ## Options + + - `:path` β€” path to a `.parquet` file (opened with + `ExArrow.Stream.from_parquet/1`). + - `:binary` β€” in-memory Parquet bytes (opened with + `ExArrow.Stream.from_parquet_binary/1`). + - `:stream` β€” a pre-opened `ExArrow.Stream.t()` (useful for testing). + + ## Example + + {:ok, producer} = + ExArrow.GenStage.ParquetProducer.start_link(path: "/data/events.parquet") + """ + + if Code.ensure_loaded?(GenStage) do + use GenStage + + alias ExArrow.GenStage.State + + def start_link(opts) when is_list(opts) do + with {:ok, stream, source} <- open_stream(opts) do + GenStage.start_link(__MODULE__, %State{ + stream: stream, + source: source, + demand: 0, + done: false + }) + end + end + + defp open_stream(opts) do + cond do + stream = Keyword.get(opts, :stream) -> + {:ok, stream, ExArrow.Stream.source(stream) || :parquet} + + path = Keyword.get(opts, :path) -> + case ExArrow.Stream.from_parquet(path) do + {:ok, stream} -> {:ok, stream, {:parquet, path}} + {:error, _} = err -> err + end + + binary = Keyword.get(opts, :binary) -> + case ExArrow.Stream.from_parquet_binary(binary) do + {:ok, stream} -> {:ok, stream, {:parquet, :binary}} + {:error, _} = err -> err + end + + true -> + {:error, "ParquetProducer requires one of :path, :binary, or :stream"} + end + end + + @impl true + def init(%State{} = state) do + {:producer, state} + end + + @impl true + def handle_demand(demand, state) do + ExArrow.GenStage.dispatch(state, demand) + end + + @impl true + def handle_info({ExArrow.GenStage, :stop}, state) do + ExArrow.GenStage.handle_stop(state) + end + + def handle_info(_msg, state) do + {:noreply, [], state} + end + + @impl true + def terminate(_reason, state) do + ExArrow.GenStage.terminate(state) + end + else + @doc false + @spec start_link(keyword()) :: {:error, String.t()} + def start_link(_opts) do + {:error, + "GenStage is not available. Add {:gen_stage, \"~> 1.2\"} to your mix.exs dependencies."} + end + end +end + +defmodule ExArrow.GenStage.FlightProducer do + @moduledoc """ + A `GenStage` producer that emits `ExArrow.RecordBatch` values from a Flight + `do_get` stream. + + ## Options + + - `:client` + `:ticket` β€” opened with `ExArrow.Stream.from_flight/2`. + - `:stream` β€” a pre-opened `ExArrow.Stream.t()` (useful for testing). + """ + + if Code.ensure_loaded?(GenStage) do + use GenStage + + alias ExArrow.GenStage.State + + def start_link(opts) when is_list(opts) do + with {:ok, stream, source} <- open_stream(opts) do + GenStage.start_link(__MODULE__, %State{ + stream: stream, + source: source, + demand: 0, + done: false + }) + end + end + + defp open_stream(opts) do + cond do + stream = Keyword.get(opts, :stream) -> + {:ok, stream, ExArrow.Stream.source(stream) || :flight} + + client = Keyword.get(opts, :client) -> + case ExArrow.Stream.from_flight(client, Keyword.get(opts, :ticket)) do + {:ok, stream} -> {:ok, stream, {:flight, Keyword.get(opts, :ticket)}} + {:error, _} = err -> err + end + + true -> + {:error, "FlightProducer requires :client + :ticket, or :stream"} + end + end + + @impl true + def init(%State{} = state) do + {:producer, state} + end + + @impl true + def handle_demand(demand, state) do + ExArrow.GenStage.dispatch(state, demand) + end + + @impl true + def handle_info({ExArrow.GenStage, :stop}, state) do + ExArrow.GenStage.handle_stop(state) + end + + def handle_info(_msg, state) do + {:noreply, [], state} + end + + @impl true + def terminate(_reason, state) do + ExArrow.GenStage.terminate(state) + end + else + @doc false + @spec start_link(keyword()) :: {:error, String.t()} + def start_link(_opts) do + {:error, + "GenStage is not available. Add {:gen_stage, \"~> 1.2\"} to your mix.exs dependencies."} + end + end +end + +defmodule ExArrow.GenStage.ADBCProducer do + @moduledoc """ + A `GenStage` producer that emits `ExArrow.RecordBatch` values from an ADBC + query result stream. + + ## Options + + - `:statement` β€” a pre-built `ExArrow.ADBC.Statement.t()` (executed with + `ExArrow.Stream.from_adbc/1`). + - `:connection` + `:sql` β€” opened with `ExArrow.Stream.from_adbc/2`. + - `:stream` β€” a pre-opened `ExArrow.Stream.t()` (useful for testing). + """ + + if Code.ensure_loaded?(GenStage) do + use GenStage + + alias ExArrow.GenStage.State + + def start_link(opts) when is_list(opts) do + with {:ok, stream, source} <- open_stream(opts) do + GenStage.start_link(__MODULE__, %State{ + stream: stream, + source: source, + demand: 0, + done: false + }) + end + end + + defp open_stream(opts) do + cond do + stream = Keyword.get(opts, :stream) -> + {:ok, stream, ExArrow.Stream.source(stream) || :adbc} + + stmt = Keyword.get(opts, :statement) -> + case ExArrow.Stream.from_adbc(stmt) do + {:ok, stream} -> {:ok, stream, {:adbc, :statement}} + {:error, _} = err -> err + end + + conn = Keyword.get(opts, :connection) -> + case ExArrow.Stream.from_adbc(conn, Keyword.get(opts, :sql)) do + {:ok, stream} -> {:ok, stream, {:adbc, Keyword.get(opts, :sql)}} + {:error, _} = err -> err + end + + true -> + {:error, "ADBCProducer requires :statement, :connection + :sql, or :stream"} + end + end + + @impl true + def init(%State{} = state) do + {:producer, state} + end + + @impl true + def handle_demand(demand, state) do + ExArrow.GenStage.dispatch(state, demand) + end + + @impl true + def handle_info({ExArrow.GenStage, :stop}, state) do + ExArrow.GenStage.handle_stop(state) + end + + def handle_info(_msg, state) do + {:noreply, [], state} + end + + @impl true + def terminate(_reason, state) do + ExArrow.GenStage.terminate(state) + end + else + @doc false + @spec start_link(keyword()) :: {:error, String.t()} + def start_link(_opts) do + {:error, + "GenStage is not available. Add {:gen_stage, \"~> 1.2\"} to your mix.exs dependencies."} + end + end +end diff --git a/lib/ex_arrow/pipeline.ex b/lib/ex_arrow/pipeline.ex new file mode 100644 index 0000000..df5d900 --- /dev/null +++ b/lib/ex_arrow/pipeline.ex @@ -0,0 +1,237 @@ +defmodule ExArrow.Pipeline do + @moduledoc """ + A thin pipeline abstraction over ExArrow streams. + + `ExArrow.Pipeline` provides a stable, composable API for transforming and + sinking Arrow streams without exposing the underlying execution mechanism. + Internally it threads an Elixir `Stream` of `ExArrow.RecordBatch` values + alongside the stream's schema, so transformations stay lazy and batches are + never converted to row maps. + + Every function accepts and returns `{:ok, pipeline} | {:error, reason}`, so + pipelines compose with `|>/2` directly from `ExArrow.Stream.from_*/1` + constructors (which return `{:ok, stream} | {:error, _}`). Errors short- + circuit through every stage. + + ## Example + + ExArrow.Stream.from_flight_sql(client, "SELECT * FROM events") + |> ExArrow.Pipeline.map_batches(fn batch -> + {:ok, slim} = ExArrow.Batch.select(batch, ["id", "score"]) + slim + end) + |> ExArrow.Pipeline.write_parquet("/data/events.parquet") + + ## Laziness + + `map_batches/2` and `each_batch/2` return a lazy pipeline β€” no batches are + consumed until a sink (`write_parquet/2`, `write_flight/3`, or + `write_dataframe/1`) runs. This means a pipeline built but never sunk does + no work. + + ## Telemetry + + `map_batches/2` and `each_batch/2` emit `[:ex_arrow, :pipeline, :batch]` for + every batch that flows through the stage, carrying `rows`, `columns`, and + `batch_count` measurements and `%{source: :pipeline}` metadata. The sinks + emit their respective `[:ex_arrow, :parquet, :write]` and + `[:ex_arrow, :flight, :query]` events. + """ + + alias ExArrow.Flight.Client, as: FlightClient + alias ExArrow.IPC + alias ExArrow.Parquet.Writer + alias ExArrow.RecordBatch + alias ExArrow.Stream, as: ArrowStream + + @enforce_keys [:schema, :enum] + defstruct [:schema, :enum] + + @type t :: %__MODULE__{schema: ExArrow.Schema.t() | nil, enum: Enumerable.t()} + @type input :: ExArrow.Stream.t() | t() + @type threaded :: {:ok, input()} | {:error, term()} + + @doc """ + Wrap an `ExArrow.Stream.t()` (or thread an existing pipeline) and apply + `fun` to each batch. + + The resulting pipeline is lazy: `fun` runs only when a sink consumes the + pipeline. `fun` receives an `ExArrow.RecordBatch.t()` and should return an + `ExArrow.RecordBatch.t()` (or any term if the downstream sink expects it). + + ## Example + + ExArrow.Stream.from_parquet("/data/events.parquet") + |> ExArrow.Pipeline.map_batches(&ExArrow.Batch.schema/1) + """ + @spec map_batches(threaded(), (RecordBatch.t() -> term())) :: threaded() + def map_batches(threaded, fun) + + def map_batches({:error, _} = err, _fun), do: err + + def map_batches({:ok, input}, fun) when is_function(fun, 1) do + with {:ok, pipeline} <- wrap(input) do + enum = + Stream.map(pipeline.enum, fn batch -> + result = fun.(batch) + emit_pipeline_telemetry(batch) + result + end) + + {:ok, %{pipeline | enum: enum}} + end + end + + @doc """ + Run `fun` for its side effects on each batch without changing the pipeline. + + Lazy: the side effect runs only when a sink consumes the pipeline. The + pipeline's batches pass through unchanged. + """ + @spec each_batch(threaded(), (RecordBatch.t() -> term())) :: threaded() + def each_batch(threaded, fun) + + def each_batch({:error, _} = err, _fun), do: err + + def each_batch({:ok, input}, fun) when is_function(fun, 1) do + with {:ok, pipeline} <- wrap(input) do + enum = + Stream.map(pipeline.enum, fn batch -> + fun.(batch) + emit_pipeline_telemetry(batch) + batch + end) + + {:ok, %{pipeline | enum: enum}} + end + end + + @doc """ + Consume the pipeline and write every batch to a Parquet file at `path`. + + Triggers evaluation of all upstream stages. Emits a + `[:ex_arrow, :parquet, :write]` telemetry event. Returns `:ok` or + `{:error, message}`. + + > #### Memory {: .warning} + > This sink materialises all batches into a BEAM list before writing. + > For very large streams this can be a memory concern. The underlying + > `ExArrow.Parquet.Writer.to_file/3` currently requires a list; a + > streaming write path may be added in a future release. For datasets + > that do not fit in memory, iterate with `ExArrow.Stream.next/1` and + > write batches incrementally. + """ + @spec write_parquet(threaded(), Path.t()) :: :ok | {:error, String.t()} + def write_parquet(threaded, path) when is_binary(path) do + with {:ok, pipeline} <- unwrap(threaded), + batches = Enum.to_list(pipeline.enum), + {:ok, schema} <- resolve_schema(pipeline, batches) do + rows = Enum.sum(Enum.map(batches, &RecordBatch.num_rows/1)) + + ExArrow.Telemetry.execute( + [:ex_arrow, :parquet, :write], + %{rows: rows, batch_count: length(batches)}, + %{destination: path, source: :pipeline} + ) + + Writer.to_file(path, schema, batches) + end + end + + @doc """ + Consume the pipeline and upload every batch to a Flight server via `client`. + + `opts` are forwarded to `ExArrow.Flight.Client.do_put/4`. Emits a + `[:ex_arrow, :flight, :query]` telemetry event. Returns `:ok` or + `{:error, reason}`. + + > #### Memory {: .warning} + > This sink materialises all batches into a BEAM list before uploading. + > See `write_parquet/2` for details and alternatives. + """ + @spec write_flight(threaded(), ExArrow.Flight.Client.t(), keyword()) :: + :ok | {:error, term()} + def write_flight(threaded, client, opts \\ []) + + def write_flight({:error, _} = err, _client, _opts), do: err + + def write_flight({:ok, input}, client, opts) do + with {:ok, pipeline} <- wrap(input), + batches = Enum.to_list(pipeline.enum), + {:ok, schema} <- resolve_schema(pipeline, batches) do + rows = Enum.sum(Enum.map(batches, &RecordBatch.num_rows/1)) + + ExArrow.Telemetry.execute( + [:ex_arrow, :flight, :query], + %{rows: rows, batch_count: length(batches)}, + %{destination: Keyword.get(opts, :descriptor), source: :pipeline} + ) + + FlightClient.do_put(client, schema, batches, opts) + end + end + + @doc """ + Consume the pipeline and convert it into an `Explorer.DataFrame`. + + Requires the optional `{:explorer, "~> 0.11"}` dependency. Returns + `{:ok, dataframe}` or `{:error, message}`. + + > #### Memory {: .warning} + > This sink materialises all batches into a BEAM list before conversion. + > See `write_parquet/2` for details and alternatives. + """ + @spec write_dataframe(threaded()) :: + {:ok, Explorer.DataFrame.t()} | {:error, String.t()} + def write_dataframe(threaded) do + with {:ok, pipeline} <- unwrap(threaded), + batches = Enum.to_list(pipeline.enum), + {:ok, schema} <- resolve_schema(pipeline, batches), + {:ok, ipc_bin} <- IPC.Writer.to_binary(schema, batches), + {:ok, stream} <- IPC.Reader.from_binary(ipc_bin) do + ExArrow.DataFrame.from_arrow(stream) + end + end + + # Internals + + defp wrap(%__MODULE__{} = pipeline), do: {:ok, pipeline} + + defp wrap(source) do + if ArrowStream.stream?(source) do + case ArrowStream.schema(source) do + {:ok, schema} -> {:ok, %__MODULE__{schema: schema, enum: source}} + {:error, _} = err -> err + end + else + {:error, "expected an ExArrow.Stream or ExArrow.Pipeline, got: #{inspect(source)}"} + end + end + + defp unwrap({:ok, input}), do: wrap(input) + defp unwrap({:error, _} = err), do: err + + # The schema is derived from the first emitted batch so that transformations + # which change the schema (select/drop/rename) are reflected accurately. The + # schema captured at wrap time is only used as a fallback when the pipeline + # produced zero batches (so an empty Parquet file still carries a schema). + defp resolve_schema(%__MODULE__{schema: captured}, []) do + if captured, do: {:ok, captured}, else: {:error, "no schema and no batches"} + end + + defp resolve_schema(_pipeline, [first | _]) do + if ExArrow.RecordBatch.record_batch?(first) do + {:ok, RecordBatch.schema(first)} + else + {:error, "pipeline produced a non-batch value; cannot derive schema: #{inspect(first)}"} + end + end + + defp emit_pipeline_telemetry(batch) do + if ExArrow.RecordBatch.record_batch?(batch) do + measurements = ExArrow.Telemetry.batch_measurements(batch) + + ExArrow.Telemetry.execute([:ex_arrow, :pipeline, :batch], measurements, %{source: :pipeline}) + end + end +end diff --git a/lib/ex_arrow/sink.ex b/lib/ex_arrow/sink.ex new file mode 100644 index 0000000..fb45888 --- /dev/null +++ b/lib/ex_arrow/sink.ex @@ -0,0 +1,191 @@ +defmodule ExArrow.Sink.Helpers do + @moduledoc false + + alias ExArrow.RecordBatch + alias ExArrow.Stream + + @spec normalise_source(term(), String.t()) :: + {:ok, ExArrow.Schema.t(), [RecordBatch.t()]} | {:error, String.t()} + def normalise_source(source, sink_name) do + cond do + Stream.stream?(source) -> + with {:ok, schema} <- Stream.schema(source) do + {:ok, schema, Stream.to_list(source)} + end + + is_tuple(source) and tuple_size(source) == 2 -> + {schema, batches} = source + + if is_list(batches) do + {:ok, schema, batches} + else + {:error, "expected {schema, [batches]}, got: #{inspect(source)}"} + end + + is_list(source) -> + normalise_list(source, sink_name) + + true -> + {:error, "unsupported source for #{sink_name} sink: #{inspect(source)}"} + end + end + + defp normalise_list([], sink_name) do + {:error, "cannot write an empty batch list to #{sink_name}"} + end + + defp normalise_list([first | _] = batches, _sink_name) do + if RecordBatch.record_batch?(first) do + {:ok, RecordBatch.schema(first), batches} + else + {:error, "expected a list of ExArrow.RecordBatch, got: #{inspect(first)}"} + end + end +end + +defmodule ExArrow.Sink.Parquet do + @moduledoc """ + Write an Arrow stream or batch list to a Parquet file. + + A thin sink that consumes an `ExArrow.Stream.t()` (or a + `{schema, [batches]}` tuple) and writes it to a Parquet file via + `ExArrow.Parquet.Writer.to_file/3`. Emits a + `[:ex_arrow, :parquet, :write]` telemetry event. + + ## Example + + {:ok, stream} = ExArrow.Stream.from_flight_sql(client, "SELECT * FROM t") + ExArrow.Sink.Parquet.write(stream, "/data/out.parquet") + """ + + alias ExArrow.Parquet.Writer + alias ExArrow.RecordBatch + alias ExArrow.Sink.Helpers + alias ExArrow.Stream + + @doc """ + Write `source` to a Parquet file at `path`. + + `source` may be: + + - an `ExArrow.Stream.t()` β€” schema is read with `ExArrow.Stream.schema/1` + and batches with `ExArrow.Stream.to_list/1`. + - a `{schema, [batches]}` tuple. + - a list of `ExArrow.RecordBatch.t()` β€” schema is taken from the first batch. + + Returns `:ok` or `{:error, message}`. + """ + @spec write(Stream.t() | {ExArrow.Schema.t(), [RecordBatch.t()]} | [RecordBatch.t()], Path.t()) :: + :ok | {:error, String.t()} + def write(source, path) when is_binary(path) do + with {:ok, schema, batches} <- Helpers.normalise_source(source, "Parquet") do + rows = Enum.sum(Enum.map(batches, &RecordBatch.num_rows/1)) + + ExArrow.Telemetry.execute( + [:ex_arrow, :parquet, :write], + %{rows: rows, batch_count: length(batches)}, + %{destination: path, source: :sink} + ) + + Writer.to_file(path, schema, batches) + end + end +end + +defmodule ExArrow.Sink.Flight do + @moduledoc """ + Upload an Arrow stream or batch list to a Flight server. + + Wraps `ExArrow.Flight.Client.do_put/4`. Emits a + `[:ex_arrow, :flight, :query]` telemetry event with the destination + descriptor in the metadata. + + ## Example + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + ExArrow.Sink.Flight.write(stream, client, descriptor: {:cmd, "events"}) + """ + + alias ExArrow.Flight.Client + alias ExArrow.RecordBatch + alias ExArrow.Sink.Helpers + alias ExArrow.Stream + + @doc """ + Upload `source` to a Flight server via `client`. + + Accepts the same `source` shapes as `ExArrow.Sink.Parquet.write/2`. `opts` + are forwarded to `ExArrow.Flight.Client.do_put/4` (e.g. `:descriptor`). + + Returns `:ok` or `{:error, reason}`. + """ + @spec write( + Stream.t() | {ExArrow.Schema.t(), [RecordBatch.t()]} | [RecordBatch.t()], + Client.t(), + keyword() + ) :: :ok | {:error, term()} + def write(source, client, opts \\ []) + + def write([], _client, _opts), do: :ok + + def write(source, client, opts) do + with {:ok, schema, batches} <- Helpers.normalise_source(source, "Flight") do + emit_flight_telemetry(batches, Keyword.get(opts, :descriptor)) + Client.do_put(client, schema, batches, opts) + end + end + + defp emit_flight_telemetry(batches, descriptor) do + rows = Enum.sum(Enum.map(batches, &RecordBatch.num_rows/1)) + + ExArrow.Telemetry.execute( + [:ex_arrow, :flight, :query], + %{rows: rows, batch_count: length(batches)}, + %{destination: descriptor, source: :sink} + ) + end +end + +defmodule ExArrow.Sink.DataFrame do + @moduledoc """ + Convert an Arrow stream or batch into an Explorer DataFrame. + + Delegates to `ExArrow.DataFrame.from_arrow/1`, which accepts either an + `ExArrow.RecordBatch.t()` or an `ExArrow.Stream.t()`. Requires the optional + `{:explorer, "~> 0.11"}` dependency. + + ## Example + + {:ok, df} = ExArrow.Sink.DataFrame.write(stream) + """ + + @doc """ + Convert `source` to an `Explorer.DataFrame`. + + Returns `{:ok, dataframe}` or `{:error, message}`. + """ + @spec write(ExArrow.RecordBatch.t() | ExArrow.Stream.t()) :: + {:ok, Explorer.DataFrame.t()} | {:error, String.t()} + def write(source), do: ExArrow.DataFrame.from_arrow(source) +end + +defmodule ExArrow.Sink.Nx do + @moduledoc """ + Convert an Arrow batch into an `Nx.Tensor`. + + Delegates to `ExArrow.to_nx/1`. Requires the optional + `{:nx, "~> 0.12"}` dependency. + + ## Example + + {:ok, tensor} = ExArrow.Sink.Nx.write(batch) + """ + + @doc """ + Convert `batch` to an `Nx.Tensor`. + + Returns `{:ok, tensor}` or `{:error, message}`. + """ + @spec write(ExArrow.RecordBatch.t()) :: {:ok, Nx.Tensor.t()} | {:error, String.t()} + def write(batch), do: ExArrow.to_nx(batch) +end diff --git a/lib/ex_arrow/stream.ex b/lib/ex_arrow/stream.ex index aeda199..d096fdf 100644 --- a/lib/ex_arrow/stream.ex +++ b/lib/ex_arrow/stream.ex @@ -55,17 +55,180 @@ defmodule ExArrow.Stream do Stopping enumeration early (e.g. `Enum.take/2`) is safe β€” the resource will be released when the stream variable goes out of scope. """ + alias ExArrow.ADBC.Connection, as: ADBCConnection + alias ExArrow.ADBC.Statement, as: ADBCStatement alias ExArrow.RecordBatch alias ExArrow.Schema - @opaque t :: %__MODULE__{resource: reference(), backend: :ipc | :adbc | :parquet | :flight_sql} - defstruct [:resource, backend: :ipc] + @opaque t :: %__MODULE__{ + resource: reference(), + backend: :ipc | :adbc | :parquet | :flight_sql, + source: term() + } + defstruct [:resource, :source, backend: :ipc] + + @doc """ + Returns the origin metadata attached to this stream by the `from_*/` + constructors. The value is backend-specific (e.g. `{:parquet, path}` or + `{:flight_sql, sql}`) and is forwarded to telemetry events as `:source`. + """ + @spec source(t()) :: term() + def source(%__MODULE__{source: source}), do: source @doc false @spec stream?(term()) :: boolean() def stream?(%__MODULE__{}), do: true def stream?(_), do: false + # Source constructors + + @doc """ + Open an Arrow IPC stream from an in-memory `binary`. + + Delegates to `ExArrow.IPC.Reader.from_binary/1` and tags the resulting + stream with `source: {:ipc, :binary}` for telemetry. Returns + `{:ok, stream}` or `{:error, message}`. + + ## Example + + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_bytes) + """ + @spec from_ipc(binary()) :: {:ok, t()} | {:error, String.t()} + def from_ipc(binary) when is_binary(binary) do + with {:ok, stream} <- ExArrow.IPC.Reader.from_binary(binary) do + {:ok, %{stream | source: {:ipc, :binary}}} + end + end + + @doc """ + Open an Arrow IPC stream from a file at `path`. + + Delegates to `ExArrow.IPC.Reader.from_file/1` and tags the stream with + `source: {:ipc, path}`. Returns `{:ok, stream}` or `{:error, message}`. + """ + @spec from_ipc_file(Path.t()) :: {:ok, t()} | {:error, String.t()} + def from_ipc_file(path) when is_binary(path) do + with {:ok, stream} <- ExArrow.IPC.Reader.from_file(path) do + {:ok, %{stream | source: {:ipc, path}}} + end + end + + @doc """ + Open a Parquet file at `path` for lazy row-group streaming. + + Delegates to `ExArrow.Parquet.Reader.from_file/1`, tags the stream with + `source: {:parquet, path}`, and emits a `[:ex_arrow, :parquet, :read]` + telemetry event on successful open. Returns `{:ok, stream}` or + `{:error, message}`. + + ## Example + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + """ + @spec from_parquet(Path.t()) :: {:ok, t()} | {:error, String.t()} + def from_parquet(path) when is_binary(path) do + with {:ok, stream} <- ExArrow.Parquet.Reader.from_file(path) do + ExArrow.Telemetry.execute([:ex_arrow, :parquet, :read], %{}, %{source: path}) + {:ok, %{stream | source: {:parquet, path}}} + end + end + + @doc """ + Open a Parquet stream from an in-memory `binary`. + + Delegates to `ExArrow.Parquet.Reader.from_binary/1` and emits a + `[:ex_arrow, :parquet, :read]` telemetry event on successful open with + `source: :binary`. + """ + @spec from_parquet_binary(binary()) :: {:ok, t()} | {:error, String.t()} + def from_parquet_binary(binary) when is_binary(binary) do + with {:ok, stream} <- ExArrow.Parquet.Reader.from_binary(binary) do + ExArrow.Telemetry.execute([:ex_arrow, :parquet, :read], %{}, %{source: :binary}) + {:ok, %{stream | source: {:parquet, :binary}}} + end + end + + @doc """ + Retrieve data for `ticket` from a Flight server as a stream of record + batches. + + Delegates to `ExArrow.Flight.Client.do_get/2` and emits a + `[:ex_arrow, :flight, :query]` telemetry event on success. The stream is + tagged with `source: {:flight, ticket}`. + + ## Example + + {:ok, stream} = ExArrow.Stream.from_flight(client, "sales_2024") + """ + @spec from_flight(ExArrow.Flight.Client.t(), term()) :: + {:ok, t()} | {:error, term()} + def from_flight(client, ticket) do + with {:ok, stream} <- ExArrow.Flight.Client.do_get(client, ticket) do + ExArrow.Telemetry.execute([:ex_arrow, :flight, :query], %{}, %{source: ticket}) + {:ok, %{stream | source: {:flight, ticket}}} + end + end + + @doc """ + Execute a SQL query against a Flight SQL server and return a lazy stream of + record batches. + + Delegates to `ExArrow.FlightSQL.Client.stream_query/2`, emits a + `[:ex_arrow, :flight_sql, :query]` telemetry event on success, and tags the + stream with `source: {:flight_sql, sql}`. + + ## Example + + {:ok, stream} = ExArrow.Stream.from_flight_sql(client, "SELECT * FROM events") + """ + @spec from_flight_sql(ExArrow.FlightSQL.Client.t(), String.t()) :: + {:ok, t()} | {:error, term()} + # sobelow_skip ["SQL.Query"] + # False positive: SQL is forwarded to a remote Flight SQL server over gRPC. + def from_flight_sql(client, sql) when is_binary(sql) do + with {:ok, stream} <- ExArrow.FlightSQL.Client.stream_query(client, sql) do + ExArrow.Telemetry.execute([:ex_arrow, :flight_sql, :query], %{}, %{source: sql}) + {:ok, %{stream | source: {:flight_sql, sql}}} + end + end + + @doc """ + Execute an ADBC statement and return its result as a stream of record + batches. + + Accepts either a prepared `ExArrow.ADBC.Statement.t()` (built with + `ExArrow.ADBC.Statement.new/2` or `new/3`) or a `{connection, sql}` pair, + in which case a one-shot statement is created, executed, and discarded. + + The stream is tagged with `source: {:adbc, sql}` (or `:statement` when a + pre-built statement is passed). + + ## Examples + + {:ok, stream} = ExArrow.Stream.from_adbc(stmt) + + {:ok, stream} = ExArrow.Stream.from_adbc(conn, "SELECT * FROM events") + """ + @spec from_adbc(ADBCStatement.t()) :: + {:ok, t()} | {:error, term()} + def from_adbc(statement) do + with {:ok, stream} <- ADBCStatement.execute(statement) do + {:ok, %{stream | source: {:adbc, :statement}}} + end + end + + @spec from_adbc(ADBCConnection.t(), String.t()) :: + {:ok, t()} | {:error, term()} + # sobelow_skip ["SQL.Query"] + # False positive: SQL is forwarded to an ADBC driver and never executed + # locally in this process. + def from_adbc(connection, sql) when is_binary(sql) do + with {:ok, stmt} <- ADBCStatement.new(connection, sql), + {:ok, stream} <- ADBCStatement.execute(stmt) do + {:ok, %{stream | source: {:adbc, sql}}} + end + end + @doc """ Returns the schema of this stream (without consuming it). Returns `{:error, message}` if the stream is invalid (e.g. poisoned lock). @@ -115,40 +278,47 @@ defmodule ExArrow.Stream do | nil | {:error, String.t()} | {:error, {atom(), non_neg_integer(), String.t()}} - def next(%__MODULE__{resource: ref, backend: :adbc}) do + def next(%__MODULE__{resource: ref, backend: :adbc} = stream) do case native().adbc_stream_next(ref) do :done -> nil - {:ok, batch_ref} -> RecordBatch.from_ref(batch_ref) + {:ok, batch_ref} -> emit_batch(stream, RecordBatch.from_ref(batch_ref)) {:error, msg} -> {:error, msg} end end - def next(%__MODULE__{resource: ref, backend: :ipc}) do + def next(%__MODULE__{resource: ref, backend: :ipc} = stream) do case native().ipc_stream_next(ref) do :done -> nil - {:ok, batch_ref} -> RecordBatch.from_ref(batch_ref) + {:ok, batch_ref} -> emit_batch(stream, RecordBatch.from_ref(batch_ref)) {:error, msg} -> {:error, msg} end end - def next(%__MODULE__{resource: ref, backend: :parquet}) do + def next(%__MODULE__{resource: ref, backend: :parquet} = stream) do case native().parquet_stream_next(ref) do :done -> nil - {:ok, batch_ref} -> RecordBatch.from_ref(batch_ref) + {:ok, batch_ref} -> emit_batch(stream, RecordBatch.from_ref(batch_ref)) {:error, msg} -> {:error, msg} end end - def next(%__MODULE__{resource: ref, backend: :flight_sql}) do + def next(%__MODULE__{resource: ref, backend: :flight_sql} = stream) do case native().flight_sql_stream_next(ref) do :done -> nil - {:ok, batch_ref} -> RecordBatch.from_ref(batch_ref) + {:ok, batch_ref} -> emit_batch(stream, RecordBatch.from_ref(batch_ref)) # Pass the structured triple through so callers retain the gRPC code and status. {:error, {_code, _grpc_status, _msg} = triple} -> {:error, triple} {:error, msg} -> {:error, msg} end end + defp emit_batch(%__MODULE__{source: source}, batch) do + measurements = ExArrow.Telemetry.batch_measurements(batch) + metadata = %{source: source, schema: nil} + ExArrow.Telemetry.execute([:ex_arrow, :stream, :batch], measurements, metadata) + batch + end + @doc """ Collects all remaining batches from the stream into a list. diff --git a/lib/ex_arrow/telemetry.ex b/lib/ex_arrow/telemetry.ex new file mode 100644 index 0000000..f03fd5c --- /dev/null +++ b/lib/ex_arrow/telemetry.ex @@ -0,0 +1,151 @@ +defmodule ExArrow.Telemetry do + @moduledoc """ + Telemetry integration for ExArrow. + + ExArrow emits structured telemetry events for every transport and pipeline + operation. Events follow the `[:ex_arrow, component, action]` naming + convention and carry a consistent set of measurements and metadata so a + single handler can observe the whole system. + + This module is the single emission point used by `ExArrow.Stream`, + `ExArrow.Pipeline`, `ExArrow.Flow`, and the GenStage / Broadway integrations. + Application code should attach handlers with `:telemetry.attach/4` (or + `telemetry_metrics`) rather than calling `execute/3` directly. + + ## Optional dependency + + Telemetry is an optional dependency. Add `{:telemetry, "~> 1.0"}` to your + `mix.exs` to receive events. When the library is absent `execute/3` and + `span/3` degrade to no-ops, so ExArrow itself never crashes because of a + missing handler. + + ## Events + + | Event | When it is emitted | + |------------------------------------|------------------------------------------| + | `[:ex_arrow, :flight, :query]` | A Flight `do_get` stream is opened | + | `[:ex_arrow, :flight_sql, :query]` | A Flight SQL query stream is opened | + | `[:ex_arrow, :parquet, :read]` | A Parquet reader stream is opened | + | `[:ex_arrow, :parquet, :write]` | Batches are written to Parquet | + | `[:ex_arrow, :stream, :batch]` | A single batch is yielded from a stream | + | `[:ex_arrow, :pipeline, :batch]` | A pipeline stage processes a batch | + + ## Measurements + + Every event may carry any subset of the following measurements. Specific + events populate the subset that is meaningful for the operation. + + | Measurement | Type | Meaning | + |---------------|-----------------|------------------------------------------| + | `:rows` | non_neg_integer | Rows in the batch | + | `:columns` | non_neg_integer | Columns in the batch | + | `:duration` | non_neg_integer | Native monotonic time in nanoseconds | + | `:batch_count`| non_neg_integer | Number of batches in a batched operation | + + ## Metadata + + | Field | Type | Meaning | + |----------------|-----------|------------------------------------------| + | `:source` | term() | Origin of the data (path, URI, SQL) | + | `:destination` | term() | Target of a write (path, ticket, host) | + | `:schema` | term() | `ExArrow.Schema` handle when available | + | `:driver` | String.t() | Driver name (e.g. `"adbc_driver_sqlite"`)| + + ## Attaching a handler + + :telemetry.attach( + "ex-arrow-logger", + [:ex_arrow, :stream, :batch], + fn _event, measurements, metadata, _config -> + rows = measurements[:rows] + source = inspect(metadata[:source]) + IO.puts("batch: " <> Integer.to_string(rows) <> " rows from " <> source) + end, + nil + ) + + {:ok, stream} = ExArrow.Stream.from_parquet("/data/events.parquet") + Enum.each(stream, fn _batch -> :ok end) + """ + + @telemetry_available Code.ensure_loaded?(:telemetry) + + @type event_name :: [atom(), ...] + @type measurements :: map() + @type metadata :: map() + + if @telemetry_available do + @doc """ + Emit a telemetry event. + + No-ops when a handler is not attached or when the `:telemetry` application + is not running. `measurements` and `metadata` are passed straight through + to the handler. + + ## Example + + ExArrow.Telemetry.execute( + [:ex_arrow, :stream, :batch], + %{rows: 100, columns: 3}, + %{source: "/data/events.parquet"} + ) + """ + @spec execute(event_name(), measurements(), metadata()) :: :ok + def execute(event_name, measurements, metadata) when is_list(event_name) do + :telemetry.execute(event_name, measurements, metadata) + :ok + end + + @doc """ + Wrap `fun` in a telemetry span under the given event name. + + Emits `event ++ [:start]` and `event ++ [:stop]` (or `event ++ [:exception]`) + events, matching the convention used by `:telemetry.span/3`. Returns the + result of `fun`. The `start` metadata is merged into the `stop` metadata. + + ## Example + + ExArrow.Telemetry.span([:ex_arrow, :flight, :query], %{source: sql}, fn -> + ExArrow.Flight.Client.do_get(client, ticket) + end) + """ + @spec span(event_name(), metadata(), (-> {result, metadata()})) :: result + when result: term() + def span(event_name, start_metadata, fun) when is_function(fun, 0) do + :telemetry.span(event_name, start_metadata, fun) + end + else + @doc false + @spec execute(event_name(), measurements(), metadata()) :: :ok + def execute(_event_name, _measurements, _metadata), do: :ok + + @doc false + @spec span(event_name(), metadata(), (-> {result, metadata()})) :: result + when result: term() + def span(_event_name, _start_metadata, fun) when is_function(fun, 0) do + {result, _metadata} = fun() + result + end + end + + @doc """ + Build a measurements map for a single batch. + + Convenience used by `ExArrow.Stream` and `ExArrow.Pipeline` so every emitter + reports the same shape of measurements for a batch. + """ + @spec batch_measurements(ExArrow.RecordBatch.t(), keyword()) :: measurements() + def batch_measurements(batch, extra \\ []) do + alias ExArrow.RecordBatch + + rows = RecordBatch.num_rows(batch) + columns = RecordBatch.num_columns(batch) + + %{ + rows: rows, + columns: columns, + batch_count: 1 + } + |> Map.merge(Map.new(extra)) + end +end diff --git a/livebook/00_quickstart.livemd b/livebook/00_quickstart.livemd index 6272a8f..1ccd4ac 100644 --- a/livebook/00_quickstart.livemd +++ b/livebook/00_quickstart.livemd @@ -22,7 +22,7 @@ local? = File.exists?(Path.join(__DIR__, "../native/ex_arrow_native/Cargo.toml") } else { - {:ex_arrow, "~> 0.6.3"}, + {:ex_arrow, "~> 0.7.0"}, [], [adbc: [drivers: [:sqlite]]] } @@ -257,4 +257,71 @@ Connect to a Flight SQL server, prepare a query, bind parameters, and execute: IO.puts("Uncomment the cells above when a Flight SQL server is available.") ``` +--- + +## 9. Streaming pipelines (v0.7.0) + +v0.7.0 introduces first-class streaming and a pipeline DSL. The unit of +execution is the Arrow `RecordBatch`. + +**Stream constructors** β€” one entry point per source: + +```elixir +# From an IPC binary (built-in fixture) +{:ok, ipc_bytes} = ExArrow.Native.ipc_test_fixture_binary() +{:ok, stream} = ExArrow.Stream.from_ipc(ipc_bytes) +{:ok, schema} = ExArrow.Stream.schema(stream) +ExArrow.Schema.field_names(schema) +``` + +**Batch operations** β€” lightweight, Arrow-native transforms: + +```elixir +batch = ExArrow.Stream.next(stream) + +{:ok, slim} = ExArrow.Batch.select(batch, ["id"]) +ExArrow.RecordBatch.column_names(slim) + +{:ok, renamed} = ExArrow.Batch.rename(slim, %{"id" => "user_id"}) +ExArrow.RecordBatch.column_names(renamed) + +{:ok, first1} = ExArrow.Batch.take(renamed, 1) +ExArrow.RecordBatch.num_rows(first1) +``` + +**Pipeline DSL** β€” lazy `map_batches` + sink: + +```elixir +# Re-open the stream (the previous one was consumed) +{:ok, stream2} = ExArrow.Stream.from_ipc(ipc_bytes) +out_path = Path.join(System.tmp_dir!(), "ex_arrow_quickstart_pipeline.parquet") + +{:ok, stream2} +|> ExArrow.Pipeline.map_batches(fn b -> + {:ok, s} = ExArrow.Batch.select(b, ["id"]) + s +end) +|> ExArrow.Pipeline.write_parquet(out_path) + +File.exists?(out_path) +``` + +**Telemetry** β€” attach a handler and observe every batch: + +```elixir +:telemetry.attach( + "ex-arrow-quickstart", + [:ex_arrow, :stream, :batch], + fn _event, measurements, metadata, _config -> + IO.puts("batch: #{measurements[:rows]} rows from #{inspect(metadata[:source])}") + end, + nil +) + +{:ok, stream3} = ExArrow.Stream.from_ipc(ipc_bytes) +_ = ExArrow.Stream.to_list(stream3) + +:telemetry.detach("ex-arrow-quickstart") +``` + Docs: [hexdocs.pm/ex_arrow](https://hexdocs.pm/ex_arrow). diff --git a/livebook/01_ipc.livemd b/livebook/01_ipc.livemd index 13344f6..cdca933 100644 --- a/livebook/01_ipc.livemd +++ b/livebook/01_ipc.livemd @@ -22,7 +22,7 @@ local? = File.exists?(Path.join(__DIR__, "../native/ex_arrow_native/Cargo.toml") } else { - {:ex_arrow, "~> 0.6.3"}, + {:ex_arrow, "~> 0.7.0"}, [], [adbc: [drivers: [:sqlite]]] } diff --git a/livebook/02_flight.livemd b/livebook/02_flight.livemd index 189d3c0..43f7506 100644 --- a/livebook/02_flight.livemd +++ b/livebook/02_flight.livemd @@ -21,7 +21,7 @@ local? = File.exists?(Path.join(__DIR__, "../native/ex_arrow_native/Cargo.toml") } else { - {:ex_arrow, "~> 0.6.3"}, + {:ex_arrow, "~> 0.7.0"}, [], [adbc: [drivers: [:sqlite]]] } diff --git a/livebook/03_adbc.livemd b/livebook/03_adbc.livemd index c5444e6..3edc66d 100644 --- a/livebook/03_adbc.livemd +++ b/livebook/03_adbc.livemd @@ -21,7 +21,7 @@ local? = File.exists?(Path.join(__DIR__, "../native/ex_arrow_native/Cargo.toml") } else { - {:ex_arrow, "~> 0.6.3"}, + {:ex_arrow, "~> 0.7.0"}, [], [adbc: [drivers: [:sqlite]]] } diff --git a/livebook/04_adbc_integration.livemd b/livebook/04_adbc_integration.livemd index c7a8ba9..b478173 100644 --- a/livebook/04_adbc_integration.livemd +++ b/livebook/04_adbc_integration.livemd @@ -22,7 +22,7 @@ local? = File.exists?(Path.join(__DIR__, "../native/ex_arrow_native/Cargo.toml") } else { - {:ex_arrow, "~> 0.6.3"}, + {:ex_arrow, "~> 0.7.0"}, [], [adbc: [drivers: [:sqlite]]] } diff --git a/livebook/README.md b/livebook/README.md index 873905f..ea7a27e 100644 --- a/livebook/README.md +++ b/livebook/README.md @@ -6,13 +6,13 @@ Tutorial notebooks for the **ex_arrow** library, suitable for an introductory Me | File | Purpose | |------|---------| -| **00_quickstart.livemd** | Get something running in minutes: IPC read/write, Flight echo client/server, ADBC SQLite query, Explorer/Nx interchange. | +| **00_quickstart.livemd** | Get something running in minutes: IPC read/write, Flight echo client/server, ADBC SQLite query, Explorer/Nx interchange, and v0.7.0 streaming pipelines. | | **01_ipc.livemd** | IPC deep dive: stream vs file format, reading from binary/file, writing, schema and types, optional Explorer interop. | | **02_flight.livemd** | Arrow Flight: echo server, client, do_put/do_get, list_flights, get_flight_info, get_schema, actions, Flight SQL prepared statements. | | **03_adbc.livemd** | ADBC: `:adbc_package` backend, Database β†’ Connection β†’ Statement β†’ Stream, metadata APIs (native driver), Explorer roundtrip. | | **04_adbc_integration.livemd** | **adbc_package** backend with connection pooling (NimblePool), concurrent queries. | -Together they demonstrate ExArrow functionality: IPC (stream + file), Flight (client + server + Flight SQL), and ADBC (Arrow result streams). +Together they demonstrate ExArrow functionality: IPC (stream + file), Flight (client + server + Flight SQL), ADBC (Arrow result streams), and the v0.7.0 pipeline DSL (`ExArrow.Stream`, `ExArrow.Batch`, `ExArrow.Pipeline`, telemetry). ## How to run @@ -26,7 +26,7 @@ Together they demonstrate ExArrow functionality: IPC (stream + file), Flight (cl | Where you open the notebook | `ex_arrow` source | |----------------------------|-------------------| | From `livebook/` in a git clone | Local path + `EX_ARROW_BUILD=1` (compile NIF from Rust) | -| From Livebook autosave or elsewhere | Hex `~> 0.6.3` (precompiled NIF, no Rust) | +| From Livebook autosave or elsewhere | Hex `~> 0.7.0` (precompiled NIF, no Rust) | ### ADBC in Livebook diff --git a/mix.exs b/mix.exs index 9d10d3b..0caec63 100644 --- a/mix.exs +++ b/mix.exs @@ -1,7 +1,7 @@ defmodule ExArrow.MixProject do use Mix.Project - @version "0.6.3" + @version "0.7.0" @source_url "https://github.com/thanos/ex_arrow" def project do @@ -85,6 +85,10 @@ defmodule ExArrow.MixProject do {:nx, "~> 0.12", optional: true}, {:nimble_pool, "~> 1.1", optional: true}, {:rustler, "~> 0.36 or ~> 0.38", optional: true}, + {:telemetry, "~> 1.0", optional: true}, + {:flow, "~> 1.2", optional: true}, + {:gen_stage, "~> 1.2", optional: true}, + {:broadway, "~> 1.0", optional: true}, {:ex_doc, "~> 0.40", only: :dev}, {:benchee, "~> 1.5", only: :dev}, {:benchee_html, "~> 1.0", only: :dev}, @@ -109,6 +113,12 @@ defmodule ExArrow.MixProject do "guides/02_explorer_integration.md", "guides/03_nx_integration.md", "guides/04_arrow_ecosystem.md", + "guides/05_arrow_pipelines_overview.md", + "guides/06_arrow_streams.md", + "guides/07_arrow_and_flow.md", + "guides/08_arrow_and_genstage.md", + "guides/09_arrow_and_broadway.md", + "guides/10_arrow_pipeline_patterns.md", "docs/overview.md", "docs/memory_model.md", "docs/ipc_guide.md", @@ -126,6 +136,27 @@ defmodule ExArrow.MixProject do IPC: [ExArrow.IPC.Reader, ExArrow.IPC.Writer, ExArrow.IPC.File], Parquet: [ExArrow.Parquet.Reader, ExArrow.Parquet.Writer], "Compute kernels": [ExArrow.Compute], + "Batch operations": [ExArrow.Batch], + Pipeline: [ + ExArrow.Pipeline, + ExArrow.Sink.Parquet, + ExArrow.Sink.Flight, + ExArrow.Sink.DataFrame, + ExArrow.Sink.Nx + ], + Flow: [ExArrow.Flow], + GenStage: [ + ExArrow.GenStage, + ExArrow.GenStage.ParquetProducer, + ExArrow.GenStage.FlightProducer, + ExArrow.GenStage.ADBCProducer + ], + Broadway: [ + ExArrow.Broadway.BatchBuilder, + ExArrow.Broadway.ParquetSink, + ExArrow.Broadway.FlightSink + ], + Telemetry: [ExArrow.Telemetry], Flight: [ ExArrow.Flight.Client, ExArrow.Flight.Server, diff --git a/mix.lock b/mix.lock index 377d8e4..61718bd 100644 --- a/mix.lock +++ b/mix.lock @@ -6,6 +6,7 @@ "benchee": {:hex, :benchee, "1.5.1", "b95cbc36c4b98969a5c592a246e171041eb683c56bad1cb4f49a3b081ba66087", [:mix], [{:deep_merge, "~> 1.0", [hex: :deep_merge, repo: "hexpm", optional: false]}, {:statistex, "~> 1.1", [hex: :statistex, repo: "hexpm", optional: false]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: true]}], "hexpm", "a539301f8dfd4efc5c5123bfb9d47ebde20092a863a5b5b16c2a60d2243dfce7"}, "benchee_html": {:hex, :benchee_html, "1.0.1", "1e247c0886c3fdb0d3f4b184b653a8d6fb96e4ad0d0389267fe4f36968772e24", [:mix], [{:benchee, ">= 0.99.0 and < 2.0.0", [hex: :benchee, repo: "hexpm", optional: false]}, {:benchee_json, "~> 1.0", [hex: :benchee_json, repo: "hexpm", optional: false]}], "hexpm", "b00a181af7152431901e08f3fc9f7197ed43ff50421a8347b0c80bf45d5b3fef"}, "benchee_json": {:hex, :benchee_json, "1.0.0", "cc661f4454d5995c08fe10dd1f2f72f229c8f0fb1c96f6b327a8c8fc96a91fe5", [:mix], [{:benchee, ">= 0.99.0 and < 2.0.0", [hex: :benchee, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "da05d813f9123505f870344d68fb7c86a4f0f9074df7d7b7e2bb011a63ec231c"}, + "broadway": {:hex, :broadway, "1.3.0", "f75f6376159b74f55c5ba2629dac613e4fd79d9e71148ab5fbac8fdd7c999d2a", [:mix], [{:gen_stage, "~> 1.0", [hex: :gen_stage, repo: "hexpm", optional: false]}, {:nimble_options, "~> 0.3.7 or ~> 0.4 or ~> 1.0", [hex: :nimble_options, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.3 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "bef3b4c5512d0072917b70239cbecf8f76a2587465a5b7c3e2b9ae18b4bc405b"}, "bun": {:hex, :bun, "1.6.0", "ab49eb8b956f7492e8d6c5a43c0b96e94074aa0ace698a4be7e9b3b3b281d1b3", [:mix], [{:nimble_pool, "~> 1.0", [hex: :nimble_pool, repo: "hexpm", optional: false]}], "hexpm", "6f17b0544b452077da87d14192042744826dabb089d968b8caf8b298968e7a17"}, "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"}, "castore": {:hex, :castore, "1.0.18", "5e43ef0ec7d31195dfa5a65a86e6131db999d074179d2ba5a8de11fe14570f55", [:mix], [], "hexpm", "f393e4fe6317829b158fb74d86eb681f737d2fe326aa61ccf6293c4104957e34"}, @@ -27,7 +28,9 @@ "file_system": {:hex, :file_system, "1.1.1", "31864f4685b0148f25bd3fbef2b1228457c0c89024ad67f7a81a3ffbc0bbad3a", [:mix], [], "hexpm", "7a15ff97dfe526aeefb090a7a9d3d03aa907e100e262a0f8f7746b78f8f87a5d"}, "finch": {:hex, :finch, "0.22.0", "5c48fa6f9706a78eb9036cacb67b8b996b4e66d111c543f4c29bb0f879a6806b", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:mint, "~> 1.8", [hex: :mint, repo: "hexpm", optional: false]}, {:nimble_options, "~> 0.4 or ~> 1.0", [hex: :nimble_options, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.1", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "b94e83c47780fc6813f746a1f1a34ee65cda42da4c5ea26a68f0acc4498e23dc"}, "fine": {:hex, :fine, "0.1.6", "4bf7151493443c454aac9f2fa2f34f5fefd0346a83fb5586a016c4a135c63247", [:mix], [], "hexpm", "5638eb4495488e885ebec167fa57973e5c35e1a50c344eb7666c90ec1c4e3b12"}, + "flow": {:hex, :flow, "1.2.4", "1dd58918287eb286656008777cb32714b5123d3855956f29aa141ebae456922d", [:mix], [{:gen_stage, "~> 1.0", [hex: :gen_stage, repo: "hexpm", optional: false]}], "hexpm", "874adde96368e71870f3510b91e35bc31652291858c86c0e75359cbdd35eb211"}, "fss": {:hex, :fss, "0.1.1", "9db2344dbbb5d555ce442ac7c2f82dd975b605b50d169314a20f08ed21e08642", [:mix], [], "hexpm", "78ad5955c7919c3764065b21144913df7515d52e228c09427a004afe9c1a16b0"}, + "gen_stage": {:hex, :gen_stage, "1.3.2", "7c77e5d1e97de2c6c2f78f306f463bca64bf2f4c3cdd606affc0100b89743b7b", [:mix], [], "hexpm", "0ffae547fa777b3ed889a6b9e1e64566217413d018cabd825f786e843ffe63e7"}, "hpax": {:hex, :hpax, "1.0.3", "ed67ef51ad4df91e75cc6a1494f851850c0bd98ebc0be6e81b026e765ee535aa", [:mix], [], "hexpm", "8eab6e1cfa8d5918c2ce4ba43588e894af35dbd8e91e6e55c817bca5847df34a"}, "iso8601": {:hex, :iso8601, "1.3.4", "7b1f095f86f6cf65e1e5a77872e8e8bf69bd58d4c3a415b3f77d9cc9423ecbb9", [:rebar3], [], "hexpm", "a334469c07f1c219326bc891a95f5eec8eb12dd8071a3fff56a7843cb20fae34"}, "jason": {:hex, :jason, "1.4.5", "2e3a008590b0b8d7388c20293e9dcc9cf3e5d642fd2a114e4cbbb52e595d940a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0 or ~> 3.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "b0c823996102bcd0239b3c2444eb00409b72f6a140c1950bc8b457d836b30684"}, diff --git a/test/ex_arrow/batch_property_test.exs b/test/ex_arrow/batch_property_test.exs new file mode 100644 index 0000000..9a98e60 --- /dev/null +++ b/test/ex_arrow/batch_property_test.exs @@ -0,0 +1,81 @@ +defmodule ExArrow.BatchPropertyTest do + use ExUnit.Case, async: true + use ExUnitProperties + + import ExArrow.TestFixtures + alias ExArrow.Batch + alias ExArrow.RecordBatch + + defp s64_column(batch, name) do + ref = RecordBatch.resource_ref(batch) + {:ok, {binary, "s64", _n}} = ExArrow.Native.record_batch_column_buffer(ref, name) + + for <>, do: v + end + + property "select/2 with all columns preserves row count and values" do + check all(values <- list_of(integer(-1000..1000), min_length: 1, max_length: 50)) do + batch = s64_batch(values) + + assert {:ok, selected} = Batch.select(batch, ["v"]) + assert RecordBatch.num_rows(selected) == length(values) + assert s64_column(selected, "v") == values + end + end + + property "take/2 with n >= rows returns the batch unchanged" do + check all(values <- list_of(integer(-1000..1000), min_length: 1, max_length: 50)) do + batch = s64_batch(values) + n = length(values) + + assert {:ok, taken} = Batch.take(batch, n) + assert s64_column(taken, "v") == values + + assert {:ok, taken2} = Batch.take(batch, n + 10) + assert s64_column(taken2, "v") == values + end + end + + property "take/2 with n returns the first n values" do + check all(values <- list_of(integer(-1000..1000), min_length: 1, max_length: 50)) do + batch = s64_batch(values) + n = :rand.uniform(length(values)) + + assert {:ok, taken} = Batch.take(batch, n) + assert s64_column(taken, "v") == Enum.take(values, n) + end + end + + property "drop/2 of an empty list preserves the schema" do + check all(values <- list_of(integer(-1000..1000), min_length: 1, max_length: 50)) do + batch = s64_batch(values) + + assert {:ok, rest} = Batch.drop(batch, []) + assert RecordBatch.column_names(rest) == ["v"] + assert s64_column(rest, "v") == values + end + end + + property "rename/2 round-trips with the inverse mapping" do + check all(values <- list_of(integer(-1000..1000), min_length: 1, max_length: 50)) do + batch = s64_batch(values) + + assert {:ok, r1} = Batch.rename(batch, %{"v" => "w"}) + assert {:ok, r2} = Batch.rename(r1, %{"w" => "v"}) + assert RecordBatch.column_names(r2) == ["v"] + assert s64_column(r2, "v") == values + end + end + + property "take/2 with indices selects exactly those rows (original order)" do + check all(values <- list_of(integer(-1000..1000), min_length: 2, max_length: 50)) do + batch = s64_batch(values) + n = length(values) + # pick a deterministic subset: even indices + indices = for i <- 0..(n - 1), rem(i, 2) == 0, do: i + + assert {:ok, taken} = Batch.take(batch, indices) + assert s64_column(taken, "v") == Enum.take_every(values, 2) + end + end +end diff --git a/test/ex_arrow/batch_test.exs b/test/ex_arrow/batch_test.exs new file mode 100644 index 0000000..5f05e5e --- /dev/null +++ b/test/ex_arrow/batch_test.exs @@ -0,0 +1,236 @@ +defmodule ExArrow.BatchTest do + use ExUnit.Case, async: true + + alias ExArrow.Batch + alias ExArrow.RecordBatch + alias ExArrow.Schema + + # Build a batch with two columns: id (s64) and label (utf8). + defp two_col_batch(ids, labels) do + n = length(ids) + + id_bin = + ids + |> Enum.map(&<<&1::little-signed-64>>) + |> IO.iodata_to_binary() + + label_bin = + labels + |> Enum.map(fn s -> + bytes = s |> String.to_charlist() |> IO.iodata_to_binary() + <> + end) + |> IO.iodata_to_binary() + + {:ok, batch} = + RecordBatch.from_columns(["id", "label"], [id_bin, label_bin], ["s64", "utf8"], n) + + batch + end + + # Build a numeric-only batch with two s64 columns for rename tests. + defp two_s64_batch(a, b) do + n = length(a) + + bin_fn = fn list -> + list |> Enum.map(&<<&1::little-signed-64>>) |> IO.iodata_to_binary() + end + + {:ok, batch} = + RecordBatch.from_columns(["a", "b"], [bin_fn.(a), bin_fn.(b)], ["s64", "s64"], n) + + batch + end + + # Extract an s64 column as a list of integers. + defp s64_column(batch, name) do + ref = RecordBatch.resource_ref(batch) + {:ok, {binary, "s64", _n}} = ExArrow.Native.record_batch_column_buffer(ref, name) + + for <>, do: v + end + + defp column_names(batch), do: RecordBatch.column_names(batch) + + describe "schema/1" do + @tag :nif + test "returns the batch schema" do + batch = two_col_batch([1, 2, 3], ["a", "b", "c"]) + schema = Batch.schema(batch) + assert Schema.field_names(schema) == ["id", "label"] + end + end + + describe "select/2" do + @tag :nif + test "selects columns in the requested order" do + batch = two_col_batch([1, 2], ["a", "b"]) + assert {:ok, selected} = Batch.select(batch, ["label", "id"]) + assert column_names(selected) == ["label", "id"] + end + + @tag :nif + test "returns error for unknown column" do + batch = two_col_batch([1], ["a"]) + assert {:error, msg} = Batch.select(batch, ["missing"]) + assert msg =~ "not found" + end + + @tag :nif + test "preserves row count" do + batch = two_col_batch([10, 20, 30], ["x", "y", "z"]) + assert {:ok, selected} = Batch.select(batch, ["id"]) + assert RecordBatch.num_rows(selected) == 3 + assert s64_column(selected, "id") == [10, 20, 30] + end + end + + describe "drop/2" do + @tag :nif + test "removes the named columns and keeps the rest in order" do + batch = two_col_batch([1, 2], ["a", "b"]) + assert {:ok, rest} = Batch.drop(batch, ["label"]) + assert column_names(rest) == ["id"] + assert s64_column(rest, "id") == [1, 2] + end + + @tag :nif + test "dropping all columns yields an empty-schema batch" do + batch = two_col_batch([1], ["a"]) + assert {:ok, rest} = Batch.drop(batch, ["id", "label"]) + assert RecordBatch.num_columns(rest) == 0 + assert RecordBatch.num_rows(rest) == 1 + end + + @tag :nif + test "returns an error for an unknown column name" do + batch = two_col_batch([1], ["a"]) + assert {:error, msg} = Batch.drop(batch, ["no_such_column"]) + assert msg =~ "no_such_column" + end + + @tag :nif + test "returns an error when one of several columns is unknown" do + batch = two_col_batch([1], ["a"]) + assert {:error, msg} = Batch.drop(batch, ["id", "no_such_column"]) + assert msg =~ "no_such_column" + end + end + + describe "rename/2" do + @tag :nif + test "renames the supplied columns and preserves order and types" do + batch = two_s64_batch([1, 2], [10, 20]) + assert {:ok, renamed} = Batch.rename(batch, %{"a" => "x", "b" => "y"}) + assert column_names(renamed) == ["x", "y"] + assert s64_column(renamed, "x") == [1, 2] + assert s64_column(renamed, "y") == [10, 20] + end + + @tag :nif + test "leaves un-mapped columns unchanged" do + batch = two_s64_batch([1], [10]) + assert {:ok, renamed} = Batch.rename(batch, %{"a" => "x"}) + assert column_names(renamed) == ["x", "b"] + end + + @tag :nif + test "accepts a keyword list mapping" do + batch = two_s64_batch([1], [10]) + assert {:ok, renamed} = Batch.rename(batch, a: "x") + assert column_names(renamed) == ["x", "b"] + end + + @tag :nif + test "returns an error for an unknown source column" do + batch = two_s64_batch([1], [10]) + assert {:error, msg} = Batch.rename(batch, %{"missing" => "x"}) + assert msg =~ "unknown column" + end + + @tag :nif + test "preserves schema types after rename" do + batch = two_s64_batch([1, 2], [10, 20]) + assert {:ok, renamed} = Batch.rename(batch, %{"a" => "x"}) + fields = Schema.fields(RecordBatch.schema(renamed)) + assert Enum.map(fields, &{&1.name, &1.type}) == [{"x", :int64}, {"b", :int64}] + end + + @tag :nif + test "returns an error for unsupported (utf8) columns" do + batch = two_col_batch([1, 2], ["a", "b"]) + assert {:error, msg} = Batch.rename(batch, %{"id" => "user_id"}) + assert msg =~ "unsupported column type" + end + end + + describe "take/2 β€” integer" do + @tag :nif + test "returns the first n rows" do + batch = two_col_batch([10, 20, 30, 40], ["a", "b", "c", "d"]) + assert {:ok, taken} = Batch.take(batch, 2) + assert RecordBatch.num_rows(taken) == 2 + assert s64_column(taken, "id") == [10, 20] + end + + @tag :nif + test "n larger than rows returns the batch unchanged" do + batch = two_col_batch([1, 2], ["a", "b"]) + assert {:ok, taken} = Batch.take(batch, 10) + assert RecordBatch.num_rows(taken) == 2 + assert s64_column(taken, "id") == [1, 2] + end + + @tag :nif + test "n = 0 returns an empty batch with the same columns" do + batch = two_col_batch([1, 2], ["a", "b"]) + assert {:ok, taken} = Batch.take(batch, 0) + assert RecordBatch.num_rows(taken) == 0 + assert column_names(taken) == ["id", "label"] + end + + test "negative n returns an error" do + batch = two_col_batch([1], ["a"]) + assert {:error, _} = Batch.take(batch, -1) + end + end + + describe "take/2 β€” indices" do + @tag :nif + test "selects the rows at the given indices (in original row order)" do + batch = two_col_batch([10, 20, 30, 40], ["a", "b", "c", "d"]) + assert {:ok, taken} = Batch.take(batch, [2, 0, 3]) + assert RecordBatch.num_rows(taken) == 3 + # Rows 0, 2, 3 kept in original order: ids 10, 30, 40 + assert s64_column(taken, "id") == [10, 30, 40] + end + + @tag :nif + test "empty index list returns an empty batch" do + batch = two_col_batch([1, 2], ["a", "b"]) + assert {:ok, taken} = Batch.take(batch, []) + assert RecordBatch.num_rows(taken) == 0 + end + + @tag :nif + test "out-of-range index is an error" do + batch = two_col_batch([1, 2], ["a", "b"]) + assert {:error, msg} = Batch.take(batch, [0, 5]) + assert msg =~ "out of range" + end + end + + describe "filter/2" do + @tag :nif + test "delegates to Compute.filter/2 via a boolean predicate" do + batch = two_col_batch([10, 20, 30, 40], ["a", "b", "c", "d"]) + # mask: keep rows 0 and 2 (id 10, 30) + mask = <<1, 0, 1, 0>> + {:ok, mask_batch} = RecordBatch.from_columns(["m"], [mask], ["bool"], 4) + + assert {:ok, filtered} = Batch.filter(batch, mask_batch) + assert RecordBatch.num_rows(filtered) == 2 + assert s64_column(filtered, "id") == [10, 30] + end + end +end diff --git a/test/ex_arrow/broadway_test.exs b/test/ex_arrow/broadway_test.exs new file mode 100644 index 0000000..6bdc7ed --- /dev/null +++ b/test/ex_arrow/broadway_test.exs @@ -0,0 +1,176 @@ +defmodule ExArrow.BroadwayTest do + use ExUnit.Case, async: false + + import ExArrow.TestFixtures + alias ExArrow.Broadway.BatchBuilder + alias ExArrow.Broadway.FlightSink + alias ExArrow.Broadway.ParquetSink + alias ExArrow.RecordBatch + + # A minimal fake Broadway message struct. Broadway.Message is a struct with a + # :data field; we build one without starting a Broadway pipeline. + defp msg(data), do: %{__struct__: Broadway.Message, data: data} + + describe "BatchBuilder.extract_batches/1" do + @tag :nif + test "extracts RecordBatch handles from message data" do + b1 = s64_batch([1, 2]) + b2 = s64_batch([3, 4]) + + assert {:ok, [^b1, ^b2]} = BatchBuilder.extract_batches([msg(b1), msg(b2)]) + end + + @tag :nif + test "builds batches from {names, binaries, dtypes, length} data" do + data = {["id"], [<<1::little-signed-64, 2::little-signed-64>>], ["s64"], 2} + assert {:ok, [batch]} = BatchBuilder.extract_batches([msg(data)]) + assert RecordBatch.num_rows(batch) == 2 + end + + test "returns an error for unsupported message data" do + assert {:error, msg_text} = BatchBuilder.extract_batches([msg(:nope)]) + assert msg_text =~ "unsupported Broadway message data" + end + + test "returns an error for a non-message input" do + assert {:error, msg_text} = BatchBuilder.extract_batches([:not_a_message]) + assert msg_text =~ "expected a Broadway.Message" + end + end + + describe "BatchBuilder.from_messages/1" do + @tag :nif + test "returns the shared schema and batch list" do + b1 = s64_batch([1, 2]) + b2 = s64_batch([3, 4]) + + assert {:ok, schema, [^b1, ^b2]} = + BatchBuilder.from_messages([msg(b1), msg(b2)]) + + assert ExArrow.Schema.field_names(schema) == ["v"] + end + + test "returns an error for an empty message list" do + assert {:error, msg_text} = BatchBuilder.from_messages([]) + assert msg_text =~ "no batches" + end + end + + describe "ParquetSink.write/3" do + @tag :tmp_dir + @tag :nif + test "writes batches to a Parquet file", %{tmp_dir: dir} do + path = Path.join(dir, "out.parquet") + batch = s64_batch([10, 20, 30]) + + schema = RecordBatch.schema(batch) + assert :ok = ParquetSink.write(path, schema, [batch]) + assert File.exists?(path) + + # Round-trip: read it back. + {:ok, stream} = ExArrow.Stream.from_parquet(path) + rt = ExArrow.Stream.next(stream) + assert RecordBatch.num_rows(rt) == 3 + end + + @tag :tmp_dir + @tag :nif + test "emits [:ex_arrow, :parquet, :write] telemetry", %{tmp_dir: dir} do + path = Path.join(dir, "tel.parquet") + batch = s64_batch([1, 2]) + schema = RecordBatch.schema(batch) + + ref = make_ref() + + :telemetry.attach( + {:ex_arrow_bway_pq, ref}, + [:ex_arrow, :parquet, :write], + fn _event, measurements, metadata, config -> + send(config[:pid], {:pq_write, measurements, metadata}) + end, + %{pid: self()} + ) + + ParquetSink.write(path, schema, [batch]) + + # Filter for the event matching our destination path; other async tests + # may emit on the same event name concurrently. + assert_receive {:pq_write, measurements, %{destination: ^path} = metadata}, + 200 + + assert measurements[:rows] == 2 + assert measurements[:batch_count] == 1 + assert metadata[:source] == :broadway + + :telemetry.detach({:ex_arrow_bway_pq, ref}) + end + + @tag :tmp_dir + @tag :nif + test "returns error for a bad path", %{tmp_dir: dir} do + batch = s64_batch([1]) + schema = RecordBatch.schema(batch) + bad_path = Path.join(dir, "no_such_dir/out.parquet") + assert {:error, _} = ParquetSink.write(bad_path, schema, [batch]) + end + end + + describe "FlightSink.write/4" do + setup context do + prev = Application.get_env(:ex_arrow, :flight_client_impl) + Application.put_env(:ex_arrow, :flight_client_impl, ExArrow.Flight.ClientMock) + Mox.set_mox_from_context(context) + on_exit(fn -> Application.delete_env(:ex_arrow, :flight_client_impl) end) + :ok + end + + @tag :nif + test "calls Flight.Client.do_put/4 with schema and batches" do + batch = s64_batch([1, 2]) + schema = RecordBatch.schema(batch) + client = %ExArrow.Flight.Client{resource: make_ref()} + + ExArrow.Flight.ClientMock + |> Mox.expect(:do_put, fn ^client, _schema, _batches, _opts -> :ok end) + + assert :ok = FlightSink.write(client, schema, [batch], descriptor: {:cmd, "x"}) + end + + test "no-ops on an empty batch list" do + client = %ExArrow.Flight.Client{resource: make_ref()} + schema = ExArrow.Schema.from_ref(make_ref()) + assert :ok = FlightSink.write(client, schema, []) + end + + @tag :nif + test "emits [:ex_arrow, :flight, :query] telemetry" do + batch = s64_batch([1, 2, 3]) + schema = RecordBatch.schema(batch) + client = %ExArrow.Flight.Client{resource: make_ref()} + + ExArrow.Flight.ClientMock + |> Mox.expect(:do_put, fn _, _, _, _ -> :ok end) + + ref = make_ref() + + :telemetry.attach( + {:ex_arrow_bway_flight, ref}, + [:ex_arrow, :flight, :query], + fn _event, measurements, metadata, config -> + send(config[:pid], {:flight_write, measurements, metadata}) + end, + %{pid: self()} + ) + + FlightSink.write(client, schema, [batch], descriptor: {:cmd, "events"}) + + assert_receive {:flight_write, measurements, %{destination: {:cmd, "events"}} = metadata}, + 200 + + assert measurements[:rows] == 3 + assert metadata[:source] == :broadway + + :telemetry.detach({:ex_arrow_bway_flight, ref}) + end + end +end diff --git a/test/ex_arrow/flow_test.exs b/test/ex_arrow/flow_test.exs new file mode 100644 index 0000000..2cf8cbd --- /dev/null +++ b/test/ex_arrow/flow_test.exs @@ -0,0 +1,158 @@ +defmodule ExArrow.FlowTest do + use ExUnit.Case, async: true + + alias ExArrow.IPC + # NOTE: do not alias ExArrow.Flow as `Flow` here β€” bare `Flow` must resolve to + # the Flow library so Flow.map/Flow.partition/Flow.reduce/Flow.flat_map work. + alias ExArrow.Stream + + # Build a multi-batch IPC stream so Flow has several batches to process. + defp multi_batch_stream(num_batches) do + {:ok, fixture} = ExArrow.Native.ipc_test_fixture_binary() + {:ok, reader} = ExArrow.Native.ipc_reader_from_binary(fixture) + schema_ref = ExArrow.Native.ipc_stream_schema(reader) + {:ok, batch_ref} = ExArrow.Native.ipc_stream_next(reader) + + batch_refs = for _ <- 1..num_batches, do: batch_ref + {:ok, bin} = ExArrow.Native.ipc_writer_to_binary(schema_ref, batch_refs) + + {:ok, stream} = IPC.Reader.from_binary(bin) + stream + end + + describe "from_batches/1" do + @tag :nif + test "builds a Flow from an ExArrow.Stream and yields every batch" do + stream = multi_batch_stream(4) + + rows = + stream + |> ExArrow.Flow.from_batches() + |> Flow.map(&ExArrow.RecordBatch.num_rows/1) + |> Enum.to_list() + + assert length(rows) == 4 + assert Enum.all?(rows, &(&1 > 0)) + end + + @tag :nif + test "accepts {:ok, stream} and unwraps it" do + {:ok, stream} = ExArrow.IPC.Reader.from_binary(multi_batch_stream_binary(3)) + + counts = + {:ok, stream} + |> ExArrow.Flow.from_batches() + |> Flow.map(&ExArrow.RecordBatch.num_columns/1) + |> Enum.to_list() + + assert length(counts) == 3 + end + + @tag :nif + test "raises on {:error, _} input" do + assert_raise RuntimeError, ~r/error result/, fn -> + ExArrow.Flow.from_batches({:error, "boom"}) + end + end + + @tag :nif + test "supports Flow.partition/2 and Flow.reduce/3" do + stream = multi_batch_stream(5) + + row_counts = + stream + |> ExArrow.Flow.from_batches() + |> Flow.partition() + |> Flow.reduce(fn -> [] end, fn batch, acc -> + [ExArrow.RecordBatch.num_rows(batch) | acc] + end) + |> Enum.to_list() + + assert Enum.sum(row_counts) > 0 + end + + @tag :nif + test "supports Flow.flat_map/2" do + stream = multi_batch_stream(3) + + # Expand each batch into a one-element list, then collect. + result = + stream + |> ExArrow.Flow.from_batches() + |> Flow.flat_map(fn batch -> [ExArrow.RecordBatch.num_rows(batch)] end) + |> Enum.to_list() + + assert length(result) == 3 + end + end + + describe "map_batches/2" do + @tag :nif + test "maps over batches and returns transformed values" do + stream = multi_batch_stream(3) + + rows = + stream + |> ExArrow.Flow.from_batches() + |> ExArrow.Flow.map_batches(fn batch -> ExArrow.RecordBatch.num_rows(batch) end) + |> Enum.to_list() + + assert length(rows) == 3 + end + + @tag :nif + test "emits [:ex_arrow, :pipeline, :batch] telemetry per batch" do + stream = multi_batch_stream(2) + ref = make_ref() + + :telemetry.attach( + {:ex_arrow_flow, ref}, + [:ex_arrow, :pipeline, :batch], + fn _event, measurements, _meta, config -> + send(config[:pid], {:batch_event, measurements[:rows]}) + end, + %{pid: self()} + ) + + stream + |> ExArrow.Flow.from_batches() + |> ExArrow.Flow.map_batches(fn _batch -> :ok end) + |> Enum.to_list() + + assert_received {:batch_event, _rows} + assert_received {:batch_event, _rows} + + :telemetry.detach({:ex_arrow_flow, ref}) + end + end + + describe "each_batch/2" do + @tag :nif + test "runs a side effect and preserves the batches in the flow" do + stream = multi_batch_stream(2) + pid = self() + + batches = + stream + |> ExArrow.Flow.from_batches() + |> ExArrow.Flow.each_batch(fn _batch -> send(pid, :seen) end) + |> Enum.to_list() + + assert length(batches) == 2 + assert Enum.all?(batches, &ExArrow.RecordBatch.record_batch?/1) + assert_received :seen + assert_received :seen + end + end + + # Build the IPC binary for a multi-batch stream. + defp multi_batch_stream_binary(num_batches) do + {:ok, fixture} = ExArrow.Native.ipc_test_fixture_binary() + {:ok, reader} = ExArrow.Native.ipc_reader_from_binary(fixture) + schema_ref = ExArrow.Native.ipc_stream_schema(reader) + {:ok, batch_ref} = ExArrow.Native.ipc_stream_next(reader) + batch_refs = for _ <- 1..num_batches, do: batch_ref + {:ok, bin} = ExArrow.Native.ipc_writer_to_binary(schema_ref, batch_refs) + bin + end +end diff --git a/test/ex_arrow/gen_stage_test.exs b/test/ex_arrow/gen_stage_test.exs new file mode 100644 index 0000000..8de2167 --- /dev/null +++ b/test/ex_arrow/gen_stage_test.exs @@ -0,0 +1,248 @@ +defmodule ExArrow.GenStageTest do + use ExUnit.Case, async: true + + import ExArrow.TestFixtures + alias ExArrow.GenStage.ADBCProducer + alias ExArrow.GenStage.FlightProducer + alias ExArrow.GenStage.ParquetProducer + + # A single-batch Parquet binary (Parquet merges input batches into one row + # group, so multi-batch input still yields one batch on read). + defp parquet_binary do + {:ok, fixture} = ExArrow.Native.ipc_test_fixture_binary() + {:ok, reader} = ExArrow.Native.ipc_reader_from_binary(fixture) + schema_ref = ExArrow.Native.ipc_stream_schema(reader) + {:ok, batch_ref} = ExArrow.Native.ipc_stream_next(reader) + {:ok, pq} = ExArrow.Native.parquet_writer_to_binary(schema_ref, [batch_ref]) + pq + end + + # A consumer that demands batches and forwards them to the test process. + defmodule Collector do + use GenStage + + def start_link(parent) do + GenStage.start_link(__MODULE__, parent) + end + + def init(parent), do: {:consumer, parent} + + def handle_events(batches, _from, parent) do + send(parent, {:batches, batches}) + {:noreply, [], parent} + end + end + + # A producer-consumer that transforms each batch (selects the first column). + defmodule FirstColumnTransformer do + use GenStage + + def start_link(parent) do + GenStage.start_link(__MODULE__, parent) + end + + def init(parent), do: {:producer_consumer, parent} + + def handle_events(batches, _from, parent) do + transformed = + Enum.map(batches, fn batch -> + [name | _] = ExArrow.RecordBatch.column_names(batch) + {:ok, slim} = ExArrow.Batch.select(batch, [name]) + slim + end) + + {:noreply, transformed, parent} + end + end + + # Collect all batches emitted as separate {:batches, [...]} messages. + defp collect_all(timeout_ms \\ 1000) do + collect_all([], timeout_ms) + end + + defp collect_all(acc, timeout_ms) do + receive do + {:batches, batches} -> collect_all(acc ++ batches, timeout_ms) + after + timeout_ms -> acc + end + end + + defp collect_n(parent, n) do + collect_n(parent, n, []) + end + + defp collect_n(_parent, 0, acc), do: List.flatten(Enum.reverse(acc)) + + defp collect_n(parent, n, acc) do + receive do + {:batches, batches} -> collect_n(parent, n - length(batches), [batches | acc]) + after + 1000 -> flunk("timed out waiting for batches (got #{length(List.flatten(acc))})") + end + end + + # # ── ParquetProducer ────────────────────────────────────────────────────────── + + describe "ParquetProducer" do + @tag :nif + test "emits the Parquet batch on demand and stops when exhausted" do + {:ok, producer} = ParquetProducer.start_link(binary: parquet_binary()) + + {:ok, consumer} = Collector.start_link(self()) + GenStage.sync_subscribe(consumer, to: producer, max_demand: 10) + + batches = collect_all() + assert length(batches) == 1 + assert ExArrow.RecordBatch.record_batch?(hd(batches)) + + # Producer stops on its own once exhausted. + ref = Process.monitor(producer) + assert_receive {:DOWN, ^ref, :process, ^producer, _reason}, 2000 + end + + @tag :nif + test "emits [:ex_arrow, :stream, :batch] telemetry per batch" do + {:ok, producer} = ParquetProducer.start_link(binary: parquet_binary()) + + telem_ref = make_ref() + + :telemetry.attach( + {:ex_arrow_gs, telem_ref}, + [:ex_arrow, :stream, :batch], + fn _event, _measurements, metadata, config -> + send(config[:pid], {:telem, metadata[:source]}) + end, + %{pid: self()} + ) + + {:ok, consumer} = Collector.start_link(self()) + GenStage.sync_subscribe(consumer, to: producer, max_demand: 10) + + collect_all() + assert_received {:telem, {:parquet, :binary}}, 200 + + :telemetry.detach({:ex_arrow_gs, telem_ref}) + end + + @tag :nif + test "supports a pre-opened :stream option (multi-batch IPC)" do + {:ok, producer} = ParquetProducer.start_link(stream: ipc_stream(3)) + + {:ok, consumer} = Collector.start_link(self()) + GenStage.sync_subscribe(consumer, to: producer, max_demand: 1) + + batches = collect_n(self(), 3) + assert length(batches) == 3 + end + + test "returns error when no source option is given" do + assert {:error, msg} = ParquetProducer.start_link([]) + assert msg =~ "requires one of" + end + + @tag :nif + test ":stream option with nil source falls back to :parquet label" do + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_binary(1)) + # Force source to nil by constructing a stream without source metadata + nil_source_stream = %{stream | source: nil} + {:ok, producer} = ParquetProducer.start_link(stream: nil_source_stream) + + telem_ref = make_ref() + + :telemetry.attach( + {:ex_arrow_gs_nil, telem_ref}, + [:ex_arrow, :stream, :batch], + fn _event, _measurements, metadata, config -> + send(config[:pid], {:nil_source_telem, metadata[:source]}) + end, + %{pid: self()} + ) + + {:ok, consumer} = Collector.start_link(self()) + GenStage.sync_subscribe(consumer, to: producer, max_demand: 10) + collect_all() + + assert_received {:nil_source_telem, :parquet} + + :telemetry.detach({:ex_arrow_gs_nil, telem_ref}) + end + end + + # # ── FlightProducer ─────────────────────────────────────────────────────────── + + describe "FlightProducer" do + @tag :nif + test "works with any pre-opened stream via :stream option" do + {:ok, producer} = FlightProducer.start_link(stream: ipc_stream(2)) + + {:ok, consumer} = Collector.start_link(self()) + GenStage.sync_subscribe(consumer, to: producer, max_demand: 1) + + batches = collect_n(self(), 2) + assert length(batches) == 2 + end + + test "returns error when missing :client and :ticket" do + assert {:error, msg} = FlightProducer.start_link([]) + assert msg =~ "requires" + end + end + + # # ── ADBCProducer ───────────────────────────────────────────────────────────── + + describe "ADBCProducer" do + @tag :nif + test "works with any pre-opened stream via :stream option" do + {:ok, producer} = ADBCProducer.start_link(stream: ipc_stream(2)) + + {:ok, consumer} = Collector.start_link(self()) + GenStage.sync_subscribe(consumer, to: producer, max_demand: 1) + + batches = collect_n(self(), 2) + assert length(batches) == 2 + end + + test "returns error when no source option is given" do + assert {:error, msg} = ADBCProducer.start_link([]) + assert msg =~ "requires" + end + end + + # # ── Producer-consumer pattern ──────────────────────────────────────────────── + + describe "producer-consumer pattern" do + @tag :nif + test "a producer-consumer transforms batches between producer and consumer" do + {:ok, producer} = ParquetProducer.start_link(stream: ipc_stream(2)) + {:ok, transformer} = FirstColumnTransformer.start_link(self()) + {:ok, consumer} = Collector.start_link(self()) + + GenStage.sync_subscribe(transformer, to: producer, max_demand: 1) + GenStage.sync_subscribe(consumer, to: transformer, max_demand: 1) + + batches = collect_n(self(), 2) + + # Each transformed batch has exactly one column. + assert Enum.all?(batches, fn b -> ExArrow.RecordBatch.num_columns(b) == 1 end) + end + end + + describe "dispatch/2 state preservation" do + test "done: true branch preserves source and other state fields" do + alias ExArrow.GenStage.State + + state = %State{ + stream: nil, + source: {:parquet, "/data/test.parquet"}, + demand: 5, + done: true + } + + {:noreply, [], returned_state} = ExArrow.GenStage.dispatch(state, 10) + + assert returned_state.source == {:parquet, "/data/test.parquet"} + assert returned_state.done == true + end + end +end diff --git a/test/ex_arrow/pipeline_test.exs b/test/ex_arrow/pipeline_test.exs new file mode 100644 index 0000000..8412d46 --- /dev/null +++ b/test/ex_arrow/pipeline_test.exs @@ -0,0 +1,197 @@ +defmodule ExArrow.PipelineTest do + use ExUnit.Case, async: false + + import ExArrow.TestFixtures + alias ExArrow.Batch + alias ExArrow.Pipeline + alias ExArrow.RecordBatch + + describe "map_batches/2" do + @tag :nif + test "applies a transformation lazily and writes the result" do + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_binary(3)) + path = tmp_path("map.parquet") + + result = + {:ok, stream} + |> Pipeline.map_batches(fn batch -> + {:ok, slim} = Batch.select(batch, ["id"]) + slim + end) + |> Pipeline.write_parquet(path) + + assert :ok = result + assert File.exists?(path) + + {:ok, rt} = ExArrow.Stream.from_parquet(path) + batches = ExArrow.Stream.to_list(rt) + assert Enum.all?(batches, fn b -> RecordBatch.num_columns(b) == 1 end) + + File.rm(path) + end + + @tag :nif + test "threads errors from the stream constructor" do + assert {:error, _} = + {:error, "no connection"} + |> Pipeline.map_batches(fn b -> b end) + end + + @tag :nif + test "rejects a fun that returns {:ok, batch} instead of a bare batch" do + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_binary(2)) + path = tmp_path("bad_fun.parquet") + + result = + {:ok, stream} + |> Pipeline.map_batches(fn b -> ExArrow.Batch.select(b, ["id"]) end) + |> Pipeline.write_parquet(path) + + assert {:error, msg} = result + assert msg =~ "non-batch" + File.rm(path) + end + + @tag :nif + test "emits [:ex_arrow, :pipeline, :batch] telemetry per batch" do + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_binary(2)) + path = tmp_path("tel.parquet") + ref = make_ref() + + :telemetry.attach( + {:ex_arrow_pl, ref}, + [:ex_arrow, :pipeline, :batch], + fn _event, _measurements, metadata, config -> + send(config[:pid], {:pl_batch, metadata}) + end, + %{pid: self()} + ) + + {:ok, stream} + |> Pipeline.map_batches(fn b -> b end) + |> Pipeline.write_parquet(path) + + assert_received {:pl_batch, %{source: :pipeline}}, 200 + assert_received {:pl_batch, %{source: :pipeline}}, 200 + + :telemetry.detach({:ex_arrow_pl, ref}) + File.rm(path) + end + end + + describe "each_batch/2" do + @tag :nif + test "runs a side effect and preserves batches" do + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_binary(2)) + path = tmp_path("each.parquet") + pid = self() + + :ok = + {:ok, stream} + |> Pipeline.each_batch(fn _batch -> send(pid, :saw_batch) end) + |> Pipeline.write_parquet(path) + + assert_received :saw_batch + assert_received :saw_batch + File.rm(path) + end + end + + describe "write_parquet/2" do + @tag :nif + test "writes a stream with no transformations" do + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_binary(2)) + path = tmp_path("raw.parquet") + + assert :ok = Pipeline.write_parquet({:ok, stream}, path) + assert File.exists?(path) + File.rm(path) + end + + @tag :nif + test "threads errors" do + assert {:error, _} = Pipeline.write_parquet({:error, "boom"}, "/tmp/x.parquet") + end + + @tag :nif + test "emits [:ex_arrow, :parquet, :write] telemetry" do + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_binary(1)) + path = tmp_path("wt.parquet") + ref = make_ref() + + :telemetry.attach( + {:ex_arrow_pl_pq, ref}, + [:ex_arrow, :parquet, :write], + fn _event, _measurements, metadata, config -> + send(config[:pid], {:pl_pq, metadata}) + end, + %{pid: self()} + ) + + Pipeline.write_parquet({:ok, stream}, path) + + assert_received {:pl_pq, %{destination: ^path, source: :pipeline}} + :telemetry.detach({:ex_arrow_pl_pq, ref}) + File.rm(path) + end + end + + describe "write_flight/3" do + setup context do + prev = Application.get_env(:ex_arrow, :flight_client_impl) + Application.put_env(:ex_arrow, :flight_client_impl, ExArrow.Flight.ClientMock) + Mox.set_mox_from_context(context) + on_exit(fn -> Application.delete_env(:ex_arrow, :flight_client_impl) end) + :ok + end + + @tag :nif + test "uploads the pipeline batches to Flight" do + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_binary(2)) + client = %ExArrow.Flight.Client{resource: make_ref()} + + ExArrow.Flight.ClientMock + |> Mox.expect(:do_put, fn ^client, _schema, batches, _opts -> + # batches may be transformed or raw; here they are raw + send(self(), {:uploaded, length(batches)}) + :ok + end) + + assert :ok = Pipeline.write_flight({:ok, stream}, client, descriptor: {:cmd, "x"}) + assert_received {:uploaded, 2} + end + + @tag :nif + test "threads errors" do + client = %ExArrow.Flight.Client{resource: make_ref()} + assert {:error, _} = Pipeline.write_flight({:error, "boom"}, client) + end + end + + describe "write_dataframe/1" do + @tag :nif + @tag :explorer + test "converts a pipeline to an Explorer DataFrame" do + batch = s64_batch([1, 2, 3]) + # Build a one-batch pipeline via a stream round-trip. + schema = RecordBatch.schema(batch) + {:ok, ipc_bin} = ExArrow.IPC.Writer.to_binary(schema, [batch]) + {:ok, stream} = ExArrow.Stream.from_ipc(ipc_bin) + + assert {:ok, df} = Pipeline.write_dataframe({:ok, stream}) + assert Explorer.DataFrame.n_rows(df) == 3 + end + + @tag :nif + test "threads errors" do + assert {:error, _} = Pipeline.write_dataframe({:error, "boom"}) + end + end + + defp tmp_path(name) do + Path.join( + System.tmp_dir!(), + "ex_arrow_pipeline_#{System.unique_integer([:positive])}_#{name}" + ) + end +end diff --git a/test/ex_arrow/sink_test.exs b/test/ex_arrow/sink_test.exs new file mode 100644 index 0000000..22abb10 --- /dev/null +++ b/test/ex_arrow/sink_test.exs @@ -0,0 +1,144 @@ +defmodule ExArrow.SinkTest do + use ExUnit.Case, async: false + + import ExArrow.TestFixtures + alias ExArrow.RecordBatch + alias ExArrow.Sink.DataFrame + alias ExArrow.Sink.Flight + alias ExArrow.Sink.Parquet + # NOTE: do not alias ExArrow.Sink.Nx as `Nx` β€” bare `Nx` must resolve to the + # Nx library so Nx.shape/Nx.to_number work in assertions. + + describe "Parquet.write/2" do + @tag :tmp_dir + @tag :nif + test "writes a stream to a Parquet file", %{tmp_dir: dir} do + path = Path.join(dir, "s.parquet") + stream = ipc_stream(2) + + assert :ok = Parquet.write(stream, path) + assert File.exists?(path) + + {:ok, rt_stream} = ExArrow.Stream.from_parquet(path) + batches = ExArrow.Stream.to_list(rt_stream) + assert length(batches) >= 1 + end + + @tag :tmp_dir + @tag :nif + test "writes a {schema, batches} tuple", %{tmp_dir: dir} do + path = Path.join(dir, "tuple.parquet") + batch = s64_batch([1, 2, 3]) + schema = RecordBatch.schema(batch) + + assert :ok = Parquet.write({schema, [batch]}, path) + assert File.exists?(path) + end + + @tag :tmp_dir + @tag :nif + test "writes a bare batch list", %{tmp_dir: dir} do + path = Path.join(dir, "list.parquet") + batch = s64_batch([10, 20]) + + assert :ok = Parquet.write([batch], path) + assert File.exists?(path) + end + + test "returns an error for an empty batch list" do + assert {:error, msg} = Parquet.write([], "/tmp/x.parquet") + assert msg =~ "empty batch list" + end + + @tag :tmp_dir + @tag :nif + test "emits [:ex_arrow, :parquet, :write] telemetry", %{tmp_dir: dir} do + path = Path.join(dir, "tel.parquet") + batch = s64_batch([1, 2]) + ref = make_ref() + + :telemetry.attach( + {:ex_arrow_sink_pq, ref}, + [:ex_arrow, :parquet, :write], + fn _event, _measurements, metadata, config -> + send(config[:pid], {:sink_pq, metadata}) + end, + %{pid: self()} + ) + + Parquet.write([batch], path) + + assert_received {:sink_pq, %{destination: ^path, source: :sink}} + :telemetry.detach({:ex_arrow_sink_pq, ref}) + end + + test "returns error for {schema, batches} where batches is not a list" do + schema = ExArrow.Schema.from_ref(make_ref()) + assert {:error, msg} = Parquet.write({schema, :not_a_list}, "/tmp/x.parquet") + assert msg =~ "expected {schema, [batches]}" + end + + test "returns error for unsupported source type" do + assert {:error, msg} = Parquet.write(:atom_source, "/tmp/x.parquet") + assert msg =~ "unsupported source" + end + end + + describe "Flight.write/3" do + setup context do + prev = Application.get_env(:ex_arrow, :flight_client_impl) + Application.put_env(:ex_arrow, :flight_client_impl, ExArrow.Flight.ClientMock) + Mox.set_mox_from_context(context) + on_exit(fn -> Application.delete_env(:ex_arrow, :flight_client_impl) end) + :ok + end + + @tag :nif + test "uploads a stream via Flight.Client.do_put/4" do + stream = ipc_stream(2) + client = %ExArrow.Flight.Client{resource: make_ref()} + + ExArrow.Flight.ClientMock + |> Mox.expect(:do_put, fn ^client, _schema, _batches, _opts -> :ok end) + + assert :ok = Flight.write(stream, client, descriptor: {:cmd, "x"}) + end + + @tag :nif + test "uploads a {schema, batches} tuple" do + batch = s64_batch([1, 2]) + schema = RecordBatch.schema(batch) + client = %ExArrow.Flight.Client{resource: make_ref()} + + ExArrow.Flight.ClientMock + |> Mox.expect(:do_put, fn ^client, ^schema, [^batch], descriptor: {:cmd, "y"} -> :ok end) + + assert :ok = Flight.write({schema, [batch]}, client, descriptor: {:cmd, "y"}) + end + + test "no-ops on an empty batch list" do + client = %ExArrow.Flight.Client{resource: make_ref()} + assert :ok = Flight.write([], client) + end + end + + describe "DataFrame.write/1" do + @tag :nif + @tag :explorer + test "converts a batch to an Explorer DataFrame" do + batch = s64_batch([1, 2, 3]) + assert {:ok, df} = DataFrame.write(batch) + assert Explorer.DataFrame.n_rows(df) == 3 + end + end + + describe "Nx.write/1" do + @tag :nif + @tag :nx + test "converts a batch to an Nx tensor" do + batch = s64_batch([10, 20, 30]) + assert {:ok, tensor} = ExArrow.Sink.Nx.write(batch) + assert Nx.shape(tensor) == {3} + end + end +end diff --git a/test/ex_arrow/stream_test.exs b/test/ex_arrow/stream_test.exs index b25a801..7f2d578 100644 --- a/test/ex_arrow/stream_test.exs +++ b/test/ex_arrow/stream_test.exs @@ -284,3 +284,226 @@ defmodule ExArrow.StreamTest do stream_ref end end + +defmodule ExArrow.StreamConstructorsTest do + use ExUnit.Case, async: false + + alias ExArrow.Schema + alias ExArrow.Stream + + # Build a small IPC binary once per test module. + defp ipc_binary do + {:ok, bin} = ExArrow.Native.ipc_test_fixture_binary() + bin + end + + defp parquet_binary do + {:ok, reader} = ExArrow.Native.ipc_reader_from_binary(ipc_binary()) + schema_ref = ExArrow.Native.ipc_stream_schema(reader) + {:ok, batch_ref} = ExArrow.Native.ipc_stream_next(reader) + {:ok, parquet_bin} = ExArrow.Native.parquet_writer_to_binary(schema_ref, [batch_ref]) + parquet_bin + end + + describe "from_ipc/1 and from_ipc_file/1" do + @tag :nif + test "from_ipc/1 returns a stream tagged with {:ipc, :binary}" do + assert {:ok, %Stream{backend: :ipc} = stream} = Stream.from_ipc(ipc_binary()) + assert Stream.source(stream) == {:ipc, :binary} + assert {:ok, %ExArrow.Schema{}} = Stream.schema(stream) + end + + @tag :tmp_dir + test "from_ipc_file/1 tags the stream with the path", %{tmp_dir: dir} do + path = Path.join(dir, "ipc.arrows") + File.write!(path, ipc_binary()) + + assert {:ok, %Stream{backend: :ipc} = stream} = Stream.from_ipc_file(path) + assert Stream.source(stream) == {:ipc, path} + end + + test "from_ipc/1 returns error for non-IPC binary" do + assert {:error, _} = Stream.from_ipc(<<"not ipc">>) + end + end + + describe "from_parquet/1 and from_parquet_binary/1" do + @tag :nif + test "from_parquet_binary/1 returns a stream tagged with {:parquet, :binary}" do + assert {:ok, %Stream{backend: :parquet} = stream} = + Stream.from_parquet_binary(parquet_binary()) + + assert Stream.source(stream) == {:parquet, :binary} + assert {:ok, %ExArrow.Schema{}} = Stream.schema(stream) + end + + @tag :tmp_dir + test "from_parquet/1 tags the stream with the path", %{tmp_dir: dir} do + path = Path.join(dir, "events.parquet") + File.write!(path, parquet_binary()) + + assert {:ok, %Stream{backend: :parquet} = stream} = Stream.from_parquet(path) + assert Stream.source(stream) == {:parquet, path} + assert {:ok, %ExArrow.Schema{}} = Stream.schema(stream) + end + + test "from_parquet/1 returns error for missing file" do + assert {:error, _} = + Stream.from_parquet( + "/tmp/ex_arrow_missing_#{:erlang.unique_integer([:positive])}.parquet" + ) + end + end + + describe "from_flight_sql/2 (delegation)" do + setup context do + prev = Application.get_env(:ex_arrow, :flight_sql_client_impl) + Application.put_env(:ex_arrow, :flight_sql_client_impl, ExArrow.FlightSQL.ClientMock) + Mox.set_mox_from_context(context) + on_exit(fn -> Application.delete_env(:ex_arrow, :flight_sql_client_impl) end) + :ok + end + + test "delegates to FlightSQL.Client.stream_query/2 and tags source" do + ExArrow.FlightSQL.ClientMock + |> Mox.expect(:query, fn _client, _sql, _opts -> + {:ok, %Stream{resource: make_ref(), backend: :flight_sql}} + end) + + client = %ExArrow.FlightSQL.Client{resource: make_ref()} + + assert {:ok, %Stream{backend: :flight_sql} = stream} = + Stream.from_flight_sql(client, "SELECT 1") + + assert Stream.source(stream) == {:flight_sql, "SELECT 1"} + end + + test "passes errors through unchanged" do + ExArrow.FlightSQL.ClientMock + |> Mox.expect(:query, fn _client, _sql, _opts -> + {:error, %ExArrow.FlightSQL.Error{code: :unavailable}} + end) + + client = %ExArrow.FlightSQL.Client{resource: make_ref()} + + assert {:error, %ExArrow.FlightSQL.Error{code: :unavailable}} = + Stream.from_flight_sql(client, "SELECT 1") + end + end + + describe "from_flight/2 (delegation)" do + setup context do + prev = Application.get_env(:ex_arrow, :flight_client_impl) + Application.put_env(:ex_arrow, :flight_client_impl, ExArrow.Flight.ClientMock) + Mox.set_mox_from_context(context) + on_exit(fn -> Application.delete_env(:ex_arrow, :flight_client_impl) end) + :ok + end + + test "delegates to Flight.Client.do_get/2 and tags source" do + ExArrow.Flight.ClientMock + |> Mox.expect(:do_get, fn _client, _ticket -> + {:ok, %Stream{resource: make_ref(), backend: :ipc}} + end) + + client = %ExArrow.Flight.Client{resource: make_ref()} + assert {:ok, %Stream{} = stream} = Stream.from_flight(client, "sales_2024") + assert Stream.source(stream) == {:flight, "sales_2024"} + end + end + + describe "from_adbc/1 and from_adbc/2 (delegation)" do + setup context do + prev = Application.get_env(:ex_arrow, :adbc_statement_impl) + Application.put_env(:ex_arrow, :adbc_statement_impl, ExArrow.ADBC.StatementMock) + Mox.set_mox_from_context(context) + on_exit(fn -> Application.delete_env(:ex_arrow, :adbc_statement_impl) end) + :ok + end + + test "from_adbc/1 executes a pre-built statement" do + ExArrow.ADBC.StatementMock + |> Mox.expect(:execute, fn _stmt -> {:ok, %Stream{resource: make_ref(), backend: :adbc}} end) + + stmt = %ExArrow.ADBC.Statement{resource: make_ref()} + assert {:ok, %Stream{backend: :adbc} = stream} = Stream.from_adbc(stmt) + assert Stream.source(stream) == {:adbc, :statement} + end + + test "from_adbc/2 builds and executes a one-shot statement" do + conn = %ExArrow.ADBC.Connection{resource: make_ref()} + stmt = %ExArrow.ADBC.Statement{resource: make_ref()} + + ExArrow.ADBC.StatementMock + |> Mox.expect(:new, fn ^conn -> {:ok, stmt} end) + |> Mox.expect(:set_sql, fn ^stmt, "SELECT 1" -> :ok end) + |> Mox.expect(:execute, fn ^stmt -> {:ok, %Stream{resource: make_ref(), backend: :adbc}} end) + + assert {:ok, %Stream{backend: :adbc} = stream} = Stream.from_adbc(conn, "SELECT 1") + assert Stream.source(stream) == {:adbc, "SELECT 1"} + end + + test "from_adbc/2 propagates statement-creation errors" do + conn = %ExArrow.ADBC.Connection{resource: make_ref()} + + ExArrow.ADBC.StatementMock + |> Mox.expect(:new, fn ^conn -> {:error, "no connection"} end) + + assert {:error, "no connection"} = Stream.from_adbc(conn, "SELECT 1") + end + end + + describe "telemetry on stream open" do + @tag :nif + test "from_parquet_binary/1 emits [:ex_arrow, :parquet, :read]" do + ref = make_ref() + + :telemetry.attach( + {:ex_arrow_parquet_read, ref}, + [:ex_arrow, :parquet, :read], + fn _event, _measurements, metadata, config -> + send(config[:pid], {:parquet_read, metadata}) + end, + %{pid: self()} + ) + + Stream.from_parquet_binary(parquet_binary()) + + assert_received {:parquet_read, %{source: :binary}} + + :telemetry.detach({:ex_arrow_parquet_read, ref}) + end + + test "from_parquet/1 does NOT emit telemetry on error" do + ref = make_ref() + + :telemetry.attach( + {:ex_arrow_parquet_read_err, ref}, + [:ex_arrow, :parquet, :read], + fn _event, _measurements, _metadata, config -> + send(config[:pid], :parquet_read_event) + end, + %{pid: self()} + ) + + Stream.from_parquet("/nonexistent_#{:erlang.unique_integer([:positive])}.parquet") + + refute_received :parquet_read_event, 100 + + :telemetry.detach({:ex_arrow_parquet_read_err, ref}) + end + end + + describe "schema preservation through constructors" do + @tag :nif + test "from_ipc/1 round-trips field names" do + {:ok, orig_stream} = ExArrow.IPC.Reader.from_binary(ipc_binary()) + {:ok, orig_schema} = Stream.schema(orig_stream) + expected_names = Schema.field_names(orig_schema) + + assert {:ok, stream} = Stream.from_ipc(ipc_binary()) + assert {:ok, schema} = Stream.schema(stream) + assert Schema.field_names(schema) == expected_names + end + end +end diff --git a/test/ex_arrow/telemetry_test.exs b/test/ex_arrow/telemetry_test.exs new file mode 100644 index 0000000..4f0989b --- /dev/null +++ b/test/ex_arrow/telemetry_test.exs @@ -0,0 +1,150 @@ +defmodule ExArrow.TelemetryTest do + use ExUnit.Case, async: true + + alias ExArrow.RecordBatch + alias ExArrow.Telemetry + + @telemetry_available Code.ensure_loaded?(:telemetry) + + setup context do + # Only attach a handler when the :telemetry application is available and the + # describe block declared an event to observe via @describetag event: ... + event = Map.get(context, :event) + + if @telemetry_available and event != nil do + ref = make_ref() + handler_id = {:ex_arrow_test, context.test, ref} + + :telemetry.attach( + handler_id, + event, + fn event, measurements, metadata, config -> + send(config[:test_pid], {:telemetry, event, measurements, metadata}) + end, + %{test_pid: self()} + ) + + on_exit(fn -> :telemetry.detach(handler_id) end) + {:ok, ref: ref} + else + {:ok, ref: nil} + end + end + + describe "execute/3 β€” [:ex_arrow, :stream, :batch]" do + @describetag event: [:ex_arrow, :stream, :batch] + + test "delivers measurements and metadata to attached handlers", %{ref: ref} do + Telemetry.execute( + [:ex_arrow, :stream, :batch], + %{rows: 10, columns: 3, batch_count: 1}, + %{source: {:parquet, "/tmp/x.parquet"}, schema: nil} + ) + + if @telemetry_available do + assert_received {:telemetry, [:ex_arrow, :stream, :batch], measurements, metadata} + assert measurements[:rows] == 10 + assert measurements[:columns] == 3 + assert metadata[:source] == {:parquet, "/tmp/x.parquet"} + end + + _ = ref + end + end + + describe "execute/3 β€” [:ex_arrow, :parquet, :read]" do + @describetag event: [:ex_arrow, :parquet, :read] + + test "emits with source metadata" do + Telemetry.execute([:ex_arrow, :parquet, :read], %{}, %{source: "/data/events.parquet"}) + + if @telemetry_available do + assert_received {:telemetry, [:ex_arrow, :parquet, :read], _measurements, metadata} + assert metadata[:source] == "/data/events.parquet" + end + end + end + + describe "batch_measurements/2" do + @tag :nif + test "reports rows, columns, and batch_count for a real batch" do + {:ok, ipc_bin} = ExArrow.Native.ipc_test_fixture_binary() + {:ok, stream} = ExArrow.IPC.Reader.from_binary(ipc_bin) + batch = ExArrow.Stream.next(stream) + + m = Telemetry.batch_measurements(batch) + + assert is_integer(m[:rows]) and m[:rows] > 0 + assert is_integer(m[:columns]) and m[:columns] > 0 + assert m[:batch_count] == 1 + end + + @tag :nif + test "merges extra measurements" do + {:ok, ipc_bin} = ExArrow.Native.ipc_test_fixture_binary() + {:ok, stream} = ExArrow.IPC.Reader.from_binary(ipc_bin) + batch = ExArrow.Stream.next(stream) + + m = Telemetry.batch_measurements(batch, duration: 123, bytes: 4096) + assert m[:duration] == 123 + assert m[:bytes] == 4096 + assert m[:batch_count] == 1 + end + + test "handles a manually-built batch" do + {:ok, batch} = + RecordBatch.from_columns( + ["id"], + [<<1::little-signed-64, 2::little-signed-64>>], + ["s64"], + 2 + ) + + m = Telemetry.batch_measurements(batch) + assert m[:rows] == 2 + assert m[:columns] == 1 + assert m[:batch_count] == 1 + end + end + + describe "span/3" do + @describetag event: [:ex_arrow, :pipeline, :batch, :stop] + + test "returns the result of the wrapped function" do + result = + Telemetry.span([:ex_arrow, :pipeline, :batch], %{source: :test}, fn -> + {42, %{}} + end) + + assert result == 42 + end + + @tag :nif + test "emits start and stop events" do + ref = make_ref() + + :telemetry.attach_many( + {:ex_arrow_span, ref}, + [ + [:ex_arrow, :pipeline, :batch, :start], + [:ex_arrow, :pipeline, :batch, :stop] + ], + fn event, _m, _meta, config -> + send(config[:test_pid], {:span_event, event}) + end, + %{test_pid: self()} + ) + + Telemetry.span([:ex_arrow, :pipeline, :batch], %{source: :test}, fn -> + {:ok, %{rows: 1}} + end) + + if @telemetry_available do + assert_received {:span_event, [:ex_arrow, :pipeline, :batch, :start]} + assert_received {:span_event, [:ex_arrow, :pipeline, :batch, :stop]} + end + + :telemetry.detach({:ex_arrow_span, ref}) + end + end +end diff --git a/test/support/test_fixtures.ex b/test/support/test_fixtures.ex new file mode 100644 index 0000000..b9d5aec --- /dev/null +++ b/test/support/test_fixtures.ex @@ -0,0 +1,46 @@ +defmodule ExArrow.TestFixtures do + @moduledoc false + + alias ExArrow.IPC.Reader, as: IPCReader + + # Shared test fixtures used across pipeline, gen_stage, sink, and flow tests. + + # Build a multi-batch IPC stream from the built-in NIF fixture. + @spec ipc_stream(pos_integer()) :: ExArrow.Stream.t() + def ipc_stream(num_batches) do + {:ok, fixture} = ExArrow.Native.ipc_test_fixture_binary() + {:ok, reader} = ExArrow.Native.ipc_reader_from_binary(fixture) + schema_ref = ExArrow.Native.ipc_stream_schema(reader) + {:ok, batch_ref} = ExArrow.Native.ipc_stream_next(reader) + batch_refs = for _ <- 1..num_batches, do: batch_ref + {:ok, ipc_bin} = ExArrow.Native.ipc_writer_to_binary(schema_ref, batch_refs) + {:ok, stream} = IPCReader.from_binary(ipc_bin) + stream + end + + # Build a multi-batch IPC binary (for tests that need the raw bytes). + @spec ipc_binary(pos_integer()) :: binary() + def ipc_binary(num_batches) do + {:ok, fixture} = ExArrow.Native.ipc_test_fixture_binary() + {:ok, reader} = ExArrow.Native.ipc_reader_from_binary(fixture) + schema_ref = ExArrow.Native.ipc_stream_schema(reader) + {:ok, batch_ref} = ExArrow.Native.ipc_stream_next(reader) + batch_refs = for _ <- 1..num_batches, do: batch_ref + {:ok, bin} = ExArrow.Native.ipc_writer_to_binary(schema_ref, batch_refs) + bin + end + + # Build an s64 batch with `values` in a column named `name`. + @spec s64_batch([integer()]) :: ExArrow.RecordBatch.t() + def s64_batch(values, name \\ "v") do + n = length(values) + + bin = + values + |> Enum.map(&<<&1::little-signed-64>>) + |> IO.iodata_to_binary() + + {:ok, batch} = ExArrow.RecordBatch.from_columns([name], [bin], ["s64"], n) + batch + end +end