This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new 63e7652868 feat(c/driver/postgresql): implement COPY writer for time
types (#4819)
63e7652868 is described below
commit 63e765286840aa5ff375de0a3bd68ffa2582af45
Author: David Li <[email protected]>
AuthorDate: Tue Sep 29 20:00:20 2026 -0700
feat(c/driver/postgresql): implement COPY writer for time types (#4819)
time64[us] was supported; add the other three possibilities.
Closes #3841.
Assisted-by: GPT-6 Astra
---
.../postgresql/copy/postgres_copy_writer_test.cc | 87 +++++++++++++++-------
c/driver/postgresql/copy/writer.h | 49 +++++++++++-
.../type/bind/time_ms.txtcase | 8 +-
.../type/bind/time_ns.txtcase | 8 +-
.../type/bind/time_s.txtcase | 8 +-
.../validation/queries/ingest/time_ms.txtcase | 30 +++++++-
.../validation/queries/ingest/time_ns.txtcase | 30 +++++++-
.../validation/queries/ingest/time_s.txtcase | 30 +++++++-
.../validation/queries/type/bind/time_ms.txtcase | 25 ++++++-
.../validation/queries/type/bind/time_ns.txtcase | 32 +++++++-
.../validation/queries/type/bind/time_s.txtcase | 25 ++++++-
11 files changed, 281 insertions(+), 51 deletions(-)
diff --git a/c/driver/postgresql/copy/postgres_copy_writer_test.cc
b/c/driver/postgresql/copy/postgres_copy_writer_test.cc
index f2bd5e23bf..e485d2b80c 100644
--- a/c/driver/postgresql/copy/postgres_copy_writer_test.cc
+++ b/c/driver/postgresql/copy/postgres_copy_writer_test.cc
@@ -473,33 +473,70 @@ TEST_F(PostgresCopyTest, PostgresCopyWriteDate) {
}
TEST_F(PostgresCopyTest, PostgresCopyWriteTime) {
- adbc_validation::Handle<struct ArrowSchema> schema;
- adbc_validation::Handle<struct ArrowArray> array;
- struct ArrowError na_error;
-
- const enum ArrowTimeUnit unit = NANOARROW_TIME_UNIT_MICRO;
- const auto values =
- std::vector<std::optional<int64_t>>{0, 86399000000, 49376123456,
std::nullopt};
-
- ArrowSchemaInit(&schema.value);
- ArrowSchemaSetTypeStruct(&schema.value, 1);
- ArrowSchemaSetTypeDateTime(schema->children[0], NANOARROW_TYPE_TIME64, unit,
nullptr);
- ArrowSchemaSetName(schema->children[0], "col");
- ASSERT_EQ(
- adbc_validation::MakeBatch<int64_t>(&schema.value, &array.value,
&na_error, values),
- ADBC_STATUS_OK);
+ // COPY (SELECT CAST(col AS TIME) FROM (VALUES ('00:00:00'), ('23:59:59'),
+ // (NULL)) AS drvd(col)) TO STDOUT WITH (FORMAT binary);
+ static const uint8_t expected[] = {
+ 0x50, 0x47, 0x43, 0x4f, 0x50, 0x59, 0x0a, 0xff, 0x0d, 0x0a, 0x00, 0x00,
0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0x00,
0x00, 0x00,
+ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x08, 0x00,
0x00, 0x00,
+ 0x14, 0x1d, 0xc8, 0x1d, 0xc0, 0x00, 0x01, 0xff, 0xff, 0xff, 0xff, 0xff,
0xff};
+ struct TimeTestParamType {
+ enum ArrowType type;
+ enum ArrowTimeUnit unit;
+ std::vector<std::optional<int64_t>> values;
+ const uint8_t* expected;
+ size_t expected_size;
+ };
+ const TimeTestParamType params[] = {
+ {NANOARROW_TYPE_TIME32,
+ NANOARROW_TIME_UNIT_SECOND,
+ {0, 86399, std::nullopt},
+ expected,
+ sizeof(expected)},
+ {NANOARROW_TYPE_TIME32,
+ NANOARROW_TIME_UNIT_MILLI,
+ {0, 86399000, std::nullopt},
+ expected,
+ sizeof(expected)},
+ {NANOARROW_TYPE_TIME64,
+ NANOARROW_TIME_UNIT_MICRO,
+ {0, 86399000000, 49376123456, std::nullopt},
+ kTestPgCopyTime,
+ sizeof(kTestPgCopyTime)},
+ {NANOARROW_TYPE_TIME64,
+ NANOARROW_TIME_UNIT_NANO,
+ {0, 86399000000000, 49376123456000, std::nullopt},
+ kTestPgCopyTime,
+ sizeof(kTestPgCopyTime)},
+ };
- PostgresCopyStreamWriteTester tester;
- ASSERT_EQ(tester.Init(&schema.value, &array.value, *type_resolver_),
NANOARROW_OK);
- ASSERT_EQ(tester.WriteAll(nullptr), ENODATA);
+ for (const auto& param : params) {
+ SCOPED_TRACE(param.unit);
+ adbc_validation::Handle<struct ArrowSchema> schema;
+ adbc_validation::Handle<struct ArrowArray> array;
+ struct ArrowError na_error;
- const struct ArrowBuffer buf = tester.WriteBuffer();
- // The last 2 bytes of a message can be transmitted via PQputCopyData
- // so no need to test those bytes from the Writer
- constexpr size_t buf_size = sizeof(kTestPgCopyTime) - 2;
- ASSERT_EQ(buf.size_bytes, buf_size);
- for (size_t i = 0; i < buf_size; i++) {
- ASSERT_EQ(buf.data[i], kTestPgCopyTime[i]);
+ ArrowSchemaInit(&schema.value);
+ ASSERT_EQ(ArrowSchemaSetTypeStruct(&schema.value, 1), NANOARROW_OK);
+ ASSERT_EQ(
+ ArrowSchemaSetTypeDateTime(schema->children[0], param.type,
param.unit, nullptr),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "col"), NANOARROW_OK);
+ ASSERT_EQ(adbc_validation::MakeBatch<int64_t>(&schema.value, &array.value,
&na_error,
+ param.values),
+ ADBC_STATUS_OK);
+
+ PostgresCopyStreamWriteTester tester;
+ ASSERT_EQ(tester.Init(&schema.value, &array.value, *type_resolver_),
NANOARROW_OK);
+ ASSERT_EQ(tester.WriteAll(nullptr), ENODATA);
+
+ const struct ArrowBuffer buf = tester.WriteBuffer();
+ // The end marker is sent separately by the caller.
+ const size_t buf_size = param.expected_size - 2;
+ ASSERT_EQ(buf.size_bytes, buf_size);
+ for (size_t i = 0; i < buf_size; i++) {
+ ASSERT_EQ(buf.data[i], param.expected[i]);
+ }
}
}
diff --git a/c/driver/postgresql/copy/writer.h
b/c/driver/postgresql/copy/writer.h
index 729f6060d9..4b38e710f7 100644
--- a/c/driver/postgresql/copy/writer.h
+++ b/c/driver/postgresql/copy/writer.h
@@ -925,6 +925,34 @@ class PostgresCopyListFieldWriter : public
PostgresCopyFieldWriter {
std::unique_ptr<PostgresCopyFieldWriter> child_;
};
+template <enum ArrowTimeUnit TU>
+class PostgresCopyTimeFieldWriter : public PostgresCopyFieldWriter {
+ public:
+ ArrowErrorCode Write(ArrowBuffer* buffer, int64_t index, ArrowError* error)
override {
+ constexpr int32_t field_size_bytes = sizeof(int64_t);
+ NANOARROW_RETURN_NOT_OK(WriteChecked<int32_t>(buffer, field_size_bytes,
error));
+
+ // We assume data is valid, so multiply can't overflow
+ int64_t value = ArrowArrayViewGetIntUnsafe(array_view_, index);
+ switch (TU) {
+ case NANOARROW_TIME_UNIT_SECOND:
+ value *= 1000000;
+ break;
+ case NANOARROW_TIME_UNIT_MILLI:
+ value *= 1000;
+ break;
+ case NANOARROW_TIME_UNIT_MICRO:
+ break;
+ case NANOARROW_TIME_UNIT_NANO:
+ // We truncate; at some point we should add an option to control
behavior
+ value /= 1000;
+ break;
+ }
+
+ return WriteChecked<int64_t>(buffer, value, error);
+ }
+};
+
template <enum ArrowTimeUnit TU>
class PostgresCopyTimestampFieldWriter : public PostgresCopyFieldWriter {
public:
@@ -1074,12 +1102,29 @@ static inline ArrowErrorCode MakeCopyFieldWriter(
*out = T::Create<T>(array_view);
return NANOARROW_OK;
}
+ case NANOARROW_TYPE_TIME32:
case NANOARROW_TYPE_TIME64: {
switch (schema_view.time_unit) {
- case NANOARROW_TIME_UNIT_MICRO:
- using T = PostgresCopyNetworkEndianFieldWriter<int64_t>;
+ case NANOARROW_TIME_UNIT_SECOND: {
+ using T = PostgresCopyTimeFieldWriter<NANOARROW_TIME_UNIT_SECOND>;
*out = T::Create<T>(array_view);
return NANOARROW_OK;
+ }
+ case NANOARROW_TIME_UNIT_MILLI: {
+ using T = PostgresCopyTimeFieldWriter<NANOARROW_TIME_UNIT_MILLI>;
+ *out = T::Create<T>(array_view);
+ return NANOARROW_OK;
+ }
+ case NANOARROW_TIME_UNIT_MICRO: {
+ using T = PostgresCopyTimeFieldWriter<NANOARROW_TIME_UNIT_MICRO>;
+ *out = T::Create<T>(array_view);
+ return NANOARROW_OK;
+ }
+ case NANOARROW_TIME_UNIT_NANO: {
+ using T = PostgresCopyTimeFieldWriter<NANOARROW_TIME_UNIT_NANO>;
+ *out = T::Create<T>(array_view);
+ return NANOARROW_OK;
+ }
default:
return ADBC_STATUS_NOT_IMPLEMENTED;
}
diff --git a/c/driver/postgresql/validation/queries/type/bind/time_ms.txtcase
b/c/driver/postgresql/validation/queries-materialize/type/bind/time_ms.txtcase
similarity index 91%
copy from c/driver/postgresql/validation/queries/type/bind/time_ms.txtcase
copy to
c/driver/postgresql/validation/queries-materialize/type/bind/time_ms.txtcase
index 30dfb7d9b3..dcb6e22da4 100644
--- a/c/driver/postgresql/validation/queries/type/bind/time_ms.txtcase
+++
b/c/driver/postgresql/validation/queries-materialize/type/bind/time_ms.txtcase
@@ -15,7 +15,9 @@
// specific language governing permissions and limitations
// under the License.
+// part: setup_query
-// part: metadata
-
-skip = "COPY Writer not implemented"
+CREATE TABLE test_time (
+ idx INT,
+ res TIME
+);
diff --git a/c/driver/postgresql/validation/queries/type/bind/time_ns.txtcase
b/c/driver/postgresql/validation/queries-materialize/type/bind/time_ns.txtcase
similarity index 91%
copy from c/driver/postgresql/validation/queries/type/bind/time_ns.txtcase
copy to
c/driver/postgresql/validation/queries-materialize/type/bind/time_ns.txtcase
index 30dfb7d9b3..dcb6e22da4 100644
--- a/c/driver/postgresql/validation/queries/type/bind/time_ns.txtcase
+++
b/c/driver/postgresql/validation/queries-materialize/type/bind/time_ns.txtcase
@@ -15,7 +15,9 @@
// specific language governing permissions and limitations
// under the License.
+// part: setup_query
-// part: metadata
-
-skip = "COPY Writer not implemented"
+CREATE TABLE test_time (
+ idx INT,
+ res TIME
+);
diff --git a/c/driver/postgresql/validation/queries/type/bind/time_s.txtcase
b/c/driver/postgresql/validation/queries-materialize/type/bind/time_s.txtcase
similarity index 91%
copy from c/driver/postgresql/validation/queries/type/bind/time_s.txtcase
copy to
c/driver/postgresql/validation/queries-materialize/type/bind/time_s.txtcase
index 30dfb7d9b3..dcb6e22da4 100644
--- a/c/driver/postgresql/validation/queries/type/bind/time_s.txtcase
+++
b/c/driver/postgresql/validation/queries-materialize/type/bind/time_s.txtcase
@@ -15,7 +15,9 @@
// specific language governing permissions and limitations
// under the License.
+// part: setup_query
-// part: metadata
-
-skip = "COPY Writer not implemented"
+CREATE TABLE test_time (
+ idx INT,
+ res TIME
+);
diff --git a/c/driver/postgresql/validation/queries/ingest/time_ms.txtcase
b/c/driver/postgresql/validation/queries/ingest/time_ms.txtcase
index 5ee360e824..f9d691bea0 100644
--- a/c/driver/postgresql/validation/queries/ingest/time_ms.txtcase
+++ b/c/driver/postgresql/validation/queries/ingest/time_ms.txtcase
@@ -15,10 +15,34 @@
// specific language governing permissions and limitations
// under the License.
-
// part: metadata
-skip = "COPY Writer not implemented"
-
[tags]
field-type-name = "time"
+
+// part: expected_schema
+
+{
+ "format": "+s",
+ "children": [
+ {
+ "name": "idx",
+ "format": "l",
+ "flags": ["nullable"]
+ },
+ {
+ "name": "value",
+ "format": "ttu",
+ "flags": ["nullable"]
+ }
+ ]
+}
+
+// part: expected
+
+{"idx": 0, "value": 0}
+{"idx": 1, "value": 3661000000}
+{"idx": 2, "value": 43200000000}
+{"idx": 3, "value": 49531123000}
+{"idx": 4, "value": 86399999000}
+{"idx": 5, "value": null}
diff --git a/c/driver/postgresql/validation/queries/ingest/time_ns.txtcase
b/c/driver/postgresql/validation/queries/ingest/time_ns.txtcase
index 5ee360e824..8f5eb1d340 100644
--- a/c/driver/postgresql/validation/queries/ingest/time_ns.txtcase
+++ b/c/driver/postgresql/validation/queries/ingest/time_ns.txtcase
@@ -15,10 +15,34 @@
// specific language governing permissions and limitations
// under the License.
-
// part: metadata
-skip = "COPY Writer not implemented"
-
[tags]
field-type-name = "time"
+
+// part: expected_schema
+
+{
+ "format": "+s",
+ "children": [
+ {
+ "name": "idx",
+ "format": "l",
+ "flags": ["nullable"]
+ },
+ {
+ "name": "value",
+ "format": "ttu",
+ "flags": ["nullable"]
+ }
+ ]
+}
+
+// part: expected
+
+{"idx": 0, "value": 0}
+{"idx": 1, "value": 3661000000}
+{"idx": 2, "value": 43200000000}
+{"idx": 3, "value": 49531123456}
+{"idx": 4, "value": 86399999999}
+{"idx": 5, "value": null}
diff --git a/c/driver/postgresql/validation/queries/ingest/time_s.txtcase
b/c/driver/postgresql/validation/queries/ingest/time_s.txtcase
index 5ee360e824..db24221186 100644
--- a/c/driver/postgresql/validation/queries/ingest/time_s.txtcase
+++ b/c/driver/postgresql/validation/queries/ingest/time_s.txtcase
@@ -15,10 +15,34 @@
// specific language governing permissions and limitations
// under the License.
-
// part: metadata
-skip = "COPY Writer not implemented"
-
[tags]
field-type-name = "time"
+
+// part: expected_schema
+
+{
+ "format": "+s",
+ "children": [
+ {
+ "name": "idx",
+ "format": "l",
+ "flags": ["nullable"]
+ },
+ {
+ "name": "value",
+ "format": "ttu",
+ "flags": ["nullable"]
+ }
+ ]
+}
+
+// part: expected
+
+{"idx": 0, "value": 0}
+{"idx": 1, "value": 3661000000}
+{"idx": 2, "value": 43200000000}
+{"idx": 3, "value": 49531000000}
+{"idx": 4, "value": 86399000000}
+{"idx": 5, "value": null}
diff --git a/c/driver/postgresql/validation/queries/type/bind/time_ms.txtcase
b/c/driver/postgresql/validation/queries/type/bind/time_ms.txtcase
index 30dfb7d9b3..83fd316a8c 100644
--- a/c/driver/postgresql/validation/queries/type/bind/time_ms.txtcase
+++ b/c/driver/postgresql/validation/queries/type/bind/time_ms.txtcase
@@ -15,7 +15,28 @@
// specific language governing permissions and limitations
// under the License.
-
// part: metadata
-skip = "COPY Writer not implemented"
+[tags]
+sql-type-name = "time"
+
+// part: expected_schema
+
+{
+ "format": "+s",
+ "children": [
+ {
+ "name": "res",
+ "format": "ttu",
+ "flags": ["nullable"]
+ }
+ ]
+}
+
+// part: expected
+
+{"res": 0}
+{"res": null}
+{"res": 49531123000}
+{"res": 45045500000}
+{"res": 86399999000}
diff --git a/c/driver/postgresql/validation/queries/type/bind/time_ns.txtcase
b/c/driver/postgresql/validation/queries/type/bind/time_ns.txtcase
index 30dfb7d9b3..62fad42471 100644
--- a/c/driver/postgresql/validation/queries/type/bind/time_ns.txtcase
+++ b/c/driver/postgresql/validation/queries/type/bind/time_ns.txtcase
@@ -15,7 +15,35 @@
// specific language governing permissions and limitations
// under the License.
-
// part: metadata
-skip = "COPY Writer not implemented"
+[tags]
+sql-type-name = "time"
+
+// part: setup_query
+
+CREATE TABLE test_time (
+ idx INT,
+ res TIME(6)
+);
+
+// part: expected_schema
+
+{
+ "format": "+s",
+ "children": [
+ {
+ "name": "res",
+ "format": "ttu",
+ "flags": ["nullable"]
+ }
+ ]
+}
+
+// part: expected
+
+{"res": 0}
+{"res": null}
+{"res": 49531123456}
+{"res": 45045500000}
+{"res": 86399999999}
diff --git a/c/driver/postgresql/validation/queries/type/bind/time_s.txtcase
b/c/driver/postgresql/validation/queries/type/bind/time_s.txtcase
index 30dfb7d9b3..0bc1666e1b 100644
--- a/c/driver/postgresql/validation/queries/type/bind/time_s.txtcase
+++ b/c/driver/postgresql/validation/queries/type/bind/time_s.txtcase
@@ -15,7 +15,28 @@
// specific language governing permissions and limitations
// under the License.
-
// part: metadata
-skip = "COPY Writer not implemented"
+[tags]
+sql-type-name = "time"
+
+// part: expected_schema
+
+{
+ "format": "+s",
+ "children": [
+ {
+ "name": "res",
+ "format": "ttu",
+ "flags": ["nullable"]
+ }
+ ]
+}
+
+// part: expected
+
+{"res": 0}
+{"res": null}
+{"res": 49531000000}
+{"res": 45045000000}
+{"res": 86399000000}
