add python wrapper

Author: rok

python/pyarrow/_dataset_parquet.pyx

@@ -620,6 +620,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions):
             "coerce_timestamps",
             "allow_truncated_timestamps",
             "use_compliant_nested_type",
+            "experimental_vector_encoding",
         }
 
         setters = set()
@@ -676,6 +677,9 @@ cdef class ParquetFileWriteOptions(FileWriteOptions):
             writer_engine_version="V2",
             use_compliant_nested_type=(
                 self._properties["use_compliant_nested_type"]
+            ),
+            experimental_vector_encoding=(
+                self._properties["experimental_vector_encoding"]
             )
         )
 
@@ -705,6 +709,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions):
             coerce_timestamps=None,
             allow_truncated_timestamps=False,
             use_compliant_nested_type=True,
+            experimental_vector_encoding=False,
             encryption_properties=None,
             write_batch_size=None,
             dictionary_pagesize_limit=None,

python/pyarrow/_parquet.pyx

@@ -2281,7 +2281,8 @@ cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties(
         writer_engine_version=None,
         use_compliant_nested_type=True,
         store_schema=True,
-        write_time_adjusted_to_utc=False) except *:
+        write_time_adjusted_to_utc=False,
+        experimental_vector_encoding=False) except *:
     """Arrow writer properties"""
     cdef:
         shared_ptr[ArrowWriterProperties] arrow_properties
@@ -2322,6 +2323,11 @@ cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties(
     else:
         arrow_props.disable_compliant_nested_types()
 
+    # experimental_vector_encoding
+
+    if experimental_vector_encoding:
+        arrow_props.enable_experimental_vector_encoding()
+
     # writer_engine_version
 
     if writer_engine_version == "V1":
@@ -2396,7 +2402,8 @@ cdef class ParquetWriter(_Weakrefable):
                   store_decimal_as_integer=False,
                   use_content_defined_chunking=False,
                   write_time_adjusted_to_utc=False,
-                  bloom_filter_options=None):
+                  bloom_filter_options=None,
+                  experimental_vector_encoding=False):
         cdef:
             shared_ptr[WriterProperties] properties
             shared_ptr[ArrowWriterProperties] arrow_properties
@@ -2443,6 +2450,7 @@ cdef class ParquetWriter(_Weakrefable):
             use_compliant_nested_type=use_compliant_nested_type,
             store_schema=store_schema,
             write_time_adjusted_to_utc=write_time_adjusted_to_utc,
+            experimental_vector_encoding=experimental_vector_encoding,
         )
 
         pool = maybe_unbox_memory_pool(memory_pool)

python/pyarrow/includes/libparquet.pxd

@@ -528,6 +528,7 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil:
             Builder* store_schema()
             Builder* enable_compliant_nested_types()
             Builder* disable_compliant_nested_types()
+            Builder* enable_experimental_vector_encoding()
             Builder* set_engine_version(ArrowWriterEngineVersion version)
             Builder* set_time_adjusted_to_utc(c_bool adjusted)
             shared_ptr[ArrowWriterProperties] build()

python/pyarrow/parquet/core.py

@@ -876,6 +876,10 @@ def _sanitize_table(table, new_schema, flavor):
 dictionary_pagesize_limit : int, default None
     Specify the dictionary page size limit per row group. If None, use the
     default 1MB.
+experimental_vector_encoding : bool, default False
+    EXPERIMENTAL: Encode supported fixed-size list values, including the storage
+    of fixed-shape tensor extension arrays, using Parquet VECTOR repetition.
+    Unsupported fixed-size list values continue using standard LIST encoding.
 store_schema : bool, default True
     By default, the Arrow schema is serialized and stored in the Parquet
     file metadata (in the "ARROW:schema" key). When reading the file,
@@ -1083,6 +1087,7 @@ def __init__(self, where, schema, filesystem=None,
                  store_decimal_as_integer=False,
                  write_time_adjusted_to_utc=False,
                  max_rows_per_page=None,
+                 experimental_vector_encoding=False,
                  **options):
         if use_deprecated_int96_timestamps is None:
             # Use int96 timestamps for Spark
@@ -1138,6 +1143,7 @@ def __init__(self, where, schema, filesystem=None,
             store_decimal_as_integer=store_decimal_as_integer,
             write_time_adjusted_to_utc=write_time_adjusted_to_utc,
             max_rows_per_page=max_rows_per_page,
+            experimental_vector_encoding=experimental_vector_encoding,
             **options)
         self.is_open = True
 
@@ -2017,6 +2023,7 @@ def write_table(table, where, row_group_size=None, version='2.6',
                 write_time_adjusted_to_utc=False,
                 max_rows_per_page=None,
                 bloom_filter_options=None,
+                experimental_vector_encoding=False,
                 **kwargs):
     # Implementor's note: when adding keywords here / updating defaults, also
     # update it in write_to_dataset and _dataset_parquet.pyx ParquetFileWriteOptions
@@ -2051,6 +2058,7 @@ def write_table(table, where, row_group_size=None, version='2.6',
                 write_time_adjusted_to_utc=write_time_adjusted_to_utc,
                 max_rows_per_page=max_rows_per_page,
                 bloom_filter_options=bloom_filter_options,
+                experimental_vector_encoding=experimental_vector_encoding,
                 **kwargs) as writer:
             writer.write_table(table, row_group_size=row_group_size)
     except Exception:

python/pyarrow/tests/parquet/test_data_types.py

@@ -581,6 +581,28 @@ def test_json_extension_type(storage_type):
         store_schema=False)
 
 
+def test_fixed_shape_tensor_vector_encoding():
+    tensor_type = pa.fixed_shape_tensor(pa.float32(), [2, 3])
+    storage = pa.array([
+        [1, 2, 3, 4, 5, 6],
+        None,
+        [7, 8, 9, 10, 11, 12],
+    ], pa.list_(pa.float32(), 6))
+    arr = pa.ExtensionArray.from_storage(tensor_type, storage)
+    table = pa.table({"tensor": arr})
+
+    sink = pa.BufferOutputStream()
+    pq.write_table(table, sink, experimental_vector_encoding=True)
+    buf = sink.getvalue()
+
+    metadata = pq.read_metadata(pa.BufferReader(buf))
+    assert "vector" in str(metadata.schema).lower()
+
+    result = pq.read_table(pa.BufferReader(buf))
+    assert result.schema == table.schema
+    assert result.equals(table)
+
+
 def test_uuid_extension_type():
     data = [
         b'\xe4`\xf9p\x83QGN\xac\x7f\xa4g>K\xa8\xcb',