This is an automated email from the ASF dual-hosted git repository.
lzljs3620320 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-python.git
The following commit(s) were added to refs/heads/main by this push:
new 5f18615 TableRead supports reading multiple splits parallely (#11)
5f18615 is described below
commit 5f18615fd6730bb0f99a1c5dbeee743fd95f730d
Author: yuzelin <[email protected]>
AuthorDate: Wed Sep 11 12:01:29 2024 +0800
TableRead supports reading multiple splits parallely (#11)
---
java_based_implementation/api_impl.py | 56 ++++---
.../org/apache/paimon/python/InvocationUtil.java | 5 +-
.../apache/paimon/python/ParallelBytesReader.java | 171 +++++++++++++++++++++
.../{BytesReader.java => RecordBytesIterator.java} | 48 ++----
.../tests/test_write_and_read.py | 58 ++++++-
java_based_implementation/util/constants.py | 3 +
java_based_implementation/util/java_utils.py | 4 +-
paimon_python_api/catalog.py | 2 +-
paimon_python_api/table_read.py | 3 +-
9 files changed, 286 insertions(+), 64 deletions(-)
diff --git a/java_based_implementation/api_impl.py
b/java_based_implementation/api_impl.py
index 2605cc3..05487f4 100644
--- a/java_based_implementation/api_impl.py
+++ b/java_based_implementation/api_impl.py
@@ -17,6 +17,7 @@
################################################################################
from java_based_implementation.java_gateway import get_gateway
+from java_based_implementation.util import constants
from java_based_implementation.util.java_utils import to_j_catalog_context,
check_batch_write
from paimon_python_api import (catalog, table, read_builder, table_scan,
split, table_read,
write_builder, table_write, commit_message,
table_commit)
@@ -27,31 +28,33 @@ from typing import List, Iterator
class Catalog(catalog.Catalog):
- def __init__(self, j_catalog):
+ def __init__(self, j_catalog, catalog_options: dict):
self._j_catalog = j_catalog
+ self._catalog_options = catalog_options
@staticmethod
- def create(catalog_context: dict) -> 'Catalog':
- j_catalog_context = to_j_catalog_context(catalog_context)
+ def create(catalog_options: dict) -> 'Catalog':
+ j_catalog_context = to_j_catalog_context(catalog_options)
gateway = get_gateway()
j_catalog = gateway.jvm.CatalogFactory.createCatalog(j_catalog_context)
- return Catalog(j_catalog)
+ return Catalog(j_catalog, catalog_options)
def get_table(self, identifier: str) -> 'Table':
gateway = get_gateway()
j_identifier = gateway.jvm.Identifier.fromString(identifier)
j_table = self._j_catalog.getTable(j_identifier)
- return Table(j_table)
+ return Table(j_table, self._catalog_options)
class Table(table.Table):
- def __init__(self, j_table):
+ def __init__(self, j_table, catalog_options: dict):
self._j_table = j_table
+ self._catalog_options = catalog_options
def new_read_builder(self) -> 'ReadBuilder':
j_read_builder =
get_gateway().jvm.InvocationUtil.getReadBuilder(self._j_table)
- return ReadBuilder(j_read_builder, self._j_table.rowType())
+ return ReadBuilder(j_read_builder, self._j_table.rowType(),
self._catalog_options)
def new_batch_write_builder(self) -> 'BatchWriteBuilder':
check_batch_write(self._j_table)
@@ -61,9 +64,10 @@ class Table(table.Table):
class ReadBuilder(read_builder.ReadBuilder):
- def __init__(self, j_read_builder, j_row_type):
+ def __init__(self, j_read_builder, j_row_type, catalog_options: dict):
self._j_read_builder = j_read_builder
self._j_row_type = j_row_type
+ self._catalog_options = catalog_options
def with_projection(self, projection: List[List[int]]) -> 'ReadBuilder':
self._j_read_builder.withProjection(projection)
@@ -79,7 +83,7 @@ class ReadBuilder(read_builder.ReadBuilder):
def new_read(self) -> 'TableRead':
j_table_read = self._j_read_builder.newRead()
- return TableRead(j_table_read, self._j_row_type)
+ return TableRead(j_table_read, self._j_row_type, self._catalog_options)
class TableScan(table_scan.TableScan):
@@ -113,24 +117,40 @@ class Split(split.Split):
class TableRead(table_read.TableRead):
- def __init__(self, j_table_read, j_row_type):
+ def __init__(self, j_table_read, j_row_type, catalog_options):
self._j_table_read = j_table_read
- self._j_bytes_reader =
get_gateway().jvm.InvocationUtil.createBytesReader(
- j_table_read, j_row_type)
+ self._j_row_type = j_row_type
+ self._catalog_options = catalog_options
+ self._j_bytes_reader = None
self._arrow_schema = None
- def create_reader(self, split: Split):
- self._j_bytes_reader.setSplit(split.to_j_split())
- # get schema
+ def create_reader(self, splits):
+ self._init()
+ j_splits = list(map(lambda s: s.to_j_split(), splits))
+ self._j_bytes_reader.setSplits(j_splits)
+ batch_iterator = self._batch_generator()
+ return RecordBatchReader.from_batches(self._arrow_schema,
batch_iterator)
+
+ def _init(self):
+ if self._j_bytes_reader is None:
+ # get thread num
+ max_workers = self._catalog_options.get(constants.MAX_WORKERS)
+ if max_workers is None:
+ # default is sequential
+ max_workers = 1
+ else:
+ max_workers = int(max_workers)
+ if max_workers <= 0:
+ raise ValueError("max_workers must be greater than 0")
+ self._j_bytes_reader =
get_gateway().jvm.InvocationUtil.createParallelBytesReader(
+ self._j_table_read, self._j_row_type, max_workers)
+
if self._arrow_schema is None:
schema_bytes = self._j_bytes_reader.serializeSchema()
schema_reader = RecordBatchStreamReader(BufferReader(schema_bytes))
self._arrow_schema = schema_reader.schema
schema_reader.close()
- batch_iterator = self._batch_generator()
- return RecordBatchReader.from_batches(self._arrow_schema,
batch_iterator)
-
def _batch_generator(self) -> Iterator[RecordBatch]:
while True:
next_bytes = self._j_bytes_reader.next()
diff --git
a/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/InvocationUtil.java
b/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/InvocationUtil.java
index f29773a..61d322a 100644
---
a/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/InvocationUtil.java
+++
b/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/InvocationUtil.java
@@ -39,8 +39,9 @@ public class InvocationUtil {
return table.newReadBuilder();
}
- public static BytesReader createBytesReader(TableRead tableRead, RowType
rowType) {
- return new BytesReader(tableRead, rowType);
+ public static ParallelBytesReader createParallelBytesReader(
+ TableRead tableRead, RowType rowType, int threadNum) {
+ return new ParallelBytesReader(tableRead, rowType, threadNum);
}
public static BytesWriter createBytesWriter(TableWrite tableWrite, RowType
rowType) {
diff --git
a/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/ParallelBytesReader.java
b/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/ParallelBytesReader.java
new file mode 100644
index 0000000..9f1b390
--- /dev/null
+++
b/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/ParallelBytesReader.java
@@ -0,0 +1,171 @@
+/*
+ * 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.
+ */
+
+package org.apache.paimon.python;
+
+import org.apache.paimon.arrow.ArrowUtils;
+import org.apache.paimon.arrow.vector.ArrowFormatWriter;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.reader.RecordReader;
+import org.apache.paimon.reader.RecordReaderIterator;
+import org.apache.paimon.table.source.Split;
+import org.apache.paimon.table.source.TableRead;
+import org.apache.paimon.types.RowType;
+import org.apache.paimon.utils.ThreadPoolUtils;
+
+import org.apache.paimon.shade.guava30.com.google.common.collect.Iterators;
+
+import org.apache.arrow.vector.VectorSchemaRoot;
+
+import javax.annotation.Nullable;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.function.Function;
+
+/** Parallely read Arrow bytes from multiple splits. */
+public class ParallelBytesReader {
+
+ private static final String THREAD_NAME_PREFIX = "PARALLEL_SPLITS_READER";
+ private static final int DEFAULT_WRITE_BATCH_SIZE = 1024;
+
+ private final TableRead tableRead;
+ private final RowType rowType;
+ private final int threadNum;
+
+ private final ConcurrentLinkedQueue<RecordReaderIterator<InternalRow>>
iterators;
+ private final ConcurrentLinkedQueue<ArrowFormatWriter> arrowFormatWriters;
+
+ private ThreadPoolExecutor executor;
+ private Iterator<byte[]> bytesIterator;
+
+ public ParallelBytesReader(TableRead tableRead, RowType rowType, int
threadNum) {
+ this.tableRead = tableRead;
+ this.rowType = rowType;
+ this.threadNum = threadNum;
+ this.iterators = new ConcurrentLinkedQueue<>();
+ this.arrowFormatWriters = new ConcurrentLinkedQueue<>();
+ }
+
+ public void setSplits(List<Split> splits) {
+ bytesIterator = randomlyExecute(getExecutor(), makeProcessor(),
splits);
+ }
+
+ public byte[] serializeSchema() {
+ ArrowFormatWriter arrowFormatWriter = newWriter();
+ VectorSchemaRoot vsr = arrowFormatWriter.getVectorSchemaRoot();
+ ByteArrayOutputStream out = new ByteArrayOutputStream();
+ ArrowUtils.serializeToIpc(vsr, out);
+ arrowFormatWriter.close();
+ return out.toByteArray();
+ }
+
+ @Nullable
+ public byte[] next() {
+ if (bytesIterator.hasNext()) {
+ return bytesIterator.next();
+ } else {
+ closeResources();
+ return null;
+ }
+ }
+
+ private ThreadPoolExecutor getExecutor() {
+ if (executor == null) {
+ executor = ThreadPoolUtils.createCachedThreadPool(threadNum,
THREAD_NAME_PREFIX);
+ }
+ return executor;
+ }
+
+ private Function<Split, Iterator<byte[]>> makeProcessor() {
+ return split -> {
+ try {
+ RecordReader<InternalRow> recordReader =
tableRead.createReader(split);
+ RecordReaderIterator<InternalRow> iterator =
+ new RecordReaderIterator<>(recordReader);
+ iterators.add(iterator);
+ ArrowFormatWriter arrowFormatWriter = newWriter();
+ arrowFormatWriters.add(arrowFormatWriter);
+ return new RecordBytesIterator(iterator, arrowFormatWriter);
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ };
+ }
+
+ private <U, T> Iterator<T> randomlyExecute(
+ ExecutorService executor, Function<U, Iterator<T>> processor,
Collection<U> input) {
+ List<Future<Iterator<T>>> futures = new ArrayList<>(input.size());
+ for (U u : input) {
+ futures.add(executor.submit(() -> processor.apply(u)));
+ }
+ return futuresToIterIter(futures);
+ }
+
+ private <T> Iterator<T> futuresToIterIter(List<Future<Iterator<T>>>
futures) {
+ final Queue<Future<Iterator<T>>> queue = new ArrayDeque<>(futures);
+ return Iterators.concat(
+ new Iterator<Iterator<T>>() {
+ public boolean hasNext() {
+ return !queue.isEmpty();
+ }
+
+ public Iterator<T> next() {
+ try {
+ return queue.poll().get();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ } catch (ExecutionException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ });
+ }
+
+ private void closeResources() {
+ for (RecordReaderIterator<InternalRow> iterator : iterators) {
+ try {
+ iterator.close();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+ iterators.clear();
+
+ for (ArrowFormatWriter arrowFormatWriter : arrowFormatWriters) {
+ arrowFormatWriter.close();
+ }
+ arrowFormatWriters.clear();
+ }
+
+ private ArrowFormatWriter newWriter() {
+ return new ArrowFormatWriter(rowType, DEFAULT_WRITE_BATCH_SIZE, true);
+ }
+}
diff --git
a/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/BytesReader.java
b/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/RecordBytesIterator.java
similarity index 57%
rename from
java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/BytesReader.java
rename to
java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/RecordBytesIterator.java
index 9272f98..d4965db 100644
---
a/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/BytesReader.java
+++
b/java_based_implementation/paimon-python-java-bridge/src/main/java/org/apache/paimon/python/RecordBytesIterator.java
@@ -21,54 +21,33 @@ package org.apache.paimon.python;
import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.arrow.vector.ArrowFormatWriter;
import org.apache.paimon.data.InternalRow;
-import org.apache.paimon.reader.RecordReader;
import org.apache.paimon.reader.RecordReaderIterator;
-import org.apache.paimon.table.source.Split;
-import org.apache.paimon.table.source.TableRead;
-import org.apache.paimon.types.RowType;
import org.apache.arrow.vector.VectorSchemaRoot;
-import javax.annotation.Nullable;
-
import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-
-/** Read Arrow bytes from split. */
-public class BytesReader {
+import java.util.Iterator;
- private static final int DEFAULT_WRITE_BATCH_SIZE = 2048;
+/** Cast a {@link RecordReaderIterator} to bytes iterator. */
+public class RecordBytesIterator implements Iterator<byte[]> {
- private final TableRead tableRead;
+ private final RecordReaderIterator<InternalRow> iterator;
private final ArrowFormatWriter arrowFormatWriter;
- private RecordReaderIterator<InternalRow> iterator;
private InternalRow nextRow;
- public BytesReader(TableRead tableRead, RowType rowType) {
- this.tableRead = tableRead;
- this.arrowFormatWriter = new ArrowFormatWriter(rowType,
DEFAULT_WRITE_BATCH_SIZE, true);
- }
-
- public void setSplit(Split split) throws IOException {
- RecordReader<InternalRow> recordReader = tableRead.createReader(split);
- iterator = new RecordReaderIterator<InternalRow>(recordReader);
+ public RecordBytesIterator(
+ RecordReaderIterator<InternalRow> iterator, ArrowFormatWriter
arrowFormatWriter) {
+ this.iterator = iterator;
+ this.arrowFormatWriter = arrowFormatWriter;
nextRow();
}
- public byte[] serializeSchema() {
- VectorSchemaRoot vsr = arrowFormatWriter.getVectorSchemaRoot();
- ByteArrayOutputStream out = new ByteArrayOutputStream();
- ArrowUtils.serializeToIpc(vsr, out);
- return out.toByteArray();
+ public boolean hasNext() {
+ return nextRow != null;
}
- @Nullable
- public byte[] next() throws Exception {
- if (nextRow == null) {
- return null;
- }
-
+ public byte[] next() {
int rowCount = 0;
while (nextRow != null && arrowFormatWriter.write(nextRow)) {
nextRow();
@@ -80,11 +59,6 @@ public class BytesReader {
vsr.setRowCount(rowCount);
ByteArrayOutputStream out = new ByteArrayOutputStream();
ArrowUtils.serializeToIpc(vsr, out);
- if (nextRow == null) {
- // close resource
- arrowFormatWriter.close();
- iterator.close();
- }
return out.toByteArray();
}
diff --git a/java_based_implementation/tests/test_write_and_read.py
b/java_based_implementation/tests/test_write_and_read.py
index 4511825..0ea834e 100644
--- a/java_based_implementation/tests/test_write_and_read.py
+++ b/java_based_implementation/tests/test_write_and_read.py
@@ -83,7 +83,7 @@ class TableWriteReadTest(unittest.TestCase):
table_commit.close()
# read data
- table = Table(j_table)
+ table = Table(j_table, {})
read_builder = table.new_read_builder()
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
@@ -92,7 +92,7 @@ class TableWriteReadTest(unittest.TestCase):
data_frames = [
batch.to_pandas()
for split in splits
- for batch in table_read.create_reader(split)
+ for batch in table_read.create_reader([split])
]
self.assertEqual(len(data_frames), 0)
@@ -132,7 +132,7 @@ class TableWriteReadTest(unittest.TestCase):
data_frames = [
batch.to_pandas()
for split in splits
- for batch in table_read.create_reader(split)
+ for batch in table_read.create_reader([split])
]
result = pd.concat(data_frames)
@@ -184,3 +184,55 @@ class TableWriteReadTest(unittest.TestCase):
self.assertEqual(
str(e.exception),
"Doesn't support writing dynamic bucket or cross partition table.")
+
+ def testParallelRead(self):
+ create_simple_table(self.warehouse, 'default', 'test_parallel_read',
False)
+
+ catalog = Catalog.create({'warehouse': self.warehouse, 'max-workers':
'2'})
+ table = catalog.get_table('default.test_parallel_read')
+
+ # prepare data
+ n_times = 4
+ expected_data = {
+ 'f0': [],
+ 'f1': []
+ }
+ for i in range(n_times):
+ data = {
+ 'f0': [i],
+ 'f1': [str(i * 2)]
+ }
+ expected_data['f0'].append(i)
+ expected_data['f1'].append(str(i * 2))
+
+ df = pd.DataFrame(data)
+ df['f0'] = df['f0'].astype('int32')
+ record_batch = pa.RecordBatch.from_pandas(df)
+
+ # write and commit data
+ write_builder = table.new_batch_write_builder()
+ table_write = write_builder.new_write()
+ table_commit = write_builder.new_commit()
+
+ table_write.write(record_batch)
+ commit_messages = table_write.prepare_commit()
+ table_commit.commit(commit_messages)
+
+ # read data parallely
+ read_builder = table.new_read_builder()
+ table_scan = read_builder.new_scan()
+ table_read = read_builder.new_read()
+ splits = table_scan.plan().splits()
+
+ data_frames = [
+ batch.to_pandas()
+ for batch in table_read.create_reader(splits)
+ ]
+ result = pd.concat(data_frames)
+
+ expected_df = pd.DataFrame(expected_data)
+ expected_df['f0'] = expected_df['f0'].astype('int32')
+
+ # check data (ignore index)
+ pd.testing.assert_frame_equal(
+ result.reset_index(drop=True), expected_df.reset_index(drop=True))
diff --git a/java_based_implementation/util/constants.py
b/java_based_implementation/util/constants.py
index 71ca0e6..7e8923a 100644
--- a/java_based_implementation/util/constants.py
+++ b/java_based_implementation/util/constants.py
@@ -23,3 +23,6 @@ PYPAIMON_JAVA_CLASSPATH = '_PYPAIMON_JAVA_CLASSPATH'
PYPAIMON_MAIN_CLASS = 'org.apache.paimon.python.PythonGatewayServer'
PYPAIMON_MAIN_ARGS = '_PYPAIMON_MAIN_ARGS'
PYPAIMON_TOX_TEST = '_PYPAIMON_TOX_TEST'
+
+# ------------------------ for catalog options ------------------------
+MAX_WORKERS = "max-workers"
diff --git a/java_based_implementation/util/java_utils.py
b/java_based_implementation/util/java_utils.py
index b9f2523..1f0d1ec 100644
--- a/java_based_implementation/util/java_utils.py
+++ b/java_based_implementation/util/java_utils.py
@@ -19,9 +19,9 @@
from java_based_implementation.java_gateway import get_gateway
-def to_j_catalog_context(catalog_context: dict):
+def to_j_catalog_context(catalog_options: dict):
gateway = get_gateway()
- j_options = gateway.jvm.Options(catalog_context)
+ j_options = gateway.jvm.Options(catalog_options)
return gateway.jvm.CatalogContext.create(j_options)
diff --git a/paimon_python_api/catalog.py b/paimon_python_api/catalog.py
index 1a78444..7fbacc5 100644
--- a/paimon_python_api/catalog.py
+++ b/paimon_python_api/catalog.py
@@ -28,7 +28,7 @@ class Catalog(ABC):
@staticmethod
@abstractmethod
- def create(catalog_context: dict) -> 'Catalog':
+ def create(catalog_options: dict) -> 'Catalog':
"""Create catalog from configuration."""
@abstractmethod
diff --git a/paimon_python_api/table_read.py b/paimon_python_api/table_read.py
index c17b49d..47b8aae 100644
--- a/paimon_python_api/table_read.py
+++ b/paimon_python_api/table_read.py
@@ -19,11 +19,12 @@
from abc import ABC, abstractmethod
from pyarrow import RecordBatchReader
from paimon_python_api.split import Split
+from typing import List
class TableRead(ABC):
"""To read data from data splits."""
@abstractmethod
- def create_reader(self, split: Split) -> RecordBatchReader:
+ def create_reader(self, splits: List[Split]) -> RecordBatchReader:
"""Return a reader containing batches of pyarrow format."""
