Skip to content

Commit a67461f

Browse files
feat: support fixed-size arrays in data readers (#726)
## Summary - support parsing explicitly typed fixed-size arrays from CSV `LOAD FROM`, including nested arrays, quoted strings, temporal values, and length validation - support fixed-size array values in the JSON reader and when materializing `COPY FROM` results into graph properties - support Arrow `FIXED_SIZE_LIST` columns in the Parquet reader, including nested arrays and scalar child types - add end-to-end CSV, JSON, and Parquet array coverage and document the explicit `CAST(..., T[N])` workflow Variable-length `LIST` types such as `FLOAT[]` remain unsupported by these readers. Users must explicitly specify a fixed-size array type such as `FLOAT[3]`; implicit LIST-to-ARRAY conversion is not performed. ## Testing - `python3 -m pytest -q tools/python_bind/tests/test_load_array.py` — 15 passed Fixes #441 Fixes #442 Fixes #443 --------- Co-authored-by: Longbin Lai <longbin.lai@gmail.com>
1 parent f49b5e0 commit a67461f

12 files changed

Lines changed: 1194 additions & 41 deletions

File tree

.github/workflows/build-and-test-common.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,7 @@ jobs:
305305
python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_query.py
306306
python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_tp_service.py
307307
python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_load.py
308+
python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_load_array.py
308309
python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_export.py
309310
python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_copy_temp.py
310311

.github/workflows/neug-extension-test.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,7 @@ jobs:
140140
export FLEX_DATA_DIR=${GITHUB_WORKSPACE}/example_dataset/comprehensive_graph
141141
GLOG_v=10 ${GITHUB_WORKSPACE}/build/tools/utils/bulk_loader -g ../../example_dataset/comprehensive_graph/graph.yaml -l ../../example_dataset/comprehensive_graph/import.yaml -d /tmp/comprehensive_graph
142142
NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_load.py -k "parquet or httpfs"
143+
NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_load_array.py -k "parquet"
143144
NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_export.py -k "parquet or httpfs"
144145
NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_db_import_export.py -k "copy_from_parquet"
145146
NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_sniffer.py -k "parquet"

doc/source/data_io/import_data.md

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -240,7 +240,25 @@ The following options control how CSV files are parsed during `COPY FROM`. These
240240
| `quoting` | bool | `true` | Whether to enable quote processing |
241241
| `escaping` | bool | `true` | Whether to enable escape character processing |
242242

243-
> **Array limitation:** `COPY FROM` CSV currently does not auto-detect or directly materialize fixed-size `ARRAY` columns from bracketed CSV fields such as `"[1,2,3]"`. Use Cypher literals/parameters for array properties, or a typed non-CSV ingestion path when available.
243+
`COPY FROM` supports loading `ARRAY` data from CSV,
244+
JSON, and Parquet files. NeuG does not implicitly convert input values to
245+
`ARRAY`; use `LOAD FROM` and explicitly cast the column to a fixed-size
246+
`ARRAY` type.
247+
248+
For example, given a `person_array.csv` file with an array column:
249+
250+
```csv
251+
id,name,address
252+
1,Alice,"[Beijing,Hangzhou,Shanghai]"
253+
2,Bob,"[London,Paris,Berlin]"
254+
```
255+
256+
```cypher
257+
COPY Person FROM (
258+
LOAD FROM "person_array.csv" (delim=',')
259+
RETURN id, name, CAST(address, 'STRING[3]') AS addresses
260+
);
261+
```
244262

245263
### JSON/JSONL
246264

doc/source/data_io/load_data.md

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,23 @@ LOAD FROM "person.csv" (delim=',', header=true)
4646
RETURN name, age;
4747
```
4848

49-
> **Array limitation:** `LOAD FROM` CSV currently does not parse CSV fields into fixed-size `ARRAY` values. A quoted field such as `"[1,2,3]"` is read as a `STRING`, and casting CSV columns to `INT64[3]`/other fixed-size array targets inside `LOAD FROM` is not supported yet.
49+
NeuG supports reading `ARRAY` data from CSV, JSON,
50+
and Parquet files. NeuG does not implicitly convert input values to `ARRAY`;
51+
explicitly cast the column to a fixed-size `ARRAY` type in the `RETURN`
52+
clause.
53+
54+
For example, given a `person_array.csv` file with an array column:
55+
56+
```csv
57+
id,name,address
58+
1,Alice,"[Beijing,Hangzhou,Shanghai]"
59+
2,Bob,"[London,Paris,Berlin]"
60+
```
61+
62+
```cypher
63+
LOAD FROM "person_array.csv" (delim=',')
64+
RETURN id, name, CAST(address, 'STRING[3]') AS addresses;
65+
```
5066

5167
### JSON / JSONL
5268

extension/parquet/src/arrow_column.cc

Lines changed: 184 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,181 @@
1616
#include "parquet/arrow_column.h"
1717

1818
#include <arrow/array/array_binary.h>
19+
#include <arrow/array/array_nested.h>
1920
#include <arrow/type.h>
2021

22+
#include "neug/common/columns/columns_utils.h"
2123
#include "neug/common/columns/value_columns.h"
24+
#include "neug/common/types/value.h"
2225
#include "neug/utils/exception/exception.h"
2326

2427
namespace neug {
2528

29+
static int64_t arrow_time_to_milliseconds(int64_t value,
30+
arrow::TimeUnit::type unit) {
31+
switch (unit) {
32+
case arrow::TimeUnit::SECOND:
33+
return value * Interval::MSECS_PER_SEC;
34+
case arrow::TimeUnit::MILLI:
35+
return value;
36+
case arrow::TimeUnit::MICRO:
37+
return value / Interval::MICROS_PER_MSEC;
38+
case arrow::TimeUnit::NANO:
39+
return value / (Interval::MICROS_PER_MSEC * 1000);
40+
}
41+
THROW_NOT_SUPPORTED_EXCEPTION("Unsupported Arrow time unit");
42+
}
43+
44+
// Keep this mapping in sync with arrow_value_at when adding Arrow types.
45+
static DataType arrow_type_to_neug_type(const arrow::DataType& type) {
46+
switch (type.id()) {
47+
case arrow::Type::BOOL:
48+
return DataType::BOOLEAN;
49+
case arrow::Type::INT32:
50+
return DataType::INT32;
51+
case arrow::Type::INT64:
52+
return DataType::INT64;
53+
case arrow::Type::UINT32:
54+
return DataType::UINT32;
55+
case arrow::Type::UINT64:
56+
return DataType::UINT64;
57+
case arrow::Type::FLOAT:
58+
return DataType::FLOAT;
59+
case arrow::Type::DOUBLE:
60+
return DataType::DOUBLE;
61+
case arrow::Type::STRING:
62+
case arrow::Type::LARGE_STRING:
63+
return DataType::VARCHAR;
64+
case arrow::Type::DATE32:
65+
case arrow::Type::DATE64:
66+
return DataType::DATE;
67+
case arrow::Type::TIMESTAMP:
68+
return DataType::TIMESTAMP_MS;
69+
case arrow::Type::DURATION:
70+
return DataType::INTERVAL;
71+
case arrow::Type::FIXED_SIZE_LIST: {
72+
const auto& array_type = static_cast<const arrow::FixedSizeListType&>(type);
73+
return DataType::Array(arrow_type_to_neug_type(*array_type.value_type()),
74+
array_type.list_size());
75+
}
76+
case arrow::Type::LIST:
77+
case arrow::Type::LARGE_LIST:
78+
THROW_NOT_SUPPORTED_EXCEPTION(
79+
"Parquet LIST is not supported as ARRAY. Specify a fixed-size Arrow "
80+
"list / NeuG ARRAY type instead.");
81+
default:
82+
THROW_NOT_SUPPORTED_EXCEPTION("Unsupported arrow type: " + type.ToString());
83+
}
84+
}
85+
86+
static Value arrow_value_at(const arrow::Array& array, int64_t index,
87+
const DataType& type) {
88+
if (array.IsNull(index)) {
89+
return Value(type);
90+
}
91+
switch (array.type_id()) {
92+
case arrow::Type::BOOL:
93+
return Value::BOOLEAN(
94+
static_cast<const arrow::BooleanArray&>(array).Value(index));
95+
case arrow::Type::INT32:
96+
return Value::INT32(
97+
static_cast<const arrow::Int32Array&>(array).Value(index));
98+
case arrow::Type::INT64:
99+
return Value::INT64(
100+
static_cast<const arrow::Int64Array&>(array).Value(index));
101+
case arrow::Type::UINT32:
102+
return Value::UINT32(
103+
static_cast<const arrow::UInt32Array&>(array).Value(index));
104+
case arrow::Type::UINT64:
105+
return Value::UINT64(
106+
static_cast<const arrow::UInt64Array&>(array).Value(index));
107+
case arrow::Type::FLOAT:
108+
return Value::FLOAT(
109+
static_cast<const arrow::FloatArray&>(array).Value(index));
110+
case arrow::Type::DOUBLE:
111+
return Value::DOUBLE(
112+
static_cast<const arrow::DoubleArray&>(array).Value(index));
113+
case arrow::Type::STRING:
114+
return Value::STRING(std::string(
115+
static_cast<const arrow::StringArray&>(array).GetView(index)));
116+
case arrow::Type::LARGE_STRING:
117+
return Value::STRING(std::string(
118+
static_cast<const arrow::LargeStringArray&>(array).GetView(index)));
119+
case arrow::Type::DATE32: {
120+
Date date;
121+
date.from_num_days(
122+
static_cast<const arrow::Date32Array&>(array).Value(index));
123+
return Value::DATE(date);
124+
}
125+
case arrow::Type::DATE64: {
126+
constexpr int64_t MILLIS_PER_DAY = 24LL * 60 * 60 * 1000;
127+
Date date;
128+
date.from_num_days(static_cast<int32_t>(
129+
static_cast<const arrow::Date64Array&>(array).Value(index) /
130+
MILLIS_PER_DAY));
131+
return Value::DATE(date);
132+
}
133+
case arrow::Type::TIMESTAMP: {
134+
const auto value =
135+
static_cast<const arrow::TimestampArray&>(array).Value(index);
136+
const auto& timestamp_type =
137+
static_cast<const arrow::TimestampType&>(*array.type());
138+
return Value::TIMESTAMPMS(
139+
DateTime(arrow_time_to_milliseconds(value, timestamp_type.unit())));
140+
}
141+
case arrow::Type::DURATION: {
142+
const auto value =
143+
static_cast<const arrow::DurationArray&>(array).Value(index);
144+
const auto& duration_type =
145+
static_cast<const arrow::DurationType&>(*array.type());
146+
Interval interval;
147+
interval.from_mill_seconds(
148+
arrow_time_to_milliseconds(value, duration_type.unit()));
149+
return Value::INTERVAL(interval);
150+
}
151+
case arrow::Type::FIXED_SIZE_LIST: {
152+
const auto& list = static_cast<const arrow::FixedSizeListArray&>(array);
153+
const auto& array_type =
154+
static_cast<const arrow::FixedSizeListType&>(*array.type());
155+
const auto child_type = ArrayType::GetChildType(type);
156+
std::vector<Value> values;
157+
values.reserve(array_type.list_size());
158+
const auto offset = list.value_offset(index);
159+
for (int32_t i = 0; i < array_type.list_size(); ++i) {
160+
values.push_back(arrow_value_at(*list.values(), offset + i, child_type));
161+
}
162+
return Value::ARRAY(type, std::move(values));
163+
}
164+
default:
165+
THROW_NOT_SUPPORTED_EXCEPTION("Unsupported arrow type: " +
166+
array.type()->ToString());
167+
}
168+
}
169+
170+
static std::shared_ptr<IContextColumn> convert_fixed_size_list_arrays(
171+
const std::vector<std::shared_ptr<arrow::Array>>& arrays) {
172+
const auto type = arrow_type_to_neug_type(*arrays.front()->type());
173+
auto builder = ColumnsUtils::create_builder(type);
174+
size_t size = 0;
175+
for (const auto& array : arrays) {
176+
size += array->length();
177+
}
178+
builder->reserve(size);
179+
for (const auto& array : arrays) {
180+
if (!array->type()->Equals(*arrays.front()->type())) {
181+
THROW_SCHEMA_MISMATCH("Parquet ARRAY chunks have different types");
182+
}
183+
for (int64_t i = 0; i < array->length(); ++i) {
184+
if (array->IsNull(i)) {
185+
builder->push_back_null();
186+
continue;
187+
}
188+
builder->push_back_elem(arrow_value_at(*array, i, type));
189+
}
190+
}
191+
return builder->finish();
192+
}
193+
26194
/// Convert numeric arrow arrays directly to ValueColumn<CppT>.
27195
template <typename ArrowArrayT, typename CppT>
28196
static std::shared_ptr<IContextColumn> convert_numeric_arrays(
@@ -89,7 +257,11 @@ static std::shared_ptr<IContextColumn> convert_date64_arrays(
89257
if (typed->IsNull(j)) {
90258
builder.push_back_null();
91259
} else {
92-
builder.push_back_opt(Date(typed->Value(j)));
260+
constexpr int64_t MILLIS_PER_DAY = 24LL * 60 * 60 * 1000;
261+
Date date;
262+
date.from_num_days(
263+
static_cast<int32_t>(typed->Value(j) / MILLIS_PER_DAY));
264+
builder.push_back_opt(date);
93265
}
94266
}
95267
}
@@ -106,7 +278,10 @@ static std::shared_ptr<IContextColumn> convert_timestamp_arrays(
106278
if (typed->IsNull(j)) {
107279
builder.push_back_null();
108280
} else {
109-
builder.push_back_opt(DateTime(typed->Value(j)));
281+
const auto& timestamp_type =
282+
static_cast<const arrow::TimestampType&>(*typed->type());
283+
builder.push_back_opt(DateTime(arrow_time_to_milliseconds(
284+
typed->Value(j), timestamp_type.unit())));
110285
}
111286
}
112287
}
@@ -144,6 +319,13 @@ std::shared_ptr<IContextColumn> arrow_arrays_to_value_column(
144319
return convert_date64_arrays(arrays);
145320
case arrow::Type::TIMESTAMP:
146321
return convert_timestamp_arrays(arrays);
322+
case arrow::Type::FIXED_SIZE_LIST:
323+
return convert_fixed_size_list_arrays(arrays);
324+
case arrow::Type::LIST:
325+
case arrow::Type::LARGE_LIST:
326+
THROW_NOT_SUPPORTED_EXCEPTION(
327+
"Parquet LIST columns are not supported; use a fixed-size ARRAY "
328+
"type such as FLOAT[3]");
147329
default:
148330
THROW_NOT_SUPPORTED_EXCEPTION("Unsupported arrow type: " +
149331
arrow_type->ToString());

include/neug/storages/loader/loader_utils.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -99,17 +99,17 @@ void fillVertexReaderMeta(label_t v_label, const std::string& v_label_name,
9999
const std::string& v_file,
100100
const LoadingConfig& loading_config,
101101
const std::vector<std::string>& vertex_property_names,
102-
const std::vector<DataTypeId>& vertex_property_types,
103-
DataTypeId pk_type, const std::string& pk_name,
102+
const std::vector<DataType>& vertex_property_types,
103+
DataType pk_type, const std::string& pk_name,
104104
size_t pk_ind, CsvReadConfig& config);
105105

106106
void fillEdgeReaderMeta(label_t src_label_id, label_t dst_label_id,
107107
label_t label_id, const std::string& edge_label_name,
108108
const std::string& e_file,
109109
const LoadingConfig& loading_config,
110110
const std::vector<std::string>& edge_property_names,
111-
const std::vector<DataTypeId>& edge_property_types,
112-
DataTypeId src_pk_type, DataTypeId dst_pk_type,
111+
const std::vector<DataType>& edge_property_types,
112+
DataType src_pk_type, DataType dst_pk_type,
113113
CsvReadConfig& config);
114114

115115
void set_properties_from_context_column(

src/storages/loader/csv_property_graph_loader.cc

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,12 +32,12 @@ CSVPropertyGraphLoader::createVertexChunkSupplier(
3232
DataType pk_type, const std::string& pk_name, int pk_ind,
3333
const LoadingConfig& loading_config, int thread_id) const {
3434
auto vertex_property_names = schema_.get_vertex_property_names(v_label);
35-
auto vertex_property_types = schema_.get_vertex_properties_id(v_label);
35+
auto vertex_property_types = schema_.get_vertex_properties(v_label);
3636

3737
CsvReadConfig config;
3838
fillVertexReaderMeta(v_label, v_label_name, v_file, loading_config,
39-
vertex_property_names, vertex_property_types,
40-
pk_type.id(), pk_name, pk_ind, config);
39+
vertex_property_names, vertex_property_types, pk_type,
40+
pk_name, pk_ind, config);
4141
return std::make_shared<CSVChunkSupplier>(v_file, std::move(config));
4242
}
4343

@@ -49,7 +49,7 @@ CSVPropertyGraphLoader::createEdgeChunkSupplier(
4949
auto edge_property_names =
5050
schema_.get_edge_property_names(src_label_id, dst_label_id, e_label_id);
5151
auto edge_property_types =
52-
schema_.get_edge_properties_id(src_label_id, dst_label_id, e_label_id);
52+
schema_.get_edge_properties(src_label_id, dst_label_id, e_label_id);
5353
auto src_pk_type =
5454
std::get<0>(schema_.get_vertex_primary_key(src_label_id)[0]);
5555
auto dst_pk_type =
@@ -58,7 +58,7 @@ CSVPropertyGraphLoader::createEdgeChunkSupplier(
5858
fillEdgeReaderMeta(src_label_id, dst_label_id, e_label_id,
5959
schema_.get_edge_label_name(e_label_id), e_file,
6060
loading_config_, edge_property_names, edge_property_types,
61-
src_pk_type.id(), dst_pk_type.id(), config);
61+
src_pk_type, dst_pk_type, config);
6262
return std::make_shared<CSVChunkSupplier>(e_file, std::move(config));
6363
}
6464

0 commit comments

Comments
 (0)