You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ne...@apache.org on 2020/10/25 03:10:55 UTC
[arrow] branch rust-parquet-arrow-writer updated (8ccd9c3 ->
9ba2179)
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a change to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git.
omit 8ccd9c3 ARROW-7842: [Rust] [Parquet] Arrow list reader
omit 453f978 ARROW-10334: [Rust] [Parquet] NullArray roundtrip
omit ead5e14 ARROW-10225: [Rust] [Parquet] Fix null comparison in roundtrip
omit 3a22d3d ARROW-10168: [Rust] [Parquet] Schema roundtrip - use Arrow schema from Parquet metadata when available
omit b7b45d1 ARROW-10191: [Rust] [Parquet] Add roundtrip Arrow -> Parquet tests for all supported Arrow DataTypes
omit 6e237bc ARROW-8426: [Rust] [Parquet] Add support for writing dictionary types
omit 2f81785 ARROW-10095: [Rust] Update rust-parquet-arrow-writer branch's encode_arrow_schema with ipc changes
omit 923d23b ARROW-8423: [Rust] [Parquet] Serialize Arrow schema metadata
omit 2b8dc08 ARROW-8289: [Rust] Parquet Arrow writer with nested support
omit ab62c28 ARROW-10321: [C++] Use check_cxx_source_compiles for AVX512 detect in compiler
omit 487895f ARROW-10311: [Release] Update crossbow verification process
omit f58db45 ARROW-9898: [C++][Gandiva] Fix linking issue with castINT/FLOAT functions
omit 2510f4f ARROW-10313: [C++] Faster UTF8 validation for small strings
omit 7189b91 ARROW-9475: [Java] Clean up usages of BaseAllocator, use BufferAllocator in…
omit 18495e0 ARROW-10294: [Java] Resolve problems of DecimalVector APIs on ArrowBufs
omit 1d10f22 ARROW-10236: [Rust] Add can_cast_types to arrow cast kernel, use in DataFusion
omit 35ace39 ARROW-10174: [Java] Fix reading/writing dict structs
omit ed8b1bc ARROW-10145: [C++][Dataset] Assert integer overflow in partitioning falls back to string
omit ac14e91 ARROW-9479: [JS] Fix Table.from for zero-item serialized tables, Table.empty for schemas containing compound types (List, FixedSizeList, Map)
omit 249adb4 ARROW-10270: [R] Fix CSV timestamp_parsers test on R-devel
omit ea29f65 ARROW-10292: [Rust] [DataFusion] Simplify merge
omit becf329 ARROW-10289: [Rust] Read dictionaries in IPC streams
omit 818593f ARROW-10295 [Rust] [DataFusion] Replace Rc<RefCell<>> by Box<> in accumulators.
omit c5280a5 ARROW-10293: [Rust] [DataFusion] Fixed benchmarks
omit f07a415 ARROW-10263: [C++][Compute] Improve variance kernel numerical stability
add 0aa2069 [Release] Update CHANGELOG.md for 2.0.0
add e46a3c6 [Release] Update .deb/.rpm changelogs for 2.0.0
add 5943421 [Release] Update versions for 2.0.0
add 4782866 [maven-release-plugin] prepare release apache-arrow-2.0.0
add b1f36ac [Release] Update versions for 3.0.0-SNAPSHOT
add f72575c [Release] Update .deb package names for 3.0.0
add a3a35b2 ARROW-10263: [C++][Compute] Improve variance kernel numerical stability
add 91b5f07 ARROW-10293: [Rust] [DataFusion] Fixed benchmarks
add a030fc5 ARROW-10295 [Rust] [DataFusion] Replace Rc<RefCell<>> by Box<> in accumulators.
add 34533b6 ARROW-10289: [Rust] Read dictionaries in IPC streams
add 7209ffc ARROW-10292: [Rust] [DataFusion] Simplify merge
add 9e671ac ARROW-10270: [R] Fix CSV timestamp_parsers test on R-devel
add 8f302d3 ARROW-9479: [JS] Fix Table.from for zero-item serialized tables, Table.empty for schemas containing compound types (List, FixedSizeList, Map)
add 03c7c02 ARROW-10145: [C++][Dataset] Assert integer overflow in partitioning falls back to string
add a7ef5d2 ARROW-10174: [Java] Fix reading/writing dict structs
add 3f69ad2 ARROW-10236: [Rust] Add can_cast_types to arrow cast kernel, use in DataFusion
add 22027c7 ARROW-10294: [Java] Resolve problems of DecimalVector APIs on ArrowBufs
add cb58146 ARROW-9475: [Java] Clean up usages of BaseAllocator, use BufferAllocator in…
add 3f96cc0 ARROW-10313: [C++] Faster UTF8 validation for small strings
add 36bf7a4 ARROW-9898: [C++][Gandiva] Fix linking issue with castINT/FLOAT functions
add 2fe8b86 ARROW-10311: [Release] Update crossbow verification process
add 3dbea74 ARROW-10321: [C++] Use check_cxx_source_compiles for AVX512 detect in compiler
add be21209 ARROW-5350: [Rust] Allow filtering on simple lists
add 29e9d13 ARROW-10002: [Rust] Remove trait specialization from arrow crate
add ad2a002 ARROW-10163: [Rust] [DataFusion] Add DictionaryArray coercion support
add b653922 ARROW-10323: [Release][wheel] Add missing verification setup step
add 5af12b0 ARROW-10340: [Packaging][deb][RPM] Use Python 3.8 for pygit2
add 6ad24c2 ARROW-10310: [C++][Gandiva] Add single argument round() in Gandiva
add 139fb51 ARROW-10241: [C++][Compute] Add variance kernel benchmark
add 457935e ARROW-10106: [FlightRPC][Java] Expose onIsReady() callback
add 7944265 ARROW-10203: [Doc] Give guidance on big-endian support in the contributors docs
add 0e13e28 ARROW-9164: [C++] Add embedded documentation to compute functions
add 445251a ARROW-9991: [C++] Split kernels for strings/binary
add 04660f8 ARROW-10348: [C++] Fix crash on invalid Parquet data
add 7696b89 ARROW-10338: [Rust] Use const fn for applicable methods
add 2321542 ARROW-10318: [C++] Use pimpl idiom in CSV parser
add 84e4b15 ARROW-10320 [Rust] [DataFusion] Migrated from batch iterators to batch streams.
add 760284f ARROW-10358: [R] Followups to 2.0.0 release
add f06c68a ARROW-10331: [Rust] [DataFusion] Re-organize DataFusion errors
add bb4f2a0 ARROW-9963: [Python] Recognize datetime.timezone.utc as UTC on conversion python->pyarrow
add 16412a1 ARROW-10328: [C++] Vendor fast_float number parsing library
add 625bf3f ARROW-10363: [Python] Remove CMake bug workaround in manylinux
add e2d8dc3 ARROW-10304: [C++][Compute] Optimize variance kernel for integers
add 5ad3c98 ARROW-10369: [Dev] Fix archery release utility test cases
add 843e8bb ARROW-9304: [C++] Add "AppendEmpty" builder APIs for use inside StructBuilder::AppendNull
add 103489b ARROW-10208: [C++] Fix split string kernels on sliced input
add 39a7539 ARROW-10207: [C++] Allow precomputing output string/list offsets in kernels
add 982ca83 ARROW-10365: [R] Remove duplicate setting of S3 flag on macOS
add 9cb0e9f ARROW-10249: [Rust] Support nested dictionaries inside list arrays
add b3992ea ARROW-10333: [Java] Get rid of org.apache.arrow.util in vector
add 590541d ARROW-9747: [Java][C++] Initial Support for 256-bit Decimals
add c49b5da ARROW-10302: [Python] Don't double-package plasma-store-server
add eea8d35 ARROW-10375 [Rust] Removed PrimitiveArrayOps
add da50838 ARROW-10383: [Doc] fix typos
add aa2e6d4 ARROW-10384: [C++] Fix typos
add 221de70 ARROW-10382: [Rust] Fix typos
new 6dc26ed ARROW-8289: [Rust] Parquet Arrow writer with nested support
new dcec335 ARROW-8423: [Rust] [Parquet] Serialize Arrow schema metadata
new 98d4ffc ARROW-10095: [Rust] Update rust-parquet-arrow-writer branch's encode_arrow_schema with ipc changes
new 5063854 ARROW-8426: [Rust] [Parquet] Add support for writing dictionary types
new 6677a41 ARROW-10191: [Rust] [Parquet] Add roundtrip Arrow -> Parquet tests for all supported Arrow DataTypes
new 96798dc ARROW-10168: [Rust] [Parquet] Schema roundtrip - use Arrow schema from Parquet metadata when available
new 25d96c5 ARROW-10225: [Rust] [Parquet] Fix null comparison in roundtrip
new 0b5a846 ARROW-10334: [Rust] [Parquet] NullArray roundtrip
new 9ba2179 ARROW-7842: [Rust] [Parquet] Arrow list reader
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (8ccd9c3)
\
N -- N -- N refs/heads/rust-parquet-arrow-writer (9ba2179)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 9 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
CHANGELOG.md | 584 +++++++++++++++++
LICENSE.txt | 8 +
c_glib/configure.ac | 2 +-
c_glib/meson.build | 2 +-
c_glib/test/test-decimal128.rb | 2 +-
ci/conda_env_archery.yml | 2 +
ci/scripts/PKGBUILD | 2 +-
cpp/CMakeLists.txt | 2 +-
cpp/cmake_modules/FindPython3Alt.cmake | 19 +-
cpp/cmake_modules/UseCython.cmake | 2 +-
cpp/examples/arrow/row-wise-conversion-example.cc | 2 +-
cpp/src/arrow/array/array_base.cc | 4 +
cpp/src/arrow/array/array_decimal.cc | 18 +-
cpp/src/arrow/array/array_decimal.h | 16 +
cpp/src/arrow/array/array_dict_test.cc | 66 +-
cpp/src/arrow/array/array_struct_test.cc | 10 +-
cpp/src/arrow/array/array_test.cc | 118 +++-
cpp/src/arrow/array/array_union_test.cc | 86 ++-
cpp/src/arrow/array/builder_adaptive.h | 20 +
cpp/src/arrow/array/builder_base.h | 16 +
cpp/src/arrow/array/builder_binary.cc | 14 +
cpp/src/arrow/array/builder_binary.h | 24 +-
cpp/src/arrow/array/builder_decimal.cc | 35 +
cpp/src/arrow/array/builder_decimal.h | 29 +
cpp/src/arrow/array/builder_dict.cc | 2 +-
cpp/src/arrow/array/builder_dict.h | 34 +-
cpp/src/arrow/array/builder_nested.cc | 39 +-
cpp/src/arrow/array/builder_nested.h | 52 +-
cpp/src/arrow/array/builder_primitive.h | 34 +
cpp/src/arrow/array/builder_union.h | 60 +-
cpp/src/arrow/array/concatenate.cc | 2 +-
cpp/src/arrow/array/validate.cc | 7 +
cpp/src/arrow/builder.cc | 2 +
cpp/src/arrow/c/bridge.cc | 24 +-
cpp/src/arrow/c/bridge_test.cc | 18 +-
cpp/src/arrow/compare.cc | 16 +
cpp/src/arrow/compute/api_aggregate.h | 6 +-
cpp/src/arrow/compute/api_scalar.h | 19 +
cpp/src/arrow/compute/cast.cc | 20 +-
cpp/src/arrow/compute/exec.cc | 159 +++--
cpp/src/arrow/compute/exec_test.cc | 21 +-
cpp/src/arrow/compute/function.cc | 15 +
cpp/src/arrow/compute/function.h | 60 +-
cpp/src/arrow/compute/function_test.cc | 18 +-
cpp/src/arrow/compute/kernel_test.cc | 14 +-
cpp/src/arrow/compute/kernels/aggregate_basic.cc | 34 +-
.../arrow/compute/kernels/aggregate_benchmark.cc | 48 ++
cpp/src/arrow/compute/kernels/aggregate_mode.cc | 12 +-
cpp/src/arrow/compute/kernels/aggregate_test.cc | 100 ++-
cpp/src/arrow/compute/kernels/aggregate_var_std.cc | 96 ++-
cpp/src/arrow/compute/kernels/scalar_arithmetic.cc | 82 ++-
cpp/src/arrow/compute/kernels/scalar_boolean.cc | 66 +-
cpp/src/arrow/compute/kernels/scalar_cast_test.cc | 2 +-
cpp/src/arrow/compute/kernels/scalar_compare.cc | 50 +-
cpp/src/arrow/compute/kernels/scalar_fill_null.cc | 10 +-
cpp/src/arrow/compute/kernels/scalar_nested.cc | 11 +-
cpp/src/arrow/compute/kernels/scalar_set_lookup.cc | 27 +-
cpp/src/arrow/compute/kernels/scalar_string.cc | 617 ++++++++++++++++--
.../compute/kernels/scalar_string_benchmark.cc | 6 +
.../arrow/compute/kernels/scalar_string_internal.h | 44 --
.../arrow/compute/kernels/scalar_string_test.cc | 82 +++
cpp/src/arrow/compute/kernels/scalar_validity.cc | 21 +-
cpp/src/arrow/compute/kernels/vector_hash.cc | 28 +-
cpp/src/arrow/compute/kernels/vector_nested.cc | 21 +-
cpp/src/arrow/compute/kernels/vector_selection.cc | 46 +-
cpp/src/arrow/compute/kernels/vector_sort.cc | 36 +-
cpp/src/arrow/compute/registry.cc | 2 +
cpp/src/arrow/compute/registry_test.cc | 7 +-
cpp/src/arrow/csv/parser.cc | 725 ++++++++++-----------
cpp/src/arrow/csv/parser.h | 158 ++---
cpp/src/arrow/csv/parser_benchmark.cc | 94 ++-
cpp/src/arrow/dataset/dataset_test.cc | 2 +-
cpp/src/arrow/dataset/filter.cc | 1 +
cpp/src/arrow/dbi/hiveserver2/thrift/Types.thrift | 2 +-
cpp/src/arrow/ipc/json_simple.cc | 25 +-
cpp/src/arrow/ipc/json_simple_test.cc | 52 +-
cpp/src/arrow/ipc/message.h | 2 +-
cpp/src/arrow/ipc/metadata_internal.cc | 27 +-
cpp/src/arrow/ipc/reader.h | 2 +-
cpp/src/arrow/ipc/writer.cc | 2 +-
cpp/src/arrow/json/parser_test.cc | 13 +-
cpp/src/arrow/pretty_print.cc | 9 +-
cpp/src/arrow/pretty_print_test.cc | 22 +-
cpp/src/arrow/python/arrow_to_pandas.cc | 31 +-
cpp/src/arrow/python/datetime.cc | 10 +
cpp/src/arrow/python/decimal.cc | 44 +-
cpp/src/arrow/python/decimal.h | 18 +
cpp/src/arrow/python/inference.cc | 13 +-
cpp/src/arrow/python/python_test.cc | 67 +-
cpp/src/arrow/python/python_to_arrow.cc | 6 +
cpp/src/arrow/scalar.cc | 8 +
cpp/src/arrow/scalar.h | 11 +
cpp/src/arrow/scalar_test.cc | 21 +-
cpp/src/arrow/testing/gtest_util.cc | 3 +-
cpp/src/arrow/testing/json_internal.cc | 39 +-
cpp/src/arrow/type.cc | 38 +-
cpp/src/arrow/type.h | 37 +-
cpp/src/arrow/type_fwd.h | 26 +-
cpp/src/arrow/type_test.cc | 90 ++-
cpp/src/arrow/type_traits.h | 28 +-
cpp/src/arrow/util/basic_decimal.cc | 233 +++++--
cpp/src/arrow/util/basic_decimal.h | 104 ++-
cpp/src/arrow/util/bit_run_reader.h | 4 +-
cpp/src/arrow/util/bit_util_benchmark.cc | 2 +-
cpp/src/arrow/util/bit_util_test.cc | 2 +-
cpp/src/arrow/util/byte_stream_split.h | 10 +-
cpp/src/arrow/util/decimal.cc | 135 +++-
cpp/src/arrow/util/decimal.h | 63 ++
cpp/src/arrow/util/decimal_benchmark.cc | 20 +-
cpp/src/arrow/util/decimal_test.cc | 154 +++++
cpp/src/arrow/util/hashing.h | 4 +-
cpp/src/arrow/util/spaced.h | 2 +-
cpp/src/arrow/util/utf8.h | 53 +-
cpp/src/arrow/util/utf8_util_test.cc | 38 ++
cpp/src/arrow/util/value_parsing.cc | 61 +-
cpp/src/arrow/vendored/fast_float/README.md | 9 +
cpp/src/arrow/vendored/fast_float/ascii_number.h | 313 +++++++++
.../arrow/vendored/fast_float/decimal_to_binary.h | 167 +++++
cpp/src/arrow/vendored/fast_float/fast_float.h | 47 ++
cpp/src/arrow/vendored/fast_float/fast_table.h | 691 ++++++++++++++++++++
cpp/src/arrow/vendored/fast_float/float_common.h | 263 ++++++++
cpp/src/arrow/vendored/fast_float/parse_number.h | 118 ++++
cpp/src/arrow/vendored/fast_float/thompson_tao.h | 375 +++++++++++
cpp/src/arrow/visitor.cc | 3 +
cpp/src/arrow/visitor.h | 3 +
cpp/src/arrow/visitor_inline.h | 1 +
cpp/src/gandiva/engine.cc | 2 +-
cpp/src/gandiva/function_registry_arithmetic.cc | 4 +
cpp/src/gandiva/precompiled/extended_math_ops.cc | 24 +-
.../gandiva/precompiled/extended_math_ops_test.cc | 10 +
cpp/src/gandiva/precompiled/types.h | 4 +
cpp/src/parquet/arrow/reader.cc | 4 +
cpp/src/parquet/arrow/reader_internal.cc | 4 +-
cpp/src/parquet/arrow/schema.cc | 3 +-
cpp/src/parquet/arrow/writer.cc | 1 -
cpp/src/parquet/column_writer.cc | 15 +-
csharp/Directory.Build.props | 2 +-
dev/archery/archery/integration/datagen.py | 37 +-
dev/archery/archery/release.py | 5 +-
dev/archery/archery/tests/test_release.py | 24 +-
dev/release/rat_exclude_files.txt | 56 +-
dev/release/verify-release-candidate.sh | 2 +
dev/tasks/conda-recipes/arrow-cpp/build-pyarrow.sh | 1 +
dev/tasks/conda-recipes/arrow-cpp/meta.yaml | 1 +
dev/tasks/homebrew-formulae/apache-arrow.rb | 2 +-
.../homebrew-formulae/autobrew/apache-arrow.rb | 2 +-
.../apache-arrow-archive-keyring/debian/changelog | 6 +
.../yum/apache-arrow-release.spec.in | 3 +
.../apache-arrow/debian.ubuntu-xenial/changelog | 6 +
.../apache-arrow/debian.ubuntu-xenial/control | 86 +--
.../libarrow-cuda-glib300.install} | 0
.../libarrow-cuda300.install} | 0
.../libarrow-dataset-glib300.install} | 0
.../libarrow-dataset300.install} | 0
.../libarrow-glib300.install} | 0
.../libarrow-python300.install} | 0
.../libarrow300.install} | 0
.../libgandiva-glib300.install} | 0
.../libgandiva300.install} | 0
.../libparquet-glib300.install} | 0
.../libparquet300.install} | 0
.../libplasma-glib300.install} | 0
.../libplasma300.install} | 0
.../linux-packages/apache-arrow/debian/changelog | 6 +
.../linux-packages/apache-arrow/debian/control | 102 +--
.../libarrow-cuda-glib300.install} | 0
.../libarrow-cuda300.install} | 0
.../libarrow-dataset-glib300.install} | 0
.../libarrow-dataset300.install} | 0
...light200.install => libarrow-flight300.install} | 0
.../libarrow-glib300.install} | 0
...0.install => libarrow-python-flight300.install} | 0
.../libarrow-python300.install} | 0
.../libarrow300.install} | 0
.../libgandiva-glib300.install} | 0
.../libgandiva300.install} | 0
.../libparquet-glib300.install} | 0
.../libparquet300.install} | 0
.../libplasma-glib300.install} | 0
.../libplasma300.install} | 0
.../linux-packages/apache-arrow/yum/arrow.spec.in | 3 +
dev/tasks/linux-packages/github.linux.amd64.yml | 4 +-
dev/tasks/tasks.yml | 304 ++++-----
docs/source/cpp/compute.rst | 31 +
docs/source/developers/archery.rst | 2 +-
docs/source/developers/contributing.rst | 48 ++
docs/source/developers/docker.rst | 2 +-
docs/source/format/Columnar.rst | 2 +
docs/source/java/vector.rst | 2 +-
docs/source/python/api/compute.rst | 2 +-
docs/source/python/dataset.rst | 2 +-
docs/source/python/ipc.rst | 18 -
java/adapter/avro/pom.xml | 2 +-
.../java/org/apache/arrow/AvroToArrowUtils.java | 2 +-
java/adapter/jdbc/pom.xml | 2 +-
.../arrow/adapter/jdbc/JdbcToArrowUtils.java | 2 +-
java/adapter/orc/pom.xml | 2 +-
java/algorithm/pom.xml | 2 +-
.../algorithm/deduplicate/DeduplicationUtils.java | 2 +-
.../deduplicate/VectorRunDeduplicator.java | 2 +-
.../deduplicate/TestDeduplicationUtils.java | 2 +-
java/flight/flight-core/pom.xml | 2 +-
.../apache/arrow/flight/BackpressureStrategy.java | 172 +++++
.../org/apache/arrow/flight/FlightService.java | 17 +-
.../arrow/flight/OutboundStreamListener.java | 13 +
.../arrow/flight/OutboundStreamListenerImpl.java | 5 +
.../org/apache/arrow/flight/TestBackPressure.java | 159 ++++-
.../arrow/flight/perf/PerformanceTestServer.java | 141 ++--
java/flight/flight-grpc/pom.xml | 2 +-
java/format/pom.xml | 2 +-
java/gandiva/pom.xml | 2 +-
.../arrow/gandiva/evaluator/DecimalTypeUtil.java | 2 +-
.../gandiva/evaluator/ExpressionRegistry.java | 2 +-
.../gandiva/evaluator/DecimalTypeUtilTest.java | 2 +-
.../gandiva/evaluator/ProjectorDecimalTest.java | 34 +-
java/memory/memory-core/pom.xml | 2 +-
.../java/org/apache/arrow/util/AutoCloseables.java | 0
.../java/org/apache/arrow/util/Collections2.java | 0
java/memory/memory-netty/pom.xml | 2 +-
java/memory/memory-unsafe/pom.xml | 2 +-
java/memory/pom.xml | 2 +-
java/performance/pom.xml | 6 +-
java/plasma/pom.xml | 2 +-
java/pom.xml | 4 +-
java/tools/pom.xml | 2 +-
java/vector/pom.xml | 2 +-
java/vector/src/main/codegen/data/ArrowTypes.tdd | 2 +-
.../src/main/codegen/data/ValueVectorTypes.tdd | 17 +
.../codegen/templates/AbstractFieldWriter.java | 8 +-
.../templates/AbstractPromotableFieldWriter.java | 26 +-
.../src/main/codegen/templates/ArrowType.java | 18 +-
.../src/main/codegen/templates/ComplexCopier.java | 5 +-
.../src/main/codegen/templates/ComplexWriters.java | 27 +-
.../main/codegen/templates/DenseUnionReader.java | 4 +-
.../main/codegen/templates/DenseUnionVector.java | 15 +-
.../main/codegen/templates/DenseUnionWriter.java | 12 +-
.../main/codegen/templates/HolderReaderImpl.java | 5 +
.../src/main/codegen/templates/StructWriters.java | 4 +-
.../templates/UnionFixedSizeListWriter.java | 54 ++
.../main/codegen/templates/UnionListWriter.java | 46 ++
.../src/main/codegen/templates/UnionMapWriter.java | 15 +
.../src/main/codegen/templates/UnionReader.java | 8 +-
.../src/main/codegen/templates/UnionVector.java | 31 +-
.../src/main/codegen/templates/UnionWriter.java | 50 +-
.../org/apache/arrow/vector/BaseValueVector.java | 2 +-
.../org/apache/arrow/vector/BitVectorHelper.java | 2 +-
.../java/org/apache/arrow/vector/BufferLayout.java | 5 +-
.../{DecimalVector.java => Decimal256Vector.java} | 91 +--
.../org/apache/arrow/vector/DecimalVector.java | 6 +-
.../java/org/apache/arrow/vector/TypeLayout.java | 2 +-
.../vector/complex/impl/PromotableWriter.java | 38 +-
.../apache/arrow/vector/ipc/JsonFileReader.java | 24 +-
.../apache/arrow/vector/ipc/JsonFileWriter.java | 13 +-
.../arrow/vector/ipc/message/ArrowRecordBatch.java | 2 +-
.../java/org/apache/arrow/vector/types/Types.java | 19 +
.../{ => vector}/util/DataSizeRoundingUtil.java | 2 +-
.../apache/arrow/vector/util/DecimalUtility.java | 50 +-
.../vector/validate/ValidateVectorTypeVisitor.java | 3 +-
...ecimalVector.java => TestDecimal256Vector.java} | 112 ++--
.../org/apache/arrow/vector/TestDecimalVector.java | 20 +-
.../org/apache/arrow/vector/TestTypeLayout.java | 6 +-
.../org/apache/arrow/vector/TestVectorAlloc.java | 2 +-
.../org/apache/arrow/vector/TestVectorReAlloc.java | 2 +-
.../vector/complex/impl/TestComplexCopier.java | 18 +-
.../vector/complex/writer/TestComplexWriter.java | 6 +-
.../apache/arrow/vector/types/pojo/TestSchema.java | 2 +-
.../arrow/vector/util/DecimalUtilityTest.java | 119 ++--
.../util/TestDataSizeRoundingUtil.java | 2 +-
js/package.json | 2 +-
matlab/CMakeLists.txt | 2 +-
python/manylinux1/build_arrow.sh | 8 -
python/manylinux1/scripts/install_cmake.sh | 2 +-
python/manylinux201x/build_arrow.sh | 8 -
python/manylinux201x/scripts/install_cmake.sh | 3 +-
python/pyarrow/__init__.py | 11 +-
python/pyarrow/_compute.pxd | 10 -
python/pyarrow/_compute.pyx | 254 ++++++--
python/pyarrow/_dataset.pyx | 11 +-
python/pyarrow/array.pxi | 9 +-
python/pyarrow/compute.py | 158 +++--
python/pyarrow/includes/libarrow.pxd | 52 +-
python/pyarrow/lib.pxd | 9 +
python/pyarrow/lib.pyx | 3 +-
python/pyarrow/plasma.py | 4 +
python/pyarrow/public-api.pxi | 4 +-
python/pyarrow/scalar.pxi | 23 +-
python/pyarrow/tests/strategies.py | 12 +-
python/pyarrow/tests/test_array.py | 10 +-
python/pyarrow/tests/test_compute.py | 125 +++-
python/pyarrow/tests/test_convert_builtin.py | 41 +-
python/pyarrow/tests/test_scalars.py | 23 +-
python/pyarrow/tests/test_schema.py | 1 +
python/pyarrow/tests/test_types.py | 33 +-
python/pyarrow/types.pxi | 49 +-
python/pyarrow/types.py | 17 +-
python/setup.py | 10 +-
r/.Rbuildignore | 1 +
r/.gitignore | 1 +
r/DESCRIPTION | 4 +-
r/NEWS.md | 5 +-
r/R/enums.R | 25 +-
r/R/parquet.R | 22 +-
r/man/write_parquet.Rd | 12 +-
r/tools/autobrew | 2 +-
r/tools/linuxlibs.R | 6 +
ruby/red-arrow-cuda/lib/arrow-cuda/version.rb | 2 +-
.../red-arrow-dataset/lib/arrow-dataset/version.rb | 2 +-
ruby/red-arrow/lib/arrow/version.rb | 2 +-
ruby/red-arrow/test/test-struct-array-builder.rb | 16 +-
ruby/red-arrow/test/test-struct-array.rb | 4 +-
ruby/red-gandiva/lib/gandiva/version.rb | 2 +-
ruby/red-parquet/lib/parquet/version.rb | 2 +-
ruby/red-plasma/lib/plasma/version.rb | 2 +-
rust/arrow-flight/Cargo.toml | 4 +-
rust/arrow-flight/src/utils.rs | 2 +-
rust/arrow/Cargo.toml | 2 +-
rust/arrow/examples/tensor_builder.rs | 2 +-
rust/arrow/src/array/array.rs | 271 ++++----
rust/arrow/src/array/builder.rs | 208 +++---
rust/arrow/src/array/data.rs | 33 +-
rust/arrow/src/array/equal.rs | 39 +-
rust/arrow/src/array/iterator.rs | 2 +-
rust/arrow/src/array/mod.rs | 3 +-
rust/arrow/src/array/union.rs | 6 +-
rust/arrow/src/bitmap.rs | 6 +-
rust/arrow/src/buffer.rs | 22 +-
rust/arrow/src/compute/kernels/arithmetic.rs | 8 +-
rust/arrow/src/compute/kernels/boolean.rs | 1 -
rust/arrow/src/compute/kernels/cast.rs | 4 +-
rust/arrow/src/compute/kernels/filter.rs | 373 ++++++++++-
rust/arrow/src/compute/kernels/sort.rs | 2 +-
rust/arrow/src/compute/kernels/take.rs | 2 +-
rust/arrow/src/datatypes.rs | 32 +-
rust/arrow/src/json/reader.rs | 595 +++++++++++++----
rust/arrow/src/lib.rs | 7 +-
rust/arrow/src/record_batch.rs | 2 +-
rust/arrow/src/tensor.rs | 4 +-
rust/arrow/src/util/bit_chunk_iterator.rs | 4 +-
rust/arrow/src/util/bit_util.rs | 4 +-
rust/arrow/src/util/display.rs | 2 +-
rust/arrow/test/data/list_string_dict_nested.json | 3 +
rust/benchmarks/Cargo.toml | 2 +-
rust/datafusion/Cargo.toml | 10 +-
rust/datafusion/README.md | 2 +-
rust/datafusion/examples/flight_server.rs | 2 +-
rust/datafusion/examples/simple_udaf.rs | 4 +-
rust/datafusion/examples/simple_udf.rs | 4 +-
rust/datafusion/src/datasource/csv.rs | 4 +-
rust/datafusion/src/datasource/memory.rs | 44 +-
rust/datafusion/src/datasource/parquet.rs | 11 +-
rust/datafusion/src/error.rs | 105 ++-
rust/datafusion/src/execution/context.rs | 37 +-
rust/datafusion/src/logical_plan/mod.rs | 28 +-
.../src/optimizer/projection_push_down.rs | 4 +-
rust/datafusion/src/optimizer/utils.rs | 8 +-
rust/datafusion/src/physical_plan/aggregates.rs | 10 +-
.../src/physical_plan/array_expressions.rs | 8 +-
rust/datafusion/src/physical_plan/common.rs | 46 +-
rust/datafusion/src/physical_plan/csv.rs | 39 +-
.../src/physical_plan/datetime_expressions.rs | 14 +-
.../src/physical_plan/distinct_expressions.rs | 6 +-
rust/datafusion/src/physical_plan/empty.rs | 16 +-
rust/datafusion/src/physical_plan/explain.rs | 14 +-
rust/datafusion/src/physical_plan/expressions.rs | 235 +++++--
rust/datafusion/src/physical_plan/filter.rs | 108 +--
rust/datafusion/src/physical_plan/functions.rs | 17 +-
rust/datafusion/src/physical_plan/group_scalar.rs | 14 +-
.../datafusion/src/physical_plan/hash_aggregate.rs | 232 ++++---
rust/datafusion/src/physical_plan/limit.rs | 39 +-
.../src/physical_plan/math_expressions.rs | 10 +-
rust/datafusion/src/physical_plan/memory.rs | 36 +-
rust/datafusion/src/physical_plan/merge.rs | 61 +-
rust/datafusion/src/physical_plan/mod.rs | 21 +-
rust/datafusion/src/physical_plan/parquet.rs | 48 +-
rust/datafusion/src/physical_plan/planner.rs | 40 +-
rust/datafusion/src/physical_plan/projection.rs | 74 ++-
rust/datafusion/src/physical_plan/sort.rs | 18 +-
.../src/physical_plan/string_expressions.rs | 6 +-
rust/datafusion/src/physical_plan/type_coercion.rs | 10 +-
rust/datafusion/src/scalar.rs | 17 +-
rust/datafusion/src/sql/parser.rs | 2 +-
rust/datafusion/src/sql/planner.rs | 52 +-
rust/datafusion/src/test/mod.rs | 1 -
rust/datafusion/src/variable/mod.rs | 2 +-
rust/datafusion/tests/sql.rs | 74 ++-
rust/datafusion/tests/user_defined_plan.rs | 115 ++--
rust/integration-testing/Cargo.toml | 2 +-
rust/parquet/Cargo.toml | 6 +-
rust/parquet/README.md | 4 +-
rust/parquet/src/arrow/array_reader.rs | 6 +-
rust/parquet/src/arrow/arrow_writer.rs | 2 +-
rust/parquet/src/arrow/converter.rs | 4 +-
rust/parquet_derive/Cargo.toml | 4 +-
rust/parquet_derive/README.md | 4 +-
rust/parquet_derive_test/Cargo.toml | 6 +-
395 files changed, 10971 insertions(+), 3342 deletions(-)
delete mode 100644 cpp/src/arrow/compute/kernels/scalar_string_internal.h
create mode 100644 cpp/src/arrow/vendored/fast_float/README.md
create mode 100644 cpp/src/arrow/vendored/fast_float/ascii_number.h
create mode 100644 cpp/src/arrow/vendored/fast_float/decimal_to_binary.h
create mode 100644 cpp/src/arrow/vendored/fast_float/fast_float.h
create mode 100644 cpp/src/arrow/vendored/fast_float/fast_table.h
create mode 100644 cpp/src/arrow/vendored/fast_float/float_common.h
create mode 100644 cpp/src/arrow/vendored/fast_float/parse_number.h
create mode 100644 cpp/src/arrow/vendored/fast_float/thompson_tao.h
rename dev/tasks/linux-packages/apache-arrow/{debian/libarrow-cuda-glib200.install => debian.ubuntu-xenial/libarrow-cuda-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libarrow-cuda200.install => debian.ubuntu-xenial/libarrow-cuda300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libarrow-dataset-glib200.install => debian.ubuntu-xenial/libarrow-dataset-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libarrow-dataset200.install => debian.ubuntu-xenial/libarrow-dataset300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libarrow-glib200.install => debian.ubuntu-xenial/libarrow-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libarrow-python200.install => debian.ubuntu-xenial/libarrow-python300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libarrow200.install => debian.ubuntu-xenial/libarrow300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libgandiva-glib200.install => debian.ubuntu-xenial/libgandiva-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libgandiva200.install => debian.ubuntu-xenial/libgandiva300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libparquet-glib200.install => debian.ubuntu-xenial/libparquet-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libparquet200.install => debian.ubuntu-xenial/libparquet300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libplasma-glib200.install => debian.ubuntu-xenial/libplasma-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian/libplasma200.install => debian.ubuntu-xenial/libplasma300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libarrow-cuda-glib200.install => debian/libarrow-cuda-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libarrow-cuda200.install => debian/libarrow-cuda300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libarrow-dataset-glib200.install => debian/libarrow-dataset-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libarrow-dataset200.install => debian/libarrow-dataset300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/debian/{libarrow-flight200.install => libarrow-flight300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libarrow-glib200.install => debian/libarrow-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/debian/{libarrow-python-flight200.install => libarrow-python-flight300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libarrow-python200.install => debian/libarrow-python300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libarrow200.install => debian/libarrow300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libgandiva-glib200.install => debian/libgandiva-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libgandiva200.install => debian/libgandiva300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libparquet-glib200.install => debian/libparquet-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libparquet200.install => debian/libparquet300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libplasma-glib200.install => debian/libplasma-glib300.install} (100%)
rename dev/tasks/linux-packages/apache-arrow/{debian.ubuntu-xenial/libplasma200.install => debian/libplasma300.install} (100%)
create mode 100644 java/flight/flight-core/src/main/java/org/apache/arrow/flight/BackpressureStrategy.java
rename java/{vector => memory/memory-core}/src/main/java/org/apache/arrow/util/AutoCloseables.java (100%)
rename java/{vector => memory/memory-core}/src/main/java/org/apache/arrow/util/Collections2.java (100%)
copy java/vector/src/main/java/org/apache/arrow/vector/{DecimalVector.java => Decimal256Vector.java} (84%)
rename java/vector/src/main/java/org/apache/arrow/{ => vector}/util/DataSizeRoundingUtil.java (98%)
copy java/vector/src/test/java/org/apache/arrow/vector/{TestDecimalVector.java => TestDecimal256Vector.java} (75%)
rename java/vector/src/test/java/org/apache/arrow/{ => vector}/util/TestDataSizeRoundingUtil.java (98%)
create mode 100644 rust/arrow/test/data/list_string_dict_nested.json
[arrow] 03/09: ARROW-10095: [Rust] Update rust-parquet-arrow-writer
branch's encode_arrow_schema with ipc changes
Posted by ne...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit 98d4ffcdbaa937031d30bbbb934c3dc45cdefad4
Author: Carol (Nichols || Goulding) <ca...@gmail.com>
AuthorDate: Fri Sep 25 17:54:11 2020 +0200
ARROW-10095: [Rust] Update rust-parquet-arrow-writer branch's encode_arrow_schema with ipc changes
Note that this PR is deliberately filed against the rust-parquet-arrow-writer branch, not master!!
Hi! 👋 I'm looking to help out with the rust-parquet-arrow-writer branch, and I just pulled it down and it wasn't compiling because in 75f804efbfe367175fef5a2238d9cd2d30ed3afe, `schema_to_bytes` was changed to take `IpcWriteOptions` and to return `EncodedData`. This updates `encode_arrow_schema` to use those changes, which should get this branch compiling and passing tests again.
I'm kind of guessing which JIRA ticket this should be associated with; honestly I think this commit can just be squashed with https://github.com/apache/arrow/commit/8f0ed91469f2e569472edaa3b69ffde051088555 next time this branch gets rebased.
Please let me know if I should change anything, I'm happy to!
Closes #8274 from carols10cents/update-with-ipc-changes
Authored-by: Carol (Nichols || Goulding) <ca...@gmail.com>
Signed-off-by: Neville Dipale <ne...@gmail.com>
---
rust/parquet/src/arrow/arrow_writer.rs | 2 +-
rust/parquet/src/arrow/schema.rs | 8 +++++---
2 files changed, 6 insertions(+), 4 deletions(-)
diff --git a/rust/parquet/src/arrow/arrow_writer.rs b/rust/parquet/src/arrow/arrow_writer.rs
index 1ca8d50..e0ad207 100644
--- a/rust/parquet/src/arrow/arrow_writer.rs
+++ b/rust/parquet/src/arrow/arrow_writer.rs
@@ -22,7 +22,7 @@ use std::rc::Rc;
use arrow::array as arrow_array;
use arrow::datatypes::{DataType as ArrowDataType, SchemaRef};
use arrow::record_batch::RecordBatch;
-use arrow_array::Array;
+use arrow_array::{Array, PrimitiveArrayOps};
use super::schema::add_encoded_arrow_schema_to_metadata;
use crate::column::writer::ColumnWriter;
diff --git a/rust/parquet/src/arrow/schema.rs b/rust/parquet/src/arrow/schema.rs
index d4cfe1f..d5a0ff9 100644
--- a/rust/parquet/src/arrow/schema.rs
+++ b/rust/parquet/src/arrow/schema.rs
@@ -27,6 +27,7 @@ use std::collections::{HashMap, HashSet};
use std::rc::Rc;
use arrow::datatypes::{DataType, DateUnit, Field, Schema, TimeUnit};
+use arrow::ipc::writer;
use crate::basic::{LogicalType, Repetition, Type as PhysicalType};
use crate::errors::{ParquetError::ArrowError, Result};
@@ -120,15 +121,16 @@ fn get_arrow_schema_from_metadata(encoded_meta: &str) -> Option<Schema> {
/// Encodes the Arrow schema into the IPC format, and base64 encodes it
fn encode_arrow_schema(schema: &Schema) -> String {
- let mut serialized_schema = arrow::ipc::writer::schema_to_bytes(&schema);
+ let options = writer::IpcWriteOptions::default();
+ let mut serialized_schema = arrow::ipc::writer::schema_to_bytes(&schema, &options);
// manually prepending the length to the schema as arrow uses the legacy IPC format
// TODO: change after addressing ARROW-9777
- let schema_len = serialized_schema.len();
+ let schema_len = serialized_schema.ipc_message.len();
let mut len_prefix_schema = Vec::with_capacity(schema_len + 8);
len_prefix_schema.append(&mut vec![255u8, 255, 255, 255]);
len_prefix_schema.append((schema_len as u32).to_le_bytes().to_vec().as_mut());
- len_prefix_schema.append(&mut serialized_schema);
+ len_prefix_schema.append(&mut serialized_schema.ipc_message);
base64::encode(&len_prefix_schema)
}
[arrow] 05/09: ARROW-10191: [Rust] [Parquet] Add roundtrip Arrow ->
Parquet tests for all supported Arrow DataTypes
Posted by ne...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit 6677a4103374c7c685f0b9fb202bac79cea2d461
Author: Carol (Nichols || Goulding) <ca...@gmail.com>
AuthorDate: Tue Oct 6 08:44:26 2020 -0600
ARROW-10191: [Rust] [Parquet] Add roundtrip Arrow -> Parquet tests for all supported Arrow DataTypes
Note that this PR goes to the rust-parquet-arrow-writer branch, not master.
Inspired by tests in cpp/src/parquet/arrow/arrow_reader_writer_test.cc
These perform round-trip Arrow -> Parquet -> Arrow of a single RecordBatch with a single column of values of each the supported data types and some of the unsupported ones.
Tests that currently fail are either marked with `#[should_panic]` (if the
reason they fail is because of a panic) or `#[ignore]` (if the reason
they fail is because the values don't match).
I am comparing the RecordBatch's column's data before and after the round trip directly; I'm not sure that this is appropriate or not because for some data types, the `null_bitmap` isn't matching and I'm not sure if it's supposed to or not.
So I would love advice on that front, and I would love to know if these tests are useful or not!
Closes #8330 from carols10cents/roundtrip-tests
Lead-authored-by: Carol (Nichols || Goulding) <ca...@gmail.com>
Co-authored-by: Neville Dipale <ne...@gmail.com>
Signed-off-by: Andy Grove <an...@nvidia.com>
---
rust/parquet/src/arrow/array_reader.rs | 102 +++++---
rust/parquet/src/arrow/arrow_writer.rs | 413 ++++++++++++++++++++++++++++++++-
rust/parquet/src/arrow/converter.rs | 25 +-
3 files changed, 505 insertions(+), 35 deletions(-)
diff --git a/rust/parquet/src/arrow/array_reader.rs b/rust/parquet/src/arrow/array_reader.rs
index 14bf7d2..4fbc54d 100644
--- a/rust/parquet/src/arrow/array_reader.rs
+++ b/rust/parquet/src/arrow/array_reader.rs
@@ -35,9 +35,10 @@ use crate::arrow::converter::{
BinaryArrayConverter, BinaryConverter, BoolConverter, BooleanArrayConverter,
Converter, Date32Converter, FixedLenBinaryConverter, FixedSizeArrayConverter,
Float32Converter, Float64Converter, Int16Converter, Int32Converter, Int64Converter,
- Int8Converter, Int96ArrayConverter, Int96Converter, TimestampMicrosecondConverter,
- TimestampMillisecondConverter, UInt16Converter, UInt32Converter, UInt64Converter,
- UInt8Converter, Utf8ArrayConverter, Utf8Converter,
+ Int8Converter, Int96ArrayConverter, Int96Converter, Time32MillisecondConverter,
+ Time32SecondConverter, Time64MicrosecondConverter, Time64NanosecondConverter,
+ TimestampMicrosecondConverter, TimestampMillisecondConverter, UInt16Converter,
+ UInt32Converter, UInt64Converter, UInt8Converter, Utf8ArrayConverter, Utf8Converter,
};
use crate::arrow::record_reader::RecordReader;
use crate::arrow::schema::parquet_to_arrow_field;
@@ -196,11 +197,27 @@ impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
.convert(self.record_reader.cast::<Int32Type>()),
_ => Err(general_err!("No conversion from parquet type to arrow type for date with unit {:?}", unit)),
}
- (ArrowType::Time32(_), PhysicalType::INT32) => {
- UInt32Converter::new().convert(self.record_reader.cast::<Int32Type>())
+ (ArrowType::Time32(unit), PhysicalType::INT32) => {
+ match unit {
+ TimeUnit::Second => {
+ Time32SecondConverter::new().convert(self.record_reader.cast::<Int32Type>())
+ }
+ TimeUnit::Millisecond => {
+ Time32MillisecondConverter::new().convert(self.record_reader.cast::<Int32Type>())
+ }
+ _ => Err(general_err!("Invalid or unsupported arrow array with datatype {:?}", self.get_data_type()))
+ }
}
- (ArrowType::Time64(_), PhysicalType::INT64) => {
- UInt64Converter::new().convert(self.record_reader.cast::<Int64Type>())
+ (ArrowType::Time64(unit), PhysicalType::INT64) => {
+ match unit {
+ TimeUnit::Microsecond => {
+ Time64MicrosecondConverter::new().convert(self.record_reader.cast::<Int64Type>())
+ }
+ TimeUnit::Nanosecond => {
+ Time64NanosecondConverter::new().convert(self.record_reader.cast::<Int64Type>())
+ }
+ _ => Err(general_err!("Invalid or unsupported arrow array with datatype {:?}", self.get_data_type()))
+ }
}
(ArrowType::Interval(IntervalUnit::YearMonth), PhysicalType::INT32) => {
UInt32Converter::new().convert(self.record_reader.cast::<Int32Type>())
@@ -941,10 +958,12 @@ mod tests {
use crate::util::test_common::{get_test_file, make_pages};
use arrow::array::{Array, ArrayRef, PrimitiveArray, StringArray, StructArray};
use arrow::datatypes::{
- DataType as ArrowType, Date32Type as ArrowDate32, Field, Int32Type as ArrowInt32,
+ ArrowPrimitiveType, DataType as ArrowType, Date32Type as ArrowDate32, Field,
+ Int32Type as ArrowInt32, Int64Type as ArrowInt64,
+ Time32MillisecondType as ArrowTime32MillisecondArray,
+ Time64MicrosecondType as ArrowTime64MicrosecondArray,
TimestampMicrosecondType as ArrowTimestampMicrosecondType,
TimestampMillisecondType as ArrowTimestampMillisecondType,
- UInt32Type as ArrowUInt32, UInt64Type as ArrowUInt64,
};
use rand::distributions::uniform::SampleUniform;
use rand::{thread_rng, Rng};
@@ -1101,7 +1120,7 @@ mod tests {
}
macro_rules! test_primitive_array_reader_one_type {
- ($arrow_parquet_type:ty, $physical_type:expr, $logical_type_str:expr, $result_arrow_type:ty, $result_primitive_type:ty) => {{
+ ($arrow_parquet_type:ty, $physical_type:expr, $logical_type_str:expr, $result_arrow_type:ty, $result_arrow_cast_type:ty, $result_primitive_type:ty) => {{
let message_type = format!(
"
message test_schema {{
@@ -1112,7 +1131,7 @@ mod tests {
);
let schema = parse_message_type(&message_type)
.map(|t| Rc::new(SchemaDescriptor::new(Rc::new(t))))
- .unwrap();
+ .expect("Unable to parse message type into a schema descriptor");
let column_desc = schema.column(0);
@@ -1142,24 +1161,48 @@ mod tests {
Box::new(page_iterator),
column_desc.clone(),
)
- .unwrap();
+ .expect("Unable to get array reader");
- let array = array_reader.next_batch(50).unwrap();
+ let array = array_reader
+ .next_batch(50)
+ .expect("Unable to get batch from reader");
+ let result_data_type = <$result_arrow_type>::get_data_type();
let array = array
.as_any()
.downcast_ref::<PrimitiveArray<$result_arrow_type>>()
- .unwrap();
-
- assert_eq!(
- &PrimitiveArray::<$result_arrow_type>::from(
- data[0..50]
- .iter()
- .map(|x| *x as $result_primitive_type)
- .collect::<Vec<$result_primitive_type>>()
- ),
- array
+ .expect(
+ format!(
+ "Unable to downcast {:?} to {:?}",
+ array.data_type(),
+ result_data_type
+ )
+ .as_str(),
+ );
+
+ // create expected array as primitive, and cast to result type
+ let expected = PrimitiveArray::<$result_arrow_cast_type>::from(
+ data[0..50]
+ .iter()
+ .map(|x| *x as $result_primitive_type)
+ .collect::<Vec<$result_primitive_type>>(),
);
+ let expected = Arc::new(expected) as ArrayRef;
+ let expected = arrow::compute::cast(&expected, &result_data_type)
+ .expect("Unable to cast expected array");
+ assert_eq!(expected.data_type(), &result_data_type);
+ let expected = expected
+ .as_any()
+ .downcast_ref::<PrimitiveArray<$result_arrow_type>>()
+ .expect(
+ format!(
+ "Unable to downcast expected {:?} to {:?}",
+ expected.data_type(),
+ result_data_type
+ )
+ .as_str(),
+ );
+ assert_eq!(expected, array);
}
}};
}
@@ -1171,27 +1214,31 @@ mod tests {
PhysicalType::INT32,
"DATE",
ArrowDate32,
+ ArrowInt32,
i32
);
test_primitive_array_reader_one_type!(
Int32Type,
PhysicalType::INT32,
"TIME_MILLIS",
- ArrowUInt32,
- u32
+ ArrowTime32MillisecondArray,
+ ArrowInt32,
+ i32
);
test_primitive_array_reader_one_type!(
Int64Type,
PhysicalType::INT64,
"TIME_MICROS",
- ArrowUInt64,
- u64
+ ArrowTime64MicrosecondArray,
+ ArrowInt64,
+ i64
);
test_primitive_array_reader_one_type!(
Int64Type,
PhysicalType::INT64,
"TIMESTAMP_MILLIS",
ArrowTimestampMillisecondType,
+ ArrowInt64,
i64
);
test_primitive_array_reader_one_type!(
@@ -1199,6 +1246,7 @@ mod tests {
PhysicalType::INT64,
"TIMESTAMP_MICROS",
ArrowTimestampMicrosecondType,
+ ArrowInt64,
i64
);
}
diff --git a/rust/parquet/src/arrow/arrow_writer.rs b/rust/parquet/src/arrow/arrow_writer.rs
index e0ad207..cf7b9a2 100644
--- a/rust/parquet/src/arrow/arrow_writer.rs
+++ b/rust/parquet/src/arrow/arrow_writer.rs
@@ -136,7 +136,6 @@ fn write_leaves(
| ArrowDataType::UInt16
| ArrowDataType::UInt32
| ArrowDataType::UInt64
- | ArrowDataType::Float16
| ArrowDataType::Float32
| ArrowDataType::Float64
| ArrowDataType::Timestamp(_, _)
@@ -176,6 +175,9 @@ fn write_leaves(
}
Ok(())
}
+ ArrowDataType::Float16 => Err(ParquetError::ArrowError(
+ "Float16 arrays not supported".to_string(),
+ )),
ArrowDataType::FixedSizeList(_, _)
| ArrowDataType::Null
| ArrowDataType::Boolean
@@ -493,7 +495,7 @@ mod tests {
use arrow::array::*;
use arrow::datatypes::ToByteSlice;
use arrow::datatypes::{DataType, Field, Schema};
- use arrow::record_batch::{RecordBatch, RecordBatchReader};
+ use arrow::record_batch::RecordBatch;
use crate::arrow::{ArrowReader, ParquetFileArrowReader};
use crate::file::{metadata::KeyValue, reader::SerializedFileReader};
@@ -597,7 +599,7 @@ mod tests {
let mut arrow_reader = ParquetFileArrowReader::new(Rc::new(file_reader));
let mut record_batch_reader = arrow_reader.get_record_reader(1024).unwrap();
- let batch = record_batch_reader.next_batch().unwrap().unwrap();
+ let batch = record_batch_reader.next().unwrap().unwrap();
let string_col = batch
.column(0)
.as_any()
@@ -688,4 +690,409 @@ mod tests {
writer.write(&batch).unwrap();
writer.close().unwrap();
}
+
+ const SMALL_SIZE: usize = 100;
+
+ fn roundtrip(filename: &str, expected_batch: RecordBatch) {
+ let file = get_temp_file(filename, &[]);
+
+ let mut writer = ArrowWriter::try_new(
+ file.try_clone().unwrap(),
+ expected_batch.schema(),
+ None,
+ )
+ .unwrap();
+ writer.write(&expected_batch).unwrap();
+ writer.close().unwrap();
+
+ let reader = SerializedFileReader::new(file).unwrap();
+ let mut arrow_reader = ParquetFileArrowReader::new(Rc::new(reader));
+ let mut record_batch_reader = arrow_reader.get_record_reader(1024).unwrap();
+
+ let actual_batch = record_batch_reader.next().unwrap().unwrap();
+
+ assert_eq!(expected_batch.schema(), actual_batch.schema());
+ assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
+ assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
+ for i in 0..expected_batch.num_columns() {
+ let expected_data = expected_batch.column(i).data();
+ let actual_data = actual_batch.column(i).data();
+
+ assert_eq!(expected_data.data_type(), actual_data.data_type());
+ assert_eq!(expected_data.len(), actual_data.len());
+ assert_eq!(expected_data.null_count(), actual_data.null_count());
+ assert_eq!(expected_data.offset(), actual_data.offset());
+ assert_eq!(expected_data.buffers(), actual_data.buffers());
+ assert_eq!(expected_data.child_data(), actual_data.child_data());
+ assert_eq!(expected_data.null_bitmap(), actual_data.null_bitmap());
+ }
+ }
+
+ fn one_column_roundtrip(filename: &str, values: ArrayRef, nullable: bool) {
+ let schema = Schema::new(vec![Field::new(
+ "col",
+ values.data_type().clone(),
+ nullable,
+ )]);
+ let expected_batch =
+ RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
+
+ roundtrip(filename, expected_batch);
+ }
+
+ fn values_required<A, I>(iter: I, filename: &str)
+ where
+ A: From<Vec<I::Item>> + Array + 'static,
+ I: IntoIterator,
+ {
+ let raw_values: Vec<_> = iter.into_iter().collect();
+ let values = Arc::new(A::from(raw_values));
+ one_column_roundtrip(filename, values, false);
+ }
+
+ fn values_optional<A, I>(iter: I, filename: &str)
+ where
+ A: From<Vec<Option<I::Item>>> + Array + 'static,
+ I: IntoIterator,
+ {
+ let optional_raw_values: Vec<_> = iter
+ .into_iter()
+ .enumerate()
+ .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
+ .collect();
+ let optional_values = Arc::new(A::from(optional_raw_values));
+ one_column_roundtrip(filename, optional_values, true);
+ }
+
+ fn required_and_optional<A, I>(iter: I, filename: &str)
+ where
+ A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
+ I: IntoIterator + Clone,
+ {
+ values_required::<A, I>(iter.clone(), filename);
+ values_optional::<A, I>(iter, filename);
+ }
+
+ #[test]
+ #[should_panic(expected = "Null arrays not supported")]
+ fn null_single_column() {
+ let values = Arc::new(NullArray::new(SMALL_SIZE));
+ one_column_roundtrip("null_single_column", values.clone(), true);
+ one_column_roundtrip("null_single_column", values, false);
+ }
+
+ #[test]
+ #[should_panic(
+ expected = "Attempting to write an Arrow type that is not yet implemented"
+ )]
+ fn bool_single_column() {
+ required_and_optional::<BooleanArray, _>(
+ [true, false].iter().cycle().copied().take(SMALL_SIZE),
+ "bool_single_column",
+ );
+ }
+
+ #[test]
+ fn i8_single_column() {
+ required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8, "i8_single_column");
+ }
+
+ #[test]
+ fn i16_single_column() {
+ required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16, "i16_single_column");
+ }
+
+ #[test]
+ fn i32_single_column() {
+ required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32, "i32_single_column");
+ }
+
+ #[test]
+ fn i64_single_column() {
+ required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64, "i64_single_column");
+ }
+
+ #[test]
+ fn u8_single_column() {
+ required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8, "u8_single_column");
+ }
+
+ #[test]
+ fn u16_single_column() {
+ required_and_optional::<UInt16Array, _>(
+ 0..SMALL_SIZE as u16,
+ "u16_single_column",
+ );
+ }
+
+ #[test]
+ fn u32_single_column() {
+ required_and_optional::<UInt32Array, _>(
+ 0..SMALL_SIZE as u32,
+ "u32_single_column",
+ );
+ }
+
+ #[test]
+ fn u64_single_column() {
+ required_and_optional::<UInt64Array, _>(
+ 0..SMALL_SIZE as u64,
+ "u64_single_column",
+ );
+ }
+
+ #[test]
+ fn f32_single_column() {
+ required_and_optional::<Float32Array, _>(
+ (0..SMALL_SIZE).map(|i| i as f32),
+ "f32_single_column",
+ );
+ }
+
+ #[test]
+ fn f64_single_column() {
+ required_and_optional::<Float64Array, _>(
+ (0..SMALL_SIZE).map(|i| i as f64),
+ "f64_single_column",
+ );
+ }
+
+ // The timestamp array types don't implement From<Vec<T>> because they need the timezone
+ // argument, and they also doesn't support building from a Vec<Option<T>>, so call
+ // one_column_roundtrip manually instead of calling required_and_optional for these tests.
+
+ #[test]
+ #[ignore] // Timestamp support isn't correct yet
+ fn timestamp_second_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
+ let values = Arc::new(TimestampSecondArray::from_vec(raw_values, None));
+
+ one_column_roundtrip("timestamp_second_single_column", values, false);
+ }
+
+ #[test]
+ fn timestamp_millisecond_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
+ let values = Arc::new(TimestampMillisecondArray::from_vec(raw_values, None));
+
+ one_column_roundtrip("timestamp_millisecond_single_column", values, false);
+ }
+
+ #[test]
+ fn timestamp_microsecond_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
+ let values = Arc::new(TimestampMicrosecondArray::from_vec(raw_values, None));
+
+ one_column_roundtrip("timestamp_microsecond_single_column", values, false);
+ }
+
+ #[test]
+ #[ignore] // Timestamp support isn't correct yet
+ fn timestamp_nanosecond_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
+ let values = Arc::new(TimestampNanosecondArray::from_vec(raw_values, None));
+
+ one_column_roundtrip("timestamp_nanosecond_single_column", values, false);
+ }
+
+ #[test]
+ fn date32_single_column() {
+ required_and_optional::<Date32Array, _>(
+ 0..SMALL_SIZE as i32,
+ "date32_single_column",
+ );
+ }
+
+ #[test]
+ #[ignore] // Date support isn't correct yet
+ fn date64_single_column() {
+ required_and_optional::<Date64Array, _>(
+ 0..SMALL_SIZE as i64,
+ "date64_single_column",
+ );
+ }
+
+ #[test]
+ #[ignore] // DateUnit resolution mismatch
+ fn time32_second_single_column() {
+ required_and_optional::<Time32SecondArray, _>(
+ 0..SMALL_SIZE as i32,
+ "time32_second_single_column",
+ );
+ }
+
+ #[test]
+ fn time32_millisecond_single_column() {
+ required_and_optional::<Time32MillisecondArray, _>(
+ 0..SMALL_SIZE as i32,
+ "time32_millisecond_single_column",
+ );
+ }
+
+ #[test]
+ fn time64_microsecond_single_column() {
+ required_and_optional::<Time64MicrosecondArray, _>(
+ 0..SMALL_SIZE as i64,
+ "time64_microsecond_single_column",
+ );
+ }
+
+ #[test]
+ #[ignore] // DateUnit resolution mismatch
+ fn time64_nanosecond_single_column() {
+ required_and_optional::<Time64NanosecondArray, _>(
+ 0..SMALL_SIZE as i64,
+ "time64_nanosecond_single_column",
+ );
+ }
+
+ #[test]
+ #[should_panic(expected = "Converting Duration to parquet not supported")]
+ fn duration_second_single_column() {
+ required_and_optional::<DurationSecondArray, _>(
+ 0..SMALL_SIZE as i64,
+ "duration_second_single_column",
+ );
+ }
+
+ #[test]
+ #[should_panic(expected = "Converting Duration to parquet not supported")]
+ fn duration_millisecond_single_column() {
+ required_and_optional::<DurationMillisecondArray, _>(
+ 0..SMALL_SIZE as i64,
+ "duration_millisecond_single_column",
+ );
+ }
+
+ #[test]
+ #[should_panic(expected = "Converting Duration to parquet not supported")]
+ fn duration_microsecond_single_column() {
+ required_and_optional::<DurationMicrosecondArray, _>(
+ 0..SMALL_SIZE as i64,
+ "duration_microsecond_single_column",
+ );
+ }
+
+ #[test]
+ #[should_panic(expected = "Converting Duration to parquet not supported")]
+ fn duration_nanosecond_single_column() {
+ required_and_optional::<DurationNanosecondArray, _>(
+ 0..SMALL_SIZE as i64,
+ "duration_nanosecond_single_column",
+ );
+ }
+
+ #[test]
+ #[should_panic(expected = "Currently unreachable because data type not supported")]
+ fn interval_year_month_single_column() {
+ required_and_optional::<IntervalYearMonthArray, _>(
+ 0..SMALL_SIZE as i32,
+ "interval_year_month_single_column",
+ );
+ }
+
+ #[test]
+ #[should_panic(expected = "Currently unreachable because data type not supported")]
+ fn interval_day_time_single_column() {
+ required_and_optional::<IntervalDayTimeArray, _>(
+ 0..SMALL_SIZE as i64,
+ "interval_day_time_single_column",
+ );
+ }
+
+ #[test]
+ #[ignore] // Binary support isn't correct yet - null_bitmap doesn't match
+ fn binary_single_column() {
+ let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
+ let many_vecs: Vec<_> = std::iter::repeat(one_vec).take(SMALL_SIZE).collect();
+ let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
+
+ // BinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
+ values_required::<BinaryArray, _>(many_vecs_iter, "binary_single_column");
+ }
+
+ #[test]
+ #[ignore] // Large Binary support isn't correct yet
+ fn large_binary_single_column() {
+ let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
+ let many_vecs: Vec<_> = std::iter::repeat(one_vec).take(SMALL_SIZE).collect();
+ let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
+
+ // LargeBinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
+ values_required::<LargeBinaryArray, _>(
+ many_vecs_iter,
+ "large_binary_single_column",
+ );
+ }
+
+ #[test]
+ #[ignore] // String support isn't correct yet - null_bitmap doesn't match
+ fn string_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
+ let raw_strs = raw_values.iter().map(|s| s.as_str());
+
+ required_and_optional::<StringArray, _>(raw_strs, "string_single_column");
+ }
+
+ #[test]
+ #[ignore] // Large String support isn't correct yet - null_bitmap and buffers don't match
+ fn large_string_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
+ let raw_strs = raw_values.iter().map(|s| s.as_str());
+
+ required_and_optional::<LargeStringArray, _>(
+ raw_strs,
+ "large_string_single_column",
+ );
+ }
+
+ #[test]
+ #[should_panic(
+ expected = "Reading parquet list array into arrow is not supported yet!"
+ )]
+ fn list_single_column() {
+ let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+ let a_value_offsets =
+ arrow::buffer::Buffer::from(&[0, 1, 3, 3, 6, 10].to_byte_slice());
+ let a_list_data = ArrayData::builder(DataType::List(Box::new(DataType::Int32)))
+ .len(5)
+ .add_buffer(a_value_offsets)
+ .add_child_data(a_values.data())
+ .build();
+ let a = ListArray::from(a_list_data);
+
+ let values = Arc::new(a);
+ one_column_roundtrip("list_single_column", values, false);
+ }
+
+ #[test]
+ #[should_panic(
+ expected = "Reading parquet list array into arrow is not supported yet!"
+ )]
+ fn large_list_single_column() {
+ let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+ let a_value_offsets =
+ arrow::buffer::Buffer::from(&[0i64, 1, 3, 3, 6, 10].to_byte_slice());
+ let a_list_data =
+ ArrayData::builder(DataType::LargeList(Box::new(DataType::Int32)))
+ .len(5)
+ .add_buffer(a_value_offsets)
+ .add_child_data(a_values.data())
+ .build();
+ let a = LargeListArray::from(a_list_data);
+
+ let values = Arc::new(a);
+ one_column_roundtrip("large_list_single_column", values, false);
+ }
+
+ #[test]
+ #[ignore] // Struct support isn't correct yet - null_bitmap doesn't match
+ fn struct_single_column() {
+ let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+ let struct_field_a = Field::new("f", DataType::Int32, false);
+ let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]);
+
+ let values = Arc::new(s);
+ one_column_roundtrip("struct_single_column", values, false);
+ }
}
diff --git a/rust/parquet/src/arrow/converter.rs b/rust/parquet/src/arrow/converter.rs
index da0cc6c..86f6492 100644
--- a/rust/parquet/src/arrow/converter.rs
+++ b/rust/parquet/src/arrow/converter.rs
@@ -17,12 +17,19 @@
use crate::arrow::record_reader::RecordReader;
use crate::data_type::{ByteArray, DataType, Int96};
-use arrow::array::{
- Array, ArrayRef, BinaryBuilder, BooleanArray, BooleanBufferBuilder,
- BufferBuilderTrait, FixedSizeBinaryBuilder, StringBuilder,
- TimestampNanosecondBuilder,
+// TODO: clean up imports (best done when there are few moving parts)
+use arrow::{
+ array::{
+ Array, ArrayRef, BinaryBuilder, BooleanArray, BooleanBufferBuilder,
+ BufferBuilderTrait, FixedSizeBinaryBuilder, StringBuilder,
+ TimestampNanosecondBuilder,
+ },
+ datatypes::Time32MillisecondType,
+};
+use arrow::{
+ compute::cast, datatypes::Time32SecondType, datatypes::Time64MicrosecondType,
+ datatypes::Time64NanosecondType,
};
-use arrow::compute::cast;
use std::convert::From;
use std::sync::Arc;
@@ -226,6 +233,14 @@ pub type TimestampMillisecondConverter =
CastConverter<ParquetInt64Type, TimestampMillisecondType, TimestampMillisecondType>;
pub type TimestampMicrosecondConverter =
CastConverter<ParquetInt64Type, TimestampMicrosecondType, TimestampMicrosecondType>;
+pub type Time32SecondConverter =
+ CastConverter<ParquetInt32Type, Time32SecondType, Time32SecondType>;
+pub type Time32MillisecondConverter =
+ CastConverter<ParquetInt32Type, Time32MillisecondType, Time32MillisecondType>;
+pub type Time64MicrosecondConverter =
+ CastConverter<ParquetInt64Type, Time64MicrosecondType, Time64MicrosecondType>;
+pub type Time64NanosecondConverter =
+ CastConverter<ParquetInt64Type, Time64NanosecondType, Time64NanosecondType>;
pub type UInt64Converter = CastConverter<ParquetInt64Type, UInt64Type, UInt64Type>;
pub type Float32Converter = CastConverter<ParquetFloatType, Float32Type, Float32Type>;
pub type Float64Converter = CastConverter<ParquetDoubleType, Float64Type, Float64Type>;
[arrow] 09/09: ARROW-7842: [Rust] [Parquet] Arrow list reader
Posted by ne...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit 9ba2179889f3bac3101e8759ccaea22b26782092
Author: Neville Dipale <ne...@gmail.com>
AuthorDate: Sat Oct 17 21:14:13 2020 +0200
ARROW-7842: [Rust] [Parquet] Arrow list reader
This is a port of #6770 to the parquet-writer branch.
We'll have more of a chance to test this reader,and ensure that we can roundtrip on list types.
Closes #8449 from nevi-me/ARROW-7842-cherry
Authored-by: Neville Dipale <ne...@gmail.com>
Signed-off-by: Neville Dipale <ne...@gmail.com>
---
rust/arrow/src/util/display.rs | 17 +
rust/datafusion/tests/sql.rs | 97 ++++-
rust/parquet/src/arrow/array_reader.rs | 630 ++++++++++++++++++++++++++++++++-
rust/parquet/src/arrow/arrow_writer.rs | 29 +-
rust/parquet/src/schema/visitor.rs | 12 +-
5 files changed, 755 insertions(+), 30 deletions(-)
diff --git a/rust/arrow/src/util/display.rs b/rust/arrow/src/util/display.rs
index 102ec5d..1a873f1 100644
--- a/rust/arrow/src/util/display.rs
+++ b/rust/arrow/src/util/display.rs
@@ -44,6 +44,22 @@ macro_rules! make_string {
}};
}
+macro_rules! make_string_from_list {
+ ($column: ident, $row: ident) => {{
+ let list = $column
+ .as_any()
+ .downcast_ref::<array::ListArray>()
+ .ok_or(ArrowError::InvalidArgumentError(format!(
+ "Repl error: could not convert list column to list array."
+ )))?
+ .value($row);
+ let string_values = (0..list.len())
+ .map(|i| array_value_to_string(&list.clone(), i))
+ .collect::<Result<Vec<String>>>()?;
+ Ok(format!("[{}]", string_values.join(", ")))
+ }};
+}
+
/// Get the value at the given row in an array as a String.
///
/// Note this function is quite inefficient and is unlikely to be
@@ -89,6 +105,7 @@ pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result<Str
DataType::Time64(unit) if *unit == TimeUnit::Nanosecond => {
make_string!(array::Time64NanosecondArray, column, row)
}
+ DataType::List(_) => make_string_from_list!(column, row),
DataType::Dictionary(index_type, _value_type) => match **index_type {
DataType::Int8 => dict_array_value_to_string::<Int8Type>(column, row),
DataType::Int16 => dict_array_value_to_string::<Int16Type>(column, row),
diff --git a/rust/datafusion/tests/sql.rs b/rust/datafusion/tests/sql.rs
index 9eceb76..a35476d 100644
--- a/rust/datafusion/tests/sql.rs
+++ b/rust/datafusion/tests/sql.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+use std::convert::TryFrom;
use std::env;
use std::sync::Arc;
@@ -22,7 +23,7 @@ extern crate arrow;
extern crate datafusion;
use arrow::{array::*, datatypes::TimeUnit};
-use arrow::{datatypes::Int32Type, record_batch::RecordBatch};
+use arrow::{datatypes::Int32Type, datatypes::Int64Type, record_batch::RecordBatch};
use arrow::{
datatypes::{DataType, Field, Schema, SchemaRef},
util::display::array_value_to_string,
@@ -129,6 +130,100 @@ async fn parquet_single_nan_schema() {
}
#[tokio::test]
+async fn parquet_list_columns() {
+ let mut ctx = ExecutionContext::new();
+ let testdata = env::var("PARQUET_TEST_DATA").expect("PARQUET_TEST_DATA not defined");
+ ctx.register_parquet(
+ "list_columns",
+ &format!("{}/list_columns.parquet", testdata),
+ )
+ .unwrap();
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new(
+ "int64_list",
+ DataType::List(Box::new(DataType::Int64)),
+ true,
+ ),
+ Field::new("utf8_list", DataType::List(Box::new(DataType::Utf8)), true),
+ ]));
+
+ let sql = "SELECT int64_list, utf8_list FROM list_columns";
+ let plan = ctx.create_logical_plan(&sql).unwrap();
+ let plan = ctx.optimize(&plan).unwrap();
+ let plan = ctx.create_physical_plan(&plan).unwrap();
+ let results = ctx.collect(plan).await.unwrap();
+
+ // int64_list utf8_list
+ // 0 [1, 2, 3] [abc, efg, hij]
+ // 1 [None, 1] None
+ // 2 [4] [efg, None, hij, xyz]
+
+ assert_eq!(1, results.len());
+ let batch = &results[0];
+ assert_eq!(3, batch.num_rows());
+ assert_eq!(2, batch.num_columns());
+ assert_eq!(schema, batch.schema());
+
+ let int_list_array = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+ let utf8_list_array = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+
+ assert_eq!(
+ int_list_array
+ .value(0)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ &PrimitiveArray::<Int64Type>::from(vec![Some(1), Some(2), Some(3),])
+ );
+
+ assert_eq!(
+ utf8_list_array
+ .value(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap(),
+ &StringArray::try_from(vec![Some("abc"), Some("efg"), Some("hij"),]).unwrap()
+ );
+
+ assert_eq!(
+ int_list_array
+ .value(1)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ &PrimitiveArray::<Int64Type>::from(vec![None, Some(1),])
+ );
+
+ assert!(utf8_list_array.is_null(1));
+
+ assert_eq!(
+ int_list_array
+ .value(2)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<Int64Type>>()
+ .unwrap(),
+ &PrimitiveArray::<Int64Type>::from(vec![Some(4),])
+ );
+
+ let result = utf8_list_array.value(2);
+ let result = result.as_any().downcast_ref::<StringArray>().unwrap();
+
+ assert_eq!(result.value(0), "efg");
+ assert!(result.is_null(1));
+ assert_eq!(result.value(2), "hij");
+ assert_eq!(result.value(3), "xyz");
+}
+
+#[tokio::test]
async fn csv_count_star() -> Result<()> {
let mut ctx = ExecutionContext::new();
register_aggregate_csv(&mut ctx)?;
diff --git a/rust/parquet/src/arrow/array_reader.rs b/rust/parquet/src/arrow/array_reader.rs
index 6fdf5d5..9646e74 100644
--- a/rust/parquet/src/arrow/array_reader.rs
+++ b/rust/parquet/src/arrow/array_reader.rs
@@ -25,13 +25,35 @@ use std::sync::Arc;
use std::vec::Vec;
use arrow::array::{
- ArrayDataBuilder, ArrayDataRef, ArrayRef, BooleanBufferBuilder, BufferBuilderTrait,
- Int16BufferBuilder, StructArray,
+ Array, ArrayData, ArrayDataBuilder, ArrayDataRef, ArrayRef, BinaryArray,
+ BinaryBuilder, BooleanBufferBuilder, BufferBuilderTrait, FixedSizeBinaryArray,
+ FixedSizeBinaryBuilder, GenericListArray, Int16BufferBuilder, ListBuilder,
+ OffsetSizeTrait, PrimitiveArray, PrimitiveBuilder, StringArray, StringBuilder,
+ StructArray,
};
use arrow::buffer::{Buffer, MutableBuffer};
use arrow::datatypes::{
- DataType as ArrowType, DateUnit, Field, IntervalUnit, Schema, TimeUnit,
+ BooleanType as ArrowBooleanType, DataType as ArrowType,
+ Date32Type as ArrowDate32Type, Date64Type as ArrowDate64Type, DateUnit,
+ DurationMicrosecondType as ArrowDurationMicrosecondType,
+ DurationMillisecondType as ArrowDurationMillisecondType,
+ DurationNanosecondType as ArrowDurationNanosecondType,
+ DurationSecondType as ArrowDurationSecondType, Field,
+ Float32Type as ArrowFloat32Type, Float64Type as ArrowFloat64Type,
+ Int16Type as ArrowInt16Type, Int32Type as ArrowInt32Type,
+ Int64Type as ArrowInt64Type, Int8Type as ArrowInt8Type, IntervalUnit, Schema,
+ Time32MillisecondType as ArrowTime32MillisecondType,
+ Time32SecondType as ArrowTime32SecondType,
+ Time64MicrosecondType as ArrowTime64MicrosecondType,
+ Time64NanosecondType as ArrowTime64NanosecondType, TimeUnit,
+ TimeUnit as ArrowTimeUnit, TimestampMicrosecondType as ArrowTimestampMicrosecondType,
+ TimestampMillisecondType as ArrowTimestampMillisecondType,
+ TimestampNanosecondType as ArrowTimestampNanosecondType,
+ TimestampSecondType as ArrowTimestampSecondType, ToByteSlice,
+ UInt16Type as ArrowUInt16Type, UInt32Type as ArrowUInt32Type,
+ UInt64Type as ArrowUInt64Type, UInt8Type as ArrowUInt8Type,
};
+use arrow::util::bit_util;
use crate::arrow::converter::{
BinaryArrayConverter, BinaryConverter, BoolConverter, BooleanArrayConverter,
@@ -532,6 +554,400 @@ where
}
}
+/// Implementation of list array reader.
+pub struct ListArrayReader<OffsetSize: OffsetSizeTrait> {
+ item_reader: Box<dyn ArrayReader>,
+ data_type: ArrowType,
+ item_type: ArrowType,
+ list_def_level: i16,
+ list_rep_level: i16,
+ def_level_buffer: Option<Buffer>,
+ rep_level_buffer: Option<Buffer>,
+ _marker: PhantomData<OffsetSize>,
+}
+
+impl<OffsetSize: OffsetSizeTrait> ListArrayReader<OffsetSize> {
+ /// Construct list array reader.
+ pub fn new(
+ item_reader: Box<dyn ArrayReader>,
+ data_type: ArrowType,
+ item_type: ArrowType,
+ def_level: i16,
+ rep_level: i16,
+ ) -> Self {
+ Self {
+ item_reader,
+ data_type,
+ item_type,
+ list_def_level: def_level,
+ list_rep_level: rep_level,
+ def_level_buffer: None,
+ rep_level_buffer: None,
+ _marker: PhantomData,
+ }
+ }
+}
+
+macro_rules! build_empty_list_array_with_primitive_items {
+ ($item_type:ident) => {{
+ let values_builder = PrimitiveBuilder::<$item_type>::new(0);
+ let mut builder = ListBuilder::new(values_builder);
+ let empty_list_array = builder.finish();
+ Ok(Arc::new(empty_list_array))
+ }};
+}
+
+macro_rules! build_empty_list_array_with_non_primitive_items {
+ ($builder:ident) => {{
+ let values_builder = $builder::new(0);
+ let mut builder = ListBuilder::new(values_builder);
+ let empty_list_array = builder.finish();
+ Ok(Arc::new(empty_list_array))
+ }};
+}
+
+fn build_empty_list_array(item_type: ArrowType) -> Result<ArrayRef> {
+ match item_type {
+ ArrowType::UInt8 => build_empty_list_array_with_primitive_items!(ArrowUInt8Type),
+ ArrowType::UInt16 => {
+ build_empty_list_array_with_primitive_items!(ArrowUInt16Type)
+ }
+ ArrowType::UInt32 => {
+ build_empty_list_array_with_primitive_items!(ArrowUInt32Type)
+ }
+ ArrowType::UInt64 => {
+ build_empty_list_array_with_primitive_items!(ArrowUInt64Type)
+ }
+ ArrowType::Int8 => build_empty_list_array_with_primitive_items!(ArrowInt8Type),
+ ArrowType::Int16 => build_empty_list_array_with_primitive_items!(ArrowInt16Type),
+ ArrowType::Int32 => build_empty_list_array_with_primitive_items!(ArrowInt32Type),
+ ArrowType::Int64 => build_empty_list_array_with_primitive_items!(ArrowInt64Type),
+ ArrowType::Float32 => {
+ build_empty_list_array_with_primitive_items!(ArrowFloat32Type)
+ }
+ ArrowType::Float64 => {
+ build_empty_list_array_with_primitive_items!(ArrowFloat64Type)
+ }
+ ArrowType::Boolean => {
+ build_empty_list_array_with_primitive_items!(ArrowBooleanType)
+ }
+ ArrowType::Date32(_) => {
+ build_empty_list_array_with_primitive_items!(ArrowDate32Type)
+ }
+ ArrowType::Date64(_) => {
+ build_empty_list_array_with_primitive_items!(ArrowDate64Type)
+ }
+ ArrowType::Time32(ArrowTimeUnit::Second) => {
+ build_empty_list_array_with_primitive_items!(ArrowTime32SecondType)
+ }
+ ArrowType::Time32(ArrowTimeUnit::Millisecond) => {
+ build_empty_list_array_with_primitive_items!(ArrowTime32MillisecondType)
+ }
+ ArrowType::Time64(ArrowTimeUnit::Microsecond) => {
+ build_empty_list_array_with_primitive_items!(ArrowTime64MicrosecondType)
+ }
+ ArrowType::Time64(ArrowTimeUnit::Nanosecond) => {
+ build_empty_list_array_with_primitive_items!(ArrowTime64NanosecondType)
+ }
+ ArrowType::Duration(ArrowTimeUnit::Second) => {
+ build_empty_list_array_with_primitive_items!(ArrowDurationSecondType)
+ }
+ ArrowType::Duration(ArrowTimeUnit::Millisecond) => {
+ build_empty_list_array_with_primitive_items!(ArrowDurationMillisecondType)
+ }
+ ArrowType::Duration(ArrowTimeUnit::Microsecond) => {
+ build_empty_list_array_with_primitive_items!(ArrowDurationMicrosecondType)
+ }
+ ArrowType::Duration(ArrowTimeUnit::Nanosecond) => {
+ build_empty_list_array_with_primitive_items!(ArrowDurationNanosecondType)
+ }
+ ArrowType::Timestamp(ArrowTimeUnit::Second, _) => {
+ build_empty_list_array_with_primitive_items!(ArrowTimestampSecondType)
+ }
+ ArrowType::Timestamp(ArrowTimeUnit::Millisecond, _) => {
+ build_empty_list_array_with_primitive_items!(ArrowTimestampMillisecondType)
+ }
+ ArrowType::Timestamp(ArrowTimeUnit::Microsecond, _) => {
+ build_empty_list_array_with_primitive_items!(ArrowTimestampMicrosecondType)
+ }
+ ArrowType::Timestamp(ArrowTimeUnit::Nanosecond, _) => {
+ build_empty_list_array_with_primitive_items!(ArrowTimestampNanosecondType)
+ }
+ ArrowType::Utf8 => {
+ build_empty_list_array_with_non_primitive_items!(StringBuilder)
+ }
+ ArrowType::Binary => {
+ build_empty_list_array_with_non_primitive_items!(BinaryBuilder)
+ }
+ _ => Err(ParquetError::General(format!(
+ "ListArray of type List({:?}) is not supported by array_reader",
+ item_type
+ ))),
+ }
+}
+
+macro_rules! remove_primitive_array_indices {
+ ($arr: expr, $item_type:ty, $indices:expr) => {{
+ let array_data = match $arr.as_any().downcast_ref::<PrimitiveArray<$item_type>>() {
+ Some(a) => a,
+ _ => return Err(ParquetError::General(format!("Error generating next batch for ListArray: {:?} cannot be downcast to PrimitiveArray", $arr))),
+ };
+ let mut builder = PrimitiveBuilder::<$item_type>::new($arr.len());
+ for i in 0..array_data.len() {
+ if !$indices.contains(&i) {
+ if array_data.is_null(i) {
+ builder.append_null()?;
+ } else {
+ builder.append_value(array_data.value(i))?;
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()))
+ }};
+}
+
+macro_rules! remove_array_indices_custom_builder {
+ ($arr: expr, $array_type:ty, $item_builder:ident, $indices:expr) => {{
+ let array_data = match $arr.as_any().downcast_ref::<$array_type>() {
+ Some(a) => a,
+ _ => return Err(ParquetError::General(format!("Error generating next batch for ListArray: {:?} cannot be downcast to PrimitiveArray", $arr))),
+ };
+ let mut builder = $item_builder::new(array_data.len());
+
+ for i in 0..array_data.len() {
+ if !$indices.contains(&i) {
+ if array_data.is_null(i) {
+ builder.append_null()?;
+ } else {
+ builder.append_value(array_data.value(i))?;
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()))
+ }};
+}
+
+macro_rules! remove_fixed_size_binary_array_indices {
+ ($arr: expr, $array_type:ty, $item_builder:ident, $indices:expr, $len:expr) => {{
+ let array_data = match $arr.as_any().downcast_ref::<$array_type>() {
+ Some(a) => a,
+ _ => return Err(ParquetError::General(format!("Error generating next batch for ListArray: {:?} cannot be downcast to PrimitiveArray", $arr))),
+ };
+ let mut builder = FixedSizeBinaryBuilder::new(array_data.len(), $len);
+ for i in 0..array_data.len() {
+ if !$indices.contains(&i) {
+ if array_data.is_null(i) {
+ builder.append_null()?;
+ } else {
+ builder.append_value(array_data.value(i))?;
+ }
+ }
+ }
+ Ok(Arc::new(builder.finish()))
+ }};
+}
+
+fn remove_indices(
+ arr: ArrayRef,
+ item_type: ArrowType,
+ indices: Vec<usize>,
+) -> Result<ArrayRef> {
+ match item_type {
+ ArrowType::UInt8 => remove_primitive_array_indices!(arr, ArrowUInt8Type, indices),
+ ArrowType::UInt16 => {
+ remove_primitive_array_indices!(arr, ArrowUInt16Type, indices)
+ }
+ ArrowType::UInt32 => {
+ remove_primitive_array_indices!(arr, ArrowUInt32Type, indices)
+ }
+ ArrowType::UInt64 => {
+ remove_primitive_array_indices!(arr, ArrowUInt64Type, indices)
+ }
+ ArrowType::Int8 => remove_primitive_array_indices!(arr, ArrowInt8Type, indices),
+ ArrowType::Int16 => remove_primitive_array_indices!(arr, ArrowInt16Type, indices),
+ ArrowType::Int32 => remove_primitive_array_indices!(arr, ArrowInt32Type, indices),
+ ArrowType::Int64 => remove_primitive_array_indices!(arr, ArrowInt64Type, indices),
+ ArrowType::Float32 => {
+ remove_primitive_array_indices!(arr, ArrowFloat32Type, indices)
+ }
+ ArrowType::Float64 => {
+ remove_primitive_array_indices!(arr, ArrowFloat64Type, indices)
+ }
+ ArrowType::Boolean => {
+ remove_primitive_array_indices!(arr, ArrowBooleanType, indices)
+ }
+ ArrowType::Date32(_) => {
+ remove_primitive_array_indices!(arr, ArrowDate32Type, indices)
+ }
+ ArrowType::Date64(_) => {
+ remove_primitive_array_indices!(arr, ArrowDate64Type, indices)
+ }
+ ArrowType::Time32(ArrowTimeUnit::Second) => {
+ remove_primitive_array_indices!(arr, ArrowTime32SecondType, indices)
+ }
+ ArrowType::Time32(ArrowTimeUnit::Millisecond) => {
+ remove_primitive_array_indices!(arr, ArrowTime32MillisecondType, indices)
+ }
+ ArrowType::Time64(ArrowTimeUnit::Microsecond) => {
+ remove_primitive_array_indices!(arr, ArrowTime64MicrosecondType, indices)
+ }
+ ArrowType::Time64(ArrowTimeUnit::Nanosecond) => {
+ remove_primitive_array_indices!(arr, ArrowTime64NanosecondType, indices)
+ }
+ ArrowType::Duration(ArrowTimeUnit::Second) => {
+ remove_primitive_array_indices!(arr, ArrowDurationSecondType, indices)
+ }
+ ArrowType::Duration(ArrowTimeUnit::Millisecond) => {
+ remove_primitive_array_indices!(arr, ArrowDurationMillisecondType, indices)
+ }
+ ArrowType::Duration(ArrowTimeUnit::Microsecond) => {
+ remove_primitive_array_indices!(arr, ArrowDurationMicrosecondType, indices)
+ }
+ ArrowType::Duration(ArrowTimeUnit::Nanosecond) => {
+ remove_primitive_array_indices!(arr, ArrowDurationNanosecondType, indices)
+ }
+ ArrowType::Timestamp(ArrowTimeUnit::Second, _) => {
+ remove_primitive_array_indices!(arr, ArrowTimestampSecondType, indices)
+ }
+ ArrowType::Timestamp(ArrowTimeUnit::Millisecond, _) => {
+ remove_primitive_array_indices!(arr, ArrowTimestampMillisecondType, indices)
+ }
+ ArrowType::Timestamp(ArrowTimeUnit::Microsecond, _) => {
+ remove_primitive_array_indices!(arr, ArrowTimestampMicrosecondType, indices)
+ }
+ ArrowType::Timestamp(ArrowTimeUnit::Nanosecond, _) => {
+ remove_primitive_array_indices!(arr, ArrowTimestampNanosecondType, indices)
+ }
+ ArrowType::Utf8 => {
+ remove_array_indices_custom_builder!(arr, StringArray, StringBuilder, indices)
+ }
+ ArrowType::Binary => {
+ remove_array_indices_custom_builder!(arr, BinaryArray, BinaryBuilder, indices)
+ }
+ ArrowType::FixedSizeBinary(size) => remove_fixed_size_binary_array_indices!(
+ arr,
+ FixedSizeBinaryArray,
+ FixedSizeBinaryBuilder,
+ indices,
+ size
+ ),
+ _ => Err(ParquetError::General(format!(
+ "ListArray of type List({:?}) is not supported by array_reader",
+ item_type
+ ))),
+ }
+}
+
+/// Implementation of ListArrayReader. Nested lists and lists of structs are not yet supported.
+impl<OffsetSize: OffsetSizeTrait> ArrayReader for ListArrayReader<OffsetSize> {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ /// Returns data type.
+ /// This must be a List.
+ fn get_data_type(&self) -> &ArrowType {
+ &self.data_type
+ }
+
+ fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
+ let next_batch_array = self.item_reader.next_batch(batch_size)?;
+ let item_type = self.item_reader.get_data_type().clone();
+
+ if next_batch_array.len() == 0 {
+ return build_empty_list_array(item_type);
+ }
+ let def_levels = self
+ .item_reader
+ .get_def_levels()
+ .ok_or_else(|| ArrowError("item_reader def levels are None.".to_string()))?;
+ let rep_levels = self
+ .item_reader
+ .get_rep_levels()
+ .ok_or_else(|| ArrowError("item_reader rep levels are None.".to_string()))?;
+
+ if !((def_levels.len() == rep_levels.len())
+ && (rep_levels.len() == next_batch_array.len()))
+ {
+ return Err(ArrowError(
+ "Expected item_reader def_levels and rep_levels to be same length as batch".to_string(),
+ ));
+ }
+
+ // Need to remove from the values array the nulls that represent null lists rather than null items
+ // null lists have def_level = 0
+ let mut null_list_indices: Vec<usize> = Vec::new();
+ for i in 0..def_levels.len() {
+ if def_levels[i] == 0 {
+ null_list_indices.push(i);
+ }
+ }
+ let batch_values = match null_list_indices.len() {
+ 0 => next_batch_array.clone(),
+ _ => remove_indices(next_batch_array.clone(), item_type, null_list_indices)?,
+ };
+
+ // null list has def_level = 0
+ // empty list has def_level = 1
+ // null item in a list has def_level = 2
+ // non-null item has def_level = 3
+ // first item in each list has rep_level = 0, subsequent items have rep_level = 1
+
+ let mut offsets: Vec<OffsetSize> = Vec::new();
+ let mut cur_offset = OffsetSize::zero();
+ for i in 0..rep_levels.len() {
+ if rep_levels[i] == 0 {
+ offsets.push(cur_offset)
+ }
+ if def_levels[i] > 0 {
+ cur_offset = cur_offset + OffsetSize::one();
+ }
+ }
+ offsets.push(cur_offset);
+
+ let num_bytes = bit_util::ceil(offsets.len(), 8);
+ let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, false);
+ let null_slice = null_buf.data_mut();
+ let mut list_index = 0;
+ for i in 0..rep_levels.len() {
+ if rep_levels[i] == 0 && def_levels[i] != 0 {
+ bit_util::set_bit(null_slice, list_index);
+ }
+ if rep_levels[i] == 0 {
+ list_index += 1;
+ }
+ }
+ let value_offsets = Buffer::from(&offsets.to_byte_slice());
+
+ // null list has def_level = 0
+ let null_count = def_levels.iter().filter(|x| x == &&0).count();
+
+ let list_data = ArrayData::builder(self.get_data_type().clone())
+ .len(offsets.len() - 1)
+ .add_buffer(value_offsets)
+ .add_child_data(batch_values.data())
+ .null_bit_buffer(null_buf.freeze())
+ .null_count(null_count)
+ .offset(next_batch_array.offset())
+ .build();
+
+ let result_array = GenericListArray::<OffsetSize>::from(list_data);
+ Ok(Arc::new(result_array))
+ }
+
+ fn get_def_levels(&self) -> Option<&[i16]> {
+ self.def_level_buffer
+ .as_ref()
+ .map(|buf| unsafe { buf.typed_data() })
+ }
+
+ fn get_rep_levels(&self) -> Option<&[i16]> {
+ self.rep_level_buffer
+ .as_ref()
+ .map(|buf| unsafe { buf.typed_data() })
+ }
+}
+
/// Implementation of struct array reader.
pub struct StructArrayReader {
children: Vec<Box<dyn ArrayReader>>,
@@ -875,16 +1291,94 @@ impl<'a> TypeVisitor<Option<Box<dyn ArrayReader>>, &'a ArrayReaderBuilderContext
}
/// Build array reader for list type.
- /// Currently this is not supported.
fn visit_list_with_item(
&mut self,
- _list_type: Rc<Type>,
- _item_type: &Type,
- _context: &'a ArrayReaderBuilderContext,
+ list_type: Rc<Type>,
+ item_type: Rc<Type>,
+ context: &'a ArrayReaderBuilderContext,
) -> Result<Option<Box<dyn ArrayReader>>> {
- Err(ArrowError(
- "Reading parquet list array into arrow is not supported yet!".to_string(),
- ))
+ let list_child = &list_type
+ .get_fields()
+ .first()
+ .ok_or_else(|| ArrowError("List field must have a child.".to_string()))?;
+ let mut new_context = context.clone();
+
+ new_context.path.append(vec![list_type.name().to_string()]);
+
+ match list_type.get_basic_info().repetition() {
+ Repetition::REPEATED => {
+ new_context.def_level += 1;
+ new_context.rep_level += 1;
+ }
+ Repetition::OPTIONAL => {
+ new_context.def_level += 1;
+ }
+ _ => (),
+ }
+
+ match list_child.get_basic_info().repetition() {
+ Repetition::REPEATED => {
+ new_context.def_level += 1;
+ new_context.rep_level += 1;
+ }
+ Repetition::OPTIONAL => {
+ new_context.def_level += 1;
+ }
+ _ => (),
+ }
+
+ let item_reader = self
+ .dispatch(item_type.clone(), &new_context)
+ .unwrap()
+ .unwrap();
+
+ let item_reader_type = item_reader.get_data_type().clone();
+
+ match item_reader_type {
+ ArrowType::List(_)
+ | ArrowType::FixedSizeList(_, _)
+ | ArrowType::Struct(_)
+ | ArrowType::Dictionary(_, _) => Err(ArrowError(format!(
+ "reading List({:?}) into arrow not supported yet",
+ item_type
+ ))),
+ _ => {
+ let arrow_type = self
+ .arrow_schema
+ .field_with_name(list_type.name())
+ .ok()
+ .map(|f| f.data_type().to_owned())
+ .unwrap_or_else(|| {
+ ArrowType::List(Box::new(item_reader_type.clone()))
+ });
+
+ let list_array_reader: Box<dyn ArrayReader> = match arrow_type {
+ ArrowType::List(_) => Box::new(ListArrayReader::<i32>::new(
+ item_reader,
+ arrow_type,
+ item_reader_type,
+ new_context.def_level,
+ new_context.rep_level,
+ )),
+ ArrowType::LargeList(_) => Box::new(ListArrayReader::<i64>::new(
+ item_reader,
+ arrow_type,
+ item_reader_type,
+ new_context.def_level,
+ new_context.rep_level,
+ )),
+
+ _ => {
+ return Err(ArrowError(format!(
+ "creating ListArrayReader with type {:?} should be unreachable",
+ arrow_type
+ )))
+ }
+ };
+
+ Ok(Some(list_array_reader))
+ }
+ }
}
}
@@ -1100,7 +1594,10 @@ mod tests {
DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator,
};
use crate::util::test_common::{get_test_file, make_pages};
- use arrow::array::{Array, ArrayRef, PrimitiveArray, StringArray, StructArray};
+ use arrow::array::{
+ Array, ArrayRef, LargeListArray, ListArray, PrimitiveArray, StringArray,
+ StructArray,
+ };
use arrow::datatypes::{
ArrowPrimitiveType, DataType as ArrowType, Date32Type as ArrowDate32, Field,
Int32Type as ArrowInt32, Int64Type as ArrowInt64,
@@ -1311,7 +1808,7 @@ mod tests {
.next_batch(50)
.expect("Unable to get batch from reader");
- let result_data_type = <$result_arrow_type>::get_data_type();
+ let result_data_type = <$result_arrow_type>::DATA_TYPE;
let array = array
.as_any()
.downcast_ref::<PrimitiveArray<$result_arrow_type>>()
@@ -1759,4 +2256,113 @@ mod tests {
assert_eq!(array_reader.get_data_type(), &arrow_type);
}
+
+ #[test]
+ fn test_list_array_reader() {
+ // [[1, null, 2], null, [3, 4]]
+ let array = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![
+ Some(1),
+ None,
+ Some(2),
+ None,
+ Some(3),
+ Some(4),
+ ]));
+ let item_array_reader = InMemoryArrayReader::new(
+ ArrowType::Int32,
+ array,
+ Some(vec![3, 2, 3, 0, 3, 3]),
+ Some(vec![0, 1, 1, 0, 0, 1]),
+ );
+
+ let mut list_array_reader = ListArrayReader::<i32>::new(
+ Box::new(item_array_reader),
+ ArrowType::List(Box::new(ArrowType::Int32)),
+ ArrowType::Int32,
+ 1,
+ 1,
+ );
+
+ let next_batch = list_array_reader.next_batch(1024).unwrap();
+ let list_array = next_batch.as_any().downcast_ref::<ListArray>().unwrap();
+
+ assert_eq!(3, list_array.len());
+ // This passes as I expect
+ assert_eq!(1, list_array.null_count());
+
+ assert_eq!(
+ list_array
+ .value(0)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<ArrowInt32>>()
+ .unwrap(),
+ &PrimitiveArray::<ArrowInt32>::from(vec![Some(1), None, Some(2)])
+ );
+
+ assert!(list_array.is_null(1));
+
+ assert_eq!(
+ list_array
+ .value(2)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<ArrowInt32>>()
+ .unwrap(),
+ &PrimitiveArray::<ArrowInt32>::from(vec![Some(3), Some(4)])
+ );
+ }
+
+ #[test]
+ fn test_large_list_array_reader() {
+ // [[1, null, 2], null, [3, 4]]
+ let array = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![
+ Some(1),
+ None,
+ Some(2),
+ None,
+ Some(3),
+ Some(4),
+ ]));
+ let item_array_reader = InMemoryArrayReader::new(
+ ArrowType::Int32,
+ array,
+ Some(vec![3, 2, 3, 0, 3, 3]),
+ Some(vec![0, 1, 1, 0, 0, 1]),
+ );
+
+ let mut list_array_reader = ListArrayReader::<i64>::new(
+ Box::new(item_array_reader),
+ ArrowType::LargeList(Box::new(ArrowType::Int32)),
+ ArrowType::Int32,
+ 1,
+ 1,
+ );
+
+ let next_batch = list_array_reader.next_batch(1024).unwrap();
+ let list_array = next_batch
+ .as_any()
+ .downcast_ref::<LargeListArray>()
+ .unwrap();
+
+ assert_eq!(3, list_array.len());
+
+ assert_eq!(
+ list_array
+ .value(0)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<ArrowInt32>>()
+ .unwrap(),
+ &PrimitiveArray::<ArrowInt32>::from(vec![Some(1), None, Some(2)])
+ );
+
+ assert!(list_array.is_null(1));
+
+ assert_eq!(
+ list_array
+ .value(2)
+ .as_any()
+ .downcast_ref::<PrimitiveArray<ArrowInt32>>()
+ .unwrap(),
+ &PrimitiveArray::<ArrowInt32>::from(vec![Some(3), Some(4)])
+ );
+ }
}
diff --git a/rust/parquet/src/arrow/arrow_writer.rs b/rust/parquet/src/arrow/arrow_writer.rs
index ff535dc..d5e2db4 100644
--- a/rust/parquet/src/arrow/arrow_writer.rs
+++ b/rust/parquet/src/arrow/arrow_writer.rs
@@ -22,7 +22,7 @@ use std::rc::Rc;
use arrow::array as arrow_array;
use arrow::datatypes::{DataType as ArrowDataType, SchemaRef};
use arrow::record_batch::RecordBatch;
-use arrow_array::{Array, PrimitiveArrayOps};
+use arrow_array::Array;
use super::schema::add_encoded_arrow_schema_to_metadata;
use crate::column::writer::ColumnWriter;
@@ -534,6 +534,7 @@ mod tests {
}
#[test]
+ #[ignore = "repetitions might be incorrect, will be addressed as part of ARROW-9728"]
fn arrow_writer_list() {
// define schema
let schema = Schema::new(vec![Field::new(
@@ -546,7 +547,7 @@ mod tests {
let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
// Construct a buffer for value offsets, for the nested array:
- // [[false], [true, false], null, [true, false, true], [false, true, false, true]]
+ // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
let a_value_offsets =
arrow::buffer::Buffer::from(&[0, 1, 3, 3, 6, 10].to_byte_slice());
@@ -562,6 +563,9 @@ mod tests {
let batch =
RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)]).unwrap();
+ // I think this setup is incorrect because this should pass
+ assert_eq!(batch.column(0).data().null_count(), 1);
+
let file = get_temp_file("test_arrow_writer_list.parquet", &[]);
let mut writer = ArrowWriter::try_new(file, Arc::new(schema), None).unwrap();
writer.write(&batch).unwrap();
@@ -1063,9 +1067,7 @@ mod tests {
}
#[test]
- #[should_panic(
- expected = "Reading parquet list array into arrow is not supported yet!"
- )]
+ #[ignore = "repetitions might be incorrect, will be addressed as part of ARROW-9728"]
fn list_single_column() {
let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
let a_value_offsets =
@@ -1075,16 +1077,18 @@ mod tests {
.add_buffer(a_value_offsets)
.add_child_data(a_values.data())
.build();
- let a = ListArray::from(a_list_data);
+ // I think this setup is incorrect because this should pass
+ assert_eq!(a_list_data.null_count(), 1);
+
+ let a = ListArray::from(a_list_data);
let values = Arc::new(a);
+
one_column_roundtrip("list_single_column", values, false);
}
#[test]
- #[should_panic(
- expected = "Reading parquet list array into arrow is not supported yet!"
- )]
+ #[ignore = "repetitions might be incorrect, will be addressed as part of ARROW-9728"]
fn large_list_single_column() {
let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
let a_value_offsets =
@@ -1095,14 +1099,17 @@ mod tests {
.add_buffer(a_value_offsets)
.add_child_data(a_values.data())
.build();
- let a = LargeListArray::from(a_list_data);
+ // I think this setup is incorrect because this should pass
+ assert_eq!(a_list_data.null_count(), 1);
+
+ let a = LargeListArray::from(a_list_data);
let values = Arc::new(a);
+
one_column_roundtrip("large_list_single_column", values, false);
}
#[test]
- #[ignore] // Struct support isn't correct yet - null_bitmap doesn't match
fn struct_single_column() {
let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
let struct_field_a = Field::new("f", DataType::Int32, false);
diff --git a/rust/parquet/src/schema/visitor.rs b/rust/parquet/src/schema/visitor.rs
index 6d712ce..a1866fb 100644
--- a/rust/parquet/src/schema/visitor.rs
+++ b/rust/parquet/src/schema/visitor.rs
@@ -50,7 +50,7 @@ pub trait TypeVisitor<R, C> {
{
self.visit_list_with_item(
list_type.clone(),
- list_item,
+ list_item.clone(),
context,
)
} else {
@@ -70,13 +70,13 @@ pub trait TypeVisitor<R, C> {
{
self.visit_list_with_item(
list_type.clone(),
- fields.first().unwrap(),
+ fields.first().unwrap().clone(),
context,
)
} else {
self.visit_list_with_item(
list_type.clone(),
- list_item,
+ list_item.clone(),
context,
)
}
@@ -114,7 +114,7 @@ pub trait TypeVisitor<R, C> {
fn visit_list_with_item(
&mut self,
list_type: TypePtr,
- item_type: &Type,
+ item_type: TypePtr,
context: C,
) -> Result<R>;
}
@@ -125,7 +125,7 @@ mod tests {
use crate::basic::Type as PhysicalType;
use crate::errors::Result;
use crate::schema::parser::parse_message_type;
- use crate::schema::types::{Type, TypePtr};
+ use crate::schema::types::TypePtr;
use std::rc::Rc;
struct TestVisitorContext {}
@@ -174,7 +174,7 @@ mod tests {
fn visit_list_with_item(
&mut self,
list_type: TypePtr,
- item_type: &Type,
+ item_type: TypePtr,
_context: TestVisitorContext,
) -> Result<bool> {
assert_eq!(
[arrow] 07/09: ARROW-10225: [Rust] [Parquet] Fix null comparison in
roundtrip
Posted by ne...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit 25d96c5eb79da4d9646cdf49a299246a5e4c2296
Author: Neville Dipale <ne...@gmail.com>
AuthorDate: Thu Oct 8 17:08:59 2020 +0200
ARROW-10225: [Rust] [Parquet] Fix null comparison in roundtrip
Closes #8388 from nevi-me/ARROW-10225
Authored-by: Neville Dipale <ne...@gmail.com>
Signed-off-by: Neville Dipale <ne...@gmail.com>
---
rust/parquet/src/arrow/arrow_writer.rs | 10 ++++++----
1 file changed, 6 insertions(+), 4 deletions(-)
diff --git a/rust/parquet/src/arrow/arrow_writer.rs b/rust/parquet/src/arrow/arrow_writer.rs
index 40e2553..a17e424 100644
--- a/rust/parquet/src/arrow/arrow_writer.rs
+++ b/rust/parquet/src/arrow/arrow_writer.rs
@@ -724,7 +724,11 @@ mod tests {
assert_eq!(expected_data.offset(), actual_data.offset());
assert_eq!(expected_data.buffers(), actual_data.buffers());
assert_eq!(expected_data.child_data(), actual_data.child_data());
- assert_eq!(expected_data.null_bitmap(), actual_data.null_bitmap());
+ // Null counts should be the same, not necessarily bitmaps
+ // A null bitmap is optional if an array has no nulls
+ if expected_data.null_count() != 0 {
+ assert_eq!(expected_data.null_bitmap(), actual_data.null_bitmap());
+ }
}
}
@@ -1001,7 +1005,7 @@ mod tests {
}
#[test]
- #[ignore] // Binary support isn't correct yet - null_bitmap doesn't match
+ #[ignore] // Binary support isn't correct yet - buffers don't match
fn binary_single_column() {
let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
let many_vecs: Vec<_> = std::iter::repeat(one_vec).take(SMALL_SIZE).collect();
@@ -1026,7 +1030,6 @@ mod tests {
}
#[test]
- #[ignore] // String support isn't correct yet - null_bitmap doesn't match
fn string_single_column() {
let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
let raw_strs = raw_values.iter().map(|s| s.as_str());
@@ -1035,7 +1038,6 @@ mod tests {
}
#[test]
- #[ignore] // Large string support isn't correct yet - null_bitmap doesn't match
fn large_string_single_column() {
let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
let raw_strs = raw_values.iter().map(|s| s.as_str());
[arrow] 01/09: ARROW-8289: [Rust] Parquet Arrow writer with nested
support
Posted by ne...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit 6dc26ed7cd94d90a76c278bbc64650a10ecb8f3b
Author: Neville Dipale <ne...@gmail.com>
AuthorDate: Thu Aug 13 18:47:34 2020 +0200
ARROW-8289: [Rust] Parquet Arrow writer with nested support
**Note**: I started making changes to #6785, and ended up deviating a lot, so I opted for making a new draft PR in case my approach is not suitable.
___
This is a draft to implement an arrow writer for parquet. It supports the following (no complete test coverage yet):
* writing primitives except for booleans and binary
* nested structs
* null values (via definition levels)
It does not yet support:
- Boolean arrays (have to be handled differently from numeric values)
- Binary arrays
- Dictionary arrays
- Union arrays (are they even possible?)
I have only added a test by creating a nested schema, which I tested on pyarrow.
```jupyter
# schema of test_complex.parquet
a: int32 not null
b: int32
c: struct<d: double, e: struct<f: float>> not null
child 0, d: double
child 1, e: struct<f: float>
child 0, f: float
```
This PR potentially addresses:
* https://issues.apache.org/jira/browse/ARROW-8289
* https://issues.apache.org/jira/browse/ARROW-8423
* https://issues.apache.org/jira/browse/ARROW-8424
* https://issues.apache.org/jira/browse/ARROW-8425
And I would like to propose either opening new JIRAs for the above incomplete items, or renaming the last 3 above.
___
**Help Needed**
I'm implementing the definition and repetition levels on first principle from an old Parquet blog post from the Twitter engineering blog. It's likely that I'm not getting some concepts correct, so I would appreciate help with:
* Checking if my logic is correct
* Guidance or suggestions on how to more efficiently extract levels from arrays
* Adding tests - I suspect we might need a lot of tests, so far we only test writing 1 batch, so I don't know how paging would work when writing a large enough file
I also don't know if the various encoding levels (dictionary, RLE, etc.) and compression levels are applied automagically, or if that'd be something we need to explicitly enable.
CC @sunchao @sadikovi @andygrove @paddyhoran
Might be of interest to @mcassels @maxburke
Closes #7319 from nevi-me/arrow-parquet-writer
Lead-authored-by: Neville Dipale <ne...@gmail.com>
Co-authored-by: Max Burke <ma...@urbanlogiq.com>
Co-authored-by: Andy Grove <an...@gmail.com>
Co-authored-by: Max Burke <ma...@gmail.com>
Signed-off-by: Neville Dipale <ne...@gmail.com>
---
rust/parquet/src/arrow/arrow_writer.rs | 682 +++++++++++++++++++++++++++++++++
rust/parquet/src/arrow/mod.rs | 5 +-
rust/parquet/src/schema/types.rs | 6 +-
3 files changed, 691 insertions(+), 2 deletions(-)
diff --git a/rust/parquet/src/arrow/arrow_writer.rs b/rust/parquet/src/arrow/arrow_writer.rs
new file mode 100644
index 0000000..0c1c490
--- /dev/null
+++ b/rust/parquet/src/arrow/arrow_writer.rs
@@ -0,0 +1,682 @@
+// 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.
+
+//! Contains writer which writes arrow data into parquet data.
+
+use std::rc::Rc;
+
+use arrow::array as arrow_array;
+use arrow::datatypes::{DataType as ArrowDataType, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use arrow_array::Array;
+
+use crate::column::writer::ColumnWriter;
+use crate::errors::{ParquetError, Result};
+use crate::file::properties::WriterProperties;
+use crate::{
+ data_type::*,
+ file::writer::{FileWriter, ParquetWriter, RowGroupWriter, SerializedFileWriter},
+};
+
+/// Arrow writer
+///
+/// Writes Arrow `RecordBatch`es to a Parquet writer
+pub struct ArrowWriter<W: ParquetWriter> {
+ /// Underlying Parquet writer
+ writer: SerializedFileWriter<W>,
+ /// A copy of the Arrow schema.
+ ///
+ /// The schema is used to verify that each record batch written has the correct schema
+ arrow_schema: SchemaRef,
+}
+
+impl<W: 'static + ParquetWriter> ArrowWriter<W> {
+ /// Try to create a new Arrow writer
+ ///
+ /// The writer will fail if:
+ /// * a `SerializedFileWriter` cannot be created from the ParquetWriter
+ /// * the Arrow schema contains unsupported datatypes such as Unions
+ pub fn try_new(
+ writer: W,
+ arrow_schema: SchemaRef,
+ props: Option<Rc<WriterProperties>>,
+ ) -> Result<Self> {
+ let schema = crate::arrow::arrow_to_parquet_schema(&arrow_schema)?;
+ let props = match props {
+ Some(props) => props,
+ None => Rc::new(WriterProperties::builder().build()),
+ };
+ let file_writer = SerializedFileWriter::new(
+ writer.try_clone()?,
+ schema.root_schema_ptr(),
+ props,
+ )?;
+
+ Ok(Self {
+ writer: file_writer,
+ arrow_schema,
+ })
+ }
+
+ /// Write a RecordBatch to writer
+ ///
+ /// *NOTE:* The writer currently does not support all Arrow data types
+ pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
+ // validate batch schema against writer's supplied schema
+ if self.arrow_schema != batch.schema() {
+ return Err(ParquetError::ArrowError(
+ "Record batch schema does not match writer schema".to_string(),
+ ));
+ }
+ // compute the definition and repetition levels of the batch
+ let mut levels = vec![];
+ batch.columns().iter().for_each(|array| {
+ let mut array_levels =
+ get_levels(array, 0, &vec![1i16; batch.num_rows()][..], None);
+ levels.append(&mut array_levels);
+ });
+ // reverse levels so we can use Vec::pop(&mut self)
+ levels.reverse();
+
+ let mut row_group_writer = self.writer.next_row_group()?;
+
+ // write leaves
+ for column in batch.columns() {
+ write_leaves(&mut row_group_writer, column, &mut levels)?;
+ }
+
+ self.writer.close_row_group(row_group_writer)
+ }
+
+ /// Close and finalise the underlying Parquet writer
+ pub fn close(&mut self) -> Result<()> {
+ self.writer.close()
+ }
+}
+
+/// Convenience method to get the next ColumnWriter from the RowGroupWriter
+#[inline]
+#[allow(clippy::borrowed_box)]
+fn get_col_writer(
+ row_group_writer: &mut Box<dyn RowGroupWriter>,
+) -> Result<ColumnWriter> {
+ let col_writer = row_group_writer
+ .next_column()?
+ .expect("Unable to get column writer");
+ Ok(col_writer)
+}
+
+#[allow(clippy::borrowed_box)]
+fn write_leaves(
+ mut row_group_writer: &mut Box<dyn RowGroupWriter>,
+ array: &arrow_array::ArrayRef,
+ mut levels: &mut Vec<Levels>,
+) -> Result<()> {
+ match array.data_type() {
+ ArrowDataType::Int8
+ | ArrowDataType::Int16
+ | ArrowDataType::Int32
+ | ArrowDataType::Int64
+ | ArrowDataType::UInt8
+ | ArrowDataType::UInt16
+ | ArrowDataType::UInt32
+ | ArrowDataType::UInt64
+ | ArrowDataType::Float16
+ | ArrowDataType::Float32
+ | ArrowDataType::Float64
+ | ArrowDataType::Timestamp(_, _)
+ | ArrowDataType::Date32(_)
+ | ArrowDataType::Date64(_)
+ | ArrowDataType::Time32(_)
+ | ArrowDataType::Time64(_)
+ | ArrowDataType::Duration(_)
+ | ArrowDataType::Interval(_)
+ | ArrowDataType::LargeBinary
+ | ArrowDataType::Binary
+ | ArrowDataType::Utf8
+ | ArrowDataType::LargeUtf8 => {
+ let mut col_writer = get_col_writer(&mut row_group_writer)?;
+ write_leaf(
+ &mut col_writer,
+ array,
+ levels.pop().expect("Levels exhausted"),
+ )?;
+ row_group_writer.close_column(col_writer)?;
+ Ok(())
+ }
+ ArrowDataType::List(_) | ArrowDataType::LargeList(_) => {
+ // write the child list
+ let data = array.data();
+ let child_array = arrow_array::make_array(data.child_data()[0].clone());
+ write_leaves(&mut row_group_writer, &child_array, &mut levels)?;
+ Ok(())
+ }
+ ArrowDataType::Struct(_) => {
+ let struct_array: &arrow_array::StructArray = array
+ .as_any()
+ .downcast_ref::<arrow_array::StructArray>()
+ .expect("Unable to get struct array");
+ for field in struct_array.columns() {
+ write_leaves(&mut row_group_writer, field, &mut levels)?;
+ }
+ Ok(())
+ }
+ ArrowDataType::FixedSizeList(_, _)
+ | ArrowDataType::Null
+ | ArrowDataType::Boolean
+ | ArrowDataType::FixedSizeBinary(_)
+ | ArrowDataType::Union(_)
+ | ArrowDataType::Dictionary(_, _) => Err(ParquetError::NYI(
+ "Attempting to write an Arrow type that is not yet implemented".to_string(),
+ )),
+ }
+}
+
+fn write_leaf(
+ writer: &mut ColumnWriter,
+ column: &arrow_array::ArrayRef,
+ levels: Levels,
+) -> Result<i64> {
+ let written = match writer {
+ ColumnWriter::Int32ColumnWriter(ref mut typed) => {
+ let array = arrow::compute::cast(column, &ArrowDataType::Int32)?;
+ let array = array
+ .as_any()
+ .downcast_ref::<arrow_array::Int32Array>()
+ .expect("Unable to get int32 array");
+ typed.write_batch(
+ get_numeric_array_slice::<Int32Type, _>(&array).as_slice(),
+ Some(levels.definition.as_slice()),
+ levels.repetition.as_deref(),
+ )?
+ }
+ ColumnWriter::BoolColumnWriter(ref mut _typed) => {
+ unreachable!("Currently unreachable because data type not supported")
+ }
+ ColumnWriter::Int64ColumnWriter(ref mut typed) => {
+ let array = arrow_array::Int64Array::from(column.data());
+ typed.write_batch(
+ get_numeric_array_slice::<Int64Type, _>(&array).as_slice(),
+ Some(levels.definition.as_slice()),
+ levels.repetition.as_deref(),
+ )?
+ }
+ ColumnWriter::Int96ColumnWriter(ref mut _typed) => {
+ unreachable!("Currently unreachable because data type not supported")
+ }
+ ColumnWriter::FloatColumnWriter(ref mut typed) => {
+ let array = arrow_array::Float32Array::from(column.data());
+ typed.write_batch(
+ get_numeric_array_slice::<FloatType, _>(&array).as_slice(),
+ Some(levels.definition.as_slice()),
+ levels.repetition.as_deref(),
+ )?
+ }
+ ColumnWriter::DoubleColumnWriter(ref mut typed) => {
+ let array = arrow_array::Float64Array::from(column.data());
+ typed.write_batch(
+ get_numeric_array_slice::<DoubleType, _>(&array).as_slice(),
+ Some(levels.definition.as_slice()),
+ levels.repetition.as_deref(),
+ )?
+ }
+ ColumnWriter::ByteArrayColumnWriter(ref mut typed) => match column.data_type() {
+ ArrowDataType::Binary | ArrowDataType::Utf8 => {
+ let array = arrow_array::BinaryArray::from(column.data());
+ typed.write_batch(
+ get_binary_array(&array).as_slice(),
+ Some(levels.definition.as_slice()),
+ levels.repetition.as_deref(),
+ )?
+ }
+ ArrowDataType::LargeBinary | ArrowDataType::LargeUtf8 => {
+ let array = arrow_array::LargeBinaryArray::from(column.data());
+ typed.write_batch(
+ get_large_binary_array(&array).as_slice(),
+ Some(levels.definition.as_slice()),
+ levels.repetition.as_deref(),
+ )?
+ }
+ _ => unreachable!("Currently unreachable because data type not supported"),
+ },
+ ColumnWriter::FixedLenByteArrayColumnWriter(ref mut _typed) => {
+ unreachable!("Currently unreachable because data type not supported")
+ }
+ };
+ Ok(written as i64)
+}
+
+/// A struct that represents definition and repetition levels.
+/// Repetition levels are only populated if the parent or current leaf is repeated
+#[derive(Debug)]
+struct Levels {
+ definition: Vec<i16>,
+ repetition: Option<Vec<i16>>,
+}
+
+/// Compute nested levels of the Arrow array, recursing into lists and structs
+fn get_levels(
+ array: &arrow_array::ArrayRef,
+ level: i16,
+ parent_def_levels: &[i16],
+ parent_rep_levels: Option<&[i16]>,
+) -> Vec<Levels> {
+ match array.data_type() {
+ ArrowDataType::Null => unimplemented!(),
+ ArrowDataType::Boolean
+ | ArrowDataType::Int8
+ | ArrowDataType::Int16
+ | ArrowDataType::Int32
+ | ArrowDataType::Int64
+ | ArrowDataType::UInt8
+ | ArrowDataType::UInt16
+ | ArrowDataType::UInt32
+ | ArrowDataType::UInt64
+ | ArrowDataType::Float16
+ | ArrowDataType::Float32
+ | ArrowDataType::Float64
+ | ArrowDataType::Utf8
+ | ArrowDataType::LargeUtf8
+ | ArrowDataType::Timestamp(_, _)
+ | ArrowDataType::Date32(_)
+ | ArrowDataType::Date64(_)
+ | ArrowDataType::Time32(_)
+ | ArrowDataType::Time64(_)
+ | ArrowDataType::Duration(_)
+ | ArrowDataType::Interval(_)
+ | ArrowDataType::Binary
+ | ArrowDataType::LargeBinary => vec![Levels {
+ definition: get_primitive_def_levels(array, parent_def_levels),
+ repetition: None,
+ }],
+ ArrowDataType::FixedSizeBinary(_) => unimplemented!(),
+ ArrowDataType::List(_) | ArrowDataType::LargeList(_) => {
+ let array_data = array.data();
+ let child_data = array_data.child_data().get(0).unwrap();
+ // get offsets, accounting for large offsets if present
+ let offsets: Vec<i64> = {
+ if let ArrowDataType::LargeList(_) = array.data_type() {
+ unsafe { array_data.buffers()[0].typed_data::<i64>() }.to_vec()
+ } else {
+ let offsets = unsafe { array_data.buffers()[0].typed_data::<i32>() };
+ offsets.to_vec().into_iter().map(|v| v as i64).collect()
+ }
+ };
+ let child_array = arrow_array::make_array(child_data.clone());
+
+ let mut list_def_levels = Vec::with_capacity(child_array.len());
+ let mut list_rep_levels = Vec::with_capacity(child_array.len());
+ let rep_levels: Vec<i16> = parent_rep_levels
+ .map(|l| l.to_vec())
+ .unwrap_or_else(|| vec![0i16; parent_def_levels.len()]);
+ parent_def_levels
+ .iter()
+ .zip(rep_levels)
+ .zip(offsets.windows(2))
+ .for_each(|((parent_def_level, parent_rep_level), window)| {
+ if *parent_def_level == 0 {
+ // parent is null, list element must also be null
+ list_def_levels.push(0);
+ list_rep_levels.push(0);
+ } else {
+ // parent is not null, check if list is empty or null
+ let start = window[0];
+ let end = window[1];
+ let len = end - start;
+ if len == 0 {
+ list_def_levels.push(*parent_def_level - 1);
+ list_rep_levels.push(parent_rep_level);
+ } else {
+ list_def_levels.push(*parent_def_level);
+ list_rep_levels.push(parent_rep_level);
+ for _ in 1..len {
+ list_def_levels.push(*parent_def_level);
+ list_rep_levels.push(parent_rep_level + 1);
+ }
+ }
+ }
+ });
+
+ // if datatype is a primitive, we can construct levels of the child array
+ match child_array.data_type() {
+ ArrowDataType::Null => unimplemented!(),
+ ArrowDataType::Boolean => unimplemented!(),
+ ArrowDataType::Int8
+ | ArrowDataType::Int16
+ | ArrowDataType::Int32
+ | ArrowDataType::Int64
+ | ArrowDataType::UInt8
+ | ArrowDataType::UInt16
+ | ArrowDataType::UInt32
+ | ArrowDataType::UInt64
+ | ArrowDataType::Float16
+ | ArrowDataType::Float32
+ | ArrowDataType::Float64
+ | ArrowDataType::Timestamp(_, _)
+ | ArrowDataType::Date32(_)
+ | ArrowDataType::Date64(_)
+ | ArrowDataType::Time32(_)
+ | ArrowDataType::Time64(_)
+ | ArrowDataType::Duration(_)
+ | ArrowDataType::Interval(_) => {
+ let def_levels =
+ get_primitive_def_levels(&child_array, &list_def_levels[..]);
+ vec![Levels {
+ definition: def_levels,
+ repetition: Some(list_rep_levels),
+ }]
+ }
+ ArrowDataType::Binary
+ | ArrowDataType::Utf8
+ | ArrowDataType::LargeUtf8 => unimplemented!(),
+ ArrowDataType::FixedSizeBinary(_) => unimplemented!(),
+ ArrowDataType::LargeBinary => unimplemented!(),
+ ArrowDataType::List(_) | ArrowDataType::LargeList(_) => {
+ // nested list
+ unimplemented!()
+ }
+ ArrowDataType::FixedSizeList(_, _) => unimplemented!(),
+ ArrowDataType::Struct(_) => get_levels(
+ array,
+ level + 1, // indicates a nesting level of 2 (list + struct)
+ &list_def_levels[..],
+ Some(&list_rep_levels[..]),
+ ),
+ ArrowDataType::Union(_) => unimplemented!(),
+ ArrowDataType::Dictionary(_, _) => unimplemented!(),
+ }
+ }
+ ArrowDataType::FixedSizeList(_, _) => unimplemented!(),
+ ArrowDataType::Struct(_) => {
+ let struct_array: &arrow_array::StructArray = array
+ .as_any()
+ .downcast_ref::<arrow_array::StructArray>()
+ .expect("Unable to get struct array");
+ let mut struct_def_levels = Vec::with_capacity(struct_array.len());
+ for i in 0..array.len() {
+ struct_def_levels.push(level + struct_array.is_valid(i) as i16);
+ }
+ // trying to create levels for struct's fields
+ let mut struct_levels = vec![];
+ struct_array.columns().into_iter().for_each(|col| {
+ let mut levels =
+ get_levels(col, level + 1, &struct_def_levels[..], parent_rep_levels);
+ struct_levels.append(&mut levels);
+ });
+ struct_levels
+ }
+ ArrowDataType::Union(_) => unimplemented!(),
+ ArrowDataType::Dictionary(_, _) => unimplemented!(),
+ }
+}
+
+/// Get the definition levels of the numeric array, with level 0 being null and 1 being not null
+/// In the case where the array in question is a child of either a list or struct, the levels
+/// are incremented in accordance with the `level` parameter.
+/// Parent levels are either 0 or 1, and are used to higher (correct terminology?) leaves as null
+fn get_primitive_def_levels(
+ array: &arrow_array::ArrayRef,
+ parent_def_levels: &[i16],
+) -> Vec<i16> {
+ let mut array_index = 0;
+ let max_def_level = parent_def_levels.iter().max().unwrap();
+ let mut primitive_def_levels = vec![];
+ parent_def_levels.iter().for_each(|def_level| {
+ if def_level < max_def_level {
+ primitive_def_levels.push(*def_level);
+ } else {
+ primitive_def_levels.push(def_level - array.is_null(array_index) as i16);
+ array_index += 1;
+ }
+ });
+ primitive_def_levels
+}
+
+macro_rules! def_get_binary_array_fn {
+ ($name:ident, $ty:ty) => {
+ fn $name(array: &$ty) -> Vec<ByteArray> {
+ let mut values = Vec::with_capacity(array.len() - array.null_count());
+ for i in 0..array.len() {
+ if array.is_valid(i) {
+ let bytes = ByteArray::from(array.value(i).to_vec());
+ values.push(bytes);
+ }
+ }
+ values
+ }
+ };
+}
+
+def_get_binary_array_fn!(get_binary_array, arrow_array::BinaryArray);
+def_get_binary_array_fn!(get_large_binary_array, arrow_array::LargeBinaryArray);
+
+/// Get the underlying numeric array slice, skipping any null values.
+/// If there are no null values, it might be quicker to get the slice directly instead of
+/// calling this function.
+fn get_numeric_array_slice<T, A>(array: &arrow_array::PrimitiveArray<A>) -> Vec<T::T>
+where
+ T: DataType,
+ A: arrow::datatypes::ArrowNumericType,
+ T::T: From<A::Native>,
+{
+ let mut values = Vec::with_capacity(array.len() - array.null_count());
+ for i in 0..array.len() {
+ if array.is_valid(i) {
+ values.push(array.value(i).into())
+ }
+ }
+ values
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::io::Seek;
+ use std::sync::Arc;
+
+ use arrow::array::*;
+ use arrow::datatypes::ToByteSlice;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use arrow::record_batch::{RecordBatch, RecordBatchReader};
+
+ use crate::arrow::{ArrowReader, ParquetFileArrowReader};
+ use crate::file::reader::SerializedFileReader;
+ use crate::util::test_common::get_temp_file;
+
+ #[test]
+ fn arrow_writer() {
+ // define schema
+ let schema = Schema::new(vec![
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::Int32, true),
+ ]);
+
+ // create some data
+ let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
+ let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
+
+ // build a record batch
+ let batch = RecordBatch::try_new(
+ Arc::new(schema.clone()),
+ vec![Arc::new(a), Arc::new(b)],
+ )
+ .unwrap();
+
+ let file = get_temp_file("test_arrow_writer.parquet", &[]);
+ let mut writer = ArrowWriter::try_new(file, Arc::new(schema), None).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ }
+
+ #[test]
+ fn arrow_writer_list() {
+ // define schema
+ let schema = Schema::new(vec![Field::new(
+ "a",
+ DataType::List(Box::new(DataType::Int32)),
+ false,
+ )]);
+
+ // create some data
+ let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[false], [true, false], null, [true, false, true], [false, true, false, true]]
+ let a_value_offsets =
+ arrow::buffer::Buffer::from(&[0, 1, 3, 3, 6, 10].to_byte_slice());
+
+ // Construct a list array from the above two
+ let a_list_data = ArrayData::builder(DataType::List(Box::new(DataType::Int32)))
+ .len(5)
+ .add_buffer(a_value_offsets)
+ .add_child_data(a_values.data())
+ .build();
+ let a = ListArray::from(a_list_data);
+
+ // build a record batch
+ let batch =
+ RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)]).unwrap();
+
+ let file = get_temp_file("test_arrow_writer_list.parquet", &[]);
+ let mut writer = ArrowWriter::try_new(file, Arc::new(schema), None).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ }
+
+ #[test]
+ fn arrow_writer_binary() {
+ let string_field = Field::new("a", DataType::Utf8, false);
+ let binary_field = Field::new("b", DataType::Binary, false);
+ let schema = Schema::new(vec![string_field, binary_field]);
+
+ let raw_string_values = vec!["foo", "bar", "baz", "quux"];
+ let raw_binary_values = vec![
+ b"foo".to_vec(),
+ b"bar".to_vec(),
+ b"baz".to_vec(),
+ b"quux".to_vec(),
+ ];
+ let raw_binary_value_refs = raw_binary_values
+ .iter()
+ .map(|x| x.as_slice())
+ .collect::<Vec<_>>();
+
+ let string_values = StringArray::from(raw_string_values.clone());
+ let binary_values = BinaryArray::from(raw_binary_value_refs);
+ let batch = RecordBatch::try_new(
+ Arc::new(schema.clone()),
+ vec![Arc::new(string_values), Arc::new(binary_values)],
+ )
+ .unwrap();
+
+ let mut file = get_temp_file("test_arrow_writer.parquet", &[]);
+ let mut writer =
+ ArrowWriter::try_new(file.try_clone().unwrap(), Arc::new(schema), None)
+ .unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ file.seek(std::io::SeekFrom::Start(0)).unwrap();
+ let file_reader = SerializedFileReader::new(file).unwrap();
+ let mut arrow_reader = ParquetFileArrowReader::new(Rc::new(file_reader));
+ let mut record_batch_reader = arrow_reader.get_record_reader(1024).unwrap();
+
+ let batch = record_batch_reader.next_batch().unwrap().unwrap();
+ let string_col = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let binary_col = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<BinaryArray>()
+ .unwrap();
+
+ for i in 0..batch.num_rows() {
+ assert_eq!(string_col.value(i), raw_string_values[i]);
+ assert_eq!(binary_col.value(i), raw_binary_values[i].as_slice());
+ }
+ }
+
+ #[test]
+ fn arrow_writer_complex() {
+ // define schema
+ let struct_field_d = Field::new("d", DataType::Float64, true);
+ let struct_field_f = Field::new("f", DataType::Float32, true);
+ let struct_field_g =
+ Field::new("g", DataType::List(Box::new(DataType::Int16)), false);
+ let struct_field_e = Field::new(
+ "e",
+ DataType::Struct(vec![struct_field_f.clone(), struct_field_g.clone()]),
+ true,
+ );
+ let schema = Schema::new(vec![
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::Int32, true),
+ Field::new(
+ "c",
+ DataType::Struct(vec![struct_field_d.clone(), struct_field_e.clone()]),
+ false,
+ ),
+ ]);
+
+ // create some data
+ let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
+ let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
+ let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
+ let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
+
+ let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
+ let g_value_offsets =
+ arrow::buffer::Buffer::from(&[0, 1, 3, 3, 6, 10].to_byte_slice());
+
+ // Construct a list array from the above two
+ let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
+ .len(5)
+ .add_buffer(g_value_offsets)
+ .add_child_data(g_value.data())
+ .build();
+ let g = ListArray::from(g_list_data);
+
+ let e = StructArray::from(vec![
+ (struct_field_f, Arc::new(f) as ArrayRef),
+ (struct_field_g, Arc::new(g) as ArrayRef),
+ ]);
+
+ let c = StructArray::from(vec![
+ (struct_field_d, Arc::new(d) as ArrayRef),
+ (struct_field_e, Arc::new(e) as ArrayRef),
+ ]);
+
+ // build a record batch
+ let batch = RecordBatch::try_new(
+ Arc::new(schema.clone()),
+ vec![Arc::new(a), Arc::new(b), Arc::new(c)],
+ )
+ .unwrap();
+
+ let file = get_temp_file("test_arrow_writer_complex.parquet", &[]);
+ let mut writer = ArrowWriter::try_new(file, Arc::new(schema), None).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ }
+}
diff --git a/rust/parquet/src/arrow/mod.rs b/rust/parquet/src/arrow/mod.rs
index ef1544d..8499481 100644
--- a/rust/parquet/src/arrow/mod.rs
+++ b/rust/parquet/src/arrow/mod.rs
@@ -51,10 +51,13 @@
pub(in crate::arrow) mod array_reader;
pub mod arrow_reader;
+pub mod arrow_writer;
pub(in crate::arrow) mod converter;
pub(in crate::arrow) mod record_reader;
pub mod schema;
pub use self::arrow_reader::ArrowReader;
pub use self::arrow_reader::ParquetFileArrowReader;
-pub use self::schema::{parquet_to_arrow_schema, parquet_to_arrow_schema_by_columns};
+pub use self::schema::{
+ arrow_to_parquet_schema, parquet_to_arrow_schema, parquet_to_arrow_schema_by_columns,
+};
diff --git a/rust/parquet/src/schema/types.rs b/rust/parquet/src/schema/types.rs
index 416073a..5799905 100644
--- a/rust/parquet/src/schema/types.rs
+++ b/rust/parquet/src/schema/types.rs
@@ -788,7 +788,7 @@ impl SchemaDescriptor {
result.clone()
}
- fn column_root_of(&self, i: usize) -> &Rc<Type> {
+ fn column_root_of(&self, i: usize) -> &TypePtr {
assert!(
i < self.leaves.len(),
"Index out of bound: {} not in [0, {})",
@@ -810,6 +810,10 @@ impl SchemaDescriptor {
self.schema.as_ref()
}
+ pub fn root_schema_ptr(&self) -> TypePtr {
+ self.schema.clone()
+ }
+
/// Returns schema name.
pub fn name(&self) -> &str {
self.schema.name()
[arrow] 02/09: ARROW-8423: [Rust] [Parquet] Serialize Arrow schema
metadata
Posted by ne...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit dcec33543a5aa5ca9fefab10bf94b257bf47611f
Author: Neville Dipale <ne...@gmail.com>
AuthorDate: Tue Aug 18 18:39:37 2020 +0200
ARROW-8423: [Rust] [Parquet] Serialize Arrow schema metadata
This will allow preserving Arrow-specific metadata when writing or reading Parquet files created from C++ or Rust.
If the schema can't be deserialised, the normal Parquet > Arrow schema conversion is performed.
Closes #7917 from nevi-me/ARROW-8243
Authored-by: Neville Dipale <ne...@gmail.com>
Signed-off-by: Neville Dipale <ne...@gmail.com>
---
rust/parquet/Cargo.toml | 3 +-
rust/parquet/src/arrow/arrow_writer.rs | 27 ++-
rust/parquet/src/arrow/mod.rs | 4 +
rust/parquet/src/arrow/schema.rs | 306 ++++++++++++++++++++++++++++-----
rust/parquet/src/file/properties.rs | 6 +-
5 files changed, 290 insertions(+), 56 deletions(-)
diff --git a/rust/parquet/Cargo.toml b/rust/parquet/Cargo.toml
index c1fd1b6..122c5b6 100644
--- a/rust/parquet/Cargo.toml
+++ b/rust/parquet/Cargo.toml
@@ -40,6 +40,7 @@ zstd = { version = "0.5", optional = true }
chrono = "0.4"
num-bigint = "0.3"
arrow = { path = "../arrow", version = "3.0.0-SNAPSHOT", optional = true }
+base64 = { version = "*", optional = true }
[dev-dependencies]
rand = "0.7"
@@ -52,4 +53,4 @@ arrow = { path = "../arrow", version = "3.0.0-SNAPSHOT" }
serde_json = { version = "1.0", features = ["preserve_order"] }
[features]
-default = ["arrow", "snap", "brotli", "flate2", "lz4", "zstd"]
+default = ["arrow", "snap", "brotli", "flate2", "lz4", "zstd", "base64"]
diff --git a/rust/parquet/src/arrow/arrow_writer.rs b/rust/parquet/src/arrow/arrow_writer.rs
index 0c1c490..1ca8d50 100644
--- a/rust/parquet/src/arrow/arrow_writer.rs
+++ b/rust/parquet/src/arrow/arrow_writer.rs
@@ -24,6 +24,7 @@ use arrow::datatypes::{DataType as ArrowDataType, SchemaRef};
use arrow::record_batch::RecordBatch;
use arrow_array::Array;
+use super::schema::add_encoded_arrow_schema_to_metadata;
use crate::column::writer::ColumnWriter;
use crate::errors::{ParquetError, Result};
use crate::file::properties::WriterProperties;
@@ -53,17 +54,17 @@ impl<W: 'static + ParquetWriter> ArrowWriter<W> {
pub fn try_new(
writer: W,
arrow_schema: SchemaRef,
- props: Option<Rc<WriterProperties>>,
+ props: Option<WriterProperties>,
) -> Result<Self> {
let schema = crate::arrow::arrow_to_parquet_schema(&arrow_schema)?;
- let props = match props {
- Some(props) => props,
- None => Rc::new(WriterProperties::builder().build()),
- };
+ // add serialized arrow schema
+ let mut props = props.unwrap_or_else(|| WriterProperties::builder().build());
+ add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut props);
+
let file_writer = SerializedFileWriter::new(
writer.try_clone()?,
schema.root_schema_ptr(),
- props,
+ Rc::new(props),
)?;
Ok(Self {
@@ -495,7 +496,7 @@ mod tests {
use arrow::record_batch::{RecordBatch, RecordBatchReader};
use crate::arrow::{ArrowReader, ParquetFileArrowReader};
- use crate::file::reader::SerializedFileReader;
+ use crate::file::{metadata::KeyValue, reader::SerializedFileReader};
use crate::util::test_common::get_temp_file;
#[test]
@@ -584,7 +585,7 @@ mod tests {
)
.unwrap();
- let mut file = get_temp_file("test_arrow_writer.parquet", &[]);
+ let mut file = get_temp_file("test_arrow_writer_binary.parquet", &[]);
let mut writer =
ArrowWriter::try_new(file.try_clone().unwrap(), Arc::new(schema), None)
.unwrap();
@@ -674,8 +675,16 @@ mod tests {
)
.unwrap();
+ let props = WriterProperties::builder()
+ .set_key_value_metadata(Some(vec![KeyValue {
+ key: "test_key".to_string(),
+ value: Some("test_value".to_string()),
+ }]))
+ .build();
+
let file = get_temp_file("test_arrow_writer_complex.parquet", &[]);
- let mut writer = ArrowWriter::try_new(file, Arc::new(schema), None).unwrap();
+ let mut writer =
+ ArrowWriter::try_new(file, Arc::new(schema), Some(props)).unwrap();
writer.write(&batch).unwrap();
writer.close().unwrap();
}
diff --git a/rust/parquet/src/arrow/mod.rs b/rust/parquet/src/arrow/mod.rs
index 8499481..2b012fb 100644
--- a/rust/parquet/src/arrow/mod.rs
+++ b/rust/parquet/src/arrow/mod.rs
@@ -58,6 +58,10 @@ pub mod schema;
pub use self::arrow_reader::ArrowReader;
pub use self::arrow_reader::ParquetFileArrowReader;
+pub use self::arrow_writer::ArrowWriter;
pub use self::schema::{
arrow_to_parquet_schema, parquet_to_arrow_schema, parquet_to_arrow_schema_by_columns,
};
+
+/// Schema metadata key used to store serialized Arrow IPC schema
+pub const ARROW_SCHEMA_META_KEY: &str = "ARROW:schema";
diff --git a/rust/parquet/src/arrow/schema.rs b/rust/parquet/src/arrow/schema.rs
index aebb9e7..d4cfe1f 100644
--- a/rust/parquet/src/arrow/schema.rs
+++ b/rust/parquet/src/arrow/schema.rs
@@ -26,24 +26,33 @@
use std::collections::{HashMap, HashSet};
use std::rc::Rc;
+use arrow::datatypes::{DataType, DateUnit, Field, Schema, TimeUnit};
+
use crate::basic::{LogicalType, Repetition, Type as PhysicalType};
use crate::errors::{ParquetError::ArrowError, Result};
-use crate::file::metadata::KeyValue;
+use crate::file::{metadata::KeyValue, properties::WriterProperties};
use crate::schema::types::{ColumnDescriptor, SchemaDescriptor, Type, TypePtr};
-use arrow::datatypes::TimeUnit;
-use arrow::datatypes::{DataType, DateUnit, Field, Schema};
-
-/// Convert parquet schema to arrow schema including optional metadata.
+/// Convert Parquet schema to Arrow schema including optional metadata.
+/// Attempts to decode any existing Arrow shcema metadata, falling back
+/// to converting the Parquet schema column-wise
pub fn parquet_to_arrow_schema(
parquet_schema: &SchemaDescriptor,
- metadata: &Option<Vec<KeyValue>>,
+ key_value_metadata: &Option<Vec<KeyValue>>,
) -> Result<Schema> {
- parquet_to_arrow_schema_by_columns(
- parquet_schema,
- 0..parquet_schema.columns().len(),
- metadata,
- )
+ let mut metadata = parse_key_value_metadata(key_value_metadata).unwrap_or_default();
+ let arrow_schema_metadata = metadata
+ .remove(super::ARROW_SCHEMA_META_KEY)
+ .map(|encoded| get_arrow_schema_from_metadata(&encoded));
+
+ match arrow_schema_metadata {
+ Some(Some(schema)) => Ok(schema),
+ _ => parquet_to_arrow_schema_by_columns(
+ parquet_schema,
+ 0..parquet_schema.columns().len(),
+ key_value_metadata,
+ ),
+ }
}
/// Convert parquet schema to arrow schema including optional metadata, only preserving some leaf columns.
@@ -81,6 +90,80 @@ where
.map(|fields| Schema::new_with_metadata(fields, metadata))
}
+/// Try to convert Arrow schema metadata into a schema
+fn get_arrow_schema_from_metadata(encoded_meta: &str) -> Option<Schema> {
+ let decoded = base64::decode(encoded_meta);
+ match decoded {
+ Ok(bytes) => {
+ let slice = if bytes[0..4] == [255u8; 4] {
+ &bytes[8..]
+ } else {
+ bytes.as_slice()
+ };
+ let message = arrow::ipc::get_root_as_message(slice);
+ message
+ .header_as_schema()
+ .map(arrow::ipc::convert::fb_to_schema)
+ }
+ Err(err) => {
+ // The C++ implementation returns an error if the schema can't be parsed.
+ // To prevent this, we explicitly log this, then compute the schema without the metadata
+ eprintln!(
+ "Unable to decode the encoded schema stored in {}, {:?}",
+ super::ARROW_SCHEMA_META_KEY,
+ err
+ );
+ None
+ }
+ }
+}
+
+/// Encodes the Arrow schema into the IPC format, and base64 encodes it
+fn encode_arrow_schema(schema: &Schema) -> String {
+ let mut serialized_schema = arrow::ipc::writer::schema_to_bytes(&schema);
+
+ // manually prepending the length to the schema as arrow uses the legacy IPC format
+ // TODO: change after addressing ARROW-9777
+ let schema_len = serialized_schema.len();
+ let mut len_prefix_schema = Vec::with_capacity(schema_len + 8);
+ len_prefix_schema.append(&mut vec![255u8, 255, 255, 255]);
+ len_prefix_schema.append((schema_len as u32).to_le_bytes().to_vec().as_mut());
+ len_prefix_schema.append(&mut serialized_schema);
+
+ base64::encode(&len_prefix_schema)
+}
+
+/// Mutates writer metadata by storing the encoded Arrow schema.
+/// If there is an existing Arrow schema metadata, it is replaced.
+pub(crate) fn add_encoded_arrow_schema_to_metadata(
+ schema: &Schema,
+ props: &mut WriterProperties,
+) {
+ let encoded = encode_arrow_schema(schema);
+
+ let schema_kv = KeyValue {
+ key: super::ARROW_SCHEMA_META_KEY.to_string(),
+ value: Some(encoded),
+ };
+
+ let mut meta = props.key_value_metadata.clone().unwrap_or_default();
+ // check if ARROW:schema exists, and overwrite it
+ let schema_meta = meta
+ .iter()
+ .enumerate()
+ .find(|(_, kv)| kv.key.as_str() == super::ARROW_SCHEMA_META_KEY);
+ match schema_meta {
+ Some((i, _)) => {
+ meta.remove(i);
+ meta.push(schema_kv);
+ }
+ None => {
+ meta.push(schema_kv);
+ }
+ }
+ props.key_value_metadata = Some(meta);
+}
+
/// Convert arrow schema to parquet schema
pub fn arrow_to_parquet_schema(schema: &Schema) -> Result<SchemaDescriptor> {
let fields: Result<Vec<TypePtr>> = schema
@@ -215,42 +298,48 @@ fn arrow_to_parquet_type(field: &Field) -> Result<Type> {
Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY)
.with_logical_type(LogicalType::INTERVAL)
.with_repetition(repetition)
- .with_length(3)
+ .with_length(12)
+ .build()
+ }
+ DataType::Binary | DataType::LargeBinary => {
+ Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
+ .with_repetition(repetition)
.build()
}
- DataType::Binary => Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
- .with_repetition(repetition)
- .build(),
DataType::FixedSizeBinary(length) => {
Type::primitive_type_builder(name, PhysicalType::FIXED_LEN_BYTE_ARRAY)
.with_repetition(repetition)
.with_length(*length)
.build()
}
- DataType::Utf8 => Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
- .with_logical_type(LogicalType::UTF8)
- .with_repetition(repetition)
- .build(),
- DataType::List(dtype) | DataType::FixedSizeList(dtype, _) => {
- Type::group_type_builder(name)
- .with_fields(&mut vec![Rc::new(
- Type::group_type_builder("list")
- .with_fields(&mut vec![Rc::new({
- let list_field = Field::new(
- "element",
- *dtype.clone(),
- field.is_nullable(),
- );
- arrow_to_parquet_type(&list_field)?
- })])
- .with_repetition(Repetition::REPEATED)
- .build()?,
- )])
- .with_logical_type(LogicalType::LIST)
- .with_repetition(Repetition::REQUIRED)
+ DataType::Utf8 | DataType::LargeUtf8 => {
+ Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
+ .with_logical_type(LogicalType::UTF8)
+ .with_repetition(repetition)
.build()
}
+ DataType::List(dtype)
+ | DataType::FixedSizeList(dtype, _)
+ | DataType::LargeList(dtype) => Type::group_type_builder(name)
+ .with_fields(&mut vec![Rc::new(
+ Type::group_type_builder("list")
+ .with_fields(&mut vec![Rc::new({
+ let list_field =
+ Field::new("element", *dtype.clone(), field.is_nullable());
+ arrow_to_parquet_type(&list_field)?
+ })])
+ .with_repetition(Repetition::REPEATED)
+ .build()?,
+ )])
+ .with_logical_type(LogicalType::LIST)
+ .with_repetition(Repetition::REQUIRED)
+ .build(),
DataType::Struct(fields) => {
+ if fields.is_empty() {
+ return Err(ArrowError(
+ "Parquet does not support writing empty structs".to_string(),
+ ));
+ }
// recursively convert children to types/nodes
let fields: Result<Vec<TypePtr>> = fields
.iter()
@@ -267,9 +356,6 @@ fn arrow_to_parquet_type(field: &Field) -> Result<Type> {
let dict_field = Field::new(name, *value.clone(), field.is_nullable());
arrow_to_parquet_type(&dict_field)
}
- DataType::LargeUtf8 | DataType::LargeBinary | DataType::LargeList(_) => {
- Err(ArrowError("Large arrays not supported".to_string()))
- }
}
}
/// This struct is used to group methods and data structures used to convert parquet
@@ -555,12 +641,16 @@ impl ParquetTypeConverter<'_> {
mod tests {
use super::*;
- use std::collections::HashMap;
+ use std::{collections::HashMap, convert::TryFrom, sync::Arc};
- use arrow::datatypes::{DataType, DateUnit, Field, TimeUnit};
+ use arrow::datatypes::{DataType, DateUnit, Field, IntervalUnit, TimeUnit};
- use crate::file::metadata::KeyValue;
- use crate::schema::{parser::parse_message_type, types::SchemaDescriptor};
+ use crate::file::{metadata::KeyValue, reader::SerializedFileReader};
+ use crate::{
+ arrow::{ArrowReader, ArrowWriter, ParquetFileArrowReader},
+ schema::{parser::parse_message_type, types::SchemaDescriptor},
+ util::test_common::get_temp_file,
+ };
#[test]
fn test_flat_primitives() {
@@ -1195,6 +1285,17 @@ mod tests {
}
#[test]
+ #[should_panic(expected = "Parquet does not support writing empty structs")]
+ fn test_empty_struct_field() {
+ let arrow_fields = vec![Field::new("struct", DataType::Struct(vec![]), false)];
+ let arrow_schema = Schema::new(arrow_fields);
+ let converted_arrow_schema = arrow_to_parquet_schema(&arrow_schema);
+
+ assert!(converted_arrow_schema.is_err());
+ converted_arrow_schema.unwrap();
+ }
+
+ #[test]
fn test_metadata() {
let message_type = "
message test_schema {
@@ -1216,4 +1317,123 @@ mod tests {
assert_eq!(converted_arrow_schema.metadata(), &expected_metadata);
}
+
+ #[test]
+ fn test_arrow_schema_roundtrip() -> Result<()> {
+ // This tests the roundtrip of an Arrow schema
+ // Fields that are commented out fail roundtrip tests or are unsupported by the writer
+ let metadata: HashMap<String, String> =
+ [("Key".to_string(), "Value".to_string())]
+ .iter()
+ .cloned()
+ .collect();
+
+ let schema = Schema::new_with_metadata(
+ vec![
+ Field::new("c1", DataType::Utf8, false),
+ Field::new("c2", DataType::Binary, false),
+ Field::new("c3", DataType::FixedSizeBinary(3), false),
+ Field::new("c4", DataType::Boolean, false),
+ Field::new("c5", DataType::Date32(DateUnit::Day), false),
+ Field::new("c6", DataType::Date64(DateUnit::Millisecond), false),
+ Field::new("c7", DataType::Time32(TimeUnit::Second), false),
+ Field::new("c8", DataType::Time32(TimeUnit::Millisecond), false),
+ Field::new("c13", DataType::Time64(TimeUnit::Microsecond), false),
+ Field::new("c14", DataType::Time64(TimeUnit::Nanosecond), false),
+ Field::new("c15", DataType::Timestamp(TimeUnit::Second, None), false),
+ Field::new(
+ "c16",
+ DataType::Timestamp(
+ TimeUnit::Millisecond,
+ Some(Arc::new("UTC".to_string())),
+ ),
+ false,
+ ),
+ Field::new(
+ "c17",
+ DataType::Timestamp(
+ TimeUnit::Microsecond,
+ Some(Arc::new("Africa/Johannesburg".to_string())),
+ ),
+ false,
+ ),
+ Field::new(
+ "c18",
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ false,
+ ),
+ Field::new("c19", DataType::Interval(IntervalUnit::DayTime), false),
+ Field::new("c20", DataType::Interval(IntervalUnit::YearMonth), false),
+ Field::new("c21", DataType::List(Box::new(DataType::Boolean)), false),
+ Field::new(
+ "c22",
+ DataType::FixedSizeList(Box::new(DataType::Boolean), 5),
+ false,
+ ),
+ Field::new(
+ "c23",
+ DataType::List(Box::new(DataType::List(Box::new(DataType::Struct(
+ vec![
+ Field::new("a", DataType::Int16, true),
+ Field::new("b", DataType::Float64, false),
+ ],
+ ))))),
+ true,
+ ),
+ Field::new(
+ "c24",
+ DataType::Struct(vec![
+ Field::new("a", DataType::Utf8, false),
+ Field::new("b", DataType::UInt16, false),
+ ]),
+ false,
+ ),
+ Field::new("c25", DataType::Interval(IntervalUnit::YearMonth), true),
+ Field::new("c26", DataType::Interval(IntervalUnit::DayTime), true),
+ // Field::new("c27", DataType::Duration(TimeUnit::Second), false),
+ // Field::new("c28", DataType::Duration(TimeUnit::Millisecond), false),
+ // Field::new("c29", DataType::Duration(TimeUnit::Microsecond), false),
+ // Field::new("c30", DataType::Duration(TimeUnit::Nanosecond), false),
+ // Field::new_dict(
+ // "c31",
+ // DataType::Dictionary(
+ // Box::new(DataType::Int32),
+ // Box::new(DataType::Utf8),
+ // ),
+ // true,
+ // 123,
+ // true,
+ // ),
+ Field::new("c32", DataType::LargeBinary, true),
+ Field::new("c33", DataType::LargeUtf8, true),
+ Field::new(
+ "c34",
+ DataType::LargeList(Box::new(DataType::LargeList(Box::new(
+ DataType::Struct(vec![
+ Field::new("a", DataType::Int16, true),
+ Field::new("b", DataType::Float64, true),
+ ]),
+ )))),
+ true,
+ ),
+ ],
+ metadata,
+ );
+
+ // write to an empty parquet file so that schema is serialized
+ let file = get_temp_file("test_arrow_schema_roundtrip.parquet", &[]);
+ let mut writer = ArrowWriter::try_new(
+ file.try_clone().unwrap(),
+ Arc::new(schema.clone()),
+ None,
+ )?;
+ writer.close()?;
+
+ // read file back
+ let parquet_reader = SerializedFileReader::try_from(file)?;
+ let mut arrow_reader = ParquetFileArrowReader::new(Rc::new(parquet_reader));
+ let read_schema = arrow_reader.get_schema()?;
+ assert_eq!(schema, read_schema);
+ Ok(())
+ }
}
diff --git a/rust/parquet/src/file/properties.rs b/rust/parquet/src/file/properties.rs
index 188d6ec..b62ce7b 100644
--- a/rust/parquet/src/file/properties.rs
+++ b/rust/parquet/src/file/properties.rs
@@ -89,8 +89,8 @@ pub type WriterPropertiesPtr = Rc<WriterProperties>;
/// Writer properties.
///
-/// It is created as an immutable data structure, use [`WriterPropertiesBuilder`] to
-/// assemble the properties.
+/// All properties except the key-value metadata are immutable,
+/// use [`WriterPropertiesBuilder`] to assemble these properties.
#[derive(Debug, Clone)]
pub struct WriterProperties {
data_pagesize_limit: usize,
@@ -99,7 +99,7 @@ pub struct WriterProperties {
max_row_group_size: usize,
writer_version: WriterVersion,
created_by: String,
- key_value_metadata: Option<Vec<KeyValue>>,
+ pub(crate) key_value_metadata: Option<Vec<KeyValue>>,
default_column_properties: ColumnProperties,
column_properties: HashMap<ColumnPath, ColumnProperties>,
}
[arrow] 04/09: ARROW-8426: [Rust] [Parquet] Add support for writing
dictionary types
Posted by ne...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit 50638549480645853e7414ee78ff0ff66e7fe4bc
Author: Carol (Nichols || Goulding) <ca...@gmail.com>
AuthorDate: Sat Oct 3 02:34:38 2020 +0200
ARROW-8426: [Rust] [Parquet] Add support for writing dictionary types
In this commit, I:
- Extracted a `build_field` function for some code shared between
`schema_to_fb` and `schema_to_fb_offset` that needed to change
- Uncommented the dictionary field from the Arrow schema roundtrip test
and add a dictionary field to the IPC roundtrip test
- If a field is a dictionary field, call `add_dictionary` with the
dictionary field information on the flatbuffer field, building the
dictionary as [the C++ code does][cpp-dictionary] and describe with the
same comment
- When getting the field type for a dictionary field, use the `value_type`
as [the C++ code does][cpp-value-type] and describe with the same
comment
The tests pass because the Parquet -> Arrow conversion for dictionaries
is [already supported][parquet-to-arrow].
[cpp-dictionary]: https://github.com/apache/arrow/blob/477c1021ac013f22389baf9154fb9ad0cf814bec/cpp/src/arrow/ipc/metadata_internal.cc#L426-L440
[cpp-value-type]: https://github.com/apache/arrow/blob/477c1021ac013f22389baf9154fb9ad0cf814bec/cpp/src/arrow/ipc/metadata_internal.cc#L662-L667
[parquet-to-arrow]: https://github.com/apache/arrow/blob/477c1021ac013f22389baf9154fb9ad0cf814bec/rust/arrow/src/ipc/convert.rs#L120-L127
Closes #8291 from carols10cents/rust-parquet-arrow-writer
Authored-by: Carol (Nichols || Goulding) <ca...@gmail.com>
Signed-off-by: Neville Dipale <ne...@gmail.com>
---
rust/arrow/src/datatypes.rs | 4 +-
rust/arrow/src/ipc/convert.rs | 105 ++++++++++++++++++++++++++++++---------
rust/parquet/src/arrow/schema.rs | 20 ++++----
3 files changed, 93 insertions(+), 36 deletions(-)
diff --git a/rust/arrow/src/datatypes.rs b/rust/arrow/src/datatypes.rs
index 3756542..6f0dc24 100644
--- a/rust/arrow/src/datatypes.rs
+++ b/rust/arrow/src/datatypes.rs
@@ -189,8 +189,8 @@ pub struct Field {
name: String,
data_type: DataType,
nullable: bool,
- dict_id: i64,
- dict_is_ordered: bool,
+ pub(crate) dict_id: i64,
+ pub(crate) dict_is_ordered: bool,
}
pub trait ArrowNativeType:
diff --git a/rust/arrow/src/ipc/convert.rs b/rust/arrow/src/ipc/convert.rs
index 7a5795d..8f429bf 100644
--- a/rust/arrow/src/ipc/convert.rs
+++ b/rust/arrow/src/ipc/convert.rs
@@ -34,18 +34,8 @@ pub fn schema_to_fb(schema: &Schema) -> FlatBufferBuilder {
let mut fields = vec![];
for field in schema.fields() {
- let fb_field_name = fbb.create_string(field.name().as_str());
- let field_type = get_fb_field_type(field.data_type(), &mut fbb);
- let mut field_builder = ipc::FieldBuilder::new(&mut fbb);
- field_builder.add_name(fb_field_name);
- field_builder.add_type_type(field_type.type_type);
- field_builder.add_nullable(field.is_nullable());
- match field_type.children {
- None => {}
- Some(children) => field_builder.add_children(children),
- };
- field_builder.add_type_(field_type.type_);
- fields.push(field_builder.finish());
+ let fb_field = build_field(&mut fbb, field);
+ fields.push(fb_field);
}
let mut custom_metadata = vec![];
@@ -80,18 +70,8 @@ pub fn schema_to_fb_offset<'a: 'b, 'b>(
) -> WIPOffset<ipc::Schema<'b>> {
let mut fields = vec![];
for field in schema.fields() {
- let fb_field_name = fbb.create_string(field.name().as_str());
- let field_type = get_fb_field_type(field.data_type(), fbb);
- let mut field_builder = ipc::FieldBuilder::new(fbb);
- field_builder.add_name(fb_field_name);
- field_builder.add_type_type(field_type.type_type);
- field_builder.add_nullable(field.is_nullable());
- match field_type.children {
- None => {}
- Some(children) => field_builder.add_children(children),
- };
- field_builder.add_type_(field_type.type_);
- fields.push(field_builder.finish());
+ let fb_field = build_field(fbb, field);
+ fields.push(fb_field);
}
let mut custom_metadata = vec![];
@@ -333,6 +313,38 @@ pub(crate) struct FBFieldType<'b> {
pub(crate) children: Option<WIPOffset<Vector<'b, ForwardsUOffset<ipc::Field<'b>>>>>,
}
+/// Create an IPC Field from an Arrow Field
+pub(crate) fn build_field<'a: 'b, 'b>(
+ fbb: &mut FlatBufferBuilder<'a>,
+ field: &Field,
+) -> WIPOffset<ipc::Field<'b>> {
+ let fb_field_name = fbb.create_string(field.name().as_str());
+ let field_type = get_fb_field_type(field.data_type(), fbb);
+
+ let fb_dictionary = if let Dictionary(index_type, _) = field.data_type() {
+ Some(get_fb_dictionary(
+ index_type,
+ field.dict_id,
+ field.dict_is_ordered,
+ fbb,
+ ))
+ } else {
+ None
+ };
+
+ let mut field_builder = ipc::FieldBuilder::new(fbb);
+ field_builder.add_name(fb_field_name);
+ fb_dictionary.map(|dictionary| field_builder.add_dictionary(dictionary));
+ field_builder.add_type_type(field_type.type_type);
+ field_builder.add_nullable(field.is_nullable());
+ match field_type.children {
+ None => {}
+ Some(children) => field_builder.add_children(children),
+ };
+ field_builder.add_type_(field_type.type_);
+ field_builder.finish()
+}
+
/// Get the IPC type of a data type
pub(crate) fn get_fb_field_type<'a: 'b, 'b>(
data_type: &DataType,
@@ -609,10 +621,45 @@ pub(crate) fn get_fb_field_type<'a: 'b, 'b>(
children: Some(fbb.create_vector(&children[..])),
}
}
+ Dictionary(_, value_type) => {
+ // In this library, the dictionary "type" is a logical construct. Here we
+ // pass through to the value type, as we've already captured the index
+ // type in the DictionaryEncoding metadata in the parent field
+ get_fb_field_type(value_type, fbb)
+ }
t => unimplemented!("Type {:?} not supported", t),
}
}
+/// Create an IPC dictionary encoding
+pub(crate) fn get_fb_dictionary<'a: 'b, 'b>(
+ index_type: &DataType,
+ dict_id: i64,
+ dict_is_ordered: bool,
+ fbb: &mut FlatBufferBuilder<'a>,
+) -> WIPOffset<ipc::DictionaryEncoding<'b>> {
+ // We assume that the dictionary index type (as an integer) has already been
+ // validated elsewhere, and can safely assume we are dealing with signed
+ // integers
+ let mut index_builder = ipc::IntBuilder::new(fbb);
+ index_builder.add_is_signed(true);
+ match *index_type {
+ Int8 => index_builder.add_bitWidth(8),
+ Int16 => index_builder.add_bitWidth(16),
+ Int32 => index_builder.add_bitWidth(32),
+ Int64 => index_builder.add_bitWidth(64),
+ _ => {}
+ }
+ let index_builder = index_builder.finish();
+
+ let mut builder = ipc::DictionaryEncodingBuilder::new(fbb);
+ builder.add_id(dict_id);
+ builder.add_indexType(index_builder);
+ builder.add_isOrdered(dict_is_ordered);
+
+ builder.finish()
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -714,6 +761,16 @@ mod tests {
false,
),
Field::new("struct<>", DataType::Struct(vec![]), true),
+ Field::new_dict(
+ "dictionary<int32, utf8>",
+ DataType::Dictionary(
+ Box::new(DataType::Int32),
+ Box::new(DataType::Utf8),
+ ),
+ true,
+ 123,
+ true,
+ ),
],
md,
);
diff --git a/rust/parquet/src/arrow/schema.rs b/rust/parquet/src/arrow/schema.rs
index d5a0ff9..4a92a46 100644
--- a/rust/parquet/src/arrow/schema.rs
+++ b/rust/parquet/src/arrow/schema.rs
@@ -1396,16 +1396,16 @@ mod tests {
// Field::new("c28", DataType::Duration(TimeUnit::Millisecond), false),
// Field::new("c29", DataType::Duration(TimeUnit::Microsecond), false),
// Field::new("c30", DataType::Duration(TimeUnit::Nanosecond), false),
- // Field::new_dict(
- // "c31",
- // DataType::Dictionary(
- // Box::new(DataType::Int32),
- // Box::new(DataType::Utf8),
- // ),
- // true,
- // 123,
- // true,
- // ),
+ Field::new_dict(
+ "c31",
+ DataType::Dictionary(
+ Box::new(DataType::Int32),
+ Box::new(DataType::Utf8),
+ ),
+ true,
+ 123,
+ true,
+ ),
Field::new("c32", DataType::LargeBinary, true),
Field::new("c33", DataType::LargeUtf8, true),
Field::new(
[arrow] 06/09: ARROW-10168: [Rust] [Parquet] Schema roundtrip - use
Arrow schema from Parquet metadata when available
Posted by ne...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit 96798dc089fc2b1add212c2a77b1bf33501736be
Author: Carol (Nichols || Goulding) <ca...@gmail.com>
AuthorDate: Thu Oct 8 00:16:42 2020 +0200
ARROW-10168: [Rust] [Parquet] Schema roundtrip - use Arrow schema from Parquet metadata when available
@nevi-me This is one commit on top of https://github.com/apache/arrow/pull/8330 that I'm opening to get some feedback from you on about whether this will help with ARROW-10168. I *think* this will bring the Rust implementation more in line with C++, but I'm not certain.
I tried removing the `#[ignore]` attributes from the `LargeArray` and `LargeUtf8` tests, but they're still failing because the schemas don't match yet-- it looks like [this code](https://github.com/apache/arrow/blob/b2842ab2eb0d7a7a633049a5591e1eaa254d4446/rust/parquet/src/arrow/array_reader.rs#L595-L638) will need to be changed as well.
That `build_array_reader` function's code looks very similar to the code I've changed here, is there a possibility for the code to be shared or is there a reason they're separate?
Closes #8354 from carols10cents/schema-roundtrip
Lead-authored-by: Carol (Nichols || Goulding) <ca...@gmail.com>
Co-authored-by: Neville Dipale <ne...@gmail.com>
Signed-off-by: Neville Dipale <ne...@gmail.com>
---
rust/arrow/src/ipc/convert.rs | 4 +-
rust/parquet/src/arrow/array_reader.rs | 106 +++++++++++++----
rust/parquet/src/arrow/arrow_reader.rs | 36 ++++--
rust/parquet/src/arrow/arrow_writer.rs | 4 +-
rust/parquet/src/arrow/converter.rs | 52 +++++++-
rust/parquet/src/arrow/mod.rs | 3 +-
rust/parquet/src/arrow/record_reader.rs | 1 +
rust/parquet/src/arrow/schema.rs | 205 +++++++++++++++++++++++++++-----
8 files changed, 338 insertions(+), 73 deletions(-)
diff --git a/rust/arrow/src/ipc/convert.rs b/rust/arrow/src/ipc/convert.rs
index 8f429bf..a02b6c4 100644
--- a/rust/arrow/src/ipc/convert.rs
+++ b/rust/arrow/src/ipc/convert.rs
@@ -334,7 +334,9 @@ pub(crate) fn build_field<'a: 'b, 'b>(
let mut field_builder = ipc::FieldBuilder::new(fbb);
field_builder.add_name(fb_field_name);
- fb_dictionary.map(|dictionary| field_builder.add_dictionary(dictionary));
+ if let Some(dictionary) = fb_dictionary {
+ field_builder.add_dictionary(dictionary)
+ }
field_builder.add_type_type(field_type.type_type);
field_builder.add_nullable(field.is_nullable());
match field_type.children {
diff --git a/rust/parquet/src/arrow/array_reader.rs b/rust/parquet/src/arrow/array_reader.rs
index 4fbc54d..40df284 100644
--- a/rust/parquet/src/arrow/array_reader.rs
+++ b/rust/parquet/src/arrow/array_reader.rs
@@ -29,16 +29,20 @@ use arrow::array::{
Int16BufferBuilder, StructArray,
};
use arrow::buffer::{Buffer, MutableBuffer};
-use arrow::datatypes::{DataType as ArrowType, DateUnit, Field, IntervalUnit, TimeUnit};
+use arrow::datatypes::{
+ DataType as ArrowType, DateUnit, Field, IntervalUnit, Schema, TimeUnit,
+};
use crate::arrow::converter::{
BinaryArrayConverter, BinaryConverter, BoolConverter, BooleanArrayConverter,
Converter, Date32Converter, FixedLenBinaryConverter, FixedSizeArrayConverter,
Float32Converter, Float64Converter, Int16Converter, Int32Converter, Int64Converter,
- Int8Converter, Int96ArrayConverter, Int96Converter, Time32MillisecondConverter,
- Time32SecondConverter, Time64MicrosecondConverter, Time64NanosecondConverter,
- TimestampMicrosecondConverter, TimestampMillisecondConverter, UInt16Converter,
- UInt32Converter, UInt64Converter, UInt8Converter, Utf8ArrayConverter, Utf8Converter,
+ Int8Converter, Int96ArrayConverter, Int96Converter, LargeBinaryArrayConverter,
+ LargeBinaryConverter, LargeUtf8ArrayConverter, LargeUtf8Converter,
+ Time32MillisecondConverter, Time32SecondConverter, Time64MicrosecondConverter,
+ Time64NanosecondConverter, TimestampMicrosecondConverter,
+ TimestampMillisecondConverter, UInt16Converter, UInt32Converter, UInt64Converter,
+ UInt8Converter, Utf8ArrayConverter, Utf8Converter,
};
use crate::arrow::record_reader::RecordReader;
use crate::arrow::schema::parquet_to_arrow_field;
@@ -612,6 +616,7 @@ impl ArrayReader for StructArrayReader {
/// Create array reader from parquet schema, column indices, and parquet file reader.
pub fn build_array_reader<T>(
parquet_schema: SchemaDescPtr,
+ arrow_schema: Schema,
column_indices: T,
file_reader: Rc<dyn FileReader>,
) -> Result<Box<dyn ArrayReader>>
@@ -650,13 +655,19 @@ where
fields: filtered_root_fields,
};
- ArrayReaderBuilder::new(Rc::new(proj), Rc::new(leaves), file_reader)
- .build_array_reader()
+ ArrayReaderBuilder::new(
+ Rc::new(proj),
+ Rc::new(arrow_schema),
+ Rc::new(leaves),
+ file_reader,
+ )
+ .build_array_reader()
}
/// Used to build array reader.
struct ArrayReaderBuilder {
root_schema: TypePtr,
+ arrow_schema: Rc<Schema>,
// Key: columns that need to be included in final array builder
// Value: column index in schema
columns_included: Rc<HashMap<*const Type, usize>>,
@@ -790,11 +801,13 @@ impl<'a> ArrayReaderBuilder {
/// Construct array reader builder.
fn new(
root_schema: TypePtr,
+ arrow_schema: Rc<Schema>,
columns_included: Rc<HashMap<*const Type, usize>>,
file_reader: Rc<dyn FileReader>,
) -> Self {
Self {
root_schema,
+ arrow_schema,
columns_included,
file_reader,
}
@@ -835,6 +848,12 @@ impl<'a> ArrayReaderBuilder {
self.file_reader.clone(),
)?);
+ let arrow_type = self
+ .arrow_schema
+ .field_with_name(cur_type.name())
+ .ok()
+ .map(|f| f.data_type());
+
match cur_type.get_physical_type() {
PhysicalType::BOOLEAN => Ok(Box::new(PrimitiveArrayReader::<BoolType>::new(
page_iterator,
@@ -866,21 +885,43 @@ impl<'a> ArrayReaderBuilder {
)),
PhysicalType::BYTE_ARRAY => {
if cur_type.get_basic_info().logical_type() == LogicalType::UTF8 {
- let converter = Utf8Converter::new(Utf8ArrayConverter {});
- Ok(Box::new(ComplexObjectArrayReader::<
- ByteArrayType,
- Utf8Converter,
- >::new(
- page_iterator, column_desc, converter
- )?))
+ if let Some(ArrowType::LargeUtf8) = arrow_type {
+ let converter =
+ LargeUtf8Converter::new(LargeUtf8ArrayConverter {});
+ Ok(Box::new(ComplexObjectArrayReader::<
+ ByteArrayType,
+ LargeUtf8Converter,
+ >::new(
+ page_iterator, column_desc, converter
+ )?))
+ } else {
+ let converter = Utf8Converter::new(Utf8ArrayConverter {});
+ Ok(Box::new(ComplexObjectArrayReader::<
+ ByteArrayType,
+ Utf8Converter,
+ >::new(
+ page_iterator, column_desc, converter
+ )?))
+ }
} else {
- let converter = BinaryConverter::new(BinaryArrayConverter {});
- Ok(Box::new(ComplexObjectArrayReader::<
- ByteArrayType,
- BinaryConverter,
- >::new(
- page_iterator, column_desc, converter
- )?))
+ if let Some(ArrowType::LargeBinary) = arrow_type {
+ let converter =
+ LargeBinaryConverter::new(LargeBinaryArrayConverter {});
+ Ok(Box::new(ComplexObjectArrayReader::<
+ ByteArrayType,
+ LargeBinaryConverter,
+ >::new(
+ page_iterator, column_desc, converter
+ )?))
+ } else {
+ let converter = BinaryConverter::new(BinaryArrayConverter {});
+ Ok(Box::new(ComplexObjectArrayReader::<
+ ByteArrayType,
+ BinaryConverter,
+ >::new(
+ page_iterator, column_desc, converter
+ )?))
+ }
}
}
PhysicalType::FIXED_LEN_BYTE_ARRAY => {
@@ -918,11 +959,15 @@ impl<'a> ArrayReaderBuilder {
for child in cur_type.get_fields() {
if let Some(child_reader) = self.dispatch(child.clone(), context)? {
- fields.push(Field::new(
- child.name(),
- child_reader.get_data_type().clone(),
- child.is_optional(),
- ));
+ let field = match self.arrow_schema.field_with_name(child.name()) {
+ Ok(f) => f.to_owned(),
+ _ => Field::new(
+ child.name(),
+ child_reader.get_data_type().clone(),
+ child.is_optional(),
+ ),
+ };
+ fields.push(field);
children_reader.push(child_reader);
}
}
@@ -945,6 +990,7 @@ impl<'a> ArrayReaderBuilder {
mod tests {
use super::*;
use crate::arrow::converter::Utf8Converter;
+ use crate::arrow::schema::parquet_to_arrow_schema;
use crate::basic::{Encoding, Type as PhysicalType};
use crate::column::page::{Page, PageReader};
use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
@@ -1591,8 +1637,16 @@ mod tests {
let file = get_test_file("nulls.snappy.parquet");
let file_reader = Rc::new(SerializedFileReader::new(file).unwrap());
+ let file_metadata = file_reader.metadata().file_metadata();
+ let arrow_schema = parquet_to_arrow_schema(
+ file_metadata.schema_descr(),
+ file_metadata.key_value_metadata(),
+ )
+ .unwrap();
+
let array_reader = build_array_reader(
file_reader.metadata().file_metadata().schema_descr_ptr(),
+ arrow_schema,
vec![0usize].into_iter(),
file_reader,
)
diff --git a/rust/parquet/src/arrow/arrow_reader.rs b/rust/parquet/src/arrow/arrow_reader.rs
index b654de1..88af583 100644
--- a/rust/parquet/src/arrow/arrow_reader.rs
+++ b/rust/parquet/src/arrow/arrow_reader.rs
@@ -19,7 +19,9 @@
use crate::arrow::array_reader::{build_array_reader, ArrayReader, StructArrayReader};
use crate::arrow::schema::parquet_to_arrow_schema;
-use crate::arrow::schema::parquet_to_arrow_schema_by_columns;
+use crate::arrow::schema::{
+ parquet_to_arrow_schema_by_columns, parquet_to_arrow_schema_by_root_columns,
+};
use crate::errors::{ParquetError, Result};
use crate::file::reader::FileReader;
use arrow::datatypes::{DataType as ArrowType, Schema, SchemaRef};
@@ -40,7 +42,12 @@ pub trait ArrowReader {
/// Read parquet schema and convert it into arrow schema.
/// This schema only includes columns identified by `column_indices`.
- fn get_schema_by_columns<T>(&mut self, column_indices: T) -> Result<Schema>
+ /// To select leaf columns (i.e. `a.b.c` instead of `a`), set `leaf_columns = true`
+ fn get_schema_by_columns<T>(
+ &mut self,
+ column_indices: T,
+ leaf_columns: bool,
+ ) -> Result<Schema>
where
T: IntoIterator<Item = usize>;
@@ -84,16 +91,28 @@ impl ArrowReader for ParquetFileArrowReader {
)
}
- fn get_schema_by_columns<T>(&mut self, column_indices: T) -> Result<Schema>
+ fn get_schema_by_columns<T>(
+ &mut self,
+ column_indices: T,
+ leaf_columns: bool,
+ ) -> Result<Schema>
where
T: IntoIterator<Item = usize>,
{
let file_metadata = self.file_reader.metadata().file_metadata();
- parquet_to_arrow_schema_by_columns(
- file_metadata.schema_descr(),
- column_indices,
- file_metadata.key_value_metadata(),
- )
+ if leaf_columns {
+ parquet_to_arrow_schema_by_columns(
+ file_metadata.schema_descr(),
+ column_indices,
+ file_metadata.key_value_metadata(),
+ )
+ } else {
+ parquet_to_arrow_schema_by_root_columns(
+ file_metadata.schema_descr(),
+ column_indices,
+ file_metadata.key_value_metadata(),
+ )
+ }
}
fn get_record_reader(
@@ -123,6 +142,7 @@ impl ArrowReader for ParquetFileArrowReader {
.metadata()
.file_metadata()
.schema_descr_ptr(),
+ self.get_schema()?,
column_indices,
self.file_reader.clone(),
)?;
diff --git a/rust/parquet/src/arrow/arrow_writer.rs b/rust/parquet/src/arrow/arrow_writer.rs
index cf7b9a2..40e2553 100644
--- a/rust/parquet/src/arrow/arrow_writer.rs
+++ b/rust/parquet/src/arrow/arrow_writer.rs
@@ -1012,7 +1012,7 @@ mod tests {
}
#[test]
- #[ignore] // Large Binary support isn't correct yet
+ #[ignore] // Large binary support isn't correct yet - buffers don't match
fn large_binary_single_column() {
let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
let many_vecs: Vec<_> = std::iter::repeat(one_vec).take(SMALL_SIZE).collect();
@@ -1035,7 +1035,7 @@ mod tests {
}
#[test]
- #[ignore] // Large String support isn't correct yet - null_bitmap and buffers don't match
+ #[ignore] // Large string support isn't correct yet - null_bitmap doesn't match
fn large_string_single_column() {
let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
let raw_strs = raw_values.iter().map(|s| s.as_str());
diff --git a/rust/parquet/src/arrow/converter.rs b/rust/parquet/src/arrow/converter.rs
index 86f6492..1aceba2 100644
--- a/rust/parquet/src/arrow/converter.rs
+++ b/rust/parquet/src/arrow/converter.rs
@@ -21,8 +21,8 @@ use crate::data_type::{ByteArray, DataType, Int96};
use arrow::{
array::{
Array, ArrayRef, BinaryBuilder, BooleanArray, BooleanBufferBuilder,
- BufferBuilderTrait, FixedSizeBinaryBuilder, StringBuilder,
- TimestampNanosecondBuilder,
+ BufferBuilderTrait, FixedSizeBinaryBuilder, LargeBinaryBuilder,
+ LargeStringBuilder, StringBuilder, TimestampNanosecondBuilder,
},
datatypes::Time32MillisecondType,
};
@@ -38,8 +38,8 @@ use arrow::datatypes::{ArrowPrimitiveType, DataType as ArrowDataType};
use arrow::array::ArrayDataBuilder;
use arrow::array::{
- BinaryArray, FixedSizeBinaryArray, PrimitiveArray, StringArray,
- TimestampNanosecondArray,
+ BinaryArray, FixedSizeBinaryArray, LargeBinaryArray, LargeStringArray,
+ PrimitiveArray, StringArray, TimestampNanosecondArray,
};
use std::marker::PhantomData;
@@ -200,6 +200,27 @@ impl Converter<Vec<Option<ByteArray>>, StringArray> for Utf8ArrayConverter {
}
}
+pub struct LargeUtf8ArrayConverter {}
+
+impl Converter<Vec<Option<ByteArray>>, LargeStringArray> for LargeUtf8ArrayConverter {
+ fn convert(&self, source: Vec<Option<ByteArray>>) -> Result<LargeStringArray> {
+ let data_size = source
+ .iter()
+ .map(|x| x.as_ref().map(|b| b.len()).unwrap_or(0))
+ .sum();
+
+ let mut builder = LargeStringBuilder::with_capacity(source.len(), data_size);
+ for v in source {
+ match v {
+ Some(array) => builder.append_value(array.as_utf8()?),
+ None => builder.append_null(),
+ }?
+ }
+
+ Ok(builder.finish())
+ }
+}
+
pub struct BinaryArrayConverter {}
impl Converter<Vec<Option<ByteArray>>, BinaryArray> for BinaryArrayConverter {
@@ -216,6 +237,22 @@ impl Converter<Vec<Option<ByteArray>>, BinaryArray> for BinaryArrayConverter {
}
}
+pub struct LargeBinaryArrayConverter {}
+
+impl Converter<Vec<Option<ByteArray>>, LargeBinaryArray> for LargeBinaryArrayConverter {
+ fn convert(&self, source: Vec<Option<ByteArray>>) -> Result<LargeBinaryArray> {
+ let mut builder = LargeBinaryBuilder::new(source.len());
+ for v in source {
+ match v {
+ Some(array) => builder.append_value(array.data()),
+ None => builder.append_null(),
+ }?
+ }
+
+ Ok(builder.finish())
+ }
+}
+
pub type BoolConverter<'a> = ArrayRefConverter<
&'a mut RecordReader<BoolType>,
BooleanArray,
@@ -246,8 +283,15 @@ pub type Float32Converter = CastConverter<ParquetFloatType, Float32Type, Float32
pub type Float64Converter = CastConverter<ParquetDoubleType, Float64Type, Float64Type>;
pub type Utf8Converter =
ArrayRefConverter<Vec<Option<ByteArray>>, StringArray, Utf8ArrayConverter>;
+pub type LargeUtf8Converter =
+ ArrayRefConverter<Vec<Option<ByteArray>>, LargeStringArray, LargeUtf8ArrayConverter>;
pub type BinaryConverter =
ArrayRefConverter<Vec<Option<ByteArray>>, BinaryArray, BinaryArrayConverter>;
+pub type LargeBinaryConverter = ArrayRefConverter<
+ Vec<Option<ByteArray>>,
+ LargeBinaryArray,
+ LargeBinaryArrayConverter,
+>;
pub type Int96Converter =
ArrayRefConverter<Vec<Option<Int96>>, TimestampNanosecondArray, Int96ArrayConverter>;
pub type FixedLenBinaryConverter = ArrayRefConverter<
diff --git a/rust/parquet/src/arrow/mod.rs b/rust/parquet/src/arrow/mod.rs
index 2b012fb..9793457 100644
--- a/rust/parquet/src/arrow/mod.rs
+++ b/rust/parquet/src/arrow/mod.rs
@@ -35,7 +35,7 @@
//!
//! println!("Converted arrow schema is: {}", arrow_reader.get_schema().unwrap());
//! println!("Arrow schema after projection is: {}",
-//! arrow_reader.get_schema_by_columns(vec![2, 4, 6]).unwrap());
+//! arrow_reader.get_schema_by_columns(vec![2, 4, 6], true).unwrap());
//!
//! let mut record_batch_reader = arrow_reader.get_record_reader(2048).unwrap();
//!
@@ -61,6 +61,7 @@ pub use self::arrow_reader::ParquetFileArrowReader;
pub use self::arrow_writer::ArrowWriter;
pub use self::schema::{
arrow_to_parquet_schema, parquet_to_arrow_schema, parquet_to_arrow_schema_by_columns,
+ parquet_to_arrow_schema_by_root_columns,
};
/// Schema metadata key used to store serialized Arrow IPC schema
diff --git a/rust/parquet/src/arrow/record_reader.rs b/rust/parquet/src/arrow/record_reader.rs
index ccfdaf8..b30ab77 100644
--- a/rust/parquet/src/arrow/record_reader.rs
+++ b/rust/parquet/src/arrow/record_reader.rs
@@ -86,6 +86,7 @@ impl<'a, T> FatPtr<'a, T> {
self.ptr
}
+ #[allow(clippy::wrong_self_convention)]
fn to_slice_mut(&mut self) -> &mut [T] {
self.ptr
}
diff --git a/rust/parquet/src/arrow/schema.rs b/rust/parquet/src/arrow/schema.rs
index 4a92a46..0cd41fe 100644
--- a/rust/parquet/src/arrow/schema.rs
+++ b/rust/parquet/src/arrow/schema.rs
@@ -56,7 +56,61 @@ pub fn parquet_to_arrow_schema(
}
}
-/// Convert parquet schema to arrow schema including optional metadata, only preserving some leaf columns.
+/// Convert parquet schema to arrow schema including optional metadata,
+/// only preserving some root columns.
+/// This is useful if we have columns `a.b`, `a.c.e` and `a.d`,
+/// and want `a` with all its child fields
+pub fn parquet_to_arrow_schema_by_root_columns<T>(
+ parquet_schema: &SchemaDescriptor,
+ column_indices: T,
+ key_value_metadata: &Option<Vec<KeyValue>>,
+) -> Result<Schema>
+where
+ T: IntoIterator<Item = usize>,
+{
+ // Reconstruct the index ranges of the parent columns
+ // An Arrow struct gets represented by 1+ columns based on how many child fields the
+ // struct has. This means that getting fields 1 and 2 might return the struct twice,
+ // if field 1 is the struct having say 3 fields, and field 2 is a primitive.
+ //
+ // The below gets the parent columns, and counts the number of child fields in each parent,
+ // such that we would end up with:
+ // - field 1 - columns: [0, 1, 2]
+ // - field 2 - columns: [3]
+ let mut parent_columns = vec![];
+ let mut curr_name = "";
+ let mut prev_name = "";
+ let mut indices = vec![];
+ (0..(parquet_schema.num_columns())).for_each(|i| {
+ let p_type = parquet_schema.get_column_root(i);
+ curr_name = p_type.get_basic_info().name();
+ if prev_name == "" {
+ // first index
+ indices.push(i);
+ prev_name = curr_name;
+ } else if curr_name != prev_name {
+ prev_name = curr_name;
+ parent_columns.push((curr_name.to_string(), indices.clone()));
+ indices = vec![i];
+ } else {
+ indices.push(i);
+ }
+ });
+ // push the last column if indices has values
+ if !indices.is_empty() {
+ parent_columns.push((curr_name.to_string(), indices));
+ }
+
+ // gather the required leaf columns
+ let leaf_columns = column_indices
+ .into_iter()
+ .flat_map(|i| parent_columns[i].1.clone());
+
+ parquet_to_arrow_schema_by_columns(parquet_schema, leaf_columns, key_value_metadata)
+}
+
+/// Convert parquet schema to arrow schema including optional metadata,
+/// only preserving some leaf columns.
pub fn parquet_to_arrow_schema_by_columns<T>(
parquet_schema: &SchemaDescriptor,
column_indices: T,
@@ -65,27 +119,56 @@ pub fn parquet_to_arrow_schema_by_columns<T>(
where
T: IntoIterator<Item = usize>,
{
+ let mut metadata = parse_key_value_metadata(key_value_metadata).unwrap_or_default();
+ let arrow_schema_metadata = metadata
+ .remove(super::ARROW_SCHEMA_META_KEY)
+ .map(|encoded| get_arrow_schema_from_metadata(&encoded))
+ .unwrap_or_default();
+
+ // add the Arrow metadata to the Parquet metadata
+ if let Some(arrow_schema) = &arrow_schema_metadata {
+ arrow_schema.metadata().iter().for_each(|(k, v)| {
+ metadata.insert(k.clone(), v.clone());
+ });
+ }
+
let mut base_nodes = Vec::new();
let mut base_nodes_set = HashSet::new();
let mut leaves = HashSet::new();
+ enum FieldType<'a> {
+ Parquet(&'a Type),
+ Arrow(Field),
+ }
+
for c in column_indices {
- let column = parquet_schema.column(c).self_type() as *const Type;
- let root = parquet_schema.get_column_root(c);
- let root_raw_ptr = root as *const Type;
-
- leaves.insert(column);
- if !base_nodes_set.contains(&root_raw_ptr) {
- base_nodes.push(root);
- base_nodes_set.insert(root_raw_ptr);
+ let column = parquet_schema.column(c);
+ let name = column.name();
+
+ if let Some(field) = arrow_schema_metadata
+ .as_ref()
+ .and_then(|schema| schema.field_with_name(name).ok().cloned())
+ {
+ base_nodes.push(FieldType::Arrow(field));
+ } else {
+ let column = column.self_type() as *const Type;
+ let root = parquet_schema.get_column_root(c);
+ let root_raw_ptr = root as *const Type;
+
+ leaves.insert(column);
+ if !base_nodes_set.contains(&root_raw_ptr) {
+ base_nodes.push(FieldType::Parquet(root));
+ base_nodes_set.insert(root_raw_ptr);
+ }
}
}
- let metadata = parse_key_value_metadata(key_value_metadata).unwrap_or_default();
-
base_nodes
.into_iter()
- .map(|t| ParquetTypeConverter::new(t, &leaves).to_field())
+ .map(|t| match t {
+ FieldType::Parquet(t) => ParquetTypeConverter::new(t, &leaves).to_field(),
+ FieldType::Arrow(f) => Ok(Some(f)),
+ })
.collect::<Result<Vec<Option<Field>>>>()
.map(|result| result.into_iter().filter_map(|f| f).collect::<Vec<Field>>())
.map(|fields| Schema::new_with_metadata(fields, metadata))
@@ -1367,21 +1450,21 @@ mod tests {
Field::new("c19", DataType::Interval(IntervalUnit::DayTime), false),
Field::new("c20", DataType::Interval(IntervalUnit::YearMonth), false),
Field::new("c21", DataType::List(Box::new(DataType::Boolean)), false),
- Field::new(
- "c22",
- DataType::FixedSizeList(Box::new(DataType::Boolean), 5),
- false,
- ),
- Field::new(
- "c23",
- DataType::List(Box::new(DataType::List(Box::new(DataType::Struct(
- vec![
- Field::new("a", DataType::Int16, true),
- Field::new("b", DataType::Float64, false),
- ],
- ))))),
- true,
- ),
+ // Field::new(
+ // "c22",
+ // DataType::FixedSizeList(Box::new(DataType::Boolean), 5),
+ // false,
+ // ),
+ // Field::new(
+ // "c23",
+ // DataType::List(Box::new(DataType::LargeList(Box::new(
+ // DataType::Struct(vec![
+ // Field::new("a", DataType::Int16, true),
+ // Field::new("b", DataType::Float64, false),
+ // ]),
+ // )))),
+ // true,
+ // ),
Field::new(
"c24",
DataType::Struct(vec![
@@ -1408,12 +1491,66 @@ mod tests {
),
Field::new("c32", DataType::LargeBinary, true),
Field::new("c33", DataType::LargeUtf8, true),
+ // Field::new(
+ // "c34",
+ // DataType::LargeList(Box::new(DataType::List(Box::new(
+ // DataType::Struct(vec![
+ // Field::new("a", DataType::Int16, true),
+ // Field::new("b", DataType::Float64, true),
+ // ]),
+ // )))),
+ // true,
+ // ),
+ ],
+ metadata,
+ );
+
+ // write to an empty parquet file so that schema is serialized
+ let file = get_temp_file("test_arrow_schema_roundtrip.parquet", &[]);
+ let mut writer = ArrowWriter::try_new(
+ file.try_clone().unwrap(),
+ Arc::new(schema.clone()),
+ None,
+ )?;
+ writer.close()?;
+
+ // read file back
+ let parquet_reader = SerializedFileReader::try_from(file)?;
+ let mut arrow_reader = ParquetFileArrowReader::new(Rc::new(parquet_reader));
+ let read_schema = arrow_reader.get_schema()?;
+ assert_eq!(schema, read_schema);
+
+ // read all fields by columns
+ let partial_read_schema =
+ arrow_reader.get_schema_by_columns(0..(schema.fields().len()), false)?;
+ assert_eq!(schema, partial_read_schema);
+
+ Ok(())
+ }
+
+ #[test]
+ #[ignore = "Roundtrip of lists currently fails because we don't check their types correctly in the Arrow schema"]
+ fn test_arrow_schema_roundtrip_lists() -> Result<()> {
+ let metadata: HashMap<String, String> =
+ [("Key".to_string(), "Value".to_string())]
+ .iter()
+ .cloned()
+ .collect();
+
+ let schema = Schema::new_with_metadata(
+ vec![
+ Field::new("c21", DataType::List(Box::new(DataType::Boolean)), false),
Field::new(
- "c34",
- DataType::LargeList(Box::new(DataType::LargeList(Box::new(
+ "c22",
+ DataType::FixedSizeList(Box::new(DataType::Boolean), 5),
+ false,
+ ),
+ Field::new(
+ "c23",
+ DataType::List(Box::new(DataType::LargeList(Box::new(
DataType::Struct(vec![
Field::new("a", DataType::Int16, true),
- Field::new("b", DataType::Float64, true),
+ Field::new("b", DataType::Float64, false),
]),
)))),
true,
@@ -1423,7 +1560,7 @@ mod tests {
);
// write to an empty parquet file so that schema is serialized
- let file = get_temp_file("test_arrow_schema_roundtrip.parquet", &[]);
+ let file = get_temp_file("test_arrow_schema_roundtrip_lists.parquet", &[]);
let mut writer = ArrowWriter::try_new(
file.try_clone().unwrap(),
Arc::new(schema.clone()),
@@ -1436,6 +1573,12 @@ mod tests {
let mut arrow_reader = ParquetFileArrowReader::new(Rc::new(parquet_reader));
let read_schema = arrow_reader.get_schema()?;
assert_eq!(schema, read_schema);
+
+ // read all fields by columns
+ let partial_read_schema =
+ arrow_reader.get_schema_by_columns(0..(schema.fields().len()), false)?;
+ assert_eq!(schema, partial_read_schema);
+
Ok(())
}
}
[arrow] 08/09: ARROW-10334: [Rust] [Parquet] NullArray roundtrip
Posted by ne...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
nevime pushed a commit to branch rust-parquet-arrow-writer
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit 0b5a84671dc46d509f966a29f406aec4c71d8e4d
Author: Neville Dipale <ne...@gmail.com>
AuthorDate: Sat Oct 17 21:04:39 2020 +0200
ARROW-10334: [Rust] [Parquet] NullArray roundtrip
This allows writing an Arrow NullArray to Parquet.
Support was added a few years ago in Parquet, and the C++ implementation supports writing null arrays.
The array is stored as an int32 which has all values set as null.
In order to implement this, we introduce a `null -> int32` cast, which creates a null int32 of same length.
Semantically, the write is the same as writing an int32 that's all null, but we create a null writer to preserve the data type.
Closes #8484 from nevi-me/ARROW-10334
Authored-by: Neville Dipale <ne...@gmail.com>
Signed-off-by: Neville Dipale <ne...@gmail.com>
---
rust/arrow/src/array/null.rs | 8 +-
rust/arrow/src/compute/kernels/cast.rs | 81 ++++++++++++-------
rust/parquet/src/arrow/array_reader.rs | 142 ++++++++++++++++++++++++++++-----
rust/parquet/src/arrow/arrow_writer.rs | 32 +++++---
rust/parquet/src/arrow/schema.rs | 6 +-
5 files changed, 205 insertions(+), 64 deletions(-)
diff --git a/rust/arrow/src/array/null.rs b/rust/arrow/src/array/null.rs
index 190d2fa..08c7cf1 100644
--- a/rust/arrow/src/array/null.rs
+++ b/rust/arrow/src/array/null.rs
@@ -113,7 +113,7 @@ impl From<ArrayDataRef> for NullArray {
impl fmt::Debug for NullArray {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "NullArray")
+ write!(f, "NullArray({})", self.len())
}
}
@@ -146,4 +146,10 @@ mod tests {
assert_eq!(array2.null_count(), 16);
assert_eq!(array2.offset(), 8);
}
+
+ #[test]
+ fn test_debug_null_array() {
+ let array = NullArray::new(1024 * 1024);
+ assert_eq!(format!("{:?}", array), "NullArray(1048576)");
+ }
}
diff --git a/rust/arrow/src/compute/kernels/cast.rs b/rust/arrow/src/compute/kernels/cast.rs
index 4e1bc85..ab34c6a 100644
--- a/rust/arrow/src/compute/kernels/cast.rs
+++ b/rust/arrow/src/compute/kernels/cast.rs
@@ -200,8 +200,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
(Timestamp(_, _), Date32(_)) => true,
(Timestamp(_, _), Date64(_)) => true,
// date64 to timestamp might not make sense,
-
- // end temporal casts
+ (Null, Int32) => true,
(_, _) => false,
}
}
@@ -729,25 +728,31 @@ pub fn cast(array: &ArrayRef, to_type: &DataType) -> Result<ArrayRef> {
// single integer operation, but need to avoid integer
// math rounding down to zero
- if to_size > from_size {
- let time_array = Date64Array::from(array.data());
- Ok(Arc::new(multiply(
- &time_array,
- &Date64Array::from(vec![to_size / from_size; array.len()]),
- )?) as ArrayRef)
- } else if to_size < from_size {
- let time_array = Date64Array::from(array.data());
- Ok(Arc::new(divide(
- &time_array,
- &Date64Array::from(vec![from_size / to_size; array.len()]),
- )?) as ArrayRef)
- } else {
- cast_array_data::<Date64Type>(array, to_type.clone())
+ match to_size.cmp(&from_size) {
+ std::cmp::Ordering::Less => {
+ let time_array = Date64Array::from(array.data());
+ Ok(Arc::new(divide(
+ &time_array,
+ &Date64Array::from(vec![from_size / to_size; array.len()]),
+ )?) as ArrayRef)
+ }
+ std::cmp::Ordering::Equal => {
+ cast_array_data::<Date64Type>(array, to_type.clone())
+ }
+ std::cmp::Ordering::Greater => {
+ let time_array = Date64Array::from(array.data());
+ Ok(Arc::new(multiply(
+ &time_array,
+ &Date64Array::from(vec![to_size / from_size; array.len()]),
+ )?) as ArrayRef)
+ }
}
}
// date64 to timestamp might not make sense,
- // end temporal casts
+ // null to primitive/flat types
+ (Null, Int32) => Ok(Arc::new(Int32Array::from(vec![None; array.len()]))),
+
(_, _) => Err(ArrowError::ComputeError(format!(
"Casting from {:?} to {:?} not supported",
from_type, to_type,
@@ -2476,44 +2481,44 @@ mod tests {
// Test casting TO StringArray
let cast_type = Utf8;
- let cast_array = cast(&array, &cast_type).expect("cast to UTF-8 succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast to UTF-8 failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
// Test casting TO Dictionary (with different index sizes)
let cast_type = Dictionary(Box::new(Int16), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
let cast_type = Dictionary(Box::new(Int32), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
let cast_type = Dictionary(Box::new(Int64), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
let cast_type = Dictionary(Box::new(UInt8), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
let cast_type = Dictionary(Box::new(UInt16), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
let cast_type = Dictionary(Box::new(UInt32), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
let cast_type = Dictionary(Box::new(UInt64), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
}
@@ -2598,11 +2603,11 @@ mod tests {
let expected = vec!["1", "null", "3"];
// Test casting TO PrimitiveArray, different dictionary type
- let cast_array = cast(&array, &Utf8).expect("cast to UTF-8 succeeded");
+ let cast_array = cast(&array, &Utf8).expect("cast to UTF-8 failed");
assert_eq!(array_to_strings(&cast_array), expected);
assert_eq!(cast_array.data_type(), &Utf8);
- let cast_array = cast(&array, &Int64).expect("cast to int64 succeeded");
+ let cast_array = cast(&array, &Int64).expect("cast to int64 failed");
assert_eq!(array_to_strings(&cast_array), expected);
assert_eq!(cast_array.data_type(), &Int64);
}
@@ -2621,13 +2626,13 @@ mod tests {
// Cast to a dictionary (same value type, Int32)
let cast_type = Dictionary(Box::new(UInt8), Box::new(Int32));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
// Cast to a dictionary (different value type, Int8)
let cast_type = Dictionary(Box::new(UInt8), Box::new(Int8));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
}
@@ -2646,11 +2651,25 @@ mod tests {
// Cast to a dictionary (same value type, Utf8)
let cast_type = Dictionary(Box::new(UInt8), Box::new(Utf8));
- let cast_array = cast(&array, &cast_type).expect("cast succeeded");
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
assert_eq!(cast_array.data_type(), &cast_type);
assert_eq!(array_to_strings(&cast_array), expected);
}
+ #[test]
+ fn test_cast_null_array_to_int32() {
+ let array = Arc::new(NullArray::new(6)) as ArrayRef;
+
+ let expected = Int32Array::from(vec![None; 6]);
+
+ // Cast to a dictionary (same value type, Utf8)
+ let cast_type = DataType::Int32;
+ let cast_array = cast(&array, &cast_type).expect("cast failed");
+ let cast_array = as_primitive_array::<Int32Type>(&cast_array);
+ assert_eq!(cast_array.data_type(), &cast_type);
+ assert_eq!(cast_array, &expected);
+ }
+
/// Print the `DictionaryArray` `array` as a vector of strings
fn array_to_strings(array: &ArrayRef) -> Vec<String> {
(0..array.len())
@@ -2768,7 +2787,7 @@ mod tests {
)),
Arc::new(TimestampNanosecondArray::from_vec(
vec![1000, 2000],
- Some(tz_name.clone()),
+ Some(tz_name),
)),
Arc::new(Date32Array::from(vec![1000, 2000])),
Arc::new(Date64Array::from(vec![1000, 2000])),
diff --git a/rust/parquet/src/arrow/array_reader.rs b/rust/parquet/src/arrow/array_reader.rs
index 40df284..6fdf5d5 100644
--- a/rust/parquet/src/arrow/array_reader.rs
+++ b/rust/parquet/src/arrow/array_reader.rs
@@ -82,6 +82,97 @@ pub trait ArrayReader {
fn get_rep_levels(&self) -> Option<&[i16]>;
}
+/// A NullArrayReader reads Parquet columns stored as null int32s with an Arrow
+/// NullArray type.
+pub struct NullArrayReader<T: DataType> {
+ data_type: ArrowType,
+ pages: Box<dyn PageIterator>,
+ def_levels_buffer: Option<Buffer>,
+ rep_levels_buffer: Option<Buffer>,
+ column_desc: ColumnDescPtr,
+ record_reader: RecordReader<T>,
+ _type_marker: PhantomData<T>,
+}
+
+impl<T: DataType> NullArrayReader<T> {
+ /// Construct null array reader.
+ pub fn new(
+ mut pages: Box<dyn PageIterator>,
+ column_desc: ColumnDescPtr,
+ ) -> Result<Self> {
+ let mut record_reader = RecordReader::<T>::new(column_desc.clone());
+ if let Some(page_reader) = pages.next() {
+ record_reader.set_page_reader(page_reader?)?;
+ }
+
+ Ok(Self {
+ data_type: ArrowType::Null,
+ pages,
+ def_levels_buffer: None,
+ rep_levels_buffer: None,
+ column_desc,
+ record_reader,
+ _type_marker: PhantomData,
+ })
+ }
+}
+
+/// Implementation of primitive array reader.
+impl<T: DataType> ArrayReader for NullArrayReader<T> {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ /// Returns data type of primitive array.
+ fn get_data_type(&self) -> &ArrowType {
+ &self.data_type
+ }
+
+ /// Reads at most `batch_size` records into array.
+ fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
+ let mut records_read = 0usize;
+ while records_read < batch_size {
+ let records_to_read = batch_size - records_read;
+
+ // NB can be 0 if at end of page
+ let records_read_once = self.record_reader.read_records(records_to_read)?;
+ records_read += records_read_once;
+
+ // Record reader exhausted
+ if records_read_once < records_to_read {
+ if let Some(page_reader) = self.pages.next() {
+ // Read from new page reader
+ self.record_reader.set_page_reader(page_reader?)?;
+ } else {
+ // Page reader also exhausted
+ break;
+ }
+ }
+ }
+
+ // convert to arrays
+ let array = arrow::array::NullArray::new(records_read);
+
+ // save definition and repetition buffers
+ self.def_levels_buffer = self.record_reader.consume_def_levels()?;
+ self.rep_levels_buffer = self.record_reader.consume_rep_levels()?;
+ self.record_reader.reset();
+ Ok(Arc::new(array))
+ }
+
+ fn get_def_levels(&self) -> Option<&[i16]> {
+ self.def_levels_buffer
+ .as_ref()
+ .map(|buf| unsafe { buf.typed_data() })
+ }
+
+ fn get_rep_levels(&self) -> Option<&[i16]> {
+ self.rep_levels_buffer
+ .as_ref()
+ .map(|buf| unsafe { buf.typed_data() })
+ }
+}
+
/// Primitive array readers are leaves of array reader tree. They accept page iterator
/// and read them into primitive arrays.
pub struct PrimitiveArrayReader<T: DataType> {
@@ -859,10 +950,19 @@ impl<'a> ArrayReaderBuilder {
page_iterator,
column_desc,
)?)),
- PhysicalType::INT32 => Ok(Box::new(PrimitiveArrayReader::<Int32Type>::new(
- page_iterator,
- column_desc,
- )?)),
+ PhysicalType::INT32 => {
+ if let Some(ArrowType::Null) = arrow_type {
+ Ok(Box::new(NullArrayReader::<Int32Type>::new(
+ page_iterator,
+ column_desc,
+ )?))
+ } else {
+ Ok(Box::new(PrimitiveArrayReader::<Int32Type>::new(
+ page_iterator,
+ column_desc,
+ )?))
+ }
+ }
PhysicalType::INT64 => Ok(Box::new(PrimitiveArrayReader::<Int64Type>::new(
page_iterator,
column_desc,
@@ -903,25 +1003,23 @@ impl<'a> ArrayReaderBuilder {
page_iterator, column_desc, converter
)?))
}
+ } else if let Some(ArrowType::LargeBinary) = arrow_type {
+ let converter =
+ LargeBinaryConverter::new(LargeBinaryArrayConverter {});
+ Ok(Box::new(ComplexObjectArrayReader::<
+ ByteArrayType,
+ LargeBinaryConverter,
+ >::new(
+ page_iterator, column_desc, converter
+ )?))
} else {
- if let Some(ArrowType::LargeBinary) = arrow_type {
- let converter =
- LargeBinaryConverter::new(LargeBinaryArrayConverter {});
- Ok(Box::new(ComplexObjectArrayReader::<
- ByteArrayType,
- LargeBinaryConverter,
- >::new(
- page_iterator, column_desc, converter
- )?))
- } else {
- let converter = BinaryConverter::new(BinaryArrayConverter {});
- Ok(Box::new(ComplexObjectArrayReader::<
- ByteArrayType,
- BinaryConverter,
- >::new(
- page_iterator, column_desc, converter
- )?))
- }
+ let converter = BinaryConverter::new(BinaryArrayConverter {});
+ Ok(Box::new(ComplexObjectArrayReader::<
+ ByteArrayType,
+ BinaryConverter,
+ >::new(
+ page_iterator, column_desc, converter
+ )?))
}
}
PhysicalType::FIXED_LEN_BYTE_ARRAY => {
diff --git a/rust/parquet/src/arrow/arrow_writer.rs b/rust/parquet/src/arrow/arrow_writer.rs
index a17e424..ff535dc 100644
--- a/rust/parquet/src/arrow/arrow_writer.rs
+++ b/rust/parquet/src/arrow/arrow_writer.rs
@@ -128,7 +128,8 @@ fn write_leaves(
mut levels: &mut Vec<Levels>,
) -> Result<()> {
match array.data_type() {
- ArrowDataType::Int8
+ ArrowDataType::Null
+ | ArrowDataType::Int8
| ArrowDataType::Int16
| ArrowDataType::Int32
| ArrowDataType::Int64
@@ -179,7 +180,6 @@ fn write_leaves(
"Float16 arrays not supported".to_string(),
)),
ArrowDataType::FixedSizeList(_, _)
- | ArrowDataType::Null
| ArrowDataType::Boolean
| ArrowDataType::FixedSizeBinary(_)
| ArrowDataType::Union(_)
@@ -279,7 +279,10 @@ fn get_levels(
parent_rep_levels: Option<&[i16]>,
) -> Vec<Levels> {
match array.data_type() {
- ArrowDataType::Null => unimplemented!(),
+ ArrowDataType::Null => vec![Levels {
+ definition: parent_def_levels.iter().map(|v| (v - 1).max(0)).collect(),
+ repetition: None,
+ }],
ArrowDataType::Boolean
| ArrowDataType::Int8
| ArrowDataType::Int16
@@ -356,7 +359,11 @@ fn get_levels(
// if datatype is a primitive, we can construct levels of the child array
match child_array.data_type() {
- ArrowDataType::Null => unimplemented!(),
+ // TODO: The behaviour of a <list<null>> is untested
+ ArrowDataType::Null => vec![Levels {
+ definition: list_def_levels,
+ repetition: Some(list_rep_levels),
+ }],
ArrowDataType::Boolean => unimplemented!(),
ArrowDataType::Int8
| ArrowDataType::Int16
@@ -701,7 +708,7 @@ mod tests {
expected_batch.schema(),
None,
)
- .unwrap();
+ .expect("Unable to write file");
writer.write(&expected_batch).unwrap();
writer.close().unwrap();
@@ -709,7 +716,10 @@ mod tests {
let mut arrow_reader = ParquetFileArrowReader::new(Rc::new(reader));
let mut record_batch_reader = arrow_reader.get_record_reader(1024).unwrap();
- let actual_batch = record_batch_reader.next().unwrap().unwrap();
+ let actual_batch = record_batch_reader
+ .next()
+ .expect("No batch found")
+ .expect("Unable to get batch");
assert_eq!(expected_batch.schema(), actual_batch.schema());
assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
@@ -778,11 +788,15 @@ mod tests {
}
#[test]
- #[should_panic(expected = "Null arrays not supported")]
+ fn all_null_primitive_single_column() {
+ let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
+ one_column_roundtrip("all_null_primitive_single_column", values, true);
+ }
+ #[test]
fn null_single_column() {
let values = Arc::new(NullArray::new(SMALL_SIZE));
- one_column_roundtrip("null_single_column", values.clone(), true);
- one_column_roundtrip("null_single_column", values, false);
+ one_column_roundtrip("null_single_column", values, true);
+ // null arrays are always nullable, a test with non-nullable nulls fails
}
#[test]
diff --git a/rust/parquet/src/arrow/schema.rs b/rust/parquet/src/arrow/schema.rs
index 0cd41fe..10270ff 100644
--- a/rust/parquet/src/arrow/schema.rs
+++ b/rust/parquet/src/arrow/schema.rs
@@ -308,7 +308,10 @@ fn arrow_to_parquet_type(field: &Field) -> Result<Type> {
};
// create type from field
match field.data_type() {
- DataType::Null => Err(ArrowError("Null arrays not supported".to_string())),
+ DataType::Null => Type::primitive_type_builder(name, PhysicalType::INT32)
+ .with_logical_type(LogicalType::NONE)
+ .with_repetition(repetition)
+ .build(),
DataType::Boolean => Type::primitive_type_builder(name, PhysicalType::BOOLEAN)
.with_repetition(repetition)
.build(),
@@ -1501,6 +1504,7 @@ mod tests {
// )))),
// true,
// ),
+ Field::new("c35", DataType::Null, true),
],
metadata,
);