[C++][Parquet] Preserve the bithwidth of the integer dictionary indices on rountrip to Parquet?

[asfimport]

When converting from a pandas dataframe to a table, categorical variables are by default given an index type int8 (presumably because there are fewer than 128 categories) in the schema. When this is written to a parquet file, the schema changes such that the index type is int32 instead. This causes an inconsistency between the schemas of tables derived from pandas and those read from disk.

A minimal recreation of the issue is as follows:

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

df = pd.DataFrame({"A": [1, 2, 3, 4, 5], "B": ["a", "a", "b", "c", "b"]})
dtypes = {
 "A": np.dtype("int8"),
 "B": pd.CategoricalDtype(categories=["a", "b", "c"], ordered=None),
}
df = df.astype(dtypes)

tbl = pa.Table.from_pandas(
 df, 
) 
where = "tmp.parquet"
filesystem = pa.fs.LocalFileSystem()

pq.write_table(
 tbl,
 filesystem.open_output_stream(
 where,
 compression=None,
 ),
 version="2.0",
)

schema = tbl.schema

read_schema = pq.ParquetFile(
 filesystem.open_input_file(where),
).schema_arrow

By printing schema and read_schema, you can the inconsistency.

I have workarounds in place for this, but am raising the issue anyway so that you can resolve it properly.

Environment: NAME="CentOS Linux"
VERSION="7 (Core)"
Reporter: Gavin

_{Note: This issue was originally created as ARROW-14767. Please see the migration documentation for further details.}

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
Thanks for the report!

I can reproduce your issue. Looking at the written Parquet file's schema, you can see that the column "B" is indicated to be of type "String":

In [29]: parquet_metadata = pq.ParquetFile(
 ...: filesystem.open_input_file(where),
 ...: )

In [30]: parquet_metadata.schema
Out[30]: 
<pyarrow._parquet.ParquetSchema object at 0x7f60be97fe00>
required group field_id=-1 schema {
 optional int32 field_id=-1 A (Int(bitWidth=8, isSigned=true));
 optional binary field_id=-1 B (String);
}

Parquet still uses dictionary encoding for compression reasons, but it is not a separate type in their type system.

Arrow has the ability to read such dictionary encoded string columns directly into an Arrow dictionary type. But at that point it doesn't preserve the int8 type for the indices, since that is not part of the parquet spec, and rather uses Arrow's default of int32 indices.

Now, we do write the Arrow schema into the Parquet file metadata, so we can use that information to ensure a more faithful roundtrip of Arrow tables to Parquet and back.

That doesn't happen here, but in principle could be done (although that will probably require an additional cast after the fact)

[lesterfan]

I ran into this issue recently as well. In case it helps, I spent some time making a minimal C++ repro since fixing this would involve a C++ code change (i.e. maybe this could be turned into a unit test). It took me a while to realize that I needed to explicitly set store_schema() in the ArrowWriterProperties as the default in C++ is false while this is true in Python.

arrow::Status arrow_test() {
 arrow::MemoryPool* pool = arrow::default_memory_pool();

 std::shared_ptr<arrow::Array> values;
 arrow::DictionaryBuilder<arrow::StringType> dict_builder(pool);
 ARROW_RETURN_NOT_OK(dict_builder.Append("abc"));
 ARROW_RETURN_NOT_OK(dict_builder.Append("def"));
 ARROW_RETURN_NOT_OK(dict_builder.Append("abc"));
 ARROW_RETURN_NOT_OK(dict_builder.Finish(&values));

 auto dict_type = dict_builder.type();
 auto schema = arrow::schema({arrow::field("x", dict_type)});

 auto table = arrow::Table::Make(schema, {values});

 const char* filepath = "/not/a/real/filepath"; // NOTE: Replace with a real filepath
 std::shared_ptr<arrow::io::FileOutputStream> outfile;
 arrow::io::FileOutputStream::Open(filepath).ValueOrDie().swap(outfile);

 auto writer_props = parquet::WriterProperties::Builder().build();

 // NOTE: This is important; if store_schema not set, the writer won't add the
 // "ARROW:schema": <base 64 encoded schema> metadata to the parquet file,
 // and the reader will read the column as a string rather than as a dictionary.
 auto arrow_writer_props =
 parquet::ArrowWriterProperties::Builder().store_schema()->build();

 ARROW_RETURN_NOT_OK(parquet::arrow::WriteTable(
 *table, pool, outfile, 1024, writer_props, arrow_writer_props));

 std::shared_ptr<arrow::io::ReadableFile> infile;
 arrow::io::ReadableFile::Open(filepath).ValueOrDie().swap(infile);
 std::unique_ptr<parquet::arrow::FileReader> reader;
 ARROW_RETURN_NOT_OK(parquet::arrow::OpenFile(infile, pool, &reader));

 std::shared_ptr<arrow::Table> table_round_trip;
 ARROW_RETURN_NOT_OK(reader->ReadTable(&table_round_trip));

 printf("table: %s\n", table->ToString().c_str());
 printf("table_round_trip: %s\n", table_round_trip->ToString().c_str());

 return arrow::Status::OK();
}

This code gives the following output:

table: x: dictionary<values=string, indices=int8, ordered=0>
----
x:
 [

 -- dictionary:
 [
 "abc",
 "def"
 ]
 -- indices:
 [
 0,
 1,
 0
 ]
 ]

table_round_trip: x: dictionary<values=string, indices=int32, ordered=0>
----
x:
 [

 -- dictionary:
 [
 "abc",
 "def"
 ]
 -- indices:
 [
 0,
 1,
 0
 ]
 ]

ok

The expected behavior is that the indices of table_round_trip should be int8 instead of int32. (The int8 comes from the DictionaryBuilder's AdaptiveIntBuilder determining that the smallest index size which can accommodate the dictionary indices is an int8).

[lesterfan]

I spent some more time looking into this today and put up #45685 which fixes the bug locally for me with an explicit cast. I wrote down some of my thoughts in the PR description; let me know what you guys think!

[pitrou]

Issue resolved by pull request 45685
#45685

[cjrh]

@lesterfan I have a directory of parquet files. For a specific categorical column, some parquet files use int8 and some use int16. In pyarrow 19.0.1, reading the directory as a dataset succeeds. But with pyarrow 20, it fails with this error when loading data from the dataset directory:

Reader code

Either

 import pandas as pd
 df = pd.read_parquet(
 path,
 engine="pyarrow",
 )

or

 import pyarrow.dataset as dataset
 dataset = dataset.dataset(path, format="parquet")
 table = dataset.to_table()
 df = table.to_pandas()

Traceback

 ...
 File "/app/venv/lib/python3.12/site-packages/pyarrow/parquet/core.py", line 1475, in read
 table = self._dataset.to_table(
 ^^^^^^^^^^^^^^^^^^^^^^^
 File "pyarrow/_dataset.pyx", line 589, in pyarrow._dataset.Dataset.to_table
 File "pyarrow/_dataset.pyx", line 3941, in pyarrow._dataset.Scanner.to_table
 File "pyarrow/error.pxi", line 155, in pyarrow.lib.pyarrow_internal_check_status
 File "pyarrow/error.pxi", line 92, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: Integer value 731 not in range: -128 to 127

Metadata

This shows the dictionary info of the parquet files in the directory:

>>> import pyarrow.dataset as dataset
>>> ds = dataset(path)
>>> for path in ds.files:
... sch = pq.read_schema(path)
... print(path, sch.field('ExpStartDate').type)
... 
dataframes.parq/00eac90ef2f504223a74498405e060a48.parquet dictionary<values=string, indices=int8, ordered=0>
dataframes.parq/0641c30f725cd448bafc335d36cd01f6b.parquet dictionary<values=string, indices=int16, ordered=0>
dataframes.parq/0cb2799478dd54c738efe76fdc1875326.parquet dictionary<values=string, indices=int8, ordered=0>
dataframes.parq/0cff477be69be4ee093d98728d4f84452.parquet dictionary<values=string, indices=int16, ordered=0>
dataframes.parq/0d103de6323904e93aecf24589c12a370.parquet dictionary<values=string, indices=int8, ordered=0>

Is my issue related to this change? Is there a way to restore the previous behaviour of upcasting to int32 on read?

[cjrh]

Hi @pitrou, if you have a moment could you have a quick look at my question above and let me know if my question has merit?

[pitrou]

@cjrh Yes, I think it has merit. Can you open a new issue for it?