You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/09/26 07:48:39 UTC
[PR] perf: choose dictionary encoding like iceberg-java in the native Iceberg writer [datafusion-comet]
andygrove opened a new pull request, #6241:
URL: https://github.com/apache/datafusion-comet/pull/6241
## Which issue does this PR close?
Closes #6114.
## Rationale for this change
iceberg-java writes Parquet through parquet-mr, and parquet-mr decides dictionary encoding separately for each column chunk. The test runs on the chunk's first data page. If the dictionary-encoded page plus the dictionary is not smaller than the plain page, the whole chunk is written plain and gets no dictionary page. parquet-rs has no such check. It keeps a dictionary until the dictionary reaches `write.parquet.dict-size-bytes` (2 MiB by default), then falls back to plain but still writes the dictionary page. As a result, high-cardinality columns in natively written files (ids, UUIDs, timestamps) each carry a dictionary page of up to 2 MiB, and every selective read of the column chunk has to fetch it.
arrow-rs has no option for this yet. [arrow-rs#9699](https://github.com/apache/arrow-rs/issues/9699) tracks it. The open PRs, [arrow-rs#10775](https://github.com/apache/arrow-rs/pull/10775) and [#10780](https://github.com/apache/arrow-rs/pull/10780), add fallback policies, but neither is parquet-mr's first-page test. So this PR makes the decision on the Comet side.
## What changes are included in this PR?
parquet-rs fixes a column's encoding when a file opens. So the native writer now makes parquet-mr's decision before it opens a partition's first file.
- **`DictionaryChooser`** (new, `iceberg_dictionary.rs`) replays parquet-mr's first-page accounting from `FallbackValuesWriter`, `DictionaryValuesWriter` and `RunLengthBitPackingHybridEncoder`:
- the plain size and the dictionary size, including the 4-byte length prefix for binary values;
- the RLE / bit-packing hybrid length of the dictionary ids, a direct port of the encoder's run splitting;
- the dictionary limit being exceeded within the first page;
- a page that holds only nulls, which parquet-mr writes plain.
The first page is `write.parquet.page-row-limit` rows, or fewer if the page size ends it sooner. It never ends before row 100, which is where parquet-mr runs its first size check. Nested leaves (struct, list, map) are handled. Columns that parquet-rs never dictionary-encodes (boolean, and fixed-length binary under format v1) are left alone. The chooser returns the table's writer properties with `set_column_dictionary_enabled(path, false)` for every column parquet-mr would write plain.
- **`PartitionFeed`** (`iceberg_write.rs`) holds back each partition's first page of rows, records the choice, and then passes the rows to `RowPacer`. The rolling writer receives the same 1000-row units as before, so roll points do not move.
- **`PartitionWriterBuilder`** replaces the single per-task `DataFileWriterBuilder`. It builds each partition's `ParquetWriterBuilder -> RollingFileWriterBuilder -> DataFileWriterBuilder` chain with that partition's properties. It shares the location and file name generators, so cleanup tracking and file numbering do not change.
- **`FanoutFeeds`** caps the rows a fanout write holds back across all of its open partitions at `write.parquet.row-group-size-bytes`. Without the cap, a task writing many small partitions would hold every row until close. Unpartitioned and clustered writes only ever hold one partition's rows.
- **Docs.** The user guide bullet for this divergence now describes what still differs, and the user guide notes the held-back rows. The contributor guide describes the new adaptation and its tests.
These differences from parquet-mr remain, and the user guide documents them:
- **One decision per partition.** The native writer decides once per partition per task, where parquet-mr decides again for every row group.
- **Pages the page size ends.** When the page size ends a column's first page, the native page ends at the first row past parquet-mr's threshold. parquet-mr ends it at its next periodic size check, whose timing depends on every column. I compared about 3,000 column and setting pairs against parquet-mr. All five mismatches were in columns close to the cut-off, with 32–64 KiB pages. At Iceberg's default page settings there were none.
**Cost.** The choice hashes every value of a partition's first page once. I measured it with a release-mode micro-benchmark on the 20,000-row, 75-column test corpus. Choosing took 27 ms. Writing the same rows with zstd(3) took 54 ms on main (dictionary encoding on for every column) and 48 ms with the chosen properties.
- A partition with more than a page of rows pays for the choice once, then writes its high-cardinality columns faster.
- The worst case is a task whose partitions each have fewer rows than a page. Every row is then sampled, which adds about 40% to the Parquet encoding of those rows.
Not covered here: Comet's native Parquet writer for Spark writes (`ParquetWriterExec`) diverges from parquet-mr in the same way.
## How are these changes tested?
- **Chooser tests** (`iceberg_dictionary.rs`) compare the chooser with parquet-mr's own answers. I recorded those answers by writing the same generated rows through iceberg-java's `Parquet.write` with `GenericParquetWriter`, then listing which column chunks got a dictionary page. Iceberg 1.5.2, 1.8.1, 1.10.0 and 1.11.0 (parquet-mr 1.13.1, 1.15.0, 1.16.0 and 1.17.1) all give the same answers. The corpus has columns on both sides of the cut-off in every physical type, plus runs, nulls, wide strings and nested leaves. The settings covered are:
- Iceberg's defaults;
- 2000-row / 16 KiB pages (the settings of the `CometIcebergNativeSuite` page skipping test);
- a row limit below the first size check;
- a partition with fewer rows than a page;
- a 4 KiB dictionary limit.
Further cases check RLE-hybrid lengths against parquet-mr's encoder, a one-byte tie (which parquet-mr writes plain), values under a null list slot, and different batch boundaries.
- **Rust end-to-end tests** (`iceberg_write.rs`) cover three cases: a unique column is written with no dictionary page; with fanout and clustered writers, each partition makes its own choice; and fanout partitions share one cap on held rows.
- **JVM parity tests** (`CometIcebergWriteActionSuite`, two new tests) write the same rows through both writers and compare the footers column by column:
- a 21-column unpartitioned table, including nested columns and two columns just either side of the cut-off;
- a partitioned table whose partitions differ in cardinality.
On main both tests fail: native files have dictionary pages on 11 columns where iceberg-java has none. With this PR they pass on the Spark 3.4, 3.5 and 4.1 profiles locally.
- **Mutation checks.** The tests fail when the chooser is disabled, when one choice is made for the whole task, when the fanout cap is removed, and under each of the chooser mutations I tried: first-check row, page-size tolerance, dictionary-limit check, strict comparison, bit-width byte, run header, 63-group run limit, and parent and leaf null handling.
- **Also run locally:**
- on the default profile, all of `CometIcebergWriteActionSuite` (76 tests), `CometIcebergRewriteActionSuite`, `CometIcebergSystemFunctionSuite` and `CometIcebergWriteDetectionSuite` (71 tests);
- `cargo test -p datafusion-comet --lib` (510 tests), `cargo clippy --all-targets --workspace -- -D warnings`, `cargo fmt`, and prettier.
- **Not run:** the Iceberg Spark SQL tests, which exercise the native writer (`run-iceberg-tests`), and the other Spark profiles for the full suites (`run-all-spark-profiles`).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: choose dictionary encoding like iceberg-java in the native Iceberg writer [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #6241:
URL: https://github.com/apache/datafusion-comet/pull/6241#issuecomment-5995066619
@adriangb I went ahead and filed apache/arrow-rs#11384 for upstreaming this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: choose dictionary encoding like iceberg-java in the native Iceberg writer [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #6241:
URL: https://github.com/apache/datafusion-comet/pull/6241
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: choose dictionary encoding like iceberg-java in the native Iceberg writer [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #6241:
URL: https://github.com/apache/datafusion-comet/pull/6241#issuecomment-5984468970
> I'd love to see smarter encodings being chosen in arrow-rs / parquet-rs - is this sort of work something you'd be interested in upstreaming?
I would be interested in doing that, but not sure when I'll find the time. Maybe you could file upstream issues requesting this and someone can pick the work up?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: choose dictionary encoding like iceberg-java in the native Iceberg writer [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #6241:
URL: https://github.com/apache/datafusion-comet/pull/6241#discussion_r4111185721
##########
native/core/src/execution/operators/iceberg_dictionary.rs:
##########
@@ -0,0 +1,1132 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Which columns of a native Iceberg data file are dictionary-encoded.
+//!
+//! iceberg-java writes Parquet through parquet-mr, which judges dictionary encoding on the first
+//! data page of every column chunk (`FallbackValuesWriter`): unless the page's dictionary-encoded
+//! values plus its dictionary come out smaller than the page's plain encoding, the chunk is written
+//! plain from its first page on and has no dictionary page. parquet-rs never makes that
+//! comparison. It keeps a column dictionary-encoded until the dictionary reaches
+//! `dictionary_page_size_limit`, then falls back to plain for the rest of the chunk but still
+//! writes the dictionary page. A high-cardinality column therefore carries a dictionary page of up
+//! to `write.parquet.dict-size-bytes` that every selective read of the chunk has to fetch, where
+//! iceberg-java writes none (apache/datafusion-comet#6114).
+//!
+//! parquet-rs cannot change a column's encoding once a file is open, so the native writer takes
+//! parquet-mr's decision before opening one: [`DictionaryChooser`] replays parquet-mr's size
+//! accounting over the rows each column's first data page would hold, and turns dictionary encoding
+//! off for the columns parquet-mr would write plain. The accounting follows `FallbackValuesWriter`,
+//! `DictionaryValuesWriter` and `RunLengthBitPackingHybridEncoder` in parquet-column, and the page
+//! boundary follows `ColumnWriteStoreBase.sizeCheck`.
+
+use std::borrow::Cow;
+
+use arrow::array::{downcast_primitive_array, Array, AsArray, OffsetSizeTrait, RecordBatch};
+use arrow::datatypes::{DataType, Schema as ArrowSchema, ToByteSlice};
+use datafusion::common::hash_map::Entry;
+use datafusion::common::HashMap;
+use datafusion::error::{DataFusionError, Result as DFResult};
+use parquet::arrow::ArrowSchemaConverter;
+use parquet::basic::Type as PhysicalType;
+use parquet::file::properties::WriterProperties;
+use parquet::schema::types::ColumnPath;
+
+/// The row count of parquet-mr's first page size check. Iceberg passes
+/// `write.parquet.row-group-check-min-record-count` to parquet-mr as the minimum row count between
+/// page size checks, and the native writer declines tables that change it from this default, so no
+/// first page ends before this row even when `write.parquet.page-row-limit` is lower.
+const FIRST_PAGE_SIZE_CHECK_ROWS: usize = 100;
+
+/// parquet-mr ends a page once less than this share of the page size is left
+/// (`ColumnWriteStoreBase.THRESHOLD_TOLERANCE_RATIO`).
+const PAGE_SIZE_TOLERANCE_RATIO: f32 = 0.1;
+
+/// Decides, from the rows a partition's data starts with, which of its columns keep dictionary
+/// encoding.
+///
+/// The choice is made once per partition and applies to every file and row group the task writes
+/// for it, where parquet-mr decides again for every column chunk. A column's first data page is
+/// the table's `write.parquet.page-row-limit` rows unless the page size ends it sooner. parquet-mr
+/// only notices a full page at its periodic size checks, whose timing depends on every column, so
+/// where the page size decides, this ends the page at the first row past parquet-mr's threshold and
+/// can judge a column on somewhat fewer rows than parquet-mr would.
+pub(super) struct DictionaryChooser {
+ /// The table's writer properties, which every choice starts from.
+ base: WriterProperties,
+ /// One entry per Parquet leaf column, in schema order. `None` for the columns there is nothing
+ /// to decide for: parquet-rs never dictionary-encodes booleans, nor fixed-length byte arrays
+ /// under Parquet format v1, and a column whose dictionary is already off stays off.
+ columns: Vec<Option<Column>>,
+ /// Rows the first data page of every column holds at most.
+ page_rows: usize,
+ /// Plain-encoded bytes at which parquet-mr's size check ends a page.
+ page_bytes: u64,
+ /// Memory at which a partition stops holding rows back, however few it has: the row group
+ /// size, since parquet-mr's first page never outlasts its first row group.
+ max_held_bytes: Option<usize>,
+}
+
+struct Column {
+ path: ColumnPath,
+ /// Plain-encoded width of one value, or `None` for a byte array, which parquet-mr accounts as
+ /// a 4-byte length plus the bytes.
+ width: Option<u64>,
+ /// `write.parquet.dict-size-bytes`. parquet-mr abandons a dictionary as soon as it grows past
+ /// this, which on a first page leaves no dictionary page at all.
+ dictionary_limit: u64,
+}
+
+impl DictionaryChooser {
+ /// `schema` is the Arrow schema every batch is written with, which fixes the Parquet leaf
+ /// columns the same way the parquet writer derives them.
+ pub(super) fn try_new(base: WriterProperties, schema: &ArrowSchema) -> DFResult<Self> {
+ let parquet_schema = ArrowSchemaConverter::new()
+ .with_coerce_types(base.coerce_types())
+ .convert(schema)
+ .map_err(DataFusionError::from)?;
+ let columns = parquet_schema
+ .columns()
+ .iter()
+ .map(|descr| {
+ let path = descr.path();
+ let width = match descr.physical_type() {
+ PhysicalType::INT32 | PhysicalType::FLOAT => Some(4),
+ PhysicalType::INT64 | PhysicalType::DOUBLE => Some(8),
+ PhysicalType::BYTE_ARRAY => None,
+ // INT96 is not an Iceberg type.
+ PhysicalType::BOOLEAN
+ | PhysicalType::FIXED_LEN_BYTE_ARRAY
+ | PhysicalType::INT96 => return None,
+ };
+ base.dictionary_enabled(path).then(|| Column {
+ path: path.clone(),
+ width,
+ dictionary_limit: base.column_dictionary_page_size_limit(path) as u64,
+ })
+ })
+ .collect();
+ let page_size = base.data_page_size_limit();
+ let tolerance = (page_size as f32 * PAGE_SIZE_TOLERANCE_RATIO) as usize;
+ Ok(Self {
+ columns,
+ page_rows: base
+ .data_page_row_count_limit()
+ .max(FIRST_PAGE_SIZE_CHECK_ROWS),
+ page_bytes: page_size.saturating_sub(tolerance) as u64,
+ max_held_bytes: base.max_row_group_bytes(),
+ base,
+ })
+ }
+
+ /// The table's writer properties, before any choice.
+ pub(super) fn base(&self) -> &WriterProperties {
+ &self.base
+ }
+
+ /// Whether a partition that has shown `rows` rows, held in `bytes` of memory, should hold back
+ /// more before choosing.
+ pub(super) fn wants_more(&self, rows: usize, bytes: usize) -> bool {
+ self.columns.iter().any(Option::is_some) && rows < self.page_rows && self.may_hold(bytes)
+ }
+
+ /// Whether rows taking `bytes` of memory may stay held back.
+ pub(super) fn may_hold(&self, bytes: usize) -> bool {
+ self.max_held_bytes.is_none_or(|max| bytes < max)
+ }
+
+ /// The writer properties for a partition whose rows start with `sample`: the table's, with
+ /// dictionary encoding turned off for every column parquet-mr would write plain.
+ pub(super) fn choose(&self, sample: &[RecordBatch]) -> WriterProperties {
+ let mut pages: Vec<Option<FirstPage>> = self
+ .columns
+ .iter()
+ .map(|column| {
+ column
+ .as_ref()
+ .map(|column| FirstPage::new(column, self.page_bytes))
+ })
+ .collect();
+ let mut rows_seen = 0;
+ for batch in sample {
+ let rows = batch.num_rows().min(self.page_rows - rows_seen);
+ let slots: Vec<(u32, usize)> = (0..rows).map(|row| (row as u32, row)).collect();
+ let mut walk = LeafWalk {
+ pages: &mut pages,
+ next: 0,
+ rows,
+ };
+ for column in batch.columns() {
+ walk.visit(column.as_ref(), &slots);
+ }
+ if walk.next != pages.len() {
+ // The batch does not have the leaves the schema promised, so there is no telling
+ // which page a value belongs to. Leave every column as the table configured it.
+ return self.base.clone();
+ }
+ rows_seen += rows;
+ if rows_seen == self.page_rows || pages.iter().flatten().all(FirstPage::is_closed) {
+ break;
+ }
+ }
+ let plain: Vec<&ColumnPath> = self
+ .columns
+ .iter()
+ .zip(&pages)
+ .filter_map(|(column, page)| match (column, page) {
+ (Some(column), Some(page)) if !page.keeps_dictionary() => Some(&column.path),
+ _ => None,
+ })
+ .collect();
+ if rows_seen == 0 || plain.is_empty() {
+ return self.base.clone();
+ }
+ plain
+ .into_iter()
+ .fold(self.base.clone().into_builder(), |builder, path| {
+ builder.set_column_dictionary_enabled(path.clone(), false)
+ })
+ .build()
+ }
+}
+
+/// Hands one batch's values to the first page of the leaf column they belong to, visiting the
+/// leaves in Parquet schema order.
+struct LeafWalk<'p, 'a> {
+ pages: &'p mut [Option<FirstPage<'a>>],
+ /// The leaf the next primitive array belongs to.
+ next: usize,
+ /// Rows of the batch that fall within the first page.
+ rows: usize,
+}
+
+impl<'a> LeafWalk<'_, 'a> {
+ /// `slots` holds a `(row, index)` pair for every entry of `array` whose ancestors are all
+ /// present, in the order the column chunk stores them.
+ fn visit(&mut self, array: &'a dyn Array, slots: &[(u32, usize)]) {
+ match array.data_type() {
+ DataType::Struct(_) => {
+ let slots = present(array, slots);
+ for child in array.as_struct().columns() {
+ self.visit(child.as_ref(), &slots);
+ }
+ }
+ DataType::List(_) => {
+ let list = array.as_list::<i32>();
+ self.visit(
+ list.values().as_ref(),
+ &children(array, list.value_offsets(), slots),
Review Comment:
[P2] Could the walker skip ineligible or closed subtrees before expanding child slots, and traverse active children with bounded storage? For a first batch containing 1,000 `array<boolean>` rows with 8,192 elements each, `children()` materializes millions of `(u32, usize)` tuples even though boolean leaves have no dictionary decision. `PartitionFeed::push` still invokes `choose()` when `wants_more()` returns false. The expected behavior is to return the unchanged properties without visiting these elements. Instead, the exact-head allocation probe measured about 128 MiB of additional temporary memory for roughly 1 MiB of input. These allocations are outside both `held_bytes` and pool reservations, creating substantial avoidable memory pressure during native writes. Skipping inactive subtrees and replacing eager child-index vectors with bounded traversal would address this.
Evidence: Ran `cargo run --offline --manifest-path /tmp/comet6241-review/Cargo.toml`. The inspected harness imports the checkout’s exact `iceberg_dictionary.rs`, uses Arrow/Parquet 59.3.0, constructs the boolean-list batch described above, and measures live allocation around `choose()`. Output: `rows=1000 bools_per_row=8192 held_bytes=1028204 wants_more=false additional_peak_bytes=134233872`. No executor exhaustion was attempted. The eager allocation comes from `children()` at lines 283–293, before the leaf eligibility check at line 261.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: choose dictionary encoding like iceberg-java in the native Iceberg writer [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6241:
URL: https://github.com/apache/datafusion-comet/pull/6241#discussion_r4111609186
##########
native/core/src/execution/operators/iceberg_dictionary.rs:
##########
@@ -0,0 +1,1132 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Which columns of a native Iceberg data file are dictionary-encoded.
+//!
+//! iceberg-java writes Parquet through parquet-mr, which judges dictionary encoding on the first
+//! data page of every column chunk (`FallbackValuesWriter`): unless the page's dictionary-encoded
+//! values plus its dictionary come out smaller than the page's plain encoding, the chunk is written
+//! plain from its first page on and has no dictionary page. parquet-rs never makes that
+//! comparison. It keeps a column dictionary-encoded until the dictionary reaches
+//! `dictionary_page_size_limit`, then falls back to plain for the rest of the chunk but still
+//! writes the dictionary page. A high-cardinality column therefore carries a dictionary page of up
+//! to `write.parquet.dict-size-bytes` that every selective read of the chunk has to fetch, where
+//! iceberg-java writes none (apache/datafusion-comet#6114).
+//!
+//! parquet-rs cannot change a column's encoding once a file is open, so the native writer takes
+//! parquet-mr's decision before opening one: [`DictionaryChooser`] replays parquet-mr's size
+//! accounting over the rows each column's first data page would hold, and turns dictionary encoding
+//! off for the columns parquet-mr would write plain. The accounting follows `FallbackValuesWriter`,
+//! `DictionaryValuesWriter` and `RunLengthBitPackingHybridEncoder` in parquet-column, and the page
+//! boundary follows `ColumnWriteStoreBase.sizeCheck`.
+
+use std::borrow::Cow;
+
+use arrow::array::{downcast_primitive_array, Array, AsArray, OffsetSizeTrait, RecordBatch};
+use arrow::datatypes::{DataType, Schema as ArrowSchema, ToByteSlice};
+use datafusion::common::hash_map::Entry;
+use datafusion::common::HashMap;
+use datafusion::error::{DataFusionError, Result as DFResult};
+use parquet::arrow::ArrowSchemaConverter;
+use parquet::basic::Type as PhysicalType;
+use parquet::file::properties::WriterProperties;
+use parquet::schema::types::ColumnPath;
+
+/// The row count of parquet-mr's first page size check. Iceberg passes
+/// `write.parquet.row-group-check-min-record-count` to parquet-mr as the minimum row count between
+/// page size checks, and the native writer declines tables that change it from this default, so no
+/// first page ends before this row even when `write.parquet.page-row-limit` is lower.
+const FIRST_PAGE_SIZE_CHECK_ROWS: usize = 100;
+
+/// parquet-mr ends a page once less than this share of the page size is left
+/// (`ColumnWriteStoreBase.THRESHOLD_TOLERANCE_RATIO`).
+const PAGE_SIZE_TOLERANCE_RATIO: f32 = 0.1;
+
+/// Decides, from the rows a partition's data starts with, which of its columns keep dictionary
+/// encoding.
+///
+/// The choice is made once per partition and applies to every file and row group the task writes
+/// for it, where parquet-mr decides again for every column chunk. A column's first data page is
+/// the table's `write.parquet.page-row-limit` rows unless the page size ends it sooner. parquet-mr
+/// only notices a full page at its periodic size checks, whose timing depends on every column, so
+/// where the page size decides, this ends the page at the first row past parquet-mr's threshold and
+/// can judge a column on somewhat fewer rows than parquet-mr would.
+pub(super) struct DictionaryChooser {
+ /// The table's writer properties, which every choice starts from.
+ base: WriterProperties,
+ /// One entry per Parquet leaf column, in schema order. `None` for the columns there is nothing
+ /// to decide for: parquet-rs never dictionary-encodes booleans, nor fixed-length byte arrays
+ /// under Parquet format v1, and a column whose dictionary is already off stays off.
+ columns: Vec<Option<Column>>,
+ /// Rows the first data page of every column holds at most.
+ page_rows: usize,
+ /// Plain-encoded bytes at which parquet-mr's size check ends a page.
+ page_bytes: u64,
+ /// Memory at which a partition stops holding rows back, however few it has: the row group
+ /// size, since parquet-mr's first page never outlasts its first row group.
+ max_held_bytes: Option<usize>,
+}
+
+struct Column {
+ path: ColumnPath,
+ /// Plain-encoded width of one value, or `None` for a byte array, which parquet-mr accounts as
+ /// a 4-byte length plus the bytes.
+ width: Option<u64>,
+ /// `write.parquet.dict-size-bytes`. parquet-mr abandons a dictionary as soon as it grows past
+ /// this, which on a first page leaves no dictionary page at all.
+ dictionary_limit: u64,
+}
+
+impl DictionaryChooser {
+ /// `schema` is the Arrow schema every batch is written with, which fixes the Parquet leaf
+ /// columns the same way the parquet writer derives them.
+ pub(super) fn try_new(base: WriterProperties, schema: &ArrowSchema) -> DFResult<Self> {
+ let parquet_schema = ArrowSchemaConverter::new()
+ .with_coerce_types(base.coerce_types())
+ .convert(schema)
+ .map_err(DataFusionError::from)?;
+ let columns = parquet_schema
+ .columns()
+ .iter()
+ .map(|descr| {
+ let path = descr.path();
+ let width = match descr.physical_type() {
+ PhysicalType::INT32 | PhysicalType::FLOAT => Some(4),
+ PhysicalType::INT64 | PhysicalType::DOUBLE => Some(8),
+ PhysicalType::BYTE_ARRAY => None,
+ // INT96 is not an Iceberg type.
+ PhysicalType::BOOLEAN
+ | PhysicalType::FIXED_LEN_BYTE_ARRAY
+ | PhysicalType::INT96 => return None,
+ };
+ base.dictionary_enabled(path).then(|| Column {
+ path: path.clone(),
+ width,
+ dictionary_limit: base.column_dictionary_page_size_limit(path) as u64,
+ })
+ })
+ .collect();
+ let page_size = base.data_page_size_limit();
+ let tolerance = (page_size as f32 * PAGE_SIZE_TOLERANCE_RATIO) as usize;
+ Ok(Self {
+ columns,
+ page_rows: base
+ .data_page_row_count_limit()
+ .max(FIRST_PAGE_SIZE_CHECK_ROWS),
+ page_bytes: page_size.saturating_sub(tolerance) as u64,
+ max_held_bytes: base.max_row_group_bytes(),
+ base,
+ })
+ }
+
+ /// The table's writer properties, before any choice.
+ pub(super) fn base(&self) -> &WriterProperties {
+ &self.base
+ }
+
+ /// Whether a partition that has shown `rows` rows, held in `bytes` of memory, should hold back
+ /// more before choosing.
+ pub(super) fn wants_more(&self, rows: usize, bytes: usize) -> bool {
+ self.columns.iter().any(Option::is_some) && rows < self.page_rows && self.may_hold(bytes)
+ }
+
+ /// Whether rows taking `bytes` of memory may stay held back.
+ pub(super) fn may_hold(&self, bytes: usize) -> bool {
+ self.max_held_bytes.is_none_or(|max| bytes < max)
+ }
+
+ /// The writer properties for a partition whose rows start with `sample`: the table's, with
+ /// dictionary encoding turned off for every column parquet-mr would write plain.
+ pub(super) fn choose(&self, sample: &[RecordBatch]) -> WriterProperties {
+ let mut pages: Vec<Option<FirstPage>> = self
+ .columns
+ .iter()
+ .map(|column| {
+ column
+ .as_ref()
+ .map(|column| FirstPage::new(column, self.page_bytes))
+ })
+ .collect();
+ let mut rows_seen = 0;
+ for batch in sample {
+ let rows = batch.num_rows().min(self.page_rows - rows_seen);
+ let slots: Vec<(u32, usize)> = (0..rows).map(|row| (row as u32, row)).collect();
+ let mut walk = LeafWalk {
+ pages: &mut pages,
+ next: 0,
+ rows,
+ };
+ for column in batch.columns() {
+ walk.visit(column.as_ref(), &slots);
+ }
+ if walk.next != pages.len() {
+ // The batch does not have the leaves the schema promised, so there is no telling
+ // which page a value belongs to. Leave every column as the table configured it.
+ return self.base.clone();
+ }
+ rows_seen += rows;
+ if rows_seen == self.page_rows || pages.iter().flatten().all(FirstPage::is_closed) {
+ break;
+ }
+ }
+ let plain: Vec<&ColumnPath> = self
+ .columns
+ .iter()
+ .zip(&pages)
+ .filter_map(|(column, page)| match (column, page) {
+ (Some(column), Some(page)) if !page.keeps_dictionary() => Some(&column.path),
+ _ => None,
+ })
+ .collect();
+ if rows_seen == 0 || plain.is_empty() {
+ return self.base.clone();
+ }
+ plain
+ .into_iter()
+ .fold(self.base.clone().into_builder(), |builder, path| {
+ builder.set_column_dictionary_enabled(path.clone(), false)
+ })
+ .build()
+ }
+}
+
+/// Hands one batch's values to the first page of the leaf column they belong to, visiting the
+/// leaves in Parquet schema order.
+struct LeafWalk<'p, 'a> {
+ pages: &'p mut [Option<FirstPage<'a>>],
+ /// The leaf the next primitive array belongs to.
+ next: usize,
+ /// Rows of the batch that fall within the first page.
+ rows: usize,
+}
+
+impl<'a> LeafWalk<'_, 'a> {
+ /// `slots` holds a `(row, index)` pair for every entry of `array` whose ancestors are all
+ /// present, in the order the column chunk stores them.
+ fn visit(&mut self, array: &'a dyn Array, slots: &[(u32, usize)]) {
+ match array.data_type() {
+ DataType::Struct(_) => {
+ let slots = present(array, slots);
+ for child in array.as_struct().columns() {
+ self.visit(child.as_ref(), &slots);
+ }
+ }
+ DataType::List(_) => {
+ let list = array.as_list::<i32>();
+ self.visit(
+ list.values().as_ref(),
+ &children(array, list.value_offsets(), slots),
Review Comment:
Fixed in 3398c8827. The chooser now only reads leaves that have a dictionary choice to make. Booleans, fixed-length binary and columns whose dictionary is already off never become candidates, so a schema with nothing else holds no rows back and `choose` returns without touching the batch. Each candidate is walked one column at a time through a lazy iterator that stops where its first page ends, so no child index vectors are built. I also dropped the per-value id vector. The RLE-hybrid length only depends on whether each id repeats the one before it, so the encoder state machine now runs on that directly, and the only state kept per column is its dictionary.
With your probe shape against the new head in release mode, the 1000 × 8192 `array<boolean>` batch goes from 134,233,856 bytes of additional peak allocation to 25 bytes (3 µs). A 1000 × 8192 `array<int>` batch goes from 138 MB to 52 KB, and the 20k-row, 75-column test corpus from 40 MiB to 3.2 MiB in the same time. The decisions don't change: the parquet-mr ground-truth tests pass as before, and so do the JVM parity tests and the rest of `CometIcebergWriteActionSuite`.
I couldn't find a clean way to measure peak allocation inside the lib's test binary, since the crate already installs its process-wide accounting allocator. So the new tests pin the two properties that bound it instead. `columns_with_no_choice_are_never_read` checks that boolean lists never become candidates, including under a struct. `a_column_is_read_only_as_far_as_its_first_page` checks that a `list<int>` column with 8192 entries per row pulls exactly 101 of 400 top-level rows before its page closes at row 100. Collecting the entries eagerly, as before, pulls all 400 and fails it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: choose dictionary encoding like iceberg-java in the native Iceberg writer [datafusion-comet]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on PR #6241:
URL: https://github.com/apache/datafusion-comet/pull/6241#issuecomment-5879512836
I'd love to see smarter encodings being chosen in arrow-rs / parquet-rs - is this sort of work something you'd be interested in upstreaming?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org