[python] Support zstd decompress for btree index

Author: JingsongLi

paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java

@@ -81,6 +81,7 @@
 import static org.apache.paimon.CoreOptions.ROW_TRACKING_ENABLED;
 import static org.apache.paimon.CoreOptions.TARGET_FILE_SIZE;
 import static org.apache.paimon.data.DataFormatTestUtil.internalRowToString;
+import static org.apache.paimon.globalindex.btree.BTreeIndexOptions.BTREE_INDEX_COMPRESSION;
 import static org.apache.paimon.table.SimpleTableTestBase.getResult;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -414,6 +415,7 @@ public void testBtreeIndexWrite() throws Exception {
         testBtreeIndexWriteString();
         testBtreeIndexWriteInt();
         testBtreeIndexWriteBigInt();
+        testBtreeIndexWriteLarge();
     }
 
     private void testBtreeIndexWriteString() throws Exception {
@@ -498,6 +500,75 @@ private <T> void testBtreeIndexWriteGeneric(
         assertThat(result).containsOnly("v2");
     }
 
+    private void testBtreeIndexWriteLarge() throws Exception {
+        // create table
+        RowType rowType =
+                RowType.of(
+                        new DataType[] {DataTypes.STRING(), DataTypes.STRING()},
+                        new String[] {"k", "v"});
+        Options options = new Options();
+        Path tablePath = new Path(warehouse.toString() + "/default.db/test_btree_index_large");
+        options.set(PATH, tablePath.toString());
+        options.set(ROW_TRACKING_ENABLED, true);
+        options.set(DATA_EVOLUTION_ENABLED, true);
+        options.set(GLOBAL_INDEX_ENABLED, true);
+        options.set(BTREE_INDEX_COMPRESSION, "zstd");
+        TableSchema tableSchema =
+                SchemaUtils.forceCommit(
+                        new SchemaManager(LocalFileIO.create(), tablePath),
+                        new Schema(
+                                rowType.getFields(),
+                                Collections.emptyList(),
+                                Collections.emptyList(),
+                                options.toMap(),
+                                ""));
+        AppendOnlyFileStoreTable table =
+                new AppendOnlyFileStoreTable(
+                        FileIOFinder.find(tablePath),
+                        tablePath,
+                        tableSchema,
+                        CatalogEnvironment.empty());
+
+        // write data
+        BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder();
+        try (BatchTableWrite write = writeBuilder.newWrite();
+                BatchTableCommit commit = writeBuilder.newCommit()) {
+            for (int i = 0; i < 2000; i++) {
+                write.write(
+                        GenericRow.of(
+                                BinaryString.fromString("k" + i),
+                                BinaryString.fromString("v" + i)));
+            }
+            commit.commit(write.prepareCommit());
+        }
+
+        // build index
+        BTreeGlobalIndexBuilder builder =
+                new BTreeGlobalIndexBuilder(table).withIndexType("btree").withIndexField("k");
+        try (BatchTableCommit commit = writeBuilder.newCommit()) {
+            commit.commit(builder.build(builder.scan(), IOManager.create(warehouse.toString())));
+        }
+
+        // assert index
+        List<IndexManifestEntry> indexEntries =
+                table.indexManifestFileReader().read(table.latestSnapshot().get().indexManifest);
+        assertThat(indexEntries)
+                .singleElement()
+                .matches(entry -> entry.indexFile().rowCount() == 2000);
+
+        // read index
+        PredicateBuilder predicateBuilder = new PredicateBuilder(table.rowType());
+        ReadBuilder readBuilder =
+                table.newReadBuilder()
+                        .withFilter(predicateBuilder.equal(0, BinaryString.fromString("k2")));
+        List<String> result = new ArrayList<>();
+        readBuilder
+                .newRead()
+                .createReader(readBuilder.newScan().plan())
+                .forEachRemaining(r -> result.add(r.getString(1).toString()));
+        assertThat(result).containsOnly("v2");
+    }
+
     // Helper method from TableTestBase
     protected Identifier identifier(String tableName) {
         return new Identifier(database, tableName);

paimon-python/pypaimon/globalindex/btree/sst_file_reader.py

@@ -30,6 +30,18 @@
 from typing import Optional, Callable
 from typing import BinaryIO
 
+try:
+    import lz4.block
+    LZ4_AVAILABLE = True
+except ImportError:
+    LZ4_AVAILABLE = False
+
+try:
+    import python_lzo
+    LZO_AVAILABLE = True
+except ImportError:
+    LZO_AVAILABLE = False
+
 from pypaimon.globalindex.btree.btree_file_footer import BlockHandle
 from pypaimon.globalindex.btree.block_entry import BlockEntry
 from pypaimon.globalindex.btree.block_reader import BlockReader, BlockIterator
@@ -136,9 +148,6 @@ def _read_block(self, block_handle: BlockHandle) -> BlockReader:
 
         trailer_offset = len(block_data) - 5
         compression_type = block_data[trailer_offset]
-        if compression_type != 0:
-            raise ValueError("Compression type not supported")
-
         crc32_value = struct.unpack('<I', block_data[trailer_offset + 1:trailer_offset + 5])[0]
 
         # Extract block data (without trailer)
@@ -149,6 +158,21 @@ def _read_block(self, block_handle: BlockHandle) -> BlockReader:
         if actual_crc32 != crc32_value:
             raise ValueError(f"CRC32 mismatch: expected {crc32_value}, got {actual_crc32}")
 
+        # Decompress if needed
+        if compression_type == 1:  # ZSTD
+            import zstandard as zstd
+            from io import BytesIO
+            decompressor = zstd.ZstdDecompressor()
+            memory_input = MemorySliceInput(block_bytes)
+            expected_len = memory_input.read_var_len_int()
+            compressed = block_bytes[memory_input.position():]
+            with decompressor.stream_reader(BytesIO(compressed)) as reader:
+                block_bytes = reader.read()
+            if len(block_bytes) != expected_len:
+                raise ValueError("Corrupted data, decompression failed.")
+        elif compression_type != 0:  # Not NONE
+            raise ValueError(f"Compression type {compression_type} not supported")
+
         return BlockReader.create(block_bytes, self.comparator)
 
     def create_iterator(self) -> SstFileIterator:

paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py

@@ -324,6 +324,7 @@ def test_read_btree_index_table(self):
         self._test_read_btree_index_generic("test_btree_index_string", "k2", pa.string())
         self._test_read_btree_index_generic("test_btree_index_int", 200, pa.int32())
         self._test_read_btree_index_generic("test_btree_index_bigint", 2000, pa.int64())
+        self._test_read_btree_index_large()
 
     def _test_read_btree_index_generic(self, table_name: str, k, k_type):
         table = self.catalog.get_table('default.' + table_name)
@@ -344,3 +345,19 @@ def _test_read_btree_index_generic(self, table_name: str, k, k_type):
             ("v", pa.string())
         ]))
         self.assertEqual(expected, actual)
+
+    def _test_read_btree_index_large(self):
+        table = self.catalog.get_table('default.test_btree_index_large')
+        read_builder: ReadBuilder = table.new_read_builder()
+        predicate_builder = read_builder.new_predicate_builder()
+
+        # read equal index
+        read_builder.with_filter(predicate_builder.equal('k', 'k2'))
+        table_read = read_builder.new_read()
+        splits = read_builder.new_scan().plan().splits()
+        actual = table_sort_by(table_read.to_arrow(splits), 'k')
+        expected = pa.Table.from_pydict({
+            'k': ["k2"],
+            'v': ["v2"]
+        })
+        self.assertEqual(expected, actual)