Skip to content

Commit 9c58cbc

Browse files
2sin182sin18
authored
[IO] Support reading batches from parquet files on Aliyun OSS.
This patch supports environment variables: - `S3_ENDPOINT`: Endpoint to S3/OSS service. - `S3_ADDRESSING_STYLE`: For Aliyun OSS, this environment variable must be virtual.
1 parent a90aeda commit 9c58cbc

3 files changed

Lines changed: 160 additions & 0 deletions

File tree

‎arrow/build.sh‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@
22

33
set -eo pipefail
44

5+
cd arrow/src
6+
git apply ../s3_enhancements.patch
7+
cd -
8+
59
if [[ ! -d $CACHE_DIR ]]; then
610
CACHE_DIR=arrow/cache
711
fi

‎arrow/s3_enhancements.patch‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
commit 0a6ecaa164267004a96f8df14a1ec9bbd2934c93
2+
Author: yuanman.ym <yuanman.ym@alibaba-inc.com>
3+
Date: Thu Dec 2 16:39:40 2021 +0800
4+
5+
[C++] Allow passing endpoint and addressing style by environment variables.
6+
7+
diff --git a/cpp/src/arrow/filesystem/s3fs.cc b/cpp/src/arrow/filesystem/s3fs.cc
8+
index cee0564..2779356 100644
9+
--- a/cpp/src/arrow/filesystem/s3fs.cc
10+
+++ b/cpp/src/arrow/filesystem/s3fs.cc
11+
@@ -82,6 +82,7 @@
12+
#include "arrow/util/atomic_shared_ptr.h"
13+
#include "arrow/util/checked_cast.h"
14+
#include "arrow/util/future.h"
15+
+#include "arrow/util/io_util.h"
16+
#include "arrow/util/key_value_metadata.h"
17+
#include "arrow/util/logging.h"
18+
#include "arrow/util/optional.h"
19+
@@ -90,6 +91,7 @@
20+
21+
namespace arrow {
22+
23+
+using internal::GetEnvVar;
24+
using internal::TaskGroup;
25+
using internal::Uri;
26+
using io::internal::SubmitIO;
27+
@@ -114,6 +116,12 @@ using internal::ToAwsString;
28+
using internal::ToURLEncodedAwsString;
29+
30+
static const char kSep = '/';
31+
+static const char kS3EndpointEnvVar[] = "S3_ENDPOINT";
32+
+static const char kS3UseHttpsEnvVar[] = "S3_USE_HTTPS";
33+
+static const char kS3AddressingStyle[] = "S3_ADDRESSING_STYLE";
34+
+static const char kHttpsPrefix[] = "https://";
35+
+static const char kHttpPrefix[] = "http://";
36+
+static const char kSchemeSep[] = "://";
37+
38+
namespace {
39+
40+
@@ -334,6 +342,39 @@ Result<S3Options> S3Options::FromUri(const Uri& uri, std::string* out_path) {
41+
options.ConfigureDefaultCredentials();
42+
}
43+
44+
+ auto maybe_endpoint = GetEnvVar(kS3EndpointEnvVar);
45+
+ if (maybe_endpoint.ok()) {
46+
+ std::string endpoint_str = *std::move(maybe_endpoint);
47+
+ if (endpoint_str.find(kSchemeSep) == std::string::npos) {
48+
+ int use_https = 1;
49+
+ auto maybe_use_https = GetEnvVar(kS3UseHttpsEnvVar);
50+
+ if (maybe_use_https.ok()) {
51+
+ try {
52+
+ use_https = std::stoi(*std::move(maybe_use_https));
53+
+ } catch (...) {
54+
+ // stoi throws an exception when it fails, just ignore it.
55+
+ }
56+
+ }
57+
+ if (use_https > 0) {
58+
+ endpoint_str = std::string(kHttpsPrefix) + endpoint_str;
59+
+ } else {
60+
+ endpoint_str = std::string(kHttpPrefix) + endpoint_str;
61+
+ }
62+
+ }
63+
+ Uri endpoint;
64+
+ RETURN_NOT_OK(endpoint.Parse(endpoint_str));
65+
+ options.scheme = endpoint.scheme();
66+
+ options.endpoint_override = endpoint.host();
67+
+ if (!endpoint.port_text().empty()) {
68+
+ options.endpoint_override += ":" + endpoint.port_text();
69+
+ }
70+
+ }
71+
+
72+
+ auto maybe_addressing_style = GetEnvVar(kS3AddressingStyle);
73+
+ if (maybe_addressing_style.ok()) {
74+
+ options.addressing_style = *std::move(maybe_addressing_style);
75+
+ }
76+
+
77+
bool region_set = false;
78+
for (const auto& kv : options_map) {
79+
if (kv.first == "region") {
80+
@@ -343,6 +384,8 @@ Result<S3Options> S3Options::FromUri(const Uri& uri, std::string* out_path) {
81+
options.scheme = kv.second;
82+
} else if (kv.first == "endpoint_override") {
83+
options.endpoint_override = kv.second;
84+
+ } else if (kv.first == "addressing_style") {
85+
+ options.addressing_style = kv.second;
86+
} else {
87+
return Status::Invalid("Unexpected query parameter in S3 URI: '", kv.first, "'");
88+
}
89+
@@ -573,7 +616,16 @@ class ClientBuilder {
90+
client_config_.caPath = ToAwsString(internal::global_options.tls_ca_dir_path);
91+
}
92+
93+
- const bool use_virtual_addressing = options_.endpoint_override.empty();
94+
+ bool use_virtual_addressing = true;
95+
+ if (options_.addressing_style == "virtual") {
96+
+ use_virtual_addressing = true;
97+
+ } else if (options_.addressing_style == "path") {
98+
+ use_virtual_addressing = false;
99+
+ } else if (options_.addressing_style == "auto") {
100+
+ use_virtual_addressing = options_.endpoint_override.empty();
101+
+ } else {
102+
+ return Status::Invalid("Invalid S3 addressing style '", options_.addressing_style, "'");
103+
+ }
104+
105+
/// Set proxy options if provided
106+
if (!options_.proxy_options.scheme.empty()) {
107+
diff --git a/cpp/src/arrow/filesystem/s3fs.h b/cpp/src/arrow/filesystem/s3fs.h
108+
index 1aad4dd..a276089 100644
109+
--- a/cpp/src/arrow/filesystem/s3fs.h
110+
+++ b/cpp/src/arrow/filesystem/s3fs.h
111+
@@ -85,6 +85,9 @@ struct ARROW_EXPORT S3Options {
112+
std::string endpoint_override;
113+
/// S3 connection transport, default "https"
114+
std::string scheme = "https";
115+
+ /// Addressing style which controls if the bucket name is in the hostname
116+
+ /// or part of the URL. Accept values are: path, virtual and auto.
117+
+ std::string addressing_style = "auto";
118+
119+
/// ARN of role to assume
120+
std::string role_arn;

‎docs/tutorials/data_loading.md‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,42 @@ batch = it.get_next()
110110
...
111111
```
112112

113+
### Read from files on S3/OSS/HDFS
114+
115+
```bash
116+
export S3_ENDPOINT=oss-cn-shanghai-internal.aliyuncs.com
117+
export AWS_ACCESS_KEY_ID=my_id
118+
export AWS_SECRET_ACCESS_KEY=my_secret
119+
export S3_ADDRESSING_STYLE=virtual
120+
```
121+
122+
```{eval-rst}
123+
.. note::
124+
See https://docs.w3cub.com/tensorflow~guide/deploy/s3.html for more
125+
information.
126+
.. note::
127+
Set `S3_ADDRESSING_STYLE` to `virtual` to support OSS.
128+
.. note::
129+
Set `S3_USE_HTTPS` to `0` to use `http` for S3 endpoint.
130+
```
131+
132+
```python
133+
import tensorflow as tf
134+
import hybridbackend.tensorflow as hb
135+
ds = hb.data.ParquetDataset(
136+
['s3://path/to/f1.parquet',
137+
'oss://path/to/f2.parquet',
138+
'hdfs://host:port/path/to/f3.parquet'],
139+
batch_size=1024,
140+
fields=['a', 'c'])
141+
ds = ds.apply(hb.data.to_sparse())
142+
ds = ds.prefetch(4)
143+
it = tf.data.make_one_shot_iterator(ds)
144+
batch = it.get_next()
145+
# {'a': tensora, 'c': tensorc}
146+
...
147+
```
148+
113149
## Performance
114150

115151
In benchmark for reading 20k samples from 200 columns of a Parquet file,

0 commit comments

Comments
 (0)